WeLM-v4.5-80B-A3B-Instruct-0608：同一份权重，transformers 4.x 推理正常、5.x 坍塌 的根因排查

一、现象

同一个 checkpoint（WeLM-v4.5-80B-A3B-Instruct-0608），用 AutoModelForCausalLM.from_pretrained(..., device_map="auto", low_cpu_mem_usage=True) 加载后做贪心解码：

直观对比同一个问题、同一份权重，只换 transformers 版本，输出就从「正常对话」退化成「无意义重复」：

>>> 输入：你好
你好！有什么可以帮你的吗？

>>> 输入：用一句话介绍你自己
我是 WeLM-v4.5-80B-A3B-Instruct，
一个由腾讯训练的对话助手。

>>> 输入：你好
\n\n\n\n\n\n\n\n …（无限重复）

>>> 输入：用一句话介绍你自己
</think></think></think> …
/ 重复乱码，完全不成句

注意：5.6 的 logits std 没有下降（2.98 → 3.14），不是「整体被压平」，而是 排序完全错乱——这个细节后来成了重要线索。

二、结论速览（TL;DR）

修复：加载后对每个 rotary 用 rope_init_fn 重算 inv_freq（仅在异常时）。5.x 自动修复、4.57.x 为 no-op。已实测：修复后 5.6 logits 与 4.57 逐位一致（30.0 / 你好）。

三、排查方法与约束的确立

这套 modeling 是 trust_remote_code 加载的本地代码（与 transformers 版本无关）。因此版本差异只可能来自 modeling 调用的 transformers 包内组件。排查前先钉死了两个关键约束，避免被噪声误导：

venv457 :  python 3.12.3   torch 2.11.0+cu130   transformers 4.57.1
system  :  python 3.12.3   torch 2.11.0+cu130   transformers 5.6.0

四、逐步定位（含走过的弯路与纠正）

本节如实记录排查过程，包括两次被污染数据误导的弯路，以及最终如何用干净实验纠正。

rope_type=yarn  attention_scaling=1.00000
inv_freq.shape=(32,) norm=0.00000  first4=[0.0, 0.0, 0.0, 0.0]

五、根因机理

modeling 里 RoPE 的定义（节选）：

class WeLMV4MoeRotaryEmbedding(nn.Module):
    def __init__(self, config, device=None):
        ...
        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]        # 从 transformers 导入
        self.config.partial_rotary_factor = config.qk_rope_head_dim / config.head_dim   # 0.25
        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)   # ← 不进 checkpoint

    @torch.no_grad()
    @dynamic_rope_update
    def forward(self, x, position_ids):
        inv_freq_expanded = self.inv_freq[None,:,None].float()...        # ← 读 self.inv_freq
        freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1,2)
        emb = torch.cat((freqs, freqs), dim=-1)
        cos = emb.cos() * self.attention_scaling
        sin = emb.sin() * self.attention_scaling
        return cos, sin

为什么 4.57.x 没问题

transformers 4.57.x 在 from_pretrained 完成后，会对这类不在 checkpoint 中的非持久 buffer 正确重新物化（按真实 device 重算 inv_freq）——这一步在 4.x 里是加载流程中独立的一环，与 _init_weights 无关，总会执行。transformers 5.0 重写了加载流程，把这一步搬进了基类 PreTrainedModel._init_weights（含 rope_scaling→rope_parameters、standardize_rope_params 等）。这一改动本身没问题——只要 modeling 的 _init_weights 会经过基类逻辑。本模型恰恰重写了 _init_weights 又没调 super()（详见下一小节的源码层定论），于是该步骤被静默绕过，inv_freq 留在 meta 落地的全 0 状态。这是唯一的版本差异点。

源码层定论：根因是自研 _init_weights 不调 super()，绕过了基类的 rope 重物化

进一步实地核对本机 transformers 5.8.1 与自研 modeling_welmv4_moe.py 源码后，把根因钉死在源码层。先纠正一个一度的误判：

self.register_buffer("inv_freq", inv_freq, persistent=False)
self.original_inv_freq = self.inv_freq          # ← 实际是有的

关键在于：transformers 5.x 把 inv_freq 的重物化逻辑搬进了基类 PreTrainedModel._init_weights（modeling_utils.py:2402-2411 那个 "RotaryEmbedding" in name and hasattr(original_inv_freq) 分支）。谁的 _init_weights 能跑到这段基类逻辑，谁就安全。

而自研 modeling_welmv4_moe.py:869-884 重写了 _init_weights 却没调 super()，且只处理 nn.Linear / nn.Embedding、完全没有 RotaryEmbedding 分支：

def _init_weights(self, module):
    ...
    if isinstance(module, nn.Linear):
        ...
    elif isinstance(module, nn.Embedding):
        ...
    # ✗ 没有 super()._init_weights(module)，也没有 rotary 分支
    # → 基类那段 rope inv_freq 重物化从不执行

这就解释了 4.x vs 5.x 的全部差异：4.57.x 时 inv_freq 的重物化是加载流程里独立的一步（与 _init_weights 无关，总会跑）；5.x 重构后把它搬进了 _init_weights，于是任何「重写 _init_weights 又不调 super」的自研 modeling 就静默漏掉了这一步。

def _init_weights(self, module):
    super()._init_weights(module)   # ← 关键：让 5.x 的 rope inv_freq 重物化分支跑起来
    cfg = self.config
    ...

六、决定性证明

在 5.6 单次加载内做 A/B：默认 vs 加载后重算 inv_freq。

===== [baseline] 默认加载(device_map=auto/meta) =====
[baseline] inv_freq = shape=(32,) norm=0.00000  first4=[0.0, 0.0, 0.0, 0.0]
[baseline] logits.max=19.2500  top1='\n\n'              ← 坍塌

===== [realfix] 加载后重新物化 inv_freq =====
[realfix] inv_freq = shape=(32,) norm=1.33674  first4=[1.0, 0.6636, 0.44037, 0.29223]
[realfix] logits.max=30.0000   top1='你好'              ← 完全恢复

修复后 5.6 的 top5 与 4.57 逐位一致：[30.0, 22.25, 20.75, 17.25, 16.875] → ['你好','您好','Hello','�','好的']。

[rope-fix] WeLMV4MoeRotaryEmbedding 共 49 个，重物化 inv_freq 49 个
[TEXT] OUTPUT:
<report>你好！有什么可以帮你的吗？</report>

七、修复方案与落点

核心思路：加载后重新物化 inv_freq，且仅在异常时（meta / 全 0 / NaN / 缺失）才动手，保证 4.57.x 等正常路径为 no-op、不引入任何行为变化，也无需锁版本。

def _rematerialize_rope_inv_freq(model, is_rank_0=True):
    for m in model.modules():
        if m.__class__.__name__ != "WeLMV4MoeRotaryEmbedding":
            continue
        inv = getattr(m, "inv_freq", None)
        needs_fix = (inv is None or getattr(inv,"is_meta",False)
                     or not bool(torch.isfinite(inv).all())
                     or float(inv.float().abs().sum()) == 0.0)
        if not needs_fix:
            continue                       # 4.57.x 正常 -> no-op
        inv_freq, scaling = m.rope_init_fn(m.config, device=torch.device("cpu"))
        m.register_buffer("inv_freq", inv_freq.to(torch.float32), persistent=False)
        m.attention_scaling = scaling
        m.original_inv_freq = m.inv_freq   # 5.x meta-load -> 重算覆盖全 0

八、影响范围：哪些场景会中招、哪些不会

触发条件可以归纳为一个充要组合：transformers 5.x 且 模型在 meta device 上被构造（构造 __init__ 时 inv_freq 落地为全 0），加载完权重后又没有重新物化它。任何不满足这个组合的场景都不受影响。

统一解法：加载后兜底物化，与场景无关

修复不依赖「判断当前是哪种场景」，而是在所有路径统一加一道兜底：模型加载完成后、deepspeed.initialize / ZeRO 包装之前，遍历每个 WeLMV4MoeRotaryEmbedding，仅当 inv_freq 异常（is_meta / 全 0 / NaN / 缺失）时用 rope_init_fn 在真实 device 上重算覆盖（见第七节代码）。

• 中招场景（5.x + meta/device_map、5.x + ZeRO-3 走 meta 等）：检测到全 0 → 重算修复，inv_freq 恢复 norm=1.337。
• 不中招场景（4.57.x、5.x 普通加载、ZeRO-2 非 meta 加载等）：检测到正常 → no-op，不引入任何行为变化，无需锁 transformers 版本。
• 修复已加到 OE / OE-off 两条训练入口与推理 load()，训练侧、推理侧、各 ZeRO 配置统一免疫。

补充：历史 checkpoint-216 能在 4.57 正常推理，说明它当初是用正常 RoPE 训出来的（4.57 或未触发此坑的加载方式），不是被这个 bug 训坏的权重。另：sliding window、attention_mask（cute 路径用 cu_seqlens，不消费 transformers 的 4D mask）等均已验证与 transformers 版本无关，不受影响。

附录：实验脚本清单

均在 SFT/debug/，通过 bridge 在远程对 4.57.1 / 5.6.0 两环境跑同一脚本对比。