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:
@@ -8,9 +8,11 @@ pub struct Tokenizer {
|
||||
decoder: Vec<Vec<u8>>,
|
||||
merge_ranks: HashMap<(u32, u32), usize>,
|
||||
special_tokens: HashMap<String, u32>,
|
||||
#[allow(dead_code)]
|
||||
special_token_ids: HashMap<u32, String>,
|
||||
pre_tokenize_re: Regex,
|
||||
eos_token_id: Option<u32>,
|
||||
byte_fallback: bool,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
@@ -23,7 +25,16 @@ struct TokenizerJson {
|
||||
#[derive(Deserialize)]
|
||||
struct ModelSection {
|
||||
vocab: HashMap<String, u32>,
|
||||
merges: Vec<String>,
|
||||
merges: Vec<MergeEntry>,
|
||||
#[serde(default)]
|
||||
byte_fallback: bool,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum MergeEntry {
|
||||
Str(String),
|
||||
Pair(Vec<String>),
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
@@ -40,7 +51,10 @@ impl Tokenizer {
|
||||
let tj: TokenizerJson = serde_json::from_str(&data)
|
||||
.unwrap_or_else(|e| panic!("failed to parse tokenizer.json: {e}"));
|
||||
|
||||
let byte_fallback = tj.model.byte_fallback;
|
||||
|
||||
// Build encoder: token bytes → ID
|
||||
// All HF tokenizers use GPT-2 byte-to-unicode mapping for vocab keys.
|
||||
let mut encoder = HashMap::new();
|
||||
for (token_str, &id) in &tj.model.vocab {
|
||||
let bytes = token_str_to_bytes(token_str);
|
||||
@@ -56,13 +70,23 @@ impl Tokenizer {
|
||||
decoder[id as usize] = token_str_to_bytes(token_str);
|
||||
}
|
||||
|
||||
// Parse merges
|
||||
// Parse merges (supports both "a b" string format and ["a", "b"] array format)
|
||||
let byte_fallback = tj.model.byte_fallback;
|
||||
let mut merge_ranks = HashMap::new();
|
||||
for (rank, merge_line) in tj.model.merges.iter().enumerate() {
|
||||
let parts: Vec<&str> = merge_line.splitn(2, ' ').collect();
|
||||
if parts.len() != 2 { continue; }
|
||||
let a_bytes = token_str_to_bytes(parts[0]);
|
||||
let b_bytes = token_str_to_bytes(parts[1]);
|
||||
for (rank, entry) in tj.model.merges.iter().enumerate() {
|
||||
let (a_str, b_str) = match entry {
|
||||
MergeEntry::Str(s) => {
|
||||
let parts: Vec<&str> = s.splitn(2, ' ').collect();
|
||||
if parts.len() != 2 { continue; }
|
||||
(parts[0].to_string(), parts[1].to_string())
|
||||
}
|
||||
MergeEntry::Pair(v) => {
|
||||
if v.len() != 2 { continue; }
|
||||
(v[0].clone(), v[1].clone())
|
||||
}
|
||||
};
|
||||
let a_bytes = token_str_to_bytes(&a_str);
|
||||
let b_bytes = token_str_to_bytes(&b_str);
|
||||
if let (Some(&a_id), Some(&b_id)) = (encoder.get(&a_bytes), encoder.get(&b_bytes)) {
|
||||
merge_ranks.insert((a_id, b_id), rank);
|
||||
}
|
||||
@@ -84,13 +108,14 @@ impl Tokenizer {
|
||||
}
|
||||
}
|
||||
|
||||
// GPT-2 pre-tokenization regex.
|
||||
// The original uses (?!\S) lookahead which Rust regex doesn't support.
|
||||
// Simplified: collapse trailing whitespace into one match. Functionally equivalent
|
||||
// for BPE since each whitespace chunk gets encoded independently anyway.
|
||||
let pre_tokenize_re = Regex::new(
|
||||
r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+"
|
||||
).unwrap();
|
||||
// Pre-tokenization regex
|
||||
let pre_tokenize_re = if byte_fallback {
|
||||
// Qwen-style: split on whitespace boundaries, keep Unicode words/numbers
|
||||
Regex::new(r"[\p{L}\p{N}]+|[^\s\p{L}\p{N}]|\s+").unwrap()
|
||||
} else {
|
||||
// GPT-2 style
|
||||
Regex::new(r"'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+").unwrap()
|
||||
};
|
||||
|
||||
Self {
|
||||
encoder,
|
||||
@@ -100,6 +125,7 @@ impl Tokenizer {
|
||||
special_token_ids,
|
||||
pre_tokenize_re,
|
||||
eos_token_id,
|
||||
byte_fallback,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -137,10 +163,16 @@ impl Tokenizer {
|
||||
fn encode_ordinary(&self, text: &str, out: &mut Vec<u32>) {
|
||||
for mat in self.pre_tokenize_re.find_iter(text) {
|
||||
let word = mat.as_str();
|
||||
// Try to encode the whole word first
|
||||
if let Some(&id) = self.encoder.get(word.as_bytes()) {
|
||||
out.push(id);
|
||||
continue;
|
||||
}
|
||||
// Fall back to per-byte encoding
|
||||
let word_bytes: Vec<u8> = word.bytes().collect();
|
||||
let mut token_ids: Vec<u32> = word_bytes.iter().map(|&b| {
|
||||
*self.encoder.get(&vec![b]).unwrap_or_else(|| {
|
||||
panic!("byte {b} not in vocab")
|
||||
panic!("byte {b} (0x{b:02X}) not in vocab")
|
||||
})
|
||||
}).collect();
|
||||
|
||||
@@ -204,48 +236,32 @@ fn token_str_to_bytes(s: &str) -> Vec<u8> {
|
||||
s.chars().map(|c| unicode_to_byte(c)).collect()
|
||||
}
|
||||
|
||||
/// Convert a Unicode char back to the byte it represents in GPT-2 encoding.
|
||||
fn unicode_to_byte(c: char) -> u8 {
|
||||
let u = c as u32;
|
||||
// GPT-2 byte encoder: maps bytes 0-255 to specific Unicode code points.
|
||||
// Printable ASCII bytes map to themselves. Others are shifted to 256+.
|
||||
match u {
|
||||
0x21..=0x7E => u as u8, // '!' to '~'
|
||||
0xA1..=0xAC => u as u8, // '¡' to '¬'
|
||||
0xAE..=0xFF => u as u8, // '®' to 'ÿ'
|
||||
// Shifted bytes: 0x100 + original_byte for bytes not in the above ranges
|
||||
0x100..=0x1FF => (u - 0x100) as u8 + {
|
||||
// The shift mapping: byte values 0..=32, 127..=160, 173
|
||||
// are shifted to 256..=288, 289+, etc.
|
||||
0
|
||||
},
|
||||
_ => {
|
||||
// Fallback: for the GPT-2 byte encoder, specific mappings
|
||||
byte_from_unicode_gpt2(c)
|
||||
// Build the inverse map on first use
|
||||
use std::sync::OnceLock;
|
||||
static INV_MAP: OnceLock<HashMap<u32, u8>> = OnceLock::new();
|
||||
|
||||
let map = INV_MAP.get_or_init(|| {
|
||||
let mut m = HashMap::new();
|
||||
// Build GPT-2's bytes_to_unicode forward map, then invert
|
||||
let mut n = 0u32;
|
||||
for b in 0..=255u16 {
|
||||
let byte = b as u8;
|
||||
let unicode = match byte {
|
||||
0x21..=0x7E | 0xA1..=0xAC | 0xAE..=0xFF => byte as u32,
|
||||
_ => {
|
||||
let u = 256 + n;
|
||||
n += 1;
|
||||
u
|
||||
}
|
||||
};
|
||||
m.insert(unicode, byte);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn byte_from_unicode_gpt2(c: char) -> u8 {
|
||||
// Build the inverse of GPT-2's bytes_to_unicode mapping.
|
||||
// The mapping assigns printable chars to themselves and shifts unprintable bytes.
|
||||
let u = c as u32;
|
||||
// Direct ASCII printable + Latin-1 supplement printable ranges map identity
|
||||
if (0x21..=0x7E).contains(&u) { return u as u8; }
|
||||
if (0xA1..=0xAC).contains(&u) { return u as u8; }
|
||||
if (0xAE..=0xFF).contains(&u) { return u as u8; }
|
||||
|
||||
// Shifted range: the remaining 68 bytes (0-32, 127-160, 173) get mapped to 256..=323
|
||||
static SHIFTED_BYTES: &[u8] = &[
|
||||
0, 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, 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, 173,
|
||||
];
|
||||
let shifted_start = 256u32;
|
||||
if u >= shifted_start && u < shifted_start + SHIFTED_BYTES.len() as u32 {
|
||||
return SHIFTED_BYTES[(u - shifted_start) as usize];
|
||||
}
|
||||
|
||||
// Shouldn't reach here for valid GPT-2 tokenizer
|
||||
c as u8
|
||||
m
|
||||
});
|
||||
|
||||
*map.get(&(c as u32)).unwrap_or_else(|| {
|
||||
panic!("unmapped unicode char U+{:04X} in tokenizer", c as u32)
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user