metrale_model_arch/weight_loader/
deepseek_v4.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: DeepSeek-V4 weight loader: layers, embedding, final norm and lm_head.
4//!
5//! Owner: model-arch weight loader.
6//! Invariants:
7//! - `supports_tp` is false: V4 layers shard experts by EP only.
8//! - `load_mtp_weights` returns `None`; the V4 MTP module is loaded separately
9//!   by `mtp::load_v4_mtp_module`.
10
11mod assemble;
12mod attn_sink;
13mod compute;
14mod csa_ape;
15mod load_layers;
16pub mod mtp;
17pub(crate) use mtp::DeepseekV4MtpModule;
18
19use anyhow::{Context, Result};
20use metrale_cache::kv_cache::KvCacheDtype;
21use metrale_config::ModelConfig;
22use metrale_gpu_runtime::gpu::GpuBackend;
23use metrale_model_weights::weights::WeightStore;
24
25use super::ModelWeightLoader;
26use metrale_model_layers::layer::TransformerLayer;
27use metrale_model_layers::weight_map::quant_helpers::dense_auto;
28use metrale_model_layers::weight_map::{DenseWeight, MtpWeights, dense};
29
30pub struct DeepSeekV4WeightLoader;
31
32impl ModelWeightLoader for DeepSeekV4WeightLoader {
33    fn supports_tp(&self) -> bool {
34        false
35    }
36
37    fn load_layers(
38        &self,
39        store: &WeightStore,
40        config: &ModelConfig,
41        gpu: &dyn GpuBackend,
42        layer_kv_dtypes: &[KvCacheDtype],
43    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
44        load_layers::load_all_layers(store, config, gpu, layer_kv_dtypes)
45    }
46
47    fn load_embedding(
48        &self,
49        store: &WeightStore,
50        _config: &ModelConfig,
51        _gpu: &dyn GpuBackend,
52    ) -> Result<DenseWeight> {
53        if let Ok(w) = dense(store, "embed.weight") {
54            return Ok(w);
55        }
56        if let Ok(w) = dense(store, "model.embed_tokens.weight") {
57            return Ok(w);
58        }
59        dense(store, "embed_tokens.weight")
60            .context("DeepSeek-V4: no embedding tensor found (tried embed.weight, model.embed_tokens.weight, embed_tokens.weight)")
61    }
62
63    fn load_final_norm(
64        &self,
65        store: &WeightStore,
66        _config: &ModelConfig,
67        _gpu: &dyn GpuBackend,
68    ) -> Result<DenseWeight> {
69        // 2026-09-25: Loaded as stored, with no offset: `deepseek_v4` is in
70        // `metrale_model_layers::ships_vanilla_norm_weights`, so the final norm runs
71        // `rms_norm_vanilla` (`out = x * w / rms`).
72        if let Ok(w) = dense_auto(store, "norm.weight", _gpu) {
73            return Ok(w);
74        }
75        if let Ok(w) = dense_auto(store, "model.norm.weight", _gpu) {
76            return Ok(w);
77        }
78        dense_auto(store, "final_norm.weight", _gpu)
79            .context("DeepSeek-V4: no final norm tensor found (tried norm.weight, model.norm.weight, final_norm.weight)")
80    }
81
82    fn load_lm_head(
83        &self,
84        store: &WeightStore,
85        config: &ModelConfig,
86        _gpu: &dyn GpuBackend,
87    ) -> Result<DenseWeight> {
88        if store.contains("lm_head.weight") {
89            return dense(store, "lm_head.weight");
90        }
91        if store.contains("output.weight") {
92            return dense(store, "output.weight");
93        }
94        if store.contains("head.weight") {
95            return dense(store, "head.weight");
96        }
97        if config.tie_word_embeddings
98            || store.contains("embed.weight")
99            || store.contains("model.embed_tokens.weight")
100        {
101            if let Ok(w) = dense(store, "embed.weight") {
102                return Ok(w);
103            }
104            if let Ok(w) = dense(store, "model.embed_tokens.weight") {
105                return Ok(w);
106            }
107            return dense(store, "embed_tokens.weight")
108                .context("DeepSeek-V4: tied lm_head — no embedding tensor found");
109        }
110        anyhow::bail!(
111            "DeepSeek-V4: lm_head not found (tried lm_head.weight, output.weight, head.weight, and tied embeddings)"
112        )
113    }
114
115    fn load_mtp_weights(
116        &self,
117        _store: &WeightStore,
118        _config: &ModelConfig,
119        _gpu: &dyn GpuBackend,
120    ) -> Result<Option<MtpWeights>> {
121        Ok(None)
122    }
123}
metrale_model_arch/weight_loader/deepseek_v4.rs

metrale_model_arch/weight_loader/
deepseek_v4.rs
