Skip to main content

cherenkov/model/
qwen.rs

1//! Checkpoint names are interpreted only at the architecture boundary.
2
3use super::TensorRole;
4
5pub(crate) fn tensor_role(name: &str, rank: usize) -> TensorRole {
6    if name.contains(".visual.") || name.starts_with("vision_tower.") {
7        return TensorRole::Opaque;
8    }
9
10    if name.contains(".ngram_embedding.") {
11        return TensorRole::NgramEmbedding;
12    }
13
14    if name.contains(".mlp.switch_mlp.") || name.contains(".mlp.experts.") {
15        return TensorRole::Expert;
16    }
17
18    let Some(prefix) = name.strip_suffix(".weight") else {
19        return TensorRole::Buffer;
20    };
21    let leaf = prefix.rsplit('.').next().unwrap_or(prefix);
22
23    match leaf {
24        "gate" | "shared_expert_gate" if prefix.contains(".mlp.") => TensorRole::Router,
25        "conv1d" => TensorRole::Convolution,
26        "embed_tokens" => TensorRole::Embedding,
27        "lm_head"
28        | "q_proj"
29        | "k_proj"
30        | "v_proj"
31        | "o_proj"
32        | "out_proj"
33        | "in_proj_a"
34        | "in_proj_b"
35        | "in_proj_qkv"
36        | "in_proj_z"
37        | "gate_proj"
38        | "up_proj"
39        | "down_proj"
40        | "key_proj"
41        | "value_proj"
42        | "index_qk_proj"
43        | "input_mix_weight_down"
44        | "input_mix_weight_up"
45        | "block_inject_weight"
46        | "fc_embedding"
47        | "fc_hidden"
48            if rank == 2 =>
49        {
50            TensorRole::Projection
51        }
52        _ if rank == 1 && (leaf.contains("norm") || leaf == "hc_scale") => TensorRole::Norm,
53        _ => TensorRole::Opaque,
54    }
55}