Enable Ulysses CP for Qwen VLM training

[eligotts]

Summary

• Enables cp_style="ulysses" for the supported custom Qwen3.5 VLM path while keeping ring + VLM blocked.
• Moves multimodal image embedding/merge and 3D MRoPE position construction before CP sharding, then shards merged inputs_embeds and 3D position IDs for decoder forward.
• Adds 3D-aware CP position sharding, full packed seq_lens publication for Ulysses attention metadata, and explicit Ulysses head divisibility errors.
• Relaxes multimodal packing under CP only for models that explicitly advertise Ulysses VLM CP support.
• Adds focused tests for CP sharding, packing gates, trainer wrapper forwarding, and 2-rank Ulysses VLM equivalence on a tiny packed multimodal batch.

Validation

• uv run --no-sync ruff check ... and uv run --no-sync pytest tests/unit/... (cp, vlm, model_forward, qwen3_5_moe_vlm).

---

Multi-GPU validation follow-up (NCCL, 2× H200)

Validated the hybrid DeltaNet/FLA CP path and the end-to-end VLM CP path on real multi-GPU hardware. Three distinct issues were found; two are fixed in this PR, one is a scoped follow-up.

✅ Fixed: DeltaNet conv1d needs a cross-rank halo exchange (commit: DeltaNet CP)

The Qwen3.5 GatedDeltaNet causal conv1d ran on each rank's local shard only. Under CP, ranks past the start of a sequence zero-padded their first (kernel-1) tokens instead of seeing the previous rank's tail — corrupting the recurrent state across the whole rank, so hybrid DeltaNet CP forward diverged ~88% from the unsharded reference (full-attention CP was exact). fla's cross-rank state passing was correct; the gap was the missing conv boundary exchange.

Fix: halo-exchange the last (kernel-1) pre-conv tokens across CP ranks before the local conv1d, using the differentiable torch.distributed.nn.all_gather (so gradients route back to the source rank), with every rank keeping the gathered tensor in its autograd graph so the backward collective fires in lockstep (else ranks needing no halo deadlock the rest).

Validated (2× H200): forward bit-exact vs unsharded at realistic shard sizes; DeltaNet weight + input grads within bf16 noise (<1%, vs ~15% before the differentiable gather); a 12-step CP run tracks the single-GPU reference and converges. Known residual (out of scope): part_len ≤ 1 fla chunk (≤64 tok/rank) still diverges — a narrow fla edge that can't occur at real seq_len÷CP; recommend asserting tokens_per_rank > chunk_size.

✅ Fixed: VLM CP — vision/embedding vs the global attention patch and FSDP (commit: frozen vision + embeddings)

End-to-end VLM CP surfaced two more issues, both rooted in the fact that the merge runs the vision encoder + text embedding before sharding (full/replicated, outside the FSDP root forward):

1. Vision attention. substitute_ulysses_attn patches the global flash_attention_2 dispatch, whose Ulysses wrapper is causal-only. The vision encoder's bidirectional attention then hit AssertionError: ulysses CP only supports causal attention. Fix: pin the vision encoder to sdpa under CP so it dispatches around the patched entry. (The vision encoder is unsharded — it must not do the Ulysses all-to-all at all.)
2. FSDP root ordering. The pre-shard merge calls the vision encoder and embed_tokens standalone, before the root forward; FSDP2 requires the first forward to go through the root, so a separately-sharded module here lazy-inits as its own root and collides (FSDP state has already been lazily initialized). Fix: keep these pre-shard-merge modules out of FSDP (full/replicated) — but only when frozen (ignored params get no FSDP gradient reduction). Frozen → ignore; trainable → raise a clear error pointing at LoRA / freeze_vision_encoder / the structural fix.

Net: CP + VLM works out of the box for frozen-embedding recipes (e.g. LoRA) and fails loudly (rather than silently mis-training) for full fine-tuning.

Validated (2× H200, cp=2 ulysses, random-init Qwen3.5-35B-A3B, fake data): full-FT (trainable embed) → raises the clear error; LoRA (embed+vision frozen) → guard takes the ignore branch and FSDP setup succeeds through to the training loop; a frozen-embed run → 6 steps end-to-end with images, finite loss, no crash.

⚠️ Follow-up (not in this PR): full-parameter VLM-CP + full-pipeline validation

• Full fine-tuning (trainable token embeddings) needs the structural fix: run the pre-shard merge through the root model forward (or embed after sharding) so embed_tokens stays FSDP-managed. Today it raises the clear error above instead.
• Full rl pipeline (orchestrator + vLLM inference + real multimodal rollouts via the mm_refs wire path) at real scale / seq len / cp>2 / multi-node is still unvalidated — the runs above exercise the trainer with mocked tensor batches, not the orchestration/inference half.

Reproduction harnesses (operator/model forward + gradient equivalence, isolated fla CP, full-model CP forward, standalone-trainer multimodal smoke) available on request.

---

Note

High Risk
Touches distributed training (FSDP2 ignored params, Ulysses attention, CP collectives), multimodal merge ordering, and hybrid linear-attention CP correctness; misconfiguration can deadlock NCCL or silently skip gradient reduction on trainable embeddings.

Overview
Adds Ulysses context parallelism (CP) for supported Qwen3.5 VLM training, replacing the previous hard block on VLM + CP.

The RL trainer now runs a pre-shard multimodal merge (prepare_vlm_inputs_for_context_parallel): vision encoding, image scatter into embeddings, and 3D MRoPE positions on the full sequence, then shards inputs_embeds and 3D position_ids for the decoder. forward() accepts premerged inputs_embeds and seq_lens_are_global for packed Ulysses varlen metadata.

FSDP + CP+VLM: frozen vision encoder and embed_tokens are kept out of FSDP (ignored_params) so pre-root merge calls do not lazy-init a second FSDP root; trainable modules raise a clear error. Vision attention is pinned to sdpa under CP so bidirectional vision does not hit the global Ulysses flash patch.

Hybrid DeltaNet CP: GatedDeltaNet adds a differentiable cross-rank conv1d halo exchange and can use global cu_seqlens from packed seq_lens.

CP utilities gain 3D position sharding, setup_cp_attention_params, Ulysses head-divisibility checks, and multimodal packing is allowed under CP only when the model advertises Ulysses VLM CP support.

^{Reviewed by Cursor Bugbot for commit 7d83615. 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 7d83615. Configure here.}

[cursor]

VLM CP setup ignores missing config

High Severity

cp_vlm FSDP exclusions, frozen-embed checks, and vision sdpa pinning run only when config.vlm is set, but RL CP multimodal training enters the pre-shard prepare_vlm_inputs_for_context_parallel path whenever mm_kwargs is present. Runs with a VLM checkpoint and multimodal batches but no [trainer.model.vlm] block skip those guards and can hit Ulysses-incompatible vision flash attention or FSDP lazy-init failures on standalone vision/embed_tokens calls.

Additional Locations (2)

• src/prime_rl/trainer/rl/train.py#L474-L492
• src/prime_rl/trainer/model.py#L464-L466

 

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

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7d8361586c

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Keep EP prefetch off unsharded CP VLM embeddings

When CP+VLM takes this new branch, embed_module is deliberately left as a plain, ignored embedding instead of being converted with fully_shard. In EP runs, the setup below still executes embed_module.set_modules_to_forward_prefetch([transformer_blocks[0]]), which only exists on FSDP-sharded modules, so Qwen VLM CP with expert parallelism crashes during setup_fsdp before training starts. Gate the later prefetch setup the same way, or only call it when the embedding was actually FSDP-sharded.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Use a CUDA-capable process group for Ulysses test

This worker immediately moves the model and tensors to CUDA, and the patched Ulysses path uses distributed all-to-all/all-gather collectives on those CUDA tensors. Initializing the process group with gloo makes the new GPU equivalence test fail in environments that actually run it, before it can check the logits. Use a CUDA-capable backend such as NCCL with appropriate per-rank device assignment, or skip when that backend is unavailable.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Derive CP VLM setup from the loaded model

When a renderer-backed VLM run leaves [model.vlm] unset, train.py can still receive mm_kwargs and enter the new CP multimodal branch, but this flag stays false so the pre-shard merge modules are not excluded from FSDP (and the vision attention pinning in get_model is skipped by the same config.vlm gate). That sends the CP path back through sharded vision/embedding modules before the FSDP root forward and crashes instead of enabling Ulysses VLM CP. Base this on the loaded VLM architecture/support check, or reject CP multimodal when the VLM setup block is absent.

Useful? React with 👍 / 👎.