一位贡献者正在准备一个拉取请求,为 Hugging Face 添加对 Gemma 4 模型系列的官方支持,并附带针对 ONNX 导出追踪的优化。这些更改解决了在 ONNX 追踪期间的架构特定限制,并修正了异构模型变体(包括混合专家 (MoE) 和密集配置)的层计算。

  • RMSNorm 公式已重新对齐,以使用标准的 Llama 风格缩放乘法来处理 Gemma 4,并带有门控属性以防止激活缩放错误。
  • 通过用 torch.bincount 替换 torch.histc,改进了对分组查询注意力 (GQA) 和 MoE 注意力的 ONNX 追踪兼容性,从而在 ONNX Opset 17 下实现图序列化。
  • 添加了符号张量转换保护机制,以防止在处理动态导出追踪期间的符号标量张量时出现 JIT 编译失败。
  • 感知存储的分片合并通过将由共享参数存储合并为连续的 1GB 二进制缓冲区,防止 MoE 配置中出现病态的文件分片模式。
  • 通过 MRO 类层次结构查找实现了健壮的 tokenizer 属性解析,以防止在自定义 tokenizer 子类中出现 AttributeError 异常。

已对更新后的包进行了本地验证,核心模型架构能够正确加载和编译。通过 gemma2 PyTorch 建模测试套件中的 157 个测试用例(以及 1,157 个子测试)确认了向后兼容性。