Hello everyone,
I am preparing to open a pull request to contribute upstream support for the Gemma 4 model family, alongside ONNX export tracing optimizations for custom hardware attention plugins. Before submitting, I wanted to open this topic to coordinate with the maintainers and share the technical details of the changes.
Our work focuses on resolving architecture-specific limitations during ONNX tracing, as well as correcting layer calculations for heterogeneous model variants (MoE and Dense).
Key Technical Contributions:
-
RMSNorm Formulation Realignment: Gemma 4 reverts to the standard Llama-style scale-multiplication (output = normed_states * weight) rather than the Gemma 2 style (weight + 1.0) offset. We have gated the use_gemma2_format property inside the default RMSNorm layer based on model type to prevent activation scaling errors when instantiating Gemma 4 configurations.
-
**ONNX Tracing Compatibility (GQA & MoE Attention):
**-
Replaced torch.histc with a torch.bincount-based calculation inside grouped_mm_experts_forward to enable graph serialization under ONNX Opset 17.
-
Updated sdpa_attention_forward to intercept standard attention operations during ONNX tracing, allowing custom hardware attention operator plugins to receive positional dual-RoPE and KV-cache parameters cleanly.
-
-
Symbolic Tensor Casting: Added scalar type-checking and int() conversion safeguards in sdpa_mask and flex_attention_mask to prevent JIT compilation failures when processing symbolic scalar tensors during dynamic export tracing.
-
Storage-Aware Shard Consolidation: Added an optional storage-level consolidation pass inside PreTrainedModel.save_pretrained for MoE configurations. This prevents pathological file-sharding patterns (such as generating tens of thousands of tiny individual files) by consolidating shared parameter storage into contiguous 1GB binary buffers on local unified memory architectures.
-
Robust Tokenizer Attribute Resolution: Implemented MRO class-hierarchy attribute lookup inside PreTrainedTokenizerBase._getattr_ to prevent erroneous AttributeError exceptions when instantiating custom tokenizer subclasses that inject class-level properties.
Current Verification Status:
We have built and verified the updated package locally. The core model architectures load and compile correctly. We also executed the gemma2 PyTorch modeling test suite to verify backwards-compatibility, resulting in 157 passed test cases (and 1,157 subtests passed).
I would appreciate any initial feedback or coordination from @ArthurZucker, @CyrilVallez, and vasqu regarding these changes before I open the pull request.