aprender-serve 0.34.0

Pure Rust ML inference engine built from scratch - model serving for GGUF and safetensors
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
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175

//! Verify RoPE type detection and application for Qwen2
use realizar::gguf::{MappedGGUFModel, OwnedQuantizedModel};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Load Qwen2 model
    let qwen_path = "/home/noah/.cache/huggingface/hub/models--Qwen--Qwen2-0.5B-Instruct-GGUF/snapshots/198f08841147e5196a6a69bd0053690fb1fd3857/qwen2-0_5b-instruct-q4_0.gguf";
    let qwen_mapped = MappedGGUFModel::from_path(qwen_path)?;
    let qwen_model = OwnedQuantizedModel::from_mapped(&qwen_mapped)?;

    println!("=== Qwen2 Config ===");
    println!("Architecture: {}", qwen_model.config().architecture);
    println!("Hidden dim: {}", qwen_model.config().hidden_dim);
    println!("Num heads: {}", qwen_model.config().num_heads);
    println!("Num KV heads: {}", qwen_model.config().num_kv_heads);
    println!("RoPE theta: {}", qwen_model.config().rope_theta);
    println!(
        "RoPE type: {} (0=NORM adjacent pairs, 2=NEOX split halves)",
        qwen_model.config().rope_type
    );
    println!("Epsilon: {}", qwen_model.config().eps);

    // Check raw metadata
    println!("\n=== Raw GGUF Metadata (RoPE related) ===");
    for (key, value) in &qwen_mapped.model.metadata {
        if key.contains("rope") || key.contains("eps") {
            println!("  {}: {:?}", key, value);
        }
    }

    // Also load TinyLlama for comparison
    let tiny_path = "/tmp/tinyllama.gguf";
    if std::path::Path::new(tiny_path).exists() {
        let tiny_mapped = MappedGGUFModel::from_path(tiny_path)?;
        let tiny_model = OwnedQuantizedModel::from_mapped(&tiny_mapped)?;

        println!("\n=== TinyLlama Config (comparison) ===");
        println!("Architecture: {}", tiny_model.config().architecture);
        println!("Hidden dim: {}", tiny_model.config().hidden_dim);
        println!("Num heads: {}", tiny_model.config().num_heads);
        println!("Num KV heads: {}", tiny_model.config().num_kv_heads);
        println!("RoPE theta: {}", tiny_model.config().rope_theta);
        println!(
            "RoPE type: {} (0=NORM adjacent pairs, 2=NEOX split halves)",
            tiny_model.config().rope_type
        );
        println!("Epsilon: {}", tiny_model.config().eps);

        println!("\n=== TinyLlama Raw Metadata (RoPE related) ===");
        for (key, value) in &tiny_mapped.model.metadata {
            if key.contains("rope") || key.contains("eps") {
                println!("  {}: {:?}", key, value);
            }
        }
    }

    // Manually test RoPE rotation
    println!("\n=== Manual RoPE Test ===");
    let head_dim = qwen_model.config().hidden_dim / qwen_model.config().num_heads;
    let half_dim = head_dim / 2;
    println!("Head dim: {}", head_dim);
    println!("Half dim: {}", half_dim);

    // Create test input: [1, 2, 3, 4, ...] for one head
    let mut x_neox: Vec<f32> = (0..head_dim).map(|i| (i + 1) as f32).collect();
    let mut x_norm: Vec<f32> = x_neox.clone();

    // Pre-compute cos/sin for position 0
    let theta = qwen_model.config().rope_theta;
    let mut cos_vals = vec![0.0f32; half_dim];
    let mut sin_vals = vec![0.0f32; half_dim];
    let position = 0usize;
    for i in 0..half_dim {
        let freq = 1.0 / theta.powf(2.0 * i as f32 / head_dim as f32);
        let angle = position as f32 * freq;
        let (sin_v, cos_v) = angle.sin_cos();
        cos_vals[i] = cos_v;
        sin_vals[i] = sin_v;
    }

    println!("\nPosition 0 cos/sin (first 4):");
    println!("  cos: {:?}", &cos_vals[..4.min(cos_vals.len())]);
    println!("  sin: {:?}", &sin_vals[..4.min(sin_vals.len())]);

    // Apply NEOX style (split halves)
    {
        let (first_half, second_half) = x_neox.split_at_mut(half_dim);
        for i in 0..half_dim {
            let v1 = first_half[i];
            let v2 = second_half[i];
            first_half[i] = v1 * cos_vals[i] - v2 * sin_vals[i];
            second_half[i] = v1 * sin_vals[i] + v2 * cos_vals[i];
        }
    }

    // Apply NORM style (adjacent pairs)
    for i in 0..half_dim {
        let x0 = x_norm[2 * i];
        let x1 = x_norm[2 * i + 1];
        x_norm[2 * i] = x0 * cos_vals[i] - x1 * sin_vals[i];
        x_norm[2 * i + 1] = x0 * sin_vals[i] + x1 * cos_vals[i];
    }

    println!(
        "\nInput (first 8): {:?}",
        &(0..head_dim).map(|i| (i + 1) as f32).collect::<Vec<_>>()[..8.min(head_dim)]
    );
    println!("NEOX result (first 8): {:?}", &x_neox[..8.min(head_dim)]);
    println!("NORM result (first 8): {:?}", &x_norm[..8.min(head_dim)]);

    // At position 0, both should give the same result because cos(0)=1, sin(0)=0
    let neox_norm: f32 = x_neox.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm_norm: f32 = x_norm.iter().map(|x| x * x).sum::<f32>().sqrt();
    let orig_norm: f32 = (1..=head_dim as i32)
        .map(|i| (i * i) as f32)
        .sum::<f32>()
        .sqrt();
    println!("\nNorm preservation check:");
    println!("  Original norm: {:.4}", orig_norm);
    println!("  NEOX norm: {:.4}", neox_norm);
    println!("  NORM norm: {:.4}", norm_norm);

    // Check position 5 where rotation actually happens
    println!("\n=== Position 5 Test ===");
    let mut x_neox_5: Vec<f32> = (0..head_dim).map(|i| (i + 1) as f32).collect();
    let mut x_norm_5: Vec<f32> = x_neox_5.clone();
    let position = 5usize;

    for i in 0..half_dim {
        let freq = 1.0 / theta.powf(2.0 * i as f32 / head_dim as f32);
        let angle = position as f32 * freq;
        let (sin_v, cos_v) = angle.sin_cos();
        cos_vals[i] = cos_v;
        sin_vals[i] = sin_v;
    }

    println!("Position 5 cos/sin (first 4):");
    println!("  cos: {:?}", &cos_vals[..4.min(cos_vals.len())]);
    println!("  sin: {:?}", &sin_vals[..4.min(sin_vals.len())]);

    // Apply NEOX style
    {
        let (first_half, second_half) = x_neox_5.split_at_mut(half_dim);
        for i in 0..half_dim {
            let v1 = first_half[i];
            let v2 = second_half[i];
            first_half[i] = v1 * cos_vals[i] - v2 * sin_vals[i];
            second_half[i] = v1 * sin_vals[i] + v2 * cos_vals[i];
        }
    }

    // Apply NORM style
    for i in 0..half_dim {
        let x0 = x_norm_5[2 * i];
        let x1 = x_norm_5[2 * i + 1];
        x_norm_5[2 * i] = x0 * cos_vals[i] - x1 * sin_vals[i];
        x_norm_5[2 * i + 1] = x0 * sin_vals[i] + x1 * cos_vals[i];
    }

    println!(
        "\nNEOX result (first 8): {:?}",
        &x_neox_5[..8.min(head_dim)]
    );
    println!("NORM result (first 8): {:?}", &x_norm_5[..8.min(head_dim)]);

    // Check if results differ significantly
    let diff: f32 = x_neox_5
        .iter()
        .zip(x_norm_5.iter())
        .map(|(a, b)| (a - b).abs())
        .sum::<f32>()
        / head_dim as f32;
    println!("\nMean absolute diff between NEOX and NORM: {:.6}", diff);

    Ok(())
}