phase 6+7+8: model loading, BPE tokenizer, GPT-2 inference (Milestone ①)
Phase 6 — Model Loading (xserv-model): - safetensors parser with single/sharded file support - ModelConfig with dual naming (GPT-2 n_embd/n_head + modern HF naming) - Weight loading flow: safetensors → mmap → CPU Tensor → GPU Phase 7 — BPE Tokenizer (xserv-tokenizer): - Full BPE encode/decode from tokenizer.json - GPT-2 byte-to-unicode mapping (printable ASCII identity + shifted bytes) - Pre-tokenization regex, special token handling - Chat template support structure Phase 8 — GPT-2 Complete Inference: - GPT-2 model definition: wte, wpe, 12 transformer blocks, ln_f - Forward pass: embedding → (LayerNorm → MHA → residual → LayerNorm → MLP → residual) × 12 → LN → logits - QKV split with correct [batch, heads, seq, dim] layout (fixed reshape bug) - Greedy sampling from last-position logits - Interactive CLI: xserv-cli <model-dir> [--max-tokens N] Verified: GPT-2 124M generates coherent English text on RTX 5090. "The future of AI is uncertain. The future of AI is uncertain..." "Once upon a time, the world was a place of great beauty..." Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
224
crates/xserv-model/src/gpt2.rs
Normal file
224
crates/xserv-model/src/gpt2.rs
Normal file
@@ -0,0 +1,224 @@
|
||||
use std::collections::HashMap;
|
||||
use xserv_kernels::*;
|
||||
use xserv_tensor::{DType, Device, Tensor};
|
||||
|
||||
use crate::config::ModelConfig;
|
||||
|
||||
pub struct GPT2 {
|
||||
pub config: ModelConfig,
|
||||
wte: Tensor, // [vocab_size, hidden]
|
||||
wpe: Tensor, // [max_pos, hidden]
|
||||
layers: Vec<GPT2Block>,
|
||||
ln_f_g: Tensor, // [hidden]
|
||||
ln_f_b: Tensor, // [hidden]
|
||||
}
|
||||
|
||||
struct GPT2Block {
|
||||
ln_1_g: Tensor,
|
||||
ln_1_b: Tensor,
|
||||
// Attention: combined QKV weight + bias, output weight + bias
|
||||
attn_qkv_w: Tensor, // [hidden, 3*hidden]
|
||||
attn_qkv_b: Tensor, // [3*hidden]
|
||||
attn_out_w: Tensor, // [hidden, hidden]
|
||||
attn_out_b: Tensor, // [hidden]
|
||||
ln_2_g: Tensor,
|
||||
ln_2_b: Tensor,
|
||||
mlp_fc_w: Tensor, // [hidden, 4*hidden]
|
||||
mlp_fc_b: Tensor, // [4*hidden]
|
||||
mlp_proj_w: Tensor, // [4*hidden, hidden]
|
||||
mlp_proj_b: Tensor, // [hidden]
|
||||
}
|
||||
|
||||
impl GPT2 {
|
||||
pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
|
||||
let take = |w: &mut HashMap<String, Tensor>, name: &str| -> Tensor {
|
||||
w.remove(name).unwrap_or_else(|| panic!("missing weight: {name}"))
|
||||
};
|
||||
|
||||
let wte = take(&mut w, "wte.weight");
|
||||
let wpe = take(&mut w, "wpe.weight");
|
||||
let ln_f_g = take(&mut w, "ln_f.weight");
|
||||
let ln_f_b = take(&mut w, "ln_f.bias");
|
||||
|
||||
let num_layers = config.num_layers();
|
||||
let mut layers = Vec::with_capacity(num_layers);
|
||||
for i in 0..num_layers {
|
||||
let p = format!("h.{i}");
|
||||
layers.push(GPT2Block {
|
||||
ln_1_g: take(&mut w, &format!("{p}.ln_1.weight")),
|
||||
ln_1_b: take(&mut w, &format!("{p}.ln_1.bias")),
|
||||
attn_qkv_w: take(&mut w, &format!("{p}.attn.c_attn.weight")),
|
||||
attn_qkv_b: take(&mut w, &format!("{p}.attn.c_attn.bias")),
|
||||
attn_out_w: take(&mut w, &format!("{p}.attn.c_proj.weight")),
|
||||
attn_out_b: take(&mut w, &format!("{p}.attn.c_proj.bias")),
|
||||
ln_2_g: take(&mut w, &format!("{p}.ln_2.weight")),
|
||||
ln_2_b: take(&mut w, &format!("{p}.ln_2.bias")),
|
||||
mlp_fc_w: take(&mut w, &format!("{p}.mlp.c_fc.weight")),
|
||||
mlp_fc_b: take(&mut w, &format!("{p}.mlp.c_fc.bias")),
|
||||
mlp_proj_w: take(&mut w, &format!("{p}.mlp.c_proj.weight")),
|
||||
mlp_proj_b: take(&mut w, &format!("{p}.mlp.c_proj.bias")),
|
||||
});
|
||||
}
|
||||
|
||||
Self { config, wte, wpe, layers, ln_f_g, ln_f_b }
|
||||
}
|
||||
|
||||
/// Full forward pass, returns logits [seq_len, vocab_size].
|
||||
pub fn forward(&self, token_ids: &[u32]) -> Tensor {
|
||||
let seq_len = token_ids.len();
|
||||
let hidden = self.config.hidden();
|
||||
let num_heads = self.config.num_heads();
|
||||
let head_dim = self.config.head_dim();
|
||||
|
||||
// Token + position embedding
|
||||
let tok_emb = embedding(&self.wte, token_ids);
|
||||
let pos_ids: Vec<u32> = (0..seq_len as u32).collect();
|
||||
let pos_emb = embedding(&self.wpe, &pos_ids);
|
||||
let mut x = add_tensors(&tok_emb, &pos_emb);
|
||||
|
||||
// Transformer layers
|
||||
for layer in &self.layers {
|
||||
// Pre-LN attention
|
||||
let residual = x.clone();
|
||||
let normed = layernorm(&x, &layer.ln_1_g, &layer.ln_1_b, self.config.ln_eps());
|
||||
|
||||
// QKV projection: [S, H] @ [H, 3H] + [3H] → [S, 3H]
|
||||
let qkv = linear(&normed, &layer.attn_qkv_w, Some(&layer.attn_qkv_b));
|
||||
// Split into Q, K, V and reshape for multi-head
|
||||
let (q, k, v) = split_qkv(&qkv, num_heads, head_dim, seq_len);
|
||||
// Attention: [1, H, S, D]
|
||||
let attn_out = attention(&q, &k, &v, true);
|
||||
// Merge heads: [1, H, S, D] → [S, hidden]
|
||||
let attn_out = merge_heads(&attn_out, seq_len, hidden);
|
||||
// Output projection
|
||||
let attn_out = linear(&attn_out, &layer.attn_out_w, Some(&layer.attn_out_b));
|
||||
x = add_tensors(&residual, &attn_out);
|
||||
|
||||
// Pre-LN MLP
|
||||
let residual = x.clone();
|
||||
let normed = layernorm(&x, &layer.ln_2_g, &layer.ln_2_b, self.config.ln_eps());
|
||||
let fc = linear(&normed, &layer.mlp_fc_w, Some(&layer.mlp_fc_b));
|
||||
let activated = gelu(&fc);
|
||||
let proj = linear(&activated, &layer.mlp_proj_w, Some(&layer.mlp_proj_b));
|
||||
x = add_tensors(&residual, &proj);
|
||||
}
|
||||
|
||||
// Final layer norm
|
||||
let x = layernorm(&x, &self.ln_f_g, &self.ln_f_b, self.config.ln_eps());
|
||||
|
||||
// LM head (tied with wte): [S, H] @ [H, V] → [S, V]
|
||||
// wte is [V, H], so we need wte^T
|
||||
let lm_head = self.wte.transpose(0, 1).contiguous();
|
||||
matmul_2d(&x, &lm_head)
|
||||
}
|
||||
}
|
||||
|
||||
// --- Helper ops ---
|
||||
|
||||
fn linear(x: &Tensor, weight: &Tensor, bias: Option<&Tensor>) -> Tensor {
|
||||
// GPT-2 stores weights as [in, out] (not transposed), so x @ w
|
||||
let out = matmul_2d(x, weight);
|
||||
if let Some(b) = bias {
|
||||
add_bias(&out, b)
|
||||
} else {
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
fn matmul_2d(a: &Tensor, b: &Tensor) -> Tensor {
|
||||
// a: [S, K], b: [K, N] → [S, N]
|
||||
assert_eq!(a.ndim(), 2);
|
||||
assert_eq!(b.ndim(), 2);
|
||||
matmul(a, b, GemmBackend::CuBlas)
|
||||
}
|
||||
|
||||
fn add_tensors(a: &Tensor, b: &Tensor) -> Tensor {
|
||||
// Element-wise add on GPU via a simple approach: scale(a, 1.0) + scale(b, 1.0)
|
||||
// TODO: proper add kernel. For now, go through CPU.
|
||||
assert_eq!(a.shape(), b.shape());
|
||||
assert_eq!(a.dtype(), DType::F32);
|
||||
let a_cpu = a.to_device(Device::Cpu);
|
||||
let b_cpu = b.to_device(Device::Cpu);
|
||||
let a_data = a_cpu.as_slice::<f32>();
|
||||
let b_data = b_cpu.as_slice::<f32>();
|
||||
let sum: Vec<f32> = a_data.iter().zip(b_data).map(|(x, y)| x + y).collect();
|
||||
Tensor::from_slice(&sum, a.shape()).to_device(a.device())
|
||||
}
|
||||
|
||||
fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
|
||||
// x: [S, N], bias: [N] → broadcast add
|
||||
assert_eq!(x.ndim(), 2);
|
||||
assert_eq!(bias.ndim(), 1);
|
||||
assert_eq!(x.shape()[1], bias.shape()[0]);
|
||||
let x_cpu = x.to_device(Device::Cpu);
|
||||
let b_cpu = bias.to_device(Device::Cpu);
|
||||
let x_data = x_cpu.as_slice::<f32>();
|
||||
let b_data = b_cpu.as_slice::<f32>();
|
||||
let n = bias.shape()[0];
|
||||
let result: Vec<f32> = x_data.iter().enumerate().map(|(i, &v)| v + b_data[i % n]).collect();
|
||||
Tensor::from_slice(&result, x.shape()).to_device(x.device())
|
||||
}
|
||||
|
||||
fn split_qkv(qkv: &Tensor, num_heads: usize, head_dim: usize, seq_len: usize) -> (Tensor, Tensor, Tensor) {
|
||||
// qkv: [S, 3*H] → Q, K, V each [1, num_heads, S, head_dim]
|
||||
let hidden = num_heads * head_dim;
|
||||
let qkv_cpu = qkv.to_device(Device::Cpu);
|
||||
let data = qkv_cpu.as_slice::<f32>();
|
||||
|
||||
// Split into Q, K, V and directly write in [1, num_heads, S, head_dim] layout
|
||||
let mut q_data = vec![0.0f32; num_heads * seq_len * head_dim];
|
||||
let mut k_data = vec![0.0f32; num_heads * seq_len * head_dim];
|
||||
let mut v_data = vec![0.0f32; num_heads * seq_len * head_dim];
|
||||
|
||||
for s in 0..seq_len {
|
||||
let row = &data[s * 3 * hidden..(s + 1) * 3 * hidden];
|
||||
for h in 0..num_heads {
|
||||
let src_off = h * head_dim;
|
||||
let dst_off = (h * seq_len + s) * head_dim;
|
||||
q_data[dst_off..dst_off + head_dim].copy_from_slice(&row[src_off..src_off + head_dim]);
|
||||
k_data[dst_off..dst_off + head_dim].copy_from_slice(&row[hidden + src_off..hidden + src_off + head_dim]);
|
||||
v_data[dst_off..dst_off + head_dim].copy_from_slice(&row[2 * hidden + src_off..2 * hidden + src_off + head_dim]);
|
||||
}
|
||||
}
|
||||
|
||||
let device = qkv.device();
|
||||
let q = Tensor::from_slice(&q_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
|
||||
let k = Tensor::from_slice(&k_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
|
||||
let v = Tensor::from_slice(&v_data, &[1, num_heads, seq_len, head_dim]).to_device(device);
|
||||
(q, k, v)
|
||||
}
|
||||
|
||||
fn merge_heads(x: &Tensor, seq_len: usize, hidden: usize) -> Tensor {
|
||||
// [1, num_heads, S, head_dim] → [S, hidden]
|
||||
let num_heads = x.shape()[1];
|
||||
let head_dim = x.shape()[3];
|
||||
let x_cpu = x.to_device(Device::Cpu);
|
||||
let src = x_cpu.as_slice::<f32>();
|
||||
|
||||
// src layout: [1][num_heads][seq_len][head_dim]
|
||||
// dst layout: [seq_len][hidden] where hidden = num_heads * head_dim
|
||||
let mut out = vec![0.0f32; seq_len * hidden];
|
||||
for s in 0..seq_len {
|
||||
for h in 0..num_heads {
|
||||
let src_off = (h * seq_len + s) * head_dim;
|
||||
let dst_off = s * hidden + h * head_dim;
|
||||
out[dst_off..dst_off + head_dim].copy_from_slice(&src[src_off..src_off + head_dim]);
|
||||
}
|
||||
}
|
||||
Tensor::from_slice(&out, &[seq_len, hidden]).to_device(x.device())
|
||||
}
|
||||
|
||||
/// Greedy sampling: return the argmax token ID from the last position's logits.
|
||||
pub fn sample_greedy(logits: &Tensor) -> u32 {
|
||||
assert_eq!(logits.ndim(), 2); // [S, V]
|
||||
let logits_cpu = logits.to_device(Device::Cpu);
|
||||
let data = logits_cpu.as_slice::<f32>();
|
||||
let vocab_size = logits.shape()[1];
|
||||
let seq_len = logits.shape()[0];
|
||||
let last_row = &data[(seq_len - 1) * vocab_size..seq_len * vocab_size];
|
||||
last_row.iter()
|
||||
.enumerate()
|
||||
.max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
|
||||
.map(|(idx, _)| idx as u32)
|
||||
.unwrap()
|
||||
}
|
||||
Reference in New Issue
Block a user