feat(rl): multimodal context parallelism

[hubert-marek]

Summary

Part 7/7 — the only new code in the stack (in neither #2485 nor #2889; both deferred RL MM CP). Based on #2947 (feat/rl-multimodal-packing) — needs MM packing from #2947, the VLM CP machinery from #2946 (which is trainer-agnostic, in the modeling files), and the seq_lens contract from #2944. Rebased onto the current #2947 after its packing review fixes and the finalized #2946 CP cleanup.

Mirrors the SFT VLM CP path in the RL trainer, replacing the blanket NotImplementedError:

• MRoPE multimodal batches under ulysses defer sharding to the model: input_ids/position_ids stay global so the vision encoder and image-embed merge see the full sequence; the model root forward shards embeds/positions after merge (and routed_experts with them). Text batches keep the existing pre-shard path.
• seq_lens passes with seq_lens_are_pre_shard=True under CP — micro-batch boundaries span the pre-shard sequence, and pad_to_multiple_of=cp already guarantees shard divisibility (padding is its own seq_lens segment since feat(trainer): seq_lens as the packed-sample boundary contract (SFT + RL) #2944).
• The LoRA token-count CP adjustment derives chunk_size from the sharded length even when input_ids stays global.
• Gates relaxed: VLM + CP now requires cp_style='ulysses' and a custom model implementation instead of blanket rejection (impl="auto" passes the validator; get_model rejects it at load time if it resolves to hf).
• Docs: guide on adding a custom VLM implementation (VLM training is custom-only as of feat(sft): multimodal (VLM) SFT #2945; this documents the full contract — registration, seq_lens consumption, MRoPE, dummy-vision FSDP symmetry, packing capability flag, CP hooks, KL validation).

Proposed follow-up (not in this stack): the deferral condition and sharding choreography here deliberately mirror sft/train.py; extracting the shared logic into one helper would keep the two trainer loops from drifting.

E2E

• RL MM cp=2 ulysses run (multimodal_color_codeword shape) — W&B link TBD

• Integration suite on 8×H100 at this tip: 27 passed, 6 failed, 1 error — every failure also fails (or is family-flaky) on unmodified main on the same node: 4× test_benchmark_regression (tolerances calibrated for the CI A6000 runners), test_reverse_text_moe (FileNotFoundError, env), test_reverse_text_multi_run (flaky family — main fails a sibling test). The one stack-correlated flip, test_reverse_text_sft::test_loss_goes_down, is a marginal assertion: 5 steps at batch 4, last-vs-first loss delta +0.0018, and this branch's loss/grad values match the feat(sft): Enable VLM SFT  #2485-head reference bit-for-bit (5.5987/73.2068). Suggest lengthening that CI run or comparing first-k/last-k means — flagged for review.

🤖 Generated with Claude Code

---

Note

Medium Risk
Touches distributed RL forward/CP sharding for multimodal batches; incorrect boundary or shard handling could skew logprobs/gradients, but scope is aligned with existing SFT behavior and gated to ulysses + custom models.

Overview
RL multimodal training can use context parallelism when cp_style='ulysses', instead of being rejected outright.

The RL trainer mirrors the SFT VLM CP path: for MRoPE batches (multimodal kwargs with image_grid_thw), it skips pre-forward sharding of input_ids/position_ids so vision merge and image embeds see the full sequence; the custom model shards after merge. Text-only CP batches still use the existing setup_cp_params path. seq_lens stays on the wire under CP with seq_lens_are_pre_shard=True so packed document boundaries refer to the pre-shard sequence. LoRA CP token counting uses sharded labels length for chunk_size when inputs stay global.

Config validation replaces blanket “VLM + CP forbidden” with VLM + CP requires cp_style='ulysses' (ring CP with VLM still errors).

Docs add a development guide for custom VLM implementations (registration, packed seq_lens, dummy vision for FSDP, optional CP hooks).

Tests swap the old rejection test for ring-vs-ulysses coverage and assert ulysses + VLM + cp=2 validates.

^{Reviewed by Cursor Bugbot for commit bd09c63. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 7c38ffa. Configure here.}

[hubert-marek]

Folded into #2947 now that the RL multimodal CP delta is small and directly coupled to multimodal packing.