[RFC] Allow Trainer to work with distributed Learner implementations

[gtnv]

Background

The minimum PyTorch rules are:

• Run one process per GPU.
• Select the local device before distributed initialization.
• Give every rank the same model structure and update count.
• Give every rank its own local batch. DDP does not partition data.
• Enter collectives in the same order on every rank.
• Construct an FSDP2 optimizer after fully_shard.

See the PyTorch contracts for DDP and FSDP2.

Proposed Refactoring

what needs to change in torchrl in order to make way for distributed primitives?

A good way of doing this would be to add a learner arg to Trainer. When present, trainer can keep its existing epoch and minibatch loop but will call learner.update(sub_batch) at optimization.

how to ensure functionality AND PERFORMANCE of the various layers is kept after adding the bare minimum functions for distributed RL?

every rank must run the same trainer program. hooks should be local so that any algorithm behavior is not routed.

DDP and FSDP2 use native pytorch communication.

Proposed Additions

Distributed PyTorch functions and their compatibility with the refactor

Addition	Behavior
Trainer(learner=...)	Calls learner.update inside the existing optimization loop. Reads the loss, stepper, and publishable model from the learner.
DDPLearner	Reuses the current loss-forward wrapper with DistributedDataParallel. Accepts an existing process group and does not launch processes.
FSDP2Learner	Comes from PR #3926. The caller applies fully_shard before constructing the optimizer.
Rank policy	Exposes rank, world size, device, and whether this rank may perform external side effects. It does not own the cluster.
Multi-rank coverage	Runs real DDP and FSDP2 updates, weight gathering, and restore behavior on CUDA hardware.

Mechanically, the launcher initializes torch.distributed. Each rank constructs the same model, loss, stepper, learner, and Trainer. Trainer selects a minibatch and calls learner.update. DDP or FSDP2 synchronizes gradients during backward.

Weight publication is collective-aware. DDP can publish from rank zero because every rank has a full replica. FSDP2Learner.get_weights must be called by every rank, but only rank zero sends the gathered weights through WeightSyncScheme.

The initial implementation supports fixed-size synchronous updates. It should certify one ordinary TorchRL loss before claiming that functionalized losses, multiple optimizers, or variable-token objectives work.

Future Features and Necessary additions

• Multi-rank FSDP2 validation. Run [Feature] Add Learner primitive (LocalLearner, FSDP2Learner) #3926 on real multi-GPU hardware and fix collective, clipping, gathering, and restore failures.

• Distributed loss normalization. Add global sum/count reduction before backward for variable samples, masks, and token counts.

• On-policy sampling. Partition PPO batches without replacement and preserve epoch, trajectory, and advantage-normalization semantics.

• Multiple optimizers. Define learner checkpoint and step semantics for TD3, SAC temperature updates, and other multi-optimizer algorithms.

• Distributed checkpointing. Coordinate collective model state with Trainer, replay, collector, RNG, and version state.

• Remote learners. Design one remote backend only after local DDP and FSDP2 use the same Learner.update contract successfully.

• Asynchronous RL. Add policy versions, stale-sample admission, backpressure, and recovery of queued and in-flight work.

• Model parallelism. Integrate TorchTitan-style tensor or pipeline parallel recipes without moving their partition rules into TorchRL.

• Examples. Add one DDP example, one FSDP2 example, and one end-to-end rollout-to-update example before larger scale claims.

cc @vmoens @theap06

[Kushagra481]

Hey, @gtnv can I work on this