import math import os import json from pathlib import Path from os.path import split from transformers import GenerationMixin, GenerationConfig, Cache, PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithPast import copy from typing import Dict, List, Optional, Literal,Union,Tuple,Any from types import SimpleNamespace from .configuration_tinyllm import Config try: from pytorch_tcn import TCN, TemporalConv1d HAS_PYTORCH_TCN = True except Exception: TCN = None TemporalConv1d = None HAS_PYTORCH_TCN = False from datetime import datetime import torch import torch.nn as nn import torch.nn.functional as F import warnings from torch.utils.checkpoint import checkpoint from torch.nn import AdaptiveLogSoftmaxWithLoss import math import torch try: from transformers.models.mamba2 import Mamba2Config, Mamba2Model from transformers.models.mamba2.modeling_mamba2 import Mamba2Cache HAS_MAMBA2 = True except Exception as e: Mamba2Config = None Mamba2Model = None Mamba2Cache = None HAS_MAMBA2 = False warnings.warn( f"[TinyLLM] transformers.models.mamba2 未能导入（{e}）。" "SSM/Mamba2 将被禁用；请确保 cfg.use_ssm=False 或安装带 mamba2 支持的 transformers。", RuntimeWarning, ) from transformers.models.llama.modeling_llama import apply_rotary_pos_emb from transformers.modeling_rope_utils import rope_config_validation, ROPE_INIT_FUNCTIONS def gumbel_noise(x:torch.Tensor): #-log(-log(U)) u=torch.rand_like(x).clamp(1e-9,1-1e-9) return -torch.log(-torch.log(u)) def _build_mlp_ratio_schedule(cfg: Config) -> list[float]: L = int(cfg.num_hidden_layers) # 若未提供分界，则退回统一 mlp_ratio if cfg.mlp_mid_start is None and cfg.mlp_back_start is None: return [float(cfg.mlp_ratio)] * L mid = cfg.mlp_mid_start if cfg.mlp_mid_start is not None else L back = cfg.mlp_back_start if cfg.mlp_back_start is not None else L mid = max(0, min(int(mid), L)) back = max(mid, min(int(back), L)) ratios = ( [float(cfg.mlp_ratio_front)] * mid + [float(cfg.mlp_ratio_mid)] * (back - mid) + [float(cfg.mlp_ratio_back)] * (L - back) ) ov = getattr(cfg, "mlp_ratio_overrides", None) or {} if isinstance(ov, dict): for k, v in ov.items(): try: i = int(k) if 0 <= i < L: ratios[i] = float(v) except Exception: pass return ratios class GradScale(torch.autograd.Function): @staticmethod def forward(ctx, x, scale: float): ctx.scale = scale return x @staticmethod def backward(ctx, g): return g * ctx.scale, None class RMSNorm(nn.Module): def __init__(self,dim:int,eps:float= 1e-5,affine:bool=True,dtype=None,device=None): super().__init__() self.dim = dim self.eps=eps self.affine=affine self.dtype=dtype self.device=device if self.affine: self.weight=nn.Parameter(torch.ones(dim,dtype=dtype,device=device)) else: self.register_parameter("weight",None) def forward(self,x:torch.Tensor): x_float=x.float() inv_rms=x_float.pow(2).mean(dim=-1,keepdim=True).add(self.eps).rsqrt() inv_rms=inv_rms.to(dtype=x.dtype) y=inv_rms*x if self.affine: y=y*self.weight return y class QK_RMSNorm(nn.Module): def __init__(self,head_dim:int,num_head:int,eps:float=1e-5,learnable_temp=True): super().__init__() self.head_dim=head_dim self.num_head=num_head self.eps=eps self.q_norm=RMSNorm(head_dim,eps) self.k_norm=RMSNorm(head_dim,eps) def forward(self,q:torch.Tensor,k:torch.Tensor): q=self.q_norm(q) k=self.k_norm(k) return q,k class RoPE(nn.Module): def __init__(self,head_dim:int,max_position:int,base:float=2e4,use_NTK:bool=False,train_length:int=2048): super().__init__() self.head_dim=head_dim self.max_position=max_position self.base=base assert head_dim%2==0,"RoPE head_dim非偶数" self.register_buffer("_dev", torch.empty(0), persistent=False) index=-torch.arange(0,head_dim,2,dtype=torch.float32,device=self._dev.device)/head_dim self.use_NTK=use_NTK self.train_length=train_length self.last_key=None #only for NTK self.last_cos=None self.last_sin=None self.target_alpha=float(max_position)/float(train_length) if use_NTK: self.register_buffer("cos_RoPE",None,persistent=False) self.register_buffer("sin_RoPE",None,persistent=False) self.register_buffer("RoPE_index",index,persistent=False)#shape[d/2] else: index=base**index #shape[D/2] m=torch.arange(max_position,dtype=torch.float32)#shape[max_position] freq=torch.einsum("m,i->mi",m,index)#[max_position,D/2] tempcos=torch.cos(freq).unsqueeze(0).unsqueeze(0) tempsin=torch.sin(freq).unsqueeze(0).unsqueeze(0) self.register_buffer("RoPE_index",None,persistent=False) self.register_buffer("cos_RoPE",tempcos,persistent=False) self.register_buffer("sin_RoPE",tempsin,persistent=False) @torch.no_grad() def getCosSin(self,start_pos:int,T:int): if self.use_NTK: key=(start_pos,T,self._dev.device) if self.last_key==key and self.last_cos is not None: return self.last_cos,self.last_sin p_index=torch.arange(start_pos,start_pos+T,dtype=torch.float32,device=self._dev.device)#shape [T] len_diff=max(self.max_position-self.train_length,1) p=torch.clamp((p_index-self.train_length)/(len_diff),0.0,max=1.0) theta_alpha=1+(self.target_alpha-1)*p#shape [T] base=self.base*(theta_alpha**(self.head_dim/(self.head_dim-2))) #shape [T] index=(base.unsqueeze(1)**self.RoPE_index.unsqueeze(0)) freq=p_index.unsqueeze(1)*index cos=torch.cos(freq).unsqueeze(0).unsqueeze(0) sin=torch.sin(freq).unsqueeze(0).unsqueeze(0) self.last_sin=sin self.last_cos=cos self.last_key=key return cos,sin else: cos_temp=self.cos_RoPE[...,start_pos:start_pos+T,:] sin_temp=self.sin_RoPE[...,start_pos:start_pos+T,:] return cos_temp,sin_temp def forward(self,x:torch.Tensor,start_pos:int=0): B,H,T,D=x.shape x_even=x[...,::2] x_odd=x[...,1::2] x_RoPE=torch.empty_like(x) cos,sin=self.getCosSin(start_pos,T) cos=cos.to(x.dtype) sin=sin.to(x.dtype) x_RoPE[...,::2]=x_even*cos-x_odd*sin x_RoPE[...,1::2]=x_even*sin+x_odd*cos return x_RoPE class HF_RoPEBackend(nn.Module): def __init__(self, hidden_size:int,num_heads:int,head_dim:int, max_position:int, base:float, rope_type:str, train_length:int,device:torch.device): super().__init__() assert head_dim % 2 == 0, "head_dim must be even for RoPE" self.head_dim = head_dim self.base = float(base) self.max_position = int(max_position) self.train_length = int(max(1, train_length)) factor = max(1.0, float(self.max_position) / float(self.train_length)) cfg = SimpleNamespace() cfg.max_position_embeddings = self.max_position cfg.rope_theta = self.base cfg.head_dim=head_dim cfg.rope_scaling = {"rope_type": rope_type, "factor": factor,"original_max_position_embeddings": train_length} if rope_type != "default" else None cfg.hidden_size = hidden_size cfg.num_attention_heads=num_heads rope_config_validation(cfg) init_fn = ROPE_INIT_FUNCTIONS[rope_type] rope_state = init_fn(cfg,device) if isinstance(rope_state, tuple): inv_freq, attn_scale = rope_state else: inv_freq = getattr(rope_state, "inv_freq") attn_scale = getattr(rope_state, "attention_factor", 1.0) self.register_buffer("inv_freq", inv_freq.to(torch.float32), persistent=False) self.attn_scale = float(attn_scale) if attn_scale is not None else None @torch.no_grad() def build_cos_sin(self, position_ids: torch.Tensor, dtype): device = self.inv_freq.device pos = position_ids.to(device=device, dtype=torch.float32) # [B,T] inv = self.inv_freq.to(dtype=torch.float32) freqs = torch.einsum("bt,d->btd", pos, inv) emb = torch.cat([freqs, freqs], dim=-1) # [B,T,d] cos = torch.cos(emb).to(dtype) sin = torch.sin(emb).to(dtype) if self.attn_scale is not None: cos = cos * float(self.attn_scale) sin = sin * float(self.attn_scale) return cos, sin class TCNBranchPT(nn.Module): """ 并联的 TCN 分支（pytorch-tcn），保持 [B,T,H] 形状。 """ def __init__( self, hidden_size: int, num_blocks: int = 4, expansion: float = 0.175, kernel_size: int = 9, dilations = [1, 2, 4,8], dropout: float = 0.05, use_norm: str = "weight_norm", ): super().__init__() if not HAS_PYTORCH_TCN: raise ImportError("use_tcn=True requires the optional 'pytorch-tcn' package") out_c = max(1, int(round(hidden_size * float(expansion)))) channels: List[int] = [out_c for _ in range(int(num_blocks))] self.tcn = TCN( num_inputs=hidden_size, num_channels=channels, kernel_size=kernel_size, dilations=dilations, dilation_reset=None, dropout=dropout, causal=True, use_norm=use_norm, activation="relu", use_skip_connections=False, input_shape="NLC", # [B,T,H] 输入 output_projection=hidden_size, output_activation=None, ) self.ln_in = nn.LayerNorm(hidden_size) @torch.no_grad() def new_state(self) -> list[torch.Tensor]: """ 返回一个“干净”的初始状态快照（全零/初始 buffer）。 """ self.tcn.reset_buffers() return [b.clone() if b is not None else None for b in self.tcn.get_buffers()] @torch.no_grad() def clone_state(self, state: list[torch.Tensor]) -> list[torch.Tensor]: return [ (t.clone() if t is not None else None) for t in state ] @torch.no_grad() def set_state(self, state: list[torch.Tensor]): """ 把外部 state 写回到 TCN 模块内部。 """ buf = [ (t.clone() if t is not None else None) for t in state ] self.tcn.set_buffers(buf) @torch.no_grad() def get_state(self) -> list[torch.Tensor]: """ 读取当前 TCN 模块内部状态（深拷贝）。 """ return [b.clone() if b is not None else None for b in self.tcn.get_buffers()] @torch.no_grad() def reset_state(self): self.tcn.reset_buffers() def forward(self, x: torch.Tensor, attention_mask: torch.Tensor | None = None, inference: bool = False, state_in: list[torch.Tensor] | None = None,return_state: bool = False,): # x: [B,T,H], mask: [B,T] (1=valid) z = self.ln_in(x) if attention_mask is not None: z = z * attention_mask.to(z.dtype).unsqueeze(-1) if not inference: y = self.tcn(z, inference=False) if attention_mask is not None: y = y * attention_mask.to(y.dtype).unsqueeze(-1) if return_state: return y, self.get_state() return y if state_in is not None: self.set_state(state_in) y = self.tcn(z, inference=True) # 走库的 streaming 逻辑 if attention_mask is not None: y = y * attention_mask.to(y.dtype).unsqueeze(-1) if return_state: return y, self.get_state() return y @staticmethod @torch.no_grad() def tcn_step_batch( tcn_branch: "TCNBranchPT", x_btH: torch.Tensor, mask_bt: torch.Tensor | None, # [B,T] or None states: list[list[torch.Tensor] | None], # len=B，每样本一个 state_list ) -> tuple[torch.Tensor, list[list[torch.Tensor]]]: B, T, H = x_btH.shape ys, new_states = [], [] for b in range(B): x_b = x_btH[b:b+1] m_b = None if mask_bt is None else mask_bt[b:b+1] st_in = states[b] y_b, st_out = tcn_branch( x_b, m_b, inference=True, state_in=st_in, return_state=True ) ys.append(y_b) new_states.append(st_out) y_btH = torch.cat(ys, dim=0) return y_btH, new_states class Mutihead_attention(nn.Module): def __init__(self, hidden_size:int, num_heads:int, num_kv_heads:int|None=None, max_position_embeddings:int=16384, RoPE_base:float=2e4, dropout:float=0.0, use_qk_RMSnorm:bool=True, qk_RMSeps:float=1e-5, learnable_temp:bool=True, use_sample_attention:bool=False, rope_type:str="yarn", train_length:int=4096, use_HF_RoPE:bool=True, ): super().__init__() self.hidden_size=hidden_size assert hidden_size%num_heads==0,"hidden size doesn't match n*num_head" self.num_heads=num_heads if not num_kv_heads: self.num_kv_heads=num_heads else: self.num_kv_heads=num_kv_heads self.dropout=dropout self.use_qk_RMSnorm=use_qk_RMSnorm self.qk_RMSeps=qk_RMSeps self.learnable_temp=learnable_temp self.dim_perhead=hidden_size//num_heads self.w_q=nn.Linear(hidden_size,num_heads*self.dim_perhead,bias=False) self.w_k=nn.Linear(hidden_size,self.num_kv_heads*self.dim_perhead,bias=False) self.w_v=nn.Linear(hidden_size,self.num_kv_heads*self.dim_perhead,bias=False) self.w_o=nn.Linear(num_heads*self.dim_perhead,hidden_size,bias=False) if use_HF_RoPE: self.register_buffer("_rope_dev_probe", torch.empty(0), persistent=False) self.rope_backend = HF_RoPEBackend( hidden_size=hidden_size, num_heads=num_heads, head_dim=self.dim_perhead, max_position=max_position_embeddings, base=RoPE_base, rope_type=rope_type, train_length=train_length, device=self._rope_dev_probe.device ) self.RoPE=None else: self.RoPE = RoPE(self.dim_perhead, max_position=max_position_embeddings, base=RoPE_base, use_NTK=True,train_length=train_length) self.rope_backend = None if self.use_qk_RMSnorm: self.qk_RMSnorm=QK_RMSNorm(self.dim_perhead,self.num_heads,self.qk_RMSeps) else: self.qk_RMSnorm=None self.use_sample_attention=use_sample_attention if self.learnable_temp: self.temp_param = nn.Parameter(torch.zeros(num_heads)) self.temp_range = 0.15 else: self.register_parameter("temp_param", None) @staticmethod def _reshape(x:torch.Tensor,num_head:int,head_dim:int): B,T,_=x.shape return x.view(B,T,num_head,head_dim).transpose(1,2).contiguous() def sample_attention(self,q,k,v,att_mask=None,is_causal=True,dropout=0.0): B,H,T,D=q.shape standard=1/math.sqrt(D) score=torch.matmul(q,k.transpose(-2,-1))*standard if is_causal: i=torch.arange(T,device=q.device) causal=(i[:,None]>=i[None,:]) score=score.masked_fill(~causal,float("-inf")) if att_mask is not None: score=score+att_mask score=score-score.max(dim=-1,keepdim=True).values att=torch.softmax(score,dim=-1) att=F.dropout(att,p=dropout,training=self.training) y=torch.matmul(att,v) return y def forward( self, x: torch.Tensor, attention_mask: torch.Tensor | None = None, is_causal: bool = True, past_state: Optional[dict] = None, use_cache: bool = False, img_len: int | None = None, ): """ x: [B,T,C] attention_mask: [B,T]，1=有效,0=pad；单样本/无pad时可为None past_state: 本层上一次缓存的状态字典，比如 { "attn": { "k": ... [B, Hkv, Lpast, D], "v": ... [B, Hkv, Lpast, D], "len": ... [B], # 每个样本已经累积的“真实token个数” "valid_mask": ... [B, Lpast] # True=这个位置是有效token, False=padding } } """ B, T, C = x.shape # 1) 线性投影拿到 q,k,v q = self.w_q(x) k = self.w_k(x) v = self.w_v(x) q = self._reshape(q, self.num_heads, self.dim_perhead) # [B, Hq, T, D] k = self._reshape(k, self.num_kv_heads, self.dim_perhead) # [B, Hkv, T, D] v = self._reshape(v, self.num_kv_heads, self.dim_perhead) # [B, Hkv, T, D] if self.use_qk_RMSnorm: q, k = self.qk_RMSnorm(q, k) # 2) 取出过去的 cache past_kv = None past_len = None # [B]，每条样本已经累积的真实长度 past_valid_mask = None # [B, Lpast]，True表示这个位置是有效token if past_state is not None and "attn" in past_state: kv = past_state["attn"] past_kv = (kv["k"], kv["v"]) past_len = kv.get("len", None) past_valid_mask = kv.get("valid_mask", None) # 3) 计算本 step 内的“局部位置”（忽略左填充） # 和本 step 的有效位 mask（哪些是真token，而不是pad） if attention_mask is not None: am_long = attention_mask.to(torch.long) # [B,T] local_pos = (am_long.cumsum(dim=-1) - 1).clamp_min(0) # [B,T], pad行会复用0 chunk_valid_mask = am_long.bool() # [B,T] else: # 无mask就说明整段都是实 token local_pos = torch.arange(T, device=x.device).view(1, T).expand(B, T) chunk_valid_mask = torch.ones((B, T), dtype=torch.bool, device=x.device) # 4) 计算 RoPE 的全局绝对位置：每个样本自己的offset if past_len is not None: # past_len: [B] => [B,1] broadcast offset = past_len.to(device=x.device, dtype=torch.long).view(B, 1) pos = local_pos + offset # [B,T] else: pos = local_pos # [B,T]，第一段/prefill # 5) 应用 RoPE if self.rope_backend is not None: cos, sin = self.rope_backend.build_cos_sin(pos, dtype=q.dtype) q, k = apply_rotary_pos_emb(q, k, cos, sin) else: # 老的自实现 RoPE 只支持统一的 start_pos（标量）。 # 单batch还好，multi-batch其实不严谨。 start_pos_scalar = 0 if past_len is not None: # 只能拿第0个样本的offset做近似；多batch会不准，所以尽量别走这个分支。 start_pos_scalar = int(past_len[0].item()) q = self.RoPE(q, start_pos=start_pos_scalar) k = self.RoPE(k, start_pos=start_pos_scalar) if self.learnable_temp and (self.temp_param is not None): s = 1.0 + self.temp_range * torch.tanh(self.temp_param).to(q.dtype) q = q * s.view(1, self.num_heads, 1, 1) # 6) 把新的 k,v 接到过去的 k,v 后面；同时合并 valid_mask TotalT = T if past_kv is not None: past_k, past_v = past_kv k = torch.cat([past_k, k], dim=-2) v = torch.cat([past_v, v], dim=-2) TotalT = k.size(-2) if past_valid_mask is not None: full_valid_mask = torch.cat([past_valid_mask, chunk_valid_mask], dim=-1) # [B, TotalT] else: full_valid_mask = chunk_valid_mask # [B,T] in first chunk # k_kv/v_kv 用于cache保存（head还是 Hkv） k_kv, v_kv = k, v # 7) 如果是 GQA，多头扩展到 Hq if self.num_kv_heads != self.num_heads: repeat = self.num_heads // self.num_kv_heads k = ( k.unsqueeze(2) .expand(B, self.num_kv_heads, repeat, TotalT, self.dim_perhead) .reshape(B, self.num_heads, TotalT, self.dim_perhead) ) v = ( v.unsqueeze(2) .expand(B, self.num_kv_heads, repeat, TotalT, self.dim_perhead) .reshape(B, self.num_heads, TotalT, self.dim_perhead) ) # 8) 构造注意力mask，保证： # - 不能看未来（causal） # - 永远不能看 padding（full_valid_mask==False 的地方） attn_mask_for_sdpa = None use_causal_flag = is_causal need_custom_mask = (past_kv is not None) or (attention_mask is not None) if need_custom_mask: # key_time: 每个 key 位置是“第几个有效 token”（0,1,2,...） key_time = full_valid_mask.long().cumsum(dim=-1) - 1 # [B, TotalT] big = torch.iinfo(torch.int64).max // 2 key_time = key_time.masked_fill(~full_valid_mask, big) # pad 给一个超大时间戳 # query_time: 当前 chunk 的“逻辑时间”，直接用上面算好的 pos # pos: [B, T]，已经是 local_pos (+ past_len) query_time = pos.to(dtype=torch.int64) # [B, T] Bq, Tq = query_time.shape assert Bq == B and Tq == T qt = query_time.view(B, 1, T, 1) # [B,1,T,1] kt = key_time.view(B, 1, 1, TotalT) # [B,1,1,TotalT] # 不能看： # 1) padding（full_valid_mask==False） # 2) 逻辑时间晚于自己的 token（strict future: key_time > query_time） blocked_bool = (kt > qt) | (~full_valid_mask.view(B, 1, 1, TotalT)) if (img_len is not None) and (img_len > 0) and (past_kv is None) and is_causal: q_idx = torch.arange(T, device=x.device) k_idx = torch.arange(TotalT, device=x.device) q_is_prefix = (q_idx < img_len).view(1, 1, T, 1) k_is_prefix = (k_idx < img_len).view(1, 1, 1, TotalT) blocked_bool = blocked_bool & ~(q_is_prefix & k_is_prefix) # prefix<->prefix 双向可见 blocked_bool = blocked_bool | (q_is_prefix & ~k_is_prefix) # prefix->text 禁止偷看 attn_mask_for_sdpa = torch.zeros((B, 1, T, TotalT), dtype=q.dtype, device=x.device) attn_mask_for_sdpa = attn_mask_for_sdpa.masked_fill( blocked_bool, torch.finfo(q.dtype).min ) use_causal_flag = False # 9) 真正做注意力 if self.use_sample_attention: # sample_attention 用的是加法 mask（0 / -inf） add_mask = attn_mask_for_sdpa y = self.sample_attention( q, k, v, add_mask, use_causal_flag and is_causal, self.dropout ) else: y = F.scaled_dot_product_attention( q, k, v, attn_mask=attn_mask_for_sdpa, dropout_p=(self.dropout if self.training else 0.0), is_causal=use_causal_flag, ) # [B,Hq,T,D] -> [B,T,H] y = y.transpose(1, 2).contiguous().view(B, T, self.hidden_size) # 10) 组装 present cache present = None if use_cache: # 更新每个样本的“真实token计数” cur_valid_lengths = chunk_valid_mask.sum(dim=-1).to(dtype=torch.long) # [B] if past_len is not None: new_len = past_len + cur_valid_lengths else: new_len = cur_valid_lengths present = { "attn": { "k": k_kv, # [B,Hkv,TotalT,D] "v": v_kv, # [B,Hkv,TotalT,D] "len": new_len, # [B] （真实token累计数） "valid_mask": full_valid_mask, # [B,TotalT] True=有效token } } return self.w_o(y), present class Drop_path(nn.Module): def __init__(self,drop_p:float=0.0): super().__init__() self.drop_p=drop_p def forward(self,x:torch.Tensor): if self.drop_p==0.0 or not self.training: return x keep=1-self.drop_p shape=(x.size(0),)+(1,)*(x.ndim-1) mask=torch.rand(shape,device=x.device) <=keep return x*(mask.to(dtype=x.dtype))/(keep) class MLP(nn.Module): def __init__(self,hidden_dim:int,dropout,amplify:float=4,use_swiGLU:bool=True): super().__init__() self.activate=nn.SiLU() self.dropout=nn.Dropout(dropout) self.use_swiGLU=use_swiGLU if use_swiGLU: self.amplify=int((amplify*2/3*hidden_dim+63)//64*64) self.upsamp=nn.Linear(hidden_dim,self.amplify,bias=False) self.swiGate=nn.Linear(hidden_dim,self.amplify, bias=False) self.downsamp=nn.Linear(self.amplify,hidden_dim, bias=False) else: self.amplify=int(amplify*hidden_dim) self.upsamp=nn.Linear(hidden_dim,self.amplify, bias=False) self.downsamp=nn.Linear(self.amplify,hidden_dim, bias=False) def forward(self,x:torch.Tensor): if self.use_swiGLU: return self.dropout(self.downsamp(self.activate(self.swiGate(x))*self.upsamp(x))) return self.dropout(self.downsamp(self.activate(self.upsamp(x)))) class VisionProjector(nn.Module): def __init__( self, vision_dim: int, hidden_dim: int, use_swiGLU: bool = True, dropout: float = 0.0, use_rmsnorm: bool = True, ): super().__init__() self.use_swiGLU = use_swiGLU self.act = nn.SiLU() self.drop = nn.Dropout(dropout) # 1) 1536 -> 1280 self.proj = nn.Linear(vision_dim, hidden_dim, bias=False) if use_rmsnorm: # 复用你上面的 RMSNorm self.norm = RMSNorm(hidden_dim, eps=1e-5, affine=True) else: self.norm = nn.LayerNorm(hidden_dim) def forward(self, feats: torch.Tensor): # feats: [B, Nv, Dv] x = self.proj(feats) x = self.norm(x) return x class QFormerBlock(nn.Module): def __init__( self, hidden_dim: int, num_heads: int, mlp_ratio: float = 2.0, dropout: float = 0.0, use_rms_norm: bool = False, ): super().__init__() Norm = RMSNorm if use_rms_norm else nn.LayerNorm # self-attn over queries self.ln_self = Norm(hidden_dim) self.self_attn = nn.MultiheadAttention( embed_dim=hidden_dim, num_heads=num_heads, dropout=dropout, batch_first=True, ) # cross-attn: q = query tokens, kv = vision tokens self.ln_cross_q = Norm(hidden_dim) self.ln_cross_kv = Norm(hidden_dim) self.cross_attn = nn.MultiheadAttention( embed_dim=hidden_dim, num_heads=num_heads, dropout=dropout, batch_first=True, ) # MLP self.ln_mlp = Norm(hidden_dim) self.mlp = MLP(hidden_dim=hidden_dim, amplify=mlp_ratio, dropout=dropout) self.drop = nn.Dropout(dropout) def forward( self, q: torch.Tensor, # [B, Nq, H] learnable queries 或 text-conditioned queries kv: torch.Tensor, # [B, Nv, H] vision tokens（已经 1536→1280 过了） kv_mask: torch.Tensor | None = None, # [B, Nv]，1=valid, 0=pad ) -> torch.Tensor: # 1) Self-Attention over q（完全非因果，双向） q_norm = self.ln_self(q) q_sa, _ = self.self_attn( q_norm, q_norm, q_norm, need_weights=False, # 有利于用 fused kernel ) q = q + self.drop(q_sa) # 2) Cross-Attention：q ← kv q_norm = self.ln_cross_q(q) kv_norm = self.ln_cross_kv(kv) if kv_mask is not None: # key_padding_mask: True = 要 mask 掉 key_padding_mask = (~kv_mask.bool()) # [B, Nv] else: key_padding_mask = None q_ca, _ = self.cross_attn( q_norm, kv_norm, kv_norm, key_padding_mask=key_padding_mask, need_weights=False, ) q = q + self.drop(q_ca) # 3) 小 MLP q_mlp = self.mlp(self.ln_mlp(q)) q = q + q_mlp return q # [B, Nq, H] class Top1Router(nn.Module): def __init__(self,hidden_size:int,num_expert:int,enable_detach:bool=False): super().__init__() self.decider=nn.Linear(hidden_size,num_expert,bias=True) self.num_expert=num_expert self.enableDetach = enable_detach def forward(self,x:torch.Tensor): x=x.detach() logit=self.decider(x) logit = logit - logit.mean(dim=-1, keepdim=True) std = logit.std(dim=-1, keepdim=True, unbiased=False).clamp_min(1e-3) logit=logit / std/1.2 prob=F.softmax(logit,dim=-1) if self.training: top_idx=(logit+0.4*gumbel_noise(logit)).topk(2, dim=-1).indices else: top_idx = prob.topk(2, dim=-1).indices top_p = prob.gather(1, top_idx) return top_idx,top_p,prob class MoEMLP(nn.Module): def __init__(self,hidden_dim:int,dropout,num_expert:int,amplify:float=4,use_swiGLU:bool=True,cap_factor:float=1.25,enable_detach:bool=False,n_shared:int=1): super().__init__() self.num_expert=num_expert self.expert=nn.ModuleList([MLP(hidden_dim,dropout,amplify,use_swiGLU) for i in range(num_expert)]) self.router=Top1Router(hidden_dim,num_expert,enable_detach) self.cap_factor=cap_factor self.enableDetach=enable_detach self.register_buffer("top2_alpha", torch.tensor(0.0)) self.n_shared = n_shared self.a, self.b = 0.6, 1.4 if n_shared>0: self.shared = nn.ModuleList([MLP(hidden_dim, dropout, amplify, use_swiGLU) for _ in range(n_shared)]) self.shared_gate = nn.Linear(hidden_dim, 1) nn.init.zeros_(self.shared_gate.weight) nn.init.zeros_(self.shared_gate.bias) else: self.shared = None self.shared_gamma = None def forward(self,x:torch.Tensor): B,T,H=x.shape N=B*T x=x.view(N,H) device=x.device e,prob_selected,prob=self.router(x) e1, e2 = e[:,0], e[:,1] w1, w2 = prob_selected[:,0:1], prob_selected[:,1:2] e1_orderd,e1_idx=e1.sort() sorted1_x=x[e1_idx] sorted1_p=w1[e1_idx] count1=torch.bincount(e1_orderd,minlength=self.num_expert) # [E] if self.training: capacity=max(1,math.ceil(self.cap_factor*N/self.num_expert)) else: capacity=int(1e12) keep_per_exp1 = torch.minimum(count1, torch.full_like(count1, capacity)) prefix=torch.cumsum(count1,dim=0) # [E] start=prefix-count1 end=prefix realend=start+keep_per_exp1 prob_32=prob.to(torch.float32) imp=prob_32.sum(dim=0) imp_hat=imp/imp.sum().clamp_min(1e-6) load=keep_per_exp1.to(torch.float32) load_hat=load/load.sum().clamp_min(1e-6) log_imp=(imp_hat.clamp_min(1e-9)).log() log_load=(load_hat.clamp_min(1e-9)).log() log_u=-math.log(self.num_expert) L_imp=(imp_hat*(log_imp-log_u)).sum() L_load=(load_hat*(log_load-log_u)).sum() aux_loss=L_imp+L_load aux_loss= aux_loss/ (2.0 * (-log_u) + 1e-12) sorted1_y=torch.zeros_like(sorted1_x) for j in range(self.num_expert): sj=start[j] rj=realend[j] ej=end[j] if ej==sj: continue if rj>sj: x_kept=sorted1_x[sj:rj] y_kept=self.expert[j](x_kept) p_kept=sorted1_p[sj:rj] if self.enableDetach: p_kept=p_kept.detach() else: # 只放行 router_grad_frac 的梯度（直通估计 + 缩放） p_kept = p_kept.detach() + GradScale.apply(p_kept - p_kept.detach(), 0.1) y_kept=y_kept*p_kept sorted1_y[sj:rj]=y_kept if ej>rj: pass inv1_idx=torch.argsort(e1_idx) y=sorted1_y[inv1_idx].reshape(B,T,H) if self.shared is not None: shared_y = 0 x_btH = x.reshape(B, T, H) for se in self.shared: shared_y = shared_y + se(x_btH) shared_y = shared_y / self.n_shared gamma_tok = torch.sigmoid(self.shared_gate(x_btH)) gamma_tok = self.a + (self.b - self.a) * gamma_tok y = y + gamma_tok * shared_y if self.training: return y,aux_loss return y class LoraAdapter(nn.Module): def __init__(self,input_dim:int,out_dim:int,rank:int=8,dropout:float=0.1,alpha:float=16): super().__init__() self.input_dim=input_dim self.output_dim=out_dim self.r=rank self.alpha=alpha self.amplify=alpha/rank self.A=nn.Linear(input_dim,rank,bias=False) self.B=nn.Linear(rank,out_dim,bias=False) self.dropout=nn.Dropout(dropout) nn.init.zeros_(self.B.weight) def forward(self,x:torch.Tensor): x=self.A(x) x=self.dropout(x) x=self.B(x)*self.amplify return x class LoraLinear(nn.Module): def __init__(self,base_linear:nn.Linear,mode:Literal["exclusive", "additive", "weighted"]="exclusive",cap_norm:float|None = None,global_scale:float=1.0): super().__init__() assert isinstance(base_linear,nn.Linear),"LoraLinear type wrong" self.base=base_linear self.input_dim=base_linear.in_features self.out_dim=base_linear.out_features self.adapters=nn.ModuleDict() self.active=[] self.mode=mode self.weights={} self.cap_norm=cap_norm self.global_scale=global_scale self.dtype=base_linear.weight.dtype self.device=base_linear.weight.device def set_mode(self,mode:str): self.mode=mode def set_cap_norm(self,cap): self.cap_norm=cap def set_global_scale(self,s): self.global_scale=s def register_adapter(self,name:str,rank:int=8,dropout:float=0.1,alpha:float=16.0,state_dict:dict|None=None,weight:float=1.0): if name in self.adapters: raise ValueError(f"LoRA adapter {name} already registered") adapter=LoraAdapter(self.input_dim,self.out_dim,rank,dropout,alpha).to(self.dtype).to(self.device) self.adapters[name]=adapter if state_dict is not None: miss,unexpect=adapter.load_state_dict(state_dict,strict=False) if miss: print(f"Lora {name}: missing key {miss} ") if unexpect: print(f"Lora {name}: unexpect key {unexpect} ") if name not in self.weights : self.weights[name]=weight def activate(self,name:str,exclusive:bool=False): if name not in self.adapters: raise KeyError(f"LoRA adapter {name} not registered.") if exclusive or self.mode=="exclusive": self.active=[name] else: self.active.append(name) def deactivate(self,name:str): if name not in self.adapters: raise KeyError(f"LoRA adapter {name} not registered.") self.active=[n for n in self.active if n!=name] def unload(self,name:str): self.deactivate(name) if name in self.weights: del self.weights[name] if name in self.adapters: del self.adapters[name] def set_active(self,names:list[str]): for name in names: if name not in self.adapters: raise KeyError(f"LoRA adapter {name} not registered.") self.active=names def set_weight(self,name:str,w:float): if name not in self.adapters: raise KeyError(f"LoRA adapter {name} not registered.") self.weights[name]=w def forward(self,x:torch.Tensor): y=self.base(x) if not self.active: return y z = None if self.mode=="additive": for n in self.active: adp=self.adapters[n] d=adp(x) z = d if z is None else (z + d) if self.mode=="exclusive": z=self.adapters[self.active[0]](x) if self.mode=="weighted": for n in self.active: adp=self.adapters[n] w=self.weights.get(n,1.0) if w==0.0: continue d=adp(x)*w z = d if z is None else (z + d) if z is None: return y if self.cap_norm is not None: norm=torch.linalg.vector_norm(z,dim=-1,keepdim=True) scale=torch.clamp(self.cap_norm/(norm+1e-6),max=1.0) z=z*scale if self.global_scale is not None: z=z*self.global_scale return y+z class Transformer_block(nn.Module): def __init__(self, hidden_size:int=640, num_heads:int=10, rms_eps: float=1e-5, dropout:float=0.1, num_kv_heads:int|None=None, use_qk_RMSnorm: bool = True, learnable_temp: bool = True, use_affine:bool=True, qkrms_eps:float|None=None, use_sampleatt:bool=False, rope_type:str="yarn", training_length:int=4096, use_swiGLU: bool = True, mlp_ratio: float = 4.0, RoPE_base:float=2e4, max_position_embeddings:int=16384, drop_path:float=0.1, resid_dropout:float=0.05, use_moe_layer: bool=False, moe_use_detach:bool=False, moe_cap_factor:float=1.25, moe_num_expert:int=4, use_ssm:bool=False, mamba_d_state: int = 64, mamba_d_conv: int = 4, mamba_expand = 1.5, use_HF_ROPE=True, use_tcn:bool=False ): super().__init__() self.register_buffer("_env",torch.empty(0)) self.gamma_att = nn.Parameter(torch.ones(1) * 1.0) self.gamma_mlp = nn.Parameter(torch.ones(1) * 1.0) if qkrms_eps is None: self.qkrms_eps=rms_eps else: self.qkrms_eps = qkrms_eps if num_kv_heads is None: self.num_kv_heads=num_heads else: self.num_kv_heads= num_kv_heads self.attRMS_norm=RMSNorm(hidden_size,rms_eps,use_affine,self._env.dtype,self._env.device) self.use_ssm=bool(use_ssm and HAS_MAMBA2) if self.use_ssm: cfg2 = Mamba2Config( hidden_size=hidden_size, # 1280 num_heads=10, head_dim=192, # 256 n_groups=5, # 5 state_size= mamba_d_state, # 64 conv_kernel=mamba_d_conv, # 4 expand=mamba_expand, # 1.5 num_hidden_layers=1, use_cache=True, use_bias=False, norm_before_gate=True, rms_norm=True, vocab_size=1, ) self.ssm = Mamba2Model(cfg2) else: self.selfatt=Mutihead_attention(hidden_size,num_heads,self.num_kv_heads,max_position_embeddings,RoPE_base,dropout,use_qk_RMSnorm,self.qkrms_eps,learnable_temp,use_sampleatt,rope_type,training_length,use_HF_ROPE) self.is_moe = use_moe_layer self.gamma_min = 0.6 self.gamma_max = 1.2 if self.is_moe: self.MLP = MoEMLP(hidden_size,dropout,moe_num_expert,mlp_ratio,use_swiGLU,moe_cap_factor,moe_use_detach) else: self.MLP = MLP(hidden_size, dropout, mlp_ratio, use_swiGLU) self.droppath=Drop_path(drop_path) self.MLPrms_norm=RMSNorm(hidden_size,rms_eps,use_affine,self._env.dtype,self._env.device) self.residual_dropout=nn.Dropout(resid_dropout) self.use_tcn = use_tcn if self.use_tcn: self.tcnRMS_norm = RMSNorm(hidden_size, rms_eps, use_affine, self._env.dtype, self._env.device) self.tcn_branch = TCNBranchPT( hidden_size=hidden_size) self.tcc_gate = nn.Parameter(torch.tensor(-2.1972246, dtype=self._env.dtype)) def _bounded(self, raw): # γ = γ_min + (γ_max-γ_min) * sigmoid(raw) return self.gamma_min + (self.gamma_max - self.gamma_min) * torch.sigmoid(raw) def forward( self, x: torch.Tensor, attention_mask: torch.Tensor | None = None, is_causal: bool = True, past_state: dict | None = None, use_cache: bool = False, img_len: int | None = None, ): # x: [B, T, H] B, T, H = x.shape past_state = past_state or {} present_state = {} if use_cache else None # ===================== # 1. Attention or SSM # ===================== if not self.use_ssm: attout, att_present = self.selfatt( self.attRMS_norm(x), attention_mask, is_causal, past_state=past_state, use_cache=use_cache, img_len=img_len, ) if use_cache and att_present is not None: present_state.update(att_present) else: # --------- SSM / Mamba2 分支 --------- if (attention_mask is not None) and (attention_mask.device != x.device): attention_mask = attention_mask.to(x.device) # 取上一次保存的状态 ssm_fields = past_state.get("ssm", None) if past_state is not None else None if ssm_fields is not None: ssm_cache_prev = ssm_fields.get("cache_params", None) cache_pos_prev = ssm_fields.get("cache_position", None) # shape [B] or None else: ssm_cache_prev = None cache_pos_prev = None # 对输入做和训练一致的norm/layout hs = self.attRMS_norm(x) # 你训练时为了layout对齐做的 transpose-contiguous-transpose，就保留 hs = hs.transpose(1, 2).contiguous().transpose(1, 2).contiguous() # [B,T,H] 但内存布局是你训练期喜欢的 fastpath_even_in_eval = (not use_cache) if fastpath_even_in_eval: was_mode = self.ssm.training try: # 临时切到 train()：只影响 ssm 子模块，不改变梯度开关（由 no_grad 决定） self.ssm.train(True) out = self.ssm( inputs_embeds=hs, attention_mask=attention_mask, cache_params=None, cache_position=None, use_cache=False, ) finally: # 恢复到进入前的模式（非常重要） self.ssm.train(was_mode) else: # 增量解码：保持外部模式（通常 eval） out = self.ssm( inputs_embeds=hs, attention_mask=attention_mask, cache_params=ssm_cache_prev, cache_position=cache_pos_prev, use_cache=True, ) # SSM 输出 attout = out.last_hidden_state.to(x.dtype) # [B,T,H] 这就是这一层SSM分支的本轮输出 ssm_noise=float(0.04) if self.training: s = torch.empty((), device=attout.device, dtype=attout.dtype).uniform_(1.0 - ssm_noise, 1.0) c = 1.0 / (1.0 - 0.5 * ssm_noise) attout = attout * (s * c) # 我们要把“下一轮要用的 cache_state”记下来，连同正确的下一拍位置 if use_cache: # 计算下一轮的 cache_position 向量 # # case 1: 这是 prefill（整段prompt喂进去） # - 此时 cache_pos_prev 是 None # - 我们需要根据这一轮喂了多少真实token(不算左pad)来初始化它 # case 2: 这是增量 streaming（一般 T=1） # - 我们已有 cache_pos_prev 了 # - 下一拍位置 = 旧位置 + 当前这次我们真正追加的token数 # if cache_pos_prev is None: if attention_mask is not None: # attention_mask: [B,T]，是1的地方才是真正喂给模型的有效token step_len_vec = attention_mask.to(torch.int64).sum(dim=-1) # [B], 每条样本本轮喂了多少有效token else: # 没有mask就说明没有pad，整段都有效 step_len_vec = torch.full( (B,), T, device=x.device, dtype=torch.long, ) # prefill之后的"下一拍位置"就是各自的有效长度 next_cache_pos = step_len_vec # shape [B] else: # streaming 步：通常 T==1 step_len_vec = torch.full( (B,), T, device=x.device, dtype=torch.long, ) # 注意这里是 elementwise 加法，不是广播成同一个数字 # cache_pos_prev.shape == [B] next_cache_pos = cache_pos_prev + step_len_vec # [B] present_state["ssm"] = { "cache_params": out.cache_params, # HF给我们的新状态 "cache_position": next_cache_pos, # 我们手工维护的 per-sample 时钟 } # ===================== # 2. 可选 TCN 分支 # ===================== if self.use_tcn: x_norm = self.tcnRMS_norm(x) if use_cache: tcn_states_in = past_state.get("tcn", None) if tcn_states_in is None: tcn_states_in = [None] * B if B == 1: state_in = tcn_states_in[0] if tcn_states_in is not None else None tcn_y, tcn_state_out = self.tcn_branch( x_norm, attention_mask, inference=True, state_in=state_in, return_state=True, ) present_state["tcn"] = [tcn_state_out] else: tcn_y, tcn_states_out = TCNBranchPT.tcn_step_batch( self.tcn_branch, x_norm, attention_mask, tcn_states_in, ) present_state["tcn"] = tcn_states_out else: tcn_y = self.tcn_branch( x_norm, attention_mask, inference=False, ) tcn_y = tcn_y.to(x.dtype) gate = 0.2 + (0.7 - 0.2) * torch.sigmoid(self.tcc_gate) attout = (1.0 - gate) * attout + gate * tcn_y # ===================== # 3. 残差 + MLP # ===================== attout = attout * self._bounded(self.gamma_att) x = x + self.droppath(self.residual_dropout(attout)) if self.is_moe and self.training: mlp_out, aux_loss = self.MLP(self.MLPrms_norm(x)) mlp_out = mlp_out * self._bounded(self.gamma_mlp) else: normed = self.MLPrms_norm(x) mlp_out = self.MLP(normed) mlp_out = mlp_out * self._bounded(self.gamma_mlp) aux_loss = x.new_zeros(()) x = x + self.droppath(self.residual_dropout(mlp_out)) return x, aux_loss, present_state class TinyLLM(PreTrainedModel, GenerationMixin): """tinyLLM with the standard Transformers save/load lifecycle.""" config_class = Config base_model_prefix = "" main_input_name = "input_ids" _is_stateful = True _supports_cache_class = False def __init__(self, cfg: Config): super().__init__(cfg) self._train_step = 0 self.cfg = cfg self.config = cfg self.generation_config = GenerationConfig( eos_token_id=cfg.eos_token_id, pad_token_id=cfg.pad_token_id, ) self.end_text_tok_id = cfg.eos_token_id self.eso_loss_radio = cfg.eso_loss_radio self.drop_high_loss = cfg.drop_high_loss self.adaptive_cutoffs = getattr(self.cfg, "adaptive_cutoffs", [20000, 60000]) self.adaptive_div = getattr(self.cfg, "adaptive_div", 4.0) self.adaptive_calibrate_every = 4 hidden_size = cfg.hidden_size vocab_size = cfg.vocab_size self.tok_embed = nn.Embedding(vocab_size, hidden_size) nn.init.normal_(self.tok_embed.weight, mean=0.0, std=0.02) self.embdrop = nn.Dropout(cfg.embeddingdropout) droppath_list = [(cfg.drop_path * i / max(cfg.num_hidden_layers - 1, 1)) for i in range(cfg.num_hidden_layers)] mlp_ratio_by_layer = _build_mlp_ratio_schedule(self.cfg) self.blocks = nn.ModuleList([Transformer_block( hidden_size=cfg.hidden_size, num_heads=cfg.num_attention_heads, rms_eps=cfg.rms_norm_eps, dropout=cfg.dropout, num_kv_heads=cfg.num_key_value_heads, use_qk_RMSnorm=cfg.use_qk_norm, learnable_temp=cfg.learnable_temp, use_affine=cfg.use_affine, qkrms_eps=cfg.qkrms_norm_eps, use_sampleatt=cfg.use_sampleatt, rope_type=cfg.rope_type, training_length=cfg.train_maxlength, use_swiGLU=cfg.use_swiGLU, mlp_ratio=mlp_ratio_by_layer[i], RoPE_base=cfg.RoPE_base, max_position_embeddings=cfg.max_position_embeddings, drop_path=droppath_list[i], use_tcn=(cfg.use_tcn and (i in cfg.tcn_layers)), use_ssm=(cfg.use_ssm and (i in cfg.ssm_layers)), resid_dropout=cfg.residual_dropout, use_moe_layer=(getattr(cfg, "use_moe", False) and (i in cfg.moe_layers)), moe_use_detach=cfg.moe_use_detach, moe_cap_factor=cfg.moe_cap_factor, moe_num_expert=cfg.num_expert, use_HF_ROPE=True ) for i in range(cfg.num_hidden_layers)]) self.register_buffer("_env0", torch.empty(0)) self.final_norm = RMSNorm(cfg.hidden_size, cfg.rms_norm_eps, cfg.use_affine, self._env0.dtype, self._env0.device) self.lm_bias = nn.Parameter(torch.zeros(vocab_size)) self._kd_seen = 0 self._kd_hit = 0 self.use_vision = bool(getattr(cfg, "use_vision", False)) self.max_vision_tokens = int(getattr(cfg, "max_vision_tokens", 0)) self.vision_feature_dim = getattr(cfg, "vision_feature_dim", None) self.num_query_tokens = cfg.num_query_tokens if self.use_vision: assert self.vision_feature_dim is not None, \ "cfg.vision_feature_dim must be set when cfg.use_vision=True" # ---- 1) 轻量 projector: vision_feature_dim -> hidden_size ---- self.vision_projector = VisionProjector( vision_dim=self.vision_feature_dim, hidden_dim=cfg.hidden_size, dropout=getattr(cfg, "vision_dropout", 0.0), use_rmsnorm=getattr(cfg, "vision_use_rmsnorm", True), ) # 可选：给整张图一个 type embedding（加到 patch / q 上都行） self.vision_type_embed = nn.Parameter(torch.zeros(cfg.hidden_size)) # ---- 2) Q-Former 配置（默认四层 encoder-style block） ---- self.qformer_num_layers = int(getattr(cfg, "qformer_num_layers", 4)) self.qformer_mlp_ratio = float(getattr(cfg, "qformer_mlp_ratio", 2.0)) self.qformer_dropout = float(getattr(cfg, "vision_dropout", 0.0)) self.qformer_use_rmsnorm = bool(getattr(cfg, "vision_use_rmsnorm", True)) # ---- 3) learnable queries（工业界标配）---- # [1, Nq, H]，每张图复制一份；初始化用 1/sqrt(H) 缩一下 self.query_embed = nn.Parameter( torch.randn(1, self.num_query_tokens, cfg.hidden_size) / math.sqrt(cfg.hidden_size) ) # Q-Former Block: Self-Attn + Cross-Attn + MLP（全是 encoder-style） self.qformer_blocks = nn.ModuleList( [ QFormerBlock( hidden_dim=cfg.hidden_size, num_heads=cfg.num_attention_heads, mlp_ratio=self.qformer_mlp_ratio, dropout=self.qformer_dropout, use_rms_norm=self.qformer_use_rmsnorm, ) for _ in range(self.qformer_num_layers) ] ) # 顶层再来一个 final norm（跟主干一样风格） self.qformer_final_norm = RMSNorm( cfg.hidden_size, cfg.rms_norm_eps, cfg.use_affine ) self.num_vision_views = int(getattr(cfg, "num_vision_views", 5)) self.num_vision_view_types=2 # 你现在 global_pos 实际更像“thumb 网格上的 index”（0..Nv-1）。 # 为了兼容：直接给到 1024（32*32），就算你 unshuffle 后 Nv=256 也没问题（pos 仍 <1024）。 self.vision_grid_tokens = int(getattr(cfg, "vision_grid_tokens", 1024)) # (a) view embedding：5 张图 = 5 个 learnable embedding（你要的那个） self.vision_view_embed = nn.Embedding(self.num_vision_view_types, cfg.hidden_size) # (b) global position embedding：离散 gpos -> embedding self.vision_pos_embed = nn.Embedding(self.vision_grid_tokens, cfg.hidden_size) nn.init.normal_(self.vision_view_embed.weight, mean=0.0, std=0.02) nn.init.normal_(self.vision_pos_embed.weight, mean=0.0, std=0.02) # (c) FiLM：把连续的 (x_norm, y_norm, dx, dy) -> (gamma, beta) # 工业里常见：小 MLP 输出 2H，然后做 proj_vis = proj_vis*(1+γ)+β film_hidden = int(getattr(cfg, "vision_film_hidden", max(32, cfg.hidden_size // 4))) self.vision_film_mlp = nn.Sequential( nn.Linear(4, film_hidden, bias=True), nn.SiLU(), nn.Linear(film_hidden, 2 * cfg.hidden_size, bias=True), ) # 关键：最后一层置零，保证初始化“完全不影响” nn.init.normal_(self.vision_film_mlp[-1].weight, mean=0.0, std=1e-3) nn.init.zeros_(self.vision_film_mlp[-1].bias) # (d) 三个 gate：你担心 gate 歪掉回不来 → 用 tanh gate + 小 scale # tanh(0)=0：初始化就是“零注入”，非常稳 self.vision_gate_view = nn.Parameter(torch.tensor(0.2)) self.vision_gate_pos = nn.Parameter(torch.tensor(0.2)) self.vision_gate_film = nn.Parameter(torch.tensor(0.2)) # FiLM 输出再乘一个小尺度，防止极端（比 hard-clip 更工业） self.vision_film_scale = float(getattr(cfg, "vision_film_scale", 0.10)) else: self.vision_projector = None self.vision_type_embed = None self.query_embed = None self.qformer_blocks = None self.qformer_final_norm = None def get_input_embeddings(self): return self.tok_embed def set_input_embeddings(self, value): self.tok_embed = value def get_output_embeddings(self): # Output logits are projected with the tied token embedding in forward(). return None def encode_image_with_qformer( self, vision_feats: torch.Tensor, # [B,5,1024,Dv] or [B,1024,Dv] vision_mask: torch.Tensor , # [B,5,1024] or [B,1024] global_pos: torch.Tensor, # [B,5,1024] or [B,1024] global_off: torch.Tensor, # [B,5,1024,4] or [B,1024,4] ) -> torch.Tensor: """ 只跑一次 QFormer： - 把 5 个 view 沿 token 维拼成 5120（可选 + 4 sep） - 返回 q_img: [B, Nq, H] """ def gate(p: torch.Tensor) -> torch.Tensor: return torch.tanh(p) if vision_mask is None or global_pos is None or global_off is None: raise ValueError("vision_mask/global_pos/global_off must be provided when vision_feats is provided") # ---- 0) 统一维度：老数据 [B,Nv,Dv] -> [B,1,Nv,Dv] ---- if vision_feats.ndim == 3: vision_feats = vision_feats.unsqueeze(1) if vision_mask is not None: vision_mask = vision_mask.unsqueeze(1) if global_pos is not None: global_pos = global_pos.unsqueeze(1) if global_off is not None: global_off = global_off.unsqueeze(1) B, V, Nv, Dv = vision_feats.shape if V != self.num_vision_views: # 你要的就是 5；不想 silent 就直接报 raise ValueError(f"Expected V={self.num_vision_views} views, got V={V}") vision_mask = vision_mask.to(torch.long) # ---- 1) 你不想截断：那就“若 max_vision_tokens 不够就报错” ---- total_kv = V * Nv if self.max_vision_tokens and total_kv > self.max_vision_tokens: raise ValueError(f"vision tokens {total_kv} exceed max_vision_tokens={self.max_vision_tokens}. " f"Increase max_vision_tokens.") # ---- 2) 构造每个 token 的 view_id: [B, V*Nv] ---- view_ids = torch.zeros((B, V, Nv), device=vision_feats.device, dtype=torch.long) view_ids[:, -1, :] = 1 # thumb=1（最后一张） view_ids = view_ids.reshape(B, V * Nv) # ---- 3) 先把 5 段拼起来（先拼 raw Dv / mask / meta）---- kv_raw = vision_feats.reshape(B, V * Nv, Dv) kv_mask = vision_mask.reshape(B, V * Nv).to(torch.long) gp = global_pos.reshape(B, V * Nv).to(torch.long) go = global_off.reshape(B, V * Nv, -1).to(torch.float32) # ---- 4) projector：Dv -> H（复用你已有 projector 接口）---- proj_vis = self.vision_projector(kv_raw) # [B, V*Nv, H] if self.vision_type_embed is not None: proj_vis = proj_vis + self.vision_type_embed.view(1, 1, -1) H = proj_vis.size(-1) dtype = proj_vis.dtype # ---- 6) 加 view embedding（工业标准：加法 + gate）---- g_view = gate(self.vision_gate_view).to(dtype) proj_vis = proj_vis + g_view * self.vision_view_embed(view_ids).to(dtype) # ---- 7) 加 global_pos embedding（离散 index）---- g_pos = gate(self.vision_gate_pos).to(dtype) if gp.min() < 0 or gp.max() >= self.vision_pos_embed.num_embeddings: raise ValueError(f"global_pos out of range: [{gp.min().item()}, {gp.max().item()}], " f"num_embeddings={self.vision_pos_embed.num_embeddings}") proj_vis = proj_vis + g_pos * self.vision_pos_embed(gp).to(dtype) # ---- 8) FiLM：连续 (x_norm,y_norm,dx,dy) 调制 ---- g_film = gate(self.vision_gate_film).to(dtype) s = (self.vision_film_scale * g_film).to(dtype) go_t = torch.tanh(go) # [-1,1]，x_norm/y_norm 本来就在 [-1,1]，tanh 基本不影响 film = self.vision_film_mlp(go_t).to(dtype) # [B,N,2H] gamma, beta = film.chunk(2, dim=-1) gamma = torch.tanh(gamma) * s beta = torch.tanh(beta) * s proj_vis = proj_vis * (1.0 + gamma) + beta # ---- 9) mask：padding token 清零（你原来也是这么干的思路）---- proj_vis = proj_vis * kv_mask.to(dtype).unsqueeze(-1) # ---- 10) QFormer：只跑一次 ---- q = self.query_embed.expand(B, -1, -1) # [B,Nq,H] for blk in self.qformer_blocks: q = blk(q, proj_vis, kv_mask=kv_mask) q = self.qformer_final_norm(q) # [B,Nq,H] # 如果整条 kv 全是 0，别污染：q 直接清零 has_any = (kv_mask.sum(dim=1, keepdim=True) > 0).to(dtype) # [B,1] q = q * has_any.view(B, 1, 1) return q @property def device(self): return next(self.parameters()).device def _project_subset_logits(self, x, idx_btK): """ x: [B, T, H] final hidden idx_btK: [B, T, K] teacher top-k indices (Long) return: [B, T, K] student logits on these K words """ B, T, H = x.shape K = idx_btK.size(-1) W = self.tok_embed.weight.float() # [V, H]（绑权重） # 展平 + unique 只算一次 idx_flat = idx_btK.to(device=W.device, dtype=torch.long).reshape(-1) uniq, inv = torch.unique(idx_flat, sorted=False, return_inverse=True) b = self.lm_bias.float() # [V] W_sub = W[uniq] # [U, H] b_sub = b[uniq] if b is not None else None x_flat = x.reshape(B * T, H).float() # [B*T, H] logits_sub = x_flat @ W_sub.t() # [B*T, U] if b_sub is not None: logits_sub = logits_sub + b_sub # broadcast # 把 U 映回 B*T*K 的排列 logits_btK = logits_sub.gather(1, inv.view(B * T, K)).view(B, T, K) return logits_btK def _mask_shortlist_dups(self, short_idx, logits_short): """ short_idx: [B,T,Kp] token id 列表 logits_short:[B,T,Kp] 对应 logits 返回：logits_short'，把重复 token（除首个）置为 -inf """ B, T, Kp = short_idx.shape # 展平成 [BT, Kp] si = short_idx.view(B * T, Kp) lg = logits_short.view(B * T, Kp) # 对每行排序以便找重复 vals, order = torch.sort(si, dim=-1) # [BT,Kp] dup = torch.zeros_like(vals, dtype=torch.bool) # [BT,Kp] dup[:, 1:] = (vals[:, 1:] == vals[:, :-1]) # 重复（与左邻相等）标 True # 把“排序后重复位置”映回“原列顺序” inv = torch.empty_like(order) inv.scatter_(1, order, torch.arange(Kp, device=order.device).unsqueeze(0).expand_as(order)) dup_in_orig = dup.gather(1, inv) # [BT,Kp] True 表示该列是重复的第二个或以后 dup_in_orig[:, 0] = False # 屏蔽重复列 lg = lg.masked_fill(dup_in_orig, float('-inf')) return lg.view(B, T, Kp) @torch.no_grad() def _debug_kd_batch(self, kd_idx, kd_val, kd_mask, labels, tau=1.5, K_print=5, max_print=10): """ 简要体检：命中率、label 的 rank 分布、概率归一偏差、抽样打印 kd_idx: [B,T,K] long kd_val: [B,T,K] float kd_mask: [B,T] bool/byte labels: [B,T] long """ device = kd_val.device B, T, K = kd_idx.shape m = kd_mask.bool() n_valid = m.sum().item() if n_valid == 0: print("[KD] no valid positions in mask.") return # 1) teacher 概率（温度） t_prob = torch.softmax(kd_val[m] / tau, dim=-1) # [M,K] prob_sum = t_prob.sum(dim=-1) # should ≈1 sum_err_mean = (prob_sum - 1.0).abs().mean().item() # 2) 是否排序（如果没排序取 argmax 定义 top1） # 这里的 top1 是 kd_val 最大对应的 idx argmax_in_k = kd_val[m].argmax(dim=-1) # [M] top1_idx = kd_idx[m].gather(1, argmax_in_k.view(-1, 1)).squeeze(1) # [M] # 3) 命中率（label 是否在 top-k 内） lab = labels[m] hits_any = (kd_idx[m] == lab.view(-1, 1)).any(dim=-1) # [M] hit_rate = hits_any.float().mean().item() # 4) label 的 rank（若在 top-k 内） # rank 定义：按 kd_val 降序后的名次；若无序，先对 kd_val 排序 # 索引到排序后的顺序 sorted_vals, order = torch.sort(kd_val[m], dim=-1, descending=True) # [M,K] sorted_idx = kd_idx[m].gather(1, order) # [M,K] # 找出 label 在 sorted_idx 的位置 eq = (sorted_idx == lab.view(-1, 1)) # [M,K] has_rank = eq.any(dim=-1) # [M] ranks = torch.where(has_rank, eq.float().argmax(dim=-1), torch.full_like(argmax_in_k, -1)) # 统计 rank 直方（0-based；-1 代表不在 top-k） rank_hist = torch.bincount(ranks.clamp(min=-1).add_(1), minlength=K + 1).cpu().tolist() # rank_hist[0] 是 -1（未命中）的数量，其余 1..K 是 rank 0..K-1 的数量 # 5) 抽样打印若干条 sel = torch.nonzero(m.flatten(), as_tuple=False).flatten() # 线性坐标 sel = sel[:min(max_print, sel.numel())] bt = torch.stack([sel // T, sel % T], dim=-1) # [S,2] print(f"[KD] M(valid)={n_valid}, hit@{K}={hit_rate * 100:.2f}%, " f"mean(|sum(p)-1|)={sum_err_mean:.3e}, top1!=label ratio={(top1_idx != lab).float().mean().item() * 100:.2f}%") print(f"[KD] rank hist (bin0=-1=miss, 1..K = rank0..K-1): {rank_hist}") for s in bt.cpu().tolist(): b, t = s l = labels[b, t].item() k_ids = kd_idx[b, t, :K_print].cpu().tolist() k_vals = kd_val[b, t, :K_print].cpu().tolist() print(f" (b={b}, t={t}) label={l} | topK_idx={k_ids} | topK_val={['%.3f' % v for v in k_vals]}") def _reduce_supervised_loss( self, per_token_loss: torch.Tensor, token_weights: torch.Tensor, ) -> torch.Tensor: """Reduce CE without letting long answers dominate every update. token_mean preserves the historical objective. sample_mean first normalizes each example by its own supervised-token count. hybrid is a convex combination and is the safer default for continual VLM SFT. """ token_denom = token_weights.sum().clamp_min(1.0) token_mean = (per_token_loss * token_weights).sum() / token_denom reduction = str(getattr(self.cfg, "loss_reduction", "token_mean")) if reduction == "token_mean": return token_mean sample_denom = token_weights.sum(dim=1) valid_samples = sample_denom > 0 if not valid_samples.any(): return token_mean sample_losses = ( (per_token_loss * token_weights).sum(dim=1) / sample_denom.clamp_min(1.0) ) sample_mean = sample_losses[valid_samples].mean() if reduction == "sample_mean": return sample_mean if reduction == "hybrid": alpha = float(getattr(self.cfg, "sample_mean_alpha", 0.75)) return alpha * sample_mean + (1.0 - alpha) * token_mean raise ValueError(f"Unsupported loss_reduction={reduction!r}") def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None, labels: torch.Tensor | None = None, vision_feats: torch.Tensor | None = None, # [B, Nv, D_v] vision_mask: torch.Tensor | None = None, # [B, Nv]，1=valid global_pos: torch.Tensor | None = None, # [B,Nv] or [B,V,Nv] global_off: torch.Tensor | None = None, # [B,Nv,4] or [B,V,Nv,4] is_causal: bool = True, past_states: Optional[List[dict]] = None, use_cache: bool = False, kd_idx: torch.Tensor | None = None, kd_val: torch.Tensor | None = None, kd_mask: torch.Tensor | None = None, short_idx=None, gold_col=None, short_logq: torch.Tensor | None = None, enhance_math: bool = True, math_sample_mask=None, force_checkpoint: bool = False, **kwargs, ): past_key_values = kwargs.pop("past_key_values", None) return_dict = kwargs.pop("return_dict", False) if past_states is None and past_key_values is not None: past_states = past_key_values B, T = input_ids.shape is_xlong = self.training and (T > getattr(self.cfg, "train_maxlength", T)) if force_checkpoint: is_xlong = True if vision_feats is not None: if not self.use_vision or self.vision_projector is None: raise ValueError("vision_feats was provided but cfg.use_vision=False or projector is None.") if kd_idx is not None or kd_val is not None or kd_mask is not None: raise NotImplementedError("KD with vision tokens is not implemented yet.") if past_states is not None: raise NotImplementedError( "vision_feats should only be passed on the first prefill step (past_states=None)." ) if self.training and use_cache and (labels is not None): use_cache = False warnings.warn("KV-cache should be disabled during training: use_cache has been set to False.") past_states = None x = self.tok_embed(input_ids) T_txt = T tok=x img_kwargs = {} if (vision_feats is not None) and (self.vision_projector is not None): if vision_feats.ndim == 4: Bv, V, Nv, Dv = vision_feats.shape if V != self.num_vision_views: raise ValueError(f"Expected V={self.num_vision_views}, got V={V}") elif vision_feats.ndim == 3: Bv, Nv, Dv = vision_feats.shape else: raise ValueError(f"vision_feats must be 3D or 4D, got shape={tuple(vision_feats.shape)}") if Bv != B: raise ValueError(f"vision_feats batch={Bv} does not match input_ids batch={B}.") # 不截断：超了就让 encode 里报错（或你也可以在这先报） q_img = self.encode_image_with_qformer( vision_feats=vision_feats, vision_mask=vision_mask, global_pos=global_pos, global_off=global_off, ) Bq, Nq, H = q_img.shape img_kwargs = {"img_len": int(Nq)} # attention_mask / labels 的扩展照旧，只是前缀长度用 Nq attn_img = torch.ones((B, Nq), dtype=torch.long, device=input_ids.device) attention_mask = torch.cat([attn_img, attention_mask], dim=1) x = torch.cat([q_img, tok], dim=1) if labels is not None: ignore = self.cfg.ignore_index labels_exp = torch.full((B, Nq + T_txt), ignore, dtype=labels.dtype, device=labels.device) labels_exp[:, Nq:] = labels labels = labels_exp else: x = tok x = self.embdrop(x) B, L, _ = x.shape seq_len = L ce_loss = torch.zeros((), device=input_ids.device, dtype=torch.float32) kd_ce = torch.zeros((), device=input_ids.device, dtype=torch.float32) present_states = [] if use_cache else None aux_total = torch.zeros((), device=input_ids.device, dtype=torch.float32) for i, layer in enumerate(self.blocks): is_moe_layer = (self.cfg.use_moe and (i in self.cfg.moe_layers)) if is_xlong: # 16K 特例：所有非 MoE 层都 checkpoint，最大化省显存 use_ckpt = (self.training and not use_cache and (not is_moe_layer)) else: # 普通 2K/8K：保持你原来的策略（要 cfg.use_checkpoint=True 且只在偶数层 ckpt） use_ckpt = ( self.training and self.cfg.use_checkpoint and not use_cache and (i % 2 == 0) and (not is_moe_layer) ) layer_past = past_states[i] if (past_states is not None) else None if use_ckpt: def layer_fwd( _x, _layer=layer, _layer_past=layer_past, _attn_mask=attention_mask, _is_causal=is_causal, ): x_out, aux, _present = _layer( _x, attention_mask=_attn_mask, is_causal=_is_causal, past_state=_layer_past, use_cache=False, # checkpoint 下不存 cache **img_kwargs, ) return x_out, aux x, aux_loss = checkpoint( layer_fwd, x, preserve_rng_state=True, use_reentrant=self.cfg.checkpoint_use_reentrant ) present = None else: x, aux_loss, present = layer( x, attention_mask=attention_mask, is_causal=is_causal, past_state=layer_past, use_cache=use_cache,**img_kwargs ) aux_total = aux_total + aux_loss if use_cache: present_states.append(present) aux_total *= 0.5 x = self.final_norm(x) # [B,T,H] # ===== 只有推理（无 labels） ===== if labels is None: logits = F.linear(x, self.tok_embed.weight, self.lm_bias) unk_id = getattr(self.cfg, "unk_token_id", None) if (unk_id is not None) and (0 <= int(unk_id) < logits.size(-1)): logits[..., int(unk_id)] = float("-inf") if return_dict: return CausalLMOutputWithPast( loss=None, logits=logits, past_key_values=present_states if use_cache else None, ) if use_cache: return {"logits": logits, "past_states": present_states, "past_key_values": present_states} return {"logits": logits} if not self.training: logits_full = F.linear(x.float(), self.tok_embed.weight.float(), self.lm_bias.float()) ce_loss = F.cross_entropy( logits_full.view(-1, logits_full.size(-1)), labels.view(-1).long(), ignore_index=self.cfg.ignore_index, label_smoothing=0.0, reduction="mean", ) loss = ce_loss + aux_total # eval 下 aux_total≈0，保留无妨 ret = {"loss": loss, "logits": logits_full, "ce_loss": ce_loss, "aux_loss": aux_total, "kd_loss": kd_ce} if use_cache: ret["past_states"] = present_states return ret # ===== 计算主 CE（两条路：adaptive / short-list 与 full） ===== if short_idx is not None and gold_col is not None: # --- 短清单 CE：保持不变 --- logits_short = self._project_subset_logits(x, short_idx.long()) # [B,T,Kp] logits_short = self._mask_shortlist_dups(short_idx, logits_short) logits_short_corr = logits_short if short_logq is not None: logits_short_corr = logits_short_corr - short_logq.to(logits_short.device, logits_short.dtype) ce_logp = F.log_softmax(logits_short_corr, dim=-1) # [B,T,Kp] tgt_col = gold_col.long() # [B,T]（恒为 0） valid = (labels != self.cfg.ignore_index) ce_nll = -ce_logp.gather(-1, tgt_col.unsqueeze(-1)).squeeze(-1) # [B,T] ce_loss = self._reduce_supervised_loss(ce_nll, valid.float()) loss = ce_loss + aux_total # KD（对短清单对齐），保持不变 has_kd = (kd_idx is not None) and (kd_val is not None) and (kd_mask is not None) and ( kd_idx.numel() > 0) if has_kd: kd_alpha = 0.2 kd_tau = 1.5 eq = (kd_idx.long().to(x.device).unsqueeze(-1) == short_idx.long().to(x.device).unsqueeze( -2)) # [B,T,K_kd,Kp] pos_map = eq.float().argmax(dim=-1) # [B,T,K_kd] hit = eq.any(dim=-1) # [B,T,K_kd] s_logp = F.log_softmax(logits_short / kd_tau, dim=-1) # [B,T,Kp] s_logp_kd = torch.take_along_dim(s_logp, pos_map, dim=-1) # [B,T,K_kd] t_prob = torch.softmax(kd_val.to(x.device) / kd_tau, dim=-1) # [B,T,K_kd] hit = hit.bool() if hit.dtype != torch.bool else hit s_logp_kd = s_logp_kd * hit.float() t_prob = t_prob * hit.float() denom = t_prob.sum(dim=-1, keepdim=True).clamp_min(1e-9) t_prob = t_prob / denom kd_token = -(t_prob * s_logp_kd).sum(dim=-1) mask_tok = kd_mask.bool() & valid.bool() & hit.any(dim=-1) if mask_tok.any(): kd_ce = kd_token[mask_tok].mean() kd_ce = (kd_tau * kd_tau) * kd_ce * kd_alpha kd_ce = kd_ce.to(torch.float32) loss = loss + kd_ce self._train_step += 1 if use_cache: return {"loss": loss, "past_states": present_states} return {"loss": loss, "ce_loss": ce_loss, "aux_loss": aux_total, "kd_loss": kd_ce} # ===== 全量 CE 分支（这里加入“高损失截尾”，仅作用于 CE） ===== logits_full = None is_xlong = (T > getattr(self.cfg, "train_maxlength", T)) # 例如 train_maxlength=8192 时，T=16384 就视为 xlong if self.training and is_xlong and (vision_feats is not None): raise RuntimeError( f"[TinyLLM] xlong training (T={T}) with vision_feats is not supported yet. " "当前实现的 xlong 分块 CE 没有正确处理 vision 前缀，" "请减小文本长度到 train_maxlength 以下，或者在 xlong 训练时关闭 vision_feats。" ) if self.training: if is_xlong: # --- 16K 等超长序列：分块算 CE，避免一次性 [B,T,V] --- chunk_size = 2048 # 你可以试 2048/4096，自行权衡 per_loss_chunks = [] W = self.tok_embed.weight b = self.lm_bias for start in range(0, T, chunk_size): end = min(T, start + chunk_size) x_slice = x[:, start:end, :] # [B, c, H] labels_slice = labels[:, start:end] # [B, c] # 如果这一块全是 ignore_index，就直接给 0，省点算力 if not (labels_slice != self.cfg.ignore_index).any(): loss_slice = x_slice.new_zeros( (B, end - start), dtype=torch.float32 ) else: logits_slice = F.linear( x_slice.float(), # 为了数值稳定，仍然用 float32 算 W.float(), b.float() if b is not None else None, ) # [B, c, V] loss_slice = F.cross_entropy( logits_slice.view(-1, logits_slice.size(-1)), labels_slice.contiguous().view(-1).to(torch.long), ignore_index=self.cfg.ignore_index, label_smoothing=0.0, reduction="none", ).view(B, end - start) # [B, c] per_loss_chunks.append(loss_slice) per_loss = torch.cat(per_loss_chunks, dim=1) # [B, T] # 注意：这里不再设置 logits_full，后面训练分支也不会返回 logits，节省显存 logits_full = None else: # --- 普通 2K/8K 序列：保持你原来的实现 --- logits_full = F.linear( x.float(), self.tok_embed.weight.float(), self.lm_bias.float() ) # [B,T,V] per_loss = F.cross_entropy( logits_full.view(-1, logits_full.size(-1)), labels.view(-1).to(torch.long), ignore_index=self.cfg.ignore_index, label_smoothing=0.0, reduction="none" ).view(B, seq_len) # 每 token NLL valid_bool = (labels != self.cfg.ignore_index).bool() # 先做 end_text 降权（仅在 valid 位） w = valid_bool.float() if hasattr(self, "end_text_tok_id"): eos_mask = (labels == self.end_text_tok_id) & valid_bool # w = torch.where(eos_mask, w.new_full((), float(self.eso_loss_radio)), w) # === 损失截尾（百分位，默认 warmup 后丢 top 1% 高损失 token）=== trunc_ratio = self.drop_high_loss # e.g., 0.01 = 丢 top 1% min_tokens_for_trunc = getattr(self.cfg, "loss_trunc_min_tokens", 100) # 至少多少有效 token 才截 if valid_bool.any() and (trunc_ratio > 0.0): v_losses = per_loss[valid_bool] if v_losses.numel() >= int(min_tokens_for_trunc): # 找到“要丢弃的 top-k 中最小的那个值”，作为阈值 num_trunc = max(1, int(v_losses.numel() * trunc_ratio)) top_vals = torch.topk(v_losses, num_trunc, largest=True, sorted=False).values loss_threshold = top_vals.min() # 保留：小于等于阈值的，或无效位（pad/ignore 仍为 0 权重） keep_mask = (per_loss <= loss_threshold) | (~valid_bool) w = w * keep_mask.float() loss_class_weight = getattr(self, "class_w_mathish", None) if loss_class_weight is not None and enhance_math: # w: 之前已经包含了 valid_mask、end_text 降权、截尾等信息 # 这里我们再乘一个 per-token 的数学权重 cw，以及可选的样本级 boost。 # ---- 1) 先构造 “label → 数学权重” 映射（只在 valid 位生效）---- cw = torch.ones_like(per_loss, dtype=per_loss.dtype, device=per_loss.device) idx = labels[valid_bool].to(torch.long) # 这些位置的 label 一定不是 ignore_index base_cw_valid = loss_class_weight.to(per_loss.dtype).to(per_loss.device).index_select(0, idx) # base_cw_valid: [num_valid]，>1 的那些就是你选出来的 mathish token if math_sample_mask is None: # ★ 兼容老逻辑：没传 mask，就当整 batch 都是数学样本 # → 所有 valid token 均按 base_cw_valid 加权（和原来完全一致） cw_valid = base_cw_valid else: # ★ 新逻辑：只有 isMath=1 的样本才启用 per-id 数学权重 # math_sample_mask: [B]，0/1 gate = math_sample_mask.view(-1, 1).to(per_loss.dtype) # [B,1] gate_full = gate.expand_as(per_loss) # [B,T] gate_valid = gate_full[valid_bool] # [num_valid] # gate=0 → cw=1（非数学样本不加权） # gate=1 → cw=base_cw_valid（数学样本里才用 per-id mathish 权重） cw_valid = 1.0 + gate_valid * (base_cw_valid - 1.0) cw[valid_bool] = cw_valid w = w * cw # ---- 2) （可选）再给 “整道数学样本” 一个样本级 boost，比如 ×1.2 ---- if math_sample_mask is not None: SAMPLE_BOOST = 1.0 # 如果你觉得太猛可以改成 1.1 sample_gate = 1.0 + (SAMPLE_BOOST - 1.0) * gate_full # [B,T] # ignore 位置本来 w=0，乘多少还是 0，不会出问题 w = w * sample_gate ce_loss = self._reduce_supervised_loss(per_loss, w) loss = ce_loss + aux_total # ====== KD（可选，保持不变）====== has_kd = (kd_idx is not None) and (kd_val is not None) and (kd_mask is not None) if has_kd: kd_alpha = 0.2 kd_tau = 1.5 kd_idx = kd_idx.to(x.device, dtype=torch.long) take_btK = self._project_subset_logits(x.float(), kd_idx) # [B,T,K] t_prob = torch.softmax(kd_val.to(take_btK.device) / kd_tau, dim=-1) # [B,T,K] s_logp = torch.log_softmax(take_btK / kd_tau, dim=-1) mask_tok = kd_mask.to(take_btK.device).bool() mask_tok = mask_tok & (labels != self.cfg.ignore_index) if mask_tok.any(): kd_token = -(t_prob * s_logp).sum(dim=-1) kd_ce = kd_token[mask_tok].mean() kd_ce = (kd_tau * kd_tau) * kd_ce * kd_alpha kd_ce = kd_ce.to(torch.float32) loss = loss + kd_ce self._train_step += 1 # ===== 返回 ===== if return_dict: return CausalLMOutputWithPast( loss=loss, logits=logits_full, past_key_values=present_states if use_cache else None, ) if use_cache: ret = {"loss": loss, "past_states": present_states} if logits_full is not None: ret["logits"] = logits_full return ret if logits_full is not None: return {"loss": loss, "logits": logits_full, "ce_loss": ce_loss, "aux_loss": aux_total, "kd_loss": kd_ce} return {"loss": loss, "ce_loss": ce_loss, "aux_loss": aux_total, "kd_loss": kd_ce} def _wrap_linear_with_lora( self, module: nn.Module, attr: str, adapter_name: str, rank: int = 16, dropout: float = 0.05, alpha: float = 16.0, mode: str = "exclusive", cap_norm: float | None = 1.0, global_scale: float = 1.0, ): base = getattr(module, attr, None) if base is None: return # Existing wrappers must register the new adapter as well; otherwise # switching between ARC and OPD on one resident base is impossible. if isinstance(base, LoraLinear): if adapter_name not in base.adapters: base.register_adapter(name=adapter_name, rank=rank, dropout=dropout, alpha=alpha) return lora = LoraLinear(base, mode=mode, cap_norm=cap_norm, global_scale=global_scale) lora.register_adapter( name=adapter_name, rank=rank, dropout=dropout, alpha=alpha, ) lora.activate(adapter_name, exclusive=True) setattr(module, attr, lora) def attach_lora_adapter( self, adapter_name: str, rank: int = 16, dropout: float = 0.05, alpha: float = 16.0, target: str = "attn_mlp_top_half", # 你可以扩展其它策略 ): """ 在指定层位上挂一个名为 adapter_name 的 LoRA。 可以多次调用，用不同 adapter_name 做多个 LoRA。 """ L = len(self.blocks) if target == "attn_mlp_top_half": start = L // 2 end = L elif target == "all": start, end = 0, L elif target == "attn_mlp_skip_first4": # 跳过最底下 4 层，从第 4 层开始一路到顶 # （0-based：0,1,2,3 不挂；4,...,L-1 都挂） start = min(4, L) # 防止 L < 4 的极端情况 end = L else: # 你可以自己扩展其他策略 start, end = 0, L for i in range(start, end): block = self.blocks[i] # 1) 注意力（非 SSM block 才有 selfatt） if hasattr(block, "selfatt"): attn = block.selfatt for attr in ["w_q", "w_k", "w_v", "w_o"]: self._wrap_linear_with_lora( attn, attr, adapter_name=adapter_name, rank=rank, dropout=dropout, alpha=alpha, mode="exclusive", cap_norm=1.0, global_scale=1.0, ) if isinstance(block.MLP, MLP): mlp = block.MLP for attr in ["upsamp", "downsamp", "swiGate"]: if hasattr(mlp, attr): self._wrap_linear_with_lora( mlp, attr, adapter_name=adapter_name, rank=rank, dropout=dropout, alpha=alpha, mode="exclusive", cap_norm=1.0, global_scale=1.0, ) # MoE / shared expert 先不用挂，之后要的话再细化 def set_active_lora(self, names: list[str] | None): """ 全局设置 LoRA 激活的 adapter 列表。 - names 为 None 或 [] 表示关闭所有 LoRA。 - 多个名字配合 LoraLinear.mode="additive"/"weighted" 使用。 """ for m in self.modules(): if isinstance(m, LoraLinear): if not names: m.active = [] else: # Different adapters may target different layer ranges. # Activate only names actually registered on each wrapper. m.set_active([name for name in names if name in m.adapters]) def activate_single_lora(self, name: str | None): if name is None: self.set_active_lora([]) else: self.set_active_lora([name]) def get_lora_state_dict(self, adapter_name: str | None = None): """ 只导出 LoRA adapter 的参数（可选指定某一个 adapter）。 返回的 key 仍然是模型里的完整路径，方便直接 load_state_dict(strict=False)。 """ full = self.state_dict() out = {} for k, v in full.items(): if ".adapters." not in k: continue if adapter_name is not None and f".adapters.{adapter_name}." not in k: continue out[k] = v.detach().cpu() return out def load_lora_state_dict( self, state: dict, adapter_name: str | None = None, strict: bool = False, ): """ 只加载 LoRA adapter 的参数（和 get_lora_state_dict 对称）。 - state 通常来自 torch.load(lora_xxx.pt)，里边只有 LoRA 的 key。 - adapter_name 仅用于过滤对应 adapter；如果为 None 则加载所有 ".adapters." 的参数。 """ if adapter_name is not None: filtered = { k: v for k, v in state.items() if f".adapters.{adapter_name}." in k } else: filtered = { k: v for k, v in state.items() if ".adapters." in k } # 注意 strict=False：我们只想更新已有的 LoRA 权重，其他缺失参数无所谓 missing, unexpected = self.load_state_dict(filtered, strict=strict) if strict: if missing: print(f"[LoRA-load] missing[:10] = {missing[:10]} (total {len(missing)})") if unexpected: print(f"[LoRA-load] unexpected[:10]= {unexpected[:10]} (total {len(unexpected)})") return missing, unexpected @staticmethod def _resolve_hub_file( model_or_path: str | os.PathLike, filename: str, *, subfolder: str | None = None, revision: str | None = None, token: str | bool | None = None, ) -> Path: local = Path(model_or_path) candidate = local / subfolder / filename if subfolder else local / filename if candidate.is_file(): return candidate from huggingface_hub import hf_hub_download return Path(hf_hub_download( repo_id=str(model_or_path), filename=filename, subfolder=subfolder, revision=revision, token=token, )) def save_lora_pretrained( self, save_directory: str | os.PathLike, adapter_name: str, *, rank: int, alpha: float, dropout: float, target: str, base_model_name_or_path: str | None = None, ) -> None: """Save one native tinyLLM LoRA in a Hub-friendly directory.""" from safetensors.torch import save_file output = Path(save_directory) output.mkdir(parents=True, exist_ok=True) state = self.get_lora_state_dict(adapter_name) if not state: raise ValueError(f"LoRA adapter {adapter_name!r} is not attached") save_file({key: value.contiguous() for key, value in state.items()}, str(output / "adapter_model.safetensors")) metadata = { "format": "tinyllm_lora_v1", "adapter_name": adapter_name, "rank": int(rank), "alpha": float(alpha), "dropout": float(dropout), "target": target, "base_model_name_or_path": base_model_name_or_path, } (output / "adapter_config.json").write_text( json.dumps(metadata, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) def load_lora_pretrained( self, model_or_path: str | os.PathLike, *, subfolder: str | None = None, adapter_name: str | None = None, activate: bool = True, revision: str | None = None, token: str | bool | None = None, ) -> dict: """Attach and load a native tinyLLM LoRA from disk or the Hub.""" from safetensors.torch import load_file config_file = self._resolve_hub_file( model_or_path, "adapter_config.json", subfolder=subfolder, revision=revision, token=token, ) weights_file = self._resolve_hub_file( model_or_path, "adapter_model.safetensors", subfolder=subfolder, revision=revision, token=token, ) metadata = json.loads(config_file.read_text(encoding="utf-8")) if metadata.get("format") != "tinyllm_lora_v1": raise ValueError(f"Unsupported adapter format: {metadata.get('format')!r}") name = adapter_name or metadata["adapter_name"] source_name = metadata["adapter_name"] self.attach_lora_adapter( adapter_name=name, rank=int(metadata["rank"]), alpha=float(metadata["alpha"]), dropout=float(metadata.get("dropout", 0.0)), target=str(metadata["target"]), ) state = load_file(str(weights_file), device="cpu") if name != source_name: marker = f".adapters.{source_name}." replacement = f".adapters.{name}." state = {key.replace(marker, replacement): value for key, value in state.items()} expected = { key for key in self.state_dict() if f".adapters.{name}." in key } if set(state) != expected: raise RuntimeError( "Adapter tensor mismatch: " f"missing={sorted(expected - set(state))[:8]}, " f"unexpected={sorted(set(state) - expected)[:8]}" ) self.load_state_dict(state, strict=False) if activate: self.activate_single_lora(name) return {**metadata, "adapter_name": name, "tensor_count": len(state)} def load_vision_delta( self, model_or_path: str | os.PathLike, *, subfolder: str | None = None, revision: str | None = None, token: str | bool | None = None, ) -> dict: """Load Q-Former/projector tensors from a local or Hub delta package.""" from safetensors.torch import load_file weights_file = self._resolve_hub_file( model_or_path, "vision_delta.safetensors", subfolder=subfolder, revision=revision, token=token, ) state = load_file(str(weights_file), device="cpu") expected = set(self.state_dict()) unexpected = sorted(set(state) - expected) if unexpected: raise RuntimeError(f"Unexpected vision delta tensors: {unexpected[:8]}") self.load_state_dict(state, strict=False) return {"tensor_count": len(state), "path": str(weights_file)} def prepare_inputs_for_generation( self, input_ids: torch.LongTensor, past_key_values: Optional[Cache] = None, attention_mask: Optional[torch.LongTensor] = None, vision_feats: torch.Tensor | None = None, vision_mask: torch.Tensor | None = None, global_pos: torch.Tensor | None = None, global_off: torch.Tensor | None = None, **kwargs, ): # 增量步：只喂最后一个 token，并清掉视觉相关输入 if past_key_values is not None: input_ids = input_ids[:, -1:] if attention_mask is None: attention_mask = input_ids.new_ones(input_ids.shape, dtype=torch.long) else: attention_mask = attention_mask[:, -input_ids.size(1):] # T 对齐 vision_feats = None vision_mask = None global_pos = None global_off = None # prefill：必须把 global_pos/global_off 原样传下去，否则 forward 收不到 return { "input_ids": input_ids, "attention_mask": attention_mask, "past_key_values": past_key_values, "vision_feats": vision_feats, "vision_mask": vision_mask, "global_pos": global_pos, "global_off": global_off, "use_cache": True, } def _prepare_generation_config( self, generation_config: GenerationConfig | None = None, model_kwargs: dict | None = None, **kwargs, ): """ 自己接管 GenerationMixin 的 config 准备逻辑，避免去碰 self.config._get_non_default_generation_parameters / _from_model_config 这些我们没实现的接口。 """ # 1) 合并 model_kwargs 和裸 kwargs if model_kwargs is None: model_kwargs = {} merged_kwargs = {**model_kwargs, **kwargs} # 2) 拿一个 base GenerationConfig if generation_config is None: base = getattr(self, "generation_config", None) if isinstance(base, GenerationConfig): generation_config = copy.deepcopy(base) else: generation_config = GenerationConfig( eos_token_id=getattr(self.cfg, "eos_token_id", None), pad_token_id=getattr(self.cfg, "pad_token_id", None), ) else: # 允许传 dict / GenerationConfig / 其它 config if isinstance(generation_config, dict): generation_config = GenerationConfig(**generation_config) elif isinstance(generation_config, GenerationConfig): generation_config = copy.deepcopy(generation_config) else: # 大概率是一个 PretrainedConfig；我们不上 _from_model_config，防止再踩坑 try: generation_config = GenerationConfig.from_model_config(generation_config) except Exception: generation_config = GenerationConfig( eos_token_id=getattr(self.cfg, "eos_token_id", None), pad_token_id=getattr(self.cfg, "pad_token_id", None), ) # 3) 用 generation_config 自己的字段当作“合法生成参数”列表 try: gen_param_keys = set(generation_config.to_dict().keys()) except Exception: gen_param_keys = set(vars(generation_config).keys()) cleaned_model_kwargs: dict = {} for k, v in merged_kwargs.items(): if k in gen_param_keys: # 比如 max_new_tokens / do_sample / eos_token_id / pad_token_id / num_beams ... setattr(generation_config, k, v) else: # 留下真正要传给 forward 的，比如 input_ids / attention_mask / past_key_values cleaned_model_kwargs[k] = v return generation_config, cleaned_model_kwargs