frink_models/
lora_attach.rs

1//! Attaching a parsed [`LoraAdapter`] to a [`Decoder`]: the walk from
2//! each `<base name>` in the adapter to the `WeightMatrix` (or
3//! matrices) frink holds for it, the shape checks llama.cpp makes
4//! against the base tensor (`llama-adapter.cpp:346-367`), and the
5//! refusals for what frink holds differently.
6//!
7//! llama.cpp keys an adapter on the base tensor's NAME: `get_weight(w)`
8//! looks `w->name` up in `ab_map` at every `build_lora_mm` call. A
9//! `WeightMatrix` carries no name (`frink_core::activation_tap` says
10//! why), so the name is resolved ONCE here, at attach, into the field
11//! the loader put that tensor in -- and where the loader split one
12//! file tensor into several matrices (a fused `attn_qkv.weight` into
13//! Q/K/V, Phi-3's fused `ffn_up.weight` into gate/up), the adapter's
14//! `lora_b` rows are split the same way, in the same order, with one
15//! copy of `lora_a` each: `B (A x)` over stacked rows IS the stacked
16//! `B_i (A x)`, exactly. Where the loader ALIASED one file tensor into
17//! two matrices (the ungated FFN's gate is its up, `loader.rs:1641`),
18//! the whole delta goes on both, so the two stay one tensor.
19//!
20//! What is refused, by name, rather than approximated:
21//!
22//! * a routed-expert tensor (`*_exps`): upstream applies those through
23//!   `build_lora_mm_id` and frink has no per-expert delta;
24//! * a tensor frink holds as something other than a `WeightMatrix`
25//!   (a norm vector, a bias): the delta would be dropped;
26//! * `token_embd.weight` on a base whose output head is tied to the
27//!   embedding: upstream aborts in `ggml_mul_mat` on that pair, so no
28//!   engine serves it;
29//! * an expert layer whose weights are leased from a store per use.
30//!
31//! Every fused Metal launch is fenced off a decoder with ANY adapter
32//! attached, through the one predicate they share
33//! (`Decoder::metal_can_serve_model`), because the stacks read weight
34//! bytes past the `WeightMatrix` methods that serve the delta. The
35//! per-matrix GPU launches (`apply_gpu`, `apply_gpu_multi`,
36//! `apply_gpu_batch`, CUDA's `mul_mm`) run the base on the device and
37//! add the rank-sized delta on the host, so Metal and CUDA still serve
38//! an adapted model -- on the per-matrix path.
39
40use std::path::{Path, PathBuf};
41use std::sync::Arc;
42
43use frink_core::weight_matrix::{LoraDelta, LoraScale, WeightMatrix};
44use frink_gguf::TensorSource;
45
46use crate::decoder::{Decoder, ExpertBacking};
47use crate::lora::{LoraAdapter, LoraError, LoraPair};
48
49/// One adapter after attach: what `GET /lora-adapters` lists and what
50/// `POST /lora-adapters` / a per-request `lora` list changes.
51#[derive(Debug, Clone)]
52pub struct LoraAttached {
53    pub path: PathBuf,
54    pub alpha: f32,
55    pub task_name: String,
56    pub prompt_prefix: String,
57    /// Shared by every delta this adapter attached.
58    pub scale: Arc<LoraScale>,
59    /// Base tensors this adapter decorates.
60    pub n_tensors: usize,
61}
62
63impl LoraAttached {
64    pub fn scale(&self) -> f32 {
65        self.scale.get()
66    }
67}
68
69/// One `--lora FNAME` or `--lora-scaled FNAME:SCALE`, before the file
70/// is opened. Parsed in ONE place for the CLI and the server, so the
71/// two cannot read `a:b:0.5` differently.
72#[derive(Debug, Clone, PartialEq)]
73pub struct LoraSpec {
74    pub path: PathBuf,
75    pub scale: f32,
76}
77
78impl LoraSpec {
79    /// `--lora FNAME`: scale 1, as `arg.cpp:2869` pushes it.
80    pub fn plain(path: impl Into<PathBuf>) -> Self {
81        Self {
82            path: path.into(),
83            scale: 1.0,
84        }
85    }
86
87    /// `--lora-scaled FNAME:SCALE` (`arg.cpp:2878-2885`). The LAST colon
88    /// splits, so a path with a drive letter or a colon in a directory
89    /// name still parses; llama.cpp splits on every colon and refuses
90    /// those.
91    pub fn parse_scaled(spec: &str) -> Result<Self, String> {
92        let (path, scale) = spec
93            .rsplit_once(':')
94            .ok_or_else(|| format!("lora-scaled format: FNAME:SCALE (got {spec:?})"))?;
95        let scale: f32 = scale
96            .trim()
97            .parse()
98            .map_err(|_| format!("lora-scaled format: FNAME:SCALE ({scale:?} is not a number)"))?;
99        if path.is_empty() {
100            return Err(format!("lora-scaled format: FNAME:SCALE (got {spec:?})"));
101        }
102        Ok(Self {
103            path: PathBuf::from(path),
104            scale,
105        })
106    }
107
108    /// The CLI's two flags as one list, in the order given: every
109    /// `--lora` value (comma-separated, as upstream's `parse_csv_row`)
110    /// then every `--lora-scaled` value.
111    pub fn from_flags(plain: &[String], scaled: &[String]) -> Result<Vec<Self>, String> {
112        let mut out = Vec::new();
113        for item in plain.iter().flat_map(|v| v.split(',')) {
114            let item = item.trim();
115            if !item.is_empty() {
116                out.push(Self::plain(item));
117            }
118        }
119        for item in scaled.iter().flat_map(|v| v.split(',')) {
120            let item = item.trim();
121            if !item.is_empty() {
122                out.push(Self::parse_scaled(item)?);
123            }
124        }
125        Ok(out)
126    }
127}
128
129impl Decoder {
130    /// Opens and attaches every spec in order. ONE function for the CLI
131    /// and the server, so a refusal reads the same from both.
132    pub fn attach_lora_specs(
133        &mut self,
134        base: &impl TensorSource,
135        specs: &[LoraSpec],
136    ) -> Result<(), LoraError> {
137        for spec in specs {
138            let adapter = LoraAdapter::open(&spec.path)?;
139            let id = self.attach_lora(base, adapter, spec.scale)?;
140            let a = &self.lora_adapters[id];
141            eprintln!(
142                "frink: lora adapter {id}: {} ({} tensor(s), alpha {}, scale {})",
143                spec.path.display(),
144                a.n_tensors,
145                a.alpha,
146                spec.scale
147            );
148        }
149        Ok(())
150    }
151}
152
153/// Where a base tensor's rows went.
154enum Parts<'a> {
155    /// One matrix, or several holding consecutive row blocks of the
156    /// file tensor, in file order.
157    Split(Vec<&'a mut WeightMatrix>),
158    /// Several matrices that are each the WHOLE file tensor.
159    Alias(Vec<&'a mut WeightMatrix>),
160}
161
162enum Target {
163    Embedding,
164    Output,
165    Layer(usize, LayerTarget),
166}
167
168enum LayerTarget {
169    Q,
170    K,
171    V,
172    Qkv,
173    O,
174    AttnGate,
175    FfnGate,
176    FfnUp,
177    FfnDown,
178    ShexpGate,
179    ShexpUp,
180    ShexpDown,
181    Router,
182}
183
184fn target_of(name: &str, path: &Path) -> Result<Target, LoraError> {
185    let no_projection = || LoraError::NoProjection {
186        path: path.to_path_buf(),
187        name: name.to_string(),
188    };
189    match name {
190        "token_embd.weight" => return Ok(Target::Embedding),
191        "output.weight" => return Ok(Target::Output),
192        _ => {}
193    }
194    let rest = name.strip_prefix("blk.").ok_or_else(no_projection)?;
195    let (il, tensor) = rest.split_once('.').ok_or_else(no_projection)?;
196    let il: usize = il.parse().map_err(|_| no_projection())?;
197    let t = tensor.strip_suffix(".weight").ok_or_else(no_projection)?;
198    if t.ends_with("_exps") {
199        return Err(LoraError::RoutedExperts {
200            path: path.to_path_buf(),
201            name: name.to_string(),
202        });
203    }
204    let kind = match t {
205        "attn_q" => LayerTarget::Q,
206        "attn_k" => LayerTarget::K,
207        "attn_v" => LayerTarget::V,
208        "attn_qkv" => LayerTarget::Qkv,
209        "attn_output" => LayerTarget::O,
210        "attn_gate" => LayerTarget::AttnGate,
211        "ffn_gate" => LayerTarget::FfnGate,
212        "ffn_up" => LayerTarget::FfnUp,
213        "ffn_down" => LayerTarget::FfnDown,
214        "ffn_gate_shexp" => LayerTarget::ShexpGate,
215        "ffn_up_shexp" => LayerTarget::ShexpUp,
216        "ffn_down_shexp" => LayerTarget::ShexpDown,
217        "ffn_gate_inp" => LayerTarget::Router,
218        _ => return Err(no_projection()),
219    };
220    Ok(Target::Layer(il, kind))
221}
222
223impl Decoder {
224    /// Is any adapter attached? The fact every fused Metal launch is
225    /// fenced on.
226    pub fn lora_attached(&self) -> bool {
227        !self.lora_adapters.is_empty()
228    }
229
230    /// Attaches `adapter` at `scale` and returns its id (its index in
231    /// [`Decoder::lora_adapters`]). `base` is the model's own GGUF: the
232    /// architecture and every tensor's presence and shape are checked
233    /// against it, as `llama_adapter_lora_init_impl` checks them against
234    /// the loaded model.
235    ///
236    /// On an error the decoder may already carry some of this adapter's
237    /// deltas; a caller must discard it rather than serve it, which is
238    /// what both the CLI and the server do (the load fails).
239    pub fn attach_lora(
240        &mut self,
241        base: &impl TensorSource,
242        adapter: LoraAdapter,
243        scale: f32,
244    ) -> Result<usize, LoraError> {
245        let path = adapter.path.to_path_buf();
246        let base_arch = base.metadata_str("general.architecture").unwrap_or("");
247        if adapter.arch != base_arch {
248            return Err(LoraError::ArchMismatch {
249                path,
250                adapter: adapter.arch,
251                base: base_arch.to_string(),
252            });
253        }
254        let handle = LoraScale::new(scale);
255        let n_tensors = adapter.pairs.len();
256        for (name, pair) in &adapter.pairs {
257            let info = base.find_tensor(name).ok_or_else(|| LoraError::NotInBase {
258                path: path.to_path_buf(),
259                name: name.clone(),
260            })?;
261            let target = target_of(name, &path)?;
262            // Checked HERE, in name order, rather than up front: on an
263            // adapter that also names `output.weight`, libllama's map
264            // walk reaches that pair's "does not exist in base model"
265            // first (`output.weight` sorts before `token_embd.weight`),
266            // and so does this loop.
267            if matches!(target, Target::Embedding) && base.find_tensor("output.weight").is_none() {
268                return Err(LoraError::TiedHead { path });
269            }
270            if info.shape.len() != 2 {
271                return Err(LoraError::NoProjection {
272                    path: path.to_path_buf(),
273                    name: name.clone(),
274                });
275            }
276            let (rows, cols) = (info.shape[1] as usize, info.shape[0] as usize);
277            let whole = checked_pair(&path, name, pair, &target, rows, cols)?;
278            let shape_err = || LoraError::Shape {
279                path: path.to_path_buf(),
280                name: name.clone(),
281                rows,
282                cols,
283                a: pair.a.shape,
284                b: pair.b.shape,
285            };
286            let alpha = adapter.alpha;
287            match self.parts_mut(&target, &path, name, rows)? {
288                Parts::Split(parts) => {
289                    // `B (A x)` over stacked rows is the stacked
290                    // `B_i (A x)`: each part takes its block of B's
291                    // rows and a copy of A.
292                    if parts.iter().map(|m| m.rows()).sum::<usize>() != rows
293                        || parts.iter().any(|m| m.cols() != cols)
294                    {
295                        return Err(shape_err());
296                    }
297                    let mut row0 = 0;
298                    for m in parts {
299                        let n = m.rows();
300                        let b = whole.b[row0 * whole.rank..(row0 + n) * whole.rank].to_vec();
301                        row0 += n;
302                        let delta = LoraDelta::new(
303                            whole.a.clone(),
304                            b,
305                            whole.rank,
306                            n,
307                            cols,
308                            alpha,
309                            Arc::clone(&handle),
310                        )
311                        .map_err(|_| shape_err())?;
312                        m.attach_lora(delta);
313                    }
314                }
315                Parts::Alias(parts) => {
316                    for m in parts {
317                        if m.rows() != rows || m.cols() != cols {
318                            return Err(shape_err());
319                        }
320                        let delta = LoraDelta::new(
321                            whole.a.clone(),
322                            whole.b.clone(),
323                            whole.rank,
324                            rows,
325                            cols,
326                            alpha,
327                            Arc::clone(&handle),
328                        )
329                        .map_err(|_| shape_err())?;
330                        m.attach_lora(delta);
331                    }
332                }
333            }
334        }
335        self.lora_adapters.push(LoraAttached {
336            path,
337            alpha: adapter.alpha,
338            task_name: adapter.task_name,
339            prompt_prefix: adapter.prompt_prefix,
340            scale: handle,
341            n_tensors,
342        });
343        Ok(self.lora_adapters.len() - 1)
344    }
345
346    /// Sets every adapter's scale from `scales` (`id -> scale`); an
347    /// adapter not listed goes to `0`, as `construct_lora_list`
348    /// (`server-context.cpp:1721-1732`) sets it. An unknown id is an
349    /// error naming the range, where upstream ignores it.
350    pub fn set_lora_scales(&self, scales: &[(usize, f32)]) -> Result<(), String> {
351        for &(id, _) in scales {
352            if id >= self.lora_adapters.len() {
353                return Err(format!(
354                    "lora adapter id {id} is out of range: {} adapter(s) loaded",
355                    self.lora_adapters.len()
356                ));
357            }
358        }
359        for (id, a) in self.lora_adapters.iter().enumerate() {
360            let s = scales
361                .iter()
362                .rev()
363                .find(|(i, _)| *i == id)
364                .map(|(_, s)| *s)
365                .unwrap_or(0.0);
366            a.scale.set(s);
367        }
368        Ok(())
369    }
370
371    /// The current scale of every adapter, by id.
372    pub fn lora_scales(&self) -> Vec<f32> {
373        self.lora_adapters.iter().map(LoraAttached::scale).collect()
374    }
375
376    /// The matrices holding base tensor `name`'s rows, and how.
377    fn parts_mut(
378        &mut self,
379        target: &Target,
380        path: &Path,
381        name: &str,
382        file_rows: usize,
383    ) -> Result<Parts<'_>, LoraError> {
384        let no_projection = || LoraError::NoProjection {
385            path: path.to_path_buf(),
386            name: name.to_string(),
387        };
388        let ungated = self.config.ffn_is_ungated();
389        Ok(match target {
390            Target::Embedding => Parts::Split(vec![&mut self.embedding]),
391            Target::Output => Parts::Split(vec![&mut self.output_head]),
392            Target::Layer(il, kind) => {
393                let layer = self.layers.get_mut(*il).ok_or_else(no_projection)?;
394                match kind {
395                    LayerTarget::Q => Parts::Split(vec![&mut layer.attn.q_proj]),
396                    LayerTarget::K => Parts::Split(vec![&mut layer.attn.k_proj]),
397                    LayerTarget::V => Parts::Split(vec![&mut layer.attn.v_proj]),
398                    LayerTarget::Qkv => Parts::Split(vec![
399                        &mut layer.attn.q_proj,
400                        &mut layer.attn.k_proj,
401                        &mut layer.attn.v_proj,
402                    ]),
403                    LayerTarget::O => Parts::Split(vec![&mut layer.attn.o_proj]),
404                    LayerTarget::AttnGate => match layer.attn.output_gate.as_mut() {
405                        Some(g) => Parts::Split(vec![&mut g.proj]),
406                        None => return Err(no_projection()),
407                    },
408                    LayerTarget::FfnGate | LayerTarget::FfnUp | LayerTarget::FfnDown => {
409                        let ex = match &mut layer.moe.experts {
410                            ExpertBacking::Resident(v) if v.len() == 1 => &mut v[0],
411                            ExpertBacking::Resident(_) => return Err(no_projection()),
412                            ExpertBacking::Stored { .. } => {
413                                return Err(LoraError::StoredExperts {
414                                    path: path.to_path_buf(),
415                                    name: name.to_string(),
416                                })
417                            }
418                        };
419                        match kind {
420                            LayerTarget::FfnGate => Parts::Split(vec![&mut ex.gate]),
421                            LayerTarget::FfnDown => Parts::Split(vec![&mut ex.down]),
422                            // `ffn_up` is three things across the loader:
423                            // the up matrix; Phi-3's fused gate+up (first
424                            // half gate, `loader.rs:1658`); or the ungated
425                            // FFN's up, which is ALSO its gate.
426                            _ if ungated => Parts::Alias(vec![&mut ex.up, &mut ex.gate]),
427                            _ if file_rows == ex.gate.rows() + ex.up.rows()
428                                && file_rows != ex.up.rows() =>
429                            {
430                                Parts::Split(vec![&mut ex.gate, &mut ex.up])
431                            }
432                            _ => Parts::Split(vec![&mut ex.up]),
433                        }
434                    }
435                    LayerTarget::ShexpGate | LayerTarget::ShexpUp | LayerTarget::ShexpDown => {
436                        let sh = layer
437                            .moe
438                            .shared_experts
439                            .first_mut()
440                            .ok_or_else(no_projection)?;
441                        Parts::Split(vec![match kind {
442                            LayerTarget::ShexpGate => &mut sh.gate,
443                            LayerTarget::ShexpUp => &mut sh.up,
444                            _ => &mut sh.down,
445                        }])
446                    }
447                    LayerTarget::Router => Parts::Split(vec![&mut layer.moe.router]),
448                }
449            }
450        })
451    }
452}
453
454/// A pair checked against the base tensor and brought into frink's one
455/// layout: `a` is `[rank][cols]`, `b` is `[rows][rank]`.
456struct CheckedPair {
457    a: Vec<f32>,
458    b: Vec<f32>,
459    rank: usize,
460}
461
462/// The shape checks of `llama-adapter.cpp:354-367`.
463fn checked_pair(
464    path: &Path,
465    name: &str,
466    pair: &LoraPair,
467    target: &Target,
468    rows: usize,
469    cols: usize,
470) -> Result<CheckedPair, LoraError> {
471    let shape_err = || LoraError::Shape {
472        path: path.to_path_buf(),
473        name: name.to_string(),
474        rows,
475        cols,
476        a: pair.a.shape,
477        b: pair.b.shape,
478    };
479    if matches!(target, Target::Embedding) {
480        // `:355-359`: B is `[n_embd, rank]`, A is `[n_vocab, rank]`
481        // (flipped and transposed by the converter), and the graph
482        // gathers a row of A per token and multiplies by B. In this
483        // module's one layout that row gather IS `lora_b`'s role and
484        // B transposed is `lora_a`'s.
485        if cols != pair.b.rows() || rows != pair.a.rows() {
486            return Err(shape_err());
487        }
488        let rank = pair.a.cols();
489        if rank != pair.b.cols() {
490            return Err(shape_err());
491        }
492        let n_embd = cols;
493        let mut a_eff = vec![0f32; rank * n_embd];
494        for e in 0..n_embd {
495            for k in 0..rank {
496                a_eff[k * n_embd + e] = pair.b.data[e * rank + k];
497            }
498        }
499        return Ok(CheckedPair {
500            a: a_eff,
501            b: pair.a.data.clone(),
502            rank,
503        });
504    }
505    // `:361-363`: `ne[0]` (cols) against A's, `ne[1]` (rows) against B's.
506    if cols != pair.a.cols() || rows != pair.b.rows() {
507        return Err(shape_err());
508    }
509    // `:364-366`: A's rank against B's.
510    let rank = pair.a.rows();
511    if rank != pair.b.cols() {
512        return Err(LoraError::NotTransposed {
513            path: path.to_path_buf(),
514            name: name.to_string(),
515            a: pair.a.shape,
516            b: pair.b.shape,
517        });
518    }
519    Ok(CheckedPair {
520        a: pair.a.data.clone(),
521        b: pair.b.data.clone(),
522        rank,
523    })
524}
525
526#[cfg(test)]
527mod tests {
528    use super::*;
529
530    #[test]
531    fn scaled_specs_split_on_the_last_colon() {
532        let s = LoraSpec::parse_scaled("a/b.gguf:0.5").unwrap();
533        assert_eq!(s.path, PathBuf::from("a/b.gguf"));
534        assert_eq!(s.scale, 0.5);
535        let s = LoraSpec::parse_scaled("C:/x/y.gguf:2").unwrap();
536        assert_eq!(s.path, PathBuf::from("C:/x/y.gguf"));
537        assert_eq!(s.scale, 2.0);
538        let s = LoraSpec::parse_scaled("z.gguf:-1").unwrap();
539        assert_eq!(s.scale, -1.0);
540        assert!(LoraSpec::parse_scaled("z.gguf")
541            .unwrap_err()
542            .contains("FNAME:SCALE"));
543        assert!(LoraSpec::parse_scaled("z.gguf:abc")
544            .unwrap_err()
545            .contains("not a number"));
546        assert!(LoraSpec::parse_scaled(":0.5").is_err());
547    }
548
549    #[test]
550    fn the_two_flags_become_one_ordered_list() {
551        let specs = LoraSpec::from_flags(
552            &["a.gguf,b.gguf".to_string(), "c.gguf".to_string()],
553            &["d.gguf:0.25".to_string()],
554        )
555        .unwrap();
556        assert_eq!(
557            specs,
558            vec![
559                LoraSpec::plain("a.gguf"),
560                LoraSpec::plain("b.gguf"),
561                LoraSpec::plain("c.gguf"),
562                LoraSpec {
563                    path: "d.gguf".into(),
564                    scale: 0.25
565                },
566            ]
567        );
568        assert!(LoraSpec::from_flags(&[], &["x".to_string()]).is_err());
569    }
570}
frink_models/lora_attach.rs

frink_models/
lora_attach.rs
