feat: context-parallel with nemotron v3 by adil-a · Pull Request #1441 · NVIDIA-NeMo/Automodel

adil-a · 2026-03-03T21:22:16Z

Summary

Adds context parallelism (CP) support for NemotronV3 hybrid models (interleaved attention + Mamba2 + MoE). NemotronV3 requires a dual-path CP strategy because attention layers and Mamba/SSM layers have fundamentally different parallelism requirements:

Attention layers: Use TE's sequence-parallel CP (p2p or allgather), splitting the sequence across ranks
Mamba/SSM layers: Use a hidden-parallel strategy (all-to-all), where each rank processes the full sequence but only a subset of heads

Supported configurations

Format	Backend	Sequence Packing	CP Support
THD	TE	Yes	✅ DualChunkSwap via TE p2p
THD	TE	No	✅ DualChunkSwap via TE p2p
BSHD	SDPA	No	✅ via PyTorch `context_parallel()`

Key changes

Mamba hidden-parallel CP (mamba_cp.py — new file):

MambaContextParallel class orchestrating CP for Mamba2 mixer layers
Autograd-aware all-to-all (_all_to_all_cp2hp / _all_to_all_hp2cp) to redistribute between sequence-sharded and hidden-sharded layouts
DualChunkSwap undo/redo (_undo_attention_load_balancing / _redo_attention_load_balancing) so SSM kernels see tokens in sequential order
THD packed sequence support via _deinterleave_packed_seqs / _reinterleave_packed_seqs for correct token ordering after all-to-all with multiple packed sequences
Parameter slicing (conv1d, dt_bias, A_log, D) returning views so gradients flow to full parameters

NemotronV3 model changes (layers.py, model.py):

Attention layers wired to TE CP via DotProductAttention.set_context_parallel_group()
Mamba layers wired via MambaContextParallel attached as mixer.cp
seq_idx computation from cu_seqlens for correct SSM state resets at sequence boundaries
THD format support in attention (cu_seqlens, max_seqlen, qkv_format passthrough)

Bug fixes:

Guard padding_mask for pad_token_id=None (NemotronV3 tokenizer returns None) in thd_utils.py, cp_utils.py, model_utils.py — tensor == None returns Python False (a scalar bool), not a tensor
Fix cu_seqlens global/local mismatch in Mamba CP — batch cu_seqlens are GLOBAL (pre-TE-partitioning), so derive local as cu_seqlens // cp_size instead of erroneously computing cu_seqlens * cp_size
Fix seq_idx in mixer layer — remove erroneous * cp_size since cu_seqlens is already global
Fix padding_idx=None crash in pack_dataset when packing with CP > 1 (getattr(tokenizer, "pad_token_id", 0) returns None when the attribute exists but is None)

Test infrastructure:

Consolidated run_nemotron_v3_attention_cp.py into run_attention_cp.py as --model_type nemotron_v3 with --configs arg for multi-format testing (bshd_te, thd_te, bshd_sdpa)
Fixed pre-existing test bugs in Qwen3MoE and DeepSeekV3 CP tests: rope_theta access via get_rope_config(), missing cp_size/cp_rank params, missing param grad all-reduce, DeepSeek freqs_cis format
New run_hybrid_nemotron_v3_cp.py — end-to-end hybrid model test (4 configs)
New run_mamba_cp.py — Mamba2Mixer layer test (4 configs including THD packed)
Added L2_Context_Parallel to CI workflow matrix (cicd-main.yml)
Fixed Qwen3.5MoE shell script (was incorrectly using pytest on a non-pytest script)

THD format and sequence packing with CP

In THD format, the batch is flattened to a 1D [T, H] tensor with cu_seqlens tracking sequence boundaries. For CP with DualChunkSwap:

TE partitions tokens across CP ranks using thd_get_partitioned_indices
After all-to-all in Mamba CP, tokens are in rank-major order — _deinterleave_packed_seqs rearranges to sequence-major so each sequence's tokens are contiguous
_undo_attention_load_balancing permutes from DualChunkSwap order [c0, c3, c1, c2] back to sequential [c0, c1, c2, c3] for the SSM kernel
After SSM, the reverse transforms (_redo + _reinterleave) restore the layout for attention

