spark_model/weight_loader/
gemma4.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Weight loader for Gemma-4 models.
4//!
5//! Gemma-4 is a pure-attention model (no SSM) with:
6//! - Sliding + full attention pattern (dual RoPE theta)
7//! - GeGLU activation (gate + up projection fused)
8//! - Explicit head_dim=256
9//! - Tied word embeddings
10//! - 4 layer norms per layer (input, post_attn, pre_ffn, post_ffn)
11//! - Per-layer scalar (`layer_scalar`)
12//! - BF16 attention weights, NVFP4 MLP weights (Standard triple-scale format)
13//!
14//! Weight prefix: `model.language_model.` (auto-detected from safetensors keys
15//! by the main loader when `nested_config = true`).
16
17use anyhow::Result;
18use atlas_core::config::ModelConfig;
19use spark_runtime::gpu::GpuBackend;
20use spark_runtime::kv_cache::KvCacheDtype;
21use spark_runtime::weights::WeightStore;
22
23use super::ModelWeightLoader;
24use crate::layer::TransformerLayer;
25use crate::weight_map::{DenseWeight, MtpWeights};
26
27mod loader_a;
28mod loader_b;
29
30pub struct Gemma4WeightLoader;
31
32impl ModelWeightLoader for Gemma4WeightLoader {
33    fn supports_tp(&self) -> bool {
34        // Gemma-4 has per-layer dim overrides: sliding layers use
35        // head_dim=256 explicitly, full-attention layers use whatever
36        // the on-disk weight's [out, in] shape declares. We slice
37        // each layer's q/k/v BF16 weights with the layer-specific
38        // full dims (read from store), then pass the per-rank-local
39        // dims to quantize_to_nvfp4. K=V aliasing in sliding layers
40        // works under TP because both K and V share the same per-rank
41        // slice pointer (the alias survives sharding).
42        // MLP weights stay full-replica per rank — Gemma-4 31B is
43        // dense; sharding the GeGLU MLP would double the TP win but
44        // requires NVFP4 byte-slicing in `quantized_any`-loaded paths
45        // (deferred — leaves Gemma-4 functionally correct under TP
46        // with extra MLP memory per rank).
47        true
48    }
49
50    fn load_layers(
51        &self,
52        store: &WeightStore,
53        config: &ModelConfig,
54        gpu: &dyn GpuBackend,
55        layer_kv_dtypes: &[KvCacheDtype],
56    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
57        loader_a::load_layers_impl(store, config, gpu, layer_kv_dtypes)
58    }
59
60    fn load_embedding(
61        &self,
62        store: &WeightStore,
63        config: &ModelConfig,
64        _gpu: &dyn GpuBackend,
65    ) -> Result<DenseWeight> {
66        loader_b::load_embedding_impl(store, config)
67    }
68
69    fn load_final_norm(
70        &self,
71        store: &WeightStore,
72        config: &ModelConfig,
73        _gpu: &dyn GpuBackend,
74    ) -> Result<DenseWeight> {
75        loader_b::load_final_norm_impl(store, config)
76    }
77
78    fn load_lm_head(
79        &self,
80        store: &WeightStore,
81        config: &ModelConfig,
82        _gpu: &dyn GpuBackend,
83    ) -> Result<DenseWeight> {
84        loader_b::load_lm_head_impl(store, config)
85    }
86
87    fn load_mtp_weights(
88        &self,
89        store: &WeightStore,
90        config: &ModelConfig,
91        gpu: &dyn GpuBackend,
92    ) -> Result<Option<MtpWeights>> {
93        loader_b::load_mtp_weights_impl(store, config, gpu)
94    }
95
96    fn kv_layer_dims(&self, config: &ModelConfig) -> Vec<(usize, usize)> {
97        loader_b::kv_layer_dims_impl(config)
98    }
99}
spark_model/weight_loader/gemma4.rs

spark_model/weight_loader/
gemma4.rs
