metrale_model_arch/weight_loader/
gemma4.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Weight loader for Gemma-4. Every layer is an ungated attention
4//! layer with a GELU dense FFN, plus an MoE FFN when `num_experts > 0`.
5//! Layers are sliding or full attention (RoPE θ 10,000 or 1,000,000), and
6//! each has four norms and an optional `layer_scalar`. With no `lm_head`
7//! tensor the embedding serves as the head.
8//!
9//! Owner: model-arch weight loader.
10//! Invariants: none beyond the types.
11
12use anyhow::Result;
13use metrale_cache::kv_cache::KvCacheDtype;
14use metrale_config::ModelConfig;
15use metrale_gpu_runtime::gpu::GpuBackend;
16use metrale_model_weights::weights::WeightStore;
17
18use super::ModelWeightLoader;
19use metrale_model_layers::layer::TransformerLayer;
20use metrale_model_layers::weight_map::{DenseWeight, MtpWeights};
21
22mod loader_a;
23mod loader_b;
24
25pub struct Gemma4WeightLoader;
26
27impl ModelWeightLoader for Gemma4WeightLoader {
28    fn supports_tp(&self) -> bool {
29        // 2026-09-25: Under TP the loader shards the attention projections
30        // (see `loader_a::load_layers_impl`); the dense FFN weights are loaded
31        // whole on every rank.
32        true
33    }
34
35    fn load_layers(
36        &self,
37        store: &WeightStore,
38        config: &ModelConfig,
39        gpu: &dyn GpuBackend,
40        layer_kv_dtypes: &[KvCacheDtype],
41    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
42        loader_a::load_layers_impl(store, config, gpu, layer_kv_dtypes)
43    }
44
45    fn load_embedding(
46        &self,
47        store: &WeightStore,
48        config: &ModelConfig,
49        _gpu: &dyn GpuBackend,
50    ) -> Result<DenseWeight> {
51        loader_b::load_embedding_impl(store, config)
52    }
53
54    fn load_final_norm(
55        &self,
56        store: &WeightStore,
57        config: &ModelConfig,
58        _gpu: &dyn GpuBackend,
59    ) -> Result<DenseWeight> {
60        loader_b::load_final_norm_impl(store, config)
61    }
62
63    fn load_lm_head(
64        &self,
65        store: &WeightStore,
66        config: &ModelConfig,
67        _gpu: &dyn GpuBackend,
68    ) -> Result<DenseWeight> {
69        loader_b::load_lm_head_impl(store, config)
70    }
71
72    fn load_mtp_weights(
73        &self,
74        store: &WeightStore,
75        config: &ModelConfig,
76        gpu: &dyn GpuBackend,
77    ) -> Result<Option<MtpWeights>> {
78        loader_b::load_mtp_weights_impl(store, config, gpu)
79    }
80
81    fn kv_layer_dims(&self, config: &ModelConfig) -> Vec<(usize, usize)> {
82        loader_b::kv_layer_dims_impl(config)
83    }
84}
metrale_model_arch/weight_loader/gemma4.rs

metrale_model_arch/weight_loader/
gemma4.rs
