# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # This file was automatically generated from src/transformers/models/qwen3/modular_qwen3.py. # Do NOT edit this file manually as any edits will be overwritten by the generation of # the file from the modular. If any change should be done, please apply the change to the # modular_qwen3.py file directly. One of our CI enforces this. # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 # coding=utf-8 # Copyright 2025 The Qwen team, Alibaba Group 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. from dataclasses import dataclass from typing import Callable, Optional, Union import torch.nn.functional as F import torch from torch import nn from torch import Tensor from transformers.models.qwen3.modeling_qwen3 import ACT2FN, Cache, DynamicCache, GenerationMixin, use_kernel_forward_from_hub, create_causal_mask, create_sliding_window_causal_mask, FlashAttentionKwargs, GenericForQuestionAnswering, GenericForSequenceClassification, GenericForTokenClassification, GradientCheckpointingLayer, BaseModelOutputWithPast, CausalLMOutputWithPast, ROPE_INIT_FUNCTIONS, dynamic_rope_update, ALL_ATTENTION_FUNCTIONS, PreTrainedModel, Unpack, TransformersKwargs, auto_docstring, can_return_tuple, deprecate_kwarg, check_model_inputs, Qwen3Config, Qwen3RMSNorm, Qwen3MLP, Qwen3Attention, apply_rotary_pos_emb, eager_attention_forward, Qwen3RotaryEmbedding from math import sqrt from torch.utils.checkpoint import checkpoint from torch import Tensor class LoraLinear(nn.Linear): def __init__(self, in_features, out_features, bias=True, device=None, dtype=None): super().__init__(in_features, out_features, bias=bias, device=device, dtype=dtype) def forward(self, input: Tensor, lora_dict=None) -> Tensor: base = F.linear(input, self.weight, self.bias) if lora_dict is None: return base A = lora_dict["A"] # [Lb, in, r] B = lora_dict["B"] # [Lb, r, out] C = lora_dict.get("C", None) # [Lb, out] or None Lb = A.shape[0] torch._assert(input.shape[-1] == self.in_features, "last dim must be in_features") torch._assert(input.shape[0] % Lb == 0, "input batch must be multiple of lora batch") num_beams = input.shape[0] // Lb # Flatten all middle dims (e.g., seq_len) into S for faster matmul # input: [B, ..., in] -> x: [Lb, beams, S, in] x = input.reshape(Lb, num_beams, -1, self.in_features) # [Lb, beams, S, in] @ [Lb, in, r] -> [Lb, beams, S, r] tmp = torch.matmul(x, A[:, None, :, :]) # [Lb, beams, S, r] @ [Lb, r, out] -> [Lb, beams, S, out] lora_out = torch.matmul(tmp, B[:, None, :, :]) if self.bias is None: torch._assert(C is None, "If bias is None, lora_dict['C'] must also be None") else: torch._assert(C is not None, "If bias is not None, lora_dict['C'] must also be not None") # C: [Lb, out] -> [Lb, 1, 1, out] broadcast across beams and S lora_out = lora_out + C[:, None, None, :] # Restore original middle dims: [Lb*beams, ..., out] lora_out = lora_out.reshape(*input.shape[:-1], self.out_features) # try: # lora_out = lora_out.reshape(*input.shape[:-1], self.out_features) # except RuntimeError: # res = f"input shape: {input.shape}, A shape: {A.shape}, B shape: {B.shape}, C shape: {C.shape if C is not None else 'None'}, \ # in_features: {self.in_features}, out_features: {self.out_features}, tmp_shape: {tmp.shape}, lora_out shape before reshape: {lora_out.shape}, \ # num_beams: {num_beams}, Lb: {Lb}, x shape: {x.shape}" # print(res) return base + lora_out def lora_params_numel(self, r): if not hasattr(self, "lora_params_numel_cache"): self.lora_params_numel_cache = ( self.in_features * r + self.out_features * r + (self.out_features if self.bias is not None else 0) ) return self.lora_params_numel_cache def set_generate_func(self, method): if method == "rl": def generate_func(r, scale, plain_tensor): idx = 0 A = plain_tensor[:, idx: idx + self.in_features * r].view(-1, self.in_features, r) * sqrt(scale) idx += self.in_features * r B = plain_tensor[:, idx: idx + self.out_features * r].view(-1, r, self.out_features) * sqrt(scale) idx += self.out_features * r C = plain_tensor[:, idx: idx + self.out_features].view(-1, self.out_features) * scale if self.bias is not None else None return {"A": A, "B": B, "C": C} elif method == "rr": def generate_func(r, scale, plain_tensor): idx = 0 A = plain_tensor[:, idx: idx + self.in_features * r].view(-1, self.in_features, r) * sqrt(scale) idx += self.in_features * r B = plain_tensor[:, idx: idx + self.out_features * r].view(-1, self.out_features, r).transpose(-1, -2) * sqrt(scale) idx += self.out_features * r C = plain_tensor[:, idx: idx + self.out_features].view(-1, self.out_features) * scale if self.bias is not None else None return {"A": A, "B": B, "C": C} elif method == "lr": def generate_func(r, scale, plain_tensor): idx = 0 A = plain_tensor[:, idx: idx + self.in_features * r].view(-1, r, self.in_features).transpose(-1, -2) * sqrt(scale) idx += self.in_features * r B = plain_tensor[:, idx: idx + self.out_features * r].view(-1, self.out_features, r).transpose(-1, -2) * sqrt(scale) idx += self.out_features * r C = plain_tensor[:, idx: idx + self.out_features].view(-1, self.out_features) * scale if self.bias is not None else None return {"A": A, "B": B, "C": C} elif method == "ll": def generate_func(r, scale, plain_tensor): idx = 0 A = plain_tensor[:, idx: idx + self.in_features * r].view(-1, r, self.in_features).transpose(-1, -2) * sqrt(scale) idx += self.in_features * r B = plain_tensor[:, idx: idx + self.out_features * r].view(-1, r, self.out_features) * sqrt(scale) idx += self.out_features * r C = plain_tensor[:, idx: idx + self.out_features].view(-1, self.out_features) * scale if self.bias is not None else None return {"A": A, "B": B, "C": C} else: raise NotImplementedError(f"LoRA method {method} not implemented") self.generate_func = generate_func def generate_lora_dict(self, r, scale, plain_tensor): torch._assert( plain_tensor.shape[-1] == self.lora_params_numel(r), "plain_tensor last dim does not match lora_params_numel" ) torch._assert(hasattr(self, "generate_func"), "generate_func not set") return self.generate_func(r, scale, plain_tensor) def init_lora_dict(self, r, scale, device): assert r > 0, "r must be positive" A = (torch.randn(size=(1, self.in_features, r), device=device) * sqrt(scale)).detach() A.requires_grad_() B = torch.zeros(size=(1, r, self.out_features), requires_grad=True, device=device) C = torch.zeros(size=(1, self.out_features), requires_grad=True, device=device) if self.bias is not None else None return {"A": A, "B": B, "C": C} def divide_idx(self, r, idx_start): A_numel = self.in_features * r B_numel = self.out_features * r assert self.bias is None idx_range = [idx_start, A_numel + idx_start] return idx_range, A_numel + B_numel + idx_start # class LoraLinear(nn.Linear): # def __init__(self, in_features, out_features, bias = True, device=None, dtype=None): # super().__init__(in_features, out_features, bias, device, dtype) # def forward(self, input: Tensor, lora_dict = None) -> Tensor: # if lora_dict is None: # return F.linear(input, self.weight, self.bias) # else: # assert input.shape[0] % lora_dict["A"].shape[0] == 0, f"input batch size {input.shape[0]} must be multiple of lora_dict['A'] batch size {lora_dict['A'].shape[0]}" # num_beams = input.shape[0] // lora_dict["A"].shape[0] # if self.bias is None: # assert lora_dict.get("C") is None, "If bias is None, lora_dict['C'] must also be None" # else: # assert lora_dict.get("C") is not None, "If bias is not None, lora_dict['C'] must also be not None" # return F.linear(input, self.weight, self.bias) + torch.bmm(torch.bmm(input, lora_dict["A"].repeat_interleave(num_beams, dim=0)), lora_dict["B"].repeat_interleave(num_beams, dim=0)) + (lora_dict["C"].unsqueeze(1).repeat_interleave(num_beams, dim=0) if self.bias is not None else 0) # def lora_params_numel(self, r): # if not hasattr(self, "lora_params_numel_cache"): # self.lora_params_numel_cache = self.in_features * r + self.out_features * r + (self.out_features if self.bias is not None else 0) # return self.lora_params_numel_cache # def set_generate_func(self, method): # pass # # if method == "rl": # # def generate_func(r, scale, plain_tensor): # # idx = 0 # # A = plain_tensor[:, idx: idx + self.in_features * r].view(-1, self.in_features, r) * sqrt(scale) # # idx += self.in_features * r # # B = plain_tensor[:, idx: idx + self.out_features * r].view(-1, r, self.out_features)* sqrt(scale) # # idx += self.out_features * r # # C = plain_tensor[:, idx: idx + self.out_features].view(-1, self.out_features) * scale if self.bias is not None else None # # return {"A": A, "B": B, "C": C} # # elif method == "rr": # # def generate_func(r, scale, plain_tensor): # # idx = 0 # # A = plain_tensor[:, idx: idx + self.in_features * r].view(-1, self.in_features, r) * sqrt(scale) # # idx += self.in_features * r # # B = plain_tensor[:, idx: idx + self.out_features * r].view(-1, self.out_features, r).transpose(-1, -2) * sqrt(scale) # # idx += self.out_features * r # # C = plain_tensor[:, idx: idx + self.out_features].view(-1, self.out_features) * scale if self.bias is not None else None # # return {"A": A, "B": B, "C": C} # # elif method == "lr": # # def generate_func(r, scale, plain_tensor): # # idx = 0 # # A = plain_tensor[:, idx: idx + self.in_features * r].view(-1, r, self.in_features).transpose(-1, -2) * sqrt(scale) # # idx += self.in_features * r # # B = plain_tensor[:, idx: idx + self.out_features * r].view(-1, self.out_features, r).transpose(-1, -2) * sqrt(scale) # # idx += self.out_features * r # # C = plain_tensor[:, idx: idx + self.out_features].view(-1, self.out_features) * scale if self.bias is not None else None # # return {"A": A, "B": B, "C": C} # # elif method == "ll": # # def generate_func(r, scale, plain_tensor): # # idx = 0 # # A = plain_tensor[:, idx: idx + self.in_features * r].view(-1, r, self.in_features).transpose(-1, -2) * sqrt(scale) # # idx += self.in_features * r # # B = plain_tensor[:, idx: idx + self.out_features * r].view(-1, r, self.out_features)* sqrt(scale) # # idx += self.out_features * r # # C = plain_tensor[:, idx: idx + self.out_features].view(-1, self.out_features) * scale if self.bias is not None else None # # return {"A": A, "B": B, "C": C} # # else: # # raise NotImplementedError(f"LoRA method {method} not implemented") # # self.generate_func = generate_func # # def generate_lora_dict(self, r, scale, plain_tensor): # # assert plain_tensor.shape[-1] == self.lora_params_numel(r), f"plain_tensor's last dimension {plain_tensor.shape[-1]} does not match lora_params_numel {self.lora_params_numel(r)}" # # assert hasattr(self, "generate_func"), "generate_func not set" # # return self.generate_func(r, scale, plain_tensor) # def generate_lora_dict(self, r, scale, plain_tensor): # assert plain_tensor.shape[-1] == self.lora_params_numel(r), f"plain_tensor's last dimension {plain_tensor.shape[-1]} does not match lora_params_numel {self.lora_params_numel(r)}" # idx = 0 # A = plain_tensor[:, idx: idx + self.in_features * r].view(-1, self.in_features, r) * sqrt(scale) # idx += self.in_features * r # B = plain_tensor[:, idx: idx + self.out_features * r].view(-1, r, self.out_features)* sqrt(scale) # idx += self.out_features * r # C = plain_tensor[:, idx: idx + self.out_features].view(-1, self.out_features) * scale if self.bias is not None else None # return {"A": A, "B": B, "C": C} # def init_lora_dict(self, r, scale, device): # A = (torch.randn(size=(1, self.in_features, r), device=device) * sqrt(scale)).detach() # A.requires_grad_() # B = torch.zeros(size=(1, r, self.out_features), requires_grad=True, device=device) # C = torch.zeros(size=(1, self.out_features), requires_grad=True, device=device) if self.bias is not None else None # return {"A": A, "B": B, "C": C} class LoraQwen3MLP(Qwen3MLP): def __init__(self, config): super().__init__(config) self.gate_proj = LoraLinear(self.hidden_size, self.intermediate_size, bias=False) self.up_proj = LoraLinear(self.hidden_size, self.intermediate_size, bias=False) self.down_proj = LoraLinear(self.intermediate_size, self.hidden_size, bias=False) def forward(self, x, loradict=None): if loradict is None: down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) else: down_proj = self.down_proj(self.act_fn(self.gate_proj(x, loradict["gate"])) * self.up_proj(x, loradict["up"]), loradict["down"]) return down_proj def lora_params_numel(self, r): if not hasattr(self, "lora_params_numel_cache"): self.lora_params_numel_cache = self.gate_proj.lora_params_numel(r) + self.up_proj.lora_params_numel(r) + self.down_proj.lora_params_numel(r) return self.lora_params_numel_cache def set_generate_func(self, method): self.gate_proj.set_generate_func(method) self.up_proj.set_generate_func(method) self.down_proj.set_generate_func(method) def generate_lora_dict(self, r, scale, plain_tensor): assert plain_tensor.shape[-1] == self.lora_params_numel(r), f"plain_tensor's last dimension {plain_tensor.shape[-1]} does not match lora_params_numel {self.lora_params_numel(r)}" idx = 0 gate = self.gate_proj.generate_lora_dict(r, scale, plain_tensor[:, idx: idx + self.gate_proj.lora_params_numel(r)]) idx += self.gate_proj.lora_params_numel(r) up = self.up_proj.generate_lora_dict(r, scale, plain_tensor[:, idx: idx + self.up_proj.lora_params_numel(r)]) idx += self.up_proj.lora_params_numel(r) down = self.down_proj.generate_lora_dict(r, scale, plain_tensor[:, idx: idx + self.down_proj.lora_params_numel(r)]) return {"gate": gate, "up": up, "down": down} def init_lora_dict(self, r, scale, device): gate = self.gate_proj.init_lora_dict(r, scale, device) up = self.up_proj.init_lora_dict(r, scale, device) down = self.down_proj.init_lora_dict(r, scale, device) return {"gate": gate, "up": up, "down": down} def divide_idx(self, r, idx_start): idx_range = [] idx_range_gate, next_start = self.gate_proj.divide_idx(r, idx_start) idx_range += idx_range_gate idx_range_up, next_start = self.up_proj.divide_idx(r, next_start) idx_range += idx_range_up idx_range_down, next_start = self.down_proj.divide_idx(r, next_start) idx_range += idx_range_down return idx_range, next_start class LoraQwen3Attention(Qwen3Attention): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__(self, config, layer_idx): super().__init__(config, layer_idx) self.q_proj = LoraLinear( config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias ) self.k_proj = LoraLinear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias ) self.v_proj = LoraLinear( config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias ) self.o_proj = LoraLinear( config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias ) self.bias = config.attention_bias @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") def forward( self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor], attention_mask: Optional[torch.Tensor], past_key_values: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, loradict = None, **kwargs: Unpack[FlashAttentionKwargs], ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: input_shape = hidden_states.shape[:-1] hidden_shape = (*input_shape, -1, self.head_dim) if loradict is None: query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) else: query_states = self.q_norm(self.q_proj(hidden_states, loradict["q"]).view(hidden_shape)).transpose(1, 2) key_states = self.k_norm(self.k_proj(hidden_states, loradict["k"]).view(hidden_shape)).transpose(1, 2) value_states = self.v_proj(hidden_states, loradict["v"]).view(hidden_shape).transpose(1, 2) cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_values is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs) attention_interface: Callable = eager_attention_forward if self.config._attn_implementation != "eager": attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] attn_output, attn_weights = attention_interface( self, query_states, key_states, value_states, attention_mask, dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, sliding_window=self.sliding_window, # diff with Llama **kwargs, ) attn_output = attn_output.reshape(*input_shape, -1).contiguous() if loradict is None: attn_output = self.o_proj(attn_output) else: attn_output = self.o_proj(attn_output, loradict["o"]) return attn_output, attn_weights def lora_params_numel(self, r): if not hasattr(self, "lora_params_numel_cache"): self.lora_params_numel_cache = 0 self.lora_params_numel_cache += self.q_proj.lora_params_numel(r) self.lora_params_numel_cache += self.k_proj.lora_params_numel(r) self.lora_params_numel_cache += self.v_proj.lora_params_numel(r) self.lora_params_numel_cache += self.o_proj.lora_params_numel(r) return self.lora_params_numel_cache def set_generate_func(self, method): self.q_proj.set_generate_func(method) self.k_proj.set_generate_func(method) self.v_proj.set_generate_func(method) self.o_proj.set_generate_func(method) def generate_lora_dict(self, r, scale, plain_tensor): assert plain_tensor.shape[-1] == self.lora_params_numel(r), f"plain_tensor's last dimension {plain_tensor.shape[-1]} does not match lora_params_numel {self.lora_params_numel(r)}" idx = 0 q = self.q_proj.generate_lora_dict(r, scale, plain_tensor[:, idx: idx + self.q_proj.lora_params_numel(r)]) idx += self.q_proj.lora_params_numel(r) k = self.k_proj.generate_lora_dict(r, scale, plain_tensor[:, idx: idx + self.k_proj.lora_params_numel(r)]) idx += self.k_proj.lora_params_numel(r) v = self.v_proj.generate_lora_dict(r, scale, plain_tensor[:, idx: idx + self.v_proj.lora_params_numel(r)]) idx += self.v_proj.lora_params_numel(r) o = self.o_proj.generate_lora_dict(r, scale, plain_tensor[:, idx: idx + self.o_proj.lora_params_numel(r)]) return {"q": q, "k": k, "v": v, "o": o} def init_lora_dict(self, r, scale, device): q = self.q_proj.init_lora_dict(r, scale, device) k = self.k_proj.init_lora_dict(r, scale, device) v = self.v_proj.init_lora_dict(r, scale, device) o = self.o_proj.init_lora_dict(r, scale, device) return {"q": q, "k": k, "v": v, "o": o} def divide_idx(self, r, idx_start): idx_range = [] idx_range_q, next_start = self.q_proj.divide_idx(r, idx_start) idx_range += idx_range_q idx_range_k, next_start = self.k_proj.divide_idx(r, next_start) idx_range += idx_range_k idx_range_v, next_start = self.v_proj.divide_idx(r, next_start) idx_range += idx_range_v idx_range_o, next_start = self.o_proj.divide_idx(r, next_start) idx_range += idx_range_o return idx_range, next_start class Qwen3DecoderLayer(GradientCheckpointingLayer): def __init__(self, config: Qwen3Config, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.self_attn = Qwen3Attention(config=config, layer_idx=layer_idx) self.mlp = Qwen3MLP(config) self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.attention_type = config.layer_types[layer_idx] @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, **kwargs, ) hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states return hidden_states class LoraQwen3DecoderLayer(Qwen3DecoderLayer): def __init__(self, config: Qwen3Config, layer_idx: int): super(Qwen3DecoderLayer, self).__init__() self.hidden_size = config.hidden_size self.self_attn = LoraQwen3Attention(config=config, layer_idx=layer_idx) self.mlp = LoraQwen3MLP(config) self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.attention_type = config.layer_types[layer_idx] @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC loradict: Optional[dict] = None, **kwargs: Unpack[TransformersKwargs], ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Self Attention hidden_states, _ = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, loradict=loradict['attention'] if loradict is not None else None, **kwargs, ) hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states, loradict=loradict['mlp'] if loradict is not None else None) hidden_states = residual + hidden_states return hidden_states def lora_params_numel(self, r): if not hasattr(self, "lora_params_numel_cache"): self.lora_params_numel_cache = 0 self.lora_params_numel_cache += self.self_attn.lora_params_numel(r) self.lora_params_numel_cache += self.mlp.lora_params_numel(r) return self.lora_params_numel_cache def set_generate_func(self, method): self.self_attn.set_generate_func(method) self.mlp.set_generate_func(method) def generate_lora_dict(self, r, scale, plain_tensor): assert plain_tensor.shape[-1] == self.lora_params_numel(r), f"plain_tensor's last dimension {plain_tensor.shape[-1]} does not match lora_params_numel {self.lora_params_numel(r)}" idx = 0 attention = self.self_attn.generate_lora_dict(r, scale, plain_tensor[:, idx: idx + self.self_attn.lora_params_numel(r)]) idx += self.self_attn.lora_params_numel(r) mlp = self.mlp.generate_lora_dict(r, scale, plain_tensor[:, idx: idx + self.mlp.lora_params_numel(r)]) return {"attention": attention, "mlp": mlp} def init_lora_dict(self, r, scale, device): return {"attention": self.self_attn.init_lora_dict(r, scale, device), "mlp": self.mlp.init_lora_dict(r, scale, device)} def divide_idx(self, r, idx_start): idx_range = [] idx_range_attn, next_start = self.self_attn.divide_idx(r, idx_start) idx_range += idx_range_attn idx_range_mlp, next_start = self.mlp.divide_idx(r, next_start) idx_range += idx_range_mlp return idx_range, next_start @auto_docstring class Qwen3PreTrainedModel(PreTrainedModel): config: Qwen3Config base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["Qwen3DecoderLayer"] _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn = True _supports_sdpa = True _supports_flex_attn = True _can_compile_fullgraph = True _supports_attention_backend = True _can_record_outputs = { "hidden_states": Qwen3DecoderLayer, "attentions": Qwen3Attention, } @dataclass class MemoryModelOutputWithPast(BaseModelOutputWithPast): ''' memory_states: Optional[torch.Tensor] = None # (batch_size, num_layer, num_mem_token, hidden_size) reg_loss: Optional[torch.FloatTensor] = None ''' memory_states: Optional[torch.Tensor] = None reg_loss: Optional[torch.FloatTensor] = None @auto_docstring class LoraQwen3Model(Qwen3PreTrainedModel): def __init__(self, config: Qwen3Config): super().__init__(config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size if config.num_mem_token == -1: self.use_mem_token = False else: self.use_mem_token = True self.num_mem_token = config.num_mem_token self.mem_tokens = nn.Parameter(torch.zeros((self.num_mem_token, config.hidden_size), requires_grad=True), requires_grad=True) self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) self.layers = nn.ModuleList( [LoraQwen3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = Qwen3RotaryEmbedding(config=config) self.gradient_checkpointing = False self.has_sliding_layers = "sliding_attention" in self.config.layer_types # Initialize weights and apply final processing self.post_init() def reset_mem_tokens(self): if self.use_mem_token: nn.init.zeros_(self.mem_tokens) @check_model_inputs @auto_docstring def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, loradict: Optional[dict] = None, ignore_mem_token: bool = False, use_gradient_checkpoint: bool = False, **kwargs: Unpack[TransformersKwargs], ) -> BaseModelOutputWithPast: r""" loradict (`dict` of `dict` of `torch.FloatTensor`, *optional*): A dictionary that maps each layer to its corresponding LoRA parameters. Each layer's LoRA parameters are stored in a nested dictionary. ignore_mem_token (`bool`, *optional*, defaults to `False`): Whether to ignore the memory tokens during the forward pass. If set to `True`, the memory tokens will not be used, and the model will behave like a standard transformer without memory tokens. use_gradient_checkpoint (`bool`, *optional*, defaults to `False`): Whether to use gradient checkpointing to save memory during training. If set to `True`, the model will recompute certain activations during the backward pass instead of storing them in memory. """ if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) if self.use_mem_token and not ignore_mem_token: inputs_embeds = torch.concat((inputs_embeds, self.mem_tokens.unsqueeze(0).repeat(inputs_embeds.shape[0], 1, 1)), dim=-2) if attention_mask is not None: attention_mask = torch.concat([attention_mask, torch.ones_like(attention_mask[:, [0]]).repeat(1, self.num_mem_token)], dim=-1) if use_cache and past_key_values is None: past_key_values = DynamicCache(config=self.config) 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) # It may already have been prepared by e.g. `generate` if not isinstance(causal_mask_mapping := attention_mask, dict): # Prepare mask arguments mask_kwargs = { "config": self.config, "input_embeds": inputs_embeds, "attention_mask": attention_mask, "cache_position": cache_position, "past_key_values": past_key_values, "position_ids": position_ids, } # Create the masks causal_mask_mapping = { "full_attention": create_causal_mask(**mask_kwargs), } # The sliding window alternating layers are not always activated depending on the config if self.has_sliding_layers: causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs) hidden_states = inputs_embeds # create position embeddings to be shared across the decoder layers position_embeddings = self.rotary_emb(hidden_states, position_ids) if self.use_mem_token and not ignore_mem_token: memory_states = torch.zeros((hidden_states.shape[0], self.config.num_hidden_layers, self.num_mem_token, self.config.hidden_size), dtype=hidden_states.dtype).to(self.device) for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): if use_gradient_checkpoint: hidden_states = checkpoint( decoder_layer, hidden_states, attention_mask=causal_mask_mapping[decoder_layer.attention_type], position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, loradict=loradict[i] if isinstance(loradict, dict) else None, **kwargs, use_reentrant=False) else: hidden_states = decoder_layer( hidden_states, attention_mask=causal_mask_mapping[decoder_layer.attention_type], position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, loradict=loradict[i] if isinstance(loradict, dict) else None, **kwargs, ) if self.use_mem_token and not ignore_mem_token: memory_states[:, i, :, :] = hidden_states[:, -self.num_mem_token:].to(self.device) # for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): # hidden_states = decoder_layer( # hidden_states, # attention_mask=causal_mask_mapping[decoder_layer.attention_type], # position_ids=position_ids, # past_key_values=past_key_values, # use_cache=use_cache, # cache_position=cache_position, # position_embeddings=position_embeddings, # loradict=loradict[i] if isinstance(loradict, dict) else None, # **kwargs, # ) # if self.use_mem_token and not ignore_mem_token: # memory_states[:, i, :, :] = hidden_states[:, -self.num_mem_token:].to(self.device) if self.use_mem_token and not ignore_mem_token: hidden_states = hidden_states[:, :-self.num_mem_token, :] hidden_states = self.norm(hidden_states) return MemoryModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None, memory_states=memory_states if (self.use_mem_token and not ignore_mem_token) else None, ) def lora_params_numel(self, r): if not hasattr(self, "lora_params_numel_cache"): self.lora_params_numel_cache = 0 for layer in self.layers: self.lora_params_numel_cache += layer.lora_params_numel(r) return self.lora_params_numel_cache def set_generate_func(self, method): for layer in self.layers: layer.set_generate_func(method) def generate_lora_dict(self, r, scale, plain_tensor): assert plain_tensor.shape[-1] == self.lora_params_numel(r), f"plain_tensor's last dimension {plain_tensor.shape[-1]} does not match lora_params_numel {self.lora_params_numel(r)}" idx = 0 loradict = {} for i, layer in enumerate(self.layers): layer_lora_params_numel = layer.lora_params_numel(r) loradict[i] = layer.generate_lora_dict(r, scale, plain_tensor[:, idx: idx + layer_lora_params_numel]) idx += layer_lora_params_numel return loradict def init_lora_dict(self, r, scale, device): loradict = {} for i, layer in enumerate(self.layers): loradict[i] = layer.init_lora_dict(r, scale, device) return loradict def divide_idx(self, r, idx_start): return self.layers[0].divide_idx(r, idx_start) @dataclass class MemoryCausalLMOutputWithPast(CausalLMOutputWithPast): ''' memory_states: Optional[torch.Tensor] = None # (batch_size, num_layer, num_mem_token, hidden_size) reg_loss: Optional[torch.FloatTensor] = None ''' memory_states: Optional[torch.Tensor] = None reg_loss: Optional[torch.FloatTensor] = None @auto_docstring class LoraQwen3ForCausalLM(Qwen3PreTrainedModel, GenerationMixin): _tied_weights_keys = ["lm_head.weight"] _tp_plan = {"lm_head": "colwise_rep"} _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} def __init__(self, config): super().__init__(config) self.model = LoraQwen3Model(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) # Initialize weights and apply final processing self.post_init() def reset_mem_tokens(self): self.model.reset_mem_tokens() @can_return_tuple @auto_docstring def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, loradict: Optional[dict] = None, ignore_mem_token: bool = False, use_gradient_checkpoint: bool = False, **kwargs: Unpack[TransformersKwargs], ) -> CausalLMOutputWithPast: r""" 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]`. loradict (`dict` of `dict` of `torch.FloatTensor`, *optional*): A dictionary that maps each layer to its corresponding LoRA parameters. Each layer's LoRA parameters are stored in a nested dictionary. ignore_mem_token (`bool`, *optional*, defaults to `False`): Whether to ignore the memory tokens during the forward pass. If set to `True`, the memory tokens will not be used, and the model will behave like a standard transformer without memory tokens. use_gradient_checkpoint (`bool`, *optional*, defaults to `False`): Whether to use gradient checkpointing to save memory during training. If set to `True`, the model will recompute certain activations during the backward pass instead of storing them in memory. Example: ```python >>> from transformers import AutoTokenizer, Qwen3ForCausalLM >>> model = Qwen3ForCausalLM.from_pretrained("Qwen/Qwen3-8B") >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B") >>> 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." ```""" # if loradict is not None: # print("use lora dict !!!!!!!!!!!!!!!!!!!!!!!!!!!!") # if getattr(self, "cache_lora_dict", None) is not None: # # check if the same # for k in loradict: # for subk in loradict[k]: # for param in loradict[k][subk]: # for t in loradict[k][subk][param]: # if loradict[k][subk][param][t] is not None and not torch.equal(loradict[k][subk][param][t], self.cache_lora_dict[k][subk][param][t]): # print(f"different lora dicts {k}.{subk}.{param}.{t} !!!!!!!!!!!!!!!!!!!!!!!!!!!!") # self.cache_lora_dict = loradict outputs: MemoryModelOutputWithPast = 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, cache_position=cache_position, loradict=loradict, ignore_mem_token=ignore_mem_token, use_gradient_checkpoint=use_gradient_checkpoint, **kwargs, ) hidden_states = outputs.last_hidden_state # Only compute necessary logits, and do not upcast them to float if we are not computing the loss slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(hidden_states[:, slice_indices, :]) # # ---- START DEBUG ---- # tok = kwargs.get("debug_tokenizer", None) # if tok is not None: # # Choose tokenizer from either kwarg # has_tok = tok is not None and hasattr(tok, "convert_ids_to_tokens") # with torch.no_grad(): # # Align labels to the logits slice # if labels is not None: # labels_slice = labels[:, slice_indices] # (B, T') # B, T, V = logits.shape # # Temperature = 1 (no scaling) # probs = torch.softmax(logits, dim=-1) # (B, T, V) # # Compute top-k over all time steps # topk_vals, topk_idx = torch.topk(logits, k=5, dim=-1) # (B, T, 5) # for b in range(B): # print(f"[DEBUG] Batch {b}:") # for t in range(T): # if labels is not None: # try: # lab = int(labels_slice[b, t + 1].item()) # except: # lab = -100 # # if lab == -100: # # continue # skip unlabeled positions # # Top-k for this (b, t) # ids_bt = topk_idx[b, t].tolist() # vals_bt = topk_vals[b, t].tolist() # probs_bt = probs[b, t, ids_bt].tolist() # toks_bt = tok.convert_ids_to_tokens(ids_bt) if has_tok else None # print(f" [t={t}]") # for r, (vid, vlogit, vprob) in enumerate(zip(ids_bt, vals_bt, probs_bt), start=1): # if has_tok: # print( # f" Top{r}: token_id={vid:<8} token={repr(toks_bt[r-1])} " # f"logit={vlogit:.4f} prob={vprob:.6f}" # ) # else: # print( # f" Top{r}: token_id={vid:<8} logit={vlogit:.4f} prob={vprob:.6f}" # ) # if labels is not None and lab != -100: # # Label token info for this (b, t) # lab_logit = logits[b, t, lab].item() # lab_prob = probs[b, t, lab].item() # if has_tok: # lab_tok = tok.convert_ids_to_tokens([lab])[0] # print( # f" Label: token_id={lab:<8} token={repr(lab_tok)} " # f"logit={lab_logit:.4f} prob={lab_prob:.6f}" # ) # else: # print( # f" Label: token_id={lab:<8} logit={lab_logit:.4f} prob={lab_prob:.6f}" # ) # # ---- END DEBUG ---- loss = None if labels is not None: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) return MemoryCausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, memory_states=outputs.memory_states if not ignore_mem_token else None, ) def lora_params_numel(self, r): return self.model.lora_params_numel(r) def set_generate_func(self, method): self.model.set_generate_func(method) def generate_lora_dict(self, r, scale, plain_tensor): return self.model.generate_lora_dict(r, scale, plain_tensor) def init_lora_dict(self, r, scale, device): return self.model.init_lora_dict(r, scale, device) def divide_idx(self, r, idx_start): return self.model.divide_idx(r, idx_start) if __name__ == "__main__": config_path = "./models/Qwen3-0.6B"