samai-2b / modeling_samai_pnet.py
r18 delivery: modeling_samai_pnet.py
6941319 verified about 1 month ago
63.1 kB
# -*- 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 次迭代的全序列输出 (W<Q 时非窗口位置沿用)
        h_win = h[:, -W:, :] if W < Q else h
        executed = 0
        win_mask_built = False   # 窗口迭代专用 mask (懒构建一次)

        # ---- Round5 方向②③ 前置开关 --------------------------------------
        forced = int(getattr(cfg, "ponder_forced_steps", 0) or 0)   # 方向③: >0 强制 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<Q 时才做窗口局部重算; 无 cache (训练) 一律全序列重算
            use_window = (k > 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