[main] feat(moe): Support gated delta net for Qwen3-Next (1/4) by yuzhongw-nvidia · Pull Request #1989 · NVIDIA/Megatron-LM

yuzhongw-nvidia · 2025-10-28T02:59:54Z

What does this PR do ?

Qwen3-Next functionality PRs.

Changes in this PR:

Gated Delta Net (GDN);
num_floating_point_operations

New supported arguments for Qwen3-Next

--attention-output-gate: true  # Qwen3-Next applies output gate to standard attention layers
--no-weight-decay-cond-type: qwen3_next  # Qwen3-Next applies weight decay to qk layernorm as a special case
--moe-shared-expert-gate: true  # Qwen3-Next have a shared expert gate
# GDN args
--linear-attention-type: gated_delta_net
--linear-attention-freq: 4  # 1 1 1 0 1 1 1 0 ..., or we can use ([1,1,1,0]*12) to represent
--linear-conv-kernel-dim: 4
--linear-key-head-dim: 128
--linear-value-head-dim: 128
--linear-num-key-heads: 16
--linear-num-value-heads: 32

New for Qwen3-Next, but already supported in MCore yet

--rotary-percent: 0.25  # Qwen3-Next applies RoPE only to the first 25% of position dimensions
--rotary-base: 10000000  # Qwen3-Next uses a larger rotary base
--apply-layernorm-1p: true  # Zero-Centered RMSNorm for Qwen3-Next
--max-position-embeddings: 262144  # Qwen3-Next-80B-A3B supports up to 256k sequence length
--moe-shared-expert-intermediate-size: 512  # Qwen3-Next have a shared expert
# MTP: Qwen3-Next has one MTP layer
--mtp-num-layers: 1
--mtp-loss-scaling-factor: 0.1

LM loss curve with the training dataset of Qwen3 are as below (GBS=256, seq_len=4096, TP1 is in green, TP2 is in blue).

wandb url

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share discuss a design-doc with the team.

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-10-28T02:59:58Z

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.

CODEOWNERS

jaredcasper · 2025-10-30T20:00:37Z

What tests have been done to validate the implementation of the gated delta net is correct?

megatron/training/arguments.py

yuzhongw-nvidia · 2025-11-05T10:17:50Z

Hi @jaredcasper , I have rebased main and make the following changes to this MR,

Resolve [main] feat(moe): Support gated delta net for Qwen3-Next (1/4) #1989 (comment)
Try to make Attention.get_query_key_value_tensors clear. Does it look better now?

Thanks. cc @yanring

yuzhongw-nvidia · 2025-11-05T10:36:56Z

What tests have been done to validate the implementation of the gated delta net is correct?

Hi @jaredcasper , we have do several things to validate the implementation.

We have a UT to verify the implementation of GDN and parallel GDN.
We have used the comparing tool in Megatron-Bridge to verify the alignment with HF implementation.
We have run a E2E training test and see the loss curve makes sense. (wandb)

In the next few weeks, we will further add a functional test to guarantee the long-term correctness of Qwen3-Next.

jaredcasper · 2025-11-10T22:24:19Z

In the next few weeks, we will further add a functional test to guarantee the long-term correctness of Qwen3-Next.

In the meantime do you think it makes sense to split this up and get the other features in?

yuzhongw-nvidia · 2025-11-11T02:56:38Z

In the next few weeks, we will further add a functional test to guarantee the long-term correctness of Qwen3-Next.

In the meantime do you think it makes sense to split this up and get the other features in?

Do you mean covering multiple features in a single functional test? It will be excellent, as long as we do not have too many features in one test case making it hard to maintain.

yuzhongw-nvidia · 2025-11-11T07:29:38Z

Hi @jaredcasper , I remember you mentioned that you need a thorough review to the gpt_layer_spec.py. Do you have any comments or feedback about it? Does it look good, or do you think we need to refine it?

jaredcasper · 2025-11-13T01:59:56Z

Do you mean covering multiple features in a single functional test? It will be excellent, as long as we do not have too many features in one test case making it hard to maintain.

No, I mean there are 7 changes listed in the PR, do they all need to be in one PR? I'm saying while some of the features are getting tests added for them does it make sense to open PRs for some of the others that have been tested.

jaredcasper · 2025-11-13T02:01:45Z

Hi @jaredcasper , I remember you mentioned that you need a thorough review to the gpt_layer_spec.py. Do you have any comments or feedback about it? Does it look good, or do you think we need to refine it?

I thought you were going to work on it a bit after our conversation, but if you are done refining I've asked some from my team to do a thorough review.

megatron/core/ssm/gated_delta_net.py

