[Feature] Add Learner primitive (LocalLearner, FSDP2Learner) by theap06 · Pull Request #3926 · pytorch/rl

theap06 · 2026-07-02T08:25:45Z

Summary

Introduces torchrl.trainers.Learner: a backend-agnostic entry point for taking
one optimization step on a tensordict batch with a given LossModule.
LocalLearner is the single-process reference implementation; FSDP2Learner
shards the same model with fully_shard and reuses the training step unchanged.

It plays the same role for training that Collector plays for data collection
and LLMWrapperBase plays for generation/scoring: a fixed, TensorDict-native
contract with interchangeable backends, so algorithm code does not need to know
whether the update runs on one device, under sharded training, or on a remote
training process.

Design

Learner.update() is concrete, in the base class, and touches only
self.model / self.optimizer / self.clip_grad_norm /
self.grad_accum_steps: zero_grad -> forward -> sum the loss module's
"loss"-prefixed output keys -> backward -> optional grad-norm clip ->
optimizer step. This is what lets FSDP2Learner reuse the exact same step as
LocalLearner -- sharded training only changes model construction and weight
gathering, not the step itself.
get_weights() is the one place sharding is not transparent: it must return
plain tensors (for WeightSyncScheme.send, which already accepts a
TensorDictBase), even when the learner's parameters are sharded.
FSDP2Learner does not decide sharding granularity or device mesh -- it
accepts a model the caller has already wrapped with fully_shard, exactly as
bare FSDP2 usage works. Keeping that decision in caller code avoids
FSDP2Learner becoming a second place those choices are made.

Two contracts worth reading before using this

The optimizer defines what is trained, not model. model is the
weight-sync source and the gradient-sync handle; the parameters that get clipped
and stepped are the ones in optimizer.param_groups. Many TorchRL losses hold
their trainable parameters on the loss module as TensorDictParams, and the
losses that expand their networks (SACLoss, REDQLoss, ...) hold copies of
the modules you passed in:

loss_module = SACLoss(actor, qvalue)
learner = LocalLearner(actor, Adam(loss_module.parameters()))  # correct
learner = LocalLearner(actor, Adam(actor.parameters()))        # critics never train

The second line is silent -- the critics are differentiated on every step and
never updated -- so update() now checks before its first optimizer step and
raises if any parameter received a gradient that no param group covers. Grad-norm
clipping likewise covers optimizer.param_groups, not model.parameters().

A Learner owns exactly one optimizer, so algorithms that deliberately use
several (per-network learning rates, a separate entropy-temperature optimizer)
are not expressible as a single Learner today; use one per optimizer.

Checkpointing is checkpoint() / load_checkpoint(), not state_dict(). A
bare Optimizer is not an nn.Module, so nn.Module.state_dict() structurally
cannot carry its state and a resume would reset Adam's moments. Overriding
state_dict to return it anyway would break the nn.Module contract -- a parent
module calls child.state_dict(destination=...) and discards the return value,
so nesting a Learner inside any other module would silently drop all of its
state. Separate names keep both contracts intact.

Gradients are not checkpointed, so a checkpoint taken mid-accumulation-window
resets the accumulation counter to 0 with a warning rather than resuming at a
non-zero step with empty gradients (which would step after fewer micro-batches
than grad_accum_steps and under-scale that update).

Notes

update() raises on a non-scalar summed loss, which is what a loss built with
reduction="none" produces -- previously this surfaced as a torch-internal
"grad can be implicitly created only for scalar outputs" error.
The "loss" prefix cannot be tightened to "loss_": "loss" with no
underscore is a real out_key (DQNLoss, GAILLoss, OnlineDTLoss).
With FSDP2Learner.get_weights()'s default cpu_offload=True, the full
weights are returned on rank 0 only and every other rank gets an empty
tensordict. That is deliberate (gathering the full model onto every rank does
not scale), but it means an unconditional
scheme.send(learner.get_weights()) sends nothing from non-zero ranks. Pass
cpu_offload=False to gather everywhere.
With gradient accumulation the output key set differs between accumulation
calls and step calls (grad_norm only appears on the latter).
No *Config companion yet. These classes have no Trainer wiring in this PR,
so a Hydra config would have nothing to instantiate against; the configs land
with the trainer integration (follow-up 2 below).
These are new, unreleased classes, so the checkpoint() rename carries no
deprecation shim.

