# -*- coding: utf-8 -*- """ modeling_samai_pnet.py — SamAI-Pnet-dMoE-2B 自洽模型实现 (发布版, 单文件)。 架构 = PonderNet 自适应停机 (质量维度) + 动态可变 MoE (效率维度) + MTP (训练期吞吐)。 本文件为发行版自洽模块: 不依赖训练仓其他文件, 仅需 torch + transformers>=5.0。 用法 (trust_remote_code 一行加载): import torch from transformers import AutoModelForCausalLM, AutoTokenizer model = AutoModelForCausalLM.from_pretrained( "tchbcb/samai-pnet-dmoe-2b", trust_remote_code=True, dtype="auto").eval().cuda() tok = AutoTokenizer.from_pretrained("tchbcb/samai-pnet-dmoe-2b") 核心机制: 1. Ponder 思考块: 后 8 层循环执行, halting 头输出每步停机概率 lam_k, 概率加权混合 h = Σ w_k·h_k, 累计质量 1-eps 即停; KV 经 commit pass 写回。 2. 动态可变 MoE: FFN 按宽度均分 8 个同构专家, Router(已训练) 输出阈值激活 (p_i >= tau 激活, 激活数 1~8 自适应); 每 token 仅首轮路由一次, 循环复用。 3. MTP: 训练期辅助头, 推理 generate() 自动旁路。 4. round5-B: per-step CE 训练的 halting 头 —— hard 样本自适应"多想", easy 样本早停 (hard-easy steps 分化 +0.433)。 诊断: model.ponder_log 追加每次 forward 的 {mode, executed, steps_mean, ...}。 """ from __future__ import annotations from collections import OrderedDict from dataclasses import dataclass import torch import torch.nn.functional as F from torch import nn from torch.utils.checkpoint import checkpoint as _ckpt from transformers import AutoConfig, AutoModelForCausalLM from transformers.activations import ACT2FN from transformers.cache_utils import DynamicCache from transformers.masking_utils import create_causal_mask from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import LlamaForCausalLM, LlamaRMSNorm from transformers.utils import can_return_tuple # -------------------------------------------------------------------------- # Config # -------------------------------------------------------------------------- class PonderLlamaConfig(LlamaConfig): """在 LlamaConfig 之上追加 PonderNet 超参 (其余字段与原 Llama config 完全兼容).""" model_type = "llama_ponder" def __init__(self, **kwargs): self.ponder_start_layer = kwargs.pop("ponder_start_layer", None) # None = 最后 8 层 self.ponder_window = kwargs.pop("ponder_window", 1) # prefill 思考窗口 W self.max_ponder_steps = kwargs.pop("max_ponder_steps", 8) # 思考块最多循环步数 K self.ponder_epsilon = kwargs.pop("ponder_epsilon", 0.01) # 累计概率 1-eps 即停 self.ponder_signal = kwargs.pop("ponder_signal", "entropy") # entropy|msp|drift|learned self.ponder_all_positions = kwargs.pop("ponder_all_positions", False) # 全序列位置思考 self.commit_kv = kwargs.pop("commit_kv", True) # 用混合表示写最终 KV self.ponder_prior_p = kwargs.pop("ponder_prior_p", 0.5) # 几何先验 p_g (训练) self.ponder_loss_beta = kwargs.pop("ponder_loss_beta", 0.01) # KL 正则权重 beta self.drift_threshold = kwargs.pop("drift_threshold", 0.98) # drift 相似度阈值 # ---- Round5 三方向开关 ---- self.ponder_step_ce_weight = kwargs.pop("ponder_step_ce_weight", 0.0) # 方向②: per-step CE 权重 (0=关) self.ponder_forced_steps = kwargs.pop("ponder_forced_steps", 0) # 方向③: 强制思考步数 rollout (0=关) self.ponder_head_hidden = kwargs.pop("ponder_head_hidden", 0) # 方向①: >0 时 halting 头升级 MLP super().__init__(**kwargs) if self.ponder_start_layer is None: self.ponder_start_layer = max(0, self.num_hidden_layers - 8) @classmethod def from_llama(cls, llama_config: LlamaConfig, **ponder_kwargs) -> "PonderLlamaConfig": """从标准 LlamaConfig 派生 (保留全部架构字段, 丢弃非配置项)。 兼容两种来源: 标准 llama config.json (无 ponder 字段) 与本类保存的 llama_ponder config.json (含 ponder 字段) —— 显式 ponder_kwargs 优先覆盖。""" d = llama_config.to_dict() for k in ("model_type", "architectures", "_name_or_path", "transformers_version", "torch_dtype"): d.pop(k, None) ponder_keys = [k for k in d if "ponder" in k or k in ("commit_kv", "drift_threshold")] for k in ponder_keys: if k in ponder_kwargs: # 显式参数覆盖 config 里的旧值, 避免重复 kwargs 报错 d.pop(k) return cls(**d, **ponder_kwargs) # -------------------------------------------------------------------------- # 输出对象 # -------------------------------------------------------------------------- @dataclass class PonderCausalLMOutputWithPast(CausalLMOutputWithPast): """额外带出 ponder 诊断信息.""" ponder_weights: torch.Tensor | None = None # (B, W, K) 每步停机质量分布 ponder_steps: torch.Tensor | None = None # (B, W) 每位置实际思考步数 ponder_kl: torch.Tensor | None = None # 标量, KL(w‖Geometric) 正则项 (labels 时) ponder_step_ce: torch.Tensor | None = None # 标量, per-step CE (方向②, 训练且权重>0 时) # -------------------------------------------------------------------------- # Model # -------------------------------------------------------------------------- class PonderLlamaForCausalLM(LlamaForCausalLM): """ 用法 A — 直接加载 MiniCPM5-2B 权重 (键名与 LlamaForCausalLM 完全一致): from modeling_samai_pnet import PonderLlamaForCausalLM, PonderLlamaConfig base = LlamaConfig.from_pretrained("tchbcb/MiniCPM5-2B-cpu") cfg = PonderLlamaConfig.from_llama(base) model = PonderLlamaForCausalLM.from_pretrained( "tchbcb/MiniCPM5-2B-cpu", config=cfg, torch_dtype="bfloat16") # missing keys 只有 ponder_head (仅 signal="learned" 时参与前向, 其余信号忽略) 用法 B — 一行加载: model = PonderLlamaForCausalLM.from_ponder("tchbcb/MiniCPM5-2B-cpu") generate() / chat 流程与 HF 标准完全一致; model.ponder_log 记录每次 forward 的思考步数。 """ def __init__(self, config: PonderLlamaConfig): super().__init__(config) cfg = self.config # 真 PonderNet 停机头; bias=-4 使初始 lam≈0.018 (初始倾向于多想几步) # 方向①: ponder_head_hidden>0 时升级为 MLP (Linear->GELU->Linear), # 打破线性可分上限 —— round4 结论: 线性头在冻结特征上已到极限。 # 数值初始化由下方 _init_weights 覆写统一处理: post_init 重遍历会把 # Linear 重置为 Llama 默认 (normal+零偏置), 自定义值必须写在那里才能存活 head_hidden = int(getattr(cfg, "ponder_head_hidden", 0) or 0) if head_hidden > 0: self.ponder_head = nn.Sequential( nn.Linear(cfg.hidden_size, head_hidden), nn.GELU(), nn.Linear(head_hidden, 1, bias=True), ) else: self.ponder_head = nn.Linear(cfg.hidden_size, 1, bias=True) self._ponder_log: list[dict] = [] # 诊断日志 (每次 forward 追加, 最多保留 4096 条) self.post_init() def _init_weights(self, module): """PonderNet 停机头自定义初始化 (覆盖 Llama 默认), 其余模块遵循 Llama 惯例。 post_init()/from_pretrained 的 init_weights 遍历都会调用本钩子, 因此自定义值不会被 Llama 默认的 normal+零偏置覆写。""" head = getattr(self, "ponder_head", None) if isinstance(module, nn.Linear) and head is not None: head_linears = [mm for mm in head.modules() if isinstance(mm, nn.Linear)] if any(module is mm for mm in head_linears): nn.init.zeros_(module.weight) nn.init.zeros_(module.bias) if module is head_linears[-1]: # 尾层 bias=-4 -> 初始 lam=sigmoid(-4)≈0.018, 倾向多思考 nn.init.constant_(module.bias, -4.0) return super()._init_weights(module) @classmethod def from_ponder(cls, pretrained_path: str, ponder_kwargs: dict | None = None, **kw): """从标准 Llama 权重目录 / HF repo 一行构造 Ponder 模型.""" base = LlamaConfig.from_pretrained(pretrained_path) cfg = PonderLlamaConfig.from_llama(base, **(ponder_kwargs or {})) return cls.from_pretrained(pretrained_path, config=cfg, **kw) # ------------------------------------------------------------------ # # 内部工具 # ------------------------------------------------------------------ # @property def _ponder_layers(self): s = self.config.ponder_start_layer return self.model.layers[s: self.config.num_hidden_layers] def _crop_ponder_kv(self, cache, n_remove: int): """只回滚思考块层的 KV 写入 (前段层的 KV 不许动).""" if cache is None or n_remove <= 0: return s = self.config.ponder_start_layer for li in range(s, self.config.num_hidden_layers): cache.layers[li].crop(-n_remove) def _run_block(self, h, attention_mask, position_ids, position_embeddings, cache, kwargs): """跑一遍思考块 (最后若干层). 训练且开启 gradient checkpointing 时逐层重算省激活 —— 思考块最多循环 K 次, 激活占用是普通前向的 K 倍, checkpointing 是小显存卡 (如 T4 16GB) 的关键.""" ckpt_fn = getattr(self.model, "_gradient_checkpointing_func", None) use_ckpt = (self.training and cache is None and getattr(self.model, "gradient_checkpointing", False) and ckpt_fn is not None) for layer in self._ponder_layers: if use_ckpt: h = ckpt_fn( layer.__call__, h, attention_mask=attention_mask, position_ids=position_ids, position_embeddings=position_embeddings, past_key_values=None, use_cache=False, **kwargs, ) else: h = layer( h, attention_mask=attention_mask, position_ids=position_ids, position_embeddings=position_embeddings, past_key_values=cache, use_cache=cache is not None, **kwargs, ) return h def _halting_prob(self, h_window, h_prev_window) -> torch.Tensor: """返回 (B, W) 的停机概率 lam ∈ (0, 1).""" cfg = self.config B, W = h_window.shape[:2] if cfg.ponder_signal == "learned": # ponder_head 恒为 fp32 (数值稳定), 输入 cast 到其参数 dtype, # 否则半精度模型 (fp16/bf16) 下会报 dtype 不匹配错误。 # head_ref 写法兼容 Linear 与 Sequential (方向① MLP 头) 两种结构 head_ref = next(self.ponder_head.parameters()) lam = torch.sigmoid( self.ponder_head(h_window.to(head_ref.dtype)) ).squeeze(-1) elif cfg.ponder_signal in ("entropy", "msp"): logits = self.lm_head(self.model.norm(h_window)).float() probs = F.softmax(logits, dim=-1) if cfg.ponder_signal == "msp": lam = probs.max(dim=-1).values else: # entropy: 自信 -> 熵小 -> lam 大 ent = -(probs * torch.log(probs.clamp_min(1e-9))).sum(-1) lam = 1.0 - ent / torch.log(torch.tensor(float(logits.size(-1)), device=logits.device)) elif cfg.ponder_signal == "drift": if h_prev_window is None: # 第一步没有参照, 强制继续思考 lam = torch.zeros(B, W, device=h_window.device, dtype=h_window.dtype) else: cos = F.cosine_similarity(h_window.float(), h_prev_window.float(), dim=-1) tau = cfg.drift_threshold lam = (cos - tau) / (1.0 - tau + 1e-6) # cos 越接近 1 (不再变化) 越停 else: raise ValueError(f"unknown ponder_signal: {cfg.ponder_signal}") # 统一 cast 回隐状态 dtype: 保证 remaining/mix/w_stack 混合链路 dtype 一致, # 否则 fp32 的 lam 会把整条链路升成 fp32, 之后 norm/lm_head (half 权重) 会崩 return lam.to(h_window.dtype).clamp(1e-6, 1.0 - 1e-6) def _step_ce(self, h_step, labels): """方向②: 单步表示的独立解码 CE。 让每一步思考表示都"可用"(可被 lm_head 解码), 回收 mix 摊薄导致的 CE 代价。 训练时用 gradient checkpointing 重算, 避免大词表 logits 常驻显存。""" def _fn(h_): lg = self.lm_head(self.model.norm(h_)) lg = lg[..., :-1, :].reshape(-1, lg.size(-1)).float() tg = labels[..., 1:].reshape(-1) return F.cross_entropy(lg, tg, ignore_index=-100) if self.training and torch.is_grad_enabled() and h_step.requires_grad: return _ckpt(_fn, h_step, use_reentrant=False) return _fn(h_step) def _win_mask(self, W, past_len, device, dtype, attention_mask=None): """窗口迭代的加性因果 mask: keys = [0..past_len+W), queries = 窗口 W 个位置。 必须手工构建而不能用 create_causal_mask: 后者从非思考层读取缓存长度, 而思考块层已被 crop 到 past_len, 会得到多 W 的错误 kv 尺寸。 返回 float 加性 mask (sdpa / eager 通用)。""" total = past_len + W q_idx = torch.arange(past_len, total, device=device) k_idx = torch.arange(total, device=device) allowed = q_idx.unsqueeze(-1) >= k_idx.unsqueeze(0) # (W, total) if attention_mask is not None: allowed = allowed.unsqueeze(0) & attention_mask[:, :total].to(torch.bool).unsqueeze(1) m = torch.zeros(allowed.size(0), 1, W, total, dtype=dtype, device=device) m = m.masked_fill(~allowed.unsqueeze(1), torch.finfo(dtype).min) # (B,1,W,total) else: m = torch.zeros(1, 1, W, total, dtype=dtype, device=device) m = m.masked_fill(~allowed.unsqueeze(0).unsqueeze(0), torch.finfo(dtype).min) return m def _log(self, entry: dict): self._ponder_log.append(entry) if len(self._ponder_log) > 4096: del self._ponder_log[: len(self._ponder_log) - 4096] # ------------------------------------------------------------------ # # 主前向 # ------------------------------------------------------------------ # @can_return_tuple def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: DynamicCache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, logits_to_keep: int | torch.Tensor = 0, output_ponder: bool = False, **kwargs, ) -> PonderCausalLMOutputWithPast: cfg = self.config m = self.model # ---- 输入 & 位置 ------------------------------------------------ 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 = m.embed_tokens(input_ids) B, Q = inputs_embeds.shape[:2] if use_cache is None: use_cache = False if self.training else cfg.use_cache if use_cache and past_key_values is None: past_key_values = DynamicCache(config=cfg) cache_position = kwargs.pop("cache_position", None) past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0 if position_ids is None: base = (cache_position if cache_position is not None else torch.arange(Q, device=inputs_embeds.device) + past_seen) position_ids = base.reshape(1, -1).expand(B, -1)[:, :Q].to(inputs_embeds.device) # ---- mask / rope (整次 forward 内复用) --------------------------- if cache_position is None: cache_position = torch.arange(Q, device=inputs_embeds.device) + past_seen causal_mask = create_causal_mask( config=cfg, input_embeds=inputs_embeds, attention_mask=attention_mask, cache_position=cache_position, past_key_values=past_key_values, position_ids=position_ids, ) position_embeddings = m.rotary_emb(inputs_embeds, position_ids=position_ids) # ---- Stage A: 前段层正常单遍 ------------------------------------- h = inputs_embeds for layer in m.layers[: cfg.ponder_start_layer]: h = layer( h, attention_mask=causal_mask, position_ids=position_ids, position_embeddings=position_embeddings, past_key_values=past_key_values, use_cache=use_cache, **kwargs, ) # ---- 思考窗口 W -------------------------------------------------- # decode (Q==1) 恒为 1; 训练/显式开启时对全部位置; 其余 prefill 思考末尾 W 个位置 if Q == 1 or self.training or cfg.ponder_all_positions: W = Q else: W = max(1, min(cfg.ponder_window, Q)) eps = cfg.ponder_epsilon K = cfg.max_ponder_steps device, dtype = inputs_embeds.device, inputs_embeds.dtype remaining = torch.ones(B, W, device=device, dtype=dtype) cum = torch.zeros(B, W, device=device, dtype=dtype) active = torch.ones(B, W, device=device, dtype=torch.bool) mix = torch.zeros(B, W, h.size(-1), device=device, dtype=dtype) w_list: list[torch.Tensor] = [] h_prev_win = None h_full = None # 第 1 次迭代的全序列输出 (W0 强制 rollout step_ce_weight = float(getattr(cfg, "ponder_step_ce_weight", 0.0) or 0.0) use_step_ce = (step_ce_weight > 0.0 and labels is not None and self.training) step_pairs: list[tuple[torch.Tensor, torch.Tensor]] = [] # [(w_k_detached, ce_k), ...] loop_K = forced if forced > 0 else K for k in range(1, loop_K + 1): executed = k if k > 1: # 替换式写入: 先回滚上一次迭代在思考块层的 KV, 本轮 update 再补回, # 使任意时刻 cache 长度恒等于真实 token 数 (对 W==Q 同样必需!) self._crop_ponder_kv(past_key_values, W) # 仅当有 cache 且 W 1) and (past_key_values is not None) and (W < Q) if not use_window: h = self._run_block(h, causal_mask, position_ids, position_embeddings, past_key_values, kwargs) h_win = h if W == Q else h[:, -W:, :] if W < Q: h_full = h else: if not win_mask_built: past_len = int(past_key_values.layers[cfg.ponder_start_layer].get_seq_length()) mask_win = self._win_mask(W, past_len, device, dtype, attention_mask) pos_emb_win = (position_embeddings[0][:, -W:], position_embeddings[1][:, -W:]) win_mask_built = True h_win = self._run_block(h_win, mask_win, position_ids[:, -W:], pos_emb_win, past_key_values, kwargs) if forced > 0: # 方向③: 强制步数 rollout —— 完全绕过 halting 头与 w 混合, # 最终表示即第 forced 步隐状态 (KV 恰为其 cache, 无需 commit) continue lam = self._halting_prob(h_win, h_prev_win) # (B, W) lam_eff = lam * active.to(lam.dtype) w_k = remaining * lam_eff mix = mix + w_k.unsqueeze(-1) * h_win w_list.append(w_k) remaining = remaining * (1.0 - lam_eff) cum = cum + w_k active = active & (cum < 1.0 - eps) h_prev_win = h_win if use_step_ce: # 方向②: 本步表示单独算解码 CE。权重 detach —— 梯度只流向表示, # 不经此损失推动 w (w 由 step-supervise / KL 单独管辖) ce_k = self._step_ce(h_win, labels) step_pairs.append((w_k.detach(), ce_k)) if not bool(active.any()): break if forced > 0: # ---- 方向③: 最终表示 = 第 n 步隐状态 (非概率混合) -------------- if W < Q: h_final = h_full.clone() h_final[:, -W:, :] = h_win else: h_final = h_win w_stack = F.one_hot( torch.full((B, W), forced - 1, device=device), forced ).to(dtype) # (B, W, n) 诊断用 else: # 剩余概率质量全部记给最后一步 (PonderNet 的 p_N 截断) mix = mix + remaining.unsqueeze(-1) * h_win w_stack = torch.stack(w_list, dim=-1).to(dtype) # (B, W, K_exec) w_stack = torch.cat([w_stack[..., :-1], w_stack[..., -1:] + remaining.unsqueeze(-1)], dim=-1) # ---- 组装最终 hidden -------------------------------------------- if W < Q: h_final = h_full h_final[:, -W:, :] = mix else: h_final = mix # ---- commit pass: 用"思考后"的混合表示写最终 KV ------------------ if (use_cache and cfg.commit_kv and past_key_values is not None and executed > 1 and forced == 0): self._crop_ponder_kv(past_key_values, W) if not win_mask_built: # W==Q 的场景 self._run_block(h_final, causal_mask, position_ids, position_embeddings, past_key_values, kwargs) else: self._run_block(h_final[:, -W:, :], mask_win, position_ids[:, -W:], pos_emb_win, past_key_values, kwargs) # commit 输出丢弃, KV 由 mix 重新计算 —— 后续 token 看到的就是思考后的表示 # ---- logits / loss ---------------------------------------------- slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(m.norm(h_final[:, slice_indices, :])) loss, ponder_kl, ponder_step_ce = None, None, None if labels is not None: ce_logits = logits.float() loss = F.cross_entropy( ce_logits[..., :-1, :].reshape(-1, ce_logits.size(-1)), labels[..., 1:].reshape(-1), ignore_index=-100) # PonderNet 正则: KL(w || Geometric(p_g)), 截断到 K 步并归一化。 # 必须升到 fp32 再取 log: fp16 下 p=0 处 clamp_min(1e-8) 仍会下溢为 0, # log(0)=-inf 与 p=0 相乘直接得到 NaN, 训练瞬间发散 p = (w_stack / w_stack.sum(-1, keepdim=True).clamp_min(1e-8)).float() # (B,W,K) Kd = p.size(-1) n_idx = torch.arange(1, Kd + 1, device=p.device, dtype=p.dtype) pg = cfg.ponder_prior_p geom = (1 - pg) ** (n_idx - 1) * pg geom = torch.cat([geom[..., :-1], ((1 - pg) ** (Kd - 1) * torch.ones_like(geom[..., -1:]))], dim=-1) geom = geom / geom.sum() kl = (p * (torch.log(p.clamp_min(1e-8)) - torch.log(geom.expand_as(p)))).sum(-1) ponder_kl = kl.mean() loss = loss + cfg.ponder_loss_beta * ponder_kl # 方向②: per-step CE 聚合 —— 每步 CE 按该步 w 质量 (detach) 加权, # 让被分配到概率质量的思考步都保持"可解码", 回收 mix 摊薄的 CE 代价 if use_step_ce and step_pairs: ws = torch.stack([p[0] for p in step_pairs], dim=0) # (K,B,W) ces = torch.stack([p[1] for p in step_pairs], dim=0) # (K,) wsum = ws.sum(0).clamp_min(1e-8) # (B,W) per_pos = (ws * ces.view(-1, 1, 1)).sum(0) / wsum # (B,W) ponder_step_ce = per_pos.mean() loss = loss + step_ce_weight * ponder_step_ce # ---- 诊断 -------------------------------------------------------- cum_w = torch.cumsum(w_stack, dim=-1) steps_used = ((cum_w < 1.0 - eps).sum(-1) + 1).clamp(max=w_stack.size(-1)) # (B, W) self._log({ "mode": "train" if self.training else ("decode" if Q == 1 else "prefill"), "signal": cfg.ponder_signal, "executed": executed, "steps_max": int(steps_used.max()), "steps_mean": float(steps_used.float().mean()), }) return PonderCausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=past_key_values, ponder_weights=w_stack if output_ponder else None, ponder_steps=steps_used if output_ponder else None, ponder_kl=ponder_kl, ponder_step_ce=ponder_step_ce, ) # -------------------------------------------------------------------------- # AutoClass 注册 (可重复 import 安全) # -------------------------------------------------------------------------- try: AutoConfig.register("llama_ponder", PonderLlamaConfig) AutoModelForCausalLM.register(PonderLlamaConfig, PonderLlamaForCausalLM) except Exception: pass # -------------------------------------------------------------------------- # Config # -------------------------------------------------------------------------- class PonderMoELlamaConfig(PonderLlamaConfig): """PonderLlamaConfig + 动态 MoE / MTP 超参. moe_num_experts=0 且 mtp_num_tokens=0 时行为与父类完全一致 (回归安全).""" model_type = "llama_ponder_moe" def __init__(self, **kwargs): # ---- 动态可变 MoE ---- self.moe_num_experts = int(kwargs.pop("moe_num_experts", 0) or 0) # 0=关 self.moe_route_tau = kwargs.pop("moe_route_tau", None) # 激活阈值, None -> 1/N self.moe_aux_weight = float(kwargs.pop("moe_aux_weight", 0.01)) # 负载均衡权重 self.moe_scope = kwargs.pop("moe_scope", "ponder") # ponder|all: 拆哪些层的 FFN # ---- MTP ---- self.mtp_num_tokens = int(kwargs.pop("mtp_num_tokens", 0) or 0) # 深度 D, 0=关 self.mtp_loss_weight = float(kwargs.pop("mtp_loss_weight", 1.0)) super().__init__(**kwargs) # 普通属性而非 property (transformers 5.16 config __getattribute__ 钩子不支持 property) self.moe_tau = (float(self.moe_route_tau) if self.moe_route_tau else 1.0 / max(1, self.moe_num_experts)) @classmethod def from_llama(cls, llama_config: LlamaConfig, **kwargs) -> "PonderMoELlamaConfig": """同父类, 但额外吞掉 moe_*/mtp_* 旧键, 兼容本类保存的 config.json.""" d = llama_config.to_dict() for k in ("model_type", "architectures", "_name_or_path", "transformers_version", "torch_dtype"): d.pop(k, None) drop = [k for k in d if "ponder" in k or k.startswith(("moe_", "mtp_")) or k in ("commit_kv", "drift_threshold")] for k in drop: if k in kwargs: d.pop(k) return cls(**d, **kwargs) @dataclass class PonderMoECausalLMOutputWithPast(PonderCausalLMOutputWithPast): """追加 MoE / MTP 诊断字段.""" moe_aux_loss: torch.Tensor | None = None # 标量, 负载均衡损失 (各 MoE 层均值) moe_active_mean: torch.Tensor | None = None # 标量, 平均激活专家数 ∈ [1, N] mtp_loss: torch.Tensor | None = None # 标量, 各深度 MTP CE 均值 mtp_acc: torch.Tensor | None = None # 标量, MTP top-1 命中率 (诊断) # -------------------------------------------------------------------------- # 同构专家: 原 FFN 的 1/N 宽切片 (gate/up 按行切, down 按列切), 参数守恒 # -------------------------------------------------------------------------- class ExpertMLP(nn.Module): def __init__(self, hidden_size: int, inter_size: int, act_fn): super().__init__() self.gate_proj = nn.Linear(hidden_size, inter_size, bias=False) self.up_proj = nn.Linear(hidden_size, inter_size, bias=False) self.down_proj = nn.Linear(inter_size, hidden_size, bias=False) self.act_fn = act_fn def forward(self, x): return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) # -------------------------------------------------------------------------- # 动态可变 MoE 层: Router(零初始化) + N 同构专家 + 阈值动态激活 + STE + 稀疏执行 # -------------------------------------------------------------------------- class PonderMoEMLP(nn.Module): def __init__(self, hidden_size: int, n_experts: int, tau: float, act_fn): super().__init__() self.n_experts = n_experts self.tau = float(tau) self.router = nn.Linear(hidden_size, n_experts, bias=True) nn.init.zeros_(self.router.weight) nn.init.zeros_(self.router.bias) # 零路由 -> p=1/N 均匀 -> 初始全激活 = 稠密等价 self.experts = nn.ModuleList([]) self.act_fn_holder = act_fn # 仅占位保存引用 (无参数) # 运行期状态 (诊断/冻结路由) self._last_mask = None # (B, T, N) bool, 最近一次 fresh 路由 self._last_p = None # (B, T, N) float, detach 后的路由概率 self.last_aux = None # 标量 tensor (fresh 时更新, 模型聚合) self.last_active = None # 标量 tensor, 平均激活专家数 self.n_route_calls = 0 # 结构性验证: Router 实际前向次数 self.cache_hook = None # SSD-LRU 挂钩: callable(expert_ids) self.cache_post_hook = None # SSD-LRU 执行后挂钩: 强制容量硬上界 # ---------------- 构造 ---------------- # @classmethod def from_dense(cls, dense, n_experts: int, tau: float | None = None) -> "PonderMoEMLP": """把稠密 LlamaMLP 拆成 n_experts 个同构专家 (参数守恒) + 零 Router.""" H, I = dense.hidden_size, dense.intermediate_size assert I % n_experts == 0, f"intermediate_size({I}) 必须被 N({n_experts}) 整除" m = cls(H, n_experts, tau if tau is not None else 1.0 / n_experts, dense.act_fn) sub = I // n_experts with torch.no_grad(): for i in range(n_experts): ex = ExpertMLP(H, sub, dense.act_fn) sl = slice(i * sub, (i + 1) * sub) ex.gate_proj.weight.copy_(dense.gate_proj.weight[sl]) ex.up_proj.weight.copy_(dense.up_proj.weight[sl]) ex.down_proj.weight.copy_(dense.down_proj.weight[:, sl]) m.experts.append(ex) return m @property def device(self): return self.router.weight.device # ---------------- 路由 ---------------- # def _route(self, x_flat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: """fresh 路由: 返回 (p, mask), 形状 (T, N). 零初始化下 p=1/N -> 全激活.""" logits = self.router(x_flat) p = torch.softmax(logits.float(), dim=-1).to(x_flat.dtype) mask = p >= (self.tau - 1e-6) empty = ~mask.any(-1) if bool(empty.any()): # 无一过阈 -> top-1 保底 top = p.argmax(-1) mask = mask.clone() mask[empty, top[empty]] = True return p, mask def _experts_sparse(self, x_flat, mask, p, with_ste: bool) -> torch.Tensor: """稀疏执行: 只计算被激活的 (token, expert) 对. 前向: 激活专家求和 (权重 1, 激活参数量语义) 反向: STE —— 经 p 软通路把梯度送回 Router (with_ste=True 时).""" out = torch.zeros_like(x_flat) dtype = x_flat.dtype active_ids = [i for i in range(len(self.experts)) if bool(mask[:, i].any())] if (not self.training) and self.cache_hook is not None and active_ids: self.cache_hook(active_ids) # SSD-LRU: 先搬入后执行 (GPU 上必须前置) for i in active_ids: ex = self.experts[i] sel = mask[:, i].nonzero(as_tuple=True)[0] e = ex(x_flat.index_select(0, sel)) if with_ste: w = p.index_select(0, sel).select(1, i) # (S,) w = (w - w.detach() + 1.0).to(dtype) # 前向=1, 反向 d/dp=1 else: w = torch.ones((), device=e.device, dtype=dtype) out = out.index_add(0, sel, e * w.unsqueeze(-1)) if (not self.training) and self.cache_post_hook is not None and active_ids: self.cache_post_hook(active_ids) # 执行完毕: 超容部分按 LRU 驱逐 (含活跃) return out # ---------------- 前向 ---------------- # def forward(self, x: torch.Tensor, route_mode: str = "fresh") -> torch.Tensor: B, T, H = x.shape x_flat = x.reshape(B * T, H) if route_mode == "all": # 调试/等价验证: 全专家权重 1 if (not self.training) and self.cache_hook is not None: self.cache_hook(list(range(len(self.experts)))) out = torch.zeros_like(x_flat) for ex in self.experts: out = out + ex(x_flat) if (not self.training) and self.cache_post_hook is not None: self.cache_post_hook(list(range(len(self.experts)))) return out.view(B, T, H) if route_mode == "frozen" and self._last_mask is not None: # 单次路由铁律: 复用最近一次 fresh 路由 (不触碰 Router, 不产生新路由) mask = self._last_mask[:, -T:, :].reshape(B * T, -1) out = self._experts_sparse(x_flat, mask, None, with_ste=False) return out.view(B, T, H) # ---- fresh 路由 ---- self.n_route_calls += 1 p, mask = self._route(x_flat) self._last_mask = mask.detach().view(B, T, -1) self._last_p = p.detach().view(B, T, -1) out = self._experts_sparse(x_flat, mask, p, with_ste=self.training or p.requires_grad) # 负载均衡 (Switch 式): aux = N * Σ f_i * P_i with torch.no_grad(): f = mask.float().mean(0) # (N,) 激活频率 P = p.float().mean(0) # (N,) 平均概率 self.last_aux = self.n_experts * (f * P).sum() self.last_active = mask.float().sum(-1).mean() return out.view(B, T, H) # -------------------------------------------------------------------------- # 自定义 DecoderLayer: 与 LlamaDecoderLayer 完全同构, 仅把 mlp 换成 PonderMoEMLP # 并支持 moe_route_mode 透传 (fresh=新路由 / frozen=复用 / all=全激活) # -------------------------------------------------------------------------- class PonderMoEDecoderLayer(nn.Module): def __init__(self, config, layer_idx: int | None, moe_mlp: PonderMoEMLP): super().__init__() self.hidden_size = config.hidden_size self.layer_idx = layer_idx # 注意: 只注册 self.mlp 一个名字 (注册两个会产出重复 state_dict 键) self.mlp = moe_mlp @classmethod def convert(cls, old_layer: nn.Module, config, n_experts: int, tau: float) -> "PonderMoEDecoderLayer": """接管一个现成 LlamaDecoderLayer 的全部子模块, 仅替换 mlp.""" new = cls(config, getattr(old_layer, "layer_idx", None), PonderMoEMLP.from_dense(old_layer.mlp, n_experts, tau)) new.self_attn = old_layer.self_attn new.input_layernorm = old_layer.input_layernorm new.post_attention_layernorm = old_layer.post_attention_layernorm return new def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values=None, use_cache: bool = False, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, moe_route_mode: str = "fresh", **kwargs, ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) 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, position_embeddings=position_embeddings, **kwargs, ) hidden_states = residual + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states, route_mode=moe_route_mode) hidden_states = residual + hidden_states return hidden_states # -------------------------------------------------------------------------- # MTP 多 Token 预测头 (宪法 §三.4, 训练期增强; 推理自动旁路) # 深度 k: 输入 [h_t, e_{t+k}] -> 预测 token t+k+1 (共享主模型 lm_head) # -------------------------------------------------------------------------- class PonderMTPHead(nn.Module): def __init__(self, config, depth: int): super().__init__() H = config.hidden_size self.depth = depth # 该头负责预测 t+depth+1 self.norm_h = LlamaRMSNorm(H, eps=config.rms_norm_eps) self.norm_e = LlamaRMSNorm(H, eps=config.rms_norm_eps) self.proj = nn.Linear(2 * H, H, bias=False) self.act = nn.GELU() self.out = nn.Linear(H, H, bias=False) self.norm_out = LlamaRMSNorm(H, eps=config.rms_norm_eps) def forward(self, h: torch.Tensor, e: torch.Tensor) -> torch.Tensor: """h, e: (B, L', H) 已对齐 —— 位置 t 处为 h_t 与 e_{t+depth}.""" x = self.proj(torch.cat([self.norm_h(h), self.norm_e(e)], dim=-1)) return self.norm_out(self.out(self.act(x))) # -------------------------------------------------------------------------- # SSD-LRU 专家缓存调度原型 (宪法 §二.4 推理策略层) # -------------------------------------------------------------------------- class SSDLRUExpertManager: """专家权重驻留 storage (默认 CPU, 模拟 SSD), 被路由激活时搬入计算设备, 容量满按 LRU 驱逐. 统计验证"单次路由 => 每 token 每层恰 1 次 miss, Ponder K 轮迭代与 commit 全命中" —— 缓存不击穿. 仅推理期可用 (训练期交换 .data 会破坏 autograd, install 时自动排除).""" def __init__(self, model: "PonderMoELlamaForCausalLM", capacity_per_layer: int | None = None, storage_device: str = "cpu"): self.model = model self.storage_device = storage_device self.device = next(model.parameters()).device # 捕获主设备 (搬移前) self.layers: dict[int, PonderMoEMLP] = {} self.order: dict[int, OrderedDict] = {} self.capacity: dict[int, int] = {} self.stats: dict[int, dict] = {} cfg = model.config N = int(cfg.moe_num_experts) for li, layer in enumerate(model.model.layers): if isinstance(layer, PonderMoEDecoderLayer): self.layers[li] = layer.mlp self.order[li] = OrderedDict() self.capacity[li] = int(capacity_per_layer or N) self.stats[li] = {"hit": 0, "miss": 0, "evict": 0, "fetch": 0, "calls": 0, "miss_calls": 0} @classmethod def install(cls, model, capacity_per_layer: int | None = None, storage_device: str = "cpu") -> "SSDLRUExpertManager": assert not model.training, "SSD-LRU 仅用于推理 (训练期禁止权重搬运)" mgr = cls(model, capacity_per_layer, storage_device) for li, ml in mgr.layers.items(): ml.cache_hook = mgr._make_hook(li) ml.cache_post_hook = mgr._make_post_hook(li) ml.experts.to(storage_device) # 全部专家下放 storage return mgr def _make_hook(self, li: int): def hook(active_ids: list[int]): st, order = self.stats[li], self.order[li] st["calls"] += 1 # 访问批次 (含纯命中批) had_miss = False # 第一遍: 本次活跃专家全部搬入计算设备 (先保证执行正确性) for eid in active_ids: if eid in order: st["hit"] += 1 order.move_to_end(eid) else: st["miss"] += 1 had_miss = True order[eid] = True self._move(li, eid, self.device) st["fetch"] += 1 # 第二遍: 超容时按 LRU 驱逐, 但绝不驱逐本次活跃专家 # (工作集 > 容量时允许缓存暂时超容, 执行正确性优先于驱逐压力) excess = len(order) - self.capacity[li] for k in list(order.keys()): if excess <= 0: break if k in active_ids: continue order.pop(k) self._move(li, k, self.storage_device) st["evict"] += 1 excess -= 1 if had_miss: st["miss_calls"] += 1 # 击穿判据: 实际产生 IO 的批次 return hook def _make_post_hook(self, li: int): def post_hook(active_ids: list[int]): st, order = self.stats[li], self.order[li] # 物理容量硬上界: 执行完毕后超容部分按 LRU 驱逐 # (此时专家已用完, 驱逐含本次活跃专家也安全; 欠配 => 下批重新 miss = 真实 thrash) while len(order) > self.capacity[li]: old, _ = order.popitem(last=False) self._move(li, old, self.storage_device) st["evict"] += 1 return post_hook def _move(self, li: int, eid: int, device: str): ex = self.layers[li].experts[eid] for p in ex.parameters(): p.data = p.data.to(device) def summary(self) -> dict: hit = sum(s["hit"] for s in self.stats.values()) miss = sum(s["miss"] for s in self.stats.values()) return { "per_layer": {li: dict(s) for li, s in self.stats.items()}, "hit": hit, "miss": miss, "evict": sum(s["evict"] for s in self.stats.values()), "fetch": sum(s["fetch"] for s in self.stats.values()), "calls": sum(s["calls"] for s in self.stats.values()), "miss_calls": sum(s["miss_calls"] for s in self.stats.values()), "hit_rate": hit / max(1, hit + miss), "capacity_per_layer": dict(self.capacity), } def remove(self): """卸载挂钩并把专家搬回主设备 (恢复常驻).""" for li, ml in self.layers.items(): ml.cache_hook = None ml.cache_post_hook = None ml.experts.to(self.device) # -------------------------------------------------------------------------- # Model: Ponder (质量) + 动态 MoE (效率) + MTP (吞吐, 训练期) # -------------------------------------------------------------------------- class PonderMoELlamaForCausalLM(PonderLlamaForCausalLM): """ 用法 — 从稠密基座一行构造 (FFN 自动拆专家): base = LlamaConfig.from_pretrained("tchbcb/MiniCPM5-2B-cpu") cfg = PonderMoELlamaConfig.from_llama( base, moe_num_experts=8, mtp_num_tokens=2, ponder_signal="learned") model = PonderMoELlamaForCausalLM.from_pretrained(..., config=cfg) moe_num_experts=0 且 mtp_num_tokens=0 时与 PonderLlamaForCausalLM 行为一致。 """ def __init__(self, config: PonderMoELlamaConfig): super().__init__(config) cfg = self.config self._moe_ready = False if int(getattr(cfg, "moe_num_experts", 0) or 0) > 0: self._convert_ffn_to_moe() # 在 post_init 之后: Router 零初始化不被覆盖 D = int(getattr(cfg, "mtp_num_tokens", 0) or 0) self.mtp_heads = (nn.ModuleList([PonderMTPHead(cfg, k + 1) for k in range(D)]) if D > 0 else None) self._moe_ready = True # ------------------------------------------------------------------ # # 结构改造 # ------------------------------------------------------------------ # def _convert_ffn_to_moe(self): cfg = self.config N = int(cfg.moe_num_experts) assert cfg.intermediate_size % N == 0, \ f"intermediate_size({cfg.intermediate_size}) 须被 moe_num_experts({N}) 整除" tau = cfg.moe_tau s = cfg.ponder_start_layer ids = (range(s, cfg.num_hidden_layers) if cfg.moe_scope == "ponder" else range(cfg.num_hidden_layers)) for li in ids: self.model.layers[li] = PonderMoEDecoderLayer.convert( self.model.layers[li], cfg, N, tau) def _moe_mlps(self) -> list[PonderMoEMLP]: return ([l.mlp for l in self.model.layers if isinstance(l, PonderMoEDecoderLayer)] if self._moe_ready else []) def _reset_moe_stats(self): for ml in self._moe_mlps(): ml.last_aux = None ml.last_active = None def _collect_moe_stats(self): auxs = [ml.last_aux for ml in self._moe_mlps() if ml.last_aux is not None] acts = [ml.last_active for ml in self._moe_mlps() if ml.last_active is not None] aux = torch.stack(auxs).mean() if auxs else None act = torch.stack(acts).mean() if acts else None return aux, act # ------------------------------------------------------------------ # # MTP (训练/评估且给定 labels 时; generate 自动旁路) # ------------------------------------------------------------------ # def _mtp_forward(self, h_final, input_ids, labels): m = self.model B, L, _ = h_final.shape emb = m.embed_tokens(input_ids) losses, accs = [], [] for head in self.mtp_heads: k = head.depth if L - k < 2: continue h = h_final[:, :L - k, :] # h_t, t=0..L-k-1 e = emb[:, k:, :] # e_{t+k} hh = head(h, e) # (B, L-k, H) lg = self.lm_head(hh)[:, :-1, :] # 预测 t+k+1 tg = labels[:, k + 1:] ce = F.cross_entropy(lg.reshape(-1, lg.size(-1)).float(), tg.reshape(-1), ignore_index=-100) if torch.isfinite(ce): losses.append(ce) with torch.no_grad(): valid = tg != -100 if bool(valid.any()): accs.append((((lg.argmax(-1) == tg) & valid).sum().float() / valid.sum()).detach()) mtp_loss = torch.stack(losses).mean() if losses else None mtp_acc = torch.stack(accs).mean() if accs else None return mtp_loss, mtp_acc # ------------------------------------------------------------------ # # 思考块执行 (带路由模式: 第 1 轮 fresh, 之后 frozen —— 单次路由铁律) # ------------------------------------------------------------------ # def _run_block(self, h, attention_mask, position_ids, position_embeddings, cache, kwargs, route_mode: str = "fresh"): ckpt_fn = getattr(self.model, "_gradient_checkpointing_func", None) use_ckpt = (self.training and cache is None and getattr(self.model, "gradient_checkpointing", False) and ckpt_fn is not None) for layer in self._ponder_layers: lkw = dict(kwargs) if isinstance(layer, PonderMoEDecoderLayer): lkw["moe_route_mode"] = route_mode if use_ckpt: h = ckpt_fn( layer.__call__, h, attention_mask=attention_mask, position_ids=position_ids, position_embeddings=position_embeddings, past_key_values=None, use_cache=False, **lkw, ) else: h = layer( h, attention_mask=attention_mask, position_ids=position_ids, position_embeddings=position_embeddings, past_key_values=cache, use_cache=cache is not None, **lkw, ) return h # ------------------------------------------------------------------ # # 主前向 (父类结构 + 路由模式 + MoE/MTP 损失) # ------------------------------------------------------------------ # @can_return_tuple def forward( self, input_ids: torch.LongTensor | None = None, attention_mask: torch.Tensor | None = None, position_ids: torch.LongTensor | None = None, past_key_values: DynamicCache | None = None, inputs_embeds: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, use_cache: bool | None = None, logits_to_keep: int | torch.Tensor = 0, output_ponder: bool = False, **kwargs, ) -> PonderMoECausalLMOutputWithPast: cfg = self.config m = self.model self._reset_moe_stats() # ---- 输入 & 位置 ------------------------------------------------ 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 = m.embed_tokens(input_ids) B, Q = inputs_embeds.shape[:2] if use_cache is None: use_cache = False if self.training else cfg.use_cache if use_cache and past_key_values is None: past_key_values = DynamicCache(config=cfg) cache_position = kwargs.pop("cache_position", None) past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0 if position_ids is None: base = (cache_position if cache_position is not None else torch.arange(Q, device=inputs_embeds.device) + past_seen) position_ids = base.reshape(1, -1).expand(B, -1)[:, :Q].to(inputs_embeds.device) # ---- mask / rope ------------------------------------------------- if cache_position is None: cache_position = torch.arange(Q, device=inputs_embeds.device) + past_seen causal_mask = create_causal_mask( config=cfg, input_embeds=inputs_embeds, attention_mask=attention_mask, cache_position=cache_position, past_key_values=past_key_values, position_ids=position_ids, ) position_embeddings = m.rotary_emb(inputs_embeds, position_ids=position_ids) # ---- Stage A: 前段层单遍 (scope=all 时 MoE 层每次调用都是新 token -> fresh) -- h = inputs_embeds for layer in m.layers[: cfg.ponder_start_layer]: if isinstance(layer, PonderMoEDecoderLayer): h = layer(h, attention_mask=causal_mask, position_ids=position_ids, position_embeddings=position_embeddings, past_key_values=past_key_values, use_cache=use_cache, moe_route_mode="fresh", **kwargs) else: h = layer(h, attention_mask=causal_mask, position_ids=position_ids, position_embeddings=position_embeddings, past_key_values=past_key_values, use_cache=use_cache, **kwargs) # ---- 思考窗口 W -------------------------------------------------- if Q == 1 or self.training or cfg.ponder_all_positions: W = Q else: W = max(1, min(cfg.ponder_window, Q)) eps = cfg.ponder_epsilon K = cfg.max_ponder_steps device, dtype = inputs_embeds.device, inputs_embeds.dtype remaining = torch.ones(B, W, device=device, dtype=dtype) cum = torch.zeros(B, W, device=device, dtype=dtype) active = torch.ones(B, W, device=device, dtype=torch.bool) mix = torch.zeros(B, W, h.size(-1), device=device, dtype=dtype) w_list: list[torch.Tensor] = [] h_prev_win = None h_full = None h_win = h[:, -W:, :] if W < Q else h executed = 0 win_mask_built = False forced = int(getattr(cfg, "ponder_forced_steps", 0) or 0) step_ce_weight = float(getattr(cfg, "ponder_step_ce_weight", 0.0) or 0.0) use_step_ce = (step_ce_weight > 0.0 and labels is not None and self.training) step_pairs: list[tuple[torch.Tensor, torch.Tensor]] = [] loop_K = forced if forced > 0 else K for k in range(1, loop_K + 1): executed = k if k > 1: self._crop_ponder_kv(past_key_values, W) use_window = (k > 1) and (past_key_values is not None) and (W < Q) route_mode = "fresh" if k == 1 else "frozen" # 单次路由铁律 if not use_window: h = self._run_block(h, causal_mask, position_ids, position_embeddings, past_key_values, kwargs, route_mode=route_mode) h_win = h if W == Q else h[:, -W:, :] if W < Q: h_full = h else: if not win_mask_built: past_len = int(past_key_values.layers[cfg.ponder_start_layer].get_seq_length()) mask_win = self._win_mask(W, past_len, device, dtype, attention_mask) pos_emb_win = (position_embeddings[0][:, -W:], position_embeddings[1][:, -W:]) win_mask_built = True h_win = self._run_block(h_win, mask_win, position_ids[:, -W:], pos_emb_win, past_key_values, kwargs, route_mode=route_mode) if forced > 0: continue lam = self._halting_prob(h_win, h_prev_win) lam_eff = lam * active.to(lam.dtype) w_k = remaining * lam_eff mix = mix + w_k.unsqueeze(-1) * h_win w_list.append(w_k) remaining = remaining * (1.0 - lam_eff) cum = cum + w_k active = active & (cum < 1.0 - eps) h_prev_win = h_win if use_step_ce: ce_k = self._step_ce(h_win, labels) step_pairs.append((w_k.detach(), ce_k)) if not bool(active.any()): break if forced > 0: if W < Q: h_final = h_full.clone() h_final[:, -W:, :] = h_win else: h_final = h_win w_stack = F.one_hot( torch.full((B, W), forced - 1, device=device), forced ).to(dtype) else: mix = mix + remaining.unsqueeze(-1) * h_win w_stack = torch.stack(w_list, dim=-1).to(dtype) w_stack = torch.cat([w_stack[..., :-1], w_stack[..., -1:] + remaining.unsqueeze(-1)], dim=-1) if W < Q: h_final = h_full h_final[:, -W:, :] = mix else: h_final = mix # ---- commit pass: 复用 frozen 路由写最终 KV ----------------------- if (use_cache and cfg.commit_kv and past_key_values is not None and executed > 1 and forced == 0): self._crop_ponder_kv(past_key_values, W) if not win_mask_built: self._run_block(h_final, causal_mask, position_ids, position_embeddings, past_key_values, kwargs, route_mode="frozen") else: self._run_block(h_final[:, -W:, :], mask_win, position_ids[:, -W:], pos_emb_win, past_key_values, kwargs, route_mode="frozen") # ---- logits / loss ---------------------------------------------- slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep logits = self.lm_head(m.norm(h_final[:, slice_indices, :])) loss, ponder_kl, ponder_step_ce = None, None, None if labels is not None: ce_logits = logits.float() loss = F.cross_entropy( ce_logits[..., :-1, :].reshape(-1, ce_logits.size(-1)), labels[..., 1:].reshape(-1), ignore_index=-100) # PonderNet KL 正则 (与父类一致, fp32 域防下溢) p = (w_stack / w_stack.sum(-1, keepdim=True).clamp_min(1e-8)).float() Kd = p.size(-1) n_idx = torch.arange(1, Kd + 1, device=p.device, dtype=p.dtype) pg = cfg.ponder_prior_p geom = (1 - pg) ** (n_idx - 1) * pg geom = torch.cat([geom[..., :-1], ((1 - pg) ** (Kd - 1) * torch.ones_like(geom[..., -1:]))], dim=-1) geom = geom / geom.sum() kl = (p * (torch.log(p.clamp_min(1e-8)) - torch.log(geom.expand_as(p)))).sum(-1) ponder_kl = kl.mean() loss = loss + cfg.ponder_loss_beta * ponder_kl if use_step_ce and step_pairs: ws = torch.stack([pp[0] for pp in step_pairs], dim=0) ces = torch.stack([pp[1] for pp in step_pairs], dim=0) wsum = ws.sum(0).clamp_min(1e-8) per_pos = (ws * ces.view(-1, 1, 1)).sum(0) / wsum ponder_step_ce = per_pos.mean() loss = loss + step_ce_weight * ponder_step_ce # MoE 负载均衡 (宪法 §六). checkpointing 下 last_aux 无梯度 -> 跳过加项仍报告 moe_aux, moe_active = self._collect_moe_stats() if (moe_aux is not None and cfg.moe_aux_weight > 0 and torch.is_tensor(moe_aux) and moe_aux.requires_grad): loss = loss + cfg.moe_aux_weight * moe_aux else: moe_aux, moe_active = self._collect_moe_stats() # MTP 多 Token 预测损失 (宪法 §六; 训练增强, 推理旁路) mtp_loss = mtp_acc = None if (labels is not None and self.mtp_heads is not None and input_ids is not None): mtp_loss, mtp_acc = self._mtp_forward(h_final, input_ids, labels) if mtp_loss is not None and cfg.mtp_loss_weight > 0: loss = loss + cfg.mtp_loss_weight * mtp_loss # ---- 诊断 -------------------------------------------------------- cum_w = torch.cumsum(w_stack, dim=-1) steps_used = ((cum_w < 1.0 - eps).sum(-1) + 1).clamp(max=w_stack.size(-1)) self._log({ "mode": "train" if self.training else ("decode" if Q == 1 else "prefill"), "signal": cfg.ponder_signal, "executed": executed, "steps_max": int(steps_used.max()), "steps_mean": float(steps_used.float().mean()), "moe_active": (float(moe_active) if torch.is_tensor(moe_active) else None), }) return PonderMoECausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=past_key_values, ponder_weights=w_stack if output_ponder else None, ponder_steps=steps_used if output_ponder else None, ponder_kl=ponder_kl, ponder_step_ce=ponder_step_ce, moe_aux_loss=moe_aux, moe_active_mean=moe_active, mtp_loss=mtp_loss, mtp_acc=mtp_acc, ) # -------------------------------------------------------------------------- # 发布版统一注册与别名 (auto_map 指向别名; 原类名保留以兼容旧代码) # -------------------------------------------------------------------------- SamaiPnetConfig = PonderMoELlamaConfig SamaiPnetForCausalLM = PonderMoELlamaForCausalLM # transformers 5.x 严格校验: AutoModelForCausalLM.register 要求 # model_class.config_class 与实际 config 类一致 (否则 ValueError) PonderLlamaForCausalLM.config_class = PonderLlamaConfig PonderMoELlamaForCausalLM.config_class = PonderMoELlamaConfig try: # 重复 import 安全 AutoConfig.register("llama_ponder", PonderLlamaConfig) AutoConfig.register("llama_ponder_moe", PonderMoELlamaConfig) AutoModelForCausalLM.register(PonderLlamaConfig, PonderLlamaForCausalLM) AutoModelForCausalLM.register(PonderMoELlamaConfig, PonderMoELlamaForCausalLM) except Exception: pass