laya/
weights.rs

1//! Parameter-name mapping.
2//!
3//! Two checkpoint layouts have to load: the original PyTorch export
4//! (`convaiinnovations/laya`) and the pre-converted MLX export
5//! (`aac6fef/laya-mlx`). Upstream reconciles them in `sanitize_weights`; the
6//! same mapping is applied here, but as a *view* over the memory-mapped file
7//! rather than by materialising a renamed copy of every tensor.
8
9use candle_core::{DType, Device, Shape, Tensor};
10use candle_nn::var_builder::SimpleBackend;
11use candle_nn::Init;
12
13/// Map an upstream parameter name onto the name this port uses.
14///
15/// * PyTorch packs multi-head attention projections into `in_proj_weight` /
16///   `in_proj_bias`; here they are a normal `Linear` called `in_proj`.
17/// * `scorer` and `act_head` are `Sequential`s, whose children live under
18///   `layers.` — already true in MLX exports, not in PyTorch ones.
19pub fn sanitize_name(name: &str) -> String {
20    let mut name = name
21        .replace(".in_proj_weight", ".in_proj.weight")
22        .replace(".in_proj_bias", ".in_proj.bias");
23    for prefix in ["scorer", "act_head"] {
24        let dotted = format!("{prefix}.");
25        let nested = format!("{prefix}.layers.");
26        if name.starts_with(&dotted) && !name.starts_with(&nested) {
27            name = format!("{nested}{}", &name[dotted.len()..]);
28        }
29    }
30    name
31}
32
33/// A `SimpleBackend` that renames on lookup.
34///
35/// Rather than sanitizing every key up front, the inverse mapping is computed
36/// once at load time and consulted per parameter, so the underlying mmap is
37/// never copied.
38pub struct SanitizedBackend {
39    inner: Box<dyn SimpleBackend>,
40    /// Sanitized name -> name as stored in the file, only where they differ.
41    aliases: std::collections::HashMap<String, String>,
42}
43
44impl SanitizedBackend {
45    pub fn new(
46        inner: Box<dyn SimpleBackend>,
47        stored_names: impl IntoIterator<Item = String>,
48    ) -> Self {
49        let mut aliases = std::collections::HashMap::new();
50        for stored in stored_names {
51            let sanitized = sanitize_name(&stored);
52            if sanitized != stored {
53                aliases.insert(sanitized, stored);
54            }
55        }
56        Self { inner, aliases }
57    }
58
59    fn resolve<'a>(&'a self, name: &'a str) -> &'a str {
60        self.aliases.get(name).map_or(name, String::as_str)
61    }
62}
63
64impl SimpleBackend for SanitizedBackend {
65    fn get(
66        &self,
67        shape: Shape,
68        name: &str,
69        init: Init,
70        dtype: DType,
71        device: &Device,
72    ) -> candle_core::Result<Tensor> {
73        self.inner
74            .get(shape, self.resolve(name), init, dtype, device)
75    }
76
77    fn get_unchecked(
78        &self,
79        name: &str,
80        dtype: DType,
81        device: &Device,
82    ) -> candle_core::Result<Tensor> {
83        self.inner.get_unchecked(self.resolve(name), dtype, device)
84    }
85
86    fn contains_tensor(&self, name: &str) -> bool {
87        self.inner.contains_tensor(self.resolve(name))
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn rewrites_packed_attention_projections() {
97        assert_eq!(
98            sanitize_name("head.layers.0.self_attn.in_proj_weight"),
99            "head.layers.0.self_attn.in_proj.weight"
100        );
101        assert_eq!(
102            sanitize_name("head.layers.0.self_attn.in_proj_bias"),
103            "head.layers.0.self_attn.in_proj.bias"
104        );
105    }
106
107    #[test]
108    fn nests_sequential_children() {
109        assert_eq!(sanitize_name("scorer.0.weight"), "scorer.layers.0.weight");
110        assert_eq!(sanitize_name("act_head.2.bias"), "act_head.layers.2.bias");
111    }
112
113    #[test]
114    fn already_converted_names_are_untouched() {
115        for name in [
116            "scorer.layers.1.weight",
117            "act_head.layers.0.bias",
118            "encoder.layers.3.attn.Wqkv.weight",
119            "type_emb.weight",
120            "temperature",
121        ] {
122            assert_eq!(sanitize_name(name), name);
123        }
124    }
125}
laya/weights.rs

laya/
weights.rs