Note: With CP > 1 and sequence packing, each sequence is padded to be divisible by 2 * cp_size before packing. This changes pack composition compared to CP=1 (fewer sequences per pack due to padding overhead), resulting in different num_label_tokens per step — this is expected behavior, not a correctness issue.

Convergence results

https://api.wandb.ai/links/nvidia-nemo-fw-public/f7qpoo9i

Signed-off-by: adil-a <adil.asif2000@hotmail.com>

copy-pr-bot · 2026-03-03T21:22:20Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Resolve conflicts in cp_utils.py and test_context_parallel.py by keeping additions from both sides (DualChunkSwap CP functions from this branch, attach_context_parallel_hooks and qwen3.5 tests from main). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: adil-a <adil.asif2000@hotmail.com>

Support both THD (TE, with/without packing) and BSHD (TE or SDPA, no packing) context parallelism for hybrid Mamba-Attention models. Changes: - Add 2D THD format detection in NemotronV3Attention (Qwen3MoE pattern) and NemotronV3Mamba2Mixer (unsqueeze/squeeze) - Add uses_te_cp flag to MambaContextParallel to conditionally apply DualChunkSwap reordering (skip for DTensor/SDPA CP) - Add deinterleave/reinterleave for packed multi-sequence all-to-all - Pass cu_seqlens through to MambaCP for THD-mode reordering - Update supports_cp to allow hybrid models with SDPA backend - Add 4-config functional test grid: BSHD+TE, THD+TE, THD+TE+packing, BSHD+SDPA across Mamba, Attention, and Hybrid test scripts Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: adil-a <adil.asif2000@hotmail.com>

The MoE parallelizer's apply_cp didn't handle mamba blocks, causing mixer.cp to stay None when EP is active. Add layer_type/self_attn properties to NemotronV3Block for compatibility with the MoE parallelizer, and add a "mamba" branch in apply_cp that creates MambaContextParallel. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: adil-a <adil.asif2000@hotmail.com>

Signed-off-by: adil-a <adil.asif2000@hotmail.com>

Bug fixes: - Guard padding_mask for pad_token_id=None (NemotronV3 tokenizer) in thd_utils.py, cp_utils.py, and model_utils.py - Fix cu_seqlens global/local mismatch in Mamba CP (was doubling already-global values), remove debug instrumentation - Fix seq_idx computation in NemotronV3 mixer (remove erroneous * cp_size) - Fix padding_idx=None crash in packed_sequence packing (train_ft.py) Test improvements: - Merge run_nemotron_v3_attention_cp.py into run_attention_cp.py as --model_type nemotron_v3 with --configs arg for multi-format testing - Fix pre-existing test bugs: rope_theta access (use get_rope_config), missing cp_size/cp_rank in attention forward, missing param grad all-reduce, DummyBlock missing mlp attr, DeepSeek freqs_cis format - Add hybrid NemotronV3 CP test shell script - Fix Qwen3.5MoE shell script (was using pytest on non-pytest file) - Add L2_Context_Parallel to CI workflow matrix Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

adil-a · 2026-03-29T07:16:48Z

/ok to test 192556a

Signed-off-by: adil-a <adil.asif2000@hotmail.com>

adil-a · 2026-03-29T07:28:56Z

/ok to test 6708890

hemildesai · 2026-03-30T03:34:14Z

/claude review

…d validation - Fix _repeat_group_state to handle 2D THD input by normalizing to 3D before expand() and squeezing back (triggered when n_groups < cp_size) - Fix hybrid SDPA test to validate actual gradients (q_proj.weight) instead of comparing zeros vs zeros (embed_tokens is detached) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

adil-a · 2026-03-30T18:07:31Z

/claude review

adil-a · 2026-03-30T20:14:19Z

/ok to test c147997

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: adil-a <adil.asif2000@hotmail.com>

adil-a · 2026-03-30T21:03:03Z

/ok to test 028580a

… model for TE detection - Add _ensure_pad_token_id() that defaults pad_token_id to 0 with a warning when the tokenizer doesn't define one (e.g. NemotronV3). Called in both NeMoAutoTokenizerWithBosEosEnforced and the registry dispatch path. - Revert scattered pad_token_id workarounds in cp_utils, thd_utils, and train_ft back to their original form since the tokenizer now guarantees a non-None pad_token_id. - Change _uses_te_dot_product_attention() to inspect actual model modules for TE DotProductAttention instead of only checking the YAML config string, with config fallback for pre-model call sites. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: adil-a <adil.asif2000@hotmail.com>

