[RFC] Allow Trainer to work with distributed Learner implementations

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

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 #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

Start by reading the existing Trainer and Learner implementations, then inspect PR #3926 and the referenced PyTorch DDP and FSDP2 contracts. Define the smallest synchronous Trainer(learner=...) path and validate it with one ordinary TorchRL loss. Done requires real multi-rank DDP/FSDP2 updates, weight gathering, and restore coverage on CUDA hardware.