ferrum_cli/commands/
vnext_checkpoint.rs

1use std::{
2    fs,
3    path::{Path, PathBuf},
4};
5
6use clap::Args;
7use ferrum_types::{
8    FerrumError, Result, TokenId, VNextCheckpointCaptureConfig, VNextTeacherForcingConfig,
9};
10use serde::Deserialize;
11
12const MAX_TEACHER_TOKEN_FILE_BYTES: usize = 64 * 1024;
13
14#[derive(Args, Clone, Debug, Default)]
15pub struct VNextCheckpointArgs {
16    /// Empty directory for typed vNext activation evidence.
17    #[arg(long = "vnext-checkpoint-dir", value_name = "DIR")]
18    pub output_dir: Option<PathBuf>,
19
20    /// Semantic ProgramValueId retained and captured after a selected execution
21    /// wave. Repeat for multiple layer or logits checkpoints.
22    #[arg(
23        long = "vnext-checkpoint-value",
24        value_name = "VALUE_ID",
25        action = clap::ArgAction::Append
26    )]
27    pub value_ids: Vec<String>,
28
29    /// Maximum number of real prefill waves to capture. Startup warmup is
30    /// excluded. Defaults to one when capture is configured.
31    #[arg(long = "vnext-checkpoint-prefill-waves", value_name = "N")]
32    pub maximum_prefill_waves: Option<usize>,
33
34    /// Maximum number of real decode waves to capture. Startup warmup is
35    /// excluded. Defaults to zero when capture is configured.
36    #[arg(long = "vnext-checkpoint-decode-waves", value_name = "N")]
37    pub maximum_decode_waves: Option<usize>,
38
39    /// Capture the existing product logits/token readback without retaining an
40    /// activation or changing the compiled memory plan.
41    #[arg(long = "vnext-checkpoint-product-output")]
42    pub capture_product_output: bool,
43
44    /// JSON file containing a bounded canonical token history for a same-history
45    /// numerical diagnostic. Supported only by one-shot `ferrum run`.
46    #[arg(long = "vnext-checkpoint-teacher-token-file", value_name = "JSON")]
47    pub teacher_token_file: Option<PathBuf>,
48}
49
50impl VNextCheckpointArgs {
51    pub fn to_config(&self) -> Result<Option<VNextCheckpointCaptureConfig>> {
52        let configured = self.output_dir.is_some()
53            || !self.value_ids.is_empty()
54            || self.maximum_prefill_waves.is_some()
55            || self.maximum_decode_waves.is_some()
56            || self.capture_product_output
57            || self.teacher_token_file.is_some();
58        if !configured {
59            return Ok(None);
60        }
61        let teacher_forcing = self
62            .teacher_token_file
63            .as_deref()
64            .map(load_teacher_forcing)
65            .transpose()?;
66        let output_dir = self.output_dir.clone().ok_or_else(|| {
67            FerrumError::config(
68                "--vnext-checkpoint-dir is required when checkpoint capture is configured",
69            )
70        })?;
71        if self.value_ids.is_empty() && !self.capture_product_output {
72            return Err(FerrumError::config(
73                "at least one --vnext-checkpoint-value or --vnext-checkpoint-product-output is required",
74            ));
75        }
76        if let Some(teacher_forcing) = &teacher_forcing {
77            if !self.capture_product_output {
78                return Err(FerrumError::config(
79                    "--vnext-checkpoint-teacher-token-file requires --vnext-checkpoint-product-output",
80                ));
81            }
82            if self.maximum_prefill_waves.is_some_and(|waves| waves != 1) {
83                return Err(FerrumError::config(
84                    "teacher-forced checkpoint capture requires exactly one final prefill wave",
85                ));
86            }
87            let expected_decode_waves = teacher_forcing.token_count().saturating_sub(1);
88            if self
89                .maximum_decode_waves
90                .is_some_and(|waves| waves != expected_decode_waves)
91            {
92                return Err(FerrumError::config(format!(
93                    "teacher-forced checkpoint capture requires {expected_decode_waves} decode waves for {} tokens",
94                    teacher_forcing.token_count()
95                )));
96            }
97        }
98        Ok(Some(VNextCheckpointCaptureConfig {
99            output_dir,
100            value_ids: self.value_ids.clone(),
101            maximum_prefill_waves: self.maximum_prefill_waves.unwrap_or(1),
102            maximum_decode_waves: self.maximum_decode_waves.unwrap_or_else(|| {
103                teacher_forcing
104                    .as_ref()
105                    .map_or(0, |teacher| teacher.token_count().saturating_sub(1))
106            }),
107            capture_product_output: self.capture_product_output,
108            teacher_forcing,
109        }))
110    }
111}
112
113#[derive(Deserialize)]
114#[serde(deny_unknown_fields)]
115struct TeacherTokenFile {
116    schema_version: u32,
117    encoding: String,
118    token_ids: Vec<u32>,
119}
120
121fn load_teacher_forcing(path: &Path) -> Result<VNextTeacherForcingConfig> {
122    let bytes = fs::read(path).map_err(|error| {
123        FerrumError::config(format!(
124            "cannot read vNext teacher-token file {}: {error}",
125            path.display()
126        ))
127    })?;
128    if bytes.len() > MAX_TEACHER_TOKEN_FILE_BYTES {
129        return Err(FerrumError::config(format!(
130            "vNext teacher-token file exceeds {MAX_TEACHER_TOKEN_FILE_BYTES} bytes"
131        )));
132    }
133    parse_teacher_forcing(&bytes)
134}
135
136fn parse_teacher_forcing(bytes: &[u8]) -> Result<VNextTeacherForcingConfig> {
137    let parsed: TeacherTokenFile = serde_json::from_slice(bytes).map_err(|error| {
138        FerrumError::config(format!("invalid vNext teacher-token JSON: {error}"))
139    })?;
140    if parsed.schema_version != 1 {
141        return Err(FerrumError::config(
142            "vNext teacher-token schema_version must be 1",
143        ));
144    }
145    if parsed.encoding != "u32-le" {
146        return Err(FerrumError::config(
147            "vNext teacher-token encoding must be u32-le",
148        ));
149    }
150    VNextTeacherForcingConfig::new(parsed.token_ids.into_iter().map(TokenId::new).collect())
151        .map_err(FerrumError::config)
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn absent_flags_preserve_the_product_default() {
160        assert!(VNextCheckpointArgs::default()
161            .to_config()
162            .unwrap()
163            .is_none());
164    }
165
166    #[test]
167    fn capture_requires_a_directory_and_semantic_value() {
168        let missing_directory = VNextCheckpointArgs {
169            value_ids: vec!["value.output.logits".to_owned()],
170            ..VNextCheckpointArgs::default()
171        };
172        assert!(missing_directory.to_config().is_err());
173
174        let missing_value = VNextCheckpointArgs {
175            output_dir: Some(PathBuf::from("capture")),
176            ..VNextCheckpointArgs::default()
177        };
178        assert!(missing_value.to_config().is_err());
179    }
180
181    #[test]
182    fn capture_defaults_to_one_real_prefill_wave() {
183        let config = VNextCheckpointArgs {
184            output_dir: Some(PathBuf::from("capture")),
185            value_ids: vec!["value.output.logits".to_owned()],
186            maximum_prefill_waves: None,
187            maximum_decode_waves: None,
188            capture_product_output: false,
189            teacher_token_file: None,
190        }
191        .to_config()
192        .unwrap()
193        .unwrap();
194        assert_eq!(config.maximum_prefill_waves, 1);
195        assert_eq!(config.maximum_decode_waves, 0);
196    }
197
198    #[test]
199    fn decode_capture_is_an_explicit_shared_product_option() {
200        let config = VNextCheckpointArgs {
201            output_dir: Some(PathBuf::from("capture")),
202            value_ids: vec!["value.output.greedy_token".to_owned()],
203            maximum_prefill_waves: Some(1),
204            maximum_decode_waves: Some(64),
205            capture_product_output: false,
206            teacher_token_file: None,
207        }
208        .to_config()
209        .unwrap()
210        .unwrap();
211
212        assert_eq!(config.maximum_prefill_waves, 1);
213        assert_eq!(config.maximum_decode_waves, 64);
214    }
215
216    #[test]
217    fn product_output_capture_does_not_require_retained_values() {
218        let config = VNextCheckpointArgs {
219            output_dir: Some(PathBuf::from("capture")),
220            maximum_prefill_waves: Some(1),
221            maximum_decode_waves: Some(64),
222            capture_product_output: true,
223            ..VNextCheckpointArgs::default()
224        }
225        .to_config()
226        .unwrap()
227        .unwrap();
228
229        assert!(config.value_ids.is_empty());
230        assert!(config.capture_product_output);
231    }
232
233    #[test]
234    fn teacher_token_json_is_typed_and_bounded() {
235        let parsed = parse_teacher_forcing(
236            br#"{"schema_version":1,"encoding":"u32-le","token_ids":[11690,369]}"#,
237        )
238        .unwrap();
239        assert_eq!(
240            parsed
241                .token_ids()
242                .iter()
243                .map(|token| token.get())
244                .collect::<Vec<_>>(),
245            [11690, 369]
246        );
247
248        assert!(parse_teacher_forcing(
249            br#"{"schema_version":2,"encoding":"u32-le","token_ids":[1]}"#
250        )
251        .is_err());
252        assert!(parse_teacher_forcing(
253            br#"{"schema_version":1,"encoding":"json","token_ids":[1]}"#
254        )
255        .is_err());
256        assert!(parse_teacher_forcing(
257            br#"{"schema_version":1,"encoding":"u32-le","token_ids":[],"extra":true}"#
258        )
259        .is_err());
260
261        let excessive = serde_json::to_vec(&serde_json::json!({
262            "schema_version": 1,
263            "encoding": "u32-le",
264            "token_ids": vec![0_u32; ferrum_types::MAX_VNEXT_TEACHER_FORCED_TOKENS + 1],
265        }))
266        .unwrap();
267        assert!(parse_teacher_forcing(&excessive).is_err());
268    }
269
270    #[test]
271    fn teacher_token_file_derives_exact_wave_contract() {
272        let root = std::env::temp_dir().join(format!(
273            "ferrum-vnext-teacher-token-test-{}",
274            uuid::Uuid::new_v4()
275        ));
276        std::fs::create_dir_all(&root).unwrap();
277        let token_file = root.join("tokens.json");
278        std::fs::write(
279            &token_file,
280            br#"{"schema_version":1,"encoding":"u32-le","token_ids":[7,11,13]}"#,
281        )
282        .unwrap();
283        let args = VNextCheckpointArgs {
284            output_dir: Some(root.join("capture")),
285            capture_product_output: true,
286            teacher_token_file: Some(token_file.clone()),
287            ..VNextCheckpointArgs::default()
288        };
289        let config = args.to_config().unwrap().unwrap();
290        assert_eq!(config.maximum_prefill_waves, 1);
291        assert_eq!(config.maximum_decode_waves, 2);
292        assert_eq!(config.teacher_forcing.unwrap().token_count(), 3);
293
294        let wrong_decode_count = VNextCheckpointArgs {
295            maximum_decode_waves: Some(3),
296            ..args.clone()
297        };
298        assert!(wrong_decode_count.to_config().is_err());
299        let missing_product_output = VNextCheckpointArgs {
300            capture_product_output: false,
301            value_ids: vec!["value.output.logits".to_owned()],
302            ..args
303        };
304        assert!(missing_product_output.to_config().is_err());
305        std::fs::remove_dir_all(root).unwrap();
306    }
307}
ferrum_cli/commands/vnext_checkpoint.rs

ferrum_cli/commands/
vnext_checkpoint.rs
