1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
//! **PER-TENSOR WEIGHT SCALES** -- the optional `<tensor>.scale` and
//! `<tensor>.input_scale` companions llama.cpp multiplies a projection's
//! output by, which frink does not apply and therefore refuses by name.
//!
//! # What they are
//!
//! `build_lora_mm(w, cur, w_s)` (`llama-graph.cpp:1486-1494`) is
//! `res = mul_mat(w, cur); if (w_s) res = mul(res, w_s)`, and its
//! `mul_mat_id` twin does the same per expert (`:1521-1530`). Since
//! `llama-model.cpp:1355-1440` a GENERIC pass after every architecture's
//! `load_arch_tensors` creates `{1}` `scale` and `input_scale` tensors
//! as `TENSOR_NOT_REQUIRED` beside every projection it finds -- Q, K,
//! V, O, fused QKV, the attention gate, dense gate/up/down, shared
//! experts, routed experts (`{n_expert}`), the recurrent projections,
//! the NextN head -- and `:1506-1507` does it for `output`. So ANY
//! architecture's file may carry them, and llama.cpp honours them on
//! all of them.
//!
//! # Who writes them
//!
//! Two converters. `conversion/base.py:687,780` write `.scale` (the
//! NVFP4 per-tensor `scale2`) and `.input_scale` beside NVFP4-packed
//! weights, a quantization type frink has no kernel for, so such a
//! file stops on the dtype before it reaches here. And BitNet: `src/
//! models/bitnet.cpp:27-43` created the seven per-projection scales
//! BEFORE the generic pass existed, for the exports whose ternary
//! weights were stored unscaled; the CURRENT `conversion/bitnet.py:
//! 23-32` folds the scale into the weights and writes no tensor, but
//! older files carry them and libllama still applies them -- measured,
//! `tests/sub_norm_graphs.rs`: the same fixture with and without seven
//! `2.0` scales gives DIFFERENT logits under libllama. And `talkie`
//! (`conversion/talkie.py:26-31`), which writes `attn_output.scale` and
//! `ffn_down.scale` from its `attn_gain` / `mlp_gain` on every export
//! and closed on exactly those two: [`SERVED_SCALE_TENSORS`] are the
//! companions the decoder applies (`AttnWeights::o_scale`,
//! `MoeWeights::down_scale`), for any architecture, as upstream.
//!
//! # Why a refusal by name
//!
//! Without this, such a file died on the unread-tensor gate
//! (`loader::assert_every_tensor_consumed`) -- the right outcome with
//! the wrong message, and one `FRINK_ALLOW_UNKNOWN_TENSORS=1` turns
//! into a model running every projection at the wrong magnitude. A
//! multiply per projection is not hard to add; it is not added because
//! no file frink can otherwise run needs it today, and a seam with no
//! caller is the OLMo lesson (`capability::WEIGHTED_LAYER_NORM`). When
//! one does, this module is where the census already is.
use crate::LoadError;
/// The suffixes `llama-model.cpp:1355-1440` create under every
/// projection name.
const SCALE_SUFFIXES: [&str; 2] = [".scale", ".input_scale"];
/// The two companions frink APPLIES: `blk.N.attn_output.scale` (`wo_s`,
/// multiplied onto the attention branch after `wo`) and
/// `blk.N.ffn_down.scale` (`ffn_down_s`, onto the dense FFN after
/// `down`). They are the two `talkie`'s converter writes on every
/// export (`conversion/talkie.py:26-31`), and `build_lora_mm`'s
/// `res = mul(mul_mat(w, x), s)` is the same multiply for both; every
/// other companion is still refused below, because nothing frink can
/// otherwise run writes it.
pub const SERVED_SCALE_TENSORS: [&str; 2] = ["attn_output.scale", "ffn_down.scale"];
/// Whether `name` is one of the two companions the decoder applies.
fn is_served(name: &str) -> bool {
SERVED_SCALE_TENSORS.iter().any(|s| {
name.strip_prefix("blk.")
.and_then(|rest| rest.split_once('.'))
.is_some_and(|(_, tail)| tail == *s)
})
}
/// Layer `l`'s `blk.N.<proj>.scale`, a `{1}` tensor, when the file has
/// one. Optional everywhere, as `TENSOR_NOT_REQUIRED` upstream.
pub fn load_projection_gain(
file: &impl frink_gguf::TensorSource,
arch: &str,
l: usize,
proj: &str,
) -> Result<Option<f32>, LoadError> {
let name = format!("blk.{l}.{proj}.scale");
if file.find_tensor(&name).is_none() {
return Ok(None);
}
let v = crate::loader::load_f32_vec(file, &name)?;
if v.len() != 1 {
return Err(LoadError::UnsupportedFeature(
arch.to_string(),
format!("{name} has {} entries; a per-tensor scale is one", v.len()),
));
}
Ok(Some(v[0]))
}
/// Refuses a file carrying any per-tensor weight scale, naming the
/// first few.
///
/// `names` is every tensor in the file; the loader calls this once,
/// before it reads a layer, so the refusal names the feature rather
/// than the unread tensor.
pub fn refuse_weight_scale_tensors<'a>(
arch: &str,
names: impl Iterator<Item = &'a str>,
) -> Result<(), LoadError> {
let mut scaled: Vec<&str> = names
.filter(|n| SCALE_SUFFIXES.iter().any(|s| n.ends_with(s)) && !is_served(n))
.collect();
if scaled.is_empty() {
return Ok(());
}
scaled.sort_unstable();
let shown = scaled
.iter()
.take(4)
.cloned()
.collect::<Vec<_>>()
.join(", ");
let listing = if scaled.len() > 4 {
format!("{shown}, … (+{} more)", scaled.len() - 4)
} else {
shown
};
Err(LoadError::UnsupportedFeature(
arch.to_string(),
format!(
"per-tensor weight scales ({} tensor(s): {listing}): llama.cpp multiplies each \
projection's output by its `.scale` / `.input_scale` companion \
(`build_lora_mm`, llama-graph.cpp:1492-1494), which frink does not apply; \
running the file without them would put every scaled projection at the \
wrong magnitude. Re-export without per-tensor scales (the current BitNet \
converter folds them into the weights)",
scaled.len()
),
))
}
#[cfg(test)]
mod tests {
use super::*;
/// The two `talkie` writes pass; the others beside them still do not.
#[test]
fn the_two_served_companions_pass_and_the_rest_are_still_refused() {
let ok = [
"blk.0.attn_output.scale",
"blk.3.ffn_down.scale",
"blk.0.attn_q.weight",
];
assert!(refuse_weight_scale_tensors("talkie", ok.into_iter()).is_ok());
let err = refuse_weight_scale_tensors(
"talkie",
["blk.0.attn_output.scale", "blk.0.attn_q.scale"].into_iter(),
)
.expect_err("attn_q.scale is not applied");
let msg = err.to_string();
assert!(
msg.contains("blk.0.attn_q.scale") && !msg.contains("attn_output.scale"),
"{msg}"
);
}
#[test]
fn a_file_without_scale_companions_passes() {
let names = [
"token_embd.weight",
"blk.0.attn_q.weight",
"output_norm.weight",
];
assert!(refuse_weight_scale_tensors("bitnet", names.into_iter()).is_ok());
}
/// Both suffixes, on a layer tensor and on the lm_head, and the
/// message names the tensor and the feature.
#[test]
fn either_suffix_anywhere_is_refused_by_name() {
for name in [
"blk.3.ffn_gate.scale",
"blk.0.attn_q.input_scale",
"output.scale",
] {
let err =
refuse_weight_scale_tensors("llama", [name, "blk.0.attn_q.weight"].into_iter())
.expect_err(name);
let msg = err.to_string();
assert!(msg.contains(name), "{msg}");
assert!(msg.contains("per-tensor weight scales"), "{msg}");
}
}
/// A weight whose NAME merely contains "scale" is not a scale
/// companion; only the exact suffixes are.
#[test]
fn the_suffix_is_matched_not_the_substring() {
let names = ["blk.0.attn_scale_gate.weight", "blk.0.scale_norm.weight"];
assert!(refuse_weight_scale_tensors("llama", names.into_iter()).is_ok());
}
}