Phlip79 · 2026-01-12T19:25:34Z

/ok to test c574f43

JRD971000

lgtm

yuzhongw-nvidia · 2026-01-14T08:39:36Z

/ok to test 271bb87

yuzhongw-nvidia · 2026-01-14T08:50:42Z

/ok to test 7c2a449

yanring

LGTM, thanks for the refinement!

yuzhongw-nvidia · 2026-01-14T10:05:50Z

/ok to test 8fe3705

yuzhongw-nvidia · 2026-01-14T14:58:18Z

/ok to test 476bc57

jaredcasper · 2026-01-14T18:20:35Z

megatron/training/training.py

+            key_projection_size = args.kv_channels * args.num_query_groups
+            value_projection_size = args.kv_channels * args.num_query_groups
+            standard_self_attn_term = (
+                3


Can you add explanations here since you aren't using expansion factor? Or ever better use expansion_factor to make this more readable. The transformation of this equation you've done here isn't immediately obvious, can you make it closer to the previous formulation?

I'm sure it is correct because:

I have checked carefully and make sure it is equivalent to the previous code.

I have checked in an E2E test and find it outputs the same result as before.

I also add some comments to make it easy to understand.

I do not agree with the point "use expansion_factor to make this more readable". I know most of the people are confused by it because there is a 2x in the expansion_factor:

# - 2x: GEMMs of a particular size are stacked twice in the standard Transformer model # architectures implemented in this codebase (e.g., h->ffn_h GEMM and ffn_h->h GEMM # in MLP layer).

which is only immediately obvious for MLP. Could you please tell me what does this 2x factor mean in each term? To tell you the truth, many colleagues and I finally understand it after a careful derivation, reflecting how truly unreadable it previously was. In contrast, the calculation of MLA doesn't use expansion_factor, and as you can see, it is much cleaner than MHA and GQA. Furthermore, the previous version has a chaotic calculation sequence and lacked comments, rendering it completely unreadable. That's why I argue that we need to rewrite it.

Different approaches, I thought the previous was much more readable. :)

