[dev] partial cuda graph support for thd format training.

[HaochenYuan]

What does this PR do ?

CUDA Graph requires static tensor shapes; THD packed-sequence inputs are variable-length.
The bridge: pad to upper bounds, decompose PackedSeqParams (dataclass → tensor kwargs)
across the graph boundary, vectorize RoPE (no GPU→CPU sync), and route a padding_mask
through the model so MoE aux loss ignores padding.

This feature collaborate with this TransformerEngine PR 2898.

Files

File	Change
packed_seq_params.py	New pad_thd_for_cuda_graph(): pads tokens/labels/cu_seqlens, generates padding_mask, handles CP partitioning.
transformer/transformer_layer.py	New _decompose_/_reconstruct_packed_seq_params_*kwargs (dataclass ↔ tensor kwargs across graph boundary); THD-aware get_layer_static_inputs.
transformer/module.py	_is_thd_cuda_graph() gate; THD static-input shape [max_T // (TP if SP else 1), 1, H]. Optional THD_DEBUG_CG_IO=1 shape logging.
transformer/transformer_config.py	New: thd_cuda_graph_max_num_seqs=32, cuda_graph_dynamic_microbatches=False, cuda_graph_num_microbatch_slots=None.
transformer/cuda_graphs.py	Dynamic microbatch slot inference: derive minimum safe slot count from PP/VPP order via topology probe + PP-group all-reduce.
models/common/embeddings/rope_utils.py	Replaced .tolist() + Python loop with torch.arange + torch.searchsorted per-token position lookup. Pure CUDA, graph-safe.
models/gpt/gpt_model.py	Non-pre_process PP stages: scatter padding_mask to SP-sharded shape.
transformer/moe/moe_layer.py	Same SP-aware padding_mask alignment for MoE aux-loss masking.
tokenizers/text/libraries/null_tokenizer.py	Adds pad_id property (returns eod_id) for THD padding.
pretrain_gpt.py	get_batch calls pad_thd_for_cuda_graph when max_seqlen_per_dp_cp_rank is set; forward_step forwards padding_mask. Backward-compatible: None → original path.

Data flow

get_batch → pad_thd_for_cuda_graph → (tokens, labels, loss_mask, position_ids,
                                       packed_seq_params, padding_mask)
forward_step → model.forward(padding_mask=...)
  gpt_model         : scatter padding_mask for SP (non-pre_process PP)
  transformer_layer : decompose PSP → graph capture/replay → reconstruct PSP
  moe_layer         : scatter padding_mask for aux-loss masking

Validation

tests/unit_tests/transformer/test_thd_cuda_graph.py::TestE2EBitwise — both
parametrized cases pass with --deterministic-mode:

• Moonlight-16B (MLA + MoE, TP2/CP2/PP2/EP4): noGraph vs cudaGraph bitwise identical (151s)
• Qwen3-8B (GQA, TP2/CP2/PP2): noGraph vs cudaGraph bitwise identical (100s)

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

[copy-pr-bot]

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.

[buptzyb]

What if the actual num seq is larger than this value?

[HaochenYuan]

No, it's not handled — if actual > max, we silently return the longer cu_seqlens, which will break the graph at replay with a shape mismatch. I think it's better to add an explicit assertion in _pad_cu_seqlens to fail fast with a clear error message, plus update the docstring/help text with a sizing formula.

[buptzyb]

I can think of two solutions.

1. fail the should_call_te_cudagraph condition so that it can go to the eager path for this microbatch.
2. update the data generation logic to prevent generating samples with num seq > thd_cuda_graph_max_num_seqs. If so, thd_cuda_graph_max_num_seqs can be renamed to thd_max_num_seqs (not only manages cudagraph).

[HaochenYuan]

I would prefer the solution 2, do you have any comment? @xiaoyao0115

[xiaoyao0115]

Maybe it'd be better to refactor the newly added functions in megatron/core/packed_seq_params.py

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/27631855121