Planned follow-ups (not in this PR)

Multi-rank FSDP2Learner verification on real multi-GPU hardware -- the one
thing I could not test here. The single-rank gloo/CPU tests exercise the
fully_shard/DTensor code paths but not actual cross-rank communication.
Refactor an existing recipe's hand-rolled training loop (e.g. the
reward-model recipe, once [Feature] Add RewardModelLoss objective for RLHF reward-model training #3922 lands) onto LocalLearner, as the first real
consumer, with the *Config companions.
A RemoteLearner design writeup scoping one external backend (TorchTitan is
the most PyTorch-native of the candidates and the most plausible first
integration target) before any implementation.

Tested

test/test_trainer.py: 100 passed, 19 skipped (26 of those are the
Learner/FSDP2Learner tests). Single-rank gloo/CPU only; multi-GPU is
follow-up 1.

c75e16b fixes a checkpoint bug that only surfaced once the suite actually ran:
get_state_dict(full_state_dict=True, cpu_offload=True) does not guarantee
fresh tensors — cpu_offload only copies when the shards are off-CPU, so on a
CPU or single-rank mesh the "gathered" state aliases the live shards and training
on after a checkpoint silently mutates it. FSDP2Learner.checkpoint now clones,
like the base implementation.

pytorch-bot · 2026-07-02T08:25:49Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit 132f0bc with merge base 99daa1c ():

NEW FAILURES - The following jobs have failed:

Unit-tests on Linux / tests-olddeps (3.10, 11.8, remainder) / linux-job (gh)
RuntimeError: Command docker exec -t 0c97b6dcd037c4025728d693447d9b04372b9fe099675b4cdd452b7432097837 /exec failed with exit code 2
Unit-tests on Linux / tests-stable-gpu (3.12, 13.0, 3) / linux-job (gh)
AssertionError: Tensor-likes are not close!

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

vmoens · 2026-08-10T08:01:43Z

Reviewed the rebased head 3614a44, with particular attention to the latest optimizer-coverage/checkpoint contract changes and the final FSDP2 checkpoint-cloning fix. The base and FSDP2 checkpoint paths now clone tensor state before returning it, accumulation restores safely reset partial windows, and the optimizer-coverage check runs before the first step. I ran the Learner/FSDP2-focused test selection locally: 27 passed, including the FSDP2 checkpoint round trip, accumulation/reference comparison, and plain-tensor weight gather. I did not find a blocking issue in the current code. The meaningful residual risk is unchanged from the PR description: these checks use world_size=1 gloo/CPU, so actual multi-rank collective behavior and rank-0-only checkpoint/weight gathering still need real distributed coverage.

vmoens · 2026-08-28T07:15:04Z

Coordinated this branch with #4140 and force-pushed the rebased stack.

The shared contract is now:

Learner owns execution placement and weight materialization.
OptimizationStepper is the single owner of loss reduction, backward, gradient accumulation, clipping, mixed precision, optimizer stepping, and optimizer checkpoint state.
LocalLearner and FSDP2Learner receive the model, loss module, and stepper at construction, then expose update(batch). This is the same stepper contract used by Trainer and GRPOTrainer.

The reconciliation also moved optimizer-coverage and scalar-loss validation into MixedPrecisionOptimizationStepper, added the FSDP gradient-sync seam there, and made Learner checkpoints compose stepper state instead of maintaining a second accumulation/optimizer implementation.

I replaced the prior Learner test expansion with behavior-focused parity, weight snapshot, checkpoint, optimizer coverage, and real FSDP2/DTensor tests. Relative to #4140, test/test_trainer.py now adds 233 lines instead of 526.

Local validation: the complete test/test_trainer.py suite passes (154 passed, 19 dependency-based skips), including two FSDP2 tests. The full changed-file pre-commit suite passes. I also added ci/olddeps because this feature touches the recent FSDP2 API.

Introduces torchrl.trainers.Learner: a backend-agnostic entry point for taking one optimization step on a tensordict batch with a given LossModule. LocalLearner is the single-process reference implementation. Mirrors the role Collector plays for data collection and LLMWrapperBase plays for generation/scoring, so training placement (local / sharded / remote) becomes a swappable backend behind a fixed contract instead of a hand-rolled loop per recipe. get_weights() returns a TensorDictBase, matching what WeightSyncScheme.send() already accepts, so this composes with the existing weight-sync path unchanged.

