[LLM] WS4 Post-Training Observability: Implement PostTrainingLogger by coder-jayp · Pull Request #3950 · pytorch/rl

coder-jayp · 2026-07-07T19:00:27Z

Description

This PR implements the Observability workstream (WS4) from RFC #3948. It introduces a modular PostTrainingLogger that unifies metric emission across TorchRL's post-training algorithms, making the observability stack consumable by external training loops.

Key Changes:

PostTrainingLogger: New backend-agnostic logger for tracking loss components, buffer utilization, and inference staleness. Uses getattr duck-typing to safely ingest both GRPOLossOutput and SFTLossOutput.
Recipe Refactoring: Updated GRPO and Expert Iteration (EI) recipes to delegate to the new logger.
Deprecation Handling: EI legacy metric keys are temporarily emitted in parallel with explicit DeprecationWarnings (targeted for removal in v0.15.0 per CLAUDE.md).
Docs & Tests: Added full Sphinx documentation and comprehensive unit testing in test_loggers.py.

Motivation and Context

Solves the need for component-level interoperability in TorchRL's LLM stack. Users can now drop this logger into their own existing loops without migrating their entire workflow.

close #3948

I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

New feature (non-breaking change which adds core functionality)
Documentation (update in the documentation)
Example (update in the folder of examples)

Checklist

I have read the CONTRIBUTION guide (required)
My change requires a change to the documentation.
I have updated the tests accordingly (required for a bug fix or a new feature).
I have updated the documentation accordingly.

pytorch-bot · 2026-07-07T19:00:30Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3950

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 4a2d797 with merge base d7659c7 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

github-actions · 2026-07-07T19:00:42Z

⚠️ PR Title Label Error

PR title must start with a label prefix in brackets (e.g., [BugFix]).

Current title: WS4 Post-Training Observability: Implement PostTrainingLogger

Supported Prefixes (case-sensitive)

Your PR title must start with exactly one of these prefixes:

Prefix	Label Applied	Example
`[Algorithm]`	new algo	`[Algorithm] Add new RL objective`
`[BE]`	BE	`[BE] Improve error messages`
`[Benchmark]` or `[Benchmarks]`	Benchmarks	`[Benchmark] Add collector benchmark`
`[BugFix]`	BugFix	`[BugFix] Fix memory leak in collector`
`[Example]` or `[Examples]`	Examples	`[Example] Add training script`
`[Feature]`	Feature	`[Feature] Add new optimizer`
`[Doc]` or `[Docs]`	Documentation	`[Doc] Update installation guide`
`[Refactor]`	Refactoring	`[Refactor] Clean up module imports`
`[CI]`	CI	`[CI] Fix workflow permissions`
`[Test]` or `[Tests]`	Tests	`[Tests] Add unit tests for buffer`
`[Trainer]` or `[Trainers]`	Trainers	`[Trainer] Add trainer config`
`[Environment]` or `[Environments]`	Environments	`[Environments] Add Gymnasium support`
`[Data]`	Data	`[Data] Fix replay buffer sampling`
`[LLM]`	llm/	`[LLM] Add reward model integration`
`[Minor]`	small change	`[Minor] Fix typo in error message`
`[Performance]` or `[Perf]`	Performance	`[Performance] Optimize tensor ops`
`[BC-Breaking]`	bc breaking	`[BC-Breaking] Remove deprecated API`
`[Deprecation]`	Deprecation	`[Deprecation] Mark old function`
`[Algorithm]` or `[Algorithms]`	new algo	`[Algorithm] Add new objective`
`[Quality]`	Quality	`[Quality] Fix typos and add codespell`
`[Versioning]`	versioning	`[Versioning] Bump release version`
`[WIP]`	WIP	`[WIP] Draft implementation`

Note: Common variations like singular/plural are supported (e.g., [Doc] or [Docs]).

coder-jayp · 2026-07-16T14:29:08Z

@vmoens, I would appreciate your feedback on this when you have a moment. Thanks.

coder-jayp · 2026-07-27T14:25:33Z

Hi @theap06, any feedback on this PR?

theap06

Sorry for the late reply! Left some comments.

theap06 · 2026-07-28T03:58:48Z

+
+        with torch.no_grad():
+            try:
+                reward_list = batch.get(("next", "reward"), default=None, as_list=True)


This reads from batch, not replay_buffer[:]. Old code for these three metric families read the whole buffer. Should rename the keys or read from replay_buffer

theap06 · 2026-07-28T03:59:31Z

+                    # _storage is private; max_size is a public plain attribute on Storage.
+                    storage = replay_buffer._storage  # noqa: SLF001
+                    if hasattr(storage, "max_size") and storage.max_size > 0:
+                        metrics["buffer/utilization"] = (


RayReplayBuffer never sets self._storage so it raises an AttributeError. also, the buffer/utilization never logs for grpo since both use the rayreplaybuffer

theap06 · 2026-07-28T04:00:32Z

+            )
+        except Exception:  # noqa: BLE001
+            pass
+        legacy_metrics["reward from batch"] = float(batch["next", "reward"].mean())


Only the rb_content block above this is guarded; the loss.loss_sft reads here aren't and nothing wraps the log_training_metrics either.

