forge-ml 0.1.0

A WebGPU-native machine learning runtime in Rust: train and run GPT-2 on any GPU wgpu reaches, with no CUDA and no Python
Documentation 
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

//! End-to-end GPT-2 acceptance gates (roadmap v3, Stage 6):
//!   - CPU vs WGPU last-position logits within 5e-3 abs
//!   - greedy continuations identical on both backends
//!   - (when tests/data/hf_golden.json exists) logits within 1e-2 abs of
//!     HF transformers, and identical greedy tokens
//!
//! Requires models/gpt2/ (scripts/download_gpt2.sh); skipped when absent.

use forge::{Device, Gpt2, Gpt2Config, Gpt2Tokenizer, Sampling};

const PROMPT: &str = "Hello, my dog is cute";
const GREEDY_TOKENS: usize = 15;

#[derive(serde::Deserialize)]
struct Golden {
    prompt: String,
    prompt_ids: Vec<u32>,
    logits_last: Vec<f32>,
    greedy_ids: Vec<u32>,
}

fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
    a.iter()
        .zip(b)
        .map(|(x, y)| (x - y).abs())
        .fold(0.0, f32::max)
}

fn greedy_ids(model: &Gpt2, tok: &Gpt2Tokenizer, n: usize) -> Vec<u32> {
    // Re-encode inside so both backends start identically.
    let mut ids = tok.encode(PROMPT).unwrap();
    for _ in 0..n {
        let logits = model.logits_last(&ids).unwrap();
        let next = logits
            .iter()
            .enumerate()
            .max_by(|a, b| a.1.total_cmp(b.1))
            .unwrap()
            .0 as u32;
        ids.push(next);
    }
    ids
}

#[test]
fn gpt2_cpu_wgpu_parity_and_golden() {
    if !std::path::Path::new("models/gpt2/model.safetensors").exists() {
        eprintln!("skipping: models/gpt2 not downloaded");
        return;
    }
    let tok = Gpt2Tokenizer::from_dir("models/gpt2").unwrap();
    let config = Gpt2Config::from_json("models/gpt2/config.json").unwrap();
    let prompt_ids = tok.encode(PROMPT).unwrap();

    println!("loading CPU model...");
    let cpu_model = Gpt2::from_safetensors(
        "models/gpt2/model.safetensors",
        config.clone(),
        &Device::Cpu,
    )
    .unwrap();
    let cpu_logits = cpu_model.logits_last(&prompt_ids).unwrap();

    println!("loading WGPU model...");
    let gpu_dev = Device::wgpu().unwrap();
    println!("wgpu adapter: {}", gpu_dev.describe());
    let gpu_model =
        Gpt2::from_safetensors("models/gpt2/model.safetensors", config, &gpu_dev).unwrap();
    let gpu_logits = gpu_model.logits_last(&prompt_ids).unwrap();

    let d = max_abs_diff(&cpu_logits, &gpu_logits);
    println!("CPU vs WGPU last-logits max abs diff: {d:.2e}");
    assert!(d <= 5e-3, "CPU/WGPU logit divergence {d} exceeds 5e-3");

    let cpu_greedy = greedy_ids(&cpu_model, &tok, GREEDY_TOKENS);
    let gpu_greedy = greedy_ids(&gpu_model, &tok, GREEDY_TOKENS);
    println!("CPU greedy: {:?}", tok.decode(&cpu_greedy));
    println!("GPU greedy: {:?}", tok.decode(&gpu_greedy));
    assert_eq!(
        cpu_greedy, gpu_greedy,
        "greedy tokens differ across backends"
    );

    // Sampling API smoke test (same seed => deterministic).
    let sampled = gpu_model
        .generate(
            &tok,
            PROMPT,
            10,
            Sampling::TopK {
                k: 40,
                temperature: 0.8,
                seed: 7,
            },
        )
        .unwrap();
    println!("sampled: {sampled:?}");
    assert!(sampled.starts_with(PROMPT));

    // Golden comparison against HF transformers, when available.
    let golden_path = "tests/data/hf_golden.json";
    if let Ok(text) = std::fs::read_to_string(golden_path) {
        let golden: Golden = serde_json::from_str(&text).unwrap();
        assert_eq!(golden.prompt, PROMPT, "golden file for different prompt");
        assert_eq!(prompt_ids, golden.prompt_ids, "tokenizer disagrees with HF");
        let dg = max_abs_diff(&cpu_logits, &golden.logits_last);
        println!("CPU vs HF logits max abs diff: {dg:.2e}");
        assert!(dg <= 1e-2, "CPU vs HF logit divergence {dg} exceeds 1e-2");
        let n = golden.greedy_ids.len().min(cpu_greedy.len());
        assert_eq!(
            &cpu_greedy[..n],
            &golden.greedy_ids[..n],
            "greedy tokens disagree with HF"
        );
        println!("HF golden checks passed");
    } else {
        eprintln!("note: {golden_path} absent; HF golden checks skipped");
    }
}