Feat/add gemma3 model

[zyc107109102]

PR 标题: feat: 支持 gemma2 与 gemma3_text 模型推理

概述

本 PR 以两个 commit 为 InfiniLM 增加两个 Gemma 文本模型家族:

• commit 1 — model_type="gemma2"(Gemma-2):LLaMA 家族架构,差异点是注意力/最终 logits 软上限(soft-capping)、每层四个 RMSNorm 的解码层、零中心归一化权重。
• commit 2 — model_type="gemma3_text"(Gemma-3 文本,如 gemma-3-1b):QK-norm、按层类型区分的 RoPE 频率、交替滑窗注意力(5:1 交错)——其中滑窗 KV cache 是框架此前不具备的能力。

顺带引入两个小的、可选的框架能力:

1. AttentionLayer / StaticAttentionImpl 上的可选 softcap 参数(默认 0.0f = 关闭,所有现有模型不受影响);
2. 模型自管的滑窗注意力,直接复用现有静态 KV cache。

全部改动在三个后端上与 HF transformers(eager attention、贪心解码)做了逐 token 对照:Iluvatar MR-V100(bf16)、NVIDIA RTX 4060(bf16)、CPU(fp32)。两个 commit 均可独立编译并通过各自的单元测试(bisect 友好)。

架构说明(Gemma 与 llama 的差异)

Gemma-2(commit 1)

Gemma-2 是 LLaMA 家族架构,有四处差异,均已对照 HF 参考实现逐条验证:

1. 注意力 logit 软上限。 缩放后的分数在因果掩码和 softmax 之前先做 tanh(s/cap) * cap 压缩;LM head 之后的最终 logits 做同样处理(所有已发布尺寸均为 attn_logit_softcapping=50、final_logit_softcapping=30)。注意力缩放使用 1/sqrt(query_pre_attn_scalar) 而非 1/sqrt(head_dim)。

2. 每层四个 RMSNorm,顺序为"分支先归一化再相加":

   residual = hidden
   hidden   = input_layernorm(hidden)          # 归一化主流
   hidden   = attn(hidden)
   hidden   = post_attention_layernorm(hidden) # 归一化分支
   residual = residual + hidden                # 然后加回主流
   hidden   = pre_feedforward_layernorm(residual)
   hidden   = mlp(hidden)
   hidden   = post_feedforward_layernorm(hidden)  # 归一化分支
   residual = residual + hidden
   

   实现为自定义 Gemma2DecoderLayer;TextModel 的 residual 契约(hidden = 尚未加回的分支,residual = 主流)保持不变,因此 TextModel 本体零改动。分支加法通过 InfiniCore 的 add_rms_norm 与 pre-feedforward 归一化融合(返回 (norm(residual+branch), residual+branch)),替代独立的 add + norm 两步;在 MR-V100 上实测性能中性(batch-1 解码 ITL 22.72 vs 22.76 ms,2k prefill TTFT 720.2 vs 721.1 ms),保留理由是算子序列更精简。

3. 零中心 RMSNorm 权重——checkpoint 存储 w - 1,加载期 remapper 统一 +1。

4. Embedding 缩放——embedding 输出在加载期乘以 sqrt(hidden_size)(经 _get_scale_emb);tied lm_head 用未缩放副本填充,与 HF 语义一致。激活仅实现 gelu_pytorch_tanh;head_dim(256)从 checkpoint 保留,不按 hidden/num_heads(2304/8 = 288)推导。

Gemma-3 文本(commit 2)

复用 Gemma-2 的四范式解码层、权重 remap 与 embedding 缩放(MLP 直接以别名复用)。新能力:

1. QK-norm:RoPE 之前对逐头向量归一化(qwen3 模式);其权重同样以 w - 1 存储,由共享 remapper 处理。

2. 按层类型区分的 RoPE。 layer_types 按 sliding_window_pattern 周期(默认 6 → 5:1)交替 sliding_attention / full_attention,与 HF 的生成规则一致((i+1) % pattern == 0 → full)。滑窗层使用 rope_local_base_freq(10k),全局层使用 rope_theta(1M)。

3. 滑窗注意力。 滑窗层直接管理静态 KV cache(cache 布局与更新逻辑对齐 StaticAttentionImpl):

   • decode 只读尾部 min(total, W) 个 key——无需掩码;
   • prefill 叠加一条按分数 dtype 在 host 侧构造的带状窗口掩码(fp32、bf16 位模式、fp16 因 -1e9 溢出改用 -inf);当 past + i - j >= W 时 key j 对 query i 不可见。因果部分交给 causal_softmax_,其对齐因果语义(valid_len = total - seq + i + 1)与窗口掩码精确复合,chunked prefill(past > 0, seq > 1)同样成立;
   • GQA 重排为头优先([bs, seq, heads, dim] → permute → [bs*nkv, ng*seq, dim]),与 static 后端一致。

   全注意力层委托共享的 AttentionLayer,可在任意后端运行;滑窗层要求 STATIC 后端(显式报错)。Gemma-3 无 logit 软上限,顶层直接复用 TextCausalLM,零改动。

框架改动:注意力 softcap(commit 1)

softcap 必须作用在分数矩阵与"掩码+softmax"之间——这在注意力后端内部——所以它是 AttentionLayer / StaticAttentionImpl 的构造参数(softcap,默认 0.0f = 关闭)。设计与 vLLM 一致(logits_soft_cap 即注意力后端参数)。现有模型不受影响(默认关闭,无数值与分发路径变化)。

不支持的组合一律响亮报错,绝不静默输出错误结果:

• paged / flash 后端拒绝非零 softcap(它们不落地分数矩阵,无处插入);
• 任意后端拒绝负 softcap;
• 最终 logits 的 softcap 负值在 config 层拒绝。

如果 maintainer 希望把这部分拆成独立的框架 PR,随时可以拆。

加载路径改动(python/infinilm/modeling_utils.py)

• _remap_gemma:对所有零中心归一化权重(input_layernorm、post_attention_layernorm、pre_feedforward_layernorm、post_feedforward_layernorm、self_attn.q_norm、self_attn.k_norm、model.norm)做 +1 平移;注册到 gemma2 与 gemma3_text。
• _get_scale_emb:对上述模型类型返回 sqrt(hidden_size);tied lm_head 用未缩放的 embedding 副本填充。

加固:不支持的配置一律响亮报错

配置组合	行为
gemma3 配置带 rope_scaling	拒绝加载——gemma-3-4b/12b/27b 文本配置在全局层使用 linear factor 8;尚未实现,拒绝加载而不是默默算错长上下文 RoPE
KV-cache 量化 + 滑窗层	拒绝——滑窗层直接以激活 dtype 写 cache,int8 cache 下会被静默截断
滑窗层缺失/为零的 sliding_window	拒绝
滑窗层 KV-cache 溢出	容量断言(与 StaticAttentionImpl 同一防线)
gemma2 且 sliding_window < max_position_embeddings	加载期警告(如上下文扩展版 2b:窗口 4096 < 上下文 8192)
负 softcap 值	拒绝(config 层 + AttentionLayer 双重)

按 commit 粒度做了 A/B:所有正常路径的输出与不加硬化逐位一致;每条防护都用探针验证过会带着预期信息触发。

局限(代码内与 PR 中均已注明)

1. Gemma-2 的滑窗注意力未实现。 seq_len <= sliding_window(所有已发布尺寸均为 4096)时结果正确——此时局部窗口覆盖全部因果上下文。超过后模型按全注意力运行:功能正常、输出连贯,但与 HF 分歧(HF 会对局部层开窗)——加载期警告覆盖此场景。
2. softcap 仅在 STATIC 后端支持。
3. 单块 prefill 的显存上限。 static 后端要落地 S×S 分数矩阵,且每个滑窗层各缓存一份窗口掩码;gemma-3-1b 在 32GB 卡上单块 prefill 实测可用至约 11.7k token,超过约 12k 后以干净的 OOM 失败(是分配报错,绝不是错误结果)。chunked prefill 与更短上下文不受影响。
4. rope_scaling(linear)未实现——见后续工作。
5. TP/PP 代码路径继承自共享 TextModel 模板(Gemma-2 顶层对齐 TextCausalLM 的 PP 分段),但本 PR 仅在单卡上实测。

验证

所有对照均为贪心解码(do_sample=False),参考实现为 HF attn_implementation="eager",除注明外均逐 token 一致。

小随机模型(fp32,CPU)——覆盖 head_dim ≠ hidden/heads、四范式、w-1 权重、tied embedding、窗口 < 提示长度:

• tiny gemma2:36/36
• tiny gemma3(窗口 8,局部/全局 theta 不同,32 token 提示):30/30

真实权重:

模型 / 平台	结果
gemma-2-2b-it,智凯 MR-V100(bf16)	84/84(3 条提示,最长 36 token)
gemma-2-2b-it,RTX 4060(bf16)	29/29
gemma-2-2b-it,长上下文 3,273 / 3,523 token	逐 token 一致,needle 回答正确
gemma-2-9b-it,MR-V100(bf16,零代码改动)	38/39——602 token 提示上单个 bf16 翻转后重新收敛
gemma-3-1b-it,MR-V100(bf16)	24/24(3 条提示)+ 602 token 提示越过 512 窗口 8/8——滑窗真正被激活而非绕过
gemma-3-1b-it,分块 prefill(past>0,seq>1,prefix-cache 复用)	10/10 + 10/10
gemma-3-1b-it,8,194 token 提示	与 HF 逐 token 一致
gemma-3-1b-it,11,694 token 大海捞针	答案正确("74329")
gemma-2 超窗探针,4,353 token	输出连贯、不崩溃(已文档化的分歧)

其他检查:

• softcap 效应探针:|logits|max 加帽 19.75 vs 不加帽 23.50;
• 权重加载:所有 checkpoint 无缺失/多余 key;
• 逐 commit 独立编译 + 单元测试(5 + 6+6),scripts/format.py --check 通过;
• 加固 A/B:正常路径输出逐位一致;负路径探针确认每条防护都带预期信息触发。

复现方式

# 单元测试
python -m unittest discover -s test/models/gemma2
python -m unittest discover -s test/models/gemma3

# 与 HF 的贪心对照(device: cpu / cuda)
python compare_script.py <model_dir> <max_tokens> bfloat16 cuda

后续工作(愿意继续贡献,优先级欢迎反馈)

1. Linear rope scaling(全局层位置 ÷ factor)——解锁 gemma-3-4b/12b/27b 文本配置,当前以明确报错拒绝。
2. 滑窗层 prefill key 范围截断——每个 query 只需要 [past-W+1, past+seq-1] 内的 key;截断可去掉绝大部分逐层掩码/分数显存,显著抬高分块 prefill 的长上下文上限。
3. 窗口掩码缓存共享——掩码只依赖 (seq_len, total_len),当前却按滑窗层实例各缓存一份。
4. Gemma-2 滑窗注意力(复用 gemma3 路径)。
5. Paged/flash 后端的 softcap——需要 InfiniOps 内核配合。

Closes #<issue 编号>