adil-a · 2026-04-03T04:08:27Z

/ok to test 8cf1a95

Resolve conflict in import_utils.py (take main's version with three-way transformers API version detection). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: adil-a <adil.asif2000@hotmail.com>

adil-a · 2026-04-03T06:20:40Z

/ok to test d240310

adil-a added 2 commits March 3, 2026 15:26

mamba hidden-parallel

cc60b27

Signed-off-by: adil-a <adil.asif2000@hotmail.com>

mamba un-swizzle and cp attention

df8e442

Signed-off-by: adil-a <adil.asif2000@hotmail.com>

adil-a self-assigned this Mar 3, 2026

adil-a linked an issue Mar 3, 2026 that may be closed by this pull request

[Feature Request] Context Parallelism Support for NemotronV3 #1409

Closed

adil-a mentioned this pull request Mar 23, 2026

Add CP and seq packing for nemotron-v3 #1598

Closed

adil-a and others added 4 commits March 26, 2026 19:51

bshd + torch sdpa fixes

dcb7ce8

Signed-off-by: adil-a <adil.asif2000@hotmail.com>

adil-a linked an issue Mar 29, 2026 that may be closed by this pull request

Add CP and seq packing for nemotron-v3 #1598

Closed

adil-a and others added 2 commits March 29, 2026 00:14

Merge branch 'main' into adil/nano-v3-cp

192556a

copy-pr-bot Bot temporarily deployed to nemo-ci March 29, 2026 07:17 Inactive

copy-pr-bot Bot had a problem deploying to test March 29, 2026 07:17 Error

adil-a marked this pull request as ready for review March 29, 2026 07:22

adil-a requested review from a team, HuiyingLi, ZhiyuLi-Nvidia, akoumpa, hemildesai and pthombre as code owners March 29, 2026 07:22

transformers version fix

6708890

Signed-off-by: adil-a <adil.asif2000@hotmail.com>

copy-pr-bot Bot temporarily deployed to nemo-ci March 29, 2026 17:45 Inactive

claude Bot reviewed Mar 30, 2026

View reviewed changes

Comment thread nemo_automodel/components/distributed/mamba_cp.py Outdated

claude Bot reviewed Mar 30, 2026

View reviewed changes

Comment thread tests/functional_tests/context_parallel/run_hybrid_nemotron_v3_cp.py Outdated

claude Bot reviewed Mar 30, 2026

View reviewed changes

Comment thread tests/functional_tests/context_parallel/run_attention_cp.py

claude Bot reviewed Mar 30, 2026

View reviewed changes

Comment thread tests/functional_tests/context_parallel/run_attention_cp.py

copy-pr-bot Bot temporarily deployed to test March 30, 2026 20:14 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci March 30, 2026 20:14 Inactive

Merge branch 'main' into adil/nano-v3-cp

028580a

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: adil-a <adil.asif2000@hotmail.com>

hemildesai reviewed Mar 31, 2026

View reviewed changes

Comment thread nemo_automodel/components/distributed/thd_utils.py Outdated

Comment thread nemo_automodel/recipes/llm/train_ft.py Outdated

Comment thread nemo_automodel/components/distributed/cp_utils.py Outdated

Merge main into adil/nano-v3-cp

d240310

Resolve conflict in import_utils.py (take main's version with three-way transformers API version detection). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: adil-a <adil.asif2000@hotmail.com>

hemildesai approved these changes Apr 3, 2026

View reviewed changes

thomasdhc approved these changes Apr 3, 2026

View reviewed changes

Conversation

adil-a commented Mar 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Supported configurations

Key changes

THD format and sequence packing with CP

Convergence results

Uh oh!

copy-pr-bot Bot commented Mar 3, 2026

Uh oh!

adil-a commented Mar 29, 2026

Uh oh!

adil-a commented Mar 29, 2026

Uh oh!

hemildesai commented Mar 30, 2026

Uh oh!

Uh oh!

Uh oh!

adil-a commented Mar 30, 2026

Uh oh!

Uh oh!

Uh oh!

adil-a commented Mar 30, 2026

Uh oh!

adil-a commented Mar 30, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

adil-a commented Apr 3, 2026

Uh oh!

adil-a commented Apr 3, 2026

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

adil-a commented Mar 3, 2026 •

edited

Loading