Refactors Learner.update() to live in the base class as concrete, generic step logic (zero_grad -> forward -> sum loss_* keys -> backward -> clip -> step), operating only on self.model/self.optimizer/self.clip_grad_norm/ self.grad_accum_steps. This is what lets FSDP2Learner reuse the exact same training step as LocalLearner: sharded training only changes how the model is constructed (wrapped with fully_shard by the caller) and how get_weights() gathers the result. FSDP2Learner.get_weights() gathers every DTensor leaf via full_tensor() into a plain tensor, so its output is consumable by WeightSyncScheme exactly like LocalLearner's, with no changes on the receiving side. Verified on a single-rank (world_size=1) gloo process group, which exercises the real fully_shard()/DTensor code path without a cluster: forward/backward/clip_grad_norm_/optimizer.step() dispatch correctly through DTensor, TensorDict.from_module()/.apply() handle DTensor leaves transparently, and FSDP2Learner produces bit-exact losses and gathered weights vs LocalLearner given the same seed/data/lr.

…d-sync Fixes three real gaps in FSDP2Learner identified after the initial PR: 1. get_weights() previously gathered every DTensor leaf to EVERY rank via full_tensor(), replicating the whole model in every rank's memory for no benefit -- does not scale to large sharded models. Now uses torch.distributed.checkpoint.state_dict.get_model_state_dict with StateDictOptions(full_state_dict=True, cpu_offload=True), which gathers to rank 0 only (other ranks get an empty tensordict) by documented DCP semantics. Verified: correct nested key shape via unflatten_keys('.'), matches the prior full_tensor()-based output. 2. No sharded-checkpoint path existed. Learner (base) gains real state_dict()/load_state_dict() covering model + optimizer state (a bare Optimizer is not an nn.Module, so plain nn.Module.state_dict() silently drops it -- a real, previously-latent bug for LocalLearner too). Both overrides clone their tensors: nn.Module.state_dict() and Optimizer.state_dict() return views onto live tensors, not copies, so without cloning, further training after checkpointing would silently corrupt the saved checkpoint (caught by round-trip tests). FSDP2Learner overrides these two methods again with get_state_dict/set_state_dict (DCP-aware, handles DTensor optimizer state, rank0-only via cpu_offload). Verified end to end: model weights AND Adam/SGD-momentum optimizer state round-trip correctly through save/perturb/load. 3. grad_accum_steps>1 was untested on FSDP2Learner, and update() did no FSDP2-specific optimization during accumulation. Learner.update() now toggles model.set_requires_gradient_sync(...) when the model exposes it (FSDP2-wrapped models do; LocalLearner's plain model doesn't, so this is a no-op there), deferring the cross-rank gradient reduction until the last micro-batch of an accumulation window instead of reducing on every micro-batch. Verified against a non-sharded reference: the accumulated, synced gradient after 2 microbatches matches a plain model accumulating the same 2 microbatches exactly. Still unverified (unchanged from the original PR): all of the above is tested at world_size=1 (single-rank gloo), which exercises the real fully_shard()/DTensor/DCP code paths but not actual cross-rank communication or memory behavior.

