A contributor is preparing a pull request to add upstream support for the Gemma 4 model family to Hugging Face, alongside optimizations for ONNX export tracing. The changes address architecture-specific limitations during ONNX tracing and correct layer calculations for heterogeneous model variants, including Mixture of Experts (MoE) and Dense configurations.
- RMSNorm formulation is realigned to use standard Llama-style scale-multiplication for Gemma 4, with a gated property to prevent activation scaling errors.
- ONNX tracing compatibility for Grouped Query Attention (GQA) and MoE attention is improved by replacing torch.histc with torch.bincount to enable graph serialization under ONNX Opset 17.
- Symbolic tensor casting safeguards are added to prevent JIT compilation failures when processing symbolic scalar tensors during dynamic export tracing.
- Storage-aware shard consolidation prevents pathological file-sharding patterns in MoE configurations by consolidating shared parameter storage into contiguous 1GB binary buffers.
- Robust tokenizer attribute resolution is implemented via MRO class-hierarchy lookup to prevent AttributeError exceptions in custom tokenizer subclasses.
The updated package has been verified locally, with core model architectures loading and compiling correctly. Backwards-compatibility was confirmed by passing 157 test cases (and 1,157 subtests) from the gemma2 PyTorch modeling test suite.