Allow fine-grained offloading with MC impl of full-CG. by rapatel · Pull Request #4253 · NVIDIA/Megatron-LM

rapatel · 2026-04-10T16:28:37Z

What does this PR do ?

⚠️ 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"

When your PR is ready, click Ready for Review.
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 · 2026-04-10T16:28:41Z

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.

rapatel · 2026-04-10T16:33:24Z

@lhb8125

lhb8125 · 2026-04-13T02:10:32Z

/ok to test b8a7993

…ading.

yaox12 · 2026-04-17T07:29:24Z

/ok to test e0ddd84

yaox12 · 2026-04-17T07:32:32Z

/ok to test 4919450

svcnvidia-nemo-ci · 2026-04-17T09:56:47Z

🔄 Merge queue validation started!

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

svcnvidia-nemo-ci · 2026-04-17T11:00:10Z

🔄 Merge queue validation started!

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

Single repair commit for the main slice 5812485 -> dev merge. - ssm/kda: restore dev's `False` distribute_saved_activations argument in both KDA recompute checkpoints. The kda.py conflict resolution had taken main NVIDIA#7054's config read; dev's recompute contract (NVIDIA#8203, test_gdn_quantized_recompute) passes False, as dev's GDN does. - tests/models/test_hybrid_mhc: main NVIDIA#7555 replaced mtp_layer_pattern with mtp_layer_config_list; assert the new attributes instead. - tests/transformer/moe/test_route_metadata_cpu: the dispatcher test double accepts packed_seq_params, which main NVIDIA#8458 now passes. - tests/transformer/test_attention_output_gate: the KDA test double sets two_stage_gates=False, read by main NVIDIA#7054's _apply_gated_norm. - tests/data/test_get_batch: stub dev's prepare_packed_seq_params hook in main NVIDIA#8525's config-cache test; all main assertions are kept. - ssm/kda: keep dev's direct-gate op order in _apply_gated_norm (norm, then gate reshape). main NVIDIA#7054 reshaped the gate first, which changed the compiled rounding that dev NVIDIA#8006's test_kda_gated_norm_preserves_upstream_ rounding pins; main's two-stage branch is unchanged. - tests/data/test_pretrain_gpt_dataset_config: main NVIDIA#8391 replaced pretrain_gpt's build_tokenizer(args) with get_tokenizer(); the dev test now patches get_tokenizer. Its assertion is unchanged. - ssm/kda_layer_config: run validate_kda() before TransformerConfig's __post_init__ as well as after it. Dev NVIDIA#6556's shared GDN-family assertions ran first and preempted the KDA ValueErrors that main NVIDIA#7054's test_kda config tests expect. Dev's plain TransformerConfig checks are unchanged. - tests/transformer/test_thd_cuda_graph: the dev HybridModel test double sets hybrid_layer_config_list=None, which main NVIDIA#7555 reads in HybridModel.forward. Main NVIDIA#7555 made the same change to its own hybrid test doubles. The test assertions are unchanged. - functional goldens: refresh dev_dgx_h100 golden values for the four gpt3_mcore_te_tp2_pp2 mHC cases after main NVIDIA#7521 changed their numerics (local H100 runs on merge 44a8771). - tests/unit_tests/ssm/test_kda.py: keep dev NVIDIA#6556 kda_safe_gate/kda_lower_bound on TransformerConfig (CLI flags); main NVIDIA#7054 placement check now excludes only those two fields and asserts they stay on the stack config. - tests/unit_tests/ssm/test_kda.py: main NVIDIA#7054's recompute-acceptance test builds a non-hybrid experimental_attention_variant="kda" config. Dev NVIDIA#6556 registers kda as a GDN-family variant, and dev requires linear_attention_freq for every standalone GDN-family model, so the kda case now passes linear_attention_freq=[1]. Dev's validation and the test assertion are unchanged. - docker: pin the ARM64 MoK build to main NVIDIA#8210's 7a8cfb7f. The merged MoK backend passes shared_output_gate_weight to mok.functional.forward, and main's test_mok_numerics runs that path on GB200. Dev's a2ad2d0c pin does not accept the argument (GB200 CI TypeError in test_real_mok_ddp_numerics). 7a8cfb7f is a fast-forward of a2ad2d0c (2 commits) that already contains both hunks of dev's docker/patches/mok.patch, so the patch, its bind mount and the install_mok.sh patch argument are removed as on main, and the MoE README names the new snapshot. The AMD64 (H100) image still skips MoK; the rest of dev's Dockerfile.ci.dev is unchanged. - .github/workflows/cicd-main.yml: NVIDIA-NeMo/RL 4aaa48fa (NVIDIA#4253) renamed tests/functional/L1_Functional_Tests_Megatron_4.sh to L1_Functional_Tests_Megatron_Generation.sh, so the NeMo RL dispatch pattern matched no script and the downstream pre-flight failed with "No test scripts match the requested selection". Point the pattern at the renamed script. Dev and main carry the same stale pattern. Signed-off-by: svcnemo-autobot <svcnemo-autobot@nvidia.com>

Allow fine-grained offloading with MC impl of full-CG.

da991e7

rapatel requested review from a team as code owners April 10, 2026 16:28

lhb8125 reviewed Apr 13, 2026

View reviewed changes

Comment thread megatron/core/transformer/transformer_config.py Resolved

Merge branch 'dev' into ripatel/cg_offload_assertion_dev

b8a7993

svcnvidia-nemo-ci added this to the Core 0.16 milestone Apr 13, 2026

copy-pr-bot Bot temporarily deployed to test April 13, 2026 02:11 Inactive

Add checks for CUDA graph settings with fine-grained activation offlo…

ef7a2d9

…ading.

rapatel requested review from a team as code owners April 17, 2026 04:27

Merge branch 'dev' into ripatel/cg_offload_assertion_dev

e0ddd84

yaox12 removed request for a team April 17, 2026 07:21

yaox12 enabled auto-merge April 17, 2026 07:29

yaox12 approved these changes Apr 17, 2026

View reviewed changes

Update transformer_config.py

4919450

copy-pr-bot Bot temporarily deployed to test April 17, 2026 07:33 Inactive

yaox12 added this pull request to the merge queue Apr 17, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Apr 17, 2026

yaox12 added this pull request to the merge queue Apr 17, 2026

Merged via the queue into NVIDIA:dev with commit 6efb083 Apr 17, 2026
92 of 100 checks passed

svcnemo-autobot mentioned this pull request Nov 2, 2026

chore: AUT-2689 nightly sync main into dev (2026-10-30) #8813

Merged

Navigation Menu

Repository navigation

Allow fine-grained offloading with MC impl of full-CG. - #4253

Allow fine-grained offloading with MC impl of full-CG.#4253
yaox12 merged 5 commits into
NVIDIA:devfrom
rapatel:ripatel/cg_offload_assertion_dev

rapatel commented Apr 10, 2026

Uh oh!

copy-pr-bot Bot commented Apr 10, 2026

Uh oh!

rapatel commented Apr 10, 2026

Uh oh!

Uh oh!

lhb8125 commented Apr 13, 2026

Uh oh!

yaox12 commented Apr 17, 2026

Uh oh!

yaox12 commented Apr 17, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 17, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Apr 17, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

rapatel commented Apr 10, 2026

What does this PR do ?

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

copy-pr-bot Bot commented Apr 10, 2026

Uh oh!

rapatel commented Apr 10, 2026

Uh oh!

Uh oh!

lhb8125 commented Apr 13, 2026

Uh oh!

yaox12 commented Apr 17, 2026

Uh oh!

yaox12 commented Apr 17, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 17, 2026

Uh oh!

Uh oh!

svcnvidia-nemo-ci commented Apr 17, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants