ferrotensor_graph/
model_builder.rs

1//! Model builder: constructs an [`Lfm2Model`] from a [`WeightsContext`].
2//!
3//! Loads Hugging Face safetensors checkpoints of two families into the same
4//! decoder structure, with optional LoRA adapters:
5//!
6//! * **Llama-style** (`model_type` `llama`/`mistral`, e.g. SmolLM2,
7//!   TinyLlama): every block is RoPE grouped-query attention plus a SwiGLU
8//!   MLP; weights named `self_attn.{q,k,v,o}_proj`, `mlp.{gate,up,down}_proj`,
9//!   `input_layernorm`, `post_attention_layernorm` and `model.norm`.
10//! * **LFM2** (`layer_types` present): blocks are short convolutions or
11//!   attention, named `operator_norm`, `ffn_norm`, `feed_forward.w{1,2,3}`,
12//!   with the final norm stored as `model.embedding_norm`.
13//!
14//! Base weights are frozen ([`Tensor::frozen`]) unless
15//! [`Lfm2ModelBuilder::with_trainable_base`] is set, so LoRA training skips
16//! their gradients.
17
18use crate::lfm::{Lfm2Block, Lfm2Conv, Lfm2LayerOp, Lfm2Model};
19use crate::{Attention, Linear, LoRAAdapter, RMSNorm, MLP};
20use ferrotensor_autograd::Tensor;
21use ferrotensor_backend::{ComputeEngine, UnifiedTensor, WgpuContext};
22use ferrotensor_loader::{Lfm2Config, WeightsContext};
23use std::sync::Arc;
24
25/// Configuration for LoRA adapters
26pub struct LoRaConfig {
27    /// Adapter rank.
28    pub rank: usize,
29    /// Scale numerator (`alpha / rank`).
30    pub alpha: f32,
31    /// Projection names that get adapters, e.g. `["q_proj", "v_proj"]`.
32    pub target_modules: Vec<String>, // e.g., ["q_proj", "v_proj"]
33}
34
35impl Default for LoRaConfig {
36    fn default() -> Self {
37        Self {
38            rank: 8,
39            alpha: 1.0,
40            target_modules: vec![
41                "q_proj".to_string(),
42                "k_proj".to_string(),
43                "v_proj".to_string(),
44                "o_proj".to_string(),
45                "in_proj".to_string(),
46                "out_proj".to_string(),
47            ],
48        }
49    }
50}
51
52/// Error returned when a checkpoint cannot be turned into a model.
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct BuildError(pub String);
55
56impl std::fmt::Display for BuildError {
57    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58        f.write_str(&self.0)
59    }
60}
61
62impl std::error::Error for BuildError {}
63
64type Result<T> = std::result::Result<T, BuildError>;
65
66/// Weight-naming family of a checkpoint.
67#[derive(Clone, Copy, Debug, PartialEq, Eq)]
68pub enum Architecture {
69    /// Llama / Mistral / SmolLM2 naming; every block is attention.
70    Llama,
71    /// Liquid LFM2 naming; conv and attention blocks per `layer_types`.
72    Lfm2,
73}
74
75impl Architecture {
76    /// Infer the family from a config: `layer_types` or `model_type: lfm2`
77    /// mean LFM2, everything else is treated as Llama.
78    pub fn detect(config: &Lfm2Config) -> Self {
79        let lfm = config
80            .model_type
81            .as_deref()
82            .is_some_and(|t| t.eq_ignore_ascii_case("lfm2"));
83        if lfm || !config.layer_types.is_empty() {
84            Self::Lfm2
85        } else {
86            Self::Llama
87        }
88    }
89}
90
91/// Builds an [`Lfm2Model`] from loaded weights.
92pub struct Lfm2ModelBuilder<'a> {
93    ctx: Arc<WgpuContext>,
94    engine: Arc<ComputeEngine>,
95    weights: &'a WeightsContext,
96    config: Lfm2Config,
97    lora_config: Option<LoRaConfig>,
98    trainable_base: bool,
99}
100
101impl<'a> Lfm2ModelBuilder<'a> {
102    fn transpose_2d(data: &[f32], rows: usize, cols: usize) -> Vec<f32> {
103        let mut transposed = vec![0.0; data.len()];
104        for row in 0..rows {
105            for col in 0..cols {
106                transposed[col * rows + row] = data[row * cols + col];
107            }
108        }
109        transposed
110    }
111
112    /// A builder for `config` reading tensors from `weights`.
113    pub fn new(
114        ctx: Arc<WgpuContext>,
115        engine: Arc<ComputeEngine>,
116        weights: &'a WeightsContext,
117        config: Lfm2Config,
118    ) -> Self {
119        Self {
120            ctx,
121            engine,
122            weights,
123            config,
124            lora_config: None,
125            trainable_base: false,
126        }
127    }
128
129    /// Attach LoRA adapters to the projections named in `lora_config`.
130    pub fn with_lora(mut self, lora_config: LoRaConfig) -> Self {
131        self.lora_config = Some(lora_config);
132        self
133    }
134
135    /// Let the base weights receive gradients (full or partial fine-tuning).
136    /// By default they are frozen and only LoRA adapters train.
137    pub fn with_trainable_base(mut self, trainable: bool) -> Self {
138        self.trainable_base = trainable;
139        self
140    }
141
142    fn wrap(&self, data: UnifiedTensor) -> Tensor {
143        if self.trainable_base {
144            Tensor::new(self.ctx.clone(), data)
145        } else {
146            Tensor::frozen(self.ctx.clone(), data)
147        }
148    }
149
150    fn has(&self, name: &str) -> bool {
151        self.weights.get_tensor_info(name).is_some()
152    }
153
154    fn read(&self, name: &str) -> Result<(Vec<f32>, Vec<usize>)> {
155        self.weights
156            .read_f32_tensor(name)
157            .map_err(|e| BuildError(format!("failed to load weight {name}: {e}")))
158    }
159
160    /// Load a tensor from weights, dequantizing if needed
161    fn load_tensor(&self, name: &str) -> Result<Tensor> {
162        if self.has(&format!("{name}.scales")) {
163            let ut = self
164                .weights
165                .upload_quantized(name, &self.ctx, &self.engine)
166                .map_err(|e| BuildError(format!("failed to load quantized weight {name}: {e}")))?;
167            Ok(self.wrap(ut))
168        } else {
169            self.load_raw_tensor(&format!("{name}.weight"))
170        }
171    }
172
173    fn load_raw_tensor(&self, name: &str) -> Result<Tensor> {
174        let (data, shape) = self.read(name)?;
175        Ok(self.wrap(UnifiedTensor::new(&self.ctx, &data, &shape)))
176    }
177
178    /// Load a PyTorch `[out, in]` weight transposed to `[in, out]`.
179    fn load_linear_weight_tensor(&self, prefix: &str) -> Result<Tensor> {
180        let (data, shape) = self.read(&format!("{prefix}.weight"))?;
181        let (data, shape) = if let [rows, cols] = shape[..] {
182            (Self::transpose_2d(&data, rows, cols), vec![cols, rows])
183        } else {
184            (data, shape)
185        };
186        Ok(self.wrap(UnifiedTensor::new(&self.ctx, &data, &shape)))
187    }
188
189    fn maybe_add_lora(&self, linear: Linear, module_name: &str, seed: u64) -> Linear {
190        if let Some(lora_cfg) = &self.lora_config {
191            if lora_cfg
192                .target_modules
193                .iter()
194                .any(|m| module_name.contains(m.as_str()))
195            {
196                let shape = &linear.weight.data.shape;
197                let lora = LoRAAdapter::new_seeded(
198                    self.ctx.clone(),
199                    shape[0],
200                    shape[1],
201                    lora_cfg.rank,
202                    lora_cfg.alpha,
203                    seed,
204                );
205                return Linear::with_lora(linear.weight, lora);
206            }
207        }
208        linear
209    }
210
211    /// `module_name` is matched against the LoRA target modules; `seed`
212    /// keeps adapters of equal shape distinct.
213    fn build_linear(&self, prefix: &str, module_name: &str, seed: u64) -> Result<Linear> {
214        let weight = if self.has(&format!("{prefix}.scales")) {
215            self.load_tensor(prefix)?
216        } else {
217            self.load_linear_weight_tensor(prefix)?
218        };
219        Ok(self.maybe_add_lora(Linear::new(weight), module_name, seed))
220    }
221
222    fn build_rms_norm(&self, prefix: &str) -> Result<RMSNorm> {
223        let weight = self.load_raw_tensor(&format!("{prefix}.weight"))?;
224        Ok(RMSNorm::new(weight, self.config.norm_eps))
225    }
226
227    fn build_mlp(&self, prefix: &str, arch: Architecture, seed: u64) -> Result<MLP> {
228        // LFM2 names its SwiGLU projections w1 (gate), w3 (up), w2 (down).
229        let [gate, up, down] = match arch {
230            Architecture::Llama => ["gate_proj", "up_proj", "down_proj"],
231            Architecture::Lfm2 => ["w1", "w3", "w2"],
232        };
233        Ok(MLP {
234            gate_proj: self.build_linear(&format!("{prefix}.{gate}"), gate, seed)?,
235            up_proj: self.build_linear(&format!("{prefix}.{up}"), up, seed)?,
236            down_proj: self.build_linear(&format!("{prefix}.{down}"), down, seed)?,
237        })
238    }
239
240    fn build_attention(&self, prefix: &str, arch: Architecture, seed: u64) -> Result<Attention> {
241        let n_heads = self.config.num_attention_heads;
242        let num_kv_heads = self.config.num_key_value_heads.unwrap_or(n_heads);
243        let head_dim = self
244            .config
245            .head_dim
246            .unwrap_or(self.config.hidden_size / n_heads.max(1));
247        if n_heads == 0 || num_kv_heads == 0 || !n_heads.is_multiple_of(num_kv_heads) {
248            return Err(BuildError(format!(
249                "num_attention_heads ({n_heads}) must be a positive multiple of num_key_value_heads ({num_kv_heads})"
250            )));
251        }
252        if head_dim > 128 {
253            return Err(BuildError(format!(
254                "head_dim {head_dim} is not supported (the attention kernel handles at most 128)"
255            )));
256        }
257        if self.has(&format!("{prefix}.q_proj.bias")) {
258            return Err(BuildError(
259                "attention projections with bias (e.g. Qwen2) are not supported yet".into(),
260            ));
261        }
262        if self.has(&format!("{prefix}.q_layernorm.weight")) {
263            return Err(BuildError(
264                "LFM2 per-head q/k layernorm is not supported yet".into(),
265            ));
266        }
267        let o_name = match arch {
268            Architecture::Llama => "o_proj",
269            Architecture::Lfm2 => "out_proj",
270        };
271
272        Ok(Attention {
273            q_proj: self.build_linear(&format!("{prefix}.q_proj"), "q_proj", seed)?,
274            k_proj: self.build_linear(&format!("{prefix}.k_proj"), "k_proj", seed)?,
275            v_proj: self.build_linear(&format!("{prefix}.v_proj"), "v_proj", seed)?,
276            o_proj: self.build_linear(&format!("{prefix}.{o_name}"), o_name, seed)?,
277            n_heads,
278            head_dim,
279            num_kv_heads,
280            dropout_prob: self.config.attention_dropout.unwrap_or(0.0),
281            rope_theta: self.config.rope_theta.unwrap_or(10000.0) as f32,
282        })
283    }
284
285    fn build_conv(&self, prefix: &str, seed: u64) -> Result<Lfm2Conv> {
286        let conv_weight = self.load_raw_tensor(&format!("{prefix}.conv.weight"))?;
287        // LFM uses no conv bias typically, create zero tensor
288        let conv_bias = self.wrap(UnifiedTensor::zeros(&self.ctx, &[self.config.hidden_size]));
289
290        Ok(Lfm2Conv {
291            in_proj: self.build_linear(&format!("{prefix}.in_proj"), "in_proj", seed)?,
292            conv_weight,
293            conv_bias,
294            out_proj: self.build_linear(&format!("{prefix}.out_proj"), "out_proj", seed)?,
295        })
296    }
297
298    fn build_layer(&self, layer_idx: usize, arch: Architecture) -> Result<Lfm2Block> {
299        let prefix = format!("model.layers.{layer_idx}");
300        let seed = layer_idx as u64 + 1;
301        match arch {
302            Architecture::Llama => Ok(Lfm2Block {
303                operator_norm: self.build_rms_norm(&format!("{prefix}.input_layernorm"))?,
304                op: Lfm2LayerOp::Attention(self.build_attention(
305                    &format!("{prefix}.self_attn"),
306                    arch,
307                    seed,
308                )?),
309                ffn_norm: self.build_rms_norm(&format!("{prefix}.post_attention_layernorm"))?,
310                mlp: self.build_mlp(&format!("{prefix}.mlp"), arch, seed)?,
311            }),
312            Architecture::Lfm2 => {
313                let layer_type = self
314                    .config
315                    .layer_types
316                    .get(layer_idx)
317                    .map(String::as_str)
318                    .unwrap_or("full_attention");
319                let op = match layer_type {
320                    "conv" => Lfm2LayerOp::Conv(self.build_conv(&format!("{prefix}.conv"), seed)?),
321                    "full_attention" | "attention" => Lfm2LayerOp::Attention(
322                        self.build_attention(&format!("{prefix}.self_attn"), arch, seed)?,
323                    ),
324                    other => {
325                        return Err(BuildError(format!(
326                            "unknown layer type {other:?} at layer {layer_idx}"
327                        )))
328                    }
329                };
330                Ok(Lfm2Block {
331                    operator_norm: self.build_rms_norm(&format!("{prefix}.operator_norm"))?,
332                    op,
333                    ffn_norm: self.build_rms_norm(&format!("{prefix}.ffn_norm"))?,
334                    mlp: self.build_mlp(&format!("{prefix}.feed_forward"), arch, seed)?,
335                })
336            }
337        }
338    }
339
340    /// Build the complete model, or explain which weight or setting is
341    /// missing or unsupported.
342    pub fn build(self) -> Result<Lfm2Model> {
343        let arch = Architecture::detect(&self.config);
344        let n_layers = self.config.num_hidden_layers;
345        log::info!("building {arch:?} model with {n_layers} layers");
346
347        let embed_tokens = self.load_tensor("model.embed_tokens")?;
348        let mut layers = Vec::with_capacity(n_layers);
349        for i in 0..n_layers {
350            layers.push(self.build_layer(i, arch)?);
351        }
352
353        // Hugging Face LFM2 applies `embedding_norm` after the last block.
354        let norm_f = if self.has("model.norm.weight") {
355            self.build_rms_norm("model.norm")?
356        } else if self.has("model.embedding_norm.weight") {
357            self.build_rms_norm("model.embedding_norm")?
358        } else {
359            return Err(BuildError(
360                "no final norm (model.norm.weight or model.embedding_norm.weight)".into(),
361            ));
362        };
363
364        let lm_head = if self.has("lm_head.weight") {
365            self.build_linear("lm_head", "lm_head", 0)?
366        } else {
367            // Tied embeddings: the head is embed_tokens^T, `[hidden, vocab]`.
368            let (data, shape) = self.read("model.embed_tokens.weight")?;
369            let [vocab, hidden] = shape[..] else {
370                return Err(BuildError(format!(
371                    "model.embed_tokens.weight must be 2-D, got {shape:?}"
372                )));
373            };
374            let weight = Self::transpose_2d(&data, vocab, hidden);
375            let weight = self.wrap(UnifiedTensor::new(&self.ctx, &weight, &[hidden, vocab]));
376            self.maybe_add_lora(Linear::new(weight), "lm_head", 0)
377        };
378
379        Ok(Lfm2Model::new(embed_tokens, None, layers, norm_f, lm_head))
380    }
381}
382
383/// Load the model in `model_path` (`config.json` plus safetensors).
384pub fn load_lfm2_model(
385    model_path: &str,
386    ctx: Arc<WgpuContext>,
387    engine: Arc<ComputeEngine>,
388    lora_config: Option<LoRaConfig>,
389) -> std::result::Result<Lfm2Model, String> {
390    use std::path::Path;
391
392    let model_dir = Path::new(model_path);
393    let config = ferrotensor_loader::ModelLoader::load_config(model_dir.join("config.json"))
394        .map_err(|e| format!("Failed to load config: {}", e))?;
395
396    let mut weights = WeightsContext::new();
397    weights
398        .load_model_dir(model_dir)
399        .map_err(|e| format!("Failed to load safetensors: {}", e))?;
400
401    let mut builder = Lfm2ModelBuilder::new(ctx, engine, &weights, config);
402    if let Some(lora) = lora_config {
403        builder = builder.with_lora(lora);
404    }
405
406    builder.build().map_err(|e| e.to_string())
407}
408
409#[cfg(test)]
410mod tests {
411    use super::Lfm2ModelBuilder;
412
413    #[test]
414    fn transpose_2d_swaps_rows_and_columns() {
415        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
416        let transposed = Lfm2ModelBuilder::transpose_2d(&data, 2, 3);
417        assert_eq!(transposed, vec![1.0, 4.0, 2.0, 5.0, 3.0, 6.0]);
418    }
419}
ferrotensor_graph/model_builder.rs

ferrotensor_graph/
model_builder.rs
