diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/probe.rs b/crates/spark-model/src/weight_loader/qwen4_exp/probe.rs index 4970939d4..3176cb0e9 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/probe.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/probe.rs @@ -29,9 +29,12 @@ pub struct NamespaceReport { pub stray_layer_norms: usize, } -/// A projection is present in either the dense/NVFP4 layout (`.weight`) or the EXL3 one (`.trellis`). +/// A projection is present in the dense/NVFP4 layout (`.weight`), the compressed-tensors +/// one (`.weight_packed`), or the EXL3 one (`.trellis`). fn has_linear(store: &WeightStore, stem: &str) -> bool { - store.contains(&format!("{stem}.weight")) || store.contains(&format!("{stem}.trellis")) + store.contains(&format!("{stem}.weight")) + || store.contains(&format!("{stem}.weight_packed")) + || store.contains(&format!("{stem}.trellis")) } /// The n-gram tables are deferred (served from NVMe), which `contains` does not see. diff --git a/crates/spark-model/src/weight_loader/qwen4_exp/probe_tests.rs b/crates/spark-model/src/weight_loader/qwen4_exp/probe_tests.rs index 2e13117a9..4d4cf2626 100644 --- a/crates/spark-model/src/weight_loader/qwen4_exp/probe_tests.rs +++ b/crates/spark-model/src/weight_loader/qwen4_exp/probe_tests.rs @@ -1,7 +1,8 @@ // SPDX-License-Identifier: AGPL-3.0-only //! `audit_namespace` must see every projection of an EXL3 checkpoint (`.trellis` -//! instead of `.weight`) and the deferred n-gram tables (never in the store map). +//! instead of `.weight`), of a compressed-tensors one (`.weight_packed`), and the +//! deferred n-gram tables (never in the store map). use super::*; use spark_runtime::gpu::DevicePtr; @@ -81,6 +82,44 @@ fn exl3_layout_is_counted() { assert!(r.ensure_loadable().is_ok()); } +#[test] +fn compressed_tensors_layout_is_counted() { + // A mixed-precision compressed-tensors checkpoint: the routed experts are NVFP4 and + // carry `.weight_packed`, while attention and lm_head are FP8 with a plain `.weight`. + // Probing only `.weight` therefore counts every layer's attention but no experts, + // and the load is refused with "0 of N layers have routed experts". + let c = config(); + let l0 = c.layer_prefix(0); + let l1 = c.layer_prefix(1); + let names = [ + n(format!("{l0}.linear_attn.in_proj_qkv.weight")), + n(format!("{l1}.self_attn.q_proj.weight")), + n(format!("{l1}.self_attn.indexer.index_qk_proj.weight")), + n(format!("{l0}.mlp.experts.0.gate_proj.weight_packed")), + n(format!("{l1}.mlp.experts.0.gate_proj.weight_packed")), + n(format!("{l0}.attn_hyper_connection.hc_norm.weight")), + n(format!("{l1}.attn_hyper_connection.hc_norm.weight")), + n(format!("{l0}.mlp_hyper_connection.hc_norm.weight")), + n(format!("{l1}.mlp_hyper_connection.hc_norm.weight")), + n("model.language_model.embed_tokens.weight".into()), + n("lm_head.weight".into()), + ]; + // The n-gram table stays BF16 in this layout and is deferred as numbered shards. + let base = format!("{l1}.ple.ple_embedding.ngram_embedding"); + let deferred = [ + n(format!("{base}.shard_0.weight")), + n(format!("{base}.shard_1.weight")), + ]; + let r = audit_namespace(&store(&names, &deferred), &c); + assert_eq!(r.gdn_layers, 1); + assert_eq!(r.attn_layers, 1); + assert_eq!(r.indexer_tensors, 1); + assert_eq!(r.expert_tensors, 2); + assert!(r.has_lm_head); + assert_eq!(r.ple_shards, 2); + assert!(r.ensure_loadable().is_ok()); +} + #[test] fn deferred_shards_are_counted() { let c = config();