[main] feat(moe): Support packed sequence for gated delta net (GDN) by yuzhongw-nvidia · Pull Request #2645 · NVIDIA/Megatron-LM

yuzhongw-nvidia · 2025-12-12T13:13:45Z

What does this PR do ?

Support packed sequence for gated delta net (GDN).

PR for dev: #2644 , #4230
Closes: #4043, #3798

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
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

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

(Step 1): Add PR label `Expert Review`

(Step 2): Collect the expert reviewers reviews

Attach the Expert Review label when your PR is ready for review.
GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

Add Final Review label
GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

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.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

copy-pr-bot · 2025-12-12T13:13:51Z

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.

…apping_main Conflict resolution + porting upstream deltas to relocated code: - gpt/fine_grained_callables.py: keep slim HEAD; PreProcess / PostProcess / TransformerLayerNode / _BackwardDWWrapper and build_mtp_layer_callables / build_layer_callables now live under common/. - common/utils.py: port PR NVIDIA#4511 (remove dead manual_release_grads code path) into TransformerLayerNode.backward_impl / backward_dw. - common/fine_grained_callables.py: port PR NVIDIA#2645 (packed sequence GDN) into build_mtp_layer_callables — unpack the new 5-tuple from _get_embeddings and forward packed_seq_params / padding_mask. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…VIDIA#2645) Signed-off-by: yuzhongw <yuzhongw@nvidia.com> Co-authored-by: kunlunl <kunlunl@nvidia.com> Co-authored-by: Xuesong Ye <xuesongyey@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: Daisy Gao <daisyg@nvidia.com>

MultiTokenPredictionLayer.forward calls self._checkpointed_forward( padding_mask=padding_mask, ...) (multi_token_prediction.py:1305), but _checkpointed_forward and its inner custom_forward never accepted padding_mask. With recompute_granularity == 'full' and self.training, this raised: TypeError: MultiTokenPredictionLayer._checkpointed_forward() got an unexpected keyword argument 'padding_mask' at multi_token_prediction.py:1301. The kwarg was introduced in #2645 on the call site; the _checkpointed_forward refactor in #4593 dropped padding_mask from the recompute path. Add padding_mask: * to _checkpointed_forward's signature * to custom_forward's signature so it flows into _proj_and_transformer_layer * positionally to te_checkpoint and tensor_parallel.checkpoint, matching the other tensor / None args (padding_mask is a rolled tensor, not a non-tensor closure-captured arg like attention_bias) * to the recompute_method == 'block' fallback that also calls _proj_and_transformer_layer directly Also remove the @pytest.mark.flaky_in_dev markers from test_forward_backward, test_fp8_support, and test_packed_sequences_with_full_recompute, which were added in #4931 to mask this exact failure. Closes #4933 Signed-off-by: oliver könig <okoenig@nvidia.com>

MultiTokenPredictionLayer.forward calls self._checkpointed_forward( padding_mask=padding_mask, ...) (multi_token_prediction.py:1305), but _checkpointed_forward and its inner custom_forward never accepted padding_mask. With recompute_granularity == 'full' and self.training, this raised: TypeError: MultiTokenPredictionLayer._checkpointed_forward() got an unexpected keyword argument 'padding_mask' at multi_token_prediction.py:1301. The kwarg was introduced in NVIDIA#2645 on the call site; the _checkpointed_forward refactor in NVIDIA#4593 dropped padding_mask from the recompute path. Add padding_mask: * to _checkpointed_forward's signature * to custom_forward's signature so it flows into _proj_and_transformer_layer * positionally to te_checkpoint and tensor_parallel.checkpoint, matching the other tensor / None args (padding_mask is a rolled tensor, not a non-tensor closure-captured arg like attention_bias) * to the recompute_method == 'block' fallback that also calls _proj_and_transformer_layer directly Also remove the @pytest.mark.flaky_in_dev markers from test_forward_backward, test_fp8_support, and test_packed_sequences_with_full_recompute, which were added in NVIDIA#4931 to mask this exact failure. Closes NVIDIA#4933 Signed-off-by: oliver könig <okoenig@nvidia.com>

