spark_model/weight_loader/
nllb.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! NLLB / M2M-100 model-family marker loader.
4//!
5//! NLLB-200 checkpoints declare `model_type = "m2m_100"` and use an
6//! encoder-decoder architecture with learned absolute positions and decoder
7//! cross-attention. NLLB **is** served on the GPU, but through a dedicated
8//! encoder-decoder runtime (`crate::model::nllb`, `NllbGpuModel`) that
9//! `build_model` constructs *before* consulting this loader — see
10//! `crate::factory::build`. This marker exists only because the generic
11//! `ModelWeightLoader` table needs an entry for the model type; the generic
12//! `TransformerLayer`/paged-KV pipeline is decoder-only and cannot express the
13//! encoder self-attention + decoder cross-attention that NLLB requires, so
14//! every generic entry point here fails fast. Reaching one means the dedicated
15//! serve path was bypassed (a routing bug), not that NLLB is unsupported.
16
17use anyhow::{Result, bail};
18use atlas_core::config::ModelConfig;
19use spark_runtime::gpu::GpuBackend;
20use spark_runtime::kv_cache::KvCacheDtype;
21use spark_runtime::weights::WeightStore;
22
23use crate::layer::TransformerLayer;
24use crate::weight_loader::ModelWeightLoader;
25use crate::weight_map::{DenseWeight, MtpWeights};
26
27pub struct NllbWeightLoader;
28
29impl NllbWeightLoader {
30    fn unsupported() -> anyhow::Error {
31        anyhow::anyhow!(
32            "NLLB / m2m_100 is served by the dedicated GPU encoder-decoder runtime (spark_model::model::nllb::NllbGpuModel), which build_model selects before this loader; the generic decoder-only ModelWeightLoader pipeline cannot serve it. Reaching this loader means the dedicated serve path was bypassed."
33        )
34    }
35}
36
37impl ModelWeightLoader for NllbWeightLoader {
38    fn supports_tp(&self) -> bool {
39        false
40    }
41
42    fn load_layers(
43        &self,
44        _store: &WeightStore,
45        _config: &ModelConfig,
46        _gpu: &dyn GpuBackend,
47        _layer_kv_dtypes: &[KvCacheDtype],
48    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
49        bail!(Self::unsupported())
50    }
51
52    fn load_embedding(
53        &self,
54        _store: &WeightStore,
55        _config: &ModelConfig,
56        _gpu: &dyn GpuBackend,
57    ) -> Result<DenseWeight> {
58        bail!(Self::unsupported())
59    }
60
61    fn load_final_norm(
62        &self,
63        _store: &WeightStore,
64        _config: &ModelConfig,
65        _gpu: &dyn GpuBackend,
66    ) -> Result<DenseWeight> {
67        bail!(Self::unsupported())
68    }
69
70    fn load_lm_head(
71        &self,
72        _store: &WeightStore,
73        _config: &ModelConfig,
74        _gpu: &dyn GpuBackend,
75    ) -> Result<DenseWeight> {
76        bail!(Self::unsupported())
77    }
78
79    fn load_mtp_weights(
80        &self,
81        _store: &WeightStore,
82        _config: &ModelConfig,
83        _gpu: &dyn GpuBackend,
84    ) -> Result<Option<MtpWeights>> {
85        Ok(None)
86    }
87}
spark_model/weight_loader/nllb.rs

spark_model/weight_loader/
nllb.rs
