Gemma 4

[SamuelMarks]

Add support for Gemma 4 Multimodal Architecture

Following the resounding success of jax-ml/bonsai#175 [JAX] I give this other Google repo Gemma 4!

Description

This PR introduces the complete implementation of the Gemma 4 architecture. It brings robust support for multimodality—specifically advanced audio and vision capabilities—along with an extensive suite of tests and updated inference scripts.

Key Changes

🧠 Core Architecture (gemma/gemma4/)

• Base Components: Implemented all foundational layers for Gemma 4, including attention.py, rope.py, moe.py, cache.py, decoder_layer.py, and modeling.py.
• Vision Integration: Added vision.py for processing visual features in the Gemma 4 architecture.
• Top-Level Models: Added gemma4_model.py and gemma4_multimodal.py for cohesive model instantiation and multimodal routing.

🎧 Audio Capabilities (gemma/gemma4_audio/)

• Introduced a dedicated package for audio processing.
• Includes audio_model.py, audio_attention.py, and audio_layers.py to seamlessly project and attend over audio modalities.

⚙️ Integration & Refactoring

• Configurations & Registration: Updated gemma/__init__.py, gemma/config.py, gemma/model.py, and gemma/tokenizer.py to support new Gemma 4 configs and routing.
• Backwards Compatibility: Adjusted gemma/gemma3_model.py to integrate cleanly alongside the new abstractions.

🛠 Scripts & Utilities

• Weight Conversion: Added scripts/convert_gemma4_weights.py to facilitate migrating or preparing checkpoint weights.
• Inference Pipeline: Updated scripts/run_multimodal.py to support Gemma4Config and Gemma4ForCausalLM, allowing users to quickly run multimodal inference using Gemma 4 checkpoints.

✅ Testing

• Comprehensive Test Suite: Added a robust suite of unit tests under tests/gemma4/ verifying attention, audio, cache, config, decoder_layer, layers, modeling, moe, rope, and vision.
• Root Tests: Added specialized tests/test_gemma4_attention.py and tests/test_gemma4_moe.py testing modules.
• Script Tests: Added tests/test_workflow_scripts.py to ensure end-to-end script reliability.

Testing Instructions

• Run the test suite: pytest tests/gemma4/
• Verify multimodal generation with: python scripts/run_multimodal.py --ckpt <path_to_gemma4_checkpoint> --variant 4b