metrale_model_weights/weights/
prefix_detect.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Weight-key prefix auto-detection for nested (multimodal) checkpoints.
4//!
5//! Split out of `weights.rs` for the 500-LoC cap — that file sat at exactly
6//! 500 after the Q2 keep-packed additions, and the `check_oom_guard` export
7//! fix re-tripped it; this move buys real headroom instead of shaving to the
8//! line again.
9
10use super::WeightStore;
11
12/// Resolve the weight-key prefix a nested (multimodal) checkpoint uses.
13///
14/// Lives here rather than in the server because it depends only on the store
15/// and the config, and BOTH the serve path and the integration harness need it:
16/// a nested checkpoint stores `model.language_model.layers.0.…`, and a caller
17/// that skips this step fails with "Weight 'model.layers.0.input_layernorm.
18/// weight' not found in store" after a full weight load. Keeping one copy is
19/// what stops the test from accepting a different set of checkpoints than
20/// production does.
21pub fn auto_detect_weight_prefix(store: &WeightStore, config: &mut metrale_config::ModelConfig) {
22    if config.weight_prefix.is_empty() && config.nested_config {
23        config.weight_prefix = if store.contains("language_model.model.embed_tokens.weight") {
24            "language_model.model".to_string()
25        } else if store.contains("model.language_model.embed_tokens.weight") {
26            "model.language_model".to_string()
27        } else {
28            let scanned = store
29                .names()
30                .find(|k| k.contains(".layers.0."))
31                .and_then(|k| k.split(".layers.0.").next())
32                .map(|s| s.to_string());
33            if let Some(ref prefix) = scanned {
34                tracing::info!("Auto-detected weight prefix: '{prefix}'");
35            }
36            scanned.unwrap_or_else(|| "model".to_string())
37        };
38    }
39    if !config.weight_prefix.is_empty() {
40        tracing::info!("Weight prefix: {}", config.weight_prefix);
41    }
42}
metrale_model_weights/weights/prefix_detect.rs

metrale_model_weights/weights/
prefix_detect.rs