…VIDIA#2645) Signed-off-by: yuzhongw <yuzhongw@nvidia.com> Co-authored-by: kunlunl <kunlunl@nvidia.com> Co-authored-by: Xuesong Ye <xuesongyey@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: Daisy Gao <daisyg@nvidia.com>

Megatron passes a PackedSeqParams object to every attention module even in BSHD mode (qkv_format="bshd"), but GDN.forward blanket-rejected any non-None packed_seq_params, crashing Qwen3.5/3.6 GDN training/inference in BSHD. Guard on qkv_format=="thd" specifically, matching upstream NVIDIA/Megatron-LM PR NVIDIA#2645. Genuine THD packing still raises (full support is a follow-up port). Fixes radixark/miles#1292

…ron-bridge to latest (#1762) This PR makes the following bumps megatron_core from `cefc2520158c7ceba3f9adbe4b547a6f7a118da1` (latest dev branch as of 6/8/26) to `71e418ea7d7b3a6c9a53238c543c3e0b43e11026` (latest main branch as of 6/8/26. megatron-bridge from `8382dc343b07b068a827ca20bae860633df3baa0` to `91a15142a4b4442a8d46ab539d1b923bd08570d0` (latest main 6/8) Megatron-Bridge has upstreamed code to the main branch that isn't on the dev branch that is needed to use Megatron-Bridge (NVIDIA-NeMo/Megatron-Bridge#3988) Since sequence packing with GDN is now supported on main, we can move back over to the latest commit on the main branch: NVIDIA/Megatron-LM#2645

…VIDIA#2645) Signed-off-by: yuzhongw <yuzhongw@nvidia.com> Co-authored-by: kunlunl <kunlunl@nvidia.com> Co-authored-by: Xuesong Ye <xuesongyey@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: Daisy Gao <daisyg@nvidia.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

…VIDIA#2645) Signed-off-by: yuzhongw <yuzhongw@nvidia.com> Co-authored-by: kunlunl <kunlunl@nvidia.com> Co-authored-by: Xuesong Ye <xuesongyey@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: Daisy Gao <daisyg@nvidia.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>

…VIDIA#2645) Signed-off-by: yuzhongw <yuzhongw@nvidia.com> Co-authored-by: kunlunl <kunlunl@nvidia.com> Co-authored-by: Xuesong Ye <xuesongyey@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: Daisy Gao <daisyg@nvidia.com>

…VIDIA#2645) Signed-off-by: yuzhongw <yuzhongw@nvidia.com> Co-authored-by: kunlunl <kunlunl@nvidia.com> Co-authored-by: Xuesong Ye <xuesongyey@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: Daisy Gao <daisyg@nvidia.com> Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>

… to core_r0.18.0) `MultiTokenPredictionLayer.forward` calls self._checkpointed_forward(..., padding_mask=padding_mask, ...) but `_checkpointed_forward` on this branch does not declare `padding_mask`, so every MTP run with `recompute_granularity == 'full'` aborts at the first forward: TypeError: MultiTokenPredictionLayer._checkpointed_forward() got an unexpected keyword argument 'padding_mask' Two landings that do not compose: NVIDIA#2645 added the call-site kwarg, NVIDIA#4593 refactored the method without it. Tracked as NVIDIA#4933, which is still open. `main` has since gained the parameter (as a side effect of NVIDIA#3457, "Roll input IDs for MTP labels"), but `core_r0.18.0` -- and the `core_v0.18.2` tag cut from it on 2026-07-20, currently the newest release -- still carry the bug. MTP with full activation recomputation is therefore unusable on every released megatron-core. Cherry-picking NVIDIA#3457 would drag an unrelated feature into a release branch, so this is the minimal equivalent: thread `padding_mask` through the six places `main` threads it, and nowhere else. * `_checkpointed_forward` signature, positioned after `attention_mask` as on `main` * `custom_forward`'s parameter list * the `_proj_and_transformer_layer` call inside `custom_forward` * both checkpoint entry points (`te_checkpoint`, `tensor_parallel.checkpoint`), positionally -- `padding_mask` is a tensor / `None`, so it belongs with the positional args rather than the closure, per this method's own docstring * the `recompute_method == 'block'` fallback path `_proj_and_transformer_layer` on this branch already accepts `padding_mask`, so no other change is needed -- the parameter was simply not plumbed through the recompute path. Verified identical to `main`: `padding_mask` now appears 8 times inside `_checkpointed_forward` on both. Assisted-by: Claude Opus 5 (Claude Code) Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: gaohongkui <gaohongkui@meituan.com>

…tron-core 0.18.x ## What does this PR do? Adds a compatibility shim so Multi-Token Prediction can be used together with `recompute_granularity=full` on every released megatron-core, and makes the existing signature-based branch in `patch_mtp_layer_checkpointed_forward` visible instead of silent. ### The bug megatron-core 0.18.x is internally inconsistent: `MultiTokenPredictionLayer.forward` calls self._checkpointed_forward(..., padding_mask=padding_mask) # multi_token_prediction.py while `_checkpointed_forward` itself does not declare `padding_mask`. Any MTP run with `recompute_granularity == 'full'` therefore dies at the first forward: TypeError: MultiTokenPredictionLayer._checkpointed_forward() got an unexpected keyword argument 'padding_mask' Upstream cause: two landings that do not compose -- NVIDIA/Megatron-LM#2645 added the call-site kwarg, verl-project#4593 refactored the method without it. Tracked as NVIDIA/Megatron-LM#4933 (open since 2026-05-22). It is fixed on megatron-core `main`, but `core_v0.18.2` was tagged 2026-07-20 -- two months after that issue was filed -- and is still the newest release, so **every released megatron-core hits this**. verl users cannot currently combine MTP with full activation recomputation without patching megatron themselves. ### Why the existing patch does not cover it `patch_mtp_layer_checkpointed_forward` skips any layer whose `_checkpointed_forward` does not start with `forward_func`. That gate matches megatron-core 0.14-0.17 only; 0.18+ renamed the first parameter to `hidden_states`, so on every 0.18+ install the patch silently does nothing: when `target_layers` is non-empty but `patched_count == 0`, neither log line is printed, so there is no way to tell whether it applied. Skipping is in fact *correct* on 0.18+ (megatron now keeps non-tensor args out of the checkpoint natively, via its own `custom_forward` closure) -- but it should say so. ## Changes - `_patch_padding_mask_kwarg()`: when the layer's `_checkpointed_forward` lacks `padding_mask`, rebind it to accept and drop the kwarg. `padding_mask` is dropped rather than forwarded because there is no parameter to forward it to on these versions, and verl never constructs one, so it is always `None` in practice. A non-`None` value raises `NotImplementedError` rather than being silently ignored -- ignoring it would treat padded positions as real tokens. Idempotent per layer. - `patch_mtp_layer_checkpointed_forward()`: log the skip-by-signature and shim counts, so "patch did not apply" is observable rather than inferred from missing output. The style follows the existing precedent a few lines above in the same file, which already probes `signature(self.mtp.forward)` for `padding_mask` before passing it. ## Test `tests/models/test_mtp_checkpointed_forward_shim_on_cpu.py` -- 5 CPU cases, no GPU, no distributed init, covering all three megatron-core signature generations: - megatron-core `main` signature (declares `padding_mask`) -> shim not installed, kwarg still reaches the method - 0.18.x signature -> reproduces the bare `TypeError` first, then asserts the shim accepts `padding_mask=None` and forwards the remaining args unchanged - non-`None` mask -> `NotImplementedError` instead of silent corruption - 0.14-0.17 `forward_func` signature -> left to the existing recompute patch - shim is idempotent Also verified end-to-end on 4x8 H20 with Qwen3.6-35B-A3B (MoE, 40 layers, 256 experts), TP2/PP2/CP4/EP8, 64K context, `mtp_num_layers=1`, `recompute_granularity=full`, megatron-core 0.18.2 + torch 2.11: without the shim training aborts at the first forward with the `TypeError` above; with it, MTP trains and `mtp_losses/*` is reported normally (30 steps, loss 0.447 -> 0.318). Reported by / verified with AI assistance (Claude). Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: gaohongkui <gaohongkui@meituan.com>

…tron-core 0.18.x ## What does this PR do? Adds a compatibility shim so Multi-Token Prediction can be used together with `recompute_granularity=full` on every released megatron-core, and makes the existing signature-based branch in `patch_mtp_layer_checkpointed_forward` visible instead of silent. ### The bug megatron-core 0.18.x is internally inconsistent: `MultiTokenPredictionLayer.forward` calls self._checkpointed_forward(..., padding_mask=padding_mask) # multi_token_prediction.py while `_checkpointed_forward` itself does not declare `padding_mask`. Any MTP run with `recompute_granularity == 'full'` therefore dies at the first forward: TypeError: MultiTokenPredictionLayer._checkpointed_forward() got an unexpected keyword argument 'padding_mask' Upstream cause: two landings that do not compose -- NVIDIA/Megatron-LM#2645 added the call-site kwarg, verl-project#4593 refactored the method without it. Tracked as NVIDIA/Megatron-LM#4933 (open since 2026-05-22). It is fixed on megatron-core `main`, but `core_v0.18.2` was tagged 2026-07-20 -- two months after that issue was filed -- and is still the newest release, so **every released megatron-core hits this**. verl users cannot currently combine MTP with full activation recomputation without patching megatron themselves. ### Why the existing patch does not cover it `patch_mtp_layer_checkpointed_forward` skips any layer whose `_checkpointed_forward` does not start with `forward_func`. That gate matches megatron-core 0.14-0.17 only; 0.18+ renamed the first parameter to `hidden_states`, so on every 0.18+ install the patch silently does nothing: when `target_layers` is non-empty but `patched_count == 0`, neither log line is printed, so there is no way to tell whether it applied. Skipping is in fact *correct* on 0.18+ (megatron now keeps non-tensor args out of the checkpoint natively, via its own `custom_forward` closure) -- but it should say so. ## Changes - `_patch_padding_mask_kwarg()`: when the layer's `_checkpointed_forward` lacks `padding_mask`, rebind it to accept and drop the kwarg. `padding_mask` is dropped rather than forwarded because there is no parameter to forward it to on these versions, and verl never constructs one, so it is always `None` in practice. A non-`None` value raises `NotImplementedError` rather than being silently ignored -- ignoring it would treat padded positions as real tokens. Idempotent per layer. - `patch_mtp_layer_checkpointed_forward()`: log the skip-by-signature and shim counts, so "patch did not apply" is observable rather than inferred from missing output. The style follows the existing precedent a few lines above in the same file, which already probes `signature(self.mtp.forward)` for `padding_mask` before passing it. ## Test `tests/models/test_mtp_checkpointed_forward_shim_on_cpu.py` -- 5 CPU cases, no GPU, no distributed init, covering all three megatron-core signature generations: - megatron-core `main` signature (declares `padding_mask`) -> shim not installed, kwarg still reaches the method - 0.18.x signature -> reproduces the bare `TypeError` first, then asserts the shim accepts `padding_mask=None` and forwards the remaining args unchanged - non-`None` mask -> `NotImplementedError` instead of silent corruption - 0.14-0.17 `forward_func` signature -> left to the existing recompute patch - shim is idempotent Also verified end-to-end on 4x8 H20 with Qwen3.6-35B-A3B (MoE, 40 layers, 256 experts), TP2/PP2/CP4/EP8, 64K context, `mtp_num_layers=1`, `recompute_granularity=full`, megatron-core 0.18.2 + torch 2.11: without the shim training aborts at the first forward with the `TypeError` above; with it, MTP trains and `mtp_losses/*` is reported normally (30 steps, loss 0.447 -> 0.318). Reported by / verified with AI assistance (Claude). Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: gaohongkui <gaohongkui1021@163.com>

yuzhongw-nvidia mentioned this pull request Dec 12, 2025

[dev] feat(moe): Support packed sequence for gated delta net (GDN) #2644

Merged

6 tasks

yuzhongw-nvidia force-pushed the gdn_thd branch 2 times, most recently from 2575c6d to 4f8888d Compare December 15, 2025 03:45

yuzhongw-nvidia force-pushed the gdn_thd branch from 4f8888d to 9ccf5a4 Compare January 6, 2026 09:42

yuzhongw-nvidia force-pushed the gdn_thd branch 2 times, most recently from 73d512d to ae8806c Compare January 21, 2026 11:04

copy-pr-bot Bot temporarily deployed to nemo-ci January 21, 2026 13:09 Inactive

copy-pr-bot Bot temporarily deployed to test January 21, 2026 13:10 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci January 21, 2026 13:54 Inactive

yuzhongw-nvidia force-pushed the gdn_thd branch from 83a8607 to 545a2a5 Compare January 22, 2026 03:59

copy-pr-bot Bot temporarily deployed to nemo-ci January 22, 2026 04:00 Inactive

yuzhongw-nvidia force-pushed the gdn_thd branch 2 times, most recently from 58fdd22 to e8ed23c Compare January 28, 2026 13:43

copy-pr-bot Bot temporarily deployed to nemo-ci January 29, 2026 06:15 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci January 29, 2026 06:16 Inactive

yuzhongw-nvidia force-pushed the gdn_thd branch from 80d2d1c to e94395d Compare February 26, 2026 09:07

yuzhongw-nvidia force-pushed the gdn_thd branch 3 times, most recently from 8ae8e94 to cebd475 Compare April 7, 2026 04:55

copy-pr-bot Bot temporarily deployed to test April 7, 2026 04:56 Inactive

yuzhongw-nvidia force-pushed the gdn_thd branch from cebd475 to ec781f4 Compare April 14, 2026 02:39

Merged via the queue into NVIDIA:main with commit 2d1fa8d May 14, 2026
69 of 72 checks passed

This was referenced May 14, 2026

[recipe] feat: enable THD packing by default for Qwen3.5-VL finetune NVIDIA-NeMo/Megatron-Bridge#3481

Merged

[model, perf] feat: real THD packing in qwen3_vl_step NVIDIA-NeMo/Megatron-Bridge#3838

Closed

xuantengh mentioned this pull request May 21, 2026

Fuse per-sequence AlltoAll into a unified one in GDN forward #4913

Merged

5 tasks

ko3n1g mentioned this pull request May 22, 2026

🐛 CI failure: MultiTokenPredictionLayer._checkpointed_forward() got unexpected kwarg 'padding_mask' #4933

Open

factnn mentioned this pull request May 25, 2026

Fix _checkpointed_forward missing padding_mask parameter #4966

Closed

ko3n1g mentioned this pull request May 26, 2026

fix: forward padding_mask through MTP recompute path #4983

Closed

sbhavani mentioned this pull request May 26, 2026

[ROADMAP][2026 Q2] Megatron Core Roadmap #4997

Closed

dznvidia mentioned this pull request May 28, 2026

GDN packed sequence support landing in Nemo containers #5044

Open

Zhichenzzz mentioned this pull request Jun 8, 2026

fix(gdn): only reject THD packed_seq_params, allow BSHD radixark/Megatron-LM#52

Merged

erictang000 mentioned this pull request Jun 8, 2026

[chore] move megatron-core from dev latest to main latest, bump megatron-bridge to latest NovaSky-AI/SkyRL#1762

Merged

wuxibin89 mentioned this pull request Jun 29, 2026

[megatron] fix: support FP8 padding for BSHD actor forward verl-project/verl#6884

Merged

This was referenced Aug 8, 2026

fix(mtp): thread padding_mask through _checkpointed_forward (backport to core_r0.18.0) #6367

Open

[megatron] fix: make MTP work with recompute_granularity=full on megatron-core 0.18.x verl-project/verl#7326

Closed

yapdianang mentioned this pull request Aug 17, 2026

Support packed sequences in Gated Delta Net radixark/Megatron-LM#85

Draft

Navigation Menu

Uh oh!

[main] feat(moe): Support packed sequence for gated delta net (GDN) - #2645

[main] feat(moe): Support packed sequence for gated delta net (GDN)#2645
asolergi-nv merged 6 commits into
NVIDIA:mainfrom
yuzhongw-nvidia:gdn_thd

yuzhongw-nvidia commented Dec 12, 2025 •

edited by Phlip79

Loading

Uh oh!

copy-pr-bot Bot commented Dec 12, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

Uh oh!

Conversation

yuzhongw-nvidia commented Dec 12, 2025 • edited by Phlip79 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Contribution process

Pre-checks

Code review

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

(Step 3): Final Review

(Optional Step 4): Cherry-pick into release branch

Merging your PR

Uh oh!

copy-pr-bot Bot commented Dec 12, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

yuzhongw-nvidia commented Dec 12, 2025 •

edited by Phlip79

Loading

(Step 1): Add PR label `Expert Review`