What does this PR do ?

Summary

Qwen3.5 (the successor to Qwen3-Next) stores the GDN A_log and out_norm (output layernorm) parameters in FP32 in its released checkpoints, while the rest of the model is BF16. The current GatedDeltaNet implementation allocates these in config.params_dtype, causing dtype mismatches during checkpoint conversion. This PR aligns the Megatron implementation with the upstream precision convention.

• Store A_log parameter in FP32 instead of config.params_dtype
• Cast out_norm (output layernorm) to FP32 after construction
• Keep A_log initialization in FP32 in reset_parameters
• No changes to dt_bias (remains at config.params_dtype)

Motivation

The Qwen3.5 family (both the 27B dense and three MoE variants) releases A_log and linear_attn.norm.weight in FP32, even though the rest of the checkpoint is BF16. This is a deliberate precision choice by the model authors for numerical stability.

The GDN forward pass already computes the decay gate g in FP32:

g = -self.A_log.exp() * F.softplus(alpha.float() + self.dt_bias)  # In fp32

However, the A_log parameter and out_norm weights were previously allocated in config.params_dtype (typically BF16). This creates two problems:

1. Precision loss in A_log: A_log stores log(A) where A is a per-head decay rate. The exponential A_log.exp() amplifies quantization error from BF16 storage, even though the compute is done in FP32.

2. Precision loss in out_norm: The output RMSNorm applied after the recurrence operates on per-head-dim scale factors. Storing these in BF16 introduces unnecessary rounding.

Both parameters are small (per-head or per-head-dim vectors), so the FP32 memory overhead is negligible.

Changes

megatron/core/ssm/gated_delta_net.py:

• __init__: A_log allocated with dtype=torch.float32 (was config.params_dtype)
• __init__: out_norm cast to FP32 via self.out_norm.to(torch.float32) after build_module
• reset_parameters: A_log initialization uses dtype=torch.float32 (was config.params_dtype)

Precedent

This follows the existing convention in MambaMixer, where A_log is already unconditionally stored in FP32:

# mamba_mixer.py, line 336-341
A = torch.empty(self.nheads_local_tp, dtype=torch.float32, ...)
A_log = torch.log(A)  # Keep A_log in fp32
self.A_log = nn.Parameter(A_log)

Alternative designs considered

Option A: Configurable via TransformerConfig flags

Add two boolean fields (linear_A_log_in_fp32, linear_out_norm_in_fp32) to TransformerConfig, defaulting to True, and have GatedDeltaNet read them.

• Pro: Models that need BF16 for these params (e.g., Qwen3-Next) can set the flags to False for bitwise-identical checkpoint roundtrips.
• Con: TransformerConfig already has ~197 fields. Adding per-parameter precision knobs for tiny vectors (one scalar per head) is hard to justify given the negligible memory difference. Every GDN-based model bridge must also remember to set these flags correctly.

Why we chose hardcoded FP32 instead:

• The memory overhead is negligible (these are per-head vectors, not weight matrices).
• The forward pass already computes in FP32 regardless.
• MambaMixer already hardcodes A_log to FP32 — this aligns the two SSM-family modules.
• For models that previously stored these in BF16 (e.g., Qwen3-Next), the BF16-to-FP32 upcast is lossless (all BF16 values are exactly representable in FP32), so checkpoint loading is unaffected. Export will produce FP32 values that are mathematically identical to the original BF16 values.
• Qwen3-Next is superseded by Qwen3.5, which already uses FP32 for these params.
  ⚠️ 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

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.

Merging your PR

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