…eckpoint contract Addresses review feedback on the Learner primitive. - The optimizer, not `model`, defines what is trained. Grad-norm clipping now covers `optimizer.param_groups` rather than `self.model.parameters()`, and `update()` validates before its first optimizer step that no parameter received a gradient the optimizer does not cover. The silent failure this prevents: `LocalLearner(actor, Adam(actor.parameters()))` with a loss module that owns or expands its critics leaves those critics differentiated on every step and never updated. Documented on the class, in `LocalLearner`'s `model` argument, and in the docs page. - Checkpointing moves from `state_dict`/`load_state_dict` to `checkpoint`/`load_checkpoint`. Overriding the `nn.Module` methods broke their contract: `destination`/`prefix`/`keep_vars` were ignored and the return value discarded, so nesting a `Learner` inside any parent module silently dropped all of its state. `state_dict` is now plain `nn.Module` behavior again. - `load_checkpoint` resets the accumulation counter to 0 with a warning instead of resuming mid-window: gradients are not checkpointed, so resuming at a non-zero step would step the optimizer after fewer micro-batches than `grad_accum_steps` and under-scale that update. - `update()` raises a clear error on a non-scalar summed loss (what a loss built with `reduction="none"` returns) rather than letting `.backward()` fail with a torch-internal message. - `LearnerCapabilities` is frozen, so the shared class-level default on `Learner.capabilities` cannot be mutated into every other instance. - `torch.distributed.checkpoint.state_dict` is imported lazily through a cached accessor rather than at module top, keeping `import torchrl` free of `torch.distributed.checkpoint`. - Corrects the `FSDP2Learner.get_weights` docs: it uses `get_model_state_dict` with `StateDictOptions`, not per-leaf `full_tensor()`, and the default `cpu_offload=True` returns weights on rank 0 only -- which an unconditional `scheme.send(learner.get_weights())` on every rank would silently no-op. - Records why the `"loss"` key prefix cannot be tightened to `"loss_"`: `"loss"` with no underscore is a real out_key of DQNLoss, GAILLoss and OnlineDTLoss. Tests cover each of the above. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Running the suite locally shows the FSDP2 checkpoint round-trip fails: after load_checkpoint the weights are not restored. get_state_dict(full_state_dict=True, cpu_offload=True) does not guarantee fresh tensors -- cpu_offload only copies when the shards are off-CPU, so on a CPU (or single-rank) mesh the "gathered" state aliases the live shards. Training on after checkpointing then mutates the checkpoint, which is exactly what Learner.checkpoint guards against with _clone_tensors. The previous docstring claimed the opposite; it now says why the clone is required. test/test_trainer.py: 100 passed, 19 skipped. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

With the canonical Adam(loss_module.parameters()) construction, loss- owned parameters (SACLoss log_alpha, a Lagrange multiplier, ...) are trained by update() but were absent from checkpoints: model.state_dict never sees them and Optimizer.state_dict stores moments, not values, so a resume silently kept whatever value was in memory. Learner.checkpoint now saves their values under extra_params (in param- group order) and load_checkpoint restores them positionally, raising on a param-group mismatch. FSDP2Learner had a second failure mode: DCP's get_state_dict maps optimizer state to model FQNs and raises a bare KeyError on non-model params, so checkpoint() could not even save. The extras are now hidden from the optimizer during the DCP walk and their values and optimizer state carried explicitly, with rank-0 broadcast on restore. get_weights stays model-only by design (inference workers do not need loss-owned parameters); this is now documented. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

vmoens · 2026-08-28T08:36:24Z

Rebased #3926 onto current main after #4140 merged. Because #4140 was squash-merged, I dropped its duplicated pre-squash commit sequence and replayed only the seven Learner commits. The new head is 132f0bc; GitHub reports the PR mergeable and its diff is limited to the Learner/OptimizationStepper files. Local validation: 154 trainer tests passed with 19 dependency-based skips, and the changed-file pre-commit suite passed.

gtnv · 2026-08-29T15:10:58Z

+    @contextlib.contextmanager
+    def _extras_hidden_from_optimizer(self):
+        """Hide non-model parameters from DCP and carry them explicitly."""
+        extras = self._extra_optimized_parameters()


_extra_optimized_parameters() selects optimizer parameters outside the FSDP-managed model, so under the documented setup their gradients are not synchronized and can diverge across ranks. please reject them or explicitly synchronize them before allowing updates.

gtnv · 2026-08-29T15:10:59Z

