Support qwen3-next MTP Training

[zx3xyy]

Bug: linear_qgkv not handled in weight converter

Problem

Upstream's scripts/models/qwen3-next-80B-A3B.sh enables --use-gated-attention, which causes Megatron to create parameters named linear_qgkv.weight instead of linear_qkv.weight. But the weight converters (megatron_to_hf/qwen3_next.py and mbridge/qwen3_next.py) only handle linear_qkv, causing a crash during weight sync.

Root Cause

In Megatron's megatron/core/transformer/attention.py (line 891):

if self.config.use_gated_attention:
    self.linear_qgkv = build_module(...)   # Q+G+K+V fused
else:
    self.linear_qkv = build_module(...)    # Q+K+V fused

MTP layers inherit this config, so they also use linear_qgkv.

Reproduction

Test script that loads the upstream converter and calls it with linear_qgkv parameter names (as produced by Megatron when --use-gated-attention is set):

import re, torch, traceback
from types import SimpleNamespace

# Load the upstream converter function directly
with open('slime/backends/megatron_utils/megatron_to_hf/qwen3_next.py') as f:
    source = f.read()
exec(compile(source, 'qwen3_next.py', 'exec'))

# Qwen3-Next-80B-A3B model config
args = SimpleNamespace(
    num_attention_heads=16,
    num_query_groups=2,
    kv_channels=256,
    hidden_size=2048,
)
param = torch.randn(9216, 2048)

# Test 1: linear_qkv.weight (handled by upstream - works)
try:
    result = convert_qwen3_next_to_hf(args,
        'module.module.decoder.layers.0.self_attention.linear_qkv.weight', param)
    print(f'[OK] linear_qkv.weight -> {[r[0] for r in result]}')
except Exception as e:
    print(f'[FAIL] linear_qkv.weight: {e}')

# Test 2: linear_qgkv.weight (created when --use-gated-attention is set)
try:
    result = convert_qwen3_next_to_hf(args,
        'module.module.decoder.layers.0.self_attention.linear_qgkv.weight', param)
    print(f'[OK] linear_qgkv.weight -> {[r[0] for r in result]}')
except ValueError as e:
    traceback.print_exc()

# Test 3: linear_qgkv.layer_norm_weight
try:
    result = convert_qwen3_next_to_hf(args,
        'module.module.decoder.layers.0.self_attention.linear_qgkv.layer_norm_weight',
        torch.randn(2048))
    print(f'[OK] linear_qgkv.layer_norm_weight -> {[r[0] for r in result]}')
except ValueError as e:
    traceback.print_exc()

# Test 4: MTP layer with linear_qgkv (MTP inherits use_gated_attention from config)
try:
    result = convert_qwen3_next_to_hf(args,
        'module.module.mtp.layers.0.transformer_layer.self_attention.linear_qgkv.weight',
        param)
    print(f'[OK] MTP linear_qgkv.weight -> {[r[0] for r in result]}')
except ValueError as e:
    traceback.print_exc()

Output (before fix)

[OK] linear_qkv.weight -> ['model.layers.0.self_attn.q_proj.weight', 'model.layers.0.self_attn.k_proj.weight', 'model.layers.0.self_attn.v_proj.weight']

Traceback (most recent call last):
  File "qwen3_next.py", line 194, in convert_qwen3_next_to_hf
ValueError: Unknown parameter name: module.module.decoder.layers.0.self_attention.linear_qgkv.weight

Traceback (most recent call last):
  File "qwen3_next.py", line 194, in convert_qwen3_next_to_hf
ValueError: Unknown parameter name: module.module.decoder.layers.0.self_attention.linear_qgkv.layer_norm_weight

Traceback (most recent call last):
  File "qwen3_next.py", line 194, in convert_qwen3_next_to_hf
ValueError: Unknown parameter name: module.module.decoder.layers.0.self_attention.linear_qgkv.weight

Summary

[OK]            linear_qkv.weight             -> converts correctly
[BUG CONFIRMED] linear_qgkv.weight            -> ValueError: Unknown parameter name
[BUG CONFIRMED] linear_qgkv.layer_norm_weight -> ValueError: Unknown parameter name
[BUG CONFIRMED] MTP linear_qgkv.weight        -> ValueError: Unknown parameter name

After fix

[OK] linear_qkv.weight            -> ['model.layers.0.self_attn.q_proj.weight', 'model.layers.0.self_attn.k_proj.weight', 'model.layers.0.self_attn.v_proj.weight']
[OK] linear_qgkv.weight           -> ['model.layers.0.self_attn.q_proj.weight', 'model.layers.0.self_attn.k_proj.weight', 'model.layers.0.self_attn.v_proj.weight']
[OK] linear_qgkv.layer_norm_weight -> ['model.layers.0.input_layernorm.weight']
[OK] MTP linear_qgkv.weight       -> ['mtp.layers.0.self_attn.q_proj.weight', 'mtp.layers.0.self_attn.k_proj.weight', 'mtp.layers.0.self_attn.v_proj.weight']

---

Bug: Incorrect eh_proj.weight half-swap in MTP conversion

Problem

The _convert_mtp_layer function in megatron_to_hf/qwen3_next.py and _weight_to_mcore_format / _weight_to_hf_format in mbridge/qwen3_next.py swap the two halves of eh_proj.weight along dim=1 when converting between HF and Megatron formats:

first_half, second_half = param.chunk(2, dim=1)
new_param = torch.cat([second_half, first_half], dim=1)

This swap is incorrect — both HF and Megatron use the same input order [embedding, hidden_state].

Evidence

HF (SGLang) forward (sglang/python/sglang/srt/models/qwen3_next_mtp.py):

hidden_states = self.fc(torch.cat((input_embeds, hidden_states), dim=-1))

Megatron forward (megatron/core/transformer/multi_token_prediction.py):

hidden_states = torch.cat((decoder_input, hidden_states), -1)
hidden_states, _ = self.eh_proj(hidden_states)

Both concatenate [embedding, hidden_state] — same order. No swap needed.

Validation

• v15 checkpoint (with swap): MTP loss abnormally high, model not learning
• v16 checkpoint (without swap): MTP loss normal, matches v13 behavior

Fix

Remove the half-swap from both megatron_to_hf/qwen3_next.py and mbridge/qwen3_next.py. Pass eh_proj.weight through directly without modification.

Results

With these two fixes, we can try MTP layer of Qwen3-Next-80B-A3B with a reasonable loss.

[guapisolo]

Good job! I will help review this.

[guapisolo]

Good job!

[guapisolo]

Good fix.