diff --git a/examples/vlm_finetune/qwen3_5/qwen3_5_27b_tp4pp4.yaml b/examples/vlm_finetune/qwen3_5/qwen3_5_27b_tp4pp4.yaml index d4375d6cf4..8aeffa331d 100644 --- a/examples/vlm_finetune/qwen3_5/qwen3_5_27b_tp4pp4.yaml +++ b/examples/vlm_finetune/qwen3_5/qwen3_5_27b_tp4pp4.yaml @@ -37,6 +37,12 @@ rng: model: _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained pretrained_model_name_or_path: Qwen/Qwen3.5-27B + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + attn: sdpa + linear: torch + rms_norm: torch + rope_fusion: false attn_implementation: sdpa processor: diff --git a/examples/vlm_finetune/qwen3_5/qwen3_5_4b.yaml b/examples/vlm_finetune/qwen3_5/qwen3_5_4b.yaml index 59afc5f5a2..f7fa15a2e5 100644 --- a/examples/vlm_finetune/qwen3_5/qwen3_5_4b.yaml +++ b/examples/vlm_finetune/qwen3_5/qwen3_5_4b.yaml @@ -38,6 +38,12 @@ rng: model: _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained pretrained_model_name_or_path: Qwen/Qwen3.5-4B + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + attn: sdpa + linear: torch + rms_norm: torch + rope_fusion: false attn_implementation: sdpa diff --git a/examples/vlm_finetune/qwen3_5/qwen3_5_4b_neat_packing.yaml b/examples/vlm_finetune/qwen3_5/qwen3_5_4b_neat_packing.yaml index 55848cdfb2..61046d1df1 100644 --- a/examples/vlm_finetune/qwen3_5/qwen3_5_4b_neat_packing.yaml +++ b/examples/vlm_finetune/qwen3_5/qwen3_5_4b_neat_packing.yaml @@ -37,6 +37,12 @@ rng: model: _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained pretrained_model_name_or_path: Qwen/Qwen3.5-4B + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + attn: sdpa + linear: torch + rms_norm: torch + rope_fusion: false attn_implementation: flash_attention_2 processor: diff --git a/examples/vlm_finetune/qwen3_5/qwen3_5_9b.yaml b/examples/vlm_finetune/qwen3_5/qwen3_5_9b.yaml index 780cfe6e9e..e4a7a891d3 100644 --- a/examples/vlm_finetune/qwen3_5/qwen3_5_9b.yaml +++ b/examples/vlm_finetune/qwen3_5/qwen3_5_9b.yaml @@ -37,6 +37,12 @@ rng: model: _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained pretrained_model_name_or_path: Qwen/Qwen3.5-9B + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + attn: sdpa + linear: torch + rms_norm: torch + rope_fusion: false attn_implementation: sdpa diff --git a/examples/vlm_finetune/qwen3_5/qwen3_6_27b.yaml b/examples/vlm_finetune/qwen3_5/qwen3_6_27b.yaml index 108174c5aa..5f0efbedfe 100644 --- a/examples/vlm_finetune/qwen3_5/qwen3_6_27b.yaml +++ b/examples/vlm_finetune/qwen3_5/qwen3_6_27b.yaml @@ -36,6 +36,12 @@ rng: model: _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained pretrained_model_name_or_path: Qwen/Qwen3.6-27B + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + attn: sdpa + linear: torch + rms_norm: torch + rope_fusion: false attn_implementation: sdpa torch_dtype: bfloat16 text_config: diff --git a/examples/vlm_finetune/qwen3_5/qwen3_6_27b_lora.yaml b/examples/vlm_finetune/qwen3_5/qwen3_6_27b_lora.yaml index 76dfda4339..b7f38a28b0 100644 --- a/examples/vlm_finetune/qwen3_5/qwen3_6_27b_lora.yaml +++ b/examples/vlm_finetune/qwen3_5/qwen3_6_27b_lora.yaml @@ -39,6 +39,12 @@ rng: model: _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained pretrained_model_name_or_path: Qwen/Qwen3.6-27B + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + attn: sdpa + linear: torch + rms_norm: torch + rope_fusion: false attn_implementation: sdpa torch_dtype: bfloat16 diff --git a/examples/vlm_kd/qwen3_5/qwen3_5_vl_4b_kd.yaml b/examples/vlm_kd/qwen3_5/qwen3_5_vl_4b_kd.yaml index 171f69b186..bec01009f3 100644 --- a/examples/vlm_kd/qwen3_5/qwen3_5_vl_4b_kd.yaml +++ b/examples/vlm_kd/qwen3_5/qwen3_5_vl_4b_kd.yaml @@ -41,6 +41,12 @@ rng: model: _target_: nemo_automodel.NeMoAutoModelForImageTextToText.from_pretrained pretrained_model_name_or_path: Qwen/Qwen3.5-4B + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + attn: sdpa + linear: torch + rms_norm: torch + rope_fusion: false attn_implementation: sdpa # Teacher diff --git a/nemo_automodel/_transformers/capabilities.py b/nemo_automodel/_transformers/capabilities.py index 840735dba4..5e1ccdb733 100644 --- a/nemo_automodel/_transformers/capabilities.py +++ b/nemo_automodel/_transformers/capabilities.py @@ -102,6 +102,10 @@ def _is_hybrid(model: "nn.Module") -> bool: inner = getattr(model, "language_model", None) if inner is not None: candidates.append(getattr(inner, "config", None)) + # VLM configs nest the decoder config under ``text_config``. + for c in list(candidates): + if c is not None: + candidates.append(getattr(c, "text_config", None)) for config in candidates: if config is None: continue @@ -111,6 +115,11 @@ def _is_hybrid(model: "nn.Module") -> bool: return True if getattr(config, "is_hybrid_model", False) is True: return True + # Qwen3.5 / Qwen3-Next style: per-layer ``layer_types`` mixing + # ``linear_attention`` (gated-delta / SSM) with ``full_attention``. + layer_types = getattr(config, "layer_types", None) + if layer_types and any(str(t) == "linear_attention" for t in layer_types): + return True return False diff --git a/nemo_automodel/_transformers/kernel_patches.py b/nemo_automodel/_transformers/kernel_patches.py index 11c61ddc87..b473462e4b 100644 --- a/nemo_automodel/_transformers/kernel_patches.py +++ b/nemo_automodel/_transformers/kernel_patches.py @@ -42,16 +42,10 @@ logger = logging.getLogger(__name__) -_MODEL_RUNTIME_PATCHES = { - "Qwen3_5ForCausalLM": ( - "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn", - "apply_model_runtime_patches", - ), - "Qwen3_5ForConditionalGeneration": ( - "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn", - "apply_model_runtime_patches", - ), -} +# Models build their CP/fp32-gate-aware modules at construction; no load-time +# runtime patching is registered here. (Qwen3.5 dense/MoE build the native +# CPAwareGatedDeltaNet + fp32 SSMGate in their model __init__.) +_MODEL_RUNTIME_PATCHES = {} def _assert_same_signature(original, patched): diff --git a/nemo_automodel/components/distributed/optimized_tp_plans.py b/nemo_automodel/components/distributed/optimized_tp_plans.py index ea3cfbffcb..cb62ae6f85 100644 --- a/nemo_automodel/components/distributed/optimized_tp_plans.py +++ b/nemo_automodel/components/distributed/optimized_tp_plans.py @@ -717,6 +717,10 @@ def _parallelize_falcon_h1( _get_class_qualname(Qwen3ForSequenceClassification): _parallelize_qwen_classification, # Hard-coded qualname to avoid eagerly importing transformers.models.qwen3_5. "transformers.models.qwen3_5.modeling_qwen3_5.Qwen3_5ForConditionalGeneration": _parallelize_qwen3_5_vlm, + # NeMo-native Qwen3.5 dense (custom-model port): same plan — shard self_attn + + # MLP, leave the GatedDeltaNet (linear_attn) replicated. + "nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForConditionalGeneration": _parallelize_qwen3_5_vlm, + "nemo_automodel.components.models.qwen3_5.model.Qwen3_5ForCausalLM": _parallelize_qwen3_5_vlm, # Falcon-H1 (hybrid Transformer + Mamba2). HF ships only a minimal # _tp_plan ({"lm_head": "colwise_gather_output"}) and names its MLP # "feed_forward", so the generic fallback plan leaves feed_forward replicated diff --git a/nemo_automodel/components/distributed/parallelizer.py b/nemo_automodel/components/distributed/parallelizer.py index ad680ab7a1..f52e029358 100644 --- a/nemo_automodel/components/distributed/parallelizer.py +++ b/nemo_automodel/components/distributed/parallelizer.py @@ -512,16 +512,13 @@ class Qwen3_5ParallelizationStrategy(DefaultParallelizationStrategy): """ def parallelize(self, model, device_mesh, dp_shard_cp_mesh_name="dp_shard_cp", **kwargs): - # Patch HF GatedDeltaNet for FSDP mixed-dtype support (and CP if enabled) - from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import patch_hf_model - cp_mesh_name = dp_shard_cp_mesh_name.replace("dp_shard_", "") cp_enabled = cp_mesh_name in device_mesh.mesh_dim_names and device_mesh[cp_mesh_name].size() > 1 - patch_hf_model(model, cp_enabled=cp_enabled) - # Submodules that must compute in fp32 even under fp32 master weights (bf16 - # compute). patch_hf_model declares ``_fp32_params`` on the model. - fp32_compute_module_names = tuple(getattr(model, "_keep_in_fp32_modules_strict", None) or ()) + # The Qwen3.5 model builds CPAwareGatedDeltaNet with a fp32 ``SSMGate`` + # (``_fp32_params``) at construction — no runtime patch needed. Keep those + # params in their own dtype-uniform fp32 FSDP group (true master weights). + fp32_compute_module_names = ("_fp32_params",) # Delegate TP, AC, mixed precision to the default strategy, but # override the FSDP sharding to use fully_shard_by_dtype. diff --git a/nemo_automodel/components/models/common/utils.py b/nemo_automodel/components/models/common/utils.py index e3ba879056..0970dba124 100644 --- a/nemo_automodel/components/models/common/utils.py +++ b/nemo_automodel/components/models/common/utils.py @@ -451,7 +451,9 @@ def get_rope_config(config) -> tuple[float, dict, float]: return rope_theta, rope_parameters, partial_rotary_factor -def cast_model_to_dtype(model: nn.Module, dtype: torch.dtype = torch.bfloat16) -> None: +def cast_model_to_dtype( + model: nn.Module, dtype: torch.dtype = torch.bfloat16, skip_modules: tuple[str, ...] = () +) -> None: """Cast model parameters to the target dtype, keeping fp32 modules in full precision. Respects ``_keep_in_fp32_modules`` / ``_keep_in_fp32_modules_strict`` on @@ -465,10 +467,34 @@ def cast_model_to_dtype(model: nn.Module, dtype: torch.dtype = torch.bfloat16) - Args: model: The model whose parameters should be cast. dtype: Target dtype (e.g. ``torch.bfloat16``). + skip_modules: Names of immediate submodules to leave entirely untouched + (kept at their current dtype). Unlike the ``_keep_in_fp32_modules`` + restore path, these are *detached* during the cast so ``model.to()`` + never visits them — the only reliable way to preserve an fp32 + parameter once it is FSDP2-sharded (post-shard ``.data`` reassignment + does not stick). The caller must guarantee each skipped submodule is + its own dtype-uniform FSDP group (e.g. Qwen3.5's ``_fp32_params`` + holder, sharded separately in fp32), so leaving it fp32 cannot break + FSDP's uniform-dtype rule. """ fp32_keywords = _get_fp32_module_keywords(model) - model.to(dtype) + # 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. + detached: list[tuple[nn.Module, str, nn.Module]] = [] + if skip_modules: + for _, parent in model.named_modules(): + for child_name, child in list(parent._modules.items()): + if child is not None and child_name in skip_modules: + detached.append((parent, child_name, child)) + parent._modules[child_name] = None + + try: + model.to(dtype) + finally: + for parent, child_name, child in detached: + parent._modules[child_name] = child + if fp32_keywords: if _has_dtensor_params(model): logger.warning( diff --git a/nemo_automodel/components/models/qwen3_5/decoder_layer.py b/nemo_automodel/components/models/qwen3_5/decoder_layer.py deleted file mode 100644 index 513c34cfc5..0000000000 --- a/nemo_automodel/components/models/qwen3_5/decoder_layer.py +++ /dev/null @@ -1,111 +0,0 @@ -# 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. - -"""Custom Qwen3.5 decoder layer that threads packed-sequence metadata to ``linear_attn``. - -HF's ``Qwen3_5DecoderLayer.forward`` calls ``self.linear_attn`` with only -``hidden_states``, ``cache_params``, ``cache_position`` and ``attention_mask``. -For NEAT-packed inputs the linear-attn kernel additionally needs: - -* ``cu_seqlens`` -- per-document cumulative lengths (FLA's segment-reset signal - for ``chunk_gated_delta_rule``). -* ``indices`` -- non-padding token indices in the flattened sequence (used to - unpad ``[B, T, ...]`` to ``[1, total_valid, ...]`` before the kernel and - re-pad after; required for B>1 packed batches). - -* ``position_ids`` -- needed by the CP path to undo PyTorch's load-balanced - shuffle. - -This subclass derives the packing kwargs from the indexed ``attention_mask`` and -forwards them, plus ``position_ids``, into ``linear_attn``. ``patch_hf_model`` -swaps every ``Qwen3_5DecoderLayer`` instance to this class at model build time, -so this is the *only* file that needs to know about the kwarg drop in HF's -decoder layer. -""" - -from __future__ import annotations - -import torch -from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5DecoderLayer - -from nemo_automodel.components.models.common.packing import get_unpad_data, is_indexed_packed_mask - - -class Qwen3_5DecoderLayerWithPacking(Qwen3_5DecoderLayer): - """Drop-in subclass of HF ``Qwen3_5DecoderLayer`` with packing-aware dispatch. - - All weights and ``__init__`` are inherited unchanged. Only ``forward`` is - overridden so the ``linear_attn`` call site receives ``cu_seqlens``, - ``indices`` and ``position_ids`` in addition to ``attention_mask``. - """ - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - attention_mask: torch.Tensor | None = None, - position_ids: torch.LongTensor | None = None, - past_key_values=None, - cache_position: torch.LongTensor | None = None, - **kwargs, - ) -> torch.Tensor: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - - if self.layer_type == "linear_attention": - cu_seqlens: torch.Tensor | None = None - indices: torch.Tensor | None = None - linear_attn_mask = attention_mask - packed_seq_ids = kwargs.get("_packed_seq_ids") - if is_indexed_packed_mask(attention_mask): - packing_mask = attention_mask - elif is_indexed_packed_mask(packed_seq_ids): - packing_mask = packed_seq_ids - else: - packing_mask = None - - if packing_mask is not None: - indices_t, cu_seqlens_t, _ = get_unpad_data(packing_mask) - indices = indices_t - cu_seqlens = cu_seqlens_t.to(torch.long) - linear_attn_mask = packing_mask - - hidden_states = self.linear_attn( - hidden_states=hidden_states, - cache_params=past_key_values, - cache_position=cache_position, - attention_mask=linear_attn_mask, - position_ids=position_ids, - cu_seqlens=cu_seqlens, - indices=indices, - ) - elif self.layer_type == "full_attention": - hidden_states, _ = self.self_attn( - hidden_states=hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_values=past_key_values, - cache_position=cache_position, - position_embeddings=position_embeddings, - **kwargs, - ) - - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - - return hidden_states diff --git a/nemo_automodel/components/models/qwen3_5/model.py b/nemo_automodel/components/models/qwen3_5/model.py index fa43a57d8d..bd2973ffe4 100644 --- a/nemo_automodel/components/models/qwen3_5/model.py +++ b/nemo_automodel/components/models/qwen3_5/model.py @@ -22,22 +22,29 @@ import torch import torch.nn as nn -from transformers.modeling_outputs import CausalLMOutputWithPast +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5Config, Qwen3_5TextConfig from transformers.models.qwen3_5.modeling_qwen3_5 import ( Qwen3_5DecoderLayer, Qwen3_5RMSNorm, - Qwen3_5TextModel, + Qwen3_5TextRotaryEmbedding, create_causal_mask, ) from transformers.models.qwen3_5.modeling_qwen3_5 import ( Qwen3_5ForConditionalGeneration as HFQwen3_5ForConditionalGeneration, ) +from transformers.models.qwen3_5.modeling_qwen3_5 import ( + Qwen3_5Model as HFQwen3_5Model, +) from nemo_automodel.components.models.common import BackendConfig from nemo_automodel.components.models.common.hf_checkpointing_mixin import HFCheckpointingMixin from nemo_automodel.components.models.common.mtp import MTPConfig, MTPModule, roll_tensor from nemo_automodel.components.models.common.utils import cast_model_to_dtype +from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import CPAwareGatedDeltaNet +from nemo_automodel.components.models.qwen3_next.layers import Qwen3NextRMSNorm +from nemo_automodel.components.models.qwen3_next.model import Block +from nemo_automodel.components.moe.layers import MoEConfig from nemo_automodel.components.utils.model_utils import squeeze_input_for_thd from nemo_automodel.shared.utils import dtype_from_str as get_dtype @@ -225,6 +232,323 @@ def factory(*, global_idx, depth, sublayer_idx, block_type, has_fusion, has_fina ) +class Fp32SafeQwen3_5TextRotaryEmbedding(Qwen3_5TextRotaryEmbedding): + """Ensure inv_freq stays in float32 across ``.to(dtype)`` calls.""" + + def _apply(self, fn: Any, recurse: bool = True): + inv_freq_fp32 = self.inv_freq.detach().clone().to(torch.float32) + result = super()._apply(fn, recurse=recurse) + self.register_buffer("inv_freq", inv_freq_fp32.to(device=self.inv_freq.device), persistent=False) + return result + + +def _dense_moe_config(config: Qwen3_5TextConfig, dtype: torch.dtype) -> MoEConfig: + """Trivial MoEConfig for the dense Qwen3.5 backbone. + + The dense model has no experts (``num_experts`` is 0/absent), so ``Block`` + builds a dense ``MLP`` and never consults this config; it is only required to + satisfy ``Block.__init__``'s signature. + """ + inter = config.intermediate_size + return MoEConfig( + dim=config.hidden_size, + inter_dim=inter, + moe_inter_dim=inter, + n_routed_experts=0, + n_shared_experts=0, + n_activated_experts=0, + n_expert_groups=0, + n_limited_groups=0, + train_gate=True, + gate_bias_update_factor=0.0, + aux_loss_coeff=0.0, + score_func="softmax", + route_scale=1.0, + norm_topk_prob=True, + dtype=dtype, + ) + + +class Qwen3_5DenseBlock(Block): + """Qwen3.5 dense decoder block on top of the Qwen3-Next ``Block``. + + Identical to ``Qwen3_5MoeBlock`` except the MLP degrades to a dense ``MLP`` + (no experts). The CP-aware GatedDeltaNet is built natively for + linear-attention layers, and the forward threads NEAT-packing kwargs. + """ + + def __init__(self, layer_idx, config, moe_config, backend): + super().__init__(layer_idx, config, moe_config, backend) + if self.layer_type == "linear_attention": + self.linear_attn = CPAwareGatedDeltaNet(config, layer_idx) + + def forward( + self, + x: torch.Tensor, + *, + freqs_cis: torch.Tensor, + attention_mask: torch.Tensor | None = None, + padding_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + **attn_kwargs: Any, + ) -> torch.Tensor: + if self.layer_type != "linear_attention": + attn_kwargs = dict(attn_kwargs) + attn_kwargs.pop("seq_index", None) + return super().forward( + x, + freqs_cis=freqs_cis, + attention_mask=attention_mask, + padding_mask=padding_mask, + position_ids=position_ids, + **attn_kwargs, + ) + + from nemo_automodel.components.models.common.packing import get_unpad_data, is_indexed_packed_mask + + cu_seqlens: torch.Tensor | None = None + indices: torch.Tensor | None = None + linear_attn_mask = attention_mask + packed_seq_ids = attn_kwargs.get("_packed_seq_ids") + if is_indexed_packed_mask(attention_mask): + packing_mask = attention_mask + elif is_indexed_packed_mask(packed_seq_ids): + packing_mask = packed_seq_ids + else: + packing_mask = None + + if packing_mask is not None: + indices_t, cu_seqlens_t, _ = get_unpad_data(packing_mask) + cu_seqlens = cu_seqlens_t.to(torch.long) + indices = indices_t + linear_attn_mask = packing_mask + + if linear_attn_mask is not None and padding_mask is None: + padding_mask = linear_attn_mask.bool().logical_not() + + normed_x = self.input_layernorm(x) + attn_out = self.linear_attn( + hidden_states=normed_x, + attention_mask=linear_attn_mask, + position_ids=position_ids, + seq_index=attn_kwargs.get("seq_index"), + cu_seqlens=cu_seqlens, + indices=indices, + ) + x = x + attn_out + mlp_out = self._mlp(x=self.post_attention_layernorm(x), padding_mask=padding_mask) + return x + mlp_out + + def init_weights(self, buffer_device: torch.device): + for norm in (self.input_layernorm, self.post_attention_layernorm): + norm.reset_parameters() + if self.layer_type == "full_attention": + self.self_attn.init_weights(buffer_device) + elif self.layer_type == "linear_attention": + self.linear_attn.dt_bias.data.fill_(1.0) + self.linear_attn.A_log.data.uniform_(0, 16).log_() + for linear in ( + self.linear_attn.in_proj_qkv, + self.linear_attn.in_proj_z, + self.linear_attn.in_proj_b, + self.linear_attn.in_proj_a, + self.linear_attn.out_proj, + ): + nn.init.trunc_normal_(linear.weight, mean=0.0, std=0.02) + if hasattr(self.linear_attn.norm, "reset_parameters"): + self.linear_attn.norm.reset_parameters() + else: + self.linear_attn.norm.weight.data.fill_(1.0) + self.mlp.init_weights(buffer_device) + + +class Qwen3_5DenseTextBackbone(nn.Module): + """Qwen3.5 dense text decoder rebuilt on the Qwen3-Next ``Block``. + + Native counterpart of ``Qwen3_5MoeTextModelBackend`` for the dense model: + reuses the same blocks/GatedDeltaNet/norm/rotary so dense and MoE share one + code path, with the fp32 ``SSMGate`` built at construction (no runtime patch). + """ + + def __init__(self, config: Qwen3_5TextConfig, backend: BackendConfig): + super().__init__() + self.config = config + self.backend = backend + self.padding_idx = getattr(config, "pad_token_id", None) + self.vocab_size = config.vocab_size + model_dtype = get_dtype(getattr(config, "torch_dtype", None), torch.bfloat16) + moe_config = _dense_moe_config(config, model_dtype) + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx, dtype=model_dtype) + self.layers = nn.ModuleDict( + {str(i): Qwen3_5DenseBlock(i, config, moe_config, backend) for i in range(config.num_hidden_layers)} + ) + self.norm = Qwen3NextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Fp32SafeQwen3_5TextRotaryEmbedding(config=config) + + def forward( + self, + input_ids: torch.Tensor | None = None, + *, + inputs_embeds: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + cache_position: torch.Tensor | None = None, + padding_mask: torch.Tensor | None = None, + past_key_values: Any | None = None, + use_cache: bool | None = None, + output_hidden_states: bool | None = None, + **attn_kwargs: Any, + ) -> BaseModelOutputWithPast: + del output_hidden_states # accepted for HF-forward compatibility; ignored + if past_key_values is not None or use_cache: + raise NotImplementedError("KV cache is not supported for the Qwen3.5 dense backend implementation.") + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + if cache_position is None: + cache_position = torch.arange(0, inputs_embeds.shape[1], device=inputs_embeds.device) + + if position_ids is None: + position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1) + elif position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + # [4, bs, seq] position_ids (dim-0 = [text, T, H, W]); keep [T, H, W] for M-RoPE. + if position_ids.ndim == 3 and position_ids.shape[0] == 4: + position_ids = position_ids[1:] + + if getattr(self, "_cp_enabled", False): + attention_mask = None + padding_mask = None + + if padding_mask is None and attention_mask is not None: + if attention_mask.ndim <= 2: + padding_mask = attention_mask.bool().logical_not() + else: + padding_mask = attention_mask[:, 0].diagonal(dim1=-2, dim2=-1).bool().logical_not() + + hidden_states = inputs_embeds + cos, sin = self.rotary_emb(hidden_states, position_ids) + head_dim = cos.shape[-1] // 2 + freqs_cis = torch.cat((cos[..., :head_dim], sin[..., :head_dim]), dim=-1) + + for decoder_layer in self.layers.values(): + hidden_states = decoder_layer( + x=hidden_states, + freqs_cis=freqs_cis, + attention_mask=attention_mask, + padding_mask=padding_mask, + position_ids=position_ids, + **attn_kwargs, + ) + + if self.norm is not None: + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=None) + + def get_input_embeddings(self) -> nn.Module: + return self.embed_tokens + + def set_input_embeddings(self, value: nn.Module) -> None: + self.embed_tokens = value + + @torch.no_grad() + def init_weights(self, buffer_device: torch.device | None = None) -> None: + buffer_device = buffer_device or _default_init_device() + with buffer_device: + if self.embed_tokens is not None: + nn.init.normal_(self.embed_tokens.weight) + if self.norm is not None: + self.norm.reset_parameters() + self.rotary_emb.device = buffer_device + for layer in self.layers.values(): + layer.init_weights(buffer_device=buffer_device) + + +class Qwen3_5Model(HFQwen3_5Model): + """Thin VLM wrapper exposing ``language_model`` internals as properties and + routing the forward: HF vision+scatter path when media is present, else the + NeMo dense backbone directly. Mirrors ``Qwen3_5MoeModel``.""" + + @property + def layers(self): + return self.language_model.layers + + @property + def embed_tokens(self): + return self.language_model.embed_tokens + + @property + def norm(self): + return self.language_model.norm + + def forward( + self, + input_ids=None, + attention_mask=None, + position_ids=None, + past_key_values=None, + inputs_embeds=None, + pixel_values=None, + pixel_values_videos=None, + image_grid_thw=None, + video_grid_thw=None, + cache_position=None, + **kwargs, + ): + # Media present + vision encoder: full HF VL forward (vision encode + + # multimodal scatter), which then calls self.language_model (NeMo backbone). + if (pixel_values is not None or pixel_values_videos is not None) and self.visual is not None: + embed_tokens = self.get_input_embeddings() + if inputs_embeds is None: + if embed_tokens is not None: + inputs_embeds = embed_tokens(input_ids) + elif ( + input_ids is not None + and isinstance(input_ids, torch.Tensor) + and input_ids.dtype in (torch.float16, torch.bfloat16, torch.float32) + ): + inputs_embeds = input_ids + input_ids = None + else: + raise ValueError("inputs_embeds must be provided for pipeline stages without embed_tokens") + media_tensor = pixel_values if pixel_values is not None else pixel_values_videos + if isinstance(media_tensor, torch.Tensor) and hasattr(self.visual, "rotary_pos_emb"): + self.visual.rotary_pos_emb.to(media_tensor.device) + return super().forward( + input_ids=None, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_values_videos=pixel_values_videos, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + cache_position=cache_position, + **kwargs, + ) + + # Text-only path: call the NeMo backend directly. + if ( + inputs_embeds is None + and input_ids is not None + and isinstance(input_ids, torch.Tensor) + and input_ids.dtype in (torch.float16, torch.bfloat16, torch.float32) + ): + inputs_embeds = input_ids + input_ids = None + if input_ids is None and inputs_embeds is None: + raise ValueError("Either input_ids or inputs_embeds must be provided") + return self.language_model( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + cache_position=cache_position, + **kwargs, + ) + + class Qwen3_5ForCausalLM(HFCheckpointingMixin, nn.Module): """Qwen3.5 dense causal LM with optional Megatron-style MTP head.""" @@ -270,13 +594,20 @@ def __init__( self.config = config self.backend = backend or BackendConfig() - self.model = Qwen3_5TextModel(config) + self.model = Qwen3_5DenseTextBackbone(config, self.backend) dtype = next(self.model.parameters()).dtype self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False, dtype=dtype) if getattr(config, "tie_word_embeddings", False): self.tie_weights() + # Keep the SSM-gating params (in each linear_attn ``_fp32_params`` holder) + # in fp32 storage even under a bf16 bulk 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 + self.mtp_config = build_mtp_config_from_hf( config, loss_scaling_factor=mtp_loss_scaling_factor, @@ -285,7 +616,7 @@ def __init__( self.mtp = build_qwen3_5_dense_mtp(config, self.mtp_config, dtype=dtype) if self.mtp_config.enabled else None if self.backend.enable_hf_state_dict_adapter: - self.state_dict_adapter = Qwen3_5DenseStateDictAdapter(route_linear_attn_fp32_params=False) + self.state_dict_adapter = Qwen3_5DenseStateDictAdapter(route_linear_attn_fp32_params=True) def get_input_embeddings(self) -> nn.Module: return self.model.embed_tokens @@ -384,8 +715,14 @@ def initialize_weights( ) -> None: buffer_device = buffer_device or _default_init_device() init_std = float(getattr(self.config, "initializer_range", 0.02)) + # The backbone (embed/norm/layers, incl. GatedDeltaNet-specific init) owns + # its own init_weights; init only the non-backbone modules (lm_head, MTP) + # generically so the GatedDeltaNet/SSMGate init is not clobbered. + self.model.init_weights(buffer_device=buffer_device) with buffer_device: - for module in self.modules(): + for name, module in self.named_modules(): + if name == "model" or name.startswith("model."): + continue if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=init_std) if module.bias is not None: @@ -394,43 +731,9 @@ def initialize_weights( nn.init.normal_(module.weight, mean=0.0, std=init_std) if module.padding_idx is not None: module.weight[module.padding_idx].zero_() - elif isinstance(module, Qwen3_5RMSNorm): + elif isinstance(module, (Qwen3_5RMSNorm, Qwen3NextRMSNorm)): nn.init.zeros_(module.weight) - cast_model_to_dtype(self, dtype) - - -class Qwen3_5TextModelPP(Qwen3_5TextModel): - """``Qwen3_5TextModel`` whose forward survives a pipeline-parallel split. - - The PP splitter rewrites ``self.layers`` from an ``nn.ModuleList`` into an - ``nn.ModuleDict`` keyed by original layer index, and drops ``norm`` (sets it to - ``None``) on every stage except the last. HF's text forward is not PP-aware: it - does ``self.layers[: config.num_hidden_layers]`` (a slice, which raises - ``KeyError`` on a ``ModuleDict``) and calls ``self.norm(...)`` unconditionally - (``None`` is not callable on non-last stages). - - Rather than fork HF's mRoPE / linear-attention-mask / rotary logic (which is - version-sensitive), this override briefly presents the stage's layers as a - slice-able ``nn.ModuleList`` over the *same* layer objects and swaps a dropped - ``norm`` for an ``nn.Identity`` no-op, delegates to HF's ``forward`` via - ``super()``, then restores both. ``embed_tokens`` is untouched: non-first stages - are fed ``inputs_embeds`` so it is never called. Outside PP (``layers`` is a - ``ModuleList`` and ``norm`` is present) both branches are skipped, so this is a - pure passthrough to the upstream forward. - """ - - def forward(self, *args, **kwargs): - saved_layers = self.layers - saved_norm = self.norm - try: - if isinstance(saved_layers, nn.ModuleDict): - self.layers = nn.ModuleList(saved_layers.values()) - if saved_norm is None: - self.norm = nn.Identity() - return super().forward(*args, **kwargs) - finally: - self.layers = saved_layers - self.norm = saved_norm + cast_model_to_dtype(self, dtype, skip_modules=("_fp32_params",)) class Qwen3_5ForConditionalGeneration(HFCheckpointingMixin, HFQwen3_5ForConditionalGeneration): @@ -487,15 +790,30 @@ def __init__( super().__init__(config) self.backend = backend or BackendConfig() - # Make the HF text backbone forward pipeline-split-safe. Same instance and - # weights — only the (overridden) forward changes — so this is transparent - # outside PP and survives the splitter's deepcopy (it is class-based, not a - # monkeypatched instance attribute). - self.model.language_model.__class__ = Qwen3_5TextModelPP - text_config = config.text_config + # Replace the HF text decoder with the native NeMo backbone (built on the + # shared Block + CPAwareGatedDeltaNet), and class-swap the inner VLM model to + # the routing wrapper. The HF vision tower + image/video scatter + mRoPE + + # generation helpers stay intact (inherited). The backbone's ModuleDict layers + # are pipeline-split-safe, so no Qwen3_5TextModelPP shim is needed. + self.model.__class__ = Qwen3_5Model + self.model.language_model = Qwen3_5DenseTextBackbone(text_config, self.backend) + + # Keep the SSM-gating params (per-layer ``_fp32_params`` holder) in fp32 + # storage even under a bf16 bulk 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 + param_dtype = next(self.model.language_model.parameters()).dtype dtype = get_dtype(getattr(text_config, "torch_dtype", None), param_dtype) + # ``super().__init__`` ran HF ``post_init`` (-> ``initialize_weights``) and may + # have cast the inherited ``lm_head`` to a different bulk dtype before the + # native backbone was swapped in; realign it to the backbone dtype so the + # final hidden states and ``lm_head`` agree. + if self.lm_head is not None and self.lm_head.weight.dtype != dtype: + self.lm_head = self.lm_head.to(dtype) self.mtp_config = build_mtp_config_from_hf( text_config, loss_scaling_factor=mtp_loss_scaling_factor, @@ -507,7 +825,7 @@ def __init__( if self.mtp is not None: cast_model_to_dtype(self.mtp, dtype) if self.backend.enable_hf_state_dict_adapter: - self.state_dict_adapter = Qwen3_5DenseStateDictAdapter(route_linear_attn_fp32_params=False) + self.state_dict_adapter = Qwen3_5DenseStateDictAdapter(route_linear_attn_fp32_params=True) def _pop_staged_vlm_media( self, @@ -745,13 +1063,27 @@ def initialize_weights( buffer_device: torch.device | None = None, dtype: torch.dtype = torch.bfloat16, ) -> None: + buffer_device = buffer_device or _default_init_device() + # Initialize the native text backbone (embed/norm/layers incl. the + # GatedDeltaNet/SSMGate-specific init). The HF vision tower + lm_head were + # initialized by ``super().__init__`` (or loaded from a checkpoint). + # HF's ``post_init`` (in ``super().__init__``) routes through + # ``init_weights -> initialize_weights`` *before* the backbone is swapped in, + # so ``language_model`` may still be the HF text model whose ``init_weights`` + # takes no ``buffer_device``; fall back to the no-arg HF signature then. + language_model = self.model.language_model + try: + language_model.init_weights(buffer_device=buffer_device) + except TypeError: + language_model.init_weights() mtp = getattr(self, "mtp", None) if mtp is not None: - buffer_device = buffer_device or _default_init_device() with buffer_device: for sublayer in mtp.layers: sublayer.init_weights(buffer_device=buffer_device) - cast_model_to_dtype(self, dtype) + # Keep the fp32 SSM-gating params fp32 (skip them in the dtype cast); each + # ``_fp32_params`` holder is sharded as its own fp32 FSDP group. + cast_model_to_dtype(self, dtype, skip_modules=("_fp32_params",)) ModelClass = Qwen3_5ForCausalLM 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 10e5eb3afb..80cc140d89 100644 --- a/nemo_automodel/components/models/qwen3_5/state_dict_adapter.py +++ b/nemo_automodel/components/models/qwen3_5/state_dict_adapter.py @@ -35,7 +35,8 @@ from nemo_automodel.components.checkpoint.state_dict_adapter import StateDictAdapter _FP32_PARAMS_TO_BARE = re.compile(r"(\.linear_attn)\._fp32_params\.") -_BARE_FP32_PARAM_NAMES = ("A_log",) +# Both SSM-gating params live in the fp32 ``SSMGate`` holder; route both on load. +_BARE_FP32_PARAM_NAMES = ("A_log", "dt_bias") _MTP_HF_TO_NATIVE = { "mtp.fc.weight": "mtp.layers.0.eh_proj.weight", "mtp.pre_fc_norm_embedding.weight": "mtp.layers.0.enorm.weight", diff --git a/nemo_automodel/components/models/qwen3_5_moe/cp_linear_attn.py b/nemo_automodel/components/models/qwen3_5_moe/cp_linear_attn.py index 3a7dc553b3..647c090aa7 100644 --- a/nemo_automodel/components/models/qwen3_5_moe/cp_linear_attn.py +++ b/nemo_automodel/components/models/qwen3_5_moe/cp_linear_attn.py @@ -34,18 +34,7 @@ from transformers.models.qwen3_5_moe.modeling_qwen3_5_moe import Qwen3_5MoeGatedDeltaNet from nemo_automodel.components.models.common.packing import get_unpad_data, is_indexed_packed_mask - - -def apply_model_runtime_patches(model, mesh=None): - """Apply Qwen3.5 runtime patches after model construction. - - The GatedDeltaNet wrapper is needed for both distributed training and - single-GPU packed-sequence runs, so it must run before sharding or first - forward rather than only from the FSDP parallelization strategy. - """ - cp_enabled = getattr(mesh, "cp_size", 1) > 1 - patch_hf_model(model, cp_enabled=cp_enabled) - return model +from nemo_automodel.shared.utils import dtype_from_str class _AllGatherConcatFn(Function): @@ -78,30 +67,59 @@ def backward(ctx, grad_output: torch.Tensor): return grad_local, None, None +class _SSMGateParam: + """Get-only (non-data) descriptor exposing an ``SSMGate`` param as an attribute. + + Lets ``self.A_log`` / ``self.dt_bias`` resolve to the fp32 ``SSMGate`` holder + (``self._fp32_params``) without a ``__getattr__`` monkeypatch. Being a non-data + descriptor, it does not intercept assignment, so HF's ``__init__`` doing + ``self.A_log = nn.Parameter(...)`` still routes through ``nn.Module.__setattr__`` + into ``_parameters`` (where it lives until ``install_ssm_gate`` moves it). + """ + + def __init__(self, name: str): + self.name = name + + def __get__(self, obj, owner=None): + if obj is None: + return self + return getattr(obj._fp32_params, self.name) + + class CPAwareGatedDeltaNet(Qwen3_5MoeGatedDeltaNet): """Drop-in replacement for ``Qwen3_5MoeGatedDeltaNet`` with FLA Context Parallelism. - All ``__init__`` parameters and weights are inherited unchanged from the HF - class. The only addition is ``_cp_mesh`` which is set externally by - ``apply_cp`` in the parallelizer. + The SSM-gating params (``A_log``/``dt_bias``) are moved into a fp32 ``SSMGate`` + submodule (``_fp32_params``) at construction so they keep fp32 storage (master + weights) even under a bf16 bulk dtype, and so FSDP can shard them in their own + dtype-uniform fp32 group. ``A_log``/``dt_bias`` remain readable as attributes via + get-only descriptors that resolve to the submodule — no ``__getattr__`` patch. + + ``_cp_mesh`` is set externally by the parallelizer to enable context parallelism. """ _cp_mesh: DeviceMesh | None + # Get-only (non-data) descriptors: reads resolve to the fp32 ``SSMGate`` holder, + # while writes during HF ``__init__`` (``self.A_log = nn.Parameter(...)``) still + # land in ``_parameters`` (handled by ``nn.Module.__setattr__``) before we move + # them into the holder. + A_log = _SSMGateParam("A_log") + dt_bias = _SSMGateParam("dt_bias") def __init__(self, config, layer_idx: int): super().__init__(config, layer_idx) self._cp_mesh = None + # HF created bare ``A_log``/``dt_bias`` in ``_parameters``; move them into a + # native fp32 ``SSMGate`` submodule (built directly, not relocated at runtime). + install_ssm_gate(self, fp32_dtype=_resolve_ssm_dtype(config)) def _compute_gate(self, a: torch.Tensor) -> torch.Tensor: - """Compute the gating value ``g`` using fp32 params. + """Compute the gating value ``g`` via the fp32 ``SSMGate`` submodule. - When ``_fp32_params`` exists (FSDP mixed-dtype), delegates to - the holder's forward so FSDP unshard/reshard lifecycle is natural. - Otherwise falls back to the inline computation. + Computing inside the submodule's forward keeps FSDP's unshard/reshard + lifecycle natural for the isolated fp32 group. """ - if hasattr(self, "_fp32_params"): - return self._fp32_params(a, self.dt_bias) - return -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias) + return self._fp32_params(a) def _forward_no_cp( self, @@ -589,144 +607,57 @@ def _forward_with_cp( return output -class _Fp32ParamHolder(torch.nn.Module): - """Holder for float32 params (A_log) that need a separate FSDP group. - - The ``forward`` computes the gating value ``g`` that HF's - ``Qwen3_5GatedDeltaNet.forward`` would normally compute inline. - By doing the computation *inside* this module's forward, FSDP's - unshard/reshard lifecycle works naturally — the params are - unsharded during the computation and resharded after. - """ +# SSM-gating params kept in fp32 storage (regardless of the model's bulk dtype) +# and isolated in the ``_fp32_params`` SSMGate submodule for FSDP. +_FP32_PARAM_NAMES = ("A_log", "dt_bias") - def forward(self, a: torch.Tensor, dt_bias: torch.Tensor) -> torch.Tensor: - return -self.A_log.float().exp() * F.softplus(a.float() + dt_bias) +class SSMGate(torch.nn.Module): + """Owns the fp32 SSM-gating params (``A_log``/``dt_bias``) and computes the gate. -def _make_fp32_getattr(orig_getattr): - """Create a ``__getattr__`` that resolves fp32 params from ``_fp32_params``. - - Allows ``self.A_log`` to resolve from the holder submodule so that - code outside forward (e.g. state_dict, checkpointing) can still - access the parameter by name. + Keeping these in a dedicated submodule lets FSDP shard them in their own + dtype-uniform fp32 group (true master weights), and computing the gate inside + ``forward`` keeps FSDP's unshard/reshard lifecycle natural. """ - def _getattr_with_fp32(self, name): - modules = self.__dict__.get("_modules", {}) - fp32_holder = modules.get("_fp32_params") - if fp32_holder is not None and name in fp32_holder._parameters: - return fp32_holder._parameters[name] - return orig_getattr(self, name) - - return _getattr_with_fp32 + def __init__(self, num_v_heads: int, dtype: torch.dtype = torch.float32): + super().__init__() + self.A_log = torch.nn.Parameter(torch.empty(num_v_heads, dtype=dtype)) + self.dt_bias = torch.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 patch_hf_model(model, cp_enabled=False): - """Patch HF Qwen3.5 GatedDeltaNet modules for FSDP and optional CP support. - - For FSDP compatibility, move float32 bare params (A_log) into a - ``_fp32_params`` submodule so ``fully_shard_by_dtype`` can wrap them - in a separate FSDP group. - Every ``Qwen3_5GatedDeltaNet`` instance's ``__class__`` is swapped to - ``CPAwareGatedDeltaNet`` whose ``forward()`` calls ``self._fp32_params()`` - to trigger FSDP unshard before accessing the fp32 params. When - ``cp_enabled=True``, the CP mesh is also configured. +def install_ssm_gate(mod, fp32_dtype=torch.float32): + """Move ``mod``'s HF-created bare ``A_log``/``dt_bias`` into a fp32 ``SSMGate``. - Additionally, every ``Qwen3_5DecoderLayer`` instance is class-swapped to - ``Qwen3_5DecoderLayerWithPacking`` so that NEAT-packed sequence metadata - (``cu_seqlens``, ``indices``, ``position_ids``) reaches ``linear_attn`` - via real keyword arguments instead of relying on instance-attribute - side-channels (issue #2131). + HF's GatedDeltaNet ``__init__`` creates ``A_log``/``dt_bias`` as bare params in + ``mod._parameters``. This relocates them into an :class:`SSMGate` submodule + registered as ``_fp32_params`` (casting to ``fp32_dtype``), so they keep fp32 + storage under a bf16 bulk dtype and get their own dtype-uniform FSDP group. + Attribute access (``self.A_log``/``self.dt_bias``) continues to work via the + :class:`_SSMGateParam` descriptors on ``CPAwareGatedDeltaNet`` — no + ``__getattr__`` patch. Returns the gate submodule. """ - import logging - - try: - from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5GatedDeltaNet - except ImportError: - return - - try: - from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5DecoderLayer - - from nemo_automodel.components.models.qwen3_5.decoder_layer import Qwen3_5DecoderLayerWithPacking - except (AttributeError, ImportError): - Qwen3_5DecoderLayer = None - Qwen3_5DecoderLayerWithPacking = None - - _logger = logging.getLogger(__name__) - _PATCHED_ATTR = "_fp32_getattr_patched" - patched = 0 - patched_classes = set() - for name, mod in model.named_modules(): - # Class-swap decoder layers so their forward threads packing kwargs - # into linear_attn. Doing this before the GatedDeltaNet pass means - # the swap is independent of which (if any) inner layer is patched. - if ( - Qwen3_5DecoderLayer is not None - and isinstance(mod, Qwen3_5DecoderLayer) - and not isinstance(mod, Qwen3_5DecoderLayerWithPacking) - and getattr(mod, "layer_type", None) == "linear_attention" - ): - mod.__class__ = Qwen3_5DecoderLayerWithPacking - - if not isinstance(mod, Qwen3_5GatedDeltaNet): - continue - - mod.__class__ = CPAwareGatedDeltaNet - mod._cp_mesh = None - - # Move float32 bare params into a holder submodule for FSDP. - # The CPAwareGatedDeltaNet forward calls self._fp32_params() - # to trigger FSDP unshard; __getattr__ redirects self.A_log - # to the holder so it returns the unsharded plain tensor. - holder = None - for pname in list(mod._parameters.keys()): - param = mod._parameters[pname] - if param is not None and param.dtype == torch.float32: - if holder is None: - holder = _Fp32ParamHolder() - setattr(holder, pname, param) - del mod._parameters[pname] - if holder is not None: - mod.add_module("_fp32_params", holder) - - # Guard against re-wrapping __getattr__ on repeated calls. - cls = type(mod) - if cls not in patched_classes and not getattr(cls, _PATCHED_ATTR, False): - cls.__getattr__ = _make_fp32_getattr(cls.__getattr__) - setattr(cls, _PATCHED_ATTR, True) - patched_classes.add(cls) - patched += 1 - - if patched > 0: - _logger.info( - "Patched %d GatedDeltaNet modules (cp=%s) with FSDP-safe fp32 param wrapping.", - patched, - cp_enabled, - ) - - # Declare the fp32-compute submodules so fully_shard_by_dtype keeps them in - # fp32 compute even under fp32 master weights (bf16 compute for the bulk). The - # ``_fp32_params`` holder created above is the authoritative marker; the - # GatedDeltaNet forward uses those params directly without re-upcasting. - existing = tuple(getattr(model, "_keep_in_fp32_modules_strict", None) or ()) - if "_fp32_params" not in existing: - model._keep_in_fp32_modules_strict = existing + ("_fp32_params",) - - # Attach a Qwen3.5 state_dict_adapter so saved checkpoints hide the - # ``_fp32_params`` wrapping and remain HF-loadable directly. Keep - # ``from_hf`` in bare-key mode: adapter-mediated DCP loads operate in - # the HF key namespace even when physical params live in ``_fp32_params``. - # Use dynamic import to avoid pulling ``components.checkpoint`` into - # this file's static import graph: ``cp_linear_attn`` is reached from - # ``components.distributed.parallelizer`` and the adapter inherits - # from ``components.checkpoint.state_dict_adapter``, which would - # otherwise create a forbidden ``distributed -> checkpoint`` chain - # under the import-linter ``independence`` contract. - import importlib - - adapter_module = importlib.import_module("nemo_automodel.components.models.qwen3_5.state_dict_adapter") - adapter = getattr(model, "state_dict_adapter", None) - if adapter is None: - model.state_dict_adapter = adapter_module.Qwen3_5DenseStateDictAdapter(route_linear_attn_fp32_params=False) + num_v_heads = mod._parameters["A_log"].shape[0] + gate = SSMGate(num_v_heads, dtype=fp32_dtype) + for pname in _FP32_PARAM_NAMES: + param = mod._parameters.pop(pname) + if param.dtype != fp32_dtype: + param.data = param.data.to(fp32_dtype) + setattr(gate, pname, param) # overwrite the freshly-built empty param + mod.add_module("_fp32_params", gate) + return gate + + +def _resolve_ssm_dtype(config): + """Resolve the fp32 storage dtype for the SSM-gating params from ``config``. + + Honors ``mamba_ssm_dtype`` (Qwen3.5 stores ``A_log``/``dt_bias`` in fp32); + defaults to ``torch.float32``. + """ + ssm_dtype = getattr(config, "mamba_ssm_dtype", None) + if isinstance(ssm_dtype, str): + ssm_dtype = dtype_from_str(ssm_dtype) + return ssm_dtype or torch.float32 diff --git a/nemo_automodel/components/models/qwen3_5_moe/model.py b/nemo_automodel/components/models/qwen3_5_moe/model.py index 87a5750637..6a86217c24 100644 --- a/nemo_automodel/components/models/qwen3_5_moe/model.py +++ b/nemo_automodel/components/models/qwen3_5_moe/model.py @@ -786,6 +786,16 @@ def __init__( # Expose moe_config for FSDP sync mixin self.model.moe_config = self.model.language_model.moe_config + # Keep the SSM-gating params (A_log/dt_bias) — isolated in each + # linear_attn ``_fp32_params`` holder at construction — in fp32 storage + # even when the model's bulk dtype is bf16, matching the intrinsically-fp32 + # dtype Qwen3.5 checkpoints store them in. cast_model_to_dtype() (called + # from initialize_weights) honors this list. + 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 + self.vocab_size = text_config.vocab_size pad_token_id = getattr(text_config, "pad_token_id", None) self.pad_token_id = pad_token_id if pad_token_id is not None else -1 @@ -1073,7 +1083,11 @@ def initialize_weights( for sublayer in mtp.layers: sublayer.init_weights(buffer_device=buffer_device) - cast_model_to_dtype(self, dtype) + # Skip the SSM-gating holders so they keep fp32 storage (master weights): + # cast_model_to_dtype cannot reliably restore fp32 once FSDP2-sharded, so it + # detaches them and never casts them. Each holder is its own fp32 FSDP group + # (moe/parallelizer._shard_fp32_param_holders), so this is dtype-uniform-safe. + cast_model_to_dtype(self, dtype, skip_modules=("_fp32_params",)) with buffer_device: self.model.language_model.rotary_emb.device = buffer_device 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 3c3aa856e2..66e932d1f4 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 @@ -52,6 +52,24 @@ 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) + + +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}" + class Qwen3_5MoeStateDictAdapter(StateDictAdapter): """Converts between HF Qwen3.5-MoE checkpoints and the NeMo native format. @@ -260,7 +278,8 @@ def from_hf( down_tensor, device_mesh, rank ) - return state_dict + # 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()} 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.""" @@ -289,6 +308,7 @@ def convert_single_tensor_to_hf(self, fqn: str, tensor: Any, **kwargs) -> list[t break new_fqn = map_qwen3_5_mtp_to_hf_key(new_fqn) + new_fqn = _strip_fp32_params(new_fqn) 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 5f14b8a893..ec4c64c7e8 100644 --- a/nemo_automodel/components/models/qwen3_next/layers.py +++ b/nemo_automodel/components/models/qwen3_next/layers.py @@ -118,14 +118,17 @@ def forward( k = self.k_proj(x) v = self.v_proj(x) + # Use -1 for the head dim so the (possibly tensor-parallel-sharded) local + # head count is inferred from the projection output. Without TP this equals + # the full head count; with TP (colwise q/k/v_proj) it is heads // tp_size. if qkv_format == "thd": - q = q.view(num_tokens, self.num_heads, self.head_dim * 2) - k = k.view(num_tokens, self.num_kv_heads, self.head_dim) - v = v.view(num_tokens, self.num_kv_heads, self.head_dim) + q = q.view(num_tokens, -1, self.head_dim * 2) + k = k.view(num_tokens, -1, self.head_dim) + v = v.view(num_tokens, -1, self.head_dim) else: - q = q.view(bsz, seqlen, self.num_heads, self.head_dim * 2) - k = k.view(bsz, seqlen, self.num_kv_heads, self.head_dim) - v = v.view(bsz, seqlen, self.num_kv_heads, self.head_dim) + q = q.view(bsz, seqlen, -1, self.head_dim * 2) + k = k.view(bsz, seqlen, -1, self.head_dim) + v = v.view(bsz, seqlen, -1, self.head_dim) q, gate = torch.chunk(q, 2, dim=-1) gate = gate.reshape(*x.shape[:-1], -1) diff --git a/nemo_automodel/components/models/qwen3_next/model.py b/nemo_automodel/components/models/qwen3_next/model.py index 17ee28a0be..5fc1fd015a 100644 --- a/nemo_automodel/components/models/qwen3_next/model.py +++ b/nemo_automodel/components/models/qwen3_next/model.py @@ -101,10 +101,14 @@ def forward( return x def _mlp(self, x: torch.Tensor, padding_mask: torch.Tensor | None) -> torch.Tensor: - if isinstance(self.mlp, MLP): + # ``self.mlp`` may be wrapped by activation checkpointing (submodule-level + # AC), so inspect the underlying module to pick the dense (no padding_mask) + # vs MoE (padding_mask) call signature, but invoke the wrapped module. + mlp = getattr(self.mlp, "_checkpoint_wrapped_module", self.mlp) + if isinstance(mlp, MLP): return self.mlp(x) else: - assert isinstance(self.mlp, MoE) + assert isinstance(mlp, MoE) return self.mlp(x, padding_mask) def init_weights(self, buffer_device: torch.device): diff --git a/nemo_automodel/components/moe/parallelizer.py b/nemo_automodel/components/moe/parallelizer.py index ad3512318d..8d9dc42daa 100644 --- a/nemo_automodel/components/moe/parallelizer.py +++ b/nemo_automodel/components/moe/parallelizer.py @@ -294,6 +294,44 @@ def selective_checkpointing_context_fn(): parent_layers.register_module(layer_id, block) +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. + + 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 + stubs) cannot hold fp32 holders, so an empty set is returned. + """ + if not hasattr(block, "named_modules"): + return set() + fp32_mp_policy = MixedPrecisionPolicy( + param_dtype=torch.float32, + reduce_dtype=torch.float32, + output_dtype=torch.float32, + cast_forward_inputs=False, + ) + ignored: set = set() + for name, sub in block.named_modules(): + if not name.endswith("_fp32_params"): + continue + holder_params = list(sub.parameters(recurse=False)) + if not holder_params: + continue + fully_shard( + sub, + mesh=fsdp_mesh, + reshard_after_forward=reshard_after_forward, + mp_policy=fp32_mp_policy, + offload_policy=offload_policy, + ) + ignored.update(holder_params) + return ignored + + def apply_fsdp( model: torch.nn.Module, fsdp_mesh: DeviceMesh, @@ -365,6 +403,14 @@ def apply_fsdp( ignored_params = None 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. + 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 fully_shard_default(block, ignored_params=ignored_params) if hasattr(_model, "embed_tokens") and _model.embed_tokens is not None: diff --git a/tests/unit_tests/_transformers/test_auto_model.py b/tests/unit_tests/_transformers/test_auto_model.py index c0da3faa05..21bea1fe9f 100644 --- a/tests/unit_tests/_transformers/test_auto_model.py +++ b/tests/unit_tests/_transformers/test_auto_model.py @@ -354,9 +354,13 @@ def __init__(self, architectures=None): self.config = types.SimpleNamespace(architectures=architectures) def test_apply_model_runtime_patches_dispatches_by_architecture(self): + # The registry mechanism is exercised with a temporary test entry — the + # built-in registry no longer ships any entries (Qwen3.5 builds its + # CP/fp32-gate modules at construction instead of patching at load time). + import nemo_automodel._transformers.kernel_patches as kp from nemo_automodel._transformers.kernel_patches import apply_model_runtime_patches - model = self._DummyModel(["Qwen3_5ForCausalLM"]) + model = self._DummyModel(["FakeArchForCausalLM"]) mesh = types.SimpleNamespace(cp_size=1) calls = [] @@ -365,20 +369,26 @@ def fake_hook(model, mesh): return model fake_module = types.SimpleNamespace(apply_model_runtime_patches=fake_hook) + test_registry = {"FakeArchForCausalLM": ("fake.module.path", "apply_model_runtime_patches")} - with patch( - "nemo_automodel._transformers.kernel_patches.importlib.import_module", - return_value=fake_module, - ) as mock_import: + with ( + patch.object(kp, "_MODEL_RUNTIME_PATCHES", test_registry), + patch( + "nemo_automodel._transformers.kernel_patches.importlib.import_module", + return_value=fake_module, + ) as mock_import, + ): assert apply_model_runtime_patches(model, mesh) is model - mock_import.assert_called_once_with("nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn") + mock_import.assert_called_once_with("fake.module.path") assert calls == [(model, mesh)] def test_apply_model_runtime_patches_deduplicates_hook_specs(self): + # Two architectures sharing one hook spec must invoke the hook once. + import nemo_automodel._transformers.kernel_patches as kp from nemo_automodel._transformers.kernel_patches import apply_model_runtime_patches - model = self._DummyModel(["Qwen3_5ForCausalLM", "Qwen3_5ForConditionalGeneration"]) + model = self._DummyModel(["FakeArchA", "FakeArchB"]) mesh = types.SimpleNamespace(cp_size=2) calls = [] @@ -387,10 +397,15 @@ def fake_hook(model, mesh): return model fake_module = types.SimpleNamespace(apply_model_runtime_patches=fake_hook) + shared_spec = ("fake.module.path", "apply_model_runtime_patches") + test_registry = {"FakeArchA": shared_spec, "FakeArchB": shared_spec} - with patch( - "nemo_automodel._transformers.kernel_patches.importlib.import_module", - return_value=fake_module, + with ( + patch.object(kp, "_MODEL_RUNTIME_PATCHES", test_registry), + patch( + "nemo_automodel._transformers.kernel_patches.importlib.import_module", + return_value=fake_module, + ), ): assert apply_model_runtime_patches(model, mesh) is model diff --git a/tests/unit_tests/_transformers/test_capabilities_hybrid_vlm.py b/tests/unit_tests/_transformers/test_capabilities_hybrid_vlm.py index a88d2753e5..3716df64c9 100644 --- a/tests/unit_tests/_transformers/test_capabilities_hybrid_vlm.py +++ b/tests/unit_tests/_transformers/test_capabilities_hybrid_vlm.py @@ -135,3 +135,37 @@ def test_is_hybrid_inner_language_model_with_no_config(): inner_lm = SimpleNamespace(config=None) model = _DummyModel(config=outer_cfg, language_model=inner_lm) assert _is_hybrid(model) is False + + +def test_is_hybrid_via_layer_types_linear_attention(): + """Qwen3.5 / Qwen3-Next style: per-layer ``layer_types`` mixing + ``linear_attention`` with ``full_attention`` marks the model hybrid.""" + cfg = SimpleNamespace(layer_types=["full_attention", "linear_attention", "full_attention"]) + model = _DummyModel(config=cfg) + assert _is_hybrid(model) is True + + +def test_not_hybrid_when_layer_types_all_full_attention(): + """``layer_types`` present but with no ``linear_attention`` is not hybrid.""" + cfg = SimpleNamespace(layer_types=["full_attention", "full_attention"]) + model = _DummyModel(config=cfg) + assert _is_hybrid(model) is False + + +def test_is_hybrid_via_nested_text_config_layer_types(): + """VLM configs nest the decoder config under ``text_config``; the hybrid + marker (``layer_types``) lives there, so ``_is_hybrid`` must drill into it.""" + text_cfg = SimpleNamespace(layer_types=["full_attention", "linear_attention"]) + outer_cfg = SimpleNamespace(text_config=text_cfg) + model = _DummyModel(config=outer_cfg) + assert _is_hybrid(model) is True + + +def test_is_hybrid_via_language_model_nested_text_config(): + """language_model.config.text_config carries the hybrid ``layer_types``.""" + text_cfg = SimpleNamespace(layer_types=["linear_attention"]) + inner_cfg = SimpleNamespace(text_config=text_cfg) + inner_lm = SimpleNamespace(config=inner_cfg) + outer_cfg = SimpleNamespace() + model = _DummyModel(config=outer_cfg, language_model=inner_lm) + assert _is_hybrid(model) is True 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 6eeae1468d..0e3a425291 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 @@ -240,6 +240,52 @@ def test_fp16_dtype(self): for p in model.parameters(): assert p.dtype == torch.float16 + def test_skip_modules_left_untouched(self): + """Submodules named in ``skip_modules`` keep their original dtype.""" + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4) + self._fp32_params = nn.Linear(4, 4) + + model = Model() + cast_model_to_dtype(model, torch.bfloat16, skip_modules=("_fp32_params",)) + + # Regular submodule is cast; the skipped holder stays fp32. + assert model.linear.weight.dtype == torch.bfloat16 + assert model._fp32_params.weight.dtype == torch.float32 + + def test_skip_modules_nested_and_restored(self): + """Nested skip_modules are preserved and re-attached after the cast.""" + + class Inner(nn.Module): + def __init__(self): + super().__init__() + self._fp32_params = nn.Linear(2, 2) + self.proj = nn.Linear(2, 2) + + class Model(nn.Module): + def __init__(self): + super().__init__() + self.block = Inner() + + model = Model() + cast_model_to_dtype(model, torch.bfloat16, skip_modules=("_fp32_params",)) + + assert model.block.proj.weight.dtype == torch.bfloat16 + # Holder preserved in fp32 and re-attached (still reachable on the module). + assert model.block._fp32_params.weight.dtype == torch.float32 + assert model.block._fp32_params is dict(model.block.named_modules())["_fp32_params"] + + def test_skip_modules_empty_is_noop(self): + """An empty skip_modules tuple casts everything (default behavior).""" + model = SimpleModel() + cast_model_to_dtype(model, torch.bfloat16, skip_modules=()) + + for p in model.parameters(): + assert p.dtype == torch.bfloat16 + def test_set_valued_keep_in_fp32_preserved(self): # Mirrors HF converting _keep_in_fp32_modules (list) to a set on the instance — # the gemma4_moe/diffusion_gemma case. cast_model_to_dtype must still restore it. diff --git a/tests/unit_tests/models/qwen3_5/test_cp_linear_attn_patch.py b/tests/unit_tests/models/qwen3_5/test_cp_linear_attn_patch.py deleted file mode 100644 index 8e88ce84da..0000000000 --- a/tests/unit_tests/models/qwen3_5/test_cp_linear_attn_patch.py +++ /dev/null @@ -1,1210 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Tests for Qwen3.5 dense CP + FSDP mixed-dtype patching.""" - -from __future__ import annotations - -import sys -import types -from unittest.mock import MagicMock, patch - -import pytest -import torch -import torch.nn as nn - - -class _FakeGatedDeltaNet(nn.Module): - """Mimics HF Qwen3_5GatedDeltaNet with mixed-dtype bare params.""" - - def __init__(self): - super().__init__() - self.A_log = nn.Parameter(torch.ones(4, dtype=torch.float32)) - self.dt_bias = nn.Parameter(torch.ones(4, dtype=torch.bfloat16)) - self.conv1d = nn.Conv1d(4, 4, 1) - self.norm = nn.LayerNorm(4) - # Force norm to float32 - self.norm.weight.data = self.norm.weight.data.float() - self.norm.bias.data = self.norm.bias.data.float() - self.layer_idx = 0 - - -@pytest.fixture() -def fake_model(): - """Build a minimal model with a fake GatedDeltaNet layer.""" - model = nn.Module() - model.layers = nn.ModuleList([nn.Module()]) - model.layers[0].linear_attn = _FakeGatedDeltaNet() - model.layers[0].layer_type = "linear_attention" - return model - - -class TestPatchHfModel: - @staticmethod - def _stub_qwen3_5_modules(monkeypatch): - """Stub transformers.models.qwen3_5* so cp_linear_attn can be imported.""" - for path in ( - "transformers.models.qwen3_5_moe", - "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe", - "transformers.models.qwen3_5", - "transformers.models.qwen3_5.modeling_qwen3_5", - ): - stub = types.ModuleType(path) - stub.Qwen3_5MoeGatedDeltaNet = _FakeGatedDeltaNet - stub.Qwen3_5GatedDeltaNet = _FakeGatedDeltaNet - monkeypatch.setitem(sys.modules, path, stub) - - def test_fp32_params_moved_to_holder(self, fake_model, monkeypatch): - """Float32 bare params are moved into _fp32_params submodule via real patch_hf_model.""" - self._stub_qwen3_5_modules(monkeypatch) - - # Remove cached cp_linear_attn so re-import picks up our stubs - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - - from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import patch_hf_model - - la = fake_model.layers[0].linear_attn - assert la.A_log.dtype == torch.float32 - assert la.dt_bias.dtype == torch.bfloat16 - - patch_hf_model(fake_model, cp_enabled=False) - - # A_log (float32) should be moved out of _parameters - assert "A_log" not in la._parameters - # Accessed via __getattr__ → _fp32_params - assert la.A_log.dtype == torch.float32 - # dt_bias (bfloat16) stays as a regular parameter - assert "dt_bias" in la._parameters - # _fp32_params submodule holds the moved param - assert hasattr(la, "_fp32_params") - assert la._fp32_params.A_log.dtype == torch.float32 - # __getattr__ resolves to the same tensor in _fp32_params - assert la.A_log is la._fp32_params.A_log - - def test_class_always_swapped_for_fsdp(self, fake_model, monkeypatch): - """Class is always swapped to CPAwareGatedDeltaNet for FSDP fp32 unshard support.""" - self._stub_qwen3_5_modules(monkeypatch) - - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - - from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import ( - CPAwareGatedDeltaNet, - patch_hf_model, - ) - - la = fake_model.layers[0].linear_attn - patch_hf_model(fake_model, cp_enabled=False) - assert type(la) is CPAwareGatedDeltaNet - assert la._cp_mesh is None - - def test_class_swap_when_cp_enabled(self, fake_model, monkeypatch): - """With cp_enabled=True, class is swapped to CPAwareGatedDeltaNet.""" - self._stub_qwen3_5_modules(monkeypatch) - - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - - from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import ( - CPAwareGatedDeltaNet, - patch_hf_model, - ) - - la = fake_model.layers[0].linear_attn - patch_hf_model(fake_model, cp_enabled=True) - assert type(la) is CPAwareGatedDeltaNet - assert la._cp_mesh is None - - def test_getattr_resolves_after_param_replacement(self, fake_model, monkeypatch): - """__getattr__ resolves to _fp32_params even after the underlying tensor is replaced.""" - self._stub_qwen3_5_modules(monkeypatch) - - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - - from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import patch_hf_model - - la = fake_model.layers[0].linear_attn - patch_hf_model(fake_model, cp_enabled=False) - - # Simulate FSDP replacing the parameter in _fp32_params - new_tensor = nn.Parameter(torch.zeros(4, dtype=torch.float32)) - la._fp32_params._parameters["A_log"] = new_tensor - - # __getattr__ should resolve to the NEW tensor, not the old one - assert la.A_log is new_tensor - - def test_apply_model_runtime_patches_uses_mesh_cp_size(self, fake_model, monkeypatch): - """Runtime hook maps MeshContext cp_size to patch_hf_model cp_enabled.""" - self._stub_qwen3_5_modules(monkeypatch) - - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - - import nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn as cp_linear_attn - - mesh = types.SimpleNamespace(cp_size=2) - - with patch.object(cp_linear_attn, "patch_hf_model") as mock_patch: - assert cp_linear_attn.apply_model_runtime_patches(fake_model, mesh=mesh) is fake_model - - mock_patch.assert_called_once_with(fake_model, cp_enabled=True) - - -class TestFp32ParamHolder: - """Tests for _Fp32ParamHolder forward (gate computation).""" - - @staticmethod - def _stub_and_import(monkeypatch): - for path in ( - "transformers.models.qwen3_5_moe", - "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe", - "transformers.models.qwen3_5", - "transformers.models.qwen3_5.modeling_qwen3_5", - ): - stub = types.ModuleType(path) - stub.Qwen3_5MoeGatedDeltaNet = _FakeGatedDeltaNet - stub.Qwen3_5GatedDeltaNet = _FakeGatedDeltaNet - monkeypatch.setitem(sys.modules, path, stub) - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import _Fp32ParamHolder - - return _Fp32ParamHolder - - def test_holder_forward_computes_gate(self, monkeypatch): - """_Fp32ParamHolder.forward returns g = -A_log.exp() * softplus(a + dt_bias).""" - _Fp32ParamHolder = self._stub_and_import(monkeypatch) - holder = _Fp32ParamHolder() - holder.A_log = nn.Parameter(torch.ones(4, dtype=torch.float32)) - a = torch.zeros(4) - dt_bias = torch.zeros(4) - g = holder(a, dt_bias) - # g = -exp(1) * softplus(0 + 0) = -e * softplus(0) = -e * ln(2) - expected = -torch.ones(4).float().exp() * torch.nn.functional.softplus(torch.zeros(4)) - assert torch.allclose(g, expected, atol=1e-5) - - def test_holder_forward_dtype_is_float32(self, monkeypatch): - """Gate computation happens in float32 even with bfloat16 inputs.""" - _Fp32ParamHolder = self._stub_and_import(monkeypatch) - holder = _Fp32ParamHolder() - holder.A_log = nn.Parameter(torch.ones(4, dtype=torch.float32)) - a = torch.zeros(4, dtype=torch.bfloat16) - dt_bias = torch.zeros(4, dtype=torch.bfloat16) - g = holder(a, dt_bias) - assert g.dtype == torch.float32 - - -class TestComputeGate: - """Tests for CPAwareGatedDeltaNet._compute_gate routing.""" - - @staticmethod - def _stub_and_import(monkeypatch): - for path in ( - "transformers.models.qwen3_5_moe", - "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe", - "transformers.models.qwen3_5", - "transformers.models.qwen3_5.modeling_qwen3_5", - ): - stub = types.ModuleType(path) - stub.Qwen3_5MoeGatedDeltaNet = _FakeGatedDeltaNet - stub.Qwen3_5GatedDeltaNet = _FakeGatedDeltaNet - monkeypatch.setitem(sys.modules, path, stub) - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import ( - CPAwareGatedDeltaNet, - _Fp32ParamHolder, - patch_hf_model, - ) - - return CPAwareGatedDeltaNet, _Fp32ParamHolder, patch_hf_model - - def test_compute_gate_routes_through_holder(self, fake_model, monkeypatch): - """_compute_gate calls _fp32_params.forward when holder exists.""" - CPAwareGatedDeltaNet, _Fp32ParamHolder, patch_hf_model = self._stub_and_import(monkeypatch) - patch_hf_model(fake_model, cp_enabled=False) - la = fake_model.layers[0].linear_attn - assert isinstance(la, CPAwareGatedDeltaNet) - a = torch.zeros(4) - with patch.object(la._fp32_params, "forward", return_value=torch.zeros(4)) as mock_fwd: - la._compute_gate(a) - mock_fwd.assert_called_once() - - def test_compute_gate_fallback_without_holder(self, fake_model, monkeypatch): - """_compute_gate falls back to inline computation without _fp32_params.""" - CPAwareGatedDeltaNet, _, patch_hf_model = self._stub_and_import(monkeypatch) - la = fake_model.layers[0].linear_attn - la.__class__ = CPAwareGatedDeltaNet - la._cp_mesh = None - # No _fp32_params — A_log is still in _parameters - a = torch.zeros(4) - g = la._compute_gate(a) - assert g.dtype == torch.float32 - assert g.shape == (4,) - - -class TestPatchHfModelSentinel: - """Test that __getattr__ patching is idempotent.""" - - @staticmethod - def _stub_and_import(monkeypatch): - for path in ( - "transformers.models.qwen3_5_moe", - "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe", - "transformers.models.qwen3_5", - "transformers.models.qwen3_5.modeling_qwen3_5", - ): - stub = types.ModuleType(path) - stub.Qwen3_5MoeGatedDeltaNet = _FakeGatedDeltaNet - stub.Qwen3_5GatedDeltaNet = _FakeGatedDeltaNet - monkeypatch.setitem(sys.modules, path, stub) - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import patch_hf_model - - return patch_hf_model - - def test_double_patch_does_not_rewrap_getattr(self, monkeypatch): - """Calling patch_hf_model twice does not chain __getattr__ wrappers.""" - patch_hf_model = self._stub_and_import(monkeypatch) - - model1 = nn.Module() - model1.layers = nn.ModuleList([nn.Module()]) - model1.layers[0].linear_attn = _FakeGatedDeltaNet() - model1.layers[0].layer_type = "linear_attention" - patch_hf_model(model1, cp_enabled=False) - getattr_after_first = type(model1.layers[0].linear_attn).__getattr__ - - model2 = nn.Module() - model2.layers = nn.ModuleList([nn.Module()]) - model2.layers[0].linear_attn = _FakeGatedDeltaNet() - model2.layers[0].layer_type = "linear_attention" - patch_hf_model(model2, cp_enabled=False) - getattr_after_second = type(model2.layers[0].linear_attn).__getattr__ - - # Same function, not a wrapper of a wrapper - assert getattr_after_first is getattr_after_second - - -class TestPatchHfModelStateDictAdapter: - """Verify that patch_hf_model attaches a state_dict_adapter for HF-format saves.""" - - @staticmethod - def _stub_and_import(monkeypatch): - for path in ( - "transformers.models.qwen3_5_moe", - "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe", - "transformers.models.qwen3_5", - "transformers.models.qwen3_5.modeling_qwen3_5", - ): - stub = types.ModuleType(path) - stub.Qwen3_5MoeGatedDeltaNet = _FakeGatedDeltaNet - stub.Qwen3_5GatedDeltaNet = _FakeGatedDeltaNet - monkeypatch.setitem(sys.modules, path, stub) - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import patch_hf_model - - return patch_hf_model - - def test_attaches_dense_state_dict_adapter(self, fake_model, monkeypatch): - """After patching, model.state_dict_adapter rewrites _fp32_params keys for HF saves.""" - from nemo_automodel.components.models.qwen3_5.state_dict_adapter import ( - Qwen3_5DenseStateDictAdapter, - ) - - patch_hf_model = self._stub_and_import(monkeypatch) - assert not hasattr(fake_model, "state_dict_adapter") - - patch_hf_model(fake_model, cp_enabled=False) - - adapter = getattr(fake_model, "state_dict_adapter", None) - assert isinstance(adapter, Qwen3_5DenseStateDictAdapter) - assert not adapter.route_linear_attn_fp32_params - # Smoke-check round-trip behaviour on a sample state dict. - out = adapter.to_hf({"layers.0.linear_attn._fp32_params.A_log": torch.zeros(4)}) - assert list(out.keys()) == ["layers.0.linear_attn.A_log"] - - def test_preserves_existing_qwen3_5_adapter(self, fake_model, monkeypatch): - """If a Qwen3.5 adapter already exists, CP patching preserves its load mode.""" - from nemo_automodel.components.models.qwen3_5.state_dict_adapter import ( - Qwen3_5DenseStateDictAdapter, - ) - - patch_hf_model = self._stub_and_import(monkeypatch) - adapter = Qwen3_5DenseStateDictAdapter(route_linear_attn_fp32_params=False) - fake_model.state_dict_adapter = adapter - - patch_hf_model(fake_model, cp_enabled=False) - - assert fake_model.state_dict_adapter is adapter - assert not adapter.route_linear_attn_fp32_params - - def test_does_not_overwrite_unrelated_existing_adapter(self, fake_model, monkeypatch): - """If a non-Qwen3.5 model already has a state_dict_adapter, it is preserved.""" - patch_hf_model = self._stub_and_import(monkeypatch) - sentinel = object() - fake_model.state_dict_adapter = sentinel - - patch_hf_model(fake_model, cp_enabled=False) - - assert fake_model.state_dict_adapter is sentinel - - def test_no_adapter_when_no_layers_patched(self, monkeypatch): - """Adapter is only attached when at least one GatedDeltaNet layer was patched.""" - patch_hf_model = self._stub_and_import(monkeypatch) - - model = nn.Module() # no GatedDeltaNet children at all - patch_hf_model(model, cp_enabled=False) - - assert not hasattr(model, "state_dict_adapter") - - -class TestPackingHelpers: - """Tests for the indexed-mask helpers used by Qwen3_5DecoderLayerWithPacking.""" - - def test_is_indexed_packed_mask_detection(self): - from nemo_automodel.components.models.common.packing import is_indexed_packed_mask - - assert is_indexed_packed_mask(None) is False - assert is_indexed_packed_mask(torch.ones(1, 4, dtype=torch.long)) is False - assert is_indexed_packed_mask(torch.tensor([[1, 1, 0, 0]])) is False # 0/1 only - assert is_indexed_packed_mask(torch.tensor([[1, 1, 2, 2]])) is True - # bool dtype is short-circuited (a bool 1/2 mask isn't a thing). - assert is_indexed_packed_mask(torch.tensor([[True, True, False, False]])) is False - - def test_cu_seqlens_from_indexed_mask(self): - from nemo_automodel.components.models.common.packing import get_unpad_data - - mask = torch.tensor([[1, 1, 2, 2, 2, 0], [1, 1, 1, 1, 0, 0]]) - indices, cu_seqlens, max_seqlen = get_unpad_data(mask) - # Per-doc lengths flattened across batch: [2, 3, 4] - assert cu_seqlens.tolist() == [0, 2, 5, 9] - assert max_seqlen == 4 - # Non-padding positions in flattened B*T=12 sequence - assert indices.tolist() == [0, 1, 2, 3, 4, 6, 7, 8, 9] - - def test_dense_decoder_uses_packed_seq_ids_for_sdpa_linear_attention(self): - """Linear attention gets indexed packed ids even when full attention uses a 4D SDPA mask.""" - from nemo_automodel.components.models.qwen3_5.decoder_layer import Qwen3_5DecoderLayerWithPacking - - class RecorderLinearAttn(nn.Module): - layer_idx = 0 - - def __init__(self): - super().__init__() - self.called_with = None - - def forward(self, **kwargs): - self.called_with = kwargs - return kwargs["hidden_states"] - - layer = Qwen3_5DecoderLayerWithPacking.__new__(Qwen3_5DecoderLayerWithPacking) - nn.Module.__init__(layer) - layer.layer_type = "linear_attention" - layer.input_layernorm = nn.Identity() - layer.linear_attn = RecorderLinearAttn() - layer.post_attention_layernorm = nn.Identity() - layer.mlp = nn.Identity() - - hidden_states = torch.zeros(1, 5, 4) - sdpa_mask = torch.ones(1, 1, 5, 5, dtype=torch.bool).tril() - packed_seq_ids = torch.tensor([[1, 1, 2, 2, 2]]) - - layer( - hidden_states, - position_embeddings=(torch.empty(0), torch.empty(0)), - attention_mask=sdpa_mask, - position_ids=torch.arange(5).unsqueeze(0), - _packed_seq_ids=packed_seq_ids, - ) - - called = layer.linear_attn.called_with - assert called["attention_mask"] is packed_seq_ids - assert called["cu_seqlens"].tolist() == [0, 2, 5] - assert called["indices"].tolist() == [0, 1, 2, 3, 4] - - -class TestQwen35ParallelizationStrategyRegistration: - def test_strategy_registered(self): - """Qwen3.5 model classes are in the strategy registry.""" - from nemo_automodel.components.distributed.parallelizer import PARALLELIZATION_STRATEGIES - - assert "Qwen3_5ForConditionalGeneration" in PARALLELIZATION_STRATEGIES - assert "Qwen3_5ForCausalLM" in PARALLELIZATION_STRATEGIES - - def test_strategy_type(self): - """Strategy is Qwen3_5ParallelizationStrategy.""" - from nemo_automodel.components.distributed.parallelizer import ( - PARALLELIZATION_STRATEGIES, - Qwen3_5ParallelizationStrategy, - ) - - assert isinstance(PARALLELIZATION_STRATEGIES["Qwen3_5ForCausalLM"], Qwen3_5ParallelizationStrategy) - - -class TestQwen35ParallelizationStrategyParallelize: - """Tests for Qwen3_5ParallelizationStrategy.parallelize() method.""" - - @staticmethod - def _stub_qwen3_5_modules(monkeypatch): - """Stub transformers.models.qwen3_5* so cp_linear_attn can be imported.""" - for path in ( - "transformers.models.qwen3_5_moe", - "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe", - "transformers.models.qwen3_5", - "transformers.models.qwen3_5.modeling_qwen3_5", - ): - stub = types.ModuleType(path) - stub.Qwen3_5MoeGatedDeltaNet = _FakeGatedDeltaNet - stub.Qwen3_5GatedDeltaNet = _FakeGatedDeltaNet - monkeypatch.setitem(sys.modules, path, stub) - - @pytest.fixture() - def mock_device_mesh(self): - """Create a mock device mesh with CP support.""" - from torch.distributed.device_mesh import DeviceMesh - - mesh = MagicMock(spec=DeviceMesh) - dp_shard_mesh = MagicMock() - dp_shard_mesh.size.return_value = 2 - dp_shard_mesh.ndim = 1 - tp_mesh = MagicMock() - tp_mesh.size.return_value = 1 - tp_mesh.ndim = 1 - cp_mesh = MagicMock() - cp_mesh.size.return_value = 1 - cp_mesh.ndim = 1 - - mesh.mesh_dim_names = ("dp_replicate", "dp_shard_cp", "tp") - mesh.__getitem__ = MagicMock( - side_effect=lambda key: { - "dp_replicate": MagicMock(size=MagicMock(return_value=1), ndim=1), - "dp_shard_cp": dp_shard_mesh, - "tp": tp_mesh, - "cp": cp_mesh, - ("dp_replicate", "dp_shard_cp"): dp_shard_mesh, - }[key] - ) - - return mesh, cp_mesh, tp_mesh - - @pytest.fixture() - def mock_env(self, monkeypatch): - """Mock the distributed functions used by DefaultParallelizationStrategy.""" - import nemo_automodel.components.distributed.parallelizer as par_mod - import nemo_automodel.components.distributed.parallelizer_utils as par_utils - - fully_shard_mock = MagicMock(side_effect=lambda model, **kw: model) - monkeypatch.setattr(par_mod, "fully_shard", fully_shard_mock, raising=False) - - apply_fsdp_mock = MagicMock() - monkeypatch.setattr(par_mod, "apply_fsdp2_sharding_recursively", apply_fsdp_mock, raising=False) - - # Also mock fully_shard_by_dtype which _fsdp_by_dtype calls - fsdp_by_dtype_mock = MagicMock() - monkeypatch.setattr(par_utils, "fully_shard_by_dtype", fsdp_by_dtype_mock, raising=False) - - # Mock _pre_shard_combined_projections which _fsdp_by_dtype calls - monkeypatch.setattr(par_mod, "_pre_shard_combined_projections", MagicMock(), raising=False) - - extract_mock = MagicMock(return_value=[]) - monkeypatch.setattr(par_mod, "_extract_model_layers", extract_mock, raising=False) - - get_plan_mock = MagicMock(return_value={}) - monkeypatch.setattr(par_mod, "_get_parallel_plan", get_plan_mock, raising=False) - - validate_mock = MagicMock() - monkeypatch.setattr(par_mod, "validate_tp_mesh", validate_mock, raising=False) - - parallelize_mod_mock = MagicMock() - monkeypatch.setattr(par_mod, "parallelize_module", parallelize_mod_mock, raising=False) - - checkpoint_mock = MagicMock(side_effect=lambda x: x) - monkeypatch.setattr(par_mod, "checkpoint_wrapper", checkpoint_mock, raising=False) - - return { - "apply_fsdp": apply_fsdp_mock, - "fully_shard": fully_shard_mock, - "fully_shard_by_dtype": fsdp_by_dtype_mock, - } - - def test_parallelize_calls_patch_and_delegates(self, fake_model, monkeypatch, mock_device_mesh, mock_env): - """parallelize() patches the model and delegates to super().""" - self._stub_qwen3_5_modules(monkeypatch) - - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - - from nemo_automodel.components.distributed.parallelizer import Qwen3_5ParallelizationStrategy - - mesh, cp_mesh, tp_mesh = mock_device_mesh - strategy = Qwen3_5ParallelizationStrategy() - - with patch("nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn.patch_hf_model") as mock_patch: - result = strategy.parallelize(model=fake_model, device_mesh=mesh) - - # patch_hf_model was called (cp_enabled=False because "cp" not in mesh_dim_names) - mock_patch.assert_called_once_with(fake_model, cp_enabled=False) - # super().parallelize ran fully_shard - mock_env["fully_shard"].assert_called() - assert result is fake_model - - def test_parallelize_swaps_and_restores_fsdp_global(self, fake_model, monkeypatch, mock_device_mesh, mock_env): - """The globals swap for apply_fsdp2_sharding_recursively is restored after call.""" - self._stub_qwen3_5_modules(monkeypatch) - - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - - import nemo_automodel.components.distributed.parallelizer as par_mod - from nemo_automodel.components.distributed.parallelizer import Qwen3_5ParallelizationStrategy - - original_fn = par_mod.apply_fsdp2_sharding_recursively - strategy = Qwen3_5ParallelizationStrategy() - - # Track what function was used during super().parallelize() - called_with = {} - - def spy_apply_fsdp(*args, **kwargs): - # During super().parallelize, the global should be the custom _fsdp_by_dtype - called_with["fn"] = par_mod.apply_fsdp2_sharding_recursively - - mock_env["apply_fsdp"].side_effect = spy_apply_fsdp - - with patch("nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn.patch_hf_model"): - strategy.parallelize(model=fake_model, device_mesh=mock_device_mesh[0]) - - # After call, global is restored - assert par_mod.apply_fsdp2_sharding_recursively is original_fn - - def test_parallelize_restores_global_on_error(self, fake_model, monkeypatch, mock_device_mesh, mock_env): - """Global is restored even if super().parallelize() raises.""" - self._stub_qwen3_5_modules(monkeypatch) - - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - - import nemo_automodel.components.distributed.parallelizer as par_mod - from nemo_automodel.components.distributed.parallelizer import Qwen3_5ParallelizationStrategy - - original_fn = par_mod.apply_fsdp2_sharding_recursively - strategy = Qwen3_5ParallelizationStrategy() - - mock_env["fully_shard"].side_effect = RuntimeError("boom") - - with patch("nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn.patch_hf_model"): - with pytest.raises(RuntimeError, match="boom"): - strategy.parallelize(model=fake_model, device_mesh=mock_device_mesh[0]) - - # Global still restored - assert par_mod.apply_fsdp2_sharding_recursively is original_fn - - def test_parallelize_sets_cp_mesh_when_enabled(self, fake_model, monkeypatch, mock_device_mesh, mock_env): - """When CP is enabled, _cp_mesh is set on CPAwareGatedDeltaNet modules.""" - self._stub_qwen3_5_modules(monkeypatch) - - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - - from nemo_automodel.components.distributed.parallelizer import Qwen3_5ParallelizationStrategy - from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import ( - CPAwareGatedDeltaNet, - patch_hf_model, - ) - - mesh, cp_mesh, tp_mesh = mock_device_mesh - # Enable CP by adding "cp" to mesh_dim_names and making cp_mesh.size() > 1 - mesh.mesh_dim_names = ("dp_replicate", "dp_shard_cp", "tp", "cp") - cp_mesh.size.return_value = 2 - - # Pre-patch the model so the module is CPAwareGatedDeltaNet - patch_hf_model(fake_model, cp_enabled=True) - la = fake_model.layers[0].linear_attn - assert type(la) is CPAwareGatedDeltaNet - assert la._cp_mesh is None - - strategy = Qwen3_5ParallelizationStrategy() - - with patch("nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn.patch_hf_model"): - strategy.parallelize(model=fake_model, device_mesh=mesh) - - # CP mesh should be set - assert la._cp_mesh is cp_mesh - - def test_fsdp_by_dtype_handles_module_list(self, monkeypatch, mock_device_mesh, mock_env): - """The custom _fsdp_by_dtype correctly iterates ModuleList children.""" - self._stub_qwen3_5_modules(monkeypatch) - - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - - import nemo_automodel.components.distributed.parallelizer as par_mod - from nemo_automodel.components.distributed.parallelizer import Qwen3_5ParallelizationStrategy - - # Build a model with layers in a ModuleList - model = nn.Module() - model.config = types.SimpleNamespace( - num_attention_heads=8, - num_key_value_heads=8, - hidden_size=64, - ) - model.__class__.__name__ = "Qwen3_5ForCausalLM" - inner = nn.Module() - layer = nn.Module() - layer.mlp = nn.Linear(4, 4) - inner.layers = nn.ModuleList([layer]) - model.model = inner - - mesh, cp_mesh, tp_mesh = mock_device_mesh - - # Capture what the custom _fsdp_by_dtype does - shard_by_dtype_calls = [] - with ( - patch( - "nemo_automodel.components.distributed.parallelizer_utils.fully_shard_by_dtype", - side_effect=lambda *a, **kw: shard_by_dtype_calls.append(a[0]), - ), - patch("nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn.patch_hf_model"), - patch("nemo_automodel.components.distributed.parallelizer._pre_shard_combined_projections"), - ): - # Make extract_layers return the real layers - mock_env["apply_fsdp"].side_effect = lambda module, mesh, mp, offload=None: ( - par_mod.apply_fsdp2_sharding_recursively(module, mesh, mp, offload) - ) - strategy = Qwen3_5ParallelizationStrategy() - strategy.parallelize(model=model, device_mesh=mesh) - - # fully_shard_by_dtype should have been called for the layer child - assert len(shard_by_dtype_calls) > 0 - assert layer in shard_by_dtype_calls - - -# --------------------------------------------------------------------------- -# Helpers for _forward_no_cp tests -# --------------------------------------------------------------------------- - -# Dimensions used throughout the _forward_no_cp tests. -_HIDDEN = 16 -_NUM_K_HEADS = 2 -_NUM_V_HEADS = 2 -_HEAD_K_DIM = 4 -_HEAD_V_DIM = 4 -_KEY_DIM = _NUM_K_HEADS * _HEAD_K_DIM # 8 -_VALUE_DIM = _NUM_V_HEADS * _HEAD_V_DIM # 8 -_CONV_DIM = _KEY_DIM * 2 + _VALUE_DIM # 24 -_CONV_KERNEL = 4 - - -class _IdentityNorm(nn.Module): - """Simple norm replacement that accepts (x, z) and returns x unchanged.""" - - def forward(self, x, z): - return x - - -def _build_forward_module(monkeypatch): - """Import CPAwareGatedDeltaNet with stubs and build a module ready for _forward_no_cp. - - Returns (module, CPAwareGatedDeltaNet_class). - """ - # Stub transformers modules - for path in ( - "transformers.models.qwen3_5_moe", - "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe", - "transformers.models.qwen3_5", - "transformers.models.qwen3_5.modeling_qwen3_5", - ): - stub = types.ModuleType(path) - stub.Qwen3_5MoeGatedDeltaNet = _FakeGatedDeltaNet - stub.Qwen3_5GatedDeltaNet = _FakeGatedDeltaNet - # apply_mask_to_padding_states is imported at runtime inside _forward_no_cp - stub.apply_mask_to_padding_states = lambda hidden_states, mask: hidden_states - monkeypatch.setitem(sys.modules, path, stub) - - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - - from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import CPAwareGatedDeltaNet - - # Build a _FakeGatedDeltaNet and swap its class - mod = _FakeGatedDeltaNet() - mod.__class__ = CPAwareGatedDeltaNet - mod._cp_mesh = None - - # -- Submodules expected by _forward_no_cp -- - mod.in_proj_qkv = nn.Linear(_HIDDEN, _CONV_DIM, bias=False) - mod.in_proj_z = nn.Linear(_HIDDEN, _VALUE_DIM, bias=False) - mod.in_proj_b = nn.Linear(_HIDDEN, _NUM_V_HEADS, bias=False) - mod.in_proj_a = nn.Linear(_HIDDEN, _NUM_V_HEADS, bias=False) - mod.conv1d = nn.Conv1d( - _CONV_DIM, - _CONV_DIM, - _CONV_KERNEL, - groups=_CONV_DIM, - padding=_CONV_KERNEL - 1, - ) - mod.out_proj = nn.Linear(_VALUE_DIM, _HIDDEN, bias=False) - - # Norm that accepts (x, z) -> returns x (simplified) - mod.norm = _IdentityNorm() - - # Override A_log and dt_bias to match num_v_heads (the _FakeGatedDeltaNet - # creates them with size 4, but in_proj_a output is num_v_heads=2) - mod.A_log = nn.Parameter(torch.ones(_NUM_V_HEADS, dtype=torch.float32)) - mod.dt_bias = nn.Parameter(torch.ones(_NUM_V_HEADS, dtype=torch.bfloat16)) - - # Scalar/dim attributes - mod.head_k_dim = _HEAD_K_DIM - mod.head_v_dim = _HEAD_V_DIM - mod.num_k_heads = _NUM_K_HEADS - mod.num_v_heads = _NUM_V_HEADS - mod.key_dim = _KEY_DIM - mod.value_dim = _VALUE_DIM - mod.conv_kernel_size = _CONV_KERNEL - mod.activation = "silu" - mod.layer_idx = 0 - - # chunk_gated_delta_rule — returns (output, state) with correct shape - def _fake_chunk_gdn(query, key, value, *, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel): - # output: same shape as value [B, S, H_v, D_v] - return value.clone(), None - - mod.chunk_gated_delta_rule = _fake_chunk_gdn - - # causal_conv1d_fn — None forces the fallback conv path - mod.causal_conv1d_fn = None - - return mod, CPAwareGatedDeltaNet - - -class TestForwardNoCp: - """Tests for CPAwareGatedDeltaNet._forward_no_cp covering lines 91-193.""" - - def test_basic_forward_pass(self, monkeypatch): - """_forward_no_cp produces output with correct shape (basic path, cache_params=None).""" - mod, _ = _build_forward_module(monkeypatch) - x = torch.randn(2, 8, _HIDDEN) - out = mod._forward_no_cp(x) - assert out.shape == (2, 8, _HIDDEN) - - def test_forward_no_cp_cache_params_none(self, monkeypatch): - """Training path: cache_params=None exercises the else-branch at line 133-146.""" - mod, _ = _build_forward_module(monkeypatch) - x = torch.randn(1, 4, _HIDDEN) - out = mod._forward_no_cp(x, cache_params=None, cache_position=None) - assert out.shape == (1, 4, _HIDDEN) - - def test_forward_no_cp_causal_conv1d_fn_none_fallback(self, monkeypatch): - """When causal_conv1d_fn is None, falls back to F.silu(conv1d(...)).""" - mod, _ = _build_forward_module(monkeypatch) - assert mod.causal_conv1d_fn is None # confirm fallback path - x = torch.randn(1, 6, _HIDDEN) - out = mod._forward_no_cp(x) - assert out.shape == (1, 6, _HIDDEN) - - def test_forward_no_cp_with_causal_conv1d_fn(self, monkeypatch): - """When causal_conv1d_fn is set, it is called instead of conv1d fallback.""" - mod, _ = _build_forward_module(monkeypatch) - - # Install a mock causal_conv1d_fn - def _mock_causal_conv1d_fn(*, x, weight, bias, activation, seq_idx): - return torch.nn.functional.silu(x) - - mod.causal_conv1d_fn = _mock_causal_conv1d_fn - x = torch.randn(1, 6, _HIDDEN) - out = mod._forward_no_cp(x) - assert out.shape == (1, 6, _HIDDEN) - - def test_forward_no_cp_attention_mask(self, monkeypatch): - """attention_mask is passed through apply_mask_to_padding_states.""" - mod, _ = _build_forward_module(monkeypatch) - x = torch.randn(1, 4, _HIDDEN) - mask = torch.ones(1, 4, dtype=torch.bool) - out = mod._forward_no_cp(x, attention_mask=mask) - assert out.shape == (1, 4, _HIDDEN) - - def test_forward_no_cp_gqa_repeat(self, monkeypatch): - """When num_v_heads > num_k_heads, q/k are repeat-interleaved.""" - mod, _ = _build_forward_module(monkeypatch) - # Make v_heads > k_heads to trigger the repeat_interleave branch - new_num_v_heads = 4 - mod.num_v_heads = new_num_v_heads - mod.num_k_heads = 2 - # Adjust value_dim to match new num_v_heads - new_value_dim = new_num_v_heads * _HEAD_V_DIM # 16 - mod.value_dim = new_value_dim - conv_dim = _KEY_DIM * 2 + new_value_dim # 32 - mod.in_proj_qkv = nn.Linear(_HIDDEN, conv_dim, bias=False) - mod.in_proj_z = nn.Linear(_HIDDEN, new_value_dim, bias=False) - mod.in_proj_b = nn.Linear(_HIDDEN, new_num_v_heads, bias=False) - mod.in_proj_a = nn.Linear(_HIDDEN, new_num_v_heads, bias=False) - mod.conv1d = nn.Conv1d(conv_dim, conv_dim, _CONV_KERNEL, groups=conv_dim, padding=_CONV_KERNEL - 1) - mod.out_proj = nn.Linear(new_value_dim, _HIDDEN, bias=False) - # Update A_log and dt_bias to match new num_v_heads - mod.A_log = nn.Parameter(torch.ones(new_num_v_heads, dtype=torch.float32)) - mod.dt_bias = nn.Parameter(torch.ones(new_num_v_heads, dtype=torch.bfloat16)) - - x = torch.randn(1, 4, _HIDDEN) - out = mod._forward_no_cp(x) - assert out.shape == (1, 4, _HIDDEN) - - def test_forward_no_cp_uses_compute_gate(self, monkeypatch): - """_forward_no_cp delegates gate computation to _compute_gate.""" - mod, _ = _build_forward_module(monkeypatch) - original_compute_gate = mod._compute_gate - called = [] - - def _tracking_compute_gate(a): - called.append(True) - return original_compute_gate(a) - - mod._compute_gate = _tracking_compute_gate - x = torch.randn(1, 4, _HIDDEN) - mod._forward_no_cp(x) - assert len(called) == 1, "_compute_gate should be called exactly once" - - def test_forward_no_cp_output_dtype_matches_input(self, monkeypatch): - """Output dtype follows the projection layers (float32 in this test).""" - mod, _ = _build_forward_module(monkeypatch) - x = torch.randn(1, 4, _HIDDEN, dtype=torch.float32) - out = mod._forward_no_cp(x) - assert out.dtype == torch.float32 - - -class TestForwardDispatch: - """Tests for forward() dispatching to _forward_no_cp (lines 207-213).""" - - def test_forward_delegates_when_cp_mesh_none(self, monkeypatch): - """forward() calls _forward_no_cp when _cp_mesh is None.""" - mod, _ = _build_forward_module(monkeypatch) - mod._cp_mesh = None - x = torch.randn(1, 4, _HIDDEN) - out = mod.forward(x) - assert out.shape == (1, 4, _HIDDEN) - - def test_forward_delegates_when_cp_mesh_size_1(self, monkeypatch): - """forward() calls _forward_no_cp when _cp_mesh.size() <= 1.""" - mod, _ = _build_forward_module(monkeypatch) - mock_mesh = MagicMock() - mock_mesh.size.return_value = 1 - mod._cp_mesh = mock_mesh - x = torch.randn(1, 4, _HIDDEN) - out = mod.forward(x) - assert out.shape == (1, 4, _HIDDEN) - - def test_forward_passes_cache_params_through(self, monkeypatch): - """forward() passes cache_params, cache_position, attention_mask, and packing kwargs to _forward_no_cp.""" - mod, _ = _build_forward_module(monkeypatch) - mod._cp_mesh = None - - called_with = {} - orig_fwd = mod._forward_no_cp - - def _spy( - hidden_states, - cache_params=None, - cache_position=None, - attention_mask=None, - cu_seqlens=None, - indices=None, - ): - called_with["cache_params"] = cache_params - called_with["cache_position"] = cache_position - called_with["attention_mask"] = attention_mask - called_with["cu_seqlens"] = cu_seqlens - called_with["indices"] = indices - return orig_fwd( - hidden_states, - cache_params=cache_params, - cache_position=cache_position, - attention_mask=attention_mask, - cu_seqlens=cu_seqlens, - indices=indices, - ) - - mod._forward_no_cp = _spy - - x = torch.randn(1, 4, _HIDDEN) - mask = torch.ones(1, 4, dtype=torch.bool) - mod.forward(x, attention_mask=mask) - - assert called_with["cache_params"] is None - assert called_with["cache_position"] is None - assert called_with["attention_mask"] is mask - assert called_with["cu_seqlens"] is None - assert called_with["indices"] is None - - def test_forward_ignores_extra_cp_kwargs(self, monkeypatch): - """forward() accepts position_ids, qkv_format, etc. but ignores them on no-CP path.""" - mod, _ = _build_forward_module(monkeypatch) - mod._cp_mesh = None - x = torch.randn(1, 4, _HIDDEN) - out = mod.forward( - x, - position_ids=torch.arange(4).unsqueeze(0), - qkv_format="bshd", - cu_seqlens=None, - seq_index=None, - ) - assert out.shape == (1, 4, _HIDDEN) - - -class TestMakeFp32GetattrFallback: - """Test _make_fp32_getattr fallback when attr is not in _fp32_params.""" - - @staticmethod - def _stub_and_import(monkeypatch): - for path in ( - "transformers.models.qwen3_5_moe", - "transformers.models.qwen3_5_moe.modeling_qwen3_5_moe", - "transformers.models.qwen3_5", - "transformers.models.qwen3_5.modeling_qwen3_5", - ): - stub = types.ModuleType(path) - stub.Qwen3_5MoeGatedDeltaNet = _FakeGatedDeltaNet - stub.Qwen3_5GatedDeltaNet = _FakeGatedDeltaNet - monkeypatch.setitem(sys.modules, path, stub) - cp_mod_key = "nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn" - if cp_mod_key in sys.modules: - monkeypatch.delitem(sys.modules, cp_mod_key) - from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import ( - _make_fp32_getattr, - patch_hf_model, - ) - - return _make_fp32_getattr, patch_hf_model - - def test_fallback_raises_attribute_error(self, fake_model, monkeypatch): - """Accessing a non-existent attr via patched __getattr__ raises AttributeError.""" - _make_fp32_getattr, patch_hf_model = self._stub_and_import(monkeypatch) - patch_hf_model(fake_model, cp_enabled=False) - la = fake_model.layers[0].linear_attn - # _fp32_params exists, but "nonexistent_xyz" is not in it - with pytest.raises(AttributeError): - _ = la.nonexistent_xyz - - def test_fallback_resolves_real_attrs(self, fake_model, monkeypatch): - """Patched __getattr__ still resolves normal module attributes.""" - _make_fp32_getattr, patch_hf_model = self._stub_and_import(monkeypatch) - patch_hf_model(fake_model, cp_enabled=False) - la = fake_model.layers[0].linear_attn - # conv1d is a real submodule — should resolve fine - assert la.conv1d is not None - - -# --------------------------------------------------------------------------- -# Additional coverage for packing-aware paths (PR #2147) -# --------------------------------------------------------------------------- - - -class TestIsIndexedPackedMaskExtra: - """Branches of is_indexed_packed_mask not covered by TestPackingHelpers.""" - - def test_4d_mask_returns_false(self): - """A 4D bool/float mask is never an indexed packing mask.""" - from nemo_automodel.components.models.common.packing import is_indexed_packed_mask - - mask_4d = torch.zeros(1, 1, 4, 4, dtype=torch.int64) - # Even with values > 1 set, dim() != 2 short-circuits to False. - mask_4d[..., 0, 0] = 2 - assert is_indexed_packed_mask(mask_4d) is False - - -class TestDecoderLayerFullAttentionBranch: - """Exercise the full_attention branch of Qwen3_5DecoderLayerWithPacking.forward.""" - - def test_full_attention_calls_self_attn(self): - from nemo_automodel.components.models.qwen3_5.decoder_layer import ( - Qwen3_5DecoderLayerWithPacking, - ) - - class _RecorderSelfAttn(nn.Module): - def __init__(self): - super().__init__() - self.called_with = None - - def forward(self, **kwargs): - self.called_with = kwargs - return kwargs["hidden_states"], None - - layer = Qwen3_5DecoderLayerWithPacking.__new__(Qwen3_5DecoderLayerWithPacking) - nn.Module.__init__(layer) - layer.layer_type = "full_attention" - layer.input_layernorm = nn.Identity() - layer.self_attn = _RecorderSelfAttn() - layer.post_attention_layernorm = nn.Identity() - layer.mlp = nn.Identity() - - hs = torch.zeros(1, 5, 4) - mask = torch.ones(1, 5, dtype=torch.long) - out = layer( - hs, - position_embeddings=(torch.empty(0), torch.empty(0)), - attention_mask=mask, - position_ids=torch.arange(5).unsqueeze(0), - extra_fa_kwarg="passthrough", - ) - # Output shape preserved through identity residuals. - assert out.shape == hs.shape - called = layer.self_attn.called_with - assert called["attention_mask"] is mask - # Extra kwargs are forwarded to self_attn so FA2 wiring stays intact. - assert called.get("extra_fa_kwarg") == "passthrough" - - -class TestPatchHfModelDecoderLayerSwap: - """Cover the Qwen3_5DecoderLayer class-swap branch of patch_hf_model.""" - - def test_decoder_layers_class_swapped(self, monkeypatch): - """patch_hf_model swaps Qwen3_5DecoderLayer -> Qwen3_5DecoderLayerWithPacking - for linear_attention layers and leaves full_attention layers unswapped.""" - from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5DecoderLayer - - from nemo_automodel.components.models.qwen3_5.decoder_layer import ( - Qwen3_5DecoderLayerWithPacking, - ) - - # Make two thin instances of the HF base class — one linear, one full. - linear_layer = Qwen3_5DecoderLayer.__new__(Qwen3_5DecoderLayer) - nn.Module.__init__(linear_layer) - linear_layer.layer_type = "linear_attention" - linear_layer.linear_attn = _FakeGatedDeltaNet() - - full_layer = Qwen3_5DecoderLayer.__new__(Qwen3_5DecoderLayer) - nn.Module.__init__(full_layer) - full_layer.layer_type = "full_attention" - # No GatedDeltaNet on full_attention layers. - - model = nn.Module() - model.layers = nn.ModuleList([linear_layer, full_layer]) - - from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import patch_hf_model - - patch_hf_model(model, cp_enabled=False) - - assert isinstance(linear_layer, Qwen3_5DecoderLayerWithPacking) - # full_attention layers must NOT be swapped — they don't need packing. - assert not isinstance(full_layer, Qwen3_5DecoderLayerWithPacking) - - -class TestForwardNoCpPacking: - """Cover the packed-sample branches of _forward_no_cp (PR #2147 fix).""" - - def test_forward_no_cp_with_cu_seqlens_no_unpad(self, monkeypatch): - """is_packed=True + needs_unpad=False: cu_seqlens flows into FLA; no unpad/repad.""" - mod, _ = _build_forward_module(monkeypatch) - seen = {} - - def _fake_chunk_gdn( - query, key, value, *, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel, cu_seqlens=None - ): - seen["cu_seqlens"] = cu_seqlens - seen["q_shape"] = tuple(query.shape) - return value.clone(), None - - mod.chunk_gated_delta_rule = _fake_chunk_gdn - - def _fake_conv(*, x, weight, bias, activation, seq_idx): - seen["seq_idx_shape"] = None if seq_idx is None else tuple(seq_idx.shape) - return torch.nn.functional.silu(x) - - mod.causal_conv1d_fn = _fake_conv - - # B=1 packed: indexed mask covers full T (no padding) → needs_unpad=False. - x = torch.randn(1, 6, _HIDDEN) - indexed_mask = torch.tensor([[1, 1, 2, 2, 2, 2]], dtype=torch.long) - cu = torch.tensor([0, 2, 6], dtype=torch.long) - indices = torch.arange(6) - out = mod._forward_no_cp( - x, - attention_mask=indexed_mask, - cu_seqlens=cu, - indices=indices, - ) - assert out.shape == (1, 6, _HIDDEN) - # FLA received the cu_seqlens we passed in. - assert torch.equal(seen["cu_seqlens"], cu) - # seq_idx for conv comes straight from the indexed mask (B,T shape). - assert seen["seq_idx_shape"] == (1, 6) - - def test_forward_no_cp_with_unpad(self, monkeypatch): - """needs_unpad=True path: B>1 with real padding → unpad to [1, total_valid, H], repad on exit.""" - mod, _ = _build_forward_module(monkeypatch) - seen = {} - - def _fake_chunk_gdn( - query, key, value, *, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel, cu_seqlens=None - ): - seen["q_shape"] = tuple(query.shape) - return value.clone(), None - - mod.chunk_gated_delta_rule = _fake_chunk_gdn - - def _fake_conv(*, x, weight, bias, activation, seq_idx): - seen["seq_idx_shape"] = tuple(seq_idx.shape) - return torch.nn.functional.silu(x) - - mod.causal_conv1d_fn = _fake_conv - - # B=2, T=4: row 0 has 1 token padded, row 1 fully filled. - x = torch.randn(2, 4, _HIDDEN) - indexed_mask = torch.tensor( - [[1, 1, 2, 0], [1, 2, 2, 2]], - dtype=torch.long, - ) - # 7 non-padding tokens at flattened positions [0, 1, 2, 4, 5, 6, 7] - indices = torch.tensor([0, 1, 2, 4, 5, 6, 7], dtype=torch.long) - cu = torch.tensor([0, 2, 3, 4, 7], dtype=torch.long) - - out = mod._forward_no_cp( - x, - attention_mask=indexed_mask, - cu_seqlens=cu, - indices=indices, - ) - # Repad reconstructs the [B, T, H] shape. - assert out.shape == (2, 4, _HIDDEN) - # FLA saw the unpadded layout [1, 7, ...]. - assert seen["q_shape"][:2] == (1, 7) - # Conv saw a matching unpadded seq_idx [1, 7]. - assert seen["seq_idx_shape"] == (1, 7) - # Padded positions are zeroed in the output. - assert torch.all(out[0, 3] == 0) - - def test_forward_no_cp_derives_cu_seqlens_from_mask_fallback(self, monkeypatch): - """Direct caller (no decoder-layer subclass) passes only the indexed mask; layer derives cu_seqlens.""" - mod, _ = _build_forward_module(monkeypatch) - seen = {} - - def _fake_chunk_gdn( - query, key, value, *, g, beta, initial_state, output_final_state, use_qk_l2norm_in_kernel, cu_seqlens=None - ): - seen["cu_seqlens"] = None if cu_seqlens is None else cu_seqlens.tolist() - return value.clone(), None - - mod.chunk_gated_delta_rule = _fake_chunk_gdn - - x = torch.randn(1, 5, _HIDDEN) - indexed_mask = torch.tensor([[1, 1, 2, 2, 2]], dtype=torch.long) - # cu_seqlens / indices not passed — the layer must derive them. - out = mod._forward_no_cp(x, attention_mask=indexed_mask) - assert out.shape == (1, 5, _HIDDEN) - assert seen["cu_seqlens"] == [0, 2, 5] diff --git a/tests/unit_tests/models/qwen3_5/test_qwen3_5_dense_backbone.py b/tests/unit_tests/models/qwen3_5/test_qwen3_5_dense_backbone.py new file mode 100644 index 0000000000..f864015853 --- /dev/null +++ b/tests/unit_tests/models/qwen3_5/test_qwen3_5_dense_backbone.py @@ -0,0 +1,255 @@ +# Copyright (c) 2026, 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. + +"""Unit tests for the native dense Qwen3.5 backbone and fp32 SSMGate machinery. + +Covers the custom-model building blocks that replaced the old runtime +``patch_hf_model``: the fp32 ``SSMGate`` holder + ``install_ssm_gate`` + +``_SSMGateParam`` descriptor, the ``Qwen3_5DenseTextBackbone`` forward, and the +fp32-safe rotary embedding. All tests are CPU-only. +""" + +from __future__ import annotations + +import pytest +import torch +import torch.nn as nn + +pytest.importorskip("transformers.models.qwen3_5") +pytest.importorskip("transformers.models.qwen3_5_moe") + +from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig + +from nemo_automodel.components.models.common import BackendConfig +from nemo_automodel.components.models.qwen3_5.model import ( + Fp32SafeQwen3_5TextRotaryEmbedding, + Qwen3_5DenseTextBackbone, +) +from nemo_automodel.components.models.qwen3_5_moe.cp_linear_attn import ( + SSMGate, + _resolve_ssm_dtype, + _SSMGateParam, + install_ssm_gate, +) + + +def _backend(): + return BackendConfig( + linear="torch", + attn="sdpa", + rms_norm="torch", + # BackendConfig.rope_fusion defaults to (HAVE_TE and cuda), so it is True on + # the GPU CI runner. Qwen3.5's gated-query RoPE is incompatible with TE's + # fused rope kernel (it expects a plain 4D layout), which is why every + # Qwen3.5 recipe sets rope_fusion=false. Pin it here so the test exercises + # the supported (non-fused) path deterministically on both CPU and GPU. + rope_fusion=False, + enable_deepep=False, + fake_balanced_gate=False, + enable_hf_state_dict_adapter=True, + ) + + +def _tiny_config(layer_types=("full_attention",), **kwargs): + return Qwen3_5TextConfig( + vocab_size=64, + hidden_size=16, + num_hidden_layers=len(layer_types), + num_attention_heads=2, + num_key_value_heads=2, + head_dim=8, + intermediate_size=32, + max_position_embeddings=16, + rms_norm_eps=1e-6, + pad_token_id=0, + layer_types=list(layer_types), + attn_implementation="eager", + torch_dtype="float32", + **kwargs, + ) + + +# --------------------------------------------------------------------------- +# SSMGate / install_ssm_gate / _SSMGateParam descriptor +# --------------------------------------------------------------------------- + + +class TestSSMGate: + def test_params_are_fp32_with_expected_shape(self): + gate = SSMGate(num_v_heads=4) + assert gate.A_log.dtype == torch.float32 + assert gate.dt_bias.dtype == torch.float32 + assert tuple(gate.A_log.shape) == (4,) + assert tuple(gate.dt_bias.shape) == (4,) + + def test_forward_computes_gate_in_fp32(self): + gate = SSMGate(num_v_heads=4) + nn.init.zeros_(gate.A_log) # exp(0) = 1 + nn.init.zeros_(gate.dt_bias) + a = torch.randn(2, 3, 4, dtype=torch.float32) + out = gate(a) + assert out.shape == a.shape + assert out.dtype == torch.float32 + # g = -exp(A_log) * softplus(a + dt_bias) = -softplus(a) here. + torch.testing.assert_close(out, -torch.nn.functional.softplus(a)) + + def test_forward_upcasts_bf16_input(self): + gate = SSMGate(num_v_heads=2) + out = gate(torch.randn(1, 2, 2, dtype=torch.bfloat16)) + # Gate math runs in fp32 regardless of the activation dtype. + assert out.dtype == torch.float32 + + +class TestInstallSSMGate: + def _module_with_bare_params(self, dtype=torch.bfloat16, num_v_heads=4): + mod = nn.Module() + mod.register_parameter("A_log", nn.Parameter(torch.randn(num_v_heads, dtype=dtype))) + mod.register_parameter("dt_bias", nn.Parameter(torch.randn(num_v_heads, dtype=dtype))) + return mod + + def test_moves_params_into_fp32_holder(self): + mod = self._module_with_bare_params(dtype=torch.bfloat16) + original_a_log = mod._parameters["A_log"] + + gate = install_ssm_gate(mod, fp32_dtype=torch.float32) + + # Bare params are removed from the parent and live in the holder. + assert "A_log" not in mod._parameters + assert "dt_bias" not in mod._parameters + assert mod._fp32_params is gate + assert gate.A_log.dtype == torch.float32 + assert gate.dt_bias.dtype == torch.float32 + # The original tensor values are preserved (cast to fp32 in place). + torch.testing.assert_close(gate.A_log, original_a_log.detach().float()) + + def test_already_fp32_params_kept(self): + mod = self._module_with_bare_params(dtype=torch.float32) + a_before = mod._parameters["A_log"].detach().clone() + gate = install_ssm_gate(mod, fp32_dtype=torch.float32) + torch.testing.assert_close(gate.A_log.detach(), a_before) + + +class TestSSMGateParamDescriptor: + def test_descriptor_resolves_to_holder(self): + class Mod(nn.Module): + A_log = _SSMGateParam("A_log") + dt_bias = _SSMGateParam("dt_bias") + + def __init__(self): + super().__init__() + self.register_parameter("A_log", nn.Parameter(torch.randn(3))) + self.register_parameter("dt_bias", nn.Parameter(torch.randn(3))) + install_ssm_gate(self, fp32_dtype=torch.float32) + + mod = Mod() + # Attribute reads resolve through the descriptor into the holder. + assert mod.A_log is mod._fp32_params.A_log + assert mod.dt_bias is mod._fp32_params.dt_bias + + def test_descriptor_on_class_returns_itself(self): + class Mod(nn.Module): + A_log = _SSMGateParam("A_log") + + assert isinstance(Mod.A_log, _SSMGateParam) + + +class TestResolveSSMDtype: + def test_default_is_fp32(self): + class Cfg: + pass + + assert _resolve_ssm_dtype(Cfg()) == torch.float32 + + def test_string_dtype_resolved(self): + class Cfg: + mamba_ssm_dtype = "float32" + + assert _resolve_ssm_dtype(Cfg()) == torch.float32 + + def test_explicit_dtype_passthrough(self): + class Cfg: + mamba_ssm_dtype = torch.float32 + + assert _resolve_ssm_dtype(Cfg()) == torch.float32 + + +# --------------------------------------------------------------------------- +# Fp32SafeQwen3_5TextRotaryEmbedding +# --------------------------------------------------------------------------- + + +class TestFp32SafeRotaryEmbedding: + def test_inv_freq_stays_fp32_after_dtype_cast(self): + cfg = _tiny_config() + rope = Fp32SafeQwen3_5TextRotaryEmbedding(config=cfg) + assert rope.inv_freq.dtype == torch.float32 + rope.to(torch.bfloat16) + # The fp32 inv_freq buffer must survive a bulk bf16 cast. + assert rope.inv_freq.dtype == torch.float32 + + +# --------------------------------------------------------------------------- +# Qwen3_5DenseTextBackbone +# --------------------------------------------------------------------------- + + +class TestDenseTextBackbone: + def test_builds_expected_layer_types(self): + cfg = _tiny_config(layer_types=("full_attention", "linear_attention")) + backbone = Qwen3_5DenseTextBackbone(cfg, _backend()) + types = [backbone.layers[str(i)].layer_type for i in range(cfg.num_hidden_layers)] + assert types == ["full_attention", "linear_attention"] + # The linear-attention block builds the native CP-aware GatedDeltaNet with + # an fp32 SSMGate holder; the full-attention block does not. + assert hasattr(backbone.layers["1"], "linear_attn") + assert hasattr(backbone.layers["1"].linear_attn, "_fp32_params") + + def test_forward_shape_full_attention(self): + cfg = _tiny_config(layer_types=("full_attention",)) + backbone = Qwen3_5DenseTextBackbone(cfg, _backend()) + out = backbone(input_ids=torch.tensor([[1, 2, 3, 4]], dtype=torch.long)) + assert out.last_hidden_state.shape == (1, 4, cfg.hidden_size) + assert out.past_key_values is None + + def test_forward_shape_mixed_layers(self): + # The linear_attention block runs FLA's gated-delta-rule / causal_conv1d + # kernels, which are CUDA-only (no CPU fallback when CUDA is present), so + # build the backbone and inputs on the GPU. Skip when no GPU is available. + if not torch.cuda.is_available(): + pytest.skip("linear_attention forward requires CUDA (FLA kernels)") + device = torch.device("cuda") + cfg = _tiny_config(layer_types=("full_attention", "linear_attention")) + backbone = Qwen3_5DenseTextBackbone(cfg, _backend()).to(device) + out = backbone(input_ids=torch.tensor([[1, 2, 3, 4]], dtype=torch.long, device=device)) + assert out.last_hidden_state.shape == (1, 4, cfg.hidden_size) + + def test_forward_accepts_inputs_embeds(self): + cfg = _tiny_config(layer_types=("full_attention",)) + backbone = Qwen3_5DenseTextBackbone(cfg, _backend()) + embeds = torch.randn(1, 4, cfg.hidden_size) + out = backbone(inputs_embeds=embeds) + assert out.last_hidden_state.shape == (1, 4, cfg.hidden_size) + + def test_kv_cache_not_supported(self): + cfg = _tiny_config(layer_types=("full_attention",)) + backbone = Qwen3_5DenseTextBackbone(cfg, _backend()) + with pytest.raises(NotImplementedError): + backbone(input_ids=torch.tensor([[1, 2, 3, 4]], dtype=torch.long), use_cache=True) + + def test_get_set_input_embeddings(self): + cfg = _tiny_config(layer_types=("full_attention",)) + backbone = Qwen3_5DenseTextBackbone(cfg, _backend()) + new_emb = nn.Embedding(cfg.vocab_size, cfg.hidden_size) + backbone.set_input_embeddings(new_emb) + assert backbone.get_input_embeddings() is new_emb diff --git a/tests/unit_tests/models/qwen3_5/test_qwen3_5_mtp.py b/tests/unit_tests/models/qwen3_5/test_qwen3_5_mtp.py index a7a91d7ab0..462eb52ee6 100644 --- a/tests/unit_tests/models/qwen3_5/test_qwen3_5_mtp.py +++ b/tests/unit_tests/models/qwen3_5/test_qwen3_5_mtp.py @@ -41,6 +41,9 @@ def _tiny_config(**kwargs): pad_token_id=0, layer_types=layer_types, attn_implementation="eager", + # Match HF's default fp32 init so native-vs-HF parity comparisons line up; + # the native backbone otherwise defaults its params to bf16. + torch_dtype="float32", **kwargs, ) return cfg @@ -75,6 +78,10 @@ def _backend(): attn="sdpa", rms_norm="torch", dispatcher="torch", + # rope_fusion defaults to (HAVE_TE and cuda) -> True on the GPU CI runner. + # Qwen3.5's gated-query RoPE is incompatible with TE's fused rope kernel, + # so every recipe sets rope_fusion=false. Pin it for deterministic tests. + rope_fusion=False, fake_balanced_gate=False, enable_hf_state_dict_adapter=True, ) @@ -97,11 +104,14 @@ def test_num_nextn_override_takes_precedence(self): class TestQwen3_5MTPModel: - def test_dense_adapter_keeps_unpatched_linear_attn_keys(self): + def test_dense_adapter_routes_fp32_linear_attn_keys(self): + # The native backbone builds the fp32 ``_fp32_params`` SSMGate holder at + # construction, so the adapter must rename ``_fp32_params.A_log`` keys back + # to the bare HF layout at the save/load boundary. cfg = _tiny_config(mtp_num_hidden_layers=1) model = Qwen3_5ForCausalLM(cfg, backend=_backend()) - assert not model.state_dict_adapter.route_linear_attn_fp32_params + assert model.state_dict_adapter.route_linear_attn_fp32_params def test_mtp_disabled_matches_hf_forward(self): cfg = _tiny_config(mtp_num_hidden_layers=0, use_cache=False) @@ -155,11 +165,12 @@ def test_eval_does_not_emit_mtp_hidden_states(self): class TestQwen3_5VLMMTPModel: - def test_vlm_adapter_keeps_unpatched_linear_attn_keys(self): + def test_vlm_adapter_routes_fp32_linear_attn_keys(self): + # Same native fp32 SSMGate routing as the text-only model (see above). cfg = _tiny_vlm_config(mtp_num_hidden_layers=1) model = Qwen3_5ForConditionalGeneration(cfg, backend=_backend()) - assert not model.state_dict_adapter.route_linear_attn_fp32_params + assert model.state_dict_adapter.route_linear_attn_fp32_params def test_vlm_forward_emits_mtp_hidden_states_in_training(self): cfg = _tiny_vlm_config(mtp_num_hidden_layers=1) diff --git a/tests/unit_tests/models/qwen3_5/test_qwen3_5_pp.py b/tests/unit_tests/models/qwen3_5/test_qwen3_5_pp.py deleted file mode 100644 index 5c1c79315c..0000000000 --- a/tests/unit_tests/models/qwen3_5/test_qwen3_5_pp.py +++ /dev/null @@ -1,251 +0,0 @@ -# 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. - -"""Unit tests for the pipeline-parallel-safe Qwen3.5 dense text backbone. - -These exercise ``Qwen3_5TextModelPP.forward`` — the override that lets the HF -text model survive a pipeline split (``self.layers`` rewritten from -``nn.ModuleList`` to ``nn.ModuleDict``; ``norm`` dropped to ``None`` on non-last -stages). The override is verified in isolation by stubbing HF's -``Qwen3_5TextModel.forward`` (the ``super()`` target) so no model weights or -distributed setup are required. -""" - -from unittest.mock import patch - -import torch -import torch.nn as nn -from transformers.models.qwen3_5.configuration_qwen3_5 import ( - Qwen3_5Config, - Qwen3_5TextConfig, - Qwen3_5VisionConfig, -) -from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5TextModel - -from nemo_automodel.components.models.common import BackendConfig -from nemo_automodel.components.models.qwen3_5.model import ( - Qwen3_5CausalLMOutputWithPast, - Qwen3_5ForConditionalGeneration, - Qwen3_5TextModelPP, -) - - -def _bare_pp_model() -> Qwen3_5TextModelPP: - """A ``Qwen3_5TextModelPP`` instance without HF's heavy ``__init__``.""" - model = Qwen3_5TextModelPP.__new__(Qwen3_5TextModelPP) - nn.Module.__init__(model) - return model - - -def test_split_stage_presents_sliceable_layers_and_identity_norm(): - """On a split stage (ModuleDict layers, norm=None) HF's forward must see a - slice-able ModuleList over the same layer objects and a callable norm.""" - model = _bare_pp_model() - l0, l1 = nn.Linear(2, 2), nn.Linear(2, 2) - # Splitter keys layers by their original (non-contiguous) index. - model.layers = nn.ModuleDict({"3": l0, "7": l1}) - model.norm = None - - seen = {} - - def probe(self, *args, **kwargs): - seen["layers_type"] = type(self.layers) - seen["layers_objs"] = list(self.layers) - seen["norm_type"] = type(self.norm) - seen["kwargs"] = kwargs - # HF slices ``self.layers[: num_hidden_layers]`` — must not raise. - _ = self.layers[:64] - return "HIDDEN" - - with patch.object(Qwen3_5TextModel, "forward", probe): - out = model.forward(position_ids=None, foo=1) - - assert out == "HIDDEN" - assert seen["layers_type"] is nn.ModuleList - assert seen["layers_objs"] == [l0, l1] # same objects, preserved order - assert seen["norm_type"] is nn.Identity - assert seen["kwargs"] == {"position_ids": None, "foo": 1} - - -def test_split_stage_restores_containers_after_forward(): - """After the wrapped forward, the original ModuleDict / None norm are back.""" - model = _bare_pp_model() - layers = nn.ModuleDict({"0": nn.Linear(2, 2)}) - model.layers = layers - model.norm = None - - with patch.object(Qwen3_5TextModel, "forward", lambda self, *a, **k: "X"): - model.forward() - - assert model.layers is layers - assert isinstance(model.layers, nn.ModuleDict) - assert model.norm is None - - -def test_non_split_model_is_pure_passthrough(): - """Full (non-PP) model: ModuleList layers + real norm are left untouched.""" - model = _bare_pp_model() - layers = nn.ModuleList([nn.Linear(2, 2), nn.Linear(2, 2)]) - norm = nn.LayerNorm(2) - model.layers = layers - model.norm = norm - - seen = {} - - def probe(self, *args, **kwargs): - seen["layers"] = self.layers - seen["norm"] = self.norm - return "HIDDEN" - - with patch.object(Qwen3_5TextModel, "forward", probe): - out = model.forward() - - assert out == "HIDDEN" - # Same objects passed through — no swap occurred. - assert seen["layers"] is layers - assert seen["norm"] is norm - assert model.layers is layers - assert model.norm is norm - - -def test_containers_restored_when_forward_raises(): - """The finally-block must restore layers/norm even if HF's forward raises.""" - model = _bare_pp_model() - layers = nn.ModuleDict({"0": nn.Linear(2, 2)}) - model.layers = layers - model.norm = None - - def boom(self, *args, **kwargs): - raise RuntimeError("forward blew up") - - with patch.object(Qwen3_5TextModel, "forward", boom): - try: - model.forward() - except RuntimeError: - pass - - assert model.layers is layers - assert model.norm is None - - -def test_subclass_relationship(): - """Behavioral invariants the __class__ swap in the VLM __init__ relies on.""" - assert issubclass(Qwen3_5TextModelPP, Qwen3_5TextModel) - # No extra __init__ — the swap reuses the HF-built instance as-is. - assert "__init__" not in Qwen3_5TextModelPP.__dict__ - assert "forward" in Qwen3_5TextModelPP.__dict__ - - -# --------------------------------------------------------------------------- -# Outer-forward PP-stage dispatch, exercised on a tiny real VLM (CPU). -# A real pipeline split (FSDP2 + PipelineStage) can't run as a unit test, so we -# simulate the per-stage module layout the splitter produces: stage 0 keeps -# ``embed_tokens`` but not ``lm_head``; the last stage keeps ``lm_head`` but not -# ``embed_tokens``; middle stages keep neither. -# --------------------------------------------------------------------------- - - -def _tiny_vlm_model(): - text_config = Qwen3_5TextConfig( - vocab_size=64, - hidden_size=16, - num_hidden_layers=2, - num_attention_heads=2, - num_key_value_heads=2, - head_dim=8, - intermediate_size=32, - max_position_embeddings=16, - rms_norm_eps=1e-6, - pad_token_id=0, - layer_types=["full_attention", "full_attention"], - attn_implementation="eager", - ) - vision_config = Qwen3_5VisionConfig( - depth=1, - hidden_size=16, - intermediate_size=32, - num_heads=2, - patch_size=2, - spatial_merge_size=1, - temporal_patch_size=1, - out_hidden_size=16, - ) - config = Qwen3_5Config( - architectures=["Qwen3_5ForConditionalGeneration"], - text_config=text_config.to_dict(), - vision_config=vision_config.to_dict(), - image_token_id=60, - video_token_id=61, - vision_start_token_id=62, - vision_end_token_id=63, - ) - backend = BackendConfig( - linear="torch", - attn="sdpa", - rms_norm="torch", - dispatcher="torch", - fake_balanced_gate=False, - enable_hf_state_dict_adapter=True, - ) - model = Qwen3_5ForConditionalGeneration(config, backend=backend, num_nextn_predict_layers=0) - model.eval() - return model - - -def test_init_swaps_text_backbone_class(): - """__init__ points the HF text backbone at the PP-safe subclass.""" - model = _tiny_vlm_model() - assert isinstance(model.model.language_model, Qwen3_5TextModelPP) - - -def test_full_model_forward_returns_output_dataclass(): - """Non-PP (embed + lm_head both present) falls through to the normal path.""" - model = _tiny_vlm_model() - ids = torch.tensor([[1, 2, 3, 4]], dtype=torch.long) - out = model(input_ids=ids, attention_mask=torch.ones_like(ids), use_cache=False) - assert isinstance(out, Qwen3_5CausalLMOutputWithPast) - assert out.logits.shape == (1, 4, model.config.text_config.vocab_size) - - -def test_first_stage_returns_raw_hidden_states(): - """Stage 0 (embed present, lm_head dropped) returns hidden states, not logits.""" - model = _tiny_vlm_model() - model.lm_head = None # splitter drops lm_head on non-last stages - ids = torch.tensor([[1, 2, 3, 4]], dtype=torch.long) - out = model(input_ids=ids, attention_mask=torch.ones_like(ids), use_cache=False) - assert torch.is_tensor(out) - assert out.shape == (1, 4, model.config.text_config.hidden_size) - - -def test_last_stage_consumes_hidden_states_and_applies_lm_head(): - """Last stage (embed dropped, lm_head present) takes hidden states, emits logits.""" - model = _tiny_vlm_model() - model.model.language_model.embed_tokens = None # dropped on non-first stages - dtype = next(model.parameters()).dtype - hidden = torch.randn(1, 4, model.config.text_config.hidden_size, dtype=dtype) - out = model(inputs_embeds=hidden, attention_mask=torch.ones(1, 4, dtype=torch.long), use_cache=False) - assert torch.is_tensor(out) - assert out.shape == (1, 4, model.config.text_config.vocab_size) - - -def test_middle_stage_passes_hidden_states_through(): - """Middle stage (neither embed nor lm_head) returns hidden states for the next stage.""" - model = _tiny_vlm_model() - model.model.language_model.embed_tokens = None - model.lm_head = None - dtype = next(model.parameters()).dtype - hidden = torch.randn(1, 4, model.config.text_config.hidden_size, dtype=dtype) - out = model(inputs_embeds=hidden, attention_mask=torch.ones(1, 4, dtype=torch.long), use_cache=False) - assert torch.is_tensor(out) - assert out.shape == (1, 4, model.config.text_config.hidden_size) 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 4acfd76229..270b762bbe 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 @@ -48,12 +48,18 @@ def test_passes_through_already_routed_keys(self): key = "model.language_model.layers.0.linear_attn._fp32_params.A_log" assert _route_to_fp32_holder(key) == key - def test_does_not_route_non_a_log_keys(self): - for key in ( - "model.language_model.layers.0.linear_attn.dt_bias", - "model.language_model.layers.0.linear_attn.norm.weight", - ): - assert _route_to_fp32_holder(key) == key + def test_routes_bare_dt_bias_to_holder(self): + # Both SSM-gating master weights (A_log and dt_bias) live in the fp32 + # ``_fp32_params`` holder, so dt_bias is routed the same as 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_does_not_route_other_linear_attn_keys(self): + # Non SSM-gating params (e.g. the GatedDeltaNet norm) stay in place. + key = "model.language_model.layers.0.linear_attn.norm.weight" + assert _route_to_fp32_holder(key) == key def test_does_not_route_a_log_outside_linear_attn(self): # Defensive: only linear_attn.A_log should be routed. 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 e936ce3810..bc7c2e5233 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 @@ -627,6 +627,17 @@ def test_shared_experts_key_mapping(self, adapter): assert result[0][0] == "model.language_model.layers.0.mlp.shared_expert.gate_proj.weight" assert torch.equal(result[0][1], tensor) + def test_strips_fp32_holder_segment_on_save(self, adapter): + # The fp32 SSM-gating holder is stripped back to the bare HF key on save. + tensor = torch.randn(8) + fqn = "model.language_model.layers.0.linear_attn._fp32_params.A_log" + + result = adapter.convert_single_tensor_to_hf(fqn, tensor) + + assert len(result) == 1 + assert result[0][0] == "model.language_model.layers.0.linear_attn.A_log" + assert torch.equal(result[0][1], tensor) + def test_non_expert_tensor_passthrough(self, adapter): tensor = torch.randn(64, 64) fqn = "model.language_model.layers.0.self_attn.q_proj.weight" @@ -793,3 +804,73 @@ def test_from_hf_no_ep_shard_unchanged(self, adapter, monkeypatch): # No ep_shard slicing — full transposed tensor assert local_gate.shape == (n_experts, hidden, inter) torch.testing.assert_close(local_gate, gate_up_hf.transpose(1, 2).to(adapter.dtype)) + + +class TestFp32ParamRouting: + """Routing/stripping of SSM-gating params into/out of the fp32 holder.""" + + def test_strip_fp32_params_removes_holder_segment(self): + from nemo_automodel.components.models.qwen3_5_moe.state_dict_adapter import _strip_fp32_params + + assert ( + _strip_fp32_params("model.language_model.layers.0.linear_attn._fp32_params.A_log") + == "model.language_model.layers.0.linear_attn.A_log" + ) + assert ( + _strip_fp32_params("model.language_model.layers.0.linear_attn._fp32_params.dt_bias") + == "model.language_model.layers.0.linear_attn.dt_bias" + ) + + def test_strip_fp32_params_passthrough(self): + from nemo_automodel.components.models.qwen3_5_moe.state_dict_adapter import _strip_fp32_params + + for key in ( + "model.language_model.layers.0.self_attn.q_proj.weight", + "model.language_model.layers.0.linear_attn.norm.weight", + ): + assert _strip_fp32_params(key) == key + + def test_route_fp32_params_routes_gating_keys(self): + from nemo_automodel.components.models.qwen3_5_moe.state_dict_adapter import _route_fp32_params + + assert ( + _route_fp32_params("model.language_model.layers.0.linear_attn.A_log") + == "model.language_model.layers.0.linear_attn._fp32_params.A_log" + ) + assert ( + _route_fp32_params("model.language_model.layers.0.linear_attn.dt_bias") + == "model.language_model.layers.0.linear_attn._fp32_params.dt_bias" + ) + + def test_route_fp32_params_passthrough(self): + from nemo_automodel.components.models.qwen3_5_moe.state_dict_adapter import _route_fp32_params + + # Already routed, non-gating param, and a non-linear_attn A_log all pass through. + for key in ( + "model.language_model.layers.0.linear_attn._fp32_params.A_log", + "model.language_model.layers.0.linear_attn.norm.weight", + "model.some.other.path.A_log", + ): + assert _route_fp32_params(key) == key + + def test_route_strip_round_trip(self): + from nemo_automodel.components.models.qwen3_5_moe.state_dict_adapter import ( + _route_fp32_params, + _strip_fp32_params, + ) + + bare = "model.language_model.layers.3.linear_attn.A_log" + assert _strip_fp32_params(_route_fp32_params(bare)) == bare + + def test_from_hf_routes_gating_keys_into_holder(self, adapter): + # On load, bare HF SSM-gating keys are routed into the fp32 holder. + hf_state = { + "model.language_model.layers.0.linear_attn.A_log": torch.randn(8), + "model.language_model.layers.0.linear_attn.dt_bias": torch.randn(8), + } + + out = adapter.from_hf(hf_state) + + 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