From 907e47f05c6b49fddf7c079df782f9fcffef74dd Mon Sep 17 00:00:00 2001 From: Yuhe Zhang Date: Tue, 16 Jun 2026 16:52:01 -0400 Subject: [PATCH] fix(moe): preserve fp32 A_log in Qwen3.5-MoE and Qwen3-Next GatedDeltaNet (#2484) * fix(moe): preserve fp32 A_log in Qwen3.5-MoE and Qwen3-Next GatedDeltaNet GatedDeltaNet A_log/dt_bias are intrinsically fp32, but the native MoE Qwen3.5-MoE and Qwen3-Next models sharded through moe/parallelizer.py with a single bf16 mixed-precision policy, silently rounding these params to bf16 and degrading the exponentiated decay rate. Generalize the dense Qwen3.5 _fp32_params holder mechanism into a shared helper (components/models/common/gated_delta_net_fp32.py), isolate the fp32 GDN params into per-module holders at parallelize time, and shard each holder as a dedicated fp32 FSDP unit in apply_fsdp (preserving EP ignored_params). Add an fp32-aware Qwen3NextGatedDeltaNet subclass that routes the gate through the holder, and strip the _fp32_params prefix on save so checkpoints stay HF-clean. Signed-off-by: Yuhe Zhang * fix(moe): avoid pre-reading holder-owned dt_bias in Qwen3.5-MoE gate CPAwareGatedDeltaNet._compute_gate evaluated self.dt_bias before entering the holder forward. Once isolate_fp32_params moves an intrinsically-fp32 dt_bias into the _fp32_params holder, the patched __getattr__ makes self.dt_bias resolve to the holder-owned param, reading it outside the holder forward (before FSDP unshards), which defeats the fp32 isolation. The holder forward already prefers its own dt_bias, so route through holder(a) when it owns dt_bias and only pass self.dt_bias when dt_bias stays a bare (bf16) param -- mirroring Qwen3-Next. Signed-off-by: Yuhe Zhang * test(moe): cover apply_fsdp fp32 GDN branch; exclude GPU-only Next forward Add a CPU unit test that drives the apply_fsdp fp32 branch (isolation reported, fp32 holder policy built, per-block holder sharded) to lift Codecov patch coverage reliably on every PR. Mark Qwen3NextFp32GatedDeltaNet.forward with ``# pragma: no cover`` since it is a verbatim HF GDN forward that requires CUDA conv1d/FLA kernels and is only exercised by GPU/functional runs. Signed-off-by: Yuhe Zhang * fix(qwen): preserve GDN fp32 params under FSDP2 Signed-off-by: Yuhe Zhang * fix(qwen): scope GDN fp32 isolation to linear attention Signed-off-by: Yuhe Zhang * fix(qwen): mark fp32 GDN modules explicitly Signed-off-by: Yuhe Zhang * docs(model): clarify GDN fp32 marker contract Signed-off-by: Yuhe Zhang * fix(nemotron): keep router correction bias fp32 Signed-off-by: Yuhe Zhang * fix(dtype): preserve fp32 values during model casts Signed-off-by: Yuhe Zhang * refactor(models): make qwen-next own fp32 gdn holder Signed-off-by: Yuhe Zhang * docs(skills): clarify fp32 holder ownership Signed-off-by: Yuhe Zhang * fix(checkpoint): tolerate missing TE extra state Signed-off-by: Yuhe Zhang * fix(models): preserve duplicate fp32 buffers Signed-off-by: Yuhe Zhang --------- Signed-off-by: Yuhe Zhang Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com> (cherry picked from commit 174c5d15536e3461614ce36227f84cb71f79f8b8) Signed-off-by: Alexandros Koumparoulis --- nemo_automodel/_transformers/model_init.py | 27 ++- .../models/common/gated_delta_net_fp32.py | 105 ++++++++++ .../components/models/common/utils.py | 112 +++++++++-- .../components/models/nemotron_v3/model.py | 3 + .../components/models/qwen3_5/model.py | 17 +- .../models/qwen3_5/state_dict_adapter.py | 27 ++- .../components/models/qwen3_5_moe/model.py | 14 +- .../models/qwen3_5_moe/state_dict_adapter.py | 50 +++-- .../components/models/qwen3_next/layers.py | 186 ++++++++++++++++++ .../components/models/qwen3_next/model.py | 17 +- .../models/qwen3_next/state_dict_adapter.py | 15 ++ nemo_automodel/components/moe/parallelizer.py | 14 +- nemo_automodel/recipes/vlm/finetune.py | 2 + .../nemo-automodel-model-onboarding/SKILL.md | 10 +- .../models/common/test_cast_model_to_dtype.py | 80 ++++++++ .../common/test_gated_delta_net_fp32.py | 84 ++++++++ .../nemotron_v3/test_nemotron_v3_model.py | 37 ++++ .../test_qwen3_5_state_dict_adapter.py | 73 +++++++ .../qwen3_5_moe/test_qwen3_5_moe_fp32_gdn.py | 58 ++++++ .../qwen3_5_moe/test_qwen3_5_moe_model.py | 13 +- .../test_qwen3_5_moe_state_dict_adapter.py | 62 ++++++ .../qwen3_next/test_qwen3_next_fp32_gdn.py | 104 ++++++++++ .../qwen3_next/test_qwen3_next_model.py | 6 +- .../test_qwen3_next_state_dict_adapter.py | 107 ++++++++-- tests/unit_tests/moe/test_parallelizer.py | 80 ++++++++ .../test_improved_error_messages.py | 76 +++++++ 26 files changed, 1296 insertions(+), 83 deletions(-) create mode 100644 nemo_automodel/components/models/common/gated_delta_net_fp32.py create mode 100644 tests/unit_tests/models/common/test_gated_delta_net_fp32.py create mode 100644 tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_fp32_gdn.py create mode 100644 tests/unit_tests/models/qwen3_next/test_qwen3_next_fp32_gdn.py diff --git a/nemo_automodel/_transformers/model_init.py b/nemo_automodel/_transformers/model_init.py index efea4f917e..03d34ded7d 100644 --- a/nemo_automodel/_transformers/model_init.py +++ b/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,6 +666,7 @@ 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(): @@ -669,22 +674,34 @@ def _restore_loaded_model_dtype( 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 diff --git a/nemo_automodel/components/models/common/gated_delta_net_fp32.py b/nemo_automodel/components/models/common/gated_delta_net_fp32.py new file mode 100644 index 0000000000..5d7225e9ed --- /dev/null +++ b/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 diff --git a/nemo_automodel/components/models/common/utils.py b/nemo_automodel/components/models/common/utils.py index 0970dba124..3cfca53ebc 100644 --- a/nemo_automodel/components/models/common/utils.py +++ b/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. diff --git a/nemo_automodel/components/models/nemotron_v3/model.py b/nemo_automodel/components/models/nemotron_v3/model.py index 015991acdc..db20a6c361 100644 --- a/nemo_automodel/components/models/nemotron_v3/model.py +++ b/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 diff --git a/nemo_automodel/components/models/qwen3_5/model.py b/nemo_automodel/components/models/qwen3_5/model.py index bd2973ffe4..d45eae8fdb 100644 --- a/nemo_automodel/components/models/qwen3_5/model.py +++ b/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 diff --git a/nemo_automodel/components/models/qwen3_5/state_dict_adapter.py b/nemo_automodel/components/models/qwen3_5/state_dict_adapter.py index 80cc140d89..f2db66743a 100644 --- a/nemo_automodel/components/models/qwen3_5/state_dict_adapter.py +++ b/nemo_automodel/components/models/qwen3_5/state_dict_adapter.py @@ -14,13 +14,10 @@ """State-dict adapter for Qwen3.5 dense (non-MoE) models. -Qwen3.5 dense uses HF's GatedDeltaNet linear-attention layers. For FSDP -compatibility (mixed-dtype: bf16 + fp32 ``A_log``), ``patch_hf_model`` in -``cp_linear_attn`` moves ``A_log`` from ``mod._parameters`` into a -``_fp32_params`` submodule and patches ``__getattr__`` to redirect -``mod.A_log`` reads. After patching, the model's state_dict contains keys of -the form ``...linear_attn._fp32_params.A_log`` instead of the original -``...linear_attn.A_log``. +Qwen3.5 dense keeps its GatedDeltaNet SSM-gating parameters (``A_log`` / +``dt_bias``) in a fp32 ``_fp32_params`` holder. The model's state dict therefore +contains keys of the form ``...linear_attn._fp32_params.A_log`` instead of the +original ``...linear_attn.A_log``. This adapter renames keys at save/load boundaries so that on-disk checkpoints match the original HF Qwen3.5 layout (bare ``A_log``) and are directly @@ -33,6 +30,7 @@ from typing import Any, Optional from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter +from nemo_automodel.components.models.common.gated_delta_net_fp32 import upcast_gated_delta_net_fp32_state_tensor _FP32_PARAMS_TO_BARE = re.compile(r"(\.linear_attn)\._fp32_params\.") # Both SSM-gating params live in the fp32 ``SSMGate`` holder; route both on load. @@ -78,7 +76,11 @@ def __init__(self, *, route_linear_attn_fp32_params: bool = True) -> None: self.route_linear_attn_fp32_params = route_linear_attn_fp32_params def to_hf(self, state_dict: dict[str, Any], **kwargs: Any) -> dict[str, Any]: - return {map_qwen3_5_mtp_to_hf_key(_strip_fp32_prefix(k)): v for k, v in state_dict.items()} + hf_state_dict: dict[str, Any] = {} + for key, value in state_dict.items(): + hf_key = map_qwen3_5_mtp_to_hf_key(_strip_fp32_prefix(key)) + hf_state_dict[hf_key] = upcast_gated_delta_net_fp32_state_tensor(hf_key, value) + return hf_state_dict def from_hf( self, @@ -87,10 +89,15 @@ def from_hf( **kwargs: Any, ) -> dict[str, Any]: del device_mesh, kwargs - return {self._map_from_hf_key(k): v for k, v in hf_state_dict.items()} + native_state_dict: dict[str, Any] = {} + for key, value in hf_state_dict.items(): + native_key = self._map_from_hf_key(key) + native_state_dict[native_key] = upcast_gated_delta_net_fp32_state_tensor(native_key, value) + return native_state_dict def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs: Any) -> list[tuple[str, Any]]: - return [(map_qwen3_5_mtp_to_hf_key(_strip_fp32_prefix(fqn)), tensor)] + hf_key = map_qwen3_5_mtp_to_hf_key(_strip_fp32_prefix(fqn)) + return [(hf_key, upcast_gated_delta_net_fp32_state_tensor(hf_key, tensor))] def _map_from_hf_key(self, key: str) -> str: key = map_qwen3_5_mtp_from_hf_key(key) diff --git a/nemo_automodel/components/models/qwen3_5_moe/model.py b/nemo_automodel/components/models/qwen3_5_moe/model.py index 6a86217c24..fa9574c1a0 100644 --- a/nemo_automodel/components/models/qwen3_5_moe/model.py +++ b/nemo_automodel/components/models/qwen3_5_moe/model.py @@ -202,6 +202,18 @@ def _default_init_device() -> torch.device: return torch.device("cpu") +def _qwen3_5_moe_backend(backend: BackendConfig | None = None) -> BackendConfig: + """Return a Qwen3.5-MoE backend with TE fused RoPE disabled. + + The Qwen3.5 full-attention blocks reuse Qwen3-Next attention, and VLM/packed + execution can present THD-shaped q/k tensors. TE fused RoPE expects 4D inputs + in this path, so use non-fused RoPE while preserving the rest of the backend. + """ + 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, *, @@ -726,7 +738,7 @@ def __init__( ): if not _QWEN3_5_MOE_HF_AVAILABLE: raise UnavailableError("transformers.models.qwen3_5_moe is not available.") - backend = backend or BackendConfig() + backend = _qwen3_5_moe_backend(backend) # _init_model() only overrides the top-level hf_config.torch_dtype; for # VL configs the nested text_config / vision_config keep their original diff --git a/nemo_automodel/components/models/qwen3_5_moe/state_dict_adapter.py b/nemo_automodel/components/models/qwen3_5_moe/state_dict_adapter.py index 66e932d1f4..0ccf2dd990 100644 --- a/nemo_automodel/components/models/qwen3_5_moe/state_dict_adapter.py +++ b/nemo_automodel/components/models/qwen3_5_moe/state_dict_adapter.py @@ -45,6 +45,11 @@ from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter from nemo_automodel.components.models.common import BackendConfig +from nemo_automodel.components.models.common.gated_delta_net_fp32 import ( + route_fp32_holder_key, + strip_fp32_holder_key, + upcast_gated_delta_net_fp32_state_tensor, +) from nemo_automodel.components.models.qwen3_5.state_dict_adapter import ( map_qwen3_5_mtp_from_hf_key, map_qwen3_5_mtp_to_hf_key, @@ -52,23 +57,15 @@ from nemo_automodel.components.moe import state_dict_utils from nemo_automodel.components.moe.layers import MoEConfig -# The SSM-gating params (A_log/dt_bias) live in fp32 storage inside a -# ``linear_attn._fp32_params`` holder submodule. Checkpoints stay in the bare HF -# key namespace (``linear_attn.A_log``), so strip the holder segment on save and -# route bare keys back into the holder on load. -_FP32_PARAMS_RE = re.compile(r"(\.linear_attn)\._fp32_params\.") -_FP32_HOLDER_PARAM_NAMES = ("A_log", "dt_bias") - def _strip_fp32_params(key: str) -> str: - return _FP32_PARAMS_RE.sub(r"\1.", key) + """Strip the fp32 holder segment from GDN state-dict keys.""" + return strip_fp32_holder_key(key) def _route_fp32_params(key: str) -> str: - if not key.endswith(_FP32_HOLDER_PARAM_NAMES) or "._fp32_params." in key or ".linear_attn." not in key: - return key - head, tail = key.rsplit(".linear_attn.", 1) - return f"{head}.linear_attn._fp32_params.{tail}" + """Route bare GDN fp32 params into the holder used by the native module.""" + return route_fp32_holder_key(key) class Qwen3_5MoeStateDictAdapter(StateDictAdapter): @@ -172,10 +169,15 @@ def from_hf( state_dict: dict[str, Any] = {} mtp_expert_parts: dict[str, dict[str, dict[int, torch.Tensor]]] = {} + + def store_native_key(native_key: str, tensor: Any) -> None: + native_key = route_fp32_holder_key(native_key) + state_dict[native_key] = upcast_gated_delta_net_fp32_state_tensor(native_key, tensor) + for key, value in hf_state_dict.items(): mapped_mtp_key = map_qwen3_5_mtp_from_hf_key(key) if mapped_mtp_key != key: - state_dict[mapped_mtp_key] = value + store_native_key(mapped_mtp_key, value) continue match = re.match( @@ -232,15 +234,17 @@ def from_hf( break if mapped_key.startswith("mtp."): - state_dict[mapped_key] = value + store_native_key(mapped_key, value) elif mapped_key.startswith("model.lm_head."): - state_dict[mapped_key.removeprefix("model.")] = value + store_native_key(mapped_key.removeprefix("model."), value) elif mapped_key.startswith("lm_head."): - state_dict[mapped_key] = value + store_native_key(mapped_key, value) elif key.startswith("model."): - state_dict[mapped_key] = value + store_native_key(mapped_key, value) else: - state_dict[f"{model_prefix}{mapped_key}" if not mapped_key.startswith("model.") else mapped_key] = value + store_native_key( + f"{model_prefix}{mapped_key}" if not mapped_key.startswith("model.") else mapped_key, value + ) for layer_num, parts in mtp_expert_parts.items(): expert_ids = sorted(set(parts["gate_proj"]) | set(parts["up_proj"]) | set(parts["down_proj"])) @@ -278,8 +282,7 @@ def from_hf( down_tensor, device_mesh, rank ) - # Route bare SSM-gating keys into the linear_attn ``_fp32_params`` holder. - return {_route_fp32_params(k): v for k, v in state_dict.items()} + return state_dict def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[tuple[str, Any]]: """Rename a single native key to HF format and transpose expert tensors.""" @@ -307,8 +310,13 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t new_fqn = new_fqn.replace(pattern, replacement) break + # Hide the GatedDeltaNet fp32 holder wrapping so saved checkpoints keep the + # bare HF key (``...linear_attn.A_log`` instead of + # ``...linear_attn._fp32_params.A_log``) and stay directly HF-loadable. + new_fqn = strip_fp32_holder_key(new_fqn) + new_fqn = map_qwen3_5_mtp_to_hf_key(new_fqn) - new_fqn = _strip_fp32_params(new_fqn) + value = upcast_gated_delta_net_fp32_state_tensor(new_fqn, value) if exclude_key_regex and re.match(exclude_key_regex, new_fqn): return [] diff --git a/nemo_automodel/components/models/qwen3_next/layers.py b/nemo_automodel/components/models/qwen3_next/layers.py index ec4c64c7e8..528635ca0e 100644 --- a/nemo_automodel/components/models/qwen3_next/layers.py +++ b/nemo_automodel/components/models/qwen3_next/layers.py @@ -15,8 +15,10 @@ from typing import Any import torch +import torch.nn.functional as F from torch import nn from transformers.models.qwen3_next.configuration_qwen3_next import Qwen3NextConfig +from transformers.models.qwen3_next.modeling_qwen3_next import Qwen3NextGatedDeltaNet from nemo_automodel.components.attention.utils import ( initialize_attn_module_and_func, @@ -31,6 +33,190 @@ from nemo_automodel.shared.utils import dtype_from_str as get_dtype +class _SSMGateParam: + """Get-only descriptor exposing a param from ``_fp32_params`` when present.""" + + def __init__(self, name: str): + self.name = name + + def __get__(self, obj, owner=None): + if obj is None: + return self + holder = obj._modules.get("_fp32_params") + if holder is not None: + return getattr(holder, self.name) + param = obj._parameters.get(self.name) + if param is not None: + return param + raise AttributeError(f"{type(obj).__name__!s} has no parameter {self.name!r}") + + +class Qwen3NextSSMGate(nn.Module): + """Owns Qwen3-Next fp32 SSM-gating params and computes the decay gate.""" + + def __init__(self, num_v_heads: int, dtype: torch.dtype = torch.float32): + super().__init__() + self.A_log = nn.Parameter(torch.empty(num_v_heads, dtype=dtype)) + self.dt_bias = nn.Parameter(torch.empty(num_v_heads, dtype=dtype)) + + def forward(self, a: torch.Tensor) -> torch.Tensor: + return -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) + + +def _install_ssm_gate(mod: nn.Module, fp32_dtype: torch.dtype = torch.float32) -> Qwen3NextSSMGate: + """Move HF-created bare ``A_log``/``dt_bias`` into a native fp32 holder.""" + num_v_heads = mod._parameters["A_log"].shape[0] + gate = Qwen3NextSSMGate(num_v_heads, dtype=fp32_dtype) + for pname in ("A_log", "dt_bias"): + param = mod._parameters.pop(pname) + if param.dtype != fp32_dtype: + param.data = param.data.to(fp32_dtype) + setattr(gate, pname, param) + mod.add_module("_fp32_params", gate) + return gate + + +class Qwen3NextFp32GatedDeltaNet(Qwen3NextGatedDeltaNet): + """Qwen3-Next GatedDeltaNet that computes the decay gate via an fp32 holder. + + HF's ``Qwen3NextGatedDeltaNet`` computes the gate inline as + ``g = -exp(A_log) * softplus(a + dt_bias)`` using the bare ``A_log`` / ``dt_bias`` + parameters. ``A_log`` and ``dt_bias`` are intrinsically fp32 (``A_log`` is + exponentiated, so bf16 rounding becomes a proportional error on the decay rate that + the recurrence compounds across the sequence). + + The constructor moves those params into a native ``_fp32_params`` holder so they + are fp32 resident before any dtype cast or FSDP wrapping. To keep the gate + computation in fp32 -- and to make FSDP's unshard/reshard + gradient + reduce-scatter fire for that unit -- the gate is computed inside the holder's + forward. This subclass overrides ``forward`` to route the gate through + ``self._compute_gate(a)`` while reproducing the rest of HF's forward verbatim. + """ + + A_log = _SSMGateParam("A_log") + dt_bias = _SSMGateParam("dt_bias") + + def __init__(self, config: Qwen3NextConfig, layer_idx: int): + super().__init__(config, layer_idx) + _install_ssm_gate(self) + + def _compute_gate(self, a: torch.Tensor) -> torch.Tensor: + """Compute the decay gate ``g`` in fp32, via the holder when it exists.""" + holder = self._modules.get("_fp32_params") + if holder is not None: + return holder(a) + return -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) + + def forward( # pragma: no cover - verbatim HF GDN forward; needs CUDA conv1d/FLA kernels (GPU/functional only) + self, + hidden_states: torch.Tensor, + cache_params: Any | None = None, + attention_mask: torch.Tensor | None = None, + ): + # Mirrors transformers ``Qwen3NextGatedDeltaNet.forward`` with the gate routed + # through ``self._compute_gate(a)`` so A_log/dt_bias stay fp32 under FSDP. + from transformers.models.qwen3_next.modeling_qwen3_next import apply_mask_to_padding_states + + hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask) + + batch_size, seq_len, _ = hidden_states.shape + + use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + + if use_precomputed_states: + conv_state = cache_params.layers[self.layer_idx].conv_states + recurrent_state = cache_params.layers[self.layer_idx].recurrent_states + + projected_states_qkvz = self.in_proj_qkvz(hidden_states) + projected_states_ba = self.in_proj_ba(hidden_states) + query, key, value, z, b, a = self.fix_query_key_value_ordering(projected_states_qkvz, projected_states_ba) + query, key, value = (x.reshape(x.shape[0], x.shape[1], -1) for x in (query, key, value)) + + mixed_qkv = torch.cat((query, key, value), dim=-1) + mixed_qkv = mixed_qkv.transpose(1, 2) + + if use_precomputed_states and seq_len == 1: + mixed_qkv = self.causal_conv1d_update( + mixed_qkv, + conv_state, + self.conv1d.weight.squeeze(1), + self.conv1d.bias, + self.activation, + ) + else: + if use_precomputed_states: + mixed_qkv = torch.cat([conv_state, mixed_qkv], dim=-1) + if cache_params is not None: + new_conv_state = F.pad(mixed_qkv, (self.conv_kernel_size - mixed_qkv.shape[-1], 0)) + cache_params.update_conv_state(new_conv_state, self.layer_idx) + if self.causal_conv1d_fn is not None: + mixed_qkv = self.causal_conv1d_fn( + x=mixed_qkv, + weight=self.conv1d.weight.squeeze(1), + bias=self.conv1d.bias, + activation=self.activation, + seq_idx=None, + ) + else: + mixed_qkv = F.silu(self.conv1d(mixed_qkv)[:, :, : mixed_qkv.shape[-1]]) + if use_precomputed_states: + mixed_qkv = mixed_qkv[:, :, -seq_len:] + + mixed_qkv = mixed_qkv.transpose(1, 2) + query, key, value = torch.split( + mixed_qkv, + [self.key_dim, self.key_dim, self.value_dim], + dim=-1, + ) + query = query.reshape(query.shape[0], query.shape[1], -1, self.head_k_dim) + key = key.reshape(key.shape[0], key.shape[1], -1, self.head_k_dim) + value = value.reshape(value.shape[0], value.shape[1], -1, self.head_v_dim) + + beta = b.sigmoid() + # Gate is computed in fp32 (via the _fp32_params holder when present) so the + # exponentiated decay rate keeps full precision under bf16 compute. + g = self._compute_gate(a) + if self.num_v_heads // self.num_k_heads > 1: + query = query.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) + key = key.repeat_interleave(self.num_v_heads // self.num_k_heads, dim=2) + + if use_precomputed_states and seq_len == 1: + core_attn_out, last_recurrent_state = self.recurrent_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=recurrent_state, + output_final_state=cache_params is not None, + use_qk_l2norm_in_kernel=True, + ) + else: + core_attn_out, last_recurrent_state = self.chunk_gated_delta_rule( + query, + key, + value, + g=g, + beta=beta, + initial_state=recurrent_state if use_precomputed_states else None, + output_final_state=cache_params is not None, + use_qk_l2norm_in_kernel=True, + ) + + if cache_params is not None: + cache_params.update_recurrent_state(last_recurrent_state, self.layer_idx) + + z_shape_og = z.shape + core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1]) + z = z.reshape(-1, z.shape[-1]) + core_attn_out = self.norm(core_attn_out, z) + core_attn_out = core_attn_out.reshape(z_shape_og) + core_attn_out = core_attn_out.reshape(core_attn_out.shape[0], core_attn_out.shape[1], -1) + + output = self.out_proj(core_attn_out) + return output + + class Qwen3NextRMSNorm(nn.Module): def __init__(self, dim: int, eps: float = 1e-6): super().__init__() diff --git a/nemo_automodel/components/models/qwen3_next/model.py b/nemo_automodel/components/models/qwen3_next/model.py index 5fc1fd015a..d9e8c6be15 100644 --- a/nemo_automodel/components/models/qwen3_next/model.py +++ b/nemo_automodel/components/models/qwen3_next/model.py @@ -19,7 +19,6 @@ import torch.nn as nn from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.models.qwen3_next.configuration_qwen3_next import Qwen3NextConfig -from transformers.models.qwen3_next.modeling_qwen3_next import Qwen3NextGatedDeltaNet from nemo_automodel.components.models.common import ( BackendConfig, @@ -30,7 +29,11 @@ from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin from nemo_automodel.components.models.common.utils import cast_model_to_dtype, compute_lm_head_logits from nemo_automodel.components.models.gpt_oss.rope_utils import RotaryEmbedding, position_ids_to_freqs_cis -from nemo_automodel.components.models.qwen3_next.layers import Qwen3NextAttention, Qwen3NextRMSNorm +from nemo_automodel.components.models.qwen3_next.layers import ( + Qwen3NextAttention, + Qwen3NextFp32GatedDeltaNet, + Qwen3NextRMSNorm, +) from nemo_automodel.components.models.qwen3_next.state_dict_adapter import Qwen3NextStateDictAdapter from nemo_automodel.components.moe.config import MoEConfig from nemo_automodel.components.moe.fsdp_mixin import MoEFSDPSyncMixin @@ -44,7 +47,9 @@ def __init__(self, layer_idx: int, config: Qwen3NextConfig, moe_config: MoEConfi super().__init__() self.layer_type = config.layer_types[layer_idx] if self.layer_type == "linear_attention": - self.linear_attn = Qwen3NextGatedDeltaNet(config, layer_idx) + # fp32-aware GatedDeltaNet: keeps the intrinsically-fp32 A_log/dt_bias decay + # gate in fp32 under FSDP mixed precision (see Qwen3NextFp32GatedDeltaNet). + self.linear_attn = Qwen3NextFp32GatedDeltaNet(config, layer_idx) elif self.layer_type == "full_attention": self.self_attn = Qwen3NextAttention(config, layer_idx, backend) @@ -301,6 +306,10 @@ def __init__( self.lm_head = initialize_linear_module( self.backend.linear, config.hidden_size, config.vocab_size, bias=False, dtype=model_dtype ) + keep_fp32 = list(getattr(self, "_keep_in_fp32_modules", None) or []) + if "_fp32_params" not in keep_fp32: + keep_fp32.append("_fp32_params") + self._keep_in_fp32_modules = keep_fp32 if self.backend.enable_hf_state_dict_adapter: self.state_dict_adapter = Qwen3NextStateDictAdapter( self.config, self.model.moe_config, self.backend, dtype=model_dtype @@ -372,7 +381,7 @@ def initialize_weights( b=cutoff_factor * final_out_std, ) - cast_model_to_dtype(self, dtype) + cast_model_to_dtype(self, dtype, skip_modules=("_fp32_params",)) with buffer_device: # Ensure rotary embedding uses correct device after dtype move self.model.rotary_emb.device = buffer_device diff --git a/nemo_automodel/components/models/qwen3_next/state_dict_adapter.py b/nemo_automodel/components/models/qwen3_next/state_dict_adapter.py index a054e63250..ba1d76e032 100644 --- a/nemo_automodel/components/models/qwen3_next/state_dict_adapter.py +++ b/nemo_automodel/components/models/qwen3_next/state_dict_adapter.py @@ -20,6 +20,11 @@ from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter from nemo_automodel.components.models.common import BackendConfig +from nemo_automodel.components.models.common.gated_delta_net_fp32 import ( + route_fp32_holder_key, + strip_fp32_holder_key, + upcast_gated_delta_net_fp32_state_tensor, +) from nemo_automodel.components.moe.config import MoEConfig from nemo_automodel.components.moe.state_dict_mixin import MoESplitExpertsStateDictMixin @@ -117,6 +122,11 @@ def from_hf( # First apply key mappings for shared experts (shared_expert -> shared_experts) hf_state_dict = self._apply_key_mapping(hf_state_dict, self.hf_to_internal_map) + fp32_routed_state_dict = {} + for key, value in hf_state_dict.items(): + native_key = route_fp32_holder_key(key) + fp32_routed_state_dict[native_key] = upcast_gated_delta_net_fp32_state_tensor(native_key, value) + hf_state_dict = fp32_routed_state_dict # Then convert routed experts from split to grouped format return self._from_hf_w_merged_experts(hf_state_dict, device_mesh) @@ -147,6 +157,11 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t if pattern in key: new_key = new_key.replace(pattern, replacement) break + # Hide the GatedDeltaNet fp32 holder wrapping so saved checkpoints keep the + # bare HF key (``...linear_attn.A_log`` instead of + # ``...linear_attn._fp32_params.A_log``) and stay directly HF-loadable. + new_key = strip_fp32_holder_key(new_key) + value = upcast_gated_delta_net_fp32_state_tensor(new_key, value) mapped_result.append((new_key, value)) if exclude_key_regex: diff --git a/nemo_automodel/components/moe/parallelizer.py b/nemo_automodel/components/moe/parallelizer.py index 8d9dc42daa..6a19311926 100644 --- a/nemo_automodel/components/moe/parallelizer.py +++ b/nemo_automodel/components/moe/parallelizer.py @@ -297,10 +297,10 @@ def selective_checkpointing_context_fn(): def _shard_fp32_param_holders(block, fsdp_mesh, reshard_after_forward, offload_policy): """Shard each ``_fp32_params`` holder in ``block`` as its own fp32 FSDP unit. - Qwen3.5 GatedDeltaNet keeps its SSM-gating params (A_log/dt_bias) in fp32 - storage, isolated in a ``_fp32_params`` holder submodule. FSDP2 requires a - dtype-uniform parameter group, so these are sharded separately with an fp32 - mixed-precision policy and excluded from the block's (bf16) FSDP unit. + Model implementations own the architecture-specific decision to create these + holders (for example Qwen3.5/Qwen3-Next GatedDeltaNet ``A_log``/``dt_bias``). + FSDP only treats the holder as a dtype-uniform fp32 unit and excludes its params + from the block's bf16 FSDP unit. Returns the set of holder parameters to exclude from the block's FSDP wrap. Blocks that do not expose ``named_modules`` (e.g. non-``nn.Module`` test @@ -404,10 +404,8 @@ def apply_fsdp( if isinstance(moe_module, MoE) and ep_enabled: ignored_params = set(moe_module.experts.parameters()) - # Isolate the linear_attn SSM-gating params (A_log/dt_bias), which are kept - # in fp32 storage, into their own fp32 FSDP group so the rest of the block - # stays dtype-uniform (FSDP2 requires uniform dtype within a group). Shard - # the holder on its own and exclude its params from the block's FSDP unit. + # Shard model-owned fp32 holders on their own and exclude their params from + # the block's FSDP unit to keep the block dtype-uniform. fp32_ignored = _shard_fp32_param_holders(block, fsdp_mesh, reshard_after_forward, offload_policy) if fp32_ignored: ignored_params = (ignored_params or set()) | fp32_ignored diff --git a/nemo_automodel/recipes/vlm/finetune.py b/nemo_automodel/recipes/vlm/finetune.py index 8d3b5cbb3b..7bd6c11312 100644 --- a/nemo_automodel/recipes/vlm/finetune.py +++ b/nemo_automodel/recipes/vlm/finetune.py @@ -81,6 +81,7 @@ from nemo_automodel.recipes._dist_utils import create_distributed_setup_from_config, shard_optimizers_for_megatron_fsdp from nemo_automodel.recipes._typed_config import RecipeConfig from nemo_automodel.recipes.base_recipe import BaseRecipe +from nemo_automodel.shared.te_patches import apply_te_patches if TYPE_CHECKING: from torch.optim import Optimizer @@ -613,6 +614,7 @@ def setup(self): distributed_setup=self.distributed_setup, cfg_quantization=self.cfg.get("quantization", None), ) + apply_te_patches() optimizer = self.cfg.optimizer.build(model, device_mesh=self.device_mesh, is_peft=self.peft_config is not None) allow_megatron_fsdp_sharding = getattr(self.cfg.optimizer, "supports_megatron_fsdp_sharding", True) self.optimizer = shard_optimizers_for_megatron_fsdp( diff --git a/skills/nemo-automodel-model-onboarding/SKILL.md b/skills/nemo-automodel-model-onboarding/SKILL.md index 54fc895315..6cf5de7855 100644 --- a/skills/nemo-automodel-model-onboarding/SKILL.md +++ b/skills/nemo-automodel-model-onboarding/SKILL.md @@ -362,8 +362,14 @@ Where to declare it: - **NeMo-native model class** (you own `model.py`): a class attribute, e.g. `_keep_in_fp32_modules_strict = ["e_score_correction_bias"]` (see `deepseek_v4`, `ling_v2`). -- **HF model you only patch** (e.g. Qwen3.5): set it on the instance inside `patch_hf_model`, - e.g. `model._keep_in_fp32_modules_strict = existing + ("_fp32_params",)`. +- **Custom/HF-derived model class with fp32 runtime params**: build the fp32 structure in the + model or layer constructor. For GatedDeltaNet-style `A_log` / `dt_bias`, move them into a + real `_fp32_params` holder during construction, compute the sensitive gate inside that + holder's `forward`, keep the holder out of broad dtype casts with + `cast_model_to_dtype(..., skip_modules=("_fp32_params",))`, and make the state-dict adapter + strip/route holder keys plus upcast loaded tensors to fp32. Do not use a runtime monkeypatch, + and do not infer the contract globally from a module path such as `linear_attn` or from an + `A_log` parameter name alone. Always declare the pin for these params. A normal checkpoint load also auto-records each param's original HF dtype and uses it as a fallback, but that recording is skipped on the diff --git a/tests/unit_tests/models/common/test_cast_model_to_dtype.py b/tests/unit_tests/models/common/test_cast_model_to_dtype.py index 0e3a425291..987a2b5e0e 100644 --- a/tests/unit_tests/models/common/test_cast_model_to_dtype.py +++ b/tests/unit_tests/models/common/test_cast_model_to_dtype.py @@ -19,6 +19,7 @@ from nemo_automodel.components.models.common.utils import ( _get_fp32_module_keywords, + _get_strict_fp32_module_keywords, _has_dtensor_params, _restore_fp32_buffers, _restore_fp32_modules, @@ -76,6 +77,21 @@ def __init__(self): self.mixer.scale = nn.Parameter(torch.ones(4)) +class ModelWithStrictFp32Buffer(nn.Module): + """Model that declares one strict fp32 buffer by qualified buffer name.""" + + _keep_in_fp32_modules_strict = ["router.e_score_correction_bias"] + + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4) + self.router = nn.Module() + self.router.register_buffer( + "e_score_correction_bias", + torch.tensor([1.001, -2.003, 0.3333, 17.125], dtype=torch.float32), + ) + + class ModelWithBothFp32Attrs(nn.Module): """Model with both _keep_in_fp32_modules and _keep_in_fp32_modules_strict.""" @@ -107,6 +123,28 @@ def test_keep_in_fp32_modules_strict(self): model = ModelWithStrictFp32() assert _get_fp32_module_keywords(model) == ["head"] + def test_keep_in_fp32_modules_strict_tuple(self): + class Model(nn.Module): + _keep_in_fp32_modules_strict = ("head",) + + def __init__(self): + super().__init__() + + model = Model() + assert _get_fp32_module_keywords(model) == ["head"] + assert _get_strict_fp32_module_keywords(model) == ["head"] + + def test_keep_in_fp32_modules_strict_set(self): + class Model(nn.Module): + _keep_in_fp32_modules_strict = {"head"} + + def __init__(self): + super().__init__() + + model = Model() + assert _get_fp32_module_keywords(model) == ["head"] + assert _get_strict_fp32_module_keywords(model) == ["head"] + def test_both_attributes_deduped(self): model = ModelWithBothFp32Attrs() keywords = _get_fp32_module_keywords(model) @@ -211,11 +249,43 @@ def test_strict_fp32_modules_preserved(self): assert model.head.weight.dtype == torch.float32 assert model.linear.weight.dtype == torch.bfloat16 + def test_strict_fp32_modules_preserved_from_tuple(self): + class Model(nn.Module): + _keep_in_fp32_modules_strict = ("head",) + + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4) + self.head = nn.Linear(4, 2) + + model = Model() + cast_model_to_dtype(model, torch.bfloat16) + + assert model.head.weight.dtype == torch.float32 + assert model.linear.weight.dtype == torch.bfloat16 + def test_strict_fp32_parameters_preserved(self): model = ModelWithStrictFp32Parameter() + original_scale = torch.tensor([1.001, -2.003, 0.3333, 17.125], dtype=torch.float32) + with torch.no_grad(): + model.mixer.scale.copy_(original_scale) + cast_model_to_dtype(model, torch.bfloat16) assert model.mixer.scale.dtype == torch.float32 + assert torch.equal(model.mixer.scale, original_scale) + assert not torch.equal(model.mixer.scale, original_scale.to(torch.bfloat16).float()) + assert model.linear.weight.dtype == torch.bfloat16 + + def test_strict_fp32_buffers_preserve_values(self): + model = ModelWithStrictFp32Buffer() + original_bias = model.router.e_score_correction_bias.clone() + + cast_model_to_dtype(model, torch.bfloat16) + + assert model.router.e_score_correction_bias.dtype == torch.float32 + assert torch.equal(model.router.e_score_correction_bias, original_bias) + assert not torch.equal(model.router.e_score_correction_bias, original_bias.to(torch.bfloat16).float()) assert model.linear.weight.dtype == torch.bfloat16 def test_both_fp32_attrs_preserved(self): @@ -324,6 +394,16 @@ def test_dtensor_params_only_buffers_restored(self): for p in model.parameters(): assert p.dtype == torch.bfloat16 + def test_dtensor_strict_fp32_params_restored(self): + """Strict fp32 modules are already isolated as their own FSDP units, so they can be restored.""" + model = ModelWithStrictFp32() + + with patch("nemo_automodel.components.models.common.utils._has_dtensor_params", return_value=True): + cast_model_to_dtype(model, torch.bfloat16) + + assert model.head.weight.dtype == torch.float32 + assert model.linear.weight.dtype == torch.bfloat16 + def test_dtensor_buffers_in_matching_modules_restored(self): """Buffers in fp32-keyword-matching modules are cast to fp32 even with DTensor params.""" diff --git a/tests/unit_tests/models/common/test_gated_delta_net_fp32.py b/tests/unit_tests/models/common/test_gated_delta_net_fp32.py new file mode 100644 index 0000000000..b080f3d7f3 --- /dev/null +++ b/tests/unit_tests/models/common/test_gated_delta_net_fp32.py @@ -0,0 +1,84 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for shared GatedDeltaNet fp32 checkpoint helpers.""" + +from __future__ import annotations + +import torch + +from nemo_automodel.components.models.common.gated_delta_net_fp32 import ( + FP32_GDN_PARAM_NAMES, + has_gated_delta_net_fp32_checkpoint_contract, + is_gated_delta_net_fp32_param_key, + route_fp32_holder_key, + strip_fp32_holder_key, + upcast_gated_delta_net_fp32_state_tensor, +) + + +def test_strip_fp32_holder_key(): + assert strip_fp32_holder_key("model.layers.0.linear_attn._fp32_params.A_log") == ( + "model.layers.0.linear_attn.A_log" + ) + assert strip_fp32_holder_key("model.layers.0.linear_attn._fp32_params.dt_bias") == ( + "model.layers.0.linear_attn.dt_bias" + ) + assert strip_fp32_holder_key("model.layers.0.linear_attn.A_log") == "model.layers.0.linear_attn.A_log" + assert strip_fp32_holder_key("model.layers.0.mlp.gate.weight") == "model.layers.0.mlp.gate.weight" + + +def test_route_fp32_holder_key(): + for name in FP32_GDN_PARAM_NAMES: + assert route_fp32_holder_key(f"model.layers.0.linear_attn.{name}") == ( + f"model.layers.0.linear_attn._fp32_params.{name}" + ) + + assert route_fp32_holder_key("model.layers.0.linear_attn._fp32_params.A_log") == ( + "model.layers.0.linear_attn._fp32_params.A_log" + ) + assert route_fp32_holder_key("model.layers.0.linear_attn.conv1d.weight") == ( + "model.layers.0.linear_attn.conv1d.weight" + ) + assert route_fp32_holder_key("model.layers.0.self_attn.A_log") == "model.layers.0.self_attn.A_log" + + +def test_strip_route_round_trip(): + bare = "model.layers.3.linear_attn.A_log" + routed = route_fp32_holder_key(bare) + assert routed == "model.layers.3.linear_attn._fp32_params.A_log" + assert strip_fp32_holder_key(routed) == bare + + +def test_is_gated_delta_net_fp32_param_key(): + assert is_gated_delta_net_fp32_param_key("model.layers.0.linear_attn.A_log") + assert is_gated_delta_net_fp32_param_key("model.layers.0.linear_attn._fp32_params.dt_bias") + assert not is_gated_delta_net_fp32_param_key("model.layers.0.self_attn.A_log") + assert not is_gated_delta_net_fp32_param_key("model.layers.0.linear_attn.conv1d.weight") + + +def test_has_gated_delta_net_fp32_checkpoint_contract(): + assert has_gated_delta_net_fp32_checkpoint_contract(type("Cfg", (), {"architectures": ["Qwen3NextForCausalLM"]})()) + assert has_gated_delta_net_fp32_checkpoint_contract( + type("Cfg", (), {"architectures": ["Qwen3_5MoeForConditionalGeneration"]})() + ) + assert not has_gated_delta_net_fp32_checkpoint_contract(type("Cfg", (), {"architectures": ["LlamaForCausalLM"]})()) + + +def test_upcast_gated_delta_net_fp32_state_tensor(): + tensor = torch.ones(4, dtype=torch.bfloat16) + out = upcast_gated_delta_net_fp32_state_tensor("model.layers.0.linear_attn._fp32_params.dt_bias", tensor) + assert out.dtype == torch.float32 + assert torch.equal(out, tensor.float()) + + +def test_upcast_gated_delta_net_fp32_state_tensor_leaves_unrelated_state(): + tensor = torch.ones(4, dtype=torch.bfloat16) + out = upcast_gated_delta_net_fp32_state_tensor("model.layers.0.self_attn.q_proj.weight", tensor) + assert out is tensor + + +def test_upcast_gated_delta_net_fp32_state_tensor_handles_bare_key(): + tensor = torch.ones(4, dtype=torch.bfloat16) + out = upcast_gated_delta_net_fp32_state_tensor("model.layers.0.linear_attn.A_log", tensor) + assert out.dtype == torch.float32 diff --git a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_model.py b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_model.py index 85c45aa31d..27836c1d6a 100644 --- a/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_model.py +++ b/tests/unit_tests/models/nemotron_v3/test_nemotron_v3_model.py @@ -1134,6 +1134,43 @@ def test_moe_model_init(self, config, backend): assert hasattr(model.layers["0"].mixer, "experts") assert hasattr(model.layers["0"].mixer, "shared_experts") + def test_e_score_correction_bias_stays_fp32_after_dtype_cast(self, config, backend): + """Nemotron v3 router correction bias must stay fp32 under bf16 storage casts.""" + from nemo_automodel.components.models.common.utils import cast_model_to_dtype + from nemo_automodel.components.models.nemotron_v3.model import NemotronHForCausalLM, NemotronV3Model + + backend = BackendConfig( + linear=backend.linear, + attn=backend.attn, + rms_norm=backend.rms_norm, + enable_deepep=False, + fake_balanced_gate=False, + enable_hf_state_dict_adapter=False, + ) + base_model = NemotronV3Model(config, backend=backend) + original_bias = torch.tensor([1.001, -2.003, 0.3333, 17.125], dtype=torch.float32) + base_model.layers["0"].mixer.gate.e_score_correction_bias.copy_(original_bias) + cast_model_to_dtype(base_model, torch.bfloat16) + + base_gate = base_model.layers["0"].mixer.gate + assert base_model.embed_tokens.weight.dtype == torch.bfloat16 + assert base_model.layers["0"].mixer.experts.gate_and_up_projs.dtype == torch.bfloat16 + assert base_gate.e_score_correction_bias.dtype == torch.float32 + assert torch.equal(base_gate.e_score_correction_bias, original_bias) + + model = NemotronHForCausalLM(config, backend=backend) + model.model.layers["0"].mixer.gate.e_score_correction_bias.copy_(original_bias) + cast_model_to_dtype(model, torch.bfloat16) + + gate = model.model.layers["0"].mixer.gate + state_dict = model.state_dict() + assert model.model.embed_tokens.weight.dtype == torch.bfloat16 + assert model.model.layers["0"].mixer.experts.gate_and_up_projs.dtype == torch.bfloat16 + assert gate.e_score_correction_bias.dtype == torch.float32 + assert torch.equal(gate.e_score_correction_bias, original_bias) + assert state_dict["model.layers.0.mixer.gate.e_score_correction_bias"].dtype == torch.float32 + assert torch.equal(state_dict["model.layers.0.mixer.gate.e_score_correction_bias"], original_bias) + @skip_if_no_gpu def test_moe_model_forward(self, config, backend): """Test MoE model forward pass.""" diff --git a/tests/unit_tests/models/qwen3_5/test_qwen3_5_state_dict_adapter.py b/tests/unit_tests/models/qwen3_5/test_qwen3_5_state_dict_adapter.py index 270b762bbe..96b895bbd9 100644 --- a/tests/unit_tests/models/qwen3_5/test_qwen3_5_state_dict_adapter.py +++ b/tests/unit_tests/models/qwen3_5/test_qwen3_5_state_dict_adapter.py @@ -43,6 +43,10 @@ def test_routes_bare_a_log_to_holder(self): _route_to_fp32_holder("model.language_model.layers.0.linear_attn.A_log") == "model.language_model.layers.0.linear_attn._fp32_params.A_log" ) + assert ( + _route_to_fp32_holder("model.language_model.layers.0.linear_attn.dt_bias") + == "model.language_model.layers.0.linear_attn._fp32_params.dt_bias" + ) def test_passes_through_already_routed_keys(self): key = "model.language_model.layers.0.linear_attn._fp32_params.A_log" @@ -74,7 +78,9 @@ def setup_method(self): def _sample_state_dict(self): return { "model.language_model.layers.0.linear_attn._fp32_params.A_log": torch.zeros(4), + "model.language_model.layers.0.linear_attn._fp32_params.dt_bias": torch.zeros(4), "model.language_model.layers.1.linear_attn._fp32_params.A_log": torch.ones(4), + "model.language_model.layers.1.linear_attn._fp32_params.dt_bias": torch.ones(4), "model.language_model.layers.0.self_attn.q_proj.weight": torch.zeros(2, 2), "model.language_model.embed_tokens.weight": torch.zeros(8, 2), } @@ -96,6 +102,21 @@ def test_to_hf_renames_fp32_params(self): # Number of keys preserved. assert len(out) == len(sd) + def test_to_hf_upcasts_gdn_fp32_params_saved_as_bf16(self): + sd = { + "model.language_model.layers.0.linear_attn._fp32_params.A_log": torch.zeros(4, dtype=torch.bfloat16), + "model.language_model.layers.0.linear_attn._fp32_params.dt_bias": torch.ones(4, dtype=torch.bfloat16), + "model.language_model.layers.0.self_attn.q_proj.weight": torch.zeros(2, 2, dtype=torch.bfloat16), + } + + out = self.adapter.to_hf(sd) + + assert out["model.language_model.layers.0.linear_attn.A_log"].dtype == torch.float32 + assert out["model.language_model.layers.0.linear_attn.dt_bias"].dtype == torch.float32 + q_proj_key = "model.language_model.layers.0.self_attn.q_proj.weight" + assert out[q_proj_key] is sd[q_proj_key] + assert out[q_proj_key].dtype == torch.bfloat16 + def test_to_hf_accepts_kwargs(self): # Save callsites pass exclude_key_regex, quantization, device_mesh, etc. out = self.adapter.to_hf( @@ -109,6 +130,7 @@ def test_to_hf_accepts_kwargs(self): def test_from_hf_routes_a_log_to_holder(self): hf_sd = { "model.language_model.layers.0.linear_attn.A_log": torch.zeros(4), + "model.language_model.layers.0.linear_attn.dt_bias": torch.ones(4), "model.language_model.layers.0.self_attn.q_proj.weight": torch.zeros(2, 2), } out = self.adapter.from_hf(hf_sd) @@ -116,11 +138,45 @@ def test_from_hf_routes_a_log_to_holder(self): "model.language_model.layers.0.linear_attn._fp32_params.A_log" in out and "model.language_model.layers.0.linear_attn.A_log" not in out ) + assert ( + "model.language_model.layers.0.linear_attn._fp32_params.dt_bias" in out + and "model.language_model.layers.0.linear_attn.dt_bias" not in out + ) assert "model.language_model.layers.0.self_attn.q_proj.weight" in out + def test_from_hf_upcasts_gdn_fp32_params_loaded_as_bf16(self): + hf_sd = { + "model.language_model.layers.0.linear_attn.A_log": torch.zeros(4, dtype=torch.bfloat16), + "model.language_model.layers.0.linear_attn.dt_bias": torch.ones(4, dtype=torch.bfloat16), + "model.language_model.layers.0.self_attn.q_proj.weight": torch.zeros(2, 2, dtype=torch.bfloat16), + } + + out = self.adapter.from_hf(hf_sd) + + a_log_key = "model.language_model.layers.0.linear_attn._fp32_params.A_log" + dt_bias_key = "model.language_model.layers.0.linear_attn._fp32_params.dt_bias" + q_proj_key = "model.language_model.layers.0.self_attn.q_proj.weight" + assert out[a_log_key].dtype == torch.float32 + assert out[dt_bias_key].dtype == torch.float32 + assert out[q_proj_key] is hf_sd[q_proj_key] + assert out[q_proj_key].dtype == torch.bfloat16 + + def test_from_hf_upcasts_bare_gdn_fp32_params_for_unpatched_model(self): + hf_sd = { + "model.language_model.layers.0.linear_attn.A_log": torch.zeros(4, dtype=torch.bfloat16), + "model.language_model.layers.0.linear_attn.dt_bias": torch.ones(4, dtype=torch.bfloat16), + } + adapter = Qwen3_5DenseStateDictAdapter(route_linear_attn_fp32_params=False) + + out = adapter.from_hf(hf_sd) + + assert out["model.language_model.layers.0.linear_attn.A_log"].dtype == torch.float32 + assert out["model.language_model.layers.0.linear_attn.dt_bias"].dtype == torch.float32 + def test_from_hf_keeps_bare_a_log_when_configured_for_unpatched_model(self): hf_sd = { "model.language_model.layers.0.linear_attn.A_log": torch.zeros(4), + "model.language_model.layers.0.linear_attn.dt_bias": torch.ones(4), "model.language_model.layers.0.self_attn.q_proj.weight": torch.zeros(2, 2), } adapter = Qwen3_5DenseStateDictAdapter(route_linear_attn_fp32_params=False) @@ -129,10 +185,13 @@ def test_from_hf_keeps_bare_a_log_when_configured_for_unpatched_model(self): assert "model.language_model.layers.0.linear_attn.A_log" in out assert "model.language_model.layers.0.linear_attn._fp32_params.A_log" not in out + assert "model.language_model.layers.0.linear_attn.dt_bias" in out + assert "model.language_model.layers.0.linear_attn._fp32_params.dt_bias" not in out def test_from_hf_routes_a_log_when_configured_for_patched_model(self): hf_sd = { "model.language_model.layers.0.linear_attn.A_log": torch.zeros(4), + "model.language_model.layers.0.linear_attn.dt_bias": torch.ones(4), "model.language_model.layers.0.self_attn.q_proj.weight": torch.zeros(2, 2), } adapter = Qwen3_5DenseStateDictAdapter(route_linear_attn_fp32_params=True) @@ -141,6 +200,8 @@ def test_from_hf_routes_a_log_when_configured_for_patched_model(self): assert "model.language_model.layers.0.linear_attn._fp32_params.A_log" in out assert "model.language_model.layers.0.linear_attn.A_log" not in out + assert "model.language_model.layers.0.linear_attn._fp32_params.dt_bias" in out + assert "model.language_model.layers.0.linear_attn.dt_bias" not in out def test_round_trip_is_identity(self): sd = self._sample_state_dict() @@ -155,6 +216,18 @@ def test_convert_single_tensor_to_hf(self): "model.language_model.layers.0.linear_attn._fp32_params.A_log", t ) assert out == [("model.language_model.layers.0.linear_attn.A_log", t)] + out = self.adapter.convert_single_tensor_to_hf( + "model.language_model.layers.0.linear_attn._fp32_params.dt_bias", t + ) + assert out == [("model.language_model.layers.0.linear_attn.dt_bias", t)] + + def test_convert_single_tensor_to_hf_upcasts_gdn_fp32_params(self): + t = torch.zeros(4, dtype=torch.bfloat16) + out = self.adapter.convert_single_tensor_to_hf( + "model.language_model.layers.0.linear_attn._fp32_params.A_log", t + ) + assert out[0][0] == "model.language_model.layers.0.linear_attn.A_log" + assert out[0][1].dtype == torch.float32 def test_convert_single_tensor_passthrough(self): t = torch.zeros(2, 2) diff --git a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_fp32_gdn.py b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_fp32_gdn.py new file mode 100644 index 0000000000..37e808fc6b --- /dev/null +++ b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_fp32_gdn.py @@ -0,0 +1,58 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for Qwen3.5-MoE fp32-aware GatedDeltaNet construction.""" + +from __future__ import annotations + +import pytest +import torch + +pytest.importorskip("transformers.models.qwen3_5_moe") + +from transformers.models.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeTextConfig + +from nemo_automodel.components.models.common.gated_delta_net_fp32 import HOLDER_NAME +from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import CPAwareGatedDeltaNet + + +def _text_config() -> Qwen3_5MoeTextConfig: + return Qwen3_5MoeTextConfig( + vocab_size=128, + hidden_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + intermediate_size=64, + moe_intermediate_size=32, + shared_expert_intermediate_size=32, + num_experts=4, + num_experts_per_tok=2, + max_position_embeddings=16, + rms_norm_eps=1e-6, + router_aux_loss_coef=0.01, + pad_token_id=0, + layer_types=["linear_attention"], + ) + + +def test_constructor_forces_tracked_params_fp32_under_bf16_default_dtype(): + cfg = _text_config() + cfg.torch_dtype = torch.bfloat16 + + old_default_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.bfloat16) + gdn = CPAwareGatedDeltaNet(cfg, layer_idx=0) + finally: + torch.set_default_dtype(old_default_dtype) + + assert gdn.A_log.dtype == torch.float32 + assert gdn.dt_bias.dtype == torch.float32 + + assert HOLDER_NAME in gdn._modules + assert "A_log" not in gdn._parameters + assert "dt_bias" not in gdn._parameters + assert gdn._fp32_params.A_log.dtype == torch.float32 + assert gdn._fp32_params.dt_bias.dtype == torch.float32 diff --git a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_model.py b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_model.py index cda01163ea..060b804f26 100644 --- a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_model.py +++ b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_model.py @@ -493,9 +493,14 @@ class TestQwen3_5MoeForConditionalGeneration: def test_initialization_configures_backend_components(self, vl_config, backend_config, moe_config): model = Qwen3_5MoeForConditionalGeneration(vl_config, backend=backend_config, moe_config=moe_config) - assert model.backend is backend_config + assert model.backend is not backend_config + assert model.backend.attn == backend_config.attn + assert model.backend.linear == backend_config.linear + assert model.backend.rms_norm == backend_config.rms_norm + assert model.backend.rope_fusion is False assert isinstance(model.model, Qwen3_5MoeModel) assert isinstance(model.model.language_model, Qwen3_5MoeTextModelBackend) + assert model.model.language_model.backend is model.backend assert model.model.moe_config is model.model.language_model.moe_config vision_model = getattr(model.model, "visual") @@ -1012,7 +1017,11 @@ def test_from_config_creates_model_directly(self, vl_config, moe_config, backend model = Qwen3_5MoeForConditionalGeneration.from_config(vl_config, moe_config=moe_config, backend=backend_config) assert isinstance(model, Qwen3_5MoeForConditionalGeneration) - assert model.backend is backend_config + assert model.backend is not backend_config + assert model.backend.attn == backend_config.attn + assert model.backend.linear == backend_config.linear + assert model.backend.rms_norm == backend_config.rms_norm + assert model.backend.rope_fusion is False # --------------------------------------------------------------------------- diff --git a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_state_dict_adapter.py b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_state_dict_adapter.py index bc7c2e5233..115d3b03c4 100644 --- a/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_state_dict_adapter.py +++ b/tests/unit_tests/models/qwen3_5_moe/test_qwen3_5_moe_state_dict_adapter.py @@ -874,3 +874,65 @@ def test_from_hf_routes_gating_keys_into_holder(self, adapter): assert "model.language_model.layers.0.linear_attn._fp32_params.A_log" in out assert "model.language_model.layers.0.linear_attn._fp32_params.dt_bias" in out assert "model.language_model.layers.0.linear_attn.A_log" not in out + + def test_to_hf_strips_a_log_holder(self, adapter): + sd = {"model.language_model.layers.0.linear_attn._fp32_params.A_log": torch.zeros(4)} + out = adapter.to_hf(sd) + assert "model.language_model.layers.0.linear_attn.A_log" in out + assert all("_fp32_params" not in k for k in out) + + def test_to_hf_strips_dt_bias_holder(self, adapter): + sd = {"model.language_model.layers.2.linear_attn._fp32_params.dt_bias": torch.ones(4)} + out = adapter.to_hf(sd) + assert "model.language_model.layers.2.linear_attn.dt_bias" in out + assert all("_fp32_params" not in k for k in out) + + def test_to_hf_upcasts_gdn_fp32_params_saved_as_bf16(self, adapter): + sd = { + "model.language_model.layers.0.linear_attn._fp32_params.A_log": torch.zeros(4, dtype=torch.bfloat16), + "model.language_model.layers.0.linear_attn._fp32_params.dt_bias": torch.ones(4, dtype=torch.bfloat16), + "model.language_model.layers.0.self_attn.q_proj.weight": torch.zeros(2, 2, dtype=torch.bfloat16), + } + + out = adapter.to_hf(sd) + + assert out["model.language_model.layers.0.linear_attn.A_log"].dtype == torch.float32 + assert out["model.language_model.layers.0.linear_attn.dt_bias"].dtype == torch.float32 + q_proj_key = "model.language_model.layers.0.self_attn.q_proj.weight" + assert out[q_proj_key] is sd[q_proj_key] + assert out[q_proj_key].dtype == torch.bfloat16 + + def test_convert_single_tensor_strips_holder(self, adapter): + result = adapter.convert_single_tensor_to_hf( + "model.language_model.layers.1.linear_attn._fp32_params.A_log", torch.zeros(4) + ) + assert [k for k, _ in result] == ["model.language_model.layers.1.linear_attn.A_log"] + + def test_convert_single_tensor_upcasts_gdn_fp32_params(self, adapter): + result = adapter.convert_single_tensor_to_hf( + "model.language_model.layers.1.linear_attn._fp32_params.dt_bias", + torch.zeros(4, dtype=torch.bfloat16), + ) + assert result[0][0] == "model.language_model.layers.1.linear_attn.dt_bias" + assert result[0][1].dtype == torch.float32 + + def test_bare_key_unchanged(self, adapter): + result = adapter.convert_single_tensor_to_hf("model.language_model.layers.0.linear_attn.A_log", torch.zeros(4)) + assert [k for k, _ in result] == ["model.language_model.layers.0.linear_attn.A_log"] + + def test_from_hf_routes_and_upcasts_gdn_fp32_params_loaded_as_bf16(self, adapter): + hf_state = { + "model.language_model.layers.0.linear_attn.A_log": torch.zeros(4, dtype=torch.bfloat16), + "model.language_model.layers.0.linear_attn.dt_bias": torch.ones(4, dtype=torch.bfloat16), + "model.language_model.layers.0.self_attn.q_proj.weight": torch.zeros(2, 2, dtype=torch.bfloat16), + } + + out = adapter.from_hf(hf_state) + + a_log_key = "model.language_model.layers.0.linear_attn._fp32_params.A_log" + dt_bias_key = "model.language_model.layers.0.linear_attn._fp32_params.dt_bias" + q_proj_key = "model.language_model.layers.0.self_attn.q_proj.weight" + assert out[a_log_key].dtype == torch.float32 + assert out[dt_bias_key].dtype == torch.float32 + assert out[q_proj_key] is hf_state[q_proj_key] + assert out[q_proj_key].dtype == torch.bfloat16 diff --git a/tests/unit_tests/models/qwen3_next/test_qwen3_next_fp32_gdn.py b/tests/unit_tests/models/qwen3_next/test_qwen3_next_fp32_gdn.py new file mode 100644 index 0000000000..86780190b5 --- /dev/null +++ b/tests/unit_tests/models/qwen3_next/test_qwen3_next_fp32_gdn.py @@ -0,0 +1,104 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for Qwen3-Next fp32-aware GatedDeltaNet gate routing.""" + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.models.qwen3_next.configuration_qwen3_next import Qwen3NextConfig + +from nemo_automodel.components.models.common.gated_delta_net_fp32 import HOLDER_NAME +from nemo_automodel.components.models.qwen3_next.layers import Qwen3NextFp32GatedDeltaNet + + +def _bare_gdn(num_v: int = 4, dtype: torch.dtype = torch.float32) -> Qwen3NextFp32GatedDeltaNet: + """Build a Qwen3NextFp32GatedDeltaNet shell with only A_log/dt_bias (no full init).""" + gdn = Qwen3NextFp32GatedDeltaNet.__new__(Qwen3NextFp32GatedDeltaNet) + nn.Module.__init__(gdn) + gdn.A_log = nn.Parameter(torch.ones(num_v, dtype=dtype)) + gdn.dt_bias = nn.Parameter(torch.zeros(num_v, dtype=dtype)) + return gdn + + +def _expected_gate(a: torch.Tensor) -> torch.Tensor: + return -torch.ones(a.shape[-1]).exp() * F.softplus(torch.zeros(a.shape[-1])) + + +def test_compute_gate_fallback_without_holder(): + gdn = _bare_gdn() + a = torch.zeros_like(gdn.A_log) + g = gdn._compute_gate(a) + assert g.dtype == torch.float32 + assert torch.allclose(g, _expected_gate(a), atol=1e-5) + + +def test_compute_gate_routes_through_holder(): + gdn = _constructed_gdn() + gdn._fp32_params.A_log.data.fill_(1.0) + gdn._fp32_params.dt_bias.data.zero_() + assert HOLDER_NAME in gdn._modules + assert "A_log" not in gdn._parameters + assert "dt_bias" not in gdn._parameters + assert gdn.A_log is gdn._fp32_params._parameters["A_log"] + assert gdn.dt_bias is gdn._fp32_params._parameters["dt_bias"] + + a = torch.zeros_like(gdn.A_log) + g = gdn._compute_gate(a) + assert g.dtype == torch.float32 + assert torch.allclose(g, _expected_gate(a), atol=1e-5) + + +def test_compute_gate_fp32_with_bf16_input(): + gdn = _constructed_gdn() + g = gdn._compute_gate(torch.zeros_like(gdn.A_log, dtype=torch.bfloat16)) + assert g.dtype == torch.float32 + + +def test_constructor_forces_tracked_params_fp32_under_bf16_default_dtype(): + cfg = Qwen3NextConfig( + vocab_size=128, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + max_position_embeddings=16, + rms_norm_eps=1e-6, + layer_types=["linear_attention"], + ) + cfg.torch_dtype = torch.bfloat16 + + old_default_dtype = torch.get_default_dtype() + try: + torch.set_default_dtype(torch.bfloat16) + gdn = Qwen3NextFp32GatedDeltaNet(cfg, layer_idx=0) + finally: + torch.set_default_dtype(old_default_dtype) + + assert gdn.A_log.dtype == torch.float32 + assert gdn.dt_bias.dtype == torch.float32 + assert HOLDER_NAME in gdn._modules + assert "A_log" not in gdn._parameters + assert "dt_bias" not in gdn._parameters + assert gdn._fp32_params.A_log.dtype == torch.float32 + assert gdn._fp32_params.dt_bias.dtype == torch.float32 + + +def _constructed_gdn() -> Qwen3NextFp32GatedDeltaNet: + cfg = Qwen3NextConfig( + vocab_size=128, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + max_position_embeddings=16, + rms_norm_eps=1e-6, + layer_types=["linear_attention"], + ) + return Qwen3NextFp32GatedDeltaNet(cfg, layer_idx=0) diff --git a/tests/unit_tests/models/qwen3_next/test_qwen3_next_model.py b/tests/unit_tests/models/qwen3_next/test_qwen3_next_model.py index 3d449ee42e..a8693152b2 100644 --- a/tests/unit_tests/models/qwen3_next/test_qwen3_next_model.py +++ b/tests/unit_tests/models/qwen3_next/test_qwen3_next_model.py @@ -45,8 +45,10 @@ def forward(self, hidden_states, attention_mask=None): @pytest.fixture(autouse=True) def mock_gated_deltanet(): - """Automatically mock Qwen3NextGatedDeltaNet for all tests to avoid torch.get_current_dtype() issue""" - with patch("nemo_automodel.components.models.qwen3_next.model.Qwen3NextGatedDeltaNet", MockQwen3NextGatedDeltaNet): + """Automatically mock Qwen3NextFp32GatedDeltaNet for all tests to avoid torch.get_current_dtype() issue""" + with patch( + "nemo_automodel.components.models.qwen3_next.model.Qwen3NextFp32GatedDeltaNet", MockQwen3NextGatedDeltaNet + ): yield diff --git a/tests/unit_tests/models/qwen3_next/test_qwen3_next_state_dict_adapter.py b/tests/unit_tests/models/qwen3_next/test_qwen3_next_state_dict_adapter.py index 3f1f360826..37a3aaf0eb 100644 --- a/tests/unit_tests/models/qwen3_next/test_qwen3_next_state_dict_adapter.py +++ b/tests/unit_tests/models/qwen3_next/test_qwen3_next_state_dict_adapter.py @@ -12,20 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -import torch from unittest.mock import Mock, patch -from nemo_automodel.components.moe.config import MoEConfig -from nemo_automodel.components.models.common import BackendConfig +import torch +from nemo_automodel.components.models.common import BackendConfig from nemo_automodel.components.models.qwen3_next.state_dict_adapter import Qwen3NextStateDictAdapter +from nemo_automodel.components.moe.config import MoEConfig class TestApplyKeyMapping: def _make_adapter(self): - return Qwen3NextStateDictAdapter( - config=object(), moe_config=object(), backend=object(), dtype=torch.float32 - ) + return Qwen3NextStateDictAdapter(config=object(), moe_config=object(), backend=object(), dtype=torch.float32) def test_shared_expert_to_shared_experts_mapping(self): """Test that shared_expert (singular) is mapped to shared_experts (plural)""" @@ -178,9 +176,7 @@ def test_initialization(self): moe_config = self.create_mock_moe_config() backend = self.create_mock_backend_config() - adapter = Qwen3NextStateDictAdapter( - config=config, moe_config=moe_config, backend=backend, dtype=torch.float16 - ) + adapter = Qwen3NextStateDictAdapter(config=config, moe_config=moe_config, backend=backend, dtype=torch.float16) assert adapter.config is config assert adapter.moe_config is moe_config @@ -478,7 +474,7 @@ def test_expert_tensor_conversion_with_mapping(self): tensor = torch.randn(8, 256, 1024) fqn = "model.layers.0.mlp.experts.gate_and_up_projs" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts') as mock_convert: + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts") as mock_convert: mock_convert.return_value = [ ("model.layers.0.mlp.experts.0.gate_proj.weight", torch.randn(512, 256)), ("model.layers.0.mlp.experts.0.up_proj.weight", torch.randn(512, 256)), @@ -498,7 +494,7 @@ def test_shared_expert_key_mapping(self): tensor = torch.randn(128, 256) fqn = "model.layers.0.mlp.shared_experts.gate_proj.weight" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): result = adapter.convert_single_tensor_to_hf(fqn, tensor) assert len(result) == 1 @@ -515,7 +511,7 @@ def test_non_expert_tensor_conversion(self): tensor = torch.randn(64, 64) fqn = "model.layers.0.attention.weight" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): result = adapter.convert_single_tensor_to_hf(fqn, tensor) assert len(result) == 1 @@ -531,7 +527,7 @@ def test_exclude_key_regex(self): tensor = torch.randn(64, 64) fqn = "exclude_this.weight" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts', return_value=None): + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts", return_value=None): result = adapter.convert_single_tensor_to_hf(fqn, tensor, exclude_key_regex=r"exclude.*") assert len(result) == 0 @@ -545,7 +541,7 @@ def test_expert_tensor_with_exclude_regex_and_mapping(self): tensor = torch.randn(8, 256, 1024) fqn = "model.layers.0.mlp.experts.gate_and_up_projs" - with patch.object(adapter, '_convert_single_merged_expert_to_hf_split_experts') as mock_convert: + with patch.object(adapter, "_convert_single_merged_expert_to_hf_split_experts") as mock_convert: mock_convert.return_value = [ ("model.layers.0.mlp.shared_experts.gate_proj.weight", torch.randn(128, 256)), ("exclude_me.weight", torch.randn(64, 64)), @@ -557,3 +553,86 @@ def test_expert_tensor_with_exclude_regex_and_mapping(self): assert len(result) == 1 assert result[0][0] == "model.layers.0.mlp.shared_expert.gate_proj.weight" assert "exclude_me.weight" not in [k for k, _ in result] + + +class TestFp32HolderKeyStripping: + """After FSDP isolation, A_log/dt_bias live under ``linear_attn._fp32_params.*``. + + The adapter must hide that wrapping on save so checkpoints stay HF-loadable. + """ + + def _make_adapter(self): + moe_config = Mock() + moe_config.n_routed_experts = 8 + moe_config.moe_inter_dim = 512 + backend = Mock() + backend.dispatcher = "torch" + backend.experts = "torch" + return Qwen3NextStateDictAdapter(config=Mock(), moe_config=moe_config, backend=backend, dtype=torch.float32) + + def test_convert_single_tensor_strips_a_log(self): + adapter = self._make_adapter() + result = adapter.convert_single_tensor_to_hf("model.layers.1.linear_attn._fp32_params.A_log", torch.zeros(4)) + assert [k for k, _ in result] == ["model.layers.1.linear_attn.A_log"] + + def test_convert_single_tensor_strips_dt_bias(self): + adapter = self._make_adapter() + result = adapter.convert_single_tensor_to_hf("model.layers.0.linear_attn._fp32_params.dt_bias", torch.ones(4)) + assert [k for k, _ in result] == ["model.layers.0.linear_attn.dt_bias"] + + def test_to_hf_strips_holder(self): + adapter = self._make_adapter() + out = adapter.to_hf({"model.layers.0.linear_attn._fp32_params.A_log": torch.zeros(4)}) + assert "model.layers.0.linear_attn.A_log" in out + assert all("_fp32_params" not in k for k in out) + + def test_to_hf_upcasts_gdn_fp32_params_saved_as_bf16(self): + adapter = self._make_adapter() + state_dict = { + "model.layers.0.linear_attn._fp32_params.A_log": torch.zeros(4, dtype=torch.bfloat16), + "model.layers.0.linear_attn._fp32_params.dt_bias": torch.ones(4, dtype=torch.bfloat16), + "model.layers.0.self_attn.q_proj.weight": torch.zeros(2, 2, dtype=torch.bfloat16), + } + + out = adapter.to_hf(state_dict) + + assert out["model.layers.0.linear_attn.A_log"].dtype == torch.float32 + assert out["model.layers.0.linear_attn.dt_bias"].dtype == torch.float32 + q_proj_key = "model.layers.0.self_attn.q_proj.weight" + assert out[q_proj_key] is state_dict[q_proj_key] + assert out[q_proj_key].dtype == torch.bfloat16 + + def test_bare_key_unchanged(self): + adapter = self._make_adapter() + result = adapter.convert_single_tensor_to_hf("model.layers.0.linear_attn.A_log", torch.zeros(4)) + assert [k for k, _ in result] == ["model.layers.0.linear_attn.A_log"] + + def test_convert_single_tensor_upcasts_gdn_fp32_params(self): + adapter = self._make_adapter() + result = adapter.convert_single_tensor_to_hf( + "model.layers.0.linear_attn._fp32_params.dt_bias", + torch.zeros(4, dtype=torch.bfloat16), + ) + assert result[0][0] == "model.layers.0.linear_attn.dt_bias" + assert result[0][1].dtype == torch.float32 + + def test_from_hf_routes_and_upcasts_gdn_fp32_params_loaded_as_bf16(self): + adapter = self._make_adapter() + hf_state = { + "model.layers.0.linear_attn.A_log": torch.zeros(4, dtype=torch.bfloat16), + "model.layers.0.linear_attn.dt_bias": torch.ones(4, dtype=torch.bfloat16), + "model.layers.0.self_attn.q_proj.weight": torch.zeros(2, 2, dtype=torch.bfloat16), + } + + with patch.object( + adapter, "_from_hf_w_merged_experts", side_effect=lambda state, device_mesh=None: dict(state) + ): + out = adapter.from_hf(hf_state) + + a_log_key = "model.layers.0.linear_attn._fp32_params.A_log" + dt_bias_key = "model.layers.0.linear_attn._fp32_params.dt_bias" + q_proj_key = "model.layers.0.self_attn.q_proj.weight" + assert out[a_log_key].dtype == torch.float32 + assert out[dt_bias_key].dtype == torch.float32 + assert out[q_proj_key] is hf_state[q_proj_key] + assert out[q_proj_key].dtype == torch.bfloat16 diff --git a/tests/unit_tests/moe/test_parallelizer.py b/tests/unit_tests/moe/test_parallelizer.py index 520d47cd20..5a0e4a8b59 100644 --- a/tests/unit_tests/moe/test_parallelizer.py +++ b/tests/unit_tests/moe/test_parallelizer.py @@ -649,6 +649,86 @@ def fake_shard(dim): assert model_call is not None and model_call[1]["mesh"] is fsdp_mesh +def test_shard_fp32_param_holders_shards_each_holder(monkeypatch): + """``_shard_fp32_param_holders`` fully_shards each model-owned fp32 holder.""" + P = _import_parallelizer_with_stubs(monkeypatch) + + fully_shard_mock = MagicMock() + monkeypatch.setattr(P, "fully_shard", fully_shard_mock) + + holder_param = object() + + class Holder: + def parameters(self, recurse=False): + return iter([holder_param]) + + holder = Holder() + block = type( + "Block", + (), + {"named_modules": lambda self: iter([("", self), ("linear_attn._fp32_params", holder)])}, + )() + + mesh = object() + ignored = P._shard_fp32_param_holders(block, mesh, reshard_after_forward=False, offload_policy=None) + + assert ignored == {holder_param} + holder_call = _find_call_by_first_arg(fully_shard_mock, holder) + assert holder_call is not None + _, kwargs = holder_call + assert kwargs["mesh"] is mesh + assert kwargs["reshard_after_forward"] is False + + +def test_apply_fsdp_shards_model_owned_fp32_holders(monkeypatch): + """apply_fsdp shards each model-owned ``_fp32_params`` holder per block.""" + P = _import_parallelizer_with_stubs(monkeypatch) + monkeypatch.setattr(P, "MoE", DummyMoE) + fully_shard_mock = MagicMock() + monkeypatch.setattr(P, "fully_shard", fully_shard_mock) + monkeypatch.setattr(P, "MixedPrecisionPolicy", MagicMock(return_value="FP32_MP")) + + holder_param = object() + + class Holder: + def parameters(self, recurse=False): + return iter([holder_param]) + + holder = Holder() + + class BlockWithHolder: + def __init__(self): + self.mlp = DummyMoE() + + def named_modules(self): + return iter([("", self), ("linear_attn._fp32_params", holder)]) + + block = BlockWithHolder() + model = DummyModel([block]) + fsdp_mesh = object() + mp_policy = MagicMock() + + P.apply_fsdp( + model=model, + fsdp_mesh=fsdp_mesh, + ep_enabled=False, + ep_shard_enabled=False, + ep_shard_mesh=None, + mp_policy=mp_policy, + ) + + # The holder is sharded as its own fp32 unit before the block-level shard. + holder_call = _find_call_by_first_arg(fully_shard_mock, holder) + assert holder_call is not None + _, holder_kwargs = holder_call + assert holder_kwargs["mesh"] is fsdp_mesh + assert holder_kwargs["mp_policy"] == "FP32_MP" + + block_call = _find_call_by_first_arg(fully_shard_mock, block) + assert block_call is not None + assert holder_param in block_call[1]["ignored_params"] + + def test_apply_fsdp_without_ep_enabled_has_no_ignored_params(monkeypatch): P = _import_parallelizer_with_stubs(monkeypatch) monkeypatch.setattr(P, "MoE", DummyMoE) diff --git a/tests/unit_tests/test_improved_error_messages.py b/tests/unit_tests/test_improved_error_messages.py index 230de94c24..9fe4e932f7 100644 --- a/tests/unit_tests/test_improved_error_messages.py +++ b/tests/unit_tests/test_improved_error_messages.py @@ -294,6 +294,82 @@ def __init__(self): assert model.linear.weight.dtype == torch.bfloat16 assert model.norm.weight.dtype == torch.float32 + def test_restore_loaded_model_dtype_keeps_gdn_params_fp32(self): + from nemo_automodel._transformers.model_init import _restore_loaded_model_dtype + + class LinearAttn(torch.nn.Module): + def __init__(self): + super().__init__() + self.A_log = torch.nn.Parameter(torch.zeros(4)) + self.dt_bias = torch.nn.Parameter(torch.ones(4)) + self.other = torch.nn.Parameter(torch.ones(4)) + + class DummyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.model = torch.nn.Module() + self.model.language_model = torch.nn.Module() + layer = torch.nn.Module() + layer.linear_attn = LinearAttn() + self.model.language_model.layers = torch.nn.ModuleList([layer]) + + model = DummyModel().to(torch.float32) + with patch( + "nemo_automodel.components.checkpoint.utils._get_checkpoint_tensor_dtypes", + return_value={ + "model.language_model.layers.0.linear_attn.A_log": torch.bfloat16, + "model.language_model.layers.0.linear_attn.dt_bias": torch.bfloat16, + "model.language_model.layers.0.linear_attn.other": torch.bfloat16, + }, + ): + _restore_loaded_model_dtype( + model, + "fake/model", + SimpleNamespace(architectures=["Qwen3NextForCausalLM"]), + quantization_config=None, + load_kwargs={}, + requested_dtype=torch.bfloat16, + ) + + linear_attn = model.model.language_model.layers[0].linear_attn + assert linear_attn.A_log.dtype == torch.float32 + assert linear_attn.A_log._hf_compute_dtype == torch.float32 + assert linear_attn.dt_bias.dtype == torch.float32 + assert linear_attn.dt_bias._hf_compute_dtype == torch.float32 + assert linear_attn.other.dtype == torch.bfloat16 + assert linear_attn.other._hf_compute_dtype == torch.bfloat16 + + def test_restore_loaded_model_dtype_does_not_force_non_gdn_arch_fp32(self): + from nemo_automodel._transformers.model_init import _restore_loaded_model_dtype + + class LinearAttn(torch.nn.Module): + def __init__(self): + super().__init__() + self.A_log = torch.nn.Parameter(torch.zeros(4)) + + class DummyModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.layer = torch.nn.Module() + self.layer.linear_attn = LinearAttn() + + model = DummyModel().to(torch.float32) + with patch( + "nemo_automodel.components.checkpoint.utils._get_checkpoint_tensor_dtypes", + return_value={"layer.linear_attn.A_log": torch.bfloat16}, + ): + _restore_loaded_model_dtype( + model, + "fake/model", + SimpleNamespace(architectures=["FutureLinearAttentionForCausalLM"]), + quantization_config=None, + load_kwargs={}, + requested_dtype=torch.bfloat16, + ) + + assert model.layer.linear_attn.A_log.dtype == torch.bfloat16 + assert model.layer.linear_attn.A_log._hf_compute_dtype == torch.bfloat16 + def test_restore_loaded_model_dtype_preserves_tied_weights(self): from nemo_automodel._transformers.model_init import _restore_loaded_model_dtype