From 4f96a93ff858edf833e37ac247c663711bc0dd3f Mon Sep 17 00:00:00 2001 From: SolshineCode Date: Tue, 31 Mar 2026 23:03:14 -0700 Subject: [PATCH 1/2] Add Qwen 3.5 and Nemotron 3 (nemotron_h) model support Extends DAM to support two new model architectures: Qwen 3.5 (model_type: qwen3_5/qwen3_5_text -> mergedqwen3_5): - Hybrid attention architecture with full (softmax) and linear (Gated Delta Network) attention layers, dispatched per-layer via layer_types - Full attention uses output gating, QK-normalization, and partial RoPE - Linear attention uses causal conv1d and recurrent delta rule updates - RMSNorm uses (1+weight) formulation via custom Qwen3_5DAMRMSNorm - SwiGLU MLP with all projections as DAMLinearLayer Nemotron-H (model_type: nemotron_h -> mergednemotron_h): - Three-block-type hybrid: Mamba-2 SSM, standard GQA attention, and MoE - Each block has a single mixer (not attention+MLP like standard transformers) - Mamba-2 with pure PyTorch recurrent forward (no mamba_ssm dependency) - MoE with sigmoid routing, top-k expert selection, and shared expert - Non-gated MLP experts with ReLU squared activation Both implementations follow the existing DAM patterns: - DAMLinearLayer for all nn.Linear projections - DAMEmbeddingLayer for token embeddings (conditional on dam_embedding_layer) - DAMRMSNorm for normalization (conditional on dam_layernorms) - Proper tie_weights support for weight-tied embeddings - Full ForCausalLM with prepare_inputs_for_generation Updated files: - dam/merge.py: fix_config() handles new model types - dam/model_preparation.py: AutoConfig/AutoModel registration - dam/utils.py: find_norm_layers() detects Qwen3_5RMSNorm and NemotronHRMSNorm with graceful fallback for older transformers versions Co-Authored-By: Claude Opus 4.6 (1M context) --- dam/merge.py | 6 + dam/model_preparation.py | 8 + dam/modeling/nemotron/__init__.py | 0 dam/modeling/nemotron/config.py | 253 ++++++ dam/modeling/nemotron/modeling.py | 1274 +++++++++++++++++++++++++++++ dam/modeling/qwen3_5/__init__.py | 0 dam/modeling/qwen3_5/config.py | 199 +++++ dam/modeling/qwen3_5/modeling.py | 1198 +++++++++++++++++++++++++++ dam/utils.py | 22 +- 9 files changed, 2959 insertions(+), 1 deletion(-) create mode 100644 dam/modeling/nemotron/__init__.py create mode 100644 dam/modeling/nemotron/config.py create mode 100644 dam/modeling/nemotron/modeling.py create mode 100644 dam/modeling/qwen3_5/__init__.py create mode 100644 dam/modeling/qwen3_5/config.py create mode 100644 dam/modeling/qwen3_5/modeling.py diff --git a/dam/merge.py b/dam/merge.py index c0d1efc..fbd45a3 100644 --- a/dam/merge.py +++ b/dam/merge.py @@ -24,6 +24,12 @@ def fix_config(save_path, num_models, non_linearity, merge_embedding_layers, mer elif data['model_type'] == "llama": data['model_type'] = "mergedllama" data['architectures'][0] = 'MergedLlamaForCausalLM' + elif data['model_type'] == "qwen3_5_text" or data['model_type'] == "qwen3_5": + data['model_type'] = "mergedqwen3_5" + data['architectures'][0] = 'MergedQwen3_5ForCausalLM' + elif data['model_type'] == "nemotron_h": + data['model_type'] = "mergednemotron_h" + data['architectures'][0] = 'MergedNemotronHForCausalLM' data['num_merged_models'] = num_models data['non_linearity'] = non_linearity diff --git a/dam/model_preparation.py b/dam/model_preparation.py index 4ce59bc..9f0f9d1 100644 --- a/dam/model_preparation.py +++ b/dam/model_preparation.py @@ -5,14 +5,22 @@ from modeling.mistral.modeling import MergedMistralForCausalLM from modeling.llama3.config import MergedLlamaConfig from modeling.llama3.modeling import MergedLlamaForCausalLM +from modeling.qwen3_5.config import MergedQwen3_5Config +from modeling.qwen3_5.modeling import MergedQwen3_5ForCausalLM +from modeling.nemotron.config import MergedNemotronHConfig +from modeling.nemotron.modeling import MergedNemotronHForCausalLM from glom import glom, Assign from modeling.dam import DAMLinearLayer from utils import find_linear_layers, find_embedding_layers AutoConfig.register("mergedmistral", MergedMistralConfig) AutoConfig.register("mergedllama", MergedLlamaConfig) +AutoConfig.register("mergedqwen3_5", MergedQwen3_5Config) +AutoConfig.register("mergednemotron_h", MergedNemotronHConfig) AutoModelForCausalLM.register(MergedMistralConfig, MergedMistralForCausalLM) AutoModelForCausalLM.register(MergedLlamaConfig, MergedLlamaForCausalLM) +AutoModelForCausalLM.register(MergedQwen3_5Config, MergedQwen3_5ForCausalLM) +AutoModelForCausalLM.register(MergedNemotronHConfig, MergedNemotronHForCausalLM) def print_trainable_parameters(model): """ diff --git a/dam/modeling/nemotron/__init__.py b/dam/modeling/nemotron/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/dam/modeling/nemotron/config.py b/dam/modeling/nemotron/config.py new file mode 100644 index 0000000..c401c20 --- /dev/null +++ b/dam/modeling/nemotron/config.py @@ -0,0 +1,253 @@ +# Copyright 2024-2025 NVIDIA Corporation and The HuggingFace Inc. team. 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. +"""Merged Nemotron-H model configuration for DAM.""" + +from transformers.configuration_utils import PretrainedConfig + + +class MergedNemotronHConfig(PretrainedConfig): + r""" + Configuration class for the DAM-merged Nemotron-H model. + + This extends the NemotronHConfig with DAM-specific parameters for + differentiable adaptive merging. Nemotron-H uses a hybrid architecture + with three block types: Mamba-2 SSM, standard attention, and MoE. + + Args: + vocab_size (`int`, *optional*, defaults to 131072): + Vocabulary size of the Nemotron model. + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + num_hidden_layers (`int`, *optional*, defaults to 32): + Number of hidden layers. + num_attention_heads (`int`, *optional*, defaults to 32): + Number of attention heads for attention layers. + num_key_value_heads (`int`, *optional*, defaults to 8): + Number of key_value heads for GQA. + head_dim (`int`, *optional*, defaults to 128): + Attention head dimension. + max_position_embeddings (`int`, *optional*, defaults to 4096): + Maximum sequence length. + intermediate_size (`int`, *optional*, defaults to 21504): + Dimension of the MLP representations. + mlp_hidden_act (`str`, *optional*, defaults to `"relu2"`): + Activation function for MLP layers (ReLU squared). + mlp_bias (`bool`, *optional*, defaults to `False`): + Whether MLP layers use bias. + attention_bias (`bool`, *optional*, defaults to `False`): + Whether attention projection layers use bias. + attention_dropout (`float`, *optional*, defaults to 0.0): + Attention dropout rate. + sliding_window (`int`, *optional*): + Sliding window attention size, if applicable. + layers_block_type (`list[str]`, *optional*): + Per-layer block type specification. Each entry is one of + "mamba", "attention", or "moe". + layer_norm_epsilon (`float`, *optional*, defaults to 1e-5): + Epsilon for layer normalization. + use_cache (`bool`, *optional*, defaults to `True`): + Whether to return last key/values attentions. + rope_theta (`float`, *optional*, defaults to 10000.0): + Base period of the RoPE embeddings. + ssm_state_size (`int`, *optional*, defaults to 128): + SSM state dimension for Mamba layers. + mamba_num_heads (`int`, *optional*, defaults to 128): + Number of Mamba heads. + mamba_head_dim (`int`, *optional*, defaults to 64): + Mamba head dimension. + mamba_hidden_act (`str`, *optional*, defaults to `"silu"`): + Activation function for Mamba layers. + n_groups (`int`, *optional*, defaults to 8): + Number of groups in Mamba-2. + conv_kernel (`int`, *optional*, defaults to 4): + Kernel size for conv1d in Mamba layers. + expand (`int`, *optional*, defaults to 2): + Expansion factor for Mamba layers. + use_conv_bias (`bool`, *optional*, defaults to `True`): + Whether conv1d uses bias. + chunk_size (`int`, *optional*, defaults to 128): + Chunk size for Mamba-2 SSD computation. + mamba_proj_bias (`bool`, *optional*, defaults to `False`): + Whether Mamba projection layers use bias. + n_routed_experts (`int`, *optional*, defaults to 8): + Number of routed experts in MoE layers. + n_shared_experts (`int`, *optional*, defaults to 1): + Number of shared experts in MoE layers. + moe_intermediate_size (`int`, *optional*, defaults to 7688): + Intermediate size for MoE expert MLPs. + moe_shared_expert_intermediate_size (`int`, *optional*, defaults to 7688): + Intermediate size for shared expert MLPs. + num_experts_per_tok (`int`, *optional*, defaults to 2): + Number of experts activated per token. + num_merged_models (`int`, *optional*, defaults to 3): + The number of models being merged. + init_merger_values (`list[float]`, *optional*, defaults to []): + Initial values for the merger coefficients. + use_tanh (`bool`, *optional*, defaults to False): + Whether to apply tanh non-linearity to merger coefficients. + dam_embedding_layer (`bool`, *optional*, defaults to True): + Whether embedding layers use DAM merging. + dam_layernorms (`bool`, *optional*, defaults to True): + Whether layer normalization uses DAM merging. + uses_base_model (`bool`, *optional*, defaults to True): + Whether the base model is included in the merge. + is_embedding_coef_trainable (`bool`, *optional*, defaults to False): + Whether embedding merger coefficients are trainable. + is_norm_coef_trainable (`bool`, *optional*, defaults to False): + Whether norm merger coefficients are trainable. + """ + + model_type = "mergednemotron_h" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=131072, + hidden_size=4096, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=8, + head_dim=128, + max_position_embeddings=4096, + intermediate_size=21504, + mlp_hidden_act="relu2", + mlp_bias=False, + attention_bias=False, + attention_dropout=0.0, + sliding_window=None, + layers_block_type=None, + layer_norm_epsilon=1e-5, + initializer_range=0.02, + use_cache=True, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + tie_word_embeddings=False, + rope_theta=10000.0, + # Mamba-2 parameters + ssm_state_size=128, + mamba_num_heads=128, + mamba_head_dim=64, + mamba_hidden_act="silu", + n_groups=8, + conv_kernel=4, + expand=2, + time_step_min=0.001, + time_step_max=0.1, + time_step_floor=1e-4, + use_conv_bias=True, + chunk_size=128, + mamba_proj_bias=False, + use_mamba_kernels=True, + residual_in_fp32=False, + rescale_prenorm_residual=True, + # MoE parameters + n_routed_experts=8, + n_shared_experts=1, + moe_intermediate_size=7688, + moe_shared_expert_intermediate_size=7688, + moe_latent_size=None, + moe_shared_expert_overlap=True, + num_experts_per_tok=2, + routed_scaling_factor=1.0, + norm_topk_prob=True, + use_bias=False, + # DAM-specific parameters + num_merged_models=3, + init_merger_values=[], + use_tanh=False, + model_index=None, + dam_embedding_layer=True, + dam_layernorms=True, + uses_base_model=True, + is_embedding_coef_trainable=False, + is_norm_coef_trainable=False, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.max_position_embeddings = max_position_embeddings + self.intermediate_size = intermediate_size + self.mlp_hidden_act = mlp_hidden_act + self.mlp_bias = mlp_bias + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.sliding_window = sliding_window + self.layer_norm_epsilon = layer_norm_epsilon + self.initializer_range = initializer_range + self.use_cache = use_cache + self.rope_theta = rope_theta + + # Mamba-2 + self.ssm_state_size = ssm_state_size + self.mamba_num_heads = mamba_num_heads + self.mamba_head_dim = mamba_head_dim + self.mamba_hidden_act = mamba_hidden_act + self.n_groups = n_groups + self.conv_kernel = conv_kernel + self.expand = expand + self.time_step_min = time_step_min + self.time_step_max = time_step_max + self.time_step_floor = time_step_floor + self.use_conv_bias = use_conv_bias + self.chunk_size = chunk_size + self.mamba_proj_bias = mamba_proj_bias + self.use_mamba_kernels = use_mamba_kernels + self.residual_in_fp32 = residual_in_fp32 + self.rescale_prenorm_residual = rescale_prenorm_residual + + # MoE + self.n_routed_experts = n_routed_experts + self.n_shared_experts = n_shared_experts + self.moe_intermediate_size = moe_intermediate_size + self.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size + self.moe_latent_size = moe_latent_size + self.moe_shared_expert_overlap = moe_shared_expert_overlap + self.num_experts_per_tok = num_experts_per_tok + self.routed_scaling_factor = routed_scaling_factor + self.norm_topk_prob = norm_topk_prob + self.use_bias = use_bias + + # Layer types: default to alternating mamba/moe/attention/moe pattern + if layers_block_type is None: + default_pattern = ["mamba", "moe", "attention", "moe"] + self.layers_block_type = [ + default_pattern[i % len(default_pattern)] + for i in range(num_hidden_layers) + ] + else: + self.layers_block_type = layers_block_type + + # DAM-specific + self.num_merged_models = num_merged_models + self.init_merger_values = init_merger_values + self.use_tanh = use_tanh + self.model_index = model_index + self.dam_embedding_layer = dam_embedding_layer + self.dam_layernorms = dam_layernorms + self.uses_base_model = uses_base_model + self.is_embedding_coef_trainable = is_embedding_coef_trainable + self.is_norm_coef_trainable = is_norm_coef_trainable + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/dam/modeling/nemotron/modeling.py b/dam/modeling/nemotron/modeling.py new file mode 100644 index 0000000..9288a38 --- /dev/null +++ b/dam/modeling/nemotron/modeling.py @@ -0,0 +1,1274 @@ +"""PyTorch Merged Nemotron-H model for DAM (Differentiable Adaptive Merging).""" + +import math +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from torch import nn +from torch.nn import CrossEntropyLoss + +from transformers.cache_utils import Cache, DynamicCache +from transformers.modeling_attn_mask_utils import AttentionMaskConverter +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging + +from .config import MergedNemotronHConfig +from ..dam import DAMLinearLayer, DAMEmbeddingLayer, DAMRMSNorm + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "MergedNemotronHConfig" + + +# --------------------------------------------------------------------------- +# Local RMSNorm (used when config.dam_layernorms is False) +# --------------------------------------------------------------------------- +class NemotronHRMSNorm(nn.Module): + """NemotronHRMSNorm is equivalent to T5LayerNorm.""" + + def __init__(self, hidden_size, eps=1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +# --------------------------------------------------------------------------- +# Gated RMSNorm (inside Mamba-2 blocks, NOT DAM-merged) +# --------------------------------------------------------------------------- +class NemotronHGatedRMSNorm(nn.Module): + """Gated RMSNorm used inside Mamba-2 mixer blocks. + + Applies RMSNorm to hidden_states and then gates with silu(gate). + This is internal to the Mamba computation and is NOT DAM-merged. + """ + + def __init__(self, hidden_size, eps=1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states, gate): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + hidden_states = self.weight * hidden_states.to(input_dtype) + # Apply gating: normed_hidden_states * silu(gate) + hidden_states = hidden_states * F.silu(gate) + return hidden_states + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +# --------------------------------------------------------------------------- +# Activation: ReLU squared +# --------------------------------------------------------------------------- +def relu_squared(x): + """ReLU squared activation: relu(x)^2.""" + return F.relu(x).pow(2) + + +# --------------------------------------------------------------------------- +# Rotary Position Embedding +# --------------------------------------------------------------------------- +class MergedNemotronHRotaryEmbedding(nn.Module): + """Rotary position embedding for Nemotron-H attention layers.""" + + def __init__(self, config): + super().__init__() + self.dim = config.head_dim + self.max_position_embeddings = config.max_position_embeddings + self.base = config.rope_theta + inv_freq = 1.0 / ( + self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + @torch.no_grad() + def forward(self, x, position_ids): + """Compute cos and sin for rotary embeddings. + + Args: + x: Input tensor, used only for dtype and device. + position_ids: Position indices of shape ``(batch_size, seq_len)``. + + Returns: + Tuple of ``(cos, sin)`` each of shape ``(batch_size, seq_len, head_dim)``. + """ + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q: The query tensor. + k: The key tensor. + cos: The cosine part of the rotary embedding. + sin: The sine part of the rotary embedding. + position_ids: Deprecated and unused. + unsqueeze_dim: The dimension along which to unsqueeze cos and sin so they + can be properly broadcast to the dimensions of q and k. + + Returns: + Tuple of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +# --------------------------------------------------------------------------- +# repeat_kv helper +# --------------------------------------------------------------------------- +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Expand key/value heads for grouped-query attention. + + The hidden states go from ``(batch, num_key_value_heads, seqlen, head_dim)`` + to ``(batch, num_attention_heads, seqlen, head_dim)``. + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +# --------------------------------------------------------------------------- +# Expert MLP (non-gated, relu-squared) used within MoE +# --------------------------------------------------------------------------- +class NemotronHExpertMLP(nn.Module): + """Non-gated MLP used for MoE experts and shared experts. + + Architecture: ``up_proj -> relu_squared -> down_proj``. + All linear layers use DAMLinearLayer. + """ + + def __init__(self, config: MergedNemotronHConfig, is_shared: bool = False): + super().__init__() + intermediate = ( + config.moe_shared_expert_intermediate_size if is_shared else config.moe_intermediate_size + ) + self.up_proj = DAMLinearLayer( + config.hidden_size, intermediate, bias=config.mlp_bias, num_models=config.num_merged_models + ) + self.down_proj = DAMLinearLayer( + intermediate, config.hidden_size, bias=config.mlp_bias, num_models=config.num_merged_models + ) + + def forward(self, x): + return self.down_proj(relu_squared(self.up_proj(x))) + + +# --------------------------------------------------------------------------- +# Attention (eager only, GQA) +# --------------------------------------------------------------------------- +class MergedNemotronHAttention(nn.Module): + """Multi-headed grouped-query attention for Nemotron-H. + + Uses standard eager (manual) attention. Flash attention and SDPA + variants are intentionally omitted for the DAM implementation. + """ + + def __init__(self, config: MergedNemotronHConfig, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " + "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.max_position_embeddings + self.is_causal = True + + self.q_proj = DAMLinearLayer( + self.hidden_size, + self.num_heads * self.head_dim, + bias=config.attention_bias, + num_models=config.num_merged_models, + ) + self.k_proj = DAMLinearLayer( + self.hidden_size, + self.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + num_models=config.num_merged_models, + ) + self.v_proj = DAMLinearLayer( + self.hidden_size, + self.num_key_value_heads * self.head_dim, + bias=config.attention_bias, + num_models=config.num_merged_models, + ) + self.o_proj = DAMLinearLayer( + self.num_heads * self.head_dim, + self.hidden_size, + bias=config.attention_bias, + num_models=config.num_merged_models, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + """Forward pass for attention. + + Args: + hidden_states: Input tensor of shape ``(batch, seq_len, hidden_size)``. + attention_mask: 4-D causal mask. + position_ids: Position indices. + past_key_value: Cached key/value states. + output_attentions: Whether to return attention weights. + use_cache: Whether to return updated cache. + cache_position: Cache position indices. + position_embeddings: Pre-computed ``(cos, sin)`` from the rotary embedding. + + Returns: + Tuple of ``(attn_output, attn_weights, past_key_value)``. + """ + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + if position_embeddings is not None: + cos, sin = position_embeddings + else: + # Fallback: should not happen when called from the model, but safe default + raise ValueError("position_embeddings must be provided to MergedNemotronHAttention") + + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + # Upcast to fp32 for numerical stability + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = F.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, -1) + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +# --------------------------------------------------------------------------- +# Mamba-2 Mixer (pure PyTorch recurrent implementation) +# --------------------------------------------------------------------------- +class MergedNemotronHMamba2Mixer(nn.Module): + """Mamba-2 (SSD) mixer layer for Nemotron-H. + + Implements a pure-PyTorch token-by-token recurrent forward pass so that + no external ``mamba_ssm`` library is required. + + Key components: + - ``in_proj``: DAMLinearLayer projecting to ``[gate, hidden_B_C, dt]`` + - ``conv1d``: standard depthwise Conv1d (NOT DAM-merged) + - ``A_log``, ``dt_bias``, ``D``: standard ``nn.Parameter`` (NOT DAM-merged) + - ``norm``: NemotronHGatedRMSNorm (NOT DAM-merged) + - ``out_proj``: DAMLinearLayer + """ + + def __init__(self, config: MergedNemotronHConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + + self.hidden_size = config.hidden_size + self.mamba_num_heads = config.mamba_num_heads + self.mamba_head_dim = config.mamba_head_dim + self.ssm_state_size = config.ssm_state_size + self.n_groups = config.n_groups + self.time_step_min = config.time_step_min + + # Derived sizes + self.intermediate_size = self.mamba_num_heads * self.mamba_head_dim + self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size + + # Projection size: gate + conv_input + dt + self.projection_size = self.intermediate_size + self.conv_dim + self.mamba_num_heads + + # Input projection (DAM-merged) + self.in_proj = DAMLinearLayer( + self.hidden_size, + self.projection_size, + bias=config.mamba_proj_bias, + num_models=config.num_merged_models, + ) + + # Depthwise conv1d (standard, NOT DAM-merged) + self.conv1d = nn.Conv1d( + in_channels=self.conv_dim, + out_channels=self.conv_dim, + kernel_size=config.conv_kernel, + groups=self.conv_dim, + padding=config.conv_kernel - 1, + bias=config.use_conv_bias, + ) + + # SSM parameters (standard nn.Parameter, NOT DAM-merged) + self.A_log = nn.Parameter(torch.zeros(self.mamba_num_heads)) + self.dt_bias = nn.Parameter(torch.zeros(self.mamba_num_heads)) + self.D = nn.Parameter(torch.zeros(self.mamba_num_heads)) + + # Gated RMSNorm (NOT DAM-merged, internal to Mamba computation) + self.norm = NemotronHGatedRMSNorm(self.intermediate_size, eps=config.layer_norm_epsilon) + + # Output projection (DAM-merged) + self.out_proj = DAMLinearLayer( + self.intermediate_size, + self.hidden_size, + bias=config.mamba_proj_bias, + num_models=config.num_merged_models, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + **kwargs, + ) -> torch.Tensor: + """Forward pass for Mamba-2 mixer. + + Uses a pure PyTorch recurrent (token-by-token) computation. + + Args: + hidden_states: Input of shape ``(batch, seq_len, hidden_size)``. + attention_mask: Padding mask of shape ``(batch, seq_len)`` where 1 + indicates a valid token and 0 indicates padding. + + Returns: + Output tensor of shape ``(batch, seq_len, hidden_size)``. + """ + batch_size, seq_len, _ = hidden_states.shape + + # Project input + zxbcdt = self.in_proj(hidden_states) # (B, L, projection_size) + + # Split into gate (z), conv input (xBC), and dt + z, xBC, dt = torch.split( + zxbcdt, + [self.intermediate_size, self.conv_dim, self.mamba_num_heads], + dim=-1, + ) + + # Apply depthwise conv1d (causal) + xBC = xBC.transpose(1, 2) # (B, conv_dim, L) + xBC = self.conv1d(xBC)[:, :, :seq_len] # causal: trim to seq_len + xBC = F.silu(xBC) + xBC = xBC.transpose(1, 2) # (B, L, conv_dim) + + # Split xBC into x, B, C + x, B, C = torch.split( + xBC, + [ + self.intermediate_size, + self.n_groups * self.ssm_state_size, + self.n_groups * self.ssm_state_size, + ], + dim=-1, + ) + + # Reshape for SSM computation + x = x.view(batch_size, seq_len, self.mamba_num_heads, self.mamba_head_dim) + B = B.view(batch_size, seq_len, self.n_groups, self.ssm_state_size) + C = C.view(batch_size, seq_len, self.n_groups, self.ssm_state_size) + + # Compute A from A_log + A = -torch.exp(self.A_log.float()) # (num_heads,) + + # Compute dt with bias, softplus, and clamp + dt = F.softplus(dt + self.dt_bias).clamp(min=self.time_step_min) # (B, L, num_heads) + + # Expand B and C from groups to heads + heads_per_group = self.mamba_num_heads // self.n_groups + B = B.repeat_interleave(heads_per_group, dim=2) # (B, L, num_heads, ssm_state_size) + C = C.repeat_interleave(heads_per_group, dim=2) # (B, L, num_heads, ssm_state_size) + + # Recurrent SSM computation + state = torch.zeros( + batch_size, + self.mamba_num_heads, + self.mamba_head_dim, + self.ssm_state_size, + dtype=torch.float32, + device=hidden_states.device, + ) + outputs = [] + + for t in range(seq_len): + x_t = x[:, t, :, :] # (B, num_heads, head_dim) + B_t = B[:, t, :, :] # (B, num_heads, ssm_state_size) + C_t = C[:, t, :, :] # (B, num_heads, ssm_state_size) + dt_t = dt[:, t, :] # (B, num_heads) + + # Discretize + # dA: (B, num_heads, 1, 1) + dA = torch.exp(dt_t.unsqueeze(-1).unsqueeze(-1) * A.unsqueeze(-1).unsqueeze(0)) + # dB: (B, num_heads, head_dim, ssm_state_size) + dB = dt_t.unsqueeze(-1).unsqueeze(-1) * B_t.unsqueeze(2) * x_t.unsqueeze(-1) + + state = state * dA + dB + + # Output: y = (state @ C^T) + D * x + y_t = ( + torch.einsum("bhds,bhs->bhd", state.to(x_t.dtype), C_t) + + self.D.unsqueeze(0).unsqueeze(-1) * x_t + ) # (B, num_heads, head_dim) + outputs.append(y_t) + + y = torch.stack(outputs, dim=1) # (B, L, num_heads, head_dim) + y = y.view(batch_size, seq_len, -1) # (B, L, intermediate_size) + + # Apply gated RMS norm (gate = z) + y = self.norm(y, z) + + # Output projection + y = self.out_proj(y) + + # Mask padding tokens + if attention_mask is not None: + y = y * attention_mask.unsqueeze(-1) + + return y + + +# --------------------------------------------------------------------------- +# Mixture of Experts +# --------------------------------------------------------------------------- +class MergedNemotronHMoE(nn.Module): + """Mixture of Experts layer for Nemotron-H. + + Uses sigmoid routing with top-k selection and a shared expert that + is always active. The routing gate is a DAMLinearLayer. + """ + + def __init__(self, config: MergedNemotronHConfig): + super().__init__() + self.gate = DAMLinearLayer( + config.hidden_size, + config.n_routed_experts, + bias=False, + num_models=config.num_merged_models, + ) + self.experts = nn.ModuleList( + [NemotronHExpertMLP(config, is_shared=False) for _ in range(config.n_routed_experts)] + ) + self.shared_experts = NemotronHExpertMLP(config, is_shared=True) + self.num_experts_per_tok = config.num_experts_per_tok + self.routed_scaling_factor = config.routed_scaling_factor + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Forward pass for MoE. + + Args: + hidden_states: Input tensor of shape ``(batch, seq_len, hidden_dim)``. + + Returns: + Output tensor of shape ``(batch, seq_len, hidden_dim)``. + """ + batch_size, seq_len, hidden_dim = hidden_states.shape + hidden_states_flat = hidden_states.view(-1, hidden_dim) + + # Routing with sigmoid scores + router_logits = self.gate(hidden_states_flat) + routing_weights = torch.sigmoid(router_logits) + + # Top-k selection + topk_weights, topk_indices = torch.topk( + routing_weights, self.num_experts_per_tok, dim=-1 + ) + # Normalize selected weights + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + topk_weights = topk_weights * self.routed_scaling_factor + + # Compute expert outputs + final_output = torch.zeros_like(hidden_states_flat) + for expert_idx in range(len(self.experts)): + expert_mask = (topk_indices == expert_idx).any(dim=-1) + if expert_mask.any(): + expert_input = hidden_states_flat[expert_mask] + expert_output = self.experts[expert_idx](expert_input) + # Gather the weights assigned to this expert for each selected token + weight_mask = topk_indices == expert_idx + expert_weights = (topk_weights * weight_mask.float()).sum(dim=-1) + final_output[expert_mask] += expert_output * expert_weights[expert_mask].unsqueeze(-1) + + # Always add the shared expert output + shared_output = self.shared_experts(hidden_states_flat) + final_output = final_output + shared_output + + return final_output.view(batch_size, seq_len, hidden_dim) + + +# --------------------------------------------------------------------------- +# Decoder Block (dispatches to mamba / attention / moe) +# --------------------------------------------------------------------------- +class MergedNemotronHBlock(nn.Module): + """A single Nemotron-H decoder block. + + Each block contains a single sub-module (mixer) which is one of: + - ``MergedNemotronHMamba2Mixer`` for ``"mamba"`` layers + - ``MergedNemotronHAttention`` for ``"attention"`` layers + - ``MergedNemotronHMoE`` for ``"moe"`` layers + + The block applies pre-norm with a residual connection: + ``output = hidden_states + mixer(norm(hidden_states))``. + """ + + def __init__(self, config: MergedNemotronHConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.block_type = config.layers_block_type[layer_idx] + + # Pre-norm (DAM-merged or standard) + if config.dam_layernorms: + self.norm = DAMRMSNorm( + config.hidden_size, eps=config.layer_norm_epsilon, num_models=config.num_merged_models + ) + else: + self.norm = NemotronHRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + + # Mixer sub-module + if self.block_type == "mamba": + self.mixer = MergedNemotronHMamba2Mixer(config, layer_idx) + elif self.block_type == "attention": + self.mixer = MergedNemotronHAttention(config, layer_idx) + elif self.block_type == "moe": + self.mixer = MergedNemotronHMoE(config) + else: + raise ValueError( + f"Unknown block type '{self.block_type}' at layer {layer_idx}. " + "Expected one of 'mamba', 'attention', or 'moe'." + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> Tuple[torch.FloatTensor, ...]: + """Forward pass for a single decoder block. + + Args: + hidden_states: Input of shape ``(batch, seq_len, hidden_size)``. + attention_mask: Causal mask for attention blocks, padding mask for + mamba blocks, or ``None`` for MoE blocks. + position_ids: Position indices for attention layers. + past_key_value: Cached key/value states (attention only). + output_attentions: Whether to return attention weights. + use_cache: Whether to return updated cache. + cache_position: Cache position indices. + position_embeddings: Pre-computed ``(cos, sin)`` for RoPE. + + Returns: + Tuple starting with the output hidden states. May also contain + attention weights and cache depending on the block type and flags. + """ + residual = hidden_states + hidden_states = self.norm(hidden_states) + + if self.block_type == "mamba": + hidden_states = self.mixer(hidden_states, attention_mask=attention_mask) + hidden_states = residual + hidden_states + outputs = (hidden_states,) + + elif self.block_type == "attention": + hidden_states, self_attn_weights, present_key_value = self.mixer( + hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + hidden_states = residual + hidden_states + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + if use_cache: + outputs += (present_key_value,) + + else: # moe + hidden_states = self.mixer(hidden_states) + hidden_states = residual + hidden_states + outputs = (hidden_states,) + + return outputs + + +# --------------------------------------------------------------------------- +# PreTrainedModel base +# --------------------------------------------------------------------------- +NEMOTRONH_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`MergedNemotronHConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +class MergedNemotronHPreTrainedModel(PreTrainedModel): + """Base class for all Merged Nemotron-H DAM models.""" + + config_class = MergedNemotronHConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["MergedNemotronHBlock"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = False + _supports_sdpa = False + _supports_cache_class = True + _supports_static_cache = False + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +# --------------------------------------------------------------------------- +# Input docstring +# --------------------------------------------------------------------------- +NEMOTRONH_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks) that can be used to speed up + sequential decoding. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance; + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +# --------------------------------------------------------------------------- +# Model (decoder stack) +# --------------------------------------------------------------------------- +class MergedNemotronHModel(MergedNemotronHPreTrainedModel): + """Nemotron-H hybrid transformer-mamba-moe decoder stack for DAM. + + Consists of ``config.num_hidden_layers`` blocks. Each block may be a Mamba-2 + SSM layer, a standard GQA attention layer, or a Mixture-of-Experts layer, + as specified by ``config.layers_block_type``. + + Args: + config: MergedNemotronHConfig + """ + + def __init__(self, config: MergedNemotronHConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + # Embedding layer + if config.dam_embedding_layer: + self.embed_tokens = DAMEmbeddingLayer( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + num_models=config.num_merged_models, + padding_idx=self.padding_idx, + ) + else: + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + + # Decoder layers + self.layers = nn.ModuleList( + [MergedNemotronHBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + + # Final norm + if config.dam_layernorms: + self.norm = DAMRMSNorm( + config.hidden_size, eps=config.layer_norm_epsilon, num_models=config.num_merged_models + ) + else: + self.norm = NemotronHRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + + # Rotary embeddings (shared, computed once) + self.rotary_emb = MergedNemotronHRotaryEmbedding(config) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + """Forward pass through the full Nemotron-H decoder stack. + + Args: + input_ids: Token indices of shape ``(batch_size, sequence_length)``. + attention_mask: Padding mask of shape ``(batch_size, sequence_length)``. + position_ids: Position indices. + past_key_values: Cached key/value states for attention layers. + inputs_embeds: Pre-computed embeddings, alternative to ``input_ids``. + use_cache: Whether to return updated caches. + output_attentions: Whether to return attention weights. + output_hidden_states: Whether to return all hidden states. + return_dict: Whether to return a ``BaseModelOutputWithPast``. + cache_position: Cache position indices. + + Returns: + ``BaseModelOutputWithPast`` or tuple of tensors. + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + if self.gradient_checkpointing and self.training and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + return_legacy_cache = False + if use_cache and not isinstance(past_key_values, Cache) and not self.training: + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + return_legacy_cache = True + logger.warning_once( + "We detected that you are passing `past_key_values` as a tuple and this is deprecated and " + "will be removed in v4.43. Please use an appropriate `Cache` class " + "(https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)" + ) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + # Build the 4-D causal mask for attention layers + causal_mask = self._update_causal_mask( + attention_mask, inputs_embeds, cache_position, past_key_values, use_cache, output_attentions + ) + + # Padding mask for mamba layers (just the raw attention_mask) + mamba_mask = attention_mask + + # Pre-compute rotary embeddings for attention layers + position_embeddings = self.rotary_emb(inputs_embeds, position_ids) + + hidden_states = inputs_embeds + + # Decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + # Select the appropriate mask for each block type + if decoder_layer.block_type == "attention": + layer_mask = causal_mask + elif decoder_layer.block_type == "mamba": + layer_mask = mamba_mask + else: # moe + layer_mask = None + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + layer_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=layer_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + + hidden_states = layer_outputs[0] + + if use_cache and decoder_layer.block_type == "attention": + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions and decoder_layer.block_type == "attention": + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # Add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + if return_legacy_cache and next_cache is not None: + next_cache = next_cache.to_legacy_cache() + + if not return_dict: + return tuple( + v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None + ) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + use_cache: bool, + output_attentions: bool, + ): + """Build a 4-D causal attention mask for the attention layers. + + This is only used by attention blocks; mamba blocks use the raw padding + mask and MoE blocks use no mask. + + Args: + attention_mask: 2-D or 4-D attention mask. + input_tensor: Input embeddings, used for shape and device info. + cache_position: Current cache positions. + past_key_values: Cached states. + use_cache: Whether caching is enabled. + output_attentions: Whether attention weights are requested. + + Returns: + 4-D causal mask tensor or ``None``. + """ + # cache_position must be valid here + past_seen_tokens = cache_position[0] if past_key_values is not None else 0 + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + if attention_mask is not None and attention_mask.dim() == 4: + # Already inverted form + if attention_mask.max() != 0: + raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`") + causal_mask = attention_mask + else: + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device + ) + exclude_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + causal_mask *= exclude_mask + causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() + if attention_mask.dim() == 2: + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + + return causal_mask + + +# --------------------------------------------------------------------------- +# CausalLM head +# --------------------------------------------------------------------------- +class MergedNemotronHForCausalLM(MergedNemotronHPreTrainedModel): + """Nemotron-H model with a causal language modeling head for DAM. + + Consists of the ``MergedNemotronHModel`` decoder and a linear ``lm_head`` + that projects hidden states to vocabulary logits. + """ + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config: MergedNemotronHConfig): + super().__init__(config) + self.model = MergedNemotronHModel(config) + self.vocab_size = config.vocab_size + self.num_merged_models = config.num_merged_models + self.lm_head = DAMLinearLayer( + config.hidden_size, config.vocab_size, bias=False, num_models=config.num_merged_models + ) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def tie_weights(self): + if isinstance(self.get_input_embeddings(), DAMEmbeddingLayer) and isinstance(self.lm_head, DAMLinearLayer): + self.lm_head.tie_with_embeddings(self.get_input_embeddings()) + else: + super().tie_weights() + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, CausalLMOutputWithPast]: + r"""Forward pass for the causal language model. + + Args: + input_ids: Token indices of shape ``(batch_size, sequence_length)``. + attention_mask: Padding mask. + position_ids: Position indices. + past_key_values: Cached key/value states. + inputs_embeds: Pre-computed embeddings. + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + use_cache: Whether to return updated caches. + output_attentions: Whether to return attention weights. + output_hidden_states: Whether to return all hidden states. + return_dict: Whether to return a ``CausalLMOutputWithPast``. + cache_position: Cache position indices. + + Returns: + ``CausalLMOutputWithPast`` or tuple of tensors. + + Example: + + ```python + >>> from transformers import AutoTokenizer + >>> model = MergedNemotronHForCausalLM.from_pretrained("nvidia/nemotron-h") + >>> tokenizer = AutoTokenizer.from_pretrained("nvidia/nemotron-h") + >>> prompt = "The future of AI is" + >>> inputs = tokenizer(prompt, return_tensors="pt") + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True)[0] + ``` + """ + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # Decoder outputs: (hidden_states, next_cache, all_hidden_states, all_self_attns) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + logits = logits.float() + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Ensure tensors are on the same device + shift_labels = shift_labels.to(shift_logits.device) + loss_fct = CrossEntropyLoss() + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + cache_position=None, + position_ids=None, + use_cache=True, + **kwargs, + ): + """Prepare model inputs for auto-regressive generation. + + Handles cache slicing and position_ids creation for efficient + generation with past key values. + + Args: + input_ids: Current input token ids. + past_key_values: Cached key/value states from previous steps. + attention_mask: Attention/padding mask. + inputs_embeds: Pre-computed embeddings (used on first step only). + cache_position: Position indices into the cache. + position_ids: Position indices for the current step. + use_cache: Whether to use caching. + + Returns: + Dictionary of model inputs. + """ + # If we have cache: slice input_ids to keep only unprocessed tokens + # Exception 1: when passing input_embeds, input_ids may be missing entries + # Exception 2: some generation methods do special slicing of input_ids + if past_key_values is not None: + if inputs_embeds is not None: # Exception 1 + input_ids = input_ids[:, -cache_position.shape[0] :] + elif input_ids.shape[1] != cache_position.shape[0]: # Default case (Exception 2 is a no-op) + input_ids = input_ids[:, cache_position] + + if attention_mask is not None and position_ids is None: + # Create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + # Clone to avoid stride issues with torch.compile + position_ids = position_ids.clone(memory_format=torch.contiguous_format) + + # If `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and cache_position[0] == 0: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids.contiguous()} + + model_inputs.update( + { + "position_ids": position_ids, + "cache_position": cache_position, + "past_key_values": past_key_values, + "use_cache": use_cache, + "attention_mask": attention_mask, + } + ) + return model_inputs diff --git a/dam/modeling/qwen3_5/__init__.py b/dam/modeling/qwen3_5/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/dam/modeling/qwen3_5/config.py b/dam/modeling/qwen3_5/config.py new file mode 100644 index 0000000..28ba224 --- /dev/null +++ b/dam/modeling/qwen3_5/config.py @@ -0,0 +1,199 @@ +# Copyright 2025 Qwen team, Alibaba Cloud and the HuggingFace Inc. team. 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. +"""Merged Qwen 3.5 text model configuration for DAM.""" + +from transformers.configuration_utils import PretrainedConfig + + +class MergedQwen3_5Config(PretrainedConfig): + r""" + Configuration class for the DAM-merged Qwen 3.5 text model. + + This extends the Qwen3_5TextConfig with DAM-specific parameters for + differentiable adaptive merging. Qwen 3.5 uses a hybrid architecture + with both full (softmax) attention and linear (Gated Delta Network) + attention layers. + + Args: + vocab_size (`int`, *optional*, defaults to 248320): + Vocabulary size of the Qwen 3.5 model. + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 12288): + Dimension of the MLP representations. + num_hidden_layers (`int`, *optional*, defaults to 32): + Number of hidden layers in the Transformer decoder. + num_attention_heads (`int`, *optional*, defaults to 16): + Number of attention heads for full attention layers. + num_key_value_heads (`int`, *optional*, defaults to 4): + Number of key_value heads for GQA in full attention layers. + head_dim (`int`, *optional*, defaults to 256): + Attention head dimension for full attention layers. + hidden_act (`str`, *optional*, defaults to `"silu"`): + The non-linear activation function in the decoder. + max_position_embeddings (`int`, *optional*, defaults to 32768): + The maximum sequence length. + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer. + rms_norm_eps (`float`, *optional*, defaults to 1e-6): + The epsilon used by the rms normalization layers. + use_cache (`bool`, *optional*, defaults to `True`): + Whether to return last key/values attentions. + tie_word_embeddings (`bool`, *optional*, defaults to `False`): + Whether to tie weight embeddings. + rope_theta (`float`, *optional*, defaults to 10000000.0): + The base period of the RoPE embeddings. + attention_bias (`bool`, *optional*, defaults to `False`): + Whether to use bias in attention projection layers. + attention_dropout (`float`, *optional*, defaults to 0.0): + Attention dropout rate. + layer_types (`list[str]`, *optional*): + Per-layer type specification. Each entry is either "full_attention" + or "linear_attention". If None, auto-generated from + full_attention_interval. + full_attention_interval (`int`, *optional*, defaults to 4): + Every Nth layer uses full attention when layer_types is auto-generated. + partial_rotary_factor (`float`, *optional*, defaults to 0.25): + Fraction of head_dim that gets rotary position embeddings. + linear_conv_kernel_dim (`int`, *optional*, defaults to 4): + Kernel size for causal conv1d in linear attention layers. + linear_key_head_dim (`int`, *optional*, defaults to 128): + Per-head key dimension in linear attention. + linear_value_head_dim (`int`, *optional*, defaults to 128): + Per-head value dimension in linear attention. + linear_num_key_heads (`int`, *optional*, defaults to 16): + Number of key heads in linear attention. + linear_num_value_heads (`int`, *optional*, defaults to 32): + Number of value heads in linear attention. + attn_output_gate (`bool`, *optional*, defaults to `True`): + Whether full attention uses output gating. + num_merged_models (`int`, *optional*, defaults to 3): + The number of models being merged. + init_merger_values (`list[float]`, *optional*, defaults to []): + Initial values for the merger coefficients. + use_tanh (`bool`, *optional*, defaults to False): + Whether to apply tanh non-linearity to merger coefficients. + dam_embedding_layer (`bool`, *optional*, defaults to True): + Whether embedding layers use DAM merging. + dam_layernorms (`bool`, *optional*, defaults to True): + Whether layer normalization uses DAM merging. + uses_base_model (`bool`, *optional*, defaults to True): + Whether the base model is included in the merge. + is_embedding_coef_trainable (`bool`, *optional*, defaults to False): + Whether embedding merger coefficients are trainable. + is_norm_coef_trainable (`bool`, *optional*, defaults to False): + Whether norm merger coefficients are trainable. + """ + + model_type = "mergedqwen3_5" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=248320, + hidden_size=4096, + intermediate_size=12288, + num_hidden_layers=32, + num_attention_heads=16, + num_key_value_heads=4, + head_dim=256, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=None, + bos_token_id=None, + eos_token_id=248044, + tie_word_embeddings=False, + rope_theta=10000000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + # Qwen 3.5 hybrid attention parameters + layer_types=None, + full_attention_interval=4, + partial_rotary_factor=0.25, + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=32, + attn_output_gate=True, + # DAM-specific parameters + num_merged_models=3, + init_merger_values=[], + use_tanh=False, + model_index=None, + dam_embedding_layer=True, + dam_layernorms=True, + uses_base_model=True, + is_embedding_coef_trainable=False, + is_norm_coef_trainable=False, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim + self.hidden_act = hidden_act + self.max_position_embeddings = max_position_embeddings + self.initializer_range = initializer_range + self.rms_norm_eps = rms_norm_eps + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + + # Qwen 3.5 hybrid attention + self.full_attention_interval = full_attention_interval + self.partial_rotary_factor = partial_rotary_factor + self.linear_conv_kernel_dim = linear_conv_kernel_dim + self.linear_key_head_dim = linear_key_head_dim + self.linear_value_head_dim = linear_value_head_dim + self.linear_num_key_heads = linear_num_key_heads + self.linear_num_value_heads = linear_num_value_heads + self.attn_output_gate = attn_output_gate + + # Auto-generate layer_types if not provided + if layer_types is None: + self.layer_types = [ + "linear_attention" if bool((i + 1) % full_attention_interval) else "full_attention" + for i in range(num_hidden_layers) + ] + else: + self.layer_types = layer_types + + # DAM-specific + self.num_merged_models = num_merged_models + self.init_merger_values = init_merger_values + self.use_tanh = use_tanh + self.model_index = model_index + self.dam_embedding_layer = dam_embedding_layer + self.dam_layernorms = dam_layernorms + self.uses_base_model = uses_base_model + self.is_embedding_coef_trainable = is_embedding_coef_trainable + self.is_norm_coef_trainable = is_norm_coef_trainable + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/dam/modeling/qwen3_5/modeling.py b/dam/modeling/qwen3_5/modeling.py new file mode 100644 index 0000000..ba4ab84 --- /dev/null +++ b/dam/modeling/qwen3_5/modeling.py @@ -0,0 +1,1198 @@ +"""PyTorch Merged Qwen 3.5 text model for DAM (Differentiable Adaptive Merging).""" + +import math +from typing import List, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from torch import nn +from torch.nn import CrossEntropyLoss + +from transformers.activations import ACT2FN +from transformers.cache_utils import Cache, DynamicCache +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import logging + +from .config import MergedQwen3_5Config +from ..dam import DAMLinearLayer, DAMEmbeddingLayer, DAMRMSNorm + +logger = logging.get_logger(__name__) + +_CONFIG_FOR_DOC = "MergedQwen3_5Config" + + +# --------------------------------------------------------------------------- +# RMSNorm variants +# --------------------------------------------------------------------------- + +class Qwen3_5RMSNorm(nn.Module): + """Qwen 3.5 RMS normalization using the (1 + weight) formulation. + + Weight is initialized to zeros so that at init time the norm is an identity + scaling (the ``1 +`` term). This is the non-DAM variant used when + ``config.dam_layernorms`` is ``False``. + """ + + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.zeros(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return ((1.0 + self.weight.float()) * hidden_states).to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class Qwen3_5DAMRMSNorm(DAMRMSNorm): + """DAM-merged RMS normalization with the Qwen 3.5 ``(1 + weight)`` formulation. + + Inherits from :class:`DAMRMSNorm` and overrides ``forward`` so that the + merged weight is used as ``(1 + merged_weight) * normed_input`` instead of + the standard ``merged_weight * normed_input``. + """ + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.model_index is not None: + weight = self.weights[self.model_index].to(hidden_states.device) + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + return ((1.0 + weight.float()) * hidden_states).to(input_dtype) + else: + weight = self.get_dam_weight().to(hidden_states.device) + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.eps) + return ((1.0 + weight.float()) * hidden_states).to(input_dtype) + + +class Qwen3_5GatedRMSNorm(nn.Module): + """Gated RMS normalization used by the Gated Delta Network linear attention. + + This norm is **not** DAM-merged because it is specific to each layer and + does not correspond to a weight that needs merging across models. + """ + + def __init__(self, hidden_size, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.zeros(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states, gate=None): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + hidden_states = (1.0 + self.weight.float()) * hidden_states + if gate is not None: + hidden_states = hidden_states * F.silu(gate.to(torch.float32)) + return hidden_states.to(input_dtype) + + +# --------------------------------------------------------------------------- +# Rotary embeddings (partial) +# --------------------------------------------------------------------------- + +class MergedQwen3_5RotaryEmbedding(nn.Module): + """Rotary position embedding for Qwen 3.5. + + Only ``rotary_dim`` dimensions out of the full ``head_dim`` receive + rotary encoding (controlled by ``partial_rotary_factor``). + """ + + def __init__(self, rotary_dim, max_position_embeddings=32768, base=10000000.0, device=None): + super().__init__() + self.rotary_dim = rotary_dim + self.max_position_embeddings = max_position_embeddings + self.base = base + inv_freq = 1.0 / ( + self.base ** (torch.arange(0, self.rotary_dim, 2, dtype=torch.int64).float().to(device) / self.rotary_dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + @torch.no_grad() + def forward(self, x, position_ids): + # x: [bs, num_heads, seq_len, head_dim] + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + """Apply partial rotary position embeddings to query and key tensors. + + ``cos`` and ``sin`` have shape ``[batch, seq_len, rotary_dim]``. Only the + first ``rotary_dim`` dimensions of *q* and *k* are rotated; the rest pass + through unchanged. + """ + rotary_dim = cos.shape[-1] + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + + q_rot = q[..., :rotary_dim] + q_pass = q[..., rotary_dim:] + k_rot = k[..., :rotary_dim] + k_pass = k[..., rotary_dim:] + + q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin) + k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin) + + q_embed = torch.cat([q_embed, q_pass], dim=-1) + k_embed = torch.cat([k_embed, k_pass], dim=-1) + return q_embed, k_embed + + +# --------------------------------------------------------------------------- +# Utility: repeat KV heads +# --------------------------------------------------------------------------- + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Expand key/value heads for grouped-query attention. + + ``hidden_states`` has shape ``(batch, num_kv_heads, seqlen, head_dim)`` + and is expanded to ``(batch, num_kv_heads * n_rep, seqlen, head_dim)``. + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +# --------------------------------------------------------------------------- +# MLP +# --------------------------------------------------------------------------- + +class MergedQwen3_5MLP(nn.Module): + """Standard SwiGLU MLP with DAM-merged linear layers.""" + + def __init__(self, config: MergedQwen3_5Config): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = DAMLinearLayer( + self.hidden_size, self.intermediate_size, bias=False, num_models=config.num_merged_models + ) + self.up_proj = DAMLinearLayer( + self.hidden_size, self.intermediate_size, bias=False, num_models=config.num_merged_models + ) + self.down_proj = DAMLinearLayer( + self.intermediate_size, self.hidden_size, bias=False, num_models=config.num_merged_models + ) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, hidden_state): + return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state)) + + +# --------------------------------------------------------------------------- +# Full (softmax) attention with gating and QK-norm +# --------------------------------------------------------------------------- + +class MergedQwen3_5Attention(nn.Module): + """Multi-headed full attention with output gating and QK normalization. + + Key features specific to Qwen 3.5: + * **Output gating**: ``q_proj`` produces 2x the normal output; the extra + half is a *gate* applied after attention via ``sigmoid(gate) * attn_out``. + * **QK normalization**: separate per-head RMSNorm on query and key states. + * **Partial RoPE**: only ``partial_rotary_factor`` fraction of ``head_dim`` + receives rotary embeddings. + * **Grouped-query attention**: ``num_key_value_heads`` may differ from + ``num_attention_heads``. + """ + + def __init__(self, config: MergedQwen3_5Config, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " + "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = config.head_dim + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.is_causal = True + self.partial_rotary_factor = config.partial_rotary_factor + self.attn_output_gate = config.attn_output_gate + + # Q outputs 2x for gating when attn_output_gate is True + q_output_dim = 2 * self.num_heads * self.head_dim if self.attn_output_gate else self.num_heads * self.head_dim + self.q_proj = DAMLinearLayer( + self.hidden_size, q_output_dim, bias=config.attention_bias, num_models=config.num_merged_models + ) + self.k_proj = DAMLinearLayer( + self.hidden_size, self.num_key_value_heads * self.head_dim, + bias=config.attention_bias, num_models=config.num_merged_models + ) + self.v_proj = DAMLinearLayer( + self.hidden_size, self.num_key_value_heads * self.head_dim, + bias=config.attention_bias, num_models=config.num_merged_models + ) + self.o_proj = DAMLinearLayer( + self.num_heads * self.head_dim, self.hidden_size, + bias=config.attention_bias, num_models=config.num_merged_models + ) + + # QK norms (per head_dim) + if config.dam_layernorms: + self.q_norm = Qwen3_5DAMRMSNorm( + self.head_dim, eps=config.rms_norm_eps, num_models=config.num_merged_models + ) + self.k_norm = Qwen3_5DAMRMSNorm( + self.head_dim, eps=config.rms_norm_eps, num_models=config.num_merged_models + ) + else: + self.q_norm = Qwen3_5RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Qwen3_5RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + + # Split Q into query and gate when gating is enabled + if self.attn_output_gate: + query_states, gate = torch.chunk(query_states, 2, dim=-1) + else: + gate = None + + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # Reshape to heads: (B, L, H, D) -> (B, H, L, D) + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + if gate is not None: + gate = gate.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + + # Apply QK norms + query_states = self.q_norm(query_states) + key_states = self.k_norm(key_states) + + # Apply partial RoPE + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + # KV cache + if past_key_value is not None: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + # Expand KV for GQA + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + + # Attention scores + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + # Upcast to fp32 for softmax + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + # Apply output gate + if gate is not None: + attn_output = attn_output * torch.sigmoid(gate) + + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.view(bsz, q_len, -1) + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +# --------------------------------------------------------------------------- +# Linear attention: Gated Delta Network +# --------------------------------------------------------------------------- + +class MergedQwen3_5GatedDeltaNet(nn.Module): + """Gated Delta Network linear attention layer. + + This implements a recurrent linear attention mechanism using the *delta + rule* with exponential gating. It does **not** require any external + libraries (causal_conv1d, fla, etc.) -- everything is pure PyTorch. + + Layers that are **not** DAM-merged: + * ``conv1d`` -- standard ``nn.Conv1d`` (depthwise causal convolution) + * ``A_log``, ``dt_bias`` -- ``nn.Parameter`` + * ``norm`` -- ``Qwen3_5GatedRMSNorm`` (gated, not DAM-merged) + + All linear projections are ``DAMLinearLayer``. + """ + + def __init__(self, config: MergedQwen3_5Config, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + + self.hidden_size = config.hidden_size + self.linear_num_key_heads = config.linear_num_key_heads + self.linear_num_value_heads = config.linear_num_value_heads + self.linear_key_head_dim = config.linear_key_head_dim + self.linear_value_head_dim = config.linear_value_head_dim + + self.conv_dim = ( + self.linear_num_key_heads * self.linear_key_head_dim + + self.linear_num_key_heads * self.linear_key_head_dim + + self.linear_num_value_heads * self.linear_value_head_dim + ) + self.value_total_dim = self.linear_num_value_heads * self.linear_value_head_dim + + # Linear projections (DAM-merged) + self.in_proj_qkv = DAMLinearLayer( + self.hidden_size, self.conv_dim, bias=False, num_models=config.num_merged_models + ) + self.in_proj_z = DAMLinearLayer( + self.hidden_size, self.value_total_dim, bias=False, num_models=config.num_merged_models + ) + self.in_proj_b = DAMLinearLayer( + self.hidden_size, self.linear_num_value_heads, bias=False, num_models=config.num_merged_models + ) + self.in_proj_a = DAMLinearLayer( + self.hidden_size, self.linear_num_value_heads, bias=False, num_models=config.num_merged_models + ) + self.out_proj = DAMLinearLayer( + self.value_total_dim, self.hidden_size, bias=False, num_models=config.num_merged_models + ) + + # Causal depthwise conv1d (standard, NOT DAM) + self.conv1d = nn.Conv1d( + in_channels=self.conv_dim, + out_channels=self.conv_dim, + kernel_size=config.linear_conv_kernel_dim, + padding=config.linear_conv_kernel_dim - 1, + groups=self.conv_dim, + ) + + # Learnable parameters (standard, NOT DAM) + self.A_log = nn.Parameter(torch.zeros(self.linear_num_value_heads)) + self.dt_bias = nn.Parameter(torch.zeros(self.linear_num_value_heads)) + + # Gated RMS norm (NOT DAM-merged) + self.norm = Qwen3_5GatedRMSNorm(self.value_total_dim, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + batch_size, seq_len, _ = hidden_states.shape + + # Project inputs + qkv = self.in_proj_qkv(hidden_states) # (B, L, conv_dim) + z = self.in_proj_z(hidden_states) # (B, L, value_total_dim) -- gate + beta = torch.sigmoid(self.in_proj_b(hidden_states)) # (B, L, num_v_heads) + alpha = self.in_proj_a(hidden_states) # (B, L, num_v_heads) + + # Causal depthwise conv1d + SiLU + qkv = qkv.transpose(1, 2) # (B, conv_dim, L) + qkv = self.conv1d(qkv)[:, :, :seq_len] # trim to causal (remove right padding) + qkv = F.silu(qkv) + qkv = qkv.transpose(1, 2) # (B, L, conv_dim) + + # Split into Q, K, V + q_dim = self.linear_num_key_heads * self.linear_key_head_dim + k_dim = self.linear_num_key_heads * self.linear_key_head_dim + v_dim = self.linear_num_value_heads * self.linear_value_head_dim + q, k, v = torch.split(qkv, [q_dim, k_dim, v_dim], dim=-1) + + # Reshape to heads + q = q.view(batch_size, seq_len, self.linear_num_key_heads, self.linear_key_head_dim) + k = k.view(batch_size, seq_len, self.linear_num_key_heads, self.linear_key_head_dim) + v = v.view(batch_size, seq_len, self.linear_num_value_heads, self.linear_value_head_dim) + + # L2 normalize Q and K + q = F.normalize(q, p=2, dim=-1) + k = F.normalize(k, p=2, dim=-1) + + # Compute decay gate: g = -exp(A_log) * softplus(alpha + dt_bias) + # alpha is (B, L, num_v_heads), dt_bias is (num_v_heads,) + g = -torch.exp(self.A_log.float()) * F.softplus(alpha.unsqueeze(-1) + self.dt_bias) + # g shape: (B, L, num_v_heads, 1) -- alpha was (B, L, num_v_heads), after unsqueeze(-1) + broadcast + # Actually alpha is (B, L, num_v_heads), dt_bias is (num_v_heads,), sum is (B, L, num_v_heads) + # softplus gives (B, L, num_v_heads), multiply gives (B, L, num_v_heads) + # We need to fix the shape computation: + g = -torch.exp(self.A_log.float()) * F.softplus(alpha + self.dt_bias) # (B, L, num_v_heads) + g = g.transpose(1, 2).unsqueeze(-1) # (B, num_v_heads, L, 1) + + # If num_v_heads > num_k_heads, repeat-interleave q and k + num_kv_groups = self.linear_num_value_heads // self.linear_num_key_heads + if num_kv_groups > 1: + q = q.repeat_interleave(num_kv_groups, dim=2) + k = k.repeat_interleave(num_kv_groups, dim=2) + + # Transpose for computation: (B, heads, L, dim) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + beta = beta.unsqueeze(-1).transpose(1, 2) # (B, num_v_heads, L, 1) + + # Recurrent delta rule computation + state = torch.zeros( + batch_size, self.linear_num_value_heads, + self.linear_key_head_dim, self.linear_value_head_dim, + dtype=torch.float32, device=hidden_states.device + ) + outputs = [] + + for t in range(seq_len): + q_t = q[:, :, t:t+1, :] # (B, H, 1, Dk) + k_t = k[:, :, t:t+1, :] # (B, H, 1, Dk) + v_t = v[:, :, t:t+1, :] # (B, H, 1, Dv) + b_t = beta[:, :, t:t+1, :] # (B, H, 1, 1) + g_t = g[:, :, t:t+1, :] # (B, H, 1, 1) + + # Decay + decay = torch.exp(g_t) # (B, H, 1, 1) + + # Delta rule: + # kv_product = k_t @ state -> (B, H, 1, Dk) @ (B, H, Dk, Dv) = (B, H, 1, Dv) + kv_product = torch.matmul(k_t, state.to(k_t.dtype)) + delta = (v_t - kv_product) * b_t # (B, H, 1, Dv) + + # Update state: state = state * decay + k_t^T @ delta + # k_t^T is (B, H, Dk, 1), delta is (B, H, 1, Dv) -> outer product (B, H, Dk, Dv) + state = state * decay + torch.matmul(k_t.transpose(-1, -2), delta).to(state.dtype) + + # Output: q_t @ state = (B, H, 1, Dk) @ (B, H, Dk, Dv) = (B, H, 1, Dv) + o_t = torch.matmul(q_t, state.to(q_t.dtype)) + outputs.append(o_t) + + output = torch.cat(outputs, dim=2) # (B, H, L, Dv) + output = output.transpose(1, 2).contiguous() # (B, L, H, Dv) + output = output.view(batch_size, seq_len, -1) # (B, L, H*Dv) + + # Apply gated RMS norm + output = self.norm(output, z) + + # Output projection + output = self.out_proj(output) + + # Linear attention does not produce per-head attention weights + return output, None, past_key_value + + +# --------------------------------------------------------------------------- +# Decoder layer (hybrid: dispatches to full or linear attention) +# --------------------------------------------------------------------------- + +class MergedQwen3_5DecoderLayer(nn.Module): + """Qwen 3.5 decoder layer that dispatches to either full (softmax) or + linear (Gated Delta Network) attention based on ``config.layer_types``. + + Both branches share the same pre-/post-attention layer norms and MLP. + """ + + def __init__(self, config: MergedQwen3_5Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + self.layer_type = config.layer_types[layer_idx] + + # Attention: pick based on layer type + if self.layer_type == "linear_attention": + self.linear_attn = MergedQwen3_5GatedDeltaNet(config=config, layer_idx=layer_idx) + else: + self.self_attn = MergedQwen3_5Attention(config=config, layer_idx=layer_idx) + + # MLP + self.mlp = MergedQwen3_5MLP(config) + + # Layer norms + NormClass = ( + lambda size, eps: Qwen3_5DAMRMSNorm(size, eps=eps, num_models=config.num_merged_models) + if config.dam_layernorms + else lambda size, eps: Qwen3_5RMSNorm(size, eps=eps) + ) + self.input_layernorm = ( + Qwen3_5DAMRMSNorm(config.hidden_size, eps=config.rms_norm_eps, num_models=config.num_merged_models) + if config.dam_layernorms + else Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + ) + self.post_attention_layernorm = ( + Qwen3_5DAMRMSNorm(config.hidden_size, eps=config.rms_norm_eps, num_models=config.num_merged_models) + if config.dam_layernorms + else Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): + attention mask of size `(batch_size, 1, query_sequence_length, key_sequence_length)` for + full attention, or `None` for linear attention. + position_ids (`torch.LongTensor`, *optional*): + position ids of shape `(batch_size, sequence_length)`. + past_key_value (`Cache`, *optional*): cached past key and value projection states. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned. + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. + position_embeddings (`Tuple[torch.Tensor, torch.Tensor]`, *optional*): + Tuple of (cos, sin) for rotary position embeddings. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that inject code + into the model. + """ + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + # Dispatch to appropriate attention + if self.layer_type == "linear_attention": + hidden_states, self_attn_weights, present_key_value = self.linear_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + else: + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + + hidden_states = residual + hidden_states + + # MLP + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + +# --------------------------------------------------------------------------- +# Docstrings +# --------------------------------------------------------------------------- + +QWEN3_5_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`MergedQwen3_5Config`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +QWEN3_5_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance; + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + + +# --------------------------------------------------------------------------- +# Pre-trained model base +# --------------------------------------------------------------------------- + +class MergedQwen3_5PreTrainedModel(PreTrainedModel): + config_class = MergedQwen3_5Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["MergedQwen3_5DecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn_2 = False + _supports_sdpa = False + _supports_cache_class = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +# --------------------------------------------------------------------------- +# Text model (decoder) +# --------------------------------------------------------------------------- + +class MergedQwen3_5TextModel(MergedQwen3_5PreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a + [`MergedQwen3_5DecoderLayer`] that uses either full or linear attention depending on the + layer type specified in `config.layer_types`. + + Args: + config: MergedQwen3_5Config + """ + + def __init__(self, config: MergedQwen3_5Config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = ( + DAMEmbeddingLayer( + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + num_models=config.num_merged_models, + padding_idx=self.padding_idx, + ) + if config.dam_embedding_layer + else nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + ) + + self.layers = nn.ModuleList( + [MergedQwen3_5DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + + self.norm = ( + Qwen3_5DAMRMSNorm(config.hidden_size, eps=config.rms_norm_eps, num_models=config.num_merged_models) + if config.dam_layernorms + else Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + ) + + # Rotary embeddings (partial) + self.rotary_dim = int(config.head_dim * config.partial_rotary_factor) + self.rotary_emb = MergedQwen3_5RotaryEmbedding( + self.rotary_dim, + max_position_embeddings=config.max_position_embeddings, + base=config.rope_theta, + ) + + self.gradient_checkpointing = False + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one" + ) + + if self.gradient_checkpointing and self.training and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + return_legacy_cache = False + if use_cache and not isinstance(past_key_values, Cache) and not self.training: + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + return_legacy_cache = True + logger.warning_once( + "We detected that you are passing `past_key_values` as a tuple and this is deprecated and " + "will be removed in v4.43. Please use an appropriate `Cache` class " + "(https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)" + ) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = self._update_causal_mask( + attention_mask, inputs_embeds, cache_position, past_key_values, use_cache, output_attentions + ) + + hidden_states = inputs_embeds + + # Compute rotary embeddings once for all layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # Decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for decoder_layer in self.layers: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # Add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + if return_legacy_cache: + next_cache = next_cache.to_legacy_cache() + + if not return_dict: + return tuple( + v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None + ) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + use_cache: bool, + output_attentions: bool, + ): + """Build a 4-D causal attention mask for full-attention layers. + + Linear attention layers ignore the causal mask (they have their own + causal structure via the recurrent delta rule), so this mask is only + consumed by ``MergedQwen3_5Attention``. + """ + past_seen_tokens = cache_position[0] if past_key_values is not None else 0 + + dtype, device = input_tensor.dtype, input_tensor.device + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + if attention_mask is not None and attention_mask.dim() == 4: + # Already a 4-D mask in inverted form + if attention_mask.max() != 0: + raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`") + causal_mask = attention_mask + else: + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device + ) + exclude_mask = torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + causal_mask *= exclude_mask + causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + if attention_mask.dim() == 2: + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + + return causal_mask + + +# --------------------------------------------------------------------------- +# Causal LM head +# --------------------------------------------------------------------------- + +class MergedQwen3_5ForCausalLM(MergedQwen3_5PreTrainedModel): + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config: MergedQwen3_5Config): + super().__init__(config) + self.model = MergedQwen3_5TextModel(config) + self.vocab_size = config.vocab_size + self.num_merged_models = config.num_merged_models + self.lm_head = DAMLinearLayer( + config.hidden_size, config.vocab_size, bias=False, num_models=config.num_merged_models + ) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + def tie_weights(self): + if isinstance(self.get_input_embeddings(), DAMEmbeddingLayer) and isinstance(self.lm_head, DAMLinearLayer): + self.lm_head.tie_with_embeddings(self.get_input_embeddings()) + else: + super().tie_weights() + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + ) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer + >>> from dam.modeling.qwen3_5.modeling import MergedQwen3_5ForCausalLM + + >>> model = MergedQwen3_5ForCausalLM.from_pretrained("Qwen/Qwen3.5-7B") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3.5-7B") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # Decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + ) + + hidden_states = outputs[0] + logits = self.lm_head(hidden_states) + logits = logits.float() + + loss = None + if labels is not None: + # Shift so that tokens < n predict n + shift_logits = logits[..., :-1, :].contiguous() + shift_labels = labels[..., 1:].contiguous() + # Flatten the tokens + shift_logits = shift_logits.view(-1, self.config.vocab_size) + shift_labels = shift_labels.view(-1) + # Ensure tensors are on the same device + shift_labels = shift_labels.to(shift_logits.device) + loss_fct = CrossEntropyLoss() + loss = loss_fct(shift_logits, shift_labels) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + def prepare_inputs_for_generation( + self, + input_ids, + past_key_values=None, + attention_mask=None, + inputs_embeds=None, + cache_position=None, + position_ids=None, + use_cache=True, + **kwargs, + ): + # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens + # Exception 1: when passing input_embeds, input_ids may be missing entries + # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here + if past_key_values is not None: + if inputs_embeds is not None: # Exception 1 + input_ids = input_ids[:, -cache_position.shape[0] :] + elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2) + input_ids = input_ids[:, cache_position] + + if attention_mask is not None and position_ids is None: + # create position_ids on the fly for batch generation + position_ids = attention_mask.long().cumsum(-1) - 1 + position_ids.masked_fill_(attention_mask == 0, 1) + if past_key_values: + position_ids = position_ids[:, -input_ids.shape[1] :] + + # This `clone` call is needed to avoid recapturing cuda graphs with `torch.compile`'s + # `mode="reduce-overhead`, as otherwise the input `position_ids` would have various stride + # during the decoding. Here, simply using `.contiguous()` is not sufficient as in the + # batch size = 1 case, `position_ids` is already contiguous but with varying stride which + # retriggers a capture. + position_ids = position_ids.clone(memory_format=torch.contiguous_format) + + # if `inputs_embeds` are passed, we only want to use them in the 1st generation step + if inputs_embeds is not None and cache_position[0] == 0: + model_inputs = {"inputs_embeds": inputs_embeds} + else: + model_inputs = {"input_ids": input_ids.contiguous()} # `contiguous()` needed for compilation use cases + + model_inputs.update( + { + "position_ids": position_ids, + "cache_position": cache_position, + "past_key_values": past_key_values, + "use_cache": use_cache, + "attention_mask": attention_mask, + } + ) + return model_inputs diff --git a/dam/utils.py b/dam/utils.py index f512b78..726b094 100644 --- a/dam/utils.py +++ b/dam/utils.py @@ -2,6 +2,26 @@ from transformers.models.llama.modeling_llama import LlamaRMSNorm from transformers.models.mistral.modeling_mistral import MistralRMSNorm +# Qwen 3.5 and Nemotron-H norm classes (optional, requires recent transformers) +try: + from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5RMSNorm +except ImportError: + Qwen3_5RMSNorm = None + +try: + from transformers.models.nemotron_h.modeling_nemotron_h import NemotronHRMSNorm +except ImportError: + NemotronHRMSNorm = None + +# Collect all supported RMSNorm classes +_NORM_CLASSES = [LlamaRMSNorm, MistralRMSNorm] +if Qwen3_5RMSNorm is not None: + _NORM_CLASSES.append(Qwen3_5RMSNorm) +if NemotronHRMSNorm is not None: + _NORM_CLASSES.append(NemotronHRMSNorm) +_NORM_CLASSES = tuple(_NORM_CLASSES) + + def find_linear_layers(model): return [name for name, module in model.named_modules() if isinstance(module, torch.nn.Linear)] @@ -9,4 +29,4 @@ def find_embedding_layers(model): return [name for name, module in model.named_modules() if isinstance(module, torch.nn.Embedding)] def find_norm_layers(model): - return [name for name, module in model.named_modules() if isinstance(module, LlamaRMSNorm) or isinstance(module, MistralRMSNorm)] \ No newline at end of file + return [name for name, module in model.named_modules() if isinstance(module, _NORM_CLASSES)] \ No newline at end of file From 7e3b5f027130f231db362b405c43c54f7ab17e5b Mon Sep 17 00:00:00 2001 From: SolshineCode Date: Wed, 1 Apr 2026 16:51:57 -0700 Subject: [PATCH 2/2] Fix review findings from Cursor Bugbot, Gemini, and CPU smoke tests Cursor Bugbot fixes: - Fix Mamba-2 dA tensor reshape crash (A[None,:,None,None] for correct 4D broadcast) - Remove dead code in Qwen 3.5 decay gate computation - Remove unused NormClass lambda in decoder layer Gemini review fix: - Add null safety check on to_legacy_cache() in Qwen 3.5 model CPU smoke test fixes: - Add **kwargs to tie_weights() for transformers 5.x compatibility - Guard DynamicCache.from_legacy_cache with hasattr for forward compat All 40 CPU smoke tests pass (imports, config, model instantiation, forward pass, loss computation, DAM layer verification, fix_config, norm detection, AutoConfig registration). Co-Authored-By: Claude Opus 4.6 (1M context) --- dam/modeling/nemotron/modeling.py | 22 +++++++++++---------- dam/modeling/qwen3_5/modeling.py | 32 ++++++++++++------------------- 2 files changed, 24 insertions(+), 30 deletions(-) diff --git a/dam/modeling/nemotron/modeling.py b/dam/modeling/nemotron/modeling.py index 9288a38..d3013f1 100644 --- a/dam/modeling/nemotron/modeling.py +++ b/dam/modeling/nemotron/modeling.py @@ -478,7 +478,7 @@ def forward( # Discretize # dA: (B, num_heads, 1, 1) - dA = torch.exp(dt_t.unsqueeze(-1).unsqueeze(-1) * A.unsqueeze(-1).unsqueeze(0)) + dA = torch.exp(dt_t.unsqueeze(-1).unsqueeze(-1) * A[None, :, None, None]) # dB: (B, num_heads, head_dim, ssm_state_size) dB = dt_t.unsqueeze(-1).unsqueeze(-1) * B_t.unsqueeze(2) * x_t.unsqueeze(-1) @@ -903,13 +903,15 @@ def forward( return_legacy_cache = False if use_cache and not isinstance(past_key_values, Cache) and not self.training: - past_key_values = DynamicCache.from_legacy_cache(past_key_values) - return_legacy_cache = True - logger.warning_once( - "We detected that you are passing `past_key_values` as a tuple and this is deprecated and " - "will be removed in v4.43. Please use an appropriate `Cache` class " - "(https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)" - ) + if hasattr(DynamicCache, "from_legacy_cache"): + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + return_legacy_cache = True + logger.warning_once( + "We detected that you are passing `past_key_values` as a tuple and this is deprecated. " + "Please use an appropriate `Cache` class." + ) + else: + past_key_values = DynamicCache() # fresh cache for newer transformers if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 @@ -1108,11 +1110,11 @@ def set_decoder(self, decoder): def get_decoder(self): return self.model - def tie_weights(self): + def tie_weights(self, **kwargs): if isinstance(self.get_input_embeddings(), DAMEmbeddingLayer) and isinstance(self.lm_head, DAMLinearLayer): self.lm_head.tie_with_embeddings(self.get_input_embeddings()) else: - super().tie_weights() + super().tie_weights(**kwargs) def forward( self, diff --git a/dam/modeling/qwen3_5/modeling.py b/dam/modeling/qwen3_5/modeling.py index ba4ab84..705d322 100644 --- a/dam/modeling/qwen3_5/modeling.py +++ b/dam/modeling/qwen3_5/modeling.py @@ -474,11 +474,6 @@ def forward( # Compute decay gate: g = -exp(A_log) * softplus(alpha + dt_bias) # alpha is (B, L, num_v_heads), dt_bias is (num_v_heads,) - g = -torch.exp(self.A_log.float()) * F.softplus(alpha.unsqueeze(-1) + self.dt_bias) - # g shape: (B, L, num_v_heads, 1) -- alpha was (B, L, num_v_heads), after unsqueeze(-1) + broadcast - # Actually alpha is (B, L, num_v_heads), dt_bias is (num_v_heads,), sum is (B, L, num_v_heads) - # softplus gives (B, L, num_v_heads), multiply gives (B, L, num_v_heads) - # We need to fix the shape computation: g = -torch.exp(self.A_log.float()) * F.softplus(alpha + self.dt_bias) # (B, L, num_v_heads) g = g.transpose(1, 2).unsqueeze(-1) # (B, num_v_heads, L, 1) @@ -566,11 +561,6 @@ def __init__(self, config: MergedQwen3_5Config, layer_idx: int): self.mlp = MergedQwen3_5MLP(config) # Layer norms - NormClass = ( - lambda size, eps: Qwen3_5DAMRMSNorm(size, eps=eps, num_models=config.num_merged_models) - if config.dam_layernorms - else lambda size, eps: Qwen3_5RMSNorm(size, eps=eps) - ) self.input_layernorm = ( Qwen3_5DAMRMSNorm(config.hidden_size, eps=config.rms_norm_eps, num_models=config.num_merged_models) if config.dam_layernorms @@ -873,13 +863,15 @@ def forward( return_legacy_cache = False if use_cache and not isinstance(past_key_values, Cache) and not self.training: - past_key_values = DynamicCache.from_legacy_cache(past_key_values) - return_legacy_cache = True - logger.warning_once( - "We detected that you are passing `past_key_values` as a tuple and this is deprecated and " - "will be removed in v4.43. Please use an appropriate `Cache` class " - "(https://huggingface.co/docs/transformers/v4.41.3/en/internal/generation_utils#transformers.Cache)" - ) + if hasattr(DynamicCache, "from_legacy_cache"): + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + return_legacy_cache = True + logger.warning_once( + "We detected that you are passing `past_key_values` as a tuple and this is deprecated. " + "Please use an appropriate `Cache` class." + ) + else: + past_key_values = DynamicCache() # fresh cache for newer transformers if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 @@ -947,7 +939,7 @@ def forward( all_hidden_states += (hidden_states,) next_cache = next_decoder_cache if use_cache else None - if return_legacy_cache: + if return_legacy_cache and next_cache is not None: next_cache = next_cache.to_legacy_cache() if not return_dict: @@ -1050,11 +1042,11 @@ def set_decoder(self, decoder): def get_decoder(self): return self.model - def tie_weights(self): + def tie_weights(self, **kwargs): if isinstance(self.get_input_embeddings(), DAMEmbeddingLayer) and isinstance(self.lm_head, DAMLinearLayer): self.lm_head.tie_with_embeddings(self.get_input_embeddings()) else: - super().tie_weights() + super().tie_weights(**kwargs) def forward( self,