In any case, there is a large comment explaining expansion_factor, and the creation of the variable itself, but expansion_factor is not used any more (unless I'm missing its use?). Please clean this up if you are going to remove expansion_factor. The comments explaining expansion factor can be moved down here.

Thanks for your suggestion and discussion. I refine it again and break the expansion_factor into forward_backward_expansion_factor, fma_expansion_factor, and ffn_expansion_factor and refine their comments.

yuzhongw-nvidia · 2026-01-15T04:36:21Z

/ok to test 82ba0d0

yuzhongw-nvidia · 2026-01-16T03:46:26Z

/ok to test 3529e98

yuzhongw-nvidia · 2026-01-16T06:41:09Z

/ok to test 3e1ca94

Minor refines Co-authored-by: Li Tao <lit@nvidia.com> fix CI get_transformer_block_with_experimental_attention_variant_spec Update test_mamba_moe_model.py Reopen qwen3next functional test in lightweight mode (NVIDIA#2493) Signed-off-by: oliver könig <okoenig@nvidia.com> Co-authored-by: oliver könig <okoenig@nvidia.com>

yuzhongw-nvidia · 2026-01-16T10:39:42Z

/ok to test 418be3c

Signed-off-by: oliver könig <okoenig@nvidia.com>

yuzhongw-nvidia · 2026-01-16T12:15:26Z

/ok to test fd96ce8

yuzhongw-nvidia · 2026-01-16T12:50:42Z

/ok to test df953ac

yuzhongw-nvidia · 2026-01-16T13:53:59Z

/ok to test 13d2cc9

yuzhongw-nvidia · 2026-01-16T14:56:28Z

/ok to test 06e3ea9

yuzhongw-nvidia · 2026-01-16T15:39:04Z

/ok to test f58c310

yuzhongw-nvidia requested review from a team as code owners October 28, 2025 02:59

jaredcasper reviewed Oct 30, 2025

View reviewed changes

CODEOWNERS Outdated Show resolved Hide resolved

jaredcasper reviewed Oct 30, 2025

View reviewed changes

megatron/training/arguments.py Outdated Show resolved Hide resolved

yanring added the Expert Review [deprecated] Apply this label to indicate that your PR is ready for expert review. label Nov 5, 2025

yuzhongw-nvidia force-pushed the qwen3next branch from 9c74ec3 to 4fcc94b Compare November 5, 2025 07:04

yuzhongw-nvidia requested review from a team as code owners November 5, 2025 07:04

yuzhongw-nvidia force-pushed the qwen3next branch 2 times, most recently from 99d022e to 1a63170 Compare November 5, 2025 08:03

yuzhongw-nvidia force-pushed the qwen3next branch from 065a890 to c355b9a Compare November 7, 2025 06:09

yanring added the module: moe label Nov 7, 2025

yanring assigned yuzhongw-nvidia Nov 7, 2025

yuzhongw-nvidia mentioned this pull request Nov 13, 2025

[Dev] Fix Qwen3-Next hang on Blackwell, add a flag to control torch.compile #2058

Merged

6 tasks

BestJuly reviewed Jan 9, 2026

View reviewed changes

megatron/core/ssm/gated_delta_net.py Outdated Show resolved Hide resolved

BestJuly approved these changes Jan 12, 2026

View reviewed changes

JRD971000 approved these changes Jan 13, 2026

View reviewed changes

kunlunl mentioned this pull request Jan 14, 2026

DeepSeek V3.2 support #2440

Merged

yanring approved these changes Jan 14, 2026

View reviewed changes

jaredcasper reviewed Jan 14, 2026

View reviewed changes

yuzhongw-nvidia and others added 2 commits January 16, 2026 02:38

refine the flops computation

418be3c

ci(fix): CI_COMMIT_BRANCH on forks

b218786

Signed-off-by: oliver könig <okoenig@nvidia.com>

remove device usage

f58c310

jaredcasper approved these changes Jan 16, 2026

View reviewed changes

yuzhongw-nvidia mentioned this pull request Jan 21, 2026

[Main] ci(moe): Qwen3-next specific features functional test (lightweight mode) #2706

Closed

6 tasks

chtruong814 mentioned this pull request Jan 27, 2026

fix: Remove experimental_attention_variant arg from get_gpt_layer call in mlm_compat NVIDIA-NeMo/Megatron-Bridge#2073

Merged

5 tasks

Phlip79 mentioned this pull request Feb 25, 2026

Introduce GDN to Mamba #3535

Merged

6 tasks

cuichenx mentioned this pull request Mar 20, 2026

[perf] feat: add GDN (Gated DeltaNet) FLOPs calculator NVIDIA-NeMo/Megatron-Bridge#2925

Merged

5 tasks

sbhavani mentioned this pull request Mar 23, 2026

[ROADMAP] Megatron Core Roadmap #4003

Open

Conversation

yuzhongw-nvidia commented Oct 28, 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 Oct 28, 2025

Uh oh!

Uh oh!

jaredcasper commented Oct 30, 2025

Uh oh!

Uh oh!

yuzhongw-nvidia commented Nov 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yuzhongw-nvidia commented Nov 5, 2025

Uh oh!

jaredcasper commented Nov 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yuzhongw-nvidia commented Nov 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yuzhongw-nvidia commented Nov 11, 2025

Uh oh!

jaredcasper commented Nov 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jaredcasper commented Nov 13, 2025

Uh oh!

Uh oh!

Phlip79 commented Jan 12, 2026

Uh oh!

JRD971000 left a comment

Choose a reason for hiding this comment

Uh oh!

yuzhongw-nvidia commented Jan 14, 2026

Uh oh!

yuzhongw-nvidia commented Jan 14, 2026

Uh oh!

yanring left a comment

Choose a reason for hiding this comment

Uh oh!

yuzhongw-nvidia commented Jan 14, 2026

Uh oh!

yuzhongw-nvidia commented Jan 14, 2026

Uh oh!

jaredcasper Jan 14, 2026

Choose a reason for hiding this comment

Uh oh!

yuzhongw-nvidia Jan 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jaredcasper Jan 15, 2026

Choose a reason for hiding this comment

Uh oh!

yuzhongw-nvidia Jan 16, 2026

Choose a reason for hiding this comment

Uh oh!

yuzhongw-nvidia commented Jan 15, 2026

Uh oh!

yuzhongw-nvidia commented Jan 16, 2026

Uh oh!

yuzhongw-nvidia commented Jan 16, 2026

Uh oh!

yuzhongw-nvidia commented Jan 16, 2026

Uh oh!

yuzhongw-nvidia commented Jan 16, 2026

Uh oh!

yuzhongw-nvidia commented Jan 16, 2026

Uh oh!

yuzhongw-nvidia commented Jan 16, 2026

yuzhongw-nvidia commented Oct 28, 2025 •

edited

Loading

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

yuzhongw-nvidia commented Nov 5, 2025 •

edited

Loading

jaredcasper commented Nov 10, 2025 •

edited

Loading

yuzhongw-nvidia commented Nov 11, 2025 •

edited

Loading

jaredcasper commented Nov 13, 2025 •

edited

Loading

yuzhongw-nvidia Jan 15, 2026 •

edited

Loading