Gemma 3 architecture
How the weightless Gemma 3 registers as a MAX custom architecture and reads its weightless config fields.
models/gemma3/ wires the model-agnostic core into a MAX SupportedArchitecture. It’s the pattern to copy when adding another model.
Registration
arch.py defines weightless_gemma3_arch with name="Gemma3ForCausalLM", the same name as MAX’s built-in Gemma 3. When you pass --custom-architectures weightless_fw, MAX imports the package, reads its ARCHITECTURES list, and the custom architecture takes precedence over the built-in of the same name.
The architecture wires together:
| Piece | Class / function |
|---|---|
| Pipeline model | WeightlessGemma3PipelineModel (builds the graph with the kernel as a custom extension) |
| Config | WeightlessGemma3Config |
| Weight adapter | convert_safetensor_state_dict (auto-decomposes at load) |
It declares default_encoding="bfloat16" (also supports float32), TEXT_GENERATION, a TextTokenizer, safetensors weights, and multi_gpu_supported=False.
The package hook
The package’s __init__.py builds ARCHITECTURES from the Gemma 3 module and then clears MAX’s lazy registration entry for each overridden name. Without that cleanup, MAX’s later lazy materialization would try to register the built-in architecture (without allow_override) and fail with “Refusing to override existing architecture.”
Config fields
WeightlessGemma3Config subclasses MAX’s Gemma3Config, inheriting every standard field (hidden size, layer count, head counts, rope settings, sliding window, …) and adding four weightless fields:
| Field | Default | Description |
|---|---|---|
lora_rank |
32 |
Rank of the per-(layer, family) LoRA delta. |
tie_families |
"all" |
"all" ties every family; "attn" ties only q/k/v/o. |
shared_blocks |
1 |
Number of contiguous layer bands G. |
tied_projections |
all seven families | Families reconstructed by the kernel. |
These are read from config.json with safe defaults, so an unmodified Gemma 3 repo loads without error. WEIGHTLESS_RANK and WEIGHTLESS_BLOCKS override lora_rank and shared_blocks at runtime (see runtime configuration).
band_index() returns the layer-to-band map used by W0 injection, partitioning [0, num_hidden_layers) into shared_blocks contiguous bands of near-equal size.
Adding another model
Mirror this directory: subclass the target model’s MAX config to add the weightless fields, write a safetensors adapter that decomposes into theta, build a pipeline model that injects W0 and uses WeightlessLinear, and register a SupportedArchitecture under the model’s canonical name. The core/ and kernels/ machinery is reused unchanged.
Gemma 4 31B (models/gemma4/) follows this pattern with additional handling for dual head dimensions (256 for sliding-window, 512 for full-attention), per-(family, layer_type) W0 grouping, and MultiKVCacheParams for the mixed attention types. See Gemma 4 quality for decomposition quality findings and the activation function fix.