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

yuzhongw-nvidia · 2025-12-12T13:12:46Z

What does this PR do ?

Support packed sequence for gated delta net (GDN).

Requires cudnn>=9.19.0 for Qwen3.5 or Qwen3-Next-80B-A3B THD format training on Hopper. cudnn<=9.18.0 have some numerical issues, leading to grad NaN.

PR for main: #2645

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:12:50Z

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.

yuzhongw-nvidia · 2026-01-21T13:09:25Z

/ok to test ae8806c

yuzhongw-nvidia · 2026-01-22T03:59:58Z

/ok to test 545a2a5

kunlunl · 2026-01-28T12:50:55Z

Also need to update this line https://github.com/yuzhongw-nvidia/Megatron-LM/blob/befbcd2a10b60deb4edbd0f758275b26a6df83c7/megatron/core/ssm/gated_delta_net.py#L741 to param = param[tuple(slices)]?

yuzhongw-nvidia · 2026-01-28T13:43:39Z

Also need to update this line https://github.com/yuzhongw-nvidia/Megatron-LM/blob/befbcd2a10b60deb4edbd0f758275b26a6df83c7/megatron/core/ssm/gated_delta_net.py#L741 to param = param[tuple(slices)]?

Thanks. Resolved.

svcnvidia-nemo-ci · 2026-04-07T06:10:36Z

🔄 Merge queue validation started!

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

yuzhongw-nvidia · 2026-04-07T07:04:11Z

Hello, I wonder how this pr going now? : )

try the commit [e94395d] to fine-tune qwen3.5-4B with sequence packing and see grad_norm explosion, leading to final accuracy 0.

Hi @Zhikaiiii @Code4Graph ,

We have already merged this PR, but there are still some numerical issues for fused attention + Qwen3-Next / Qwen3.5 + THD format. Before the fix, we recommend you to use flash attention (--attention-backend flash). We will post more details ASAP in the PR desc and create a TE issue for this, thanks.

Zhikaiiii · 2026-04-07T07:26:36Z

Hello, I wonder how this pr going now? : )

try the commit [e94395d] to fine-tune qwen3.5-4B with sequence packing and see grad_norm explosion, leading to final accuracy 0.

Hi @Zhikaiiii @Code4Graph ,

We have already merged this PR, but there are still some numerical issues for fused attention + Qwen3-Next / Qwen3.5 + THD format. Before the fix, we recommend you to use flash attention (--attention-backend flash). We will post more details ASAP in the PR desc and create a TE issue for this, thanks.

Thanks for the great work! Just one more question: It seems that this feature has been merged into the dev branch as an experimental feature. Is there an estimated timeline for merging this feature into the main and release branches?

yuzhongw-nvidia · 2026-04-07T07:59:57Z

Thanks for the great work! Just one more question: It seems that this feature has been merged into the dev branch as an experimental feature. Is there an estimated timeline for merging this feature into the main and release branches?

No. We will move forward with the review process for the main PR, but we do not have an ETA for now.

Code4Graph · 2026-04-07T16:14:48Z

Hello, I wonder how this pr going now? : )

try the commit [e94395d] to fine-tune qwen3.5-4B with sequence packing and see grad_norm explosion, leading to final accuracy 0.

Hi @Zhikaiiii @Code4Graph ,

We have already merged this PR, but there are still some numerical issues for fused attention + Qwen3-Next / Qwen3.5 + THD format. Before the fix, we recommend you to use flash attention (--attention-backend flash). We will post more details ASAP in the PR desc and create a TE issue for this, thanks.

Thanks for the great work!

### What does this PR do? Support CP for bshd format, since mcore is still not support thd format for GDN NVIDIA/Megatron-LM#2644 ```bash pip install transformers==5.3.0 pip install flash-linear-attention # bshd relies on mcore dev branch pip install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@dev pip install --force-reinstall git+https://github.com/ISEEKYAN/mbridge.git ``` - model: Qwen3.5-0.8B - dataset: gsm8k - red: TP=2,CP=1; gray: TP=2,CP=2 <img width="320" height="280" alt="image" src="https://github.com/user-attachments/assets/63621abd-bdd7-4fae-8746-40573931b232" />

…VIDIA#2644) Signed-off-by: yuzhongw <yuzhongw@nvidia.com> Co-authored-by: kunlunl <kunlunl@nvidia.com>

### What does this PR do? Support CP for bshd format, since mcore is still not support thd format for GDN NVIDIA/Megatron-LM#2644 ```bash pip install transformers==5.3.0 pip install flash-linear-attention # bshd relies on mcore dev branch pip install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@dev pip install --force-reinstall git+https://github.com/ISEEKYAN/mbridge.git ``` - model: Qwen3.5-0.8B - dataset: gsm8k - red: TP=2,CP=1; gray: TP=2,CP=2 <img width="320" height="280" alt="image" src="https://github.com/user-attachments/assets/63621abd-bdd7-4fae-8746-40573931b232" />

…VIDIA#2644) Signed-off-by: yuzhongw <yuzhongw@nvidia.com> Co-authored-by: kunlunl <kunlunl@nvidia.com>

### What does this PR do? Support CP for bshd format, since mcore is still not support thd format for GDN NVIDIA/Megatron-LM#2644 ```bash pip install transformers==5.3.0 pip install flash-linear-attention # bshd relies on mcore dev branch pip install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@dev pip install --force-reinstall git+https://github.com/ISEEKYAN/mbridge.git ``` - model: Qwen3.5-0.8B - dataset: gsm8k - red: TP=2,CP=1; gray: TP=2,CP=2 <img width="320" height="280" alt="image" src="https://github.com/user-attachments/assets/63621abd-bdd7-4fae-8746-40573931b232" />

