cp: fix(moe): preserve fp32 A_log in Qwen3.5-{MoE,Next GatedDeltaNet} (2484) into r0.5.0

nemo_automodel/_transformers/model_init.py

@@ -60,6 +60,10 @@ def _check_model_inputs(func):
 import nemo_automodel.components.distributed.utils as dist_utils
 from nemo_automodel._transformers.registry import ModelRegistry
 from nemo_automodel.components.distributed.init_utils import get_local_world_size_preinit, get_world_size_safe
+from nemo_automodel.components.models.common.gated_delta_net_fp32 import (
+    has_gated_delta_net_fp32_checkpoint_contract,
+    is_gated_delta_net_fp32_param_key,
+)
 from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin
 from nemo_automodel.components.utils.model_utils import resolve_trust_remote_code, skip_random_init
 from nemo_automodel.shared.utils import dtype_from_str
@@ -662,29 +666,42 @@ def _restore_loaded_model_dtype(
     if not checkpoint_dtypes:
         return
 
+    preserve_gdn_fp32_params = has_gated_delta_net_fp32_checkpoint_contract(hf_config)
     restored_dtype_by_tensor_id: dict[int, torch.dtype] = {}
     restored_count = 0
     for name, checkpoint_dtype in checkpoint_dtypes.items():
         tensor = _get_model_tensor(model, name)
         if tensor is None:
             continue
 
+        effective_checkpoint_dtype = (
+            torch.float32
+            if checkpoint_dtype.is_floating_point
+            and preserve_gdn_fp32_params
+            and is_gated_delta_net_fp32_param_key(name)
+            else checkpoint_dtype
+        )
+
         # Record the checkpoint's original dtype on the tensor as the compute-dtype
         # hint. Storage may be upcast below (fp32 master weights), which erases the
         # dtype HF intended for compute; downstream sharding (fully_shard_by_dtype)
         # reads ``_hf_compute_dtype`` to keep intrinsically-fp32 params (e.g. ``A_log``)
         # computing in fp32 while the bulk computes in mp_policy.param_dtype.
-        if checkpoint_dtype.is_floating_point and tensor.dtype.is_floating_point:
-            tensor._hf_compute_dtype = checkpoint_dtype
+        if effective_checkpoint_dtype.is_floating_point and tensor.dtype.is_floating_point:
+            tensor._hf_compute_dtype = effective_checkpoint_dtype
 
         # Pick the unification target. For an explicit floating request, take the
         # wider of (checkpoint, requested) so explicit fp32 is honored as master
         # weights while intrinsically-fp32 checkpoint params survive a bf16 request.
         # For "auto" (requested_dtype is None) or non-floating tensors, mirror the
         # checkpoint dtype exactly (preserves today's behavior).
-        target_dtype = checkpoint_dtype
-        if requested_dtype is not None and checkpoint_dtype.is_floating_point and tensor.dtype.is_floating_point:
-            target_dtype = torch.promote_types(checkpoint_dtype, requested_dtype)
+        target_dtype = effective_checkpoint_dtype
+        if (
+            requested_dtype is not None
+            and effective_checkpoint_dtype.is_floating_point
+            and tensor.dtype.is_floating_point
+        ):
+            target_dtype = torch.promote_types(effective_checkpoint_dtype, requested_dtype)
 
         if tensor.dtype == target_dtype:
             continue

nemo_automodel/components/models/common/gated_delta_net_fp32.py

@@ -0,0 +1,105 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared checkpoint helpers for fp32 GatedDeltaNet (GDN) params.
+
+GDN layers carry intrinsically-fp32 bare parameters (``A_log`` and ``dt_bias``)
+that feed the decay gate ``g = -exp(A_log) * softplus(a + dt_bias)``. Under FSDP2
+mixed precision with fp32 master weights, the bulk of a model computes in bf16
+(``param_dtype=bf16``) while these parameters must stay in fp32 -- ``A_log`` is
+exponentiated, so bf16 rounding becomes a proportional error on the decay rate
+that the recurrence compounds across the sequence.
+
+Each model owns the runtime construction of its fp32 holder. This module only
+centralizes the checkpoint contract: hide ``_fp32_params`` in saved HF-compatible
+keys, route bare HF keys back into the holder for native load, and upcast these
+params to fp32 when checkpoint tensors arrive in a lower precision.
+"""
+
+from __future__ import annotations
+
+import re
+
+import torch
+
+HOLDER_NAME = "_fp32_params"
+
+# Intrinsically-fp32 GatedDeltaNet bare params routed through the holder.
+FP32_GDN_PARAM_NAMES = ("A_log", "dt_bias")
+GDN_FP32_CHECKPOINT_ARCHITECTURES = frozenset(
+    (
+        "Qwen3NextForCausalLM",
+        "Qwen3_5ForCausalLM",
+        "Qwen3_5ForConditionalGeneration",
+        "Qwen3_5MoeForConditionalGeneration",
+    )
+)
+
+_FP32_HOLDER_KEY_RE = re.compile(r"(\.linear_attn)\._fp32_params\.")
+
+
+def strip_fp32_holder_key(key: str) -> str:
+    """Rewrite ``...linear_attn._fp32_params.X`` -> ``...linear_attn.X``.
+
+    Used by state-dict adapters so saved checkpoints hide the ``_fp32_params``
+    wrapping and stay directly HF-loadable.
+    """
+    return _FP32_HOLDER_KEY_RE.sub(r"\1.", key)
+
+
+def route_fp32_holder_key(key: str, param_names: tuple[str, ...] = FP32_GDN_PARAM_NAMES) -> str:
+    """Rewrite a bare ``...linear_attn.X`` GDN param key into the ``_fp32_params`` holder.
+
+    Inverse of :func:`strip_fp32_holder_key` for the param names in ``param_names``.
+    No-op when the key is already routed, is not under ``linear_attn``, or is not a
+    tracked fp32 GDN param.
+    """
+    if not key.endswith(param_names):
+        return key
+    if "._fp32_params." in key:
+        return key
+    if ".linear_attn." not in key:
+        return key
+    head, tail = key.rsplit(".linear_attn.", 1)
+    return f"{head}.linear_attn._fp32_params.{tail}"
+
+
+def is_gated_delta_net_fp32_param_key(key: str, param_names: tuple[str, ...] = FP32_GDN_PARAM_NAMES) -> bool:
+    """Return whether ``key`` names an intrinsically-fp32 GDN parameter."""
+    return key.endswith(param_names) and ".linear_attn." in key
+
+
+def has_gated_delta_net_fp32_checkpoint_contract(hf_config: object) -> bool:
+    """Return whether ``hf_config`` belongs to an architecture with fp32 GDN params."""
+    architectures = getattr(hf_config, "architectures", None) or ()
+    return any(arch in GDN_FP32_CHECKPOINT_ARCHITECTURES for arch in architectures)
+
+
+def upcast_gated_delta_net_fp32_state_tensor(
+    key: str, tensor: object, param_names: tuple[str, ...] = FP32_GDN_PARAM_NAMES
+) -> object:
+    """Cast loaded GDN fp32-param tensors to fp32 while leaving other state untouched.
+
+    Construction-time upcasting is not enough for checkpoint and HF load paths that
+    replace or carry tensor values from disk. This helper preserves the fp32 GDN
+    contract at adapter boundaries before tensors enter the live model state dict.
+    """
+    if not is_gated_delta_net_fp32_param_key(key, param_names):
+        return tensor
+    if getattr(tensor, "dtype", None) == torch.float32:
+        return tensor
+    is_floating_point = getattr(tensor, "is_floating_point", None)
+    if callable(is_floating_point) and is_floating_point():
+        return tensor.to(dtype=torch.float32)
+    return tensor

nemo_automodel/components/models/common/utils.py

@@ -461,8 +461,10 @@ def cast_model_to_dtype(
 
     Uses ``nn.Module.to()`` which is safe for both plain tensors and DTensors
     (FSDP2 sharded parameters).  When the model is already FSDP2-sharded
-    (parameters are DTensors), only buffers of matching modules are restored
-    to fp32 (parameters are left as-is since FSDP2 requires uniform dtype).
+    (parameters are DTensors), strict fp32 modules are restored to fp32 because
+    they are expected to be isolated as uniform fp32 FSDP units. Non-strict fp32
+    hints only restore matching buffers, since their parameters may share an
+    FSDP unit with lower-precision parameters.
 
     Args:
         model: The model whose parameters should be cast.
@@ -478,6 +480,21 @@ def cast_model_to_dtype(
             FSDP's uniform-dtype rule.
     """
     fp32_keywords = _get_fp32_module_keywords(model)
+    strict_fp32_keywords = _get_strict_fp32_module_keywords(model)
+    has_dtensor_params = _has_dtensor_params(model)
+
+    if has_dtensor_params:
+        fp32_snapshots = _snapshot_fp32_tensors(
+            model,
+            parameter_keywords=strict_fp32_keywords,
+            buffer_keywords=fp32_keywords,
+        )
+    else:
+        fp32_snapshots = _snapshot_fp32_tensors(
+            model,
+            parameter_keywords=fp32_keywords,
+            buffer_keywords=fp32_keywords,
+        )
 
     # Detach skip_modules so ``model.to(dtype)`` does not descend into them. This
     # preserves their exact dtype (e.g. fp32 master weights) through the cast.
@@ -496,16 +513,33 @@ def cast_model_to_dtype(
             parent._modules[child_name] = child
 
     if fp32_keywords:
-        if _has_dtensor_params(model):
-            logger.warning(
-                "Model parameters are DTensors (FSDP2) — skipping fp32 parameter "
-                "restoration for keywords=%s. Only buffers will be restored to fp32. "
-                "FSDP2 requires uniform dtype within each parameter group.",
-                fp32_keywords,
+        if has_dtensor_params:
+            if strict_fp32_keywords:
+                _restore_fp32_tensor_snapshots(
+                    model,
+                    parameter_snapshots=fp32_snapshots[0],
+                    buffer_snapshots={},
+                )
+
+            buffer_only_keywords = [kw for kw in fp32_keywords if kw not in strict_fp32_keywords]
+            if buffer_only_keywords:
+                logger.warning(
+                    "Model parameters are DTensors (FSDP2) — skipping fp32 parameter "
+                    "restoration for non-strict keywords=%s. Only buffers will be restored to fp32. "
+                    "FSDP2 requires uniform dtype within each parameter group.",
+                    buffer_only_keywords,
+                )
+            _restore_fp32_tensor_snapshots(
+                model,
+                parameter_snapshots={},
+                buffer_snapshots=fp32_snapshots[1],
             )
-            _restore_fp32_buffers(model, fp32_keywords)
         else:
-            _restore_fp32_modules(model, fp32_keywords)
+            _restore_fp32_tensor_snapshots(
+                model,
+                parameter_snapshots=fp32_snapshots[0],
+                buffer_snapshots=fp32_snapshots[1],
+            )
 
 
 @contextmanager
@@ -533,8 +567,8 @@ def yield_fp32_model(model: nn.Module, restore_dtype: torch.dtype | None = None)
 
     ``_keep_in_fp32_modules`` / ``_keep_in_fp32_modules_strict`` handling is delegated to
     ``cast_model_to_dtype``: on an unsharded model those modules' params and buffers are restored
-    to fp32 on exit; on a sharded model only their buffers are (sharded fp32 params are pinned at
-    shard time, since FSDP2 forbids mixed dtypes within a group).
+    to fp32 on exit; on a sharded model, strict fp32 modules are restored while non-strict modules
+    only have their buffers restored.
 
     Args:
         model: The model to run in fp32 within the context.
@@ -557,6 +591,13 @@ def yield_fp32_model(model: nn.Module, restore_dtype: torch.dtype | None = None)
         cast_model_to_dtype(model, restore_dtype)
 
 
+def _get_strict_fp32_module_keywords(model: nn.Module) -> list[str]:
+    val = getattr(model, "_keep_in_fp32_modules_strict", None)
+    if not isinstance(val, (list, set, tuple)):
+        return []
+    return list(dict.fromkeys(val))
+
+
 def _get_fp32_module_keywords(model: nn.Module) -> list[str]:
     """Collect module name patterns that must remain in fp32.
 
@@ -590,6 +631,53 @@ def _has_dtensor_params(model: nn.Module) -> bool:
     return any(isinstance(p, DTensor) for p in model.parameters())
 
 
+def _snapshot_fp32_tensors(
+    model: nn.Module,
+    *,
+    parameter_keywords: list[str],
+    buffer_keywords: list[str],
+) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
+    """Clone fp32-preserved tensors before a broad dtype cast.
+
+    Casting ``fp32 -> bf16 -> fp32`` restores the dtype but not the original
+    values. Snapshot the matching tensors first so strict fp32 state such as
+    router correction bias or recurrent-decay parameters is restored exactly.
+    """
+    parameter_snapshots = {
+        name: param.detach().to(torch.float32).clone()
+        for name, param in model.named_parameters()
+        if param.is_floating_point() and any(keyword in name for keyword in parameter_keywords)
+    }
+    buffer_snapshots = {
+        name: buf.detach().to(torch.float32).clone()
+        for name, buf in model.named_buffers(remove_duplicate=False)
+        if buf.is_floating_point() and any(keyword in name for keyword in buffer_keywords)
+    }
+    return parameter_snapshots, buffer_snapshots
+
+
+def _restore_fp32_tensor_snapshots(
+    model: nn.Module,
+    *,
+    parameter_snapshots: dict[str, torch.Tensor],
+    buffer_snapshots: dict[str, torch.Tensor],
+) -> None:
+    """Restore fp32-preserved tensors from pre-cast snapshots."""
+    named_parameters = dict(model.named_parameters())
+    for name, snapshot in parameter_snapshots.items():
+        param = named_parameters.get(name)
+        if param is None:
+            continue
+        param.data = snapshot.to(dtype=torch.float32)
+
+    for name, snapshot in buffer_snapshots.items():
+        module_name, _, buffer_name = name.rpartition(".")
+        module = model.get_submodule(module_name) if module_name else model
+        if buffer_name not in module._buffers:
+            continue
+        module._buffers[buffer_name] = snapshot.to(dtype=torch.float32)
+
+
 def _restore_fp32_modules(model: nn.Module, fp32_keywords: list[str]) -> None:
     """Cast modules or individual tensors matching *fp32_keywords* back to float32.
 

nemo_automodel/components/models/nemotron_v3/model.py

@@ -62,6 +62,8 @@ class NemotronV3Model(nn.Module):
     This is a hybrid architecture with Mamba2, Attention, MLP, and MoE layers.
     """
 
+    _keep_in_fp32_modules_strict = ["e_score_correction_bias"]
+
     def __init__(
         self,
         config,
@@ -283,6 +285,7 @@ class NemotronHForCausalLM(HFCheckpointingMixin, GenerationMixin, nn.Module, MoE
     # Hybrid Mamba2/Attention uses NemotronHybridCache, not DynamicCache.
     _is_stateful: bool = True
     main_input_name: str = "input_ids"
+    _keep_in_fp32_modules_strict = ["e_score_correction_bias"]
 
     # Skip patch_hf_model_for_pp; our forward already handles PP routing.
     _pp_keep_self_forward: bool = True

nemo_automodel/components/models/qwen3_5/model.py

@@ -75,6 +75,19 @@ def _default_init_device() -> torch.device:
     return torch.device("cpu")
 
 
+def _qwen3_5_backend(backend: BackendConfig | None = None) -> BackendConfig:
+    """Return a Qwen3.5 backend with TE fused RoPE disabled.
+
+    Qwen3.5 VLM training can feed full-attention layers in packed/THD shape via
+    the shared Qwen3-Next attention block. TE fused RoPE expects 4D inputs there,
+    so keep the non-fused RoPE path while preserving the rest of the backend
+    selection (TE Linear, attention backend, etc.).
+    """
+    resolved = copy.copy(backend) if backend is not None else BackendConfig()
+    resolved.rope_fusion = False
+    return resolved
+
+
 def build_mtp_config_from_hf(
     config: Any,
     *,
@@ -592,7 +605,7 @@ def __init__(
         super().__init__()
         del kwargs
         self.config = config
-        self.backend = backend or BackendConfig()
+        self.backend = _qwen3_5_backend(backend)
 
         self.model = Qwen3_5DenseTextBackbone(config, self.backend)
         dtype = next(self.model.parameters()).dtype
@@ -788,7 +801,7 @@ def __init__(
     ) -> None:
         del kwargs
         super().__init__(config)
-        self.backend = backend or BackendConfig()
+        self.backend = _qwen3_5_backend(backend)
 
         text_config = config.text_config
         # Replace the HF text decoder with the native NeMo backbone (built on the