phase 10: Qwen3-8B support (Milestone ②)
Qwen3 model (qwen3.rs): - RMSNorm + QK normalization (per-head q_norm/k_norm) - GQA: 32 Q heads, 8 KV heads, repeat_kv for attention - SwiGLU FFN: gate_proj → SiLU → * up_proj → down_proj - RoPE with transpose for [1,H,S,D] ↔ [S,H,D] layout - BF16 forward pass, [out,in] weight layout via linear_t - No attention bias (attention_bias=false) Tokenizer fixes: - Fixed unicode_to_byte: shifted bytes now use correct inverse lookup table - MergeEntry supports both string and array formats - Both GPT-2 and Qwen3 tokenizers work correctly (English + Chinese) KVCache refactored: - Dtype-agnostic: stores raw bytes per-head, works for F32 and BF16 - append_kv_tensor/get_kv_tensors use Tensor directly CLI updated: - Auto-detects model type from config.json (gpt2 vs qwen3) - Supports both GPT-2 (F32) and Qwen3 (BF16) Verified: Qwen3-8B generates coherent English and Chinese on single RTX 5090. 61/61 tests pass, GPT-2 performance no regression. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -1,8 +1,7 @@
|
||||
use std::io::{self, Write};
|
||||
use std::path::PathBuf;
|
||||
use xserv_model::gpt2::{sample_greedy, KVCache};
|
||||
use xserv_model::{loader, GPT2, ModelConfig};
|
||||
use xserv_tensor::Device;
|
||||
use xserv_model::{loader, KVCache, ModelConfig};
|
||||
use xserv_tensor::{DType, Device};
|
||||
use xserv_tokenizer::Tokenizer;
|
||||
|
||||
fn main() {
|
||||
@@ -25,43 +24,59 @@ fn main() {
|
||||
eprintln!("GPU: {} ({} MB free)", info.name, info.free_memory / 1024 / 1024);
|
||||
|
||||
let config = ModelConfig::from_file(&model_dir.join("config.json"));
|
||||
let model_type = config.model_type.as_deref().unwrap_or("unknown");
|
||||
eprintln!(
|
||||
"Model: {:?}, layers={}, hidden={}, heads={}, vocab={}",
|
||||
config.model_type,
|
||||
config.num_layers(),
|
||||
config.hidden(),
|
||||
config.num_heads(),
|
||||
config.vocab_size
|
||||
"Model: {model_type}, layers={}, hidden={}, heads={}/{} kv, vocab={}",
|
||||
config.num_layers(), config.hidden(), config.num_heads(),
|
||||
config.num_kv_heads(), config.vocab_size
|
||||
);
|
||||
|
||||
eprintln!("Loading weights...");
|
||||
let weights = loader::load_model_dir(&model_dir, Device::Cuda(0));
|
||||
eprintln!("Loaded {} tensors", weights.len());
|
||||
|
||||
let model = GPT2::from_weights(config.clone(), weights);
|
||||
let is_qwen3 = model_type.contains("qwen");
|
||||
let dtype = if is_qwen3 { DType::BF16 } else { DType::F32 };
|
||||
|
||||
// Build model
|
||||
enum Model {
|
||||
GPT2(xserv_model::GPT2),
|
||||
Qwen3(xserv_model::Qwen3),
|
||||
}
|
||||
let model = if is_qwen3 {
|
||||
Model::Qwen3(xserv_model::Qwen3::from_weights(config.clone(), weights))
|
||||
} else {
|
||||
Model::GPT2(xserv_model::GPT2::from_weights(config.clone(), weights))
|
||||
};
|
||||
|
||||
let tokenizer = Tokenizer::from_file(&model_dir.join("tokenizer.json"));
|
||||
eprintln!("Ready (KV cache enabled).\n");
|
||||
eprintln!("Ready (KV cache, dtype={dtype}).\n");
|
||||
|
||||
loop {
|
||||
print!("xserv> ");
|
||||
io::stdout().flush().unwrap();
|
||||
let mut input = String::new();
|
||||
if io::stdin().read_line(&mut input).unwrap() == 0 {
|
||||
break;
|
||||
}
|
||||
if io::stdin().read_line(&mut input).unwrap() == 0 { break; }
|
||||
let input = input.trim();
|
||||
if input.is_empty() { continue; }
|
||||
if input == "quit" || input == "exit" { break; }
|
||||
|
||||
let token_ids = tokenizer.encode(input);
|
||||
let kv_heads = if is_qwen3 { config.num_kv_heads() } else { config.num_heads() };
|
||||
let mut cache = KVCache::new(
|
||||
config.num_layers(), config.num_heads(), config.head_dim(),
|
||||
Device::Cuda(0),
|
||||
config.num_layers(), kv_heads, config.head_dim(), dtype, Device::Cuda(0),
|
||||
);
|
||||
|
||||
// Prefill
|
||||
let logits = model.forward_with_cache(&token_ids, &mut cache);
|
||||
let mut next = sample_greedy(&logits);
|
||||
// Prefill + decode
|
||||
let logits = match &model {
|
||||
Model::GPT2(m) => m.forward_with_cache(&token_ids, &mut cache),
|
||||
Model::Qwen3(m) => m.forward_with_cache(&token_ids, &mut cache),
|
||||
};
|
||||
let mut next = match &model {
|
||||
Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
|
||||
Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
|
||||
};
|
||||
|
||||
print!("{input}");
|
||||
io::stdout().flush().unwrap();
|
||||
|
||||
@@ -72,8 +87,14 @@ fn main() {
|
||||
|
||||
if tokenizer.eos_token_id() == Some(next) { break; }
|
||||
|
||||
let logits = model.forward_with_cache(&[next], &mut cache);
|
||||
next = sample_greedy(&logits);
|
||||
let logits = match &model {
|
||||
Model::GPT2(m) => m.forward_with_cache(&[next], &mut cache),
|
||||
Model::Qwen3(m) => m.forward_with_cache(&[next], &mut cache),
|
||||
};
|
||||
next = match &model {
|
||||
Model::GPT2(_) => xserv_model::gpt2::sample_greedy(&logits),
|
||||
Model::Qwen3(_) => xserv_model::qwen3::sample_greedy(&logits),
|
||||
};
|
||||
}
|
||||
println!();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user