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
//! Sidecar metadata for baked checkpoints + a pure VRAM auto-size estimator.
//!
//! The sidecar (`<checkpoint>.meta.json`) records the dims a checkpoint was
//! trained with so the proxy/eval load the right model without hardcoding.
//! The auto-size estimator picks the largest model config that fits a memory
//! budget, so "auto-size to VRAM" works end-to-end.
use std::path::Path;
use serde::{Deserialize, Serialize};
/// Persisted alongside a checkpoint as `<path>.meta.json`.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ModelMeta {
pub d_model: usize,
pub n_layers: usize,
pub vocab_size: usize,
pub lr_inner: f32,
pub norm_eps: f32,
/// Best held-out cross-entropy seen during training (eval signal).
pub val_ce: f32,
/// Tokenizer file this model was trained against.
pub tokenizer: String,
/// Whether this checkpoint was trained with inner-loop stabilization
/// (normalized keys + state clamp). The proxy/eval must run it the same way.
/// `#[serde(default)]` → old sidecars (no field) load as `false`.
#[serde(default)]
pub stabilize: bool,
/// Whether training used the detached-chunk final-token loss fast path.
/// This is training provenance only; inference remains unchanged.
#[serde(default)]
pub last_token_only: bool,
/// Whether this checkpoint was trained with the *learned* data-dependent
/// Gated-DeltaNet forget gate (a per-layer `w_α` projection that decides how
/// much memory to retain per token). It adds weights, so the model must be
/// built with it enabled to load such a checkpoint. `#[serde(default)]` → old
/// sidecars load as `false`, so the parameter-free path is used and existing
/// checkpoints are unaffected.
#[serde(default)]
pub learned_gate: bool,
}
impl ModelMeta {
/// Sidecar path for a checkpoint path: `foo.bin` -> `foo.meta.json`.
pub fn sidecar_path(checkpoint: &str) -> String {
let p = Path::new(checkpoint);
let stem = p.file_stem().and_then(|s| s.to_str()).unwrap_or("model");
match p.parent() {
Some(dir) if !dir.as_os_str().is_empty() => {
format!("{}/{stem}.meta.json", dir.to_string_lossy())
}
_ => format!("{stem}.meta.json"),
}
}
pub fn save(&self, checkpoint: &str) -> std::io::Result<()> {
let json = serde_json::to_string_pretty(self).expect("serialize meta");
std::fs::write(Self::sidecar_path(checkpoint), json)
}
pub fn load(checkpoint: &str) -> Option<ModelMeta> {
let txt = std::fs::read_to_string(Self::sidecar_path(checkpoint)).ok()?;
serde_json::from_str(&txt).ok()
}
}
/// One rung of the auto-size ladder.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct SizeRung {
pub d_model: usize,
pub n_layers: usize,
}
/// Default ladder, smallest → largest. Ceiling tuned for ~3.7 GB usable VRAM.
pub fn default_ladder() -> Vec<SizeRung> {
vec![
SizeRung {
d_model: 256,
n_layers: 4,
},
SizeRung {
d_model: 384,
n_layers: 6,
},
SizeRung {
d_model: 512,
n_layers: 8,
},
SizeRung {
d_model: 640,
n_layers: 8,
},
]
}
/// Estimate the training footprint (bytes) of a config: embedding + lm_head +
/// per-layer projections, times 4 (params + grad + AdamW m + v, all fp32),
/// plus a flat activation budget for one `win`-token forward.
pub fn estimate_footprint_bytes(d_model: usize, n_layers: usize, vocab: usize, win: usize) -> u64 {
let params = 2 * vocab * d_model // embedding + lm_head
+ n_layers * 3 * d_model * d_model // w_q, w_k, w_v per layer
+ n_layers * d_model // layer norms (approx)
+ d_model; // final norm
let param_bytes = params as u64 * 4 * 4; // fp32 × (param + grad + m + v)
let activation_bytes = (win as u64) * (vocab as u64) * 4 * 3; // logits + softmax scratch
param_bytes + activation_bytes
}
/// Pick the largest rung that fits `budget_bytes`. Always returns at least the
/// smallest rung, so training never refuses to start.
pub fn pick_config(budget_bytes: u64, vocab: usize, win: usize, ladder: &[SizeRung]) -> SizeRung {
let mut chosen = ladder[0];
for rung in ladder {
let fp = estimate_footprint_bytes(rung.d_model, rung.n_layers, vocab, win);
if fp <= budget_bytes {
chosen = *rung;
}
}
chosen
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn sidecar_path_swaps_extension() {
assert_eq!(
ModelMeta::sidecar_path("checkpoints/axiom_production_bpe.bin"),
"checkpoints/axiom_production_bpe.meta.json"
);
assert_eq!(ModelMeta::sidecar_path("model.bin"), "model.meta.json");
}
#[test]
fn meta_roundtrips_through_disk() {
let m = ModelMeta {
d_model: 512,
n_layers: 8,
vocab_size: 32000,
lr_inner: 1e-3,
norm_eps: 1e-6,
val_ce: 3.21,
tokenizer: "t.json".into(),
stabilize: true,
last_token_only: true,
learned_gate: true,
};
let ckpt = std::env::temp_dir().join("axiom_meta_test.bin");
let ckpt = ckpt.to_string_lossy().to_string();
m.save(&ckpt).unwrap();
assert_eq!(ModelMeta::load(&ckpt).unwrap(), m);
let _ = std::fs::remove_file(ModelMeta::sidecar_path(&ckpt));
}
#[test]
fn footprint_grows_with_size() {
let small = estimate_footprint_bytes(256, 4, 16000, 512);
let large = estimate_footprint_bytes(512, 8, 32000, 512);
assert!(large > small);
}
#[test]
fn pick_config_respects_budget() {
let ladder = default_ladder();
let tiny = pick_config(1, 32000, 512, &ladder);
assert_eq!(tiny, ladder[0]);
let huge = pick_config(u64::MAX, 32000, 512, &ladder);
assert_eq!(huge, *ladder.last().unwrap());
}
}