…DeltaNet Port the packed-sequence support from upstream NVIDIA/Megatron-LM PRs NVIDIA#2644 and NVIDIA#4230 into the GatedDeltaNet forward (adapted to this fork's GDN code): - Gate packed handling on `packed_seq_params.qkv_format == 'thd'` (not just `is not None`), so the sbhd PackedSeqParams Megatron passes uniformly is ignored and only genuine THD packing is processed. - Resolve cu_seqlens via `_resolve_cu_seqlens`, preferring *_padded and validating cu_seqlens[-1] == total_seq_len (alignment-padding handling). - Unpack the packed stream per-sequence around both CP all-to-all calls (cp2hp / hp2cp) via `_unpack_sequence`, and pass cu_seqlens into the causal_conv1d and gated_delta_rule kernels. - Guard torch_chunk_gated_delta_rule against cu_seqlens (unsupported). - Assert batch == 1 and non-deterministic mode for packed sequences. Also keep the independent fp32 numerics fix: cast A_log and dt_bias to fp32 before exp/softplus when computing g, matching the "In fp32" intent. Verified the cu_seqlens[-1] == seq_len invariant holds for the target training config (TP=4, sequence-parallel, CP=1, THD): seq_len = (L/tp)*sp*cp = L and cu_seqlens[-1] = L*cp = L. Note CP>1 + GDN is a known unsupported edge that would trip the validation. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

yuzhongw-nvidia requested review from a team as code owners December 12, 2025 13:12

yuzhongw-nvidia mentioned this pull request Dec 12, 2025

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

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

yuzhongw-nvidia requested review from a team and kunlunl January 21, 2026 13:09

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

ko3n1g added this to the Core 0.16 milestone Jan 21, 2026

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 from 545a2a5 to befbcd2 Compare January 22, 2026 07:52

yaox12 added the Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. label Jan 23, 2026

yanring mentioned this pull request Jan 26, 2026

[ROADMAP][Updated on April 07] Megatron Core MoE Roadmap #1729

Closed

48 tasks

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

yuzhongw-nvidia enabled auto-merge April 7, 2026 05:02

yuzhongw-nvidia added this pull request to the merge queue Apr 7, 2026

Merged via the queue into NVIDIA:dev with commit b6193a3 Apr 7, 2026
60 of 62 checks passed

yxs mentioned this pull request Apr 8, 2026

GDN packed sequence (THD format) + CP + SP causes shape mismatch at residual connection. #4194

Closed

CharlieFRuan mentioned this pull request Apr 16, 2026

Ensure sample_packing support for Qwen3.5 for FSDP2 and Megatron NovaSky-AI/SkyRL#1525

Open

Victarry mentioned this pull request May 15, 2026

[ROADMAP][2026 Q2] Megatron Core MoE Roadmap #4815

Closed

72 tasks

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

wuxibin89 mentioned this pull request Jun 2, 2026

[megatron, model] fix: support GDN with packed THD sequence parallel verl-project/verl#6559

Closed

6 tasks

Adiactive mentioned this pull request Jun 2, 2026

feat(megatron): Qwen3.5 dense + MoE training/inference support via megatron-bridge areal-project/AReaL#1384

Merged

15 tasks

dreamyang-liu mentioned this pull request Jun 26, 2026

GDN: THD packed-sequence support breaks under CP>1 dreamyang-liu/Megatron-LM#1

Open

gaohongkui mentioned this pull request Aug 10, 2026

[megatron] feat: add Qwen3.5 THD (packed sequence) Megatron SFT example verl-project/verl#7336

Open

6 tasks

Navigation Menu

Uh oh!

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

[dev] feat(moe): Support packed sequence for gated delta net (GDN)#2644
yuzhongw-nvidia merged 2 commits into
NVIDIA:devfrom
yuzhongw-nvidia:gdn_thd

yuzhongw-nvidia commented Dec 12, 2025 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Dec 12, 2025

Uh oh!

yuzhongw-nvidia commented Jan 21, 2026

Uh oh!

yuzhongw-nvidia commented Jan 22, 2026

Uh oh!

kunlunl commented Jan 28, 2026

Uh oh!

yuzhongw-nvidia commented Jan 28, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 7, 2026

Uh oh!

Uh oh!

yuzhongw-nvidia commented Apr 7, 2026

Uh oh!

Zhikaiiii commented Apr 7, 2026

Uh oh!

yuzhongw-nvidia commented Apr 7, 2026

Uh oh!

Code4Graph commented Apr 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Uh oh!

Conversation

yuzhongw-nvidia commented Dec 12, 2025 • edited 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!

yuzhongw-nvidia commented Jan 21, 2026

Uh oh!

yuzhongw-nvidia commented Jan 22, 2026

Uh oh!

kunlunl commented Jan 28, 2026

Uh oh!

yuzhongw-nvidia commented Jan 28, 2026

Uh oh!

svcnvidia-nemo-ci commented Apr 7, 2026

Uh oh!

Uh oh!

yuzhongw-nvidia commented Apr 7, 2026

Uh oh!

Zhikaiiii commented Apr 7, 2026

Uh oh!

yuzhongw-nvidia commented Apr 7, 2026

Uh oh!

Code4Graph commented Apr 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

yuzhongw-nvidia commented Dec 12, 2025 •

edited

Loading

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