libmir-metal 0.3.0

Metal inference backend for libmir
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

use std::{fs, path::Path};

use models::{
    execution::DecoderExecutionContract,
    layout::{DecoderConfig, ModelLayout},
    weights::TensorCatalog,
};
use serde_json::{Map, Value, json};

use super::HybridMoeModel;
use crate::engine::{Array, ModelTensors, Result, Stream, lowering};

#[test]
fn executes_a_complete_dense_hybrid_moe_model() -> Result<()> {
    for fused_experts in [false, true] {
        execute_model(fused_experts)?;
    }
    Ok(())
}

fn execute_model(fused_experts: bool) -> Result<()> {
    let root = std::env::temp_dir()
        .join(format!("libmir-metal-dense-hybrid-moe-{}-{fused_experts}", std::process::id()));
    fs::create_dir_all(&root)?;
    write_config(&root)?;
    write_weights(&root.join("model.safetensors"), fused_experts)?;

    let layout = ModelLayout::inspect(&root)?;
    let decoder = DecoderConfig::from_layout(&layout)?;
    let catalog = TensorCatalog::from_layout(&layout)?;
    let contract = DecoderExecutionContract::discover(&layout, &decoder, &catalog)?;
    let lowering = lowering::plan(&contract.semantic)?;
    let stream = Stream::new_cpu()?;
    let tensors = ModelTensors::load(&root, &stream)?;
    let model = HybridMoeModel::load_bindings(
        &tensors,
        &decoder,
        &contract.bindings,
        lowering.layers(),
        16,
        &stream,
    )?;
    let mut cache = model.new_cache(&stream)?;
    let logits = model.forward_decode(&Array::from_u32(&[1], &[1, 1])?, &mut cache, 0, &stream)?;

    assert_eq!(model.layer_count(), 1);
    assert_eq!(logits.shape()?, vec![1, 1, 8]);
    assert!(logits.to_vec_f32_on_stream(&stream)?.iter().all(|value| value.is_finite()));
    drop(tensors);
    fs::remove_dir_all(root)?;
    Ok(())
}

fn write_config(root: &Path) -> Result<()> {
    let config = json!({
        "architectures": ["HybridForCausalLM"],
        "hidden_size": 4,
        "intermediate_size": 8,
        "num_hidden_layers": 1,
        "num_attention_heads": 1,
        "num_key_value_heads": 1,
        "head_dim": 4,
        "vocab_size": 8,
        "max_position_embeddings": 32,
        "num_experts": 2,
        "num_experts_per_tok": 1,
        "moe_intermediate_size": 4,
        "hidden_act": "gelu_pytorch_tanh",
        "attention_k_eq_v": true,
        "rms_norm_eps": 0.000_001,
        "tie_word_embeddings": true
    });
    fs::write(root.join("config.json"), serde_json::to_vec(&config)?)?;
    Ok(())
}

fn write_weights(path: &Path, fused_experts: bool) -> Result<()> {
    let layer = "language_model.model.layers.0";
    let mut specs = vec![
        ("language_model.model.embed_tokens.weight".into(), vec![8, 4]),
        ("language_model.model.norm.weight".into(), vec![4]),
        (format!("{layer}.input_layernorm.weight"), vec![4]),
        (format!("{layer}.self_attn.q_proj.weight"), vec![4, 4]),
        (format!("{layer}.self_attn.k_proj.weight"), vec![4, 4]),
        (format!("{layer}.self_attn.o_proj.weight"), vec![4, 4]),
        (format!("{layer}.self_attn.q_norm.weight"), vec![4]),
        (format!("{layer}.self_attn.k_norm.weight"), vec![4]),
        (format!("{layer}.post_attention_layernorm.weight"), vec![4]),
        (format!("{layer}.pre_feedforward_layernorm.weight"), vec![4]),
        (format!("{layer}.mlp.gate_proj.weight"), vec![8, 4]),
        (format!("{layer}.mlp.up_proj.weight"), vec![8, 4]),
        (format!("{layer}.mlp.down_proj.weight"), vec![4, 8]),
        (format!("{layer}.post_feedforward_layernorm_1.weight"), vec![4]),
        (format!("{layer}.router.proj.weight"), vec![2, 4]),
        (format!("{layer}.router.scale"), vec![4]),
        (format!("{layer}.router.per_expert_scale"), vec![2]),
        (format!("{layer}.pre_feedforward_layernorm_2.weight"), vec![4]),
        (format!("{layer}.post_feedforward_layernorm_2.weight"), vec![4]),
        (format!("{layer}.post_feedforward_layernorm.weight"), vec![4]),
        (format!("{layer}.layer_scalar"), vec![1]),
    ];
    if fused_experts {
        specs.extend([
            (format!("{layer}.experts.gate_up_proj"), vec![2, 8, 4]),
            (format!("{layer}.experts.down_proj"), vec![2, 4, 4]),
        ]);
    } else {
        specs.extend([
            (format!("{layer}.experts.switch_glu.gate_proj.weight"), vec![2, 4, 4]),
            (format!("{layer}.experts.switch_glu.up_proj.weight"), vec![2, 4, 4]),
            (format!("{layer}.experts.switch_glu.down_proj.weight"), vec![2, 4, 4]),
        ]);
    }
    let mut header = Map::new();
    let mut offset = 0_usize;
    let mut payload = Vec::new();
    for (name, shape) in &specs {
        let elements = shape.iter().product::<usize>();
        let bytes = elements.checked_mul(4).ok_or(crate::engine::Error::ShapeOverflow)?;
        header.insert(
            name.clone(),
            json!({"dtype": "F32", "shape": shape, "data_offsets": [offset, offset + bytes]}),
        );
        for _ in 0..elements {
            payload.extend_from_slice(&0.125_f32.to_le_bytes());
        }
        offset += bytes;
    }
    write_safetensors(path, header, &payload)
}

fn write_safetensors(path: &Path, header: Map<String, Value>, payload: &[u8]) -> Result<()> {
    let mut header = serde_json::to_string(&Value::Object(header))?;
    while !header.len().is_multiple_of(8) {
        header.push(' ');
    }
    let mut bytes = u64::try_from(header.len())?.to_le_bytes().to_vec();
    bytes.extend_from_slice(header.as_bytes());
    bytes.extend_from_slice(payload);
    fs::write(path, bytes)?;
    Ok(())
}