theap06 · 2026-07-28T04:01:14Z

@@ -919,91 +919,55 @@ def log_training_metrics(
 ):


A few keys got renamed without any docstrings or deprecation warnings. ei_utils.py keeps everything alive via _LEGACY_KEY_MAP, worth doing the same here

theap06 · 2026-07-28T04:02:05Z

+    with torch.no_grad():
+        optim_steps = global_step // gradient_accumulation_steps
+        legacy_metrics: dict[str, Any] = {}
+        try:


All 11 warnings are made every call. Instead, worth a module-level flag so it's not 11 warnings. Warn calls every logged step.

coder-jayp · 2026-07-28T16:15:08Z

@theap06 Thanks for the detailed review. All points addressed.

…ce, GRPO deprecations, guard EI reads

…se-batch support Follow-ups to review, on top of the rebase onto main (which drops the objectives migration already merged as pytorch#3888): - Loss fields are discovered by iterating the TensorClass loss output instead of a hardcoded 11-name allowlist, so a new loss term is logged automatically. Plain objects still work via their attribute dict. - Metric-computation failures warn once per site at WARNING level instead of five blanket except-Exception blocks logging at debug: an observability component that silently emits nothing is worse than one that raises. - training/grad_norm is omitted on accumulation steps rather than logged as a literal 0.0, which polluted any aggregate over the series. - batch/reward_std is omitted for single-element batches (std would be NaN). - as_list=True returns a plain tensor for dense (padded) batches; all three read sites now handle both the ragged and the dense form. The dense form previously made every batch/* metric silently disappear. - Deprecations follow CLAUDE.md 12: the legacy GRPO keys are emitted again during the window (previously they were warned about but not emitted, and three keys plus two renames had no deprecation at all), the category is FutureWarning (DeprecationWarning is hidden by default), and the removal target is v0.16.0 (two minors after v0.14; version.txt is 0.13). - Runnable Examples block on the class docstring, verified by execution. test/test_loggers.py: 50 passed. test/objectives: 7491 passed post-rebase. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jul 7, 2026

github-actions Bot added Documentation Improvements or additions to documentation Objectives llm/ LLM-related PR, triggers LLM CI tests sota-implementations/ Record labels Jul 7, 2026

coder-jayp changed the title ~~WS4 Post-Training Observability: Implement PostTrainingLogger~~ [LLM] WS4 Post-Training Observability: Implement PostTrainingLogger Jul 7, 2026

theap06 suggested changes Jul 28, 2026

View reviewed changes

vmoens force-pushed the feat/post-training-logger branch 2 times, most recently from d8539b2 to 7bdfa15 Compare August 5, 2026 15:12

coder-jayp and others added 4 commits August 8, 2026 21:00

Implement PostTrainingLogger for WS4 observability (RFC pytorch#3948)

b4a33e7

Address reviewer feedback: batch/* namespace, _storage guard, warn-on…

0f0a396

…ce, GRPO deprecations, guard EI reads

[BugFix] Correct post-training logger metrics

4a2d797

vmoens force-pushed the feat/post-training-logger branch from 7bdfa15 to 4a2d797 Compare August 8, 2026 20:04

Navigation Menu

Uh oh!

[LLM] WS4 Post-Training Observability: Implement PostTrainingLogger - #3950

[LLM] WS4 Post-Training Observability: Implement PostTrainingLogger#3950
coder-jayp wants to merge 4 commits into
pytorch:mainfrom
coder-jayp:feat/post-training-logger

coder-jayp commented Jul 7, 2026 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Jul 7, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jul 7, 2026

Uh oh!

coder-jayp commented Jul 16, 2026

Uh oh!

coder-jayp commented Jul 27, 2026

Uh oh!

theap06 left a comment

Uh oh!

theap06 Jul 28, 2026

Uh oh!

theap06 Jul 28, 2026

Uh oh!

theap06 Jul 28, 2026

Uh oh!

theap06 Jul 28, 2026

Uh oh!

theap06 Jul 28, 2026

Uh oh!

coder-jayp commented Jul 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

coder-jayp commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Motivation and Context

Types of changes

Checklist

Uh oh!

pytorch-bot Bot commented Jul 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3950

✅ No Failures

Uh oh!

github-actions Bot commented Jul 7, 2026

⚠️ PR Title Label Error

Supported Prefixes (case-sensitive)

Uh oh!

coder-jayp commented Jul 16, 2026

Uh oh!

coder-jayp commented Jul 27, 2026

Uh oh!

theap06 left a comment

Choose a reason for hiding this comment

Uh oh!

theap06 Jul 28, 2026

Choose a reason for hiding this comment

Uh oh!

theap06 Jul 28, 2026

Choose a reason for hiding this comment

Uh oh!

theap06 Jul 28, 2026

Choose a reason for hiding this comment

Uh oh!

theap06 Jul 28, 2026

Choose a reason for hiding this comment

Uh oh!

theap06 Jul 28, 2026

Choose a reason for hiding this comment

Uh oh!

coder-jayp commented Jul 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

coder-jayp commented Jul 7, 2026 •

edited

Loading

pytorch-bot Bot commented Jul 7, 2026 •

edited

Loading