qwen4_exp: fp8-quantized n-gram (PLE) embedding rows are gathered without dequantization

快速结论：这个报错通常出现在加载 Qwen4-系列 FP8 量化模型（如 Qwen3.8-Flash-Next-FP8）并运行前向时，n-gram（PLE）embedding 表以 FP8 存储，但 gather 出来的行没有反量化，导致下游 bf16 模块 dtype 不匹配。优先排查 embedding 量化权重是否被当作普通 nn.Embedding 处理、以及 weight_scale 是否丢失。

适用环境：Issue 报告中使用的环境为 Transformers main（5.16.0.dev0），模型为 51.2B 参数 n-gram 表的 Qwen/Qwen3.8-Flash-Next-FP8；涉及 FP8（Float8_e4m3fn）与 bf16 混合精度路径。Issue 未明确列出 Python、CUDA、PyTorch 具体版本。

最快修复方案：该问题在 Issue 讨论中由 #48368 修复。请更新到包含该修复的 Transformers 版本；Issue 中原报告者给出的本地可行方案是：在 gather 之后执行反量化，即 .to(bf16) * weight_scale，并把 weight_scale 保留为 buffer。

注意事项：更新修复补丁前，本地补丁只针对 n-gram embedding；如果 ple.key_proj 也被错误转换，不会立即崩溃，而是会静默产生错误缩放的值，需要一并核对。单卡显存不足时，Issue 提到有 CPU-offload 版本可用于 CUDA-only FP8 路径，但未合并，属于未验证的可选方案。

问题场景

在 Transformers 下加载并运行 Qwen3.8-Flash-Next-FP8（或类似包含 FP8 量化 n-gram embedding 的 Qwen4 实验模型）时触发。n-gram 表位于 ple.ple_embedding.ngram_embedding，形状为 320,001,536 × 160，以 fp8 存储并带有 per-tensor 的 weight_scale（bf16）。该模块是普通的 nn.Embedding，scale 张量没有对应的参数槽位，加载时被当作 unexpected key 丢弃，Qwen4ExpTextNGramEmbedding.forward 直接返回未反量化的 fp8 行。问题出现在模型 day-0 bring-up 阶段。

报错原文

expected m1 and m2 to have the same dtype: Float8_e4m3fn != BFloat16

原因分析

最可能的原因是 n-gram embedding 的 FP8 量化权重在 gather 后没有反量化。该 embedding 模块被按普通 nn.Embedding 处理，per-tensor 的 weight_scale 因没有 parameter slot 被丢弃为 unexpected key，前向直接返回原始 fp8 行。第一个消费者 ple.key_proj（bf16）因此收到 Float8_e4m3fn 与 BFloat16 不匹配的张量而崩溃。另一种可能情况是：如果 key_proj 也被错误转换，则不会报 dtype 错误，而是静默计算出缩放错误的值。Issue 中给出的修复方向有两个：gather 后反量化（.to(bf16) * weight_scale，并将 scale 保留在 buffer 中），或让 FP8 量化器自己持有 embedding 类型的量化张量。

环境排查

确认 Transformers 版本；Issue 环境为 main / 5.16.0.dev0。
确认是否已包含 #48368 的修复。
确认模型为 Qwen3.8-Flash-Next-FP8 或同结构 FP8 量化 Qwen4 实验模型。
确认 ple.ple_embedding.ngram_embedding 的 weight_scale 是否在加载时被丢弃（unexpected key）。
确认 ple.key_proj 的 dtype 是否为 bf16；若它也处于异常量化状态，可能不报错而是计算结果异常。
确认 GPU 是否支持 FP8（Float8_e4m3fn）路径；Issue 提到相关 FP8 路径为 CUDA-only。

解决步骤

优先将 Transformers 更新到包含 #48368 的版本，该 PR 明确用于修复本 Issue。
若暂时无法更新，可参考 Issue 中的本地修复方向（可优先尝试）：在 Qwen4ExpTextNGramEmbedding.forward 中 gather 之后执行反量化，即 gathered.to(bf16) * weight_scale，并确保 weight_scale 作为 buffer 被正确加载而不是被丢弃。
检查加载日志中是否有与 weight_scale 相关的 unexpected key，确认 scale 已进入模块。
同时检查 ple.key_proj 是否也被错误量化；如果存在同类问题，需要按 PR 一并应用对应修复，避免静默错误结果。
如显存不足，Issue 提到存在可行的 CPU-offload 修改，但尚未合并，属于未验证的可选路径，不建议作为首选。

验证方法

应用修复后重新加载 Qwen/Qwen3.8-Flash-Next-FP8，模型应能正常加载并完成前向生成；不再出现 Float8_e4m3fn != BFloat16。Issue 中原报告者表示三个相关报错在本地补丁下均验证通过，模型加载和生成正确。若需进一步确认 key_proj 没有静默误算，可对比修复前后的输出数值是否一致、是否出现异常缩放。