spark_runtime/weights/
prefix_detect.rs

1// SPDX-License-Identifier: AGPL-3.0-only
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(
22    store: &WeightStore,
23    config: &mut atlas_core::config::ModelConfig,
24) {
25    if config.weight_prefix.is_empty() && config.nested_config {
26        config.weight_prefix = if store.contains("language_model.model.embed_tokens.weight") {
27            "language_model.model".to_string()
28        } else if store.contains("model.language_model.embed_tokens.weight") {
29            "model.language_model".to_string()
30        } else {
31            let scanned = store
32                .names()
33                .find(|k| k.contains(".layers.0."))
34                .and_then(|k| k.split(".layers.0.").next())
35                .map(|s| s.to_string());
36            if let Some(ref prefix) = scanned {
37                tracing::info!("Auto-detected weight prefix: '{prefix}'");
38            }
39            scanned.unwrap_or_else(|| "model".to_string())
40        };
41    }
42    if !config.weight_prefix.is_empty() {
43        tracing::info!("Weight prefix: {}", config.weight_prefix);
44    }
45}
spark_runtime/weights/prefix_detect.rs

spark_runtime/weights/
prefix_detect.rs
