feat(cp): context parallelism for hybrid + VLM SFT#2946
Conversation
hubert-marek
force-pushed
the
feat/sft-vlm
branch
from
9faf63e to
b57a384
Compare
hubert-marek
force-pushed
the
feat/sft-vlm
branch
from
b57a384 to
b21f4d5
Compare
hubert-marek
force-pushed
the
feat/sft-vlm
branch
from
b21f4d5 to
f3bffad
Compare
hubert-marek
force-pushed
the
feat/vlm-cp
branch
4 times, most recently
from
c3fc955 to
80ccc17
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 80ccc17245
ℹ️ 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".
|
|
||
| def set_context_parallel_attributes(self, cp_group: dist.ProcessGroup, cp_rank: int, cp_world_size: int) -> None: | ||
| for layer in self.layers: | ||
| if isinstance(layer, NemotronHMambaLayer): |
There was a problem hiding this comment.
With the default model.ac settings, apply_ac() checkpoint-wraps every transformer layer before setup_model_cp() runs. A CheckpointWrapper is not an instance of NemotronHMambaLayer, so this condition skips the wrapped Mamba layers and never calls layer.set_context_parallel_attributes(...); NemotronH cp_style='ulysses' runs then shard the sequence without enabling the Mamba all-to-all CP path, corrupting or breaking those experiments. Use an attribute-based check or unwrap the checkpoint wrapper before testing the layer type.
Useful? React with 👍 / 👎.
hubert-marek
force-pushed
the
feat/sft-vlm
branch
from
123a69c to
5e0124d
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b205746207
ℹ️ 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".
| self._cp_rank = cp_rank | ||
| self._cp_world_size = cp_world_size | ||
| for layer in self.layers: | ||
| if getattr(layer, "layer_type", None) == "linear_attention": |
There was a problem hiding this comment.
With the default model.ac, apply_ac() checkpoint-wraps transformer layers before SFT/RL call setup_model_cp(), and those wrappers do not expose layer_type or linear_attn. This check therefore skips every wrapped Qwen3.5 linear-attention layer (the MoE hook has the same pattern), leaving DeltaNet.cp_group unset; in the newly enabled Qwen3.5 VLM/Ulysses CP path that means the sequence is sharded without FLA state passing. Fresh evidence beyond the existing NemotronH thread is that this new Qwen3.5 hook repeats the same wrapper-sensitive check, so please unwrap the checkpoint wrapper or inspect the wrapped module.
Useful? React with 👍 / 👎.
hubert-marek
force-pushed
the
feat/sft-vlm
branch
from
5e0124d to
6db94b7
Compare
Unify model CP setup and make packed-boundary metadata CP-aware: - setup_model_cp replaces setup_hybrid_cp/setup_nemotron_h_cp: models owning linear-attention/Mamba layers expose set_context_parallel_attributes (dense + MoE Qwen3.5, NemotronH) and wire their own layers; hybrid models without the hook are rejected instead of silently misconfigured. - seq_lens stays global under CP: documents can straddle the shard cut, so input_ids/position_ids shard per rank while seq_lens keeps the full pre-shard boundaries, flagged via seq_lens_are_global through forward() / prime_forward_kwargs into both Qwen3.5 models. A cu_seqlens_are_global provenance flag rides along to fla's GatedDeltaNet instead of reading cu_seqlens back (which cost one CPU-GPU sync per linear-attention layer per microbatch); the CP conv path switches to fla's causal_conv1d with a CP context so conv state passes between ranks with document-boundary resets. - VLM CP defers sharding to the model: the vision encoder and the image-embed merge see the full sequence, then the root forward shards embeds/positions (shard_position_ids_for_cp handles 3D MRoPE) and updates ulysses attention params from global seq_lens. The SFT trainer skips pre-sharding for MRoPE batches; the vision encoder is pinned to SDPA under CP. - shard_for_cp validates divisibility (uneven shards deadlock ulysses all-to-all) and CatDataset's seq_len padding guarantees it; the LoRA token count follows the sharded loss_mask. - Qwen3.5 CP requires the custom implementation (the HF impl has no set_context_parallel_attributes); rejected in get_model. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Name the CP boundary flag after the pre-shard sequence it describes rather than the ambiguous global scope. Co-authored-by: Cursor <cursoragent@cursor.com>
Traverse wrapped model layers so Qwen3.5 DeltaNet and NemotronH Mamba receive their context-parallel topology after activation checkpointing. Co-authored-by: Cursor <cursoragent@cursor.com>
Use the resolved implementation for one model-agnostic CP guard instead of maintaining Qwen-specific config detection. Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d308a46d0f
ℹ️ 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".
| if config.cp > 1 and impl_to_use != "custom": | ||
| raise ValueError( | ||
| "Context parallelism requires model.impl='custom' " | ||
| "(or model.impl='auto' selecting a custom PrimeRL implementation)." | ||
| ) |
There was a problem hiding this comment.
When users run CP with model.impl='hf' on a non-Qwen3.5/HF-only flash-attention model, or explicitly choose the HF implementation for a supported text model, this unconditional guard raises before the trainer can install the existing HF CP patches (substitute_hf_flash_attn / substitute_hf_ulysses_attn). The config validator still permits HF CP for flash_attention_2, so this regresses an existing supported path; narrow the rejection to model families that actually lack the needed CP hook.
Useful? React with 👍 / 👎.
| inputs_embeds: Optional[torch.FloatTensor] = None, | ||
| routed_experts: Optional[torch.LongTensor] = None, | ||
| seq_lens: Optional[torch.LongTensor] = None, | ||
| seq_lens_are_pre_shard: bool = False, |
There was a problem hiding this comment.
For SFT with cp > 1, the trainer now passes global packed-document boundaries via this flag, but NemotronH never uses it: the CP Mamba path still calls mamba_cp_forward, whose scan uses seq_idx=None, so packed documents that straddle or share a CP pack leak Mamba state across document boundaries. Since CP SFT requires pack_function='cat', this corrupts normal packed NemotronH CP runs; either thread the global seq_lens into the CP Mamba scan/conv reset path or reject packed NemotronH CP until it is supported.
Useful? React with 👍 / 👎.
1 participant
Summary
Part 5/7 of the split of #2485. Based on #2945 (
feat/sft-vlm) — needs the VLM SFT path from #2945 and the seq_lens contract from #2944. Rebased onto the current #2945 after its VLM validation and cleanup fixes.Context parallelism for hybrid (linear-attention) and VLM models in SFT:
setup_model_cpreplacessetup_hybrid_cp/setup_nemotron_h_cp: models owning linear-attention/Mamba layers exposeset_context_parallel_attributes(dense + MoE Qwen3.5, NemotronH) and wire their own layers through activation-checkpoint wrappers; hybrid models without the hook are rejected instead of silently misconfigured.seq_lensstays global under CP: documents can straddle the shard cut, soinput_ids/position_idsshard per rank whileseq_lenskeeps the full pre-shard boundaries, flagged via theseq_lens_are_pre_shardtypedforward()parameter, which this PR adds to every custom model (extending feat(trainer): seq_lens as the packed-sample boundary contract (SFT + RL) #2944's universalseq_lenscontract). Acu_seqlens_are_pre_shardprovenance flag rides to fla's GatedDeltaNet instead of readingcu_seqlensback — the.item()read cost one CPU-GPU pipeline sync per linear-attention layer per microbatch (~288/step on 35B MoE cp2). The CP conv path switches to fla'scausal_conv1dwith a CP context so conv state passes between ranks with document-boundary resets.torch.compiler.disable: its boundary-state all-gather isn't dynamo-traceable. The old.item()read had been masking this — data-dependent.item()graph-breaks, so tracing always stopped just before the conv. Removing the sync let dynamo trace into fla and crash at compile time; the explicit wrapper reinstates the same graph break without the sync. Validated head-to-head at identical config (dense-text cp=2, compile on): old.item()tree 72.0k tok/s vs this branch 75.2k tok/s (~4% faster), losses matching to ~1e-3.shard_position_ids_for_cphandles 3D MRoPE) and updates ulysses attention params from globalseq_lens. The SFT trainer skips pre-sharding for MRoPE batches; the vision encoder is pinned to SDPA under CP.shard_for_cpvalidates divisibility (uneven shards deadlock ulysses all-to-all); CatDataset's pack padding from feat(trainer): seq_lens as the packed-sample boundary contract (SFT + RL) #2944 guarantees it. LoRA token count follows the shardedloss_mask. CP requires a custom implementation for every model (HF impl has no CP hook) — rejected inget_model.E2E (W&B project
pr2485)cp=2 ulysses SFT matrix, 10/10 steps each, validated in BOTH eager (
TORCHDYNAMO_DISABLE=1) and compiled (default config, exercising the graph-break fix) modes: dense-text rdu3ydne, moe-text 40nwjbkg, dense-MM p5p4e28y, moe-MM eqvyrc96; control on the pre-fix tree ioqkts7pIntegration coverage: full suite ran at feat(rl): multimodal context parallelism #2948's tip (see there); no CP-specific integration tests exist — CP coverage is the cp=2 e2e matrix above
🤖 Generated with Claude Code
Note
High Risk
Touches core training forwards, CP collectives, and compile boundaries for hybrid/VLM paths; mis-sharding or wrong
seq_lenssemantics would silently skew loss.Overview
Adds context parallelism (ulysses) for hybrid models (Qwen3.5 DeltaNet, NemotronH Mamba) and VLM SFT, where packed documents can span CP shard boundaries.
CP wiring replaces separate hybrid/Nemotron helpers with
setup_model_cp, which calls each model’sset_context_parallel_attributes. Hybrid models without that hook are rejected;cp > 1requiresmodel.impl='custom'. VLM SFT with CP must usecp_style='ulysses'(config validation).Sequence metadata: trainers pass
seq_lens_are_pre_shardsoseq_lens/cu_seqlensstay global while tokens shard per rank. Qwen3.5 MoE/dense linear layers use fla CP context for conv + gated delta rule (withtorch.compiler.disableon fla conv). SFT no longer dropsseq_lensunder CP; it shardstarget_ids/loss_maskand defers MRoPE VLM sharding to the VLM root (full vision merge, then shard embeds/positions; vision encoder forced to SDPA under CP).Utilities:
shard_for_cpchecks seq divisibility and supports 3D position sharding;setup_cp_attention_paramsis split out for VLM. HF Qwen3.5 varlen patch gains optional fla CP conv path.Reviewed by Cursor Bugbot for commit d308a46. Bugbot is set up for automated code reviews on this repo. Configure here.