[rl] Support context parallelism in the RL trainer

[wwwjn]

Summary

Context parallelism (CP) in the RL trainer.

Under CP, preprocess_inputs shards the labels along the packed token dimension, but the per-token loss arguments that the RL losses consume (advantages, generator logprobs, loss mask, temperature) reached the loss unsharded. CP > 1 therefore could not train.

• shard_per_token_tensors (in torchtitan/distributed/context_parallel.py) keeps a CP rank's contiguous token shard of each per-token tensor. This is the same layout shard_tensors produces without a load-balancing permutation.
• TrainingEngine shards the loss arguments with it when CP is enabled. It raises if a CP load balancer is configured, because the permutation would reorder the labels but not the loss arguments. Ulysses CP, the CP mode for varlen attention, already requires context_parallel_load_balancer=None.

Why the losses need no change

The GRPO and DAPO losses are masked per-token sums divided by a global loss-token count. Each CP rank sums over its own tokens. FSDP shards on (dp_shard, cp) with gradient division disabled, so summing gradients over CP reproduces the full-sequence gradient. Loss metrics are summed over the existing dp x cp loss mesh.

Using it

A recipe sets context_parallel_degree on the trainer and converts the trainer's attention to a CP kernel. Generators keep plain attention. The new CI config shows the pattern:

config.trainer.parallelism.context_parallel_degree = 2
config.trainer = apply_transforms(
    config.trainer,
    [ContextParallelTransform(
        inner_attention_map={VarlenInnerAttention: UlyssesCPVarlenInnerAttention})],
)

Tests

• New CPU tests (tests/unit_tests/cpu/distributed/test_context_parallel_per_token.py, 7 tests), all passing. They check:

  • shards match the CP token ranges and reassemble the full tensors;
  • masked sums over shards equal the unsharded sum;
  • non-tensor values pass through;
  • mismatched or indivisible token counts and a configured load balancer are rejected;
  • the engine shards with its CP mesh coordinates.

• Existing CPU suites: distributed, RL trainer, losses, batcher and trainer loop pass (189). Two FSDP mixed-precision tests fail with and without this change: the local torch build predates MixedPrecisionPolicy(param_dtype_override_fn=...).

• New 8-GPU integration test rl_grpo_cp2_fsdp2_gen_tp2: Qwen3-0.6B, trainer FSDP=2 x Ulysses varlen CP=2, two generators at TP=2. The config builds with Ulysses varlen attention in all 28 trainer layers and plain varlen attention in the generators.

• GPU numerics, CP=2 vs no CP. One RL GRPO forward/backward plus optimizer step on Qwen3-0.6B (random init), on 4 GB300 GPUs with a 2026-10-12 nightly. Both layouts use the same rank-local initial weights and two fixed 4,096-token microbatches: packed documents, padding, and per-token advantages, logprobs and masks. Deterministic mode is on.

  Metric	FSDP=2 x TP=2	Repeat, relative diff	CP=2 x TP=2, relative diff
  loss	0.3046408296	0	2.0e-7
  grad norm	2.045546532	0	2.1e-4
  entropy	11.72645378	0	8.1e-8
  ratio clipped fraction, ratio mean, logprob diff		0	0

  The no-CP repeat is bitwise identical. The CP differences are bf16 reduction-order noise; per-token loss arguments misaligned with the labels would move the loss by O(1). Peak active memory per GPU: 4.47 GiB without CP, 3.92 GiB with CP.

• GPU end to end. The integration config's CP path on one 4-GPU host (trainer FSDP=1 x Ulysses varlen CP=2, one generator at TP=2), and the same with FSDP=2 and no CP. Both run all 5 steps and exit cleanly.

  Step	CP=2 loss	CP=2 reward	CP=2 grad norm	No-CP loss	No-CP reward	No-CP grad norm
  1	-0.0039	0.67	2.41	2.4e-6	0.50	0.065
  2	0.00028	0.22	2.20	-0.00012	0.29	0.22
  3	0.0017	0.27	1.32	-0.0033	0.34	0.79
  4	-0.0020	0.20	0.27	-0.0028	0.17	0.39
  5	-0.0010	0.36	1.06	0.0036	0.37	0.77

  The two runs sample different rollouts, so their rows are not expected to match. At this 1,024-token context, CP=2 trains about half as many tokens per second as FSDP=2 on the same two GPUs. Both runs set VLLM_ALLREDUCE_USE_FLASHINFER=0 for the generator, because FlashInfer's all-reduce JIT build failed in that environment.

Relation to #3980

#3980 (pipeline parallelism in the RL trainer) adds a guard that rejects CP until per-token loss arguments are sharded. This PR does that sharding on the non-PP path. Whichever lands second should drop that guard for the non-PP case. CP combined with PP still needs the sharded loss arguments packed into the PP targets, which is out of scope here.

[tianyu-l]

any reason we don't

• remove the function to_loss_kwargs entirely

  torchtitan/torchtitan/components/data/types.py

  Line 55 in c9c8626

  def to_loss_kwargs(

• put loss kwargs in as_input_dict

Another option is to override as_input_dict method of

torchtitan/torchtitan/rl/trainer/batcher.py

Line 56 in c9c8626

def loss_kwargs(self) -> dict[str, torch.Tensor]:

directly

[wwwjn]

Sorry I didn't realize these are better places.

I think option2 sounds more natural to me because current loss_kwargs are only used for RL.
Then we'd have to override as_input_dict() in decoder in following format:

 def as_input_dict(self) -> dict[str, Any]:
        # Explicit base call: zero-argument super() fails in slots=True dataclasses.
        return {
            **TokenizedTrainingMicrobatch.as_input_dict(self),
            "loss_kwargs": {
                "generator_logprobs": self.generator_logprobs,
                "temperature": self.temperature,
                "loss_mask": self.loss_mask,
                "advantages": self.advantages,
            },
        }

We'd have to use nest dict because otherwise preprocess_inputs will need to know each key / hardcode each key and shard them across CP. What does sounds like?

[tianyu-l]

sure