+            model_state, optimizer_state = dsd.get_state_dict(
+                self.model, self.optimizer, options=options
+            )
+        return {


this override omits loss_module state outside model, so loss-owned target tensors and registered buffers are not restored. please include that state or reject loss modules that own it.

gtnv · 2026-08-29T15:10:59Z

+            for parameter, state in zip(
+                extras, checkpoint.get("extra_optim_state", [{}] * len(extras))
+            ):
+                if state:


when the saved extra optimizer state is {}, this branch leaves the current state restored by _extras_hidden_from_optimizer() in place, so an early checkpoint can inherit later moments. please clear the current entry before applying the saved state.

gtnv · 2026-08-29T15:18:50Z

+        :class:`~torchrl.trainers.OptimizationStepper` contract used by
+        :class:`~torchrl.trainers.Trainer`.
+        """
+        return self.optimization_stepper._step(self, batch)


TD3OptimizationStepper.step() directly reads clip_grad_norm and clip_norm, but Learner defines neither, so the first learner.update() increments the TD3 counter and then raises AttributeError. please define the shared stepper-context contract or reject incompatible steppers at construction.

gtnv · 2026-08-29T16:23:09Z

+        dsd = _dist_state_dict()
+        options = dsd.StateDictOptions(full_state_dict=True, cpu_offload=True)
+        with self._extras_hidden_from_optimizer() as (extras, extra_state):
+            model_state, optimizer_state = dsd.get_state_dict(


get_state_dict() gathers the full FSDP state collectively, but full_state_dict=True with cpu_offload=True returns it only on rank 0. please document that every rank must call checkpoint(), while only rank 0 should persist the returned checkpoint.

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 2, 2026

github-actions Bot added Feature New feature Documentation Improvements or additions to documentation Trainers labels Jul 2, 2026

theap06 marked this pull request as draft August 2, 2026 07:24

vmoens force-pushed the learner-primitive-clean branch 2 times, most recently from 09e5cf6 to 3614a44 Compare August 10, 2026 07:55

theap06 marked this pull request as ready for review August 17, 2026 04:11

vmoens added the user-facing User-facing changes - only include in major releases label Aug 17, 2026

gtnv reviewed Aug 26, 2026

View reviewed changes

Comment thread torchrl/trainers/learners/local.py Outdated

theap06 requested a review from gtnv August 27, 2026 18:15

vmoens force-pushed the learner-primitive-clean branch from ffdedeb to 112ee05 Compare August 28, 2026 07:14

github-actions Bot added Collectors llm/ LLM-related PR, triggers LLM CI tests sota-implementations/ Integrations/torch_geometric Integrations labels Aug 28, 2026

vmoens added the ci/olddeps Run the tests-olddeps suite (oldest supported torch) on this PR label Aug 28, 2026

vmoens force-pushed the learner-primitive-clean branch from 112ee05 to 5293e16 Compare August 28, 2026 07:14

vmoens and others added 7 commits August 28, 2026 09:35

[Refactor] Share optimization steppers across Learner and Trainer

132f0bc

vmoens force-pushed the learner-primitive-clean branch from 5293e16 to 132f0bc Compare August 28, 2026 08:35

gtnv mentioned this pull request Aug 29, 2026

[RFC] Allow Trainer to work with distributed Learner implementations #4185

Open

gtnv reviewed Aug 29, 2026

View reviewed changes

Navigation Menu

Uh oh!

[Feature] Add Learner primitive (LocalLearner, FSDP2Learner) - #3926

[Feature] Add Learner primitive (LocalLearner, FSDP2Learner)#3926
theap06 wants to merge 7 commits into
pytorch:mainfrom
theap06:learner-primitive-clean

theap06 commented Jul 2, 2026 •

edited by vmoens

Loading

Uh oh!

pytorch-bot Bot commented Jul 2, 2026 •

edited

Loading

Uh oh!

vmoens commented Aug 10, 2026

Uh oh!

Uh oh!

vmoens commented Aug 28, 2026

Uh oh!

vmoens commented Aug 28, 2026

Uh oh!

gtnv Aug 29, 2026 •

edited

Loading

Uh oh!

gtnv Aug 29, 2026 •

edited

Loading

Uh oh!

gtnv Aug 29, 2026 •

edited

Loading

Uh oh!

gtnv Aug 29, 2026 •

edited

Loading

Uh oh!

gtnv Aug 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

theap06 commented Jul 2, 2026 • edited by vmoens Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Design

Two contracts worth reading before using this

Notes

Planned follow-ups (not in this PR)

Tested

Uh oh!

pytorch-bot Bot commented Jul 2, 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/3926

❌ 2 New Failures

Uh oh!

vmoens commented Aug 10, 2026

Uh oh!

Uh oh!

vmoens commented Aug 28, 2026

Uh oh!

vmoens commented Aug 28, 2026

Uh oh!

gtnv Aug 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gtnv Aug 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gtnv Aug 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gtnv Aug 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gtnv Aug 29, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

theap06 commented Jul 2, 2026 •

edited by vmoens

Loading

pytorch-bot Bot commented Jul 2, 2026 •

edited

Loading

gtnv Aug 29, 2026 •

edited

Loading

gtnv Aug 29, 2026 •

edited

Loading

gtnv Aug 29, 2026 •

edited

Loading

gtnv Aug 29, 2026 •

edited

Loading