Gahow Wang
e1e75fc7f6
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>
88 lines
2.9 KiB
Rust
88 lines
2.9 KiB
Rust
use half::{bf16, f16};
|
|
use safetensors::SafeTensors;
|
|
use std::collections::HashMap;
|
|
use std::path::Path;
|
|
use xserv_tensor::{DType, Device, Tensor};
|
|
|
|
pub fn load_safetensors(path: &Path, device: Device) -> HashMap<String, Tensor> {
|
|
let data = std::fs::read(path)
|
|
.unwrap_or_else(|e| panic!("failed to read {}: {e}", path.display()));
|
|
let st = SafeTensors::deserialize(&data)
|
|
.unwrap_or_else(|e| panic!("failed to parse safetensors {}: {e}", path.display()));
|
|
|
|
let mut tensors = HashMap::new();
|
|
|
|
for (name, view) in st.tensors() {
|
|
let shape: Vec<usize> = view.shape().to_vec();
|
|
let raw_bytes = view.data();
|
|
let dtype = match view.dtype() {
|
|
safetensors::Dtype::F32 => DType::F32,
|
|
safetensors::Dtype::F16 => DType::F16,
|
|
safetensors::Dtype::BF16 => DType::BF16,
|
|
other => {
|
|
eprintln!("skipping tensor {name}: unsupported dtype {other:?}");
|
|
continue;
|
|
}
|
|
};
|
|
|
|
let tensor = make_tensor(raw_bytes, &shape, dtype);
|
|
let tensor = tensor.to_device(device);
|
|
tensors.insert(name.to_string(), tensor);
|
|
}
|
|
|
|
tensors
|
|
}
|
|
|
|
/// Load from a directory containing model.safetensors (or sharded files) + config.json.
|
|
pub fn load_model_dir(dir: &Path, device: Device) -> HashMap<String, Tensor> {
|
|
let single = dir.join("model.safetensors");
|
|
if single.exists() {
|
|
return load_safetensors(&single, device);
|
|
}
|
|
|
|
// Try sharded: model-00001-of-NNNNN.safetensors
|
|
let mut all_tensors = HashMap::new();
|
|
let mut entries: Vec<_> = std::fs::read_dir(dir)
|
|
.unwrap()
|
|
.filter_map(|e| e.ok())
|
|
.filter(|e| {
|
|
e.path()
|
|
.file_name()
|
|
.map(|f| f.to_string_lossy().ends_with(".safetensors"))
|
|
.unwrap_or(false)
|
|
})
|
|
.collect();
|
|
entries.sort_by_key(|e| e.file_name());
|
|
|
|
for entry in entries {
|
|
let tensors = load_safetensors(&entry.path(), device);
|
|
all_tensors.extend(tensors);
|
|
}
|
|
|
|
assert!(!all_tensors.is_empty(), "no safetensors files found in {}", dir.display());
|
|
all_tensors
|
|
}
|
|
|
|
fn make_tensor(raw_bytes: &[u8], shape: &[usize], dtype: DType) -> Tensor {
|
|
match dtype {
|
|
DType::F32 => {
|
|
let floats: &[f32] = unsafe {
|
|
std::slice::from_raw_parts(raw_bytes.as_ptr() as *const f32, raw_bytes.len() / 4)
|
|
};
|
|
Tensor::from_slice(floats, shape)
|
|
}
|
|
DType::F16 => {
|
|
let halfs: &[f16] = unsafe {
|
|
std::slice::from_raw_parts(raw_bytes.as_ptr() as *const f16, raw_bytes.len() / 2)
|
|
};
|
|
Tensor::from_slice(halfs, shape)
|
|
}
|
|
DType::BF16 => {
|
|
let bfs: &[bf16] = unsafe {
|
|
std::slice::from_raw_parts(raw_bytes.as_ptr() as *const bf16, raw_bytes.len() / 2)
|
|
};
|
|
Tensor::from_slice(bfs, shape)
|
|
}
|
|
}
|
|
}
|