Compare commits
4 Commits
main
...
0dd8851e88
| Author | SHA1 | Date | |
|---|---|---|---|
| 0dd8851e88 | |||
| 05534611ca | |||
| c7d0750c32 | |||
| 057a3c68a3 |
35
crates/xserv-model/src/bin/gptoss-logits.rs
Normal file
35
crates/xserv-model/src/bin/gptoss-logits.rs
Normal file
@@ -0,0 +1,35 @@
|
||||
//! Dump gpt-oss next-token logits for a fixed token-id sequence, to compare
|
||||
//! against the llama.cpp oracle (isolates the model forward from tokenizer
|
||||
//! differences). Usage: gptoss-logits <bf16-model-dir> <tok0> <tok1> ...
|
||||
use std::path::PathBuf;
|
||||
use half::bf16;
|
||||
use xserv_model::loader;
|
||||
use xserv_model::{GptOss, ModelConfig};
|
||||
use xserv_tensor::Device;
|
||||
|
||||
fn main() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
let model_dir = PathBuf::from(&args[1]);
|
||||
let tokens: Vec<u32> = args[2..].iter().map(|s| s.parse().expect("token id")).collect();
|
||||
assert!(!tokens.is_empty(), "need at least one token id");
|
||||
|
||||
let config = ModelConfig::from_file(&model_dir.join("config.json"));
|
||||
eprintln!("[gptoss-logits] loading {} ...", model_dir.display());
|
||||
let weights = loader::load_model_dir(&model_dir, Device::Cpu);
|
||||
let model = GptOss::from_weights(config, weights);
|
||||
eprintln!("[gptoss-logits] forward over {} tokens", tokens.len());
|
||||
|
||||
let logits = model.forward(&tokens); // [T, vocab]
|
||||
let vocab = logits.shape()[1];
|
||||
let t = logits.shape()[0];
|
||||
let host = logits.to_device(Device::Cpu);
|
||||
let data = host.as_slice::<bf16>();
|
||||
let last = &data[(t - 1) * vocab..t * vocab];
|
||||
|
||||
let mut idx: Vec<usize> = (0..vocab).collect();
|
||||
idx.sort_by(|&a, &b| last[b].to_f32().partial_cmp(&last[a].to_f32()).unwrap());
|
||||
println!("top10 next-token (id: logit):");
|
||||
for &i in &idx[..10] {
|
||||
println!(" {i}: {:.4}", last[i].to_f32());
|
||||
}
|
||||
}
|
||||
@@ -46,6 +46,28 @@ pub struct ModelConfig {
|
||||
pub rope_theta: Option<f64>,
|
||||
#[serde(default)]
|
||||
pub tie_word_embeddings: Option<bool>,
|
||||
|
||||
// Explicit head_dim (gpt-oss: 64, which is NOT hidden/num_heads). When
|
||||
// absent, head_dim() falls back to hidden/num_heads (Qwen3, GPT-2).
|
||||
#[serde(default)]
|
||||
pub head_dim: Option<usize>,
|
||||
|
||||
// MoE (gpt-oss). Absent for dense models.
|
||||
#[serde(default)]
|
||||
pub num_local_experts: Option<usize>,
|
||||
#[serde(default)]
|
||||
pub num_experts_per_tok: Option<usize>,
|
||||
// gpt-oss clamped-SwiGLU limit (config: swiglu_limit, default 7.0).
|
||||
#[serde(default)]
|
||||
pub swiglu_limit: Option<f64>,
|
||||
|
||||
// Sliding-window attention (gpt-oss: 128 on alternating layers). The
|
||||
// pattern is given by `layer_types` (e.g. "sliding_attention" /
|
||||
// "full_attention" per layer); absent for dense models.
|
||||
#[serde(default)]
|
||||
pub sliding_window: Option<usize>,
|
||||
#[serde(default)]
|
||||
pub layer_types: Option<Vec<String>>,
|
||||
}
|
||||
|
||||
impl ModelConfig {
|
||||
@@ -81,7 +103,48 @@ impl ModelConfig {
|
||||
}
|
||||
|
||||
pub fn head_dim(&self) -> usize {
|
||||
self.hidden() / self.num_heads()
|
||||
// gpt-oss sets head_dim explicitly (64 != 2880/64). Dense models omit it.
|
||||
self.head_dim.unwrap_or_else(|| self.hidden() / self.num_heads())
|
||||
}
|
||||
|
||||
// ----- MoE (gpt-oss) -----
|
||||
|
||||
/// True for MoE models (have an expert count in config).
|
||||
pub fn is_moe(&self) -> bool {
|
||||
self.num_local_experts.is_some()
|
||||
}
|
||||
|
||||
pub fn num_experts(&self) -> usize {
|
||||
self.num_local_experts.unwrap_or(0)
|
||||
}
|
||||
|
||||
pub fn experts_per_tok(&self) -> usize {
|
||||
self.num_experts_per_tok.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Clamp bound for gpt-oss SwiGLU (config `swiglu_limit`, default 7.0).
|
||||
pub fn swiglu_limit(&self) -> f32 {
|
||||
self.swiglu_limit.unwrap_or(7.0) as f32
|
||||
}
|
||||
|
||||
/// Whether layer `i` uses sliding-window attention. gpt-oss alternates per
|
||||
/// `layer_types`; if that's absent but `sliding_window` is set, fall back to
|
||||
/// the common "every other layer" pattern (even = sliding). Dense → false.
|
||||
pub fn layer_uses_sliding_window(&self, layer: usize) -> bool {
|
||||
if self.sliding_window.is_none() {
|
||||
return false;
|
||||
}
|
||||
match &self.layer_types {
|
||||
Some(types) => types
|
||||
.get(layer)
|
||||
.map(|t| t.contains("sliding"))
|
||||
.unwrap_or(false),
|
||||
None => layer % 2 == 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn sliding_window(&self) -> Option<usize> {
|
||||
self.sliding_window
|
||||
}
|
||||
|
||||
pub fn ln_eps(&self) -> f32 {
|
||||
|
||||
416
crates/xserv-model/src/gptoss.rs
Normal file
416
crates/xserv-model/src/gptoss.rs
Normal file
@@ -0,0 +1,416 @@
|
||||
//! gpt-oss-20b (MoE) forward pass — Phase 19.
|
||||
//!
|
||||
//! Correctness-first, in xserv's own style (reuses our kernels; llama.cpp is only
|
||||
//! a numerical oracle, not a code source). Differences from Qwen3 handled here:
|
||||
//! - MoE FFN: per-token top-4 router (softmax after top-k) + clamped-SwiGLU experts
|
||||
//! - attention sinks: a per-head learned logit column added to the softmax then
|
||||
//! dropped (so attention probabilities do not sum to 1)
|
||||
//! - alternating sliding-window attention (window from config on flagged layers)
|
||||
//! - q/k/v/o projection biases; head_dim 64; no q/k norm; rotate_half RoPE (θ=150000)
|
||||
//!
|
||||
//! Weights are loaded from a plain BF16 safetensors dir (MXFP4 experts are
|
||||
//! dequantized to BF16 offline by tools/gptoss_dequant.py), so the standard
|
||||
//! loader feeds us BF16 tensors and this file needs no quantization code.
|
||||
//!
|
||||
//! v1 is a self-contained non-paged forward (contiguous KV built per call) used to
|
||||
//! validate next-token agreement with llama.cpp. Paged-cache + PP + server wiring
|
||||
//! come after numerical correctness is established.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use half::bf16;
|
||||
use xserv_kernels::*;
|
||||
use xserv_tensor::{Device, Tensor};
|
||||
|
||||
use crate::config::ModelConfig;
|
||||
|
||||
pub struct GptOss {
|
||||
pub config: ModelConfig,
|
||||
embed_tokens: Tensor, // [vocab, hidden]
|
||||
layers: Vec<Block>,
|
||||
norm: Tensor, // [hidden]
|
||||
lm_head_t: Tensor, // [hidden, vocab] (pre-transposed)
|
||||
rope_cache: RopeCache,
|
||||
}
|
||||
|
||||
struct Block {
|
||||
input_norm: Tensor, // [hidden]
|
||||
post_norm: Tensor, // [hidden]
|
||||
// Attention (weights pre-transposed to [in, out]; biases [out]).
|
||||
q_proj_wt: Tensor, // [hidden, n_heads*head_dim]
|
||||
q_bias: Tensor,
|
||||
k_proj_wt: Tensor, // [hidden, n_kv*head_dim]
|
||||
k_bias: Tensor,
|
||||
v_proj_wt: Tensor,
|
||||
v_bias: Tensor,
|
||||
o_proj_wt: Tensor, // [n_heads*head_dim, hidden]
|
||||
o_bias: Tensor,
|
||||
sinks: Tensor, // [n_heads] (f32 on host)
|
||||
sliding: bool,
|
||||
// MoE.
|
||||
router_wt: Tensor, // [hidden, n_experts]
|
||||
router_bias: Tensor, // [n_experts]
|
||||
gate_up_wt: Vec<Tensor>, // per expert: [hidden, 2*inter]
|
||||
gate_up_bias: Vec<Tensor>, // [2*inter]
|
||||
down_wt: Vec<Tensor>, // per expert: [inter, hidden]
|
||||
down_bias: Vec<Tensor>, // [hidden]
|
||||
}
|
||||
|
||||
impl GptOss {
|
||||
/// Load gpt-oss from a BF16 (dequantized) HF-format weight map.
|
||||
pub fn from_weights(config: ModelConfig, mut w: HashMap<String, Tensor>) -> Self {
|
||||
crate::init_kernels();
|
||||
let dev = Device::Cuda(0);
|
||||
let take = |w: &mut HashMap<String, Tensor>, n: &str| -> Tensor {
|
||||
w.remove(n).unwrap_or_else(|| panic!("missing weight: {n}"))
|
||||
};
|
||||
let repl = |t: Tensor| t.to_device(dev);
|
||||
// pre-transpose a [out, in] linear weight to [in, out] for x@wt.
|
||||
let wt = |t: Tensor| t.to_device(dev).transpose(0, 1).contiguous();
|
||||
|
||||
let hidden = config.hidden();
|
||||
let n_experts = config.num_experts();
|
||||
let inter = config.intermediate_size.expect("intermediate_size");
|
||||
|
||||
let embed_tokens = repl(take(&mut w, "model.embed_tokens.weight"));
|
||||
let norm = repl(take(&mut w, "model.norm.weight"));
|
||||
let lm_head_t = wt(take(&mut w, "lm_head.weight"));
|
||||
|
||||
let rope_cache = yarn_rope_cache(&config);
|
||||
|
||||
let n_layers = config.num_layers();
|
||||
let mut layers = Vec::with_capacity(n_layers);
|
||||
for i in 0..n_layers {
|
||||
let p = format!("model.layers.{i}");
|
||||
// Experts are stored fused as [E, in, out]; slice per expert into [in, out].
|
||||
let gate_up = take(&mut w, &format!("{p}.mlp.experts.gate_up_proj")); // [E, hidden, 2*inter]
|
||||
let gate_up_b = take(&mut w, &format!("{p}.mlp.experts.gate_up_proj_bias")); // [E, 2*inter]
|
||||
let down = take(&mut w, &format!("{p}.mlp.experts.down_proj")); // [E, inter, hidden]
|
||||
let down_b = take(&mut w, &format!("{p}.mlp.experts.down_proj_bias")); // [E, hidden]
|
||||
let mut gate_up_wt = Vec::with_capacity(n_experts);
|
||||
let mut gate_up_bias = Vec::with_capacity(n_experts);
|
||||
let mut down_wt = Vec::with_capacity(n_experts);
|
||||
let mut down_bias = Vec::with_capacity(n_experts);
|
||||
// Experts are kept on CPU (the 32 experts per layer total ~36GB for
|
||||
// the whole model, which won't fit one GPU). Each selected expert's
|
||||
// weights (~50MB) are uploaded on demand in expert_forward; only
|
||||
// top-k experts per token are touched, so the H2D traffic is small.
|
||||
for e in 0..n_experts {
|
||||
gate_up_wt.push(slice_expert(&gate_up, e, hidden, 2 * inter)); // CPU
|
||||
gate_up_bias.push(slice_row(&gate_up_b, e, 2 * inter)); // CPU
|
||||
down_wt.push(slice_expert(&down, e, inter, hidden)); // CPU
|
||||
down_bias.push(slice_row(&down_b, e, hidden)); // CPU
|
||||
}
|
||||
|
||||
layers.push(Block {
|
||||
input_norm: repl(take(&mut w, &format!("{p}.input_layernorm.weight"))),
|
||||
post_norm: repl(take(&mut w, &format!("{p}.post_attention_layernorm.weight"))),
|
||||
q_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.q_proj.weight"))),
|
||||
q_bias: repl(take(&mut w, &format!("{p}.self_attn.q_proj.bias"))),
|
||||
k_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.k_proj.weight"))),
|
||||
k_bias: repl(take(&mut w, &format!("{p}.self_attn.k_proj.bias"))),
|
||||
v_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.v_proj.weight"))),
|
||||
v_bias: repl(take(&mut w, &format!("{p}.self_attn.v_proj.bias"))),
|
||||
o_proj_wt: wt(take(&mut w, &format!("{p}.self_attn.o_proj.weight"))),
|
||||
o_bias: repl(take(&mut w, &format!("{p}.self_attn.o_proj.bias"))),
|
||||
sinks: take(&mut w, &format!("{p}.self_attn.sinks")).to_device(Device::Cpu),
|
||||
sliding: config.layer_uses_sliding_window(i),
|
||||
router_wt: wt(take(&mut w, &format!("{p}.mlp.router.weight"))),
|
||||
router_bias: repl(take(&mut w, &format!("{p}.mlp.router.bias"))),
|
||||
gate_up_wt, gate_up_bias, down_wt, down_bias,
|
||||
});
|
||||
}
|
||||
|
||||
Self { config, embed_tokens, layers, norm, lm_head_t, rope_cache }
|
||||
}
|
||||
|
||||
/// Full prefill forward over `token_ids`; returns logits [seq_len, vocab].
|
||||
pub fn forward(&self, token_ids: &[u32]) -> Tensor {
|
||||
let t = token_ids.len();
|
||||
let hidden = self.config.hidden();
|
||||
let n_heads = self.config.num_heads();
|
||||
let n_kv = self.config.num_kv_heads();
|
||||
let head_dim = self.config.head_dim();
|
||||
let eps = self.config.rms_norm_eps.unwrap_or(1e-5) as f32;
|
||||
let positions: Vec<u32> = (0..t as u32).collect();
|
||||
|
||||
let mut x = embedding(&self.embed_tokens, token_ids); // [T, hidden]
|
||||
|
||||
for layer in &self.layers {
|
||||
let residual = x.clone();
|
||||
let normed = rmsnorm(&x, &layer.input_norm, eps);
|
||||
|
||||
// Q/K/V projections + bias.
|
||||
let q = add_bias(&matmul2(&normed, &layer.q_proj_wt), &layer.q_bias); // [T, n_heads*hd]
|
||||
let k = add_bias(&matmul2(&normed, &layer.k_proj_wt), &layer.k_bias); // [T, n_kv*hd]
|
||||
let v = add_bias(&matmul2(&normed, &layer.v_proj_wt), &layer.v_bias);
|
||||
|
||||
// RoPE (rotate_half, same convention xserv uses for Qwen3): reshape to
|
||||
// [1,H,T,D] -> [T,H,D] -> rope -> back.
|
||||
let q = reshape_heads_gpu(&q, t, n_heads, head_dim);
|
||||
let k = reshape_heads_gpu(&k, t, n_kv, head_dim);
|
||||
let q = transpose_for_rope_gpu(&q, t, n_heads, head_dim);
|
||||
let k = transpose_for_rope_gpu(&k, t, n_kv, head_dim);
|
||||
rope_inplace(&q, &self.rope_cache, &positions);
|
||||
rope_inplace(&k, &self.rope_cache, &positions);
|
||||
let q = transpose_from_rope_gpu(&q, t, n_heads, head_dim); // [1,H,T,D]
|
||||
let k = transpose_from_rope_gpu(&k, t, n_kv, head_dim);
|
||||
let v = reshape_heads_gpu(&v, t, n_kv, head_dim); // [1,H_kv,T,D]
|
||||
|
||||
// Naive attention with sinks (CPU softmax for correctness).
|
||||
let attn = attention_with_sinks(
|
||||
&q, &k, &v, &layer.sinks, n_heads, n_kv, head_dim, t,
|
||||
if layer.sliding { self.config.sliding_window() } else { None },
|
||||
); // [T, hidden]
|
||||
let attn_proj = add_bias(&matmul2(&attn, &layer.o_proj_wt), &layer.o_bias);
|
||||
x = add(&residual, &attn_proj);
|
||||
|
||||
// MoE FFN.
|
||||
let residual = x.clone();
|
||||
let normed = rmsnorm(&x, &layer.post_norm, eps);
|
||||
let moe = self.moe_ffn(&normed, layer, hidden);
|
||||
x = add(&residual, &moe);
|
||||
}
|
||||
|
||||
let x = rmsnorm(&x, &self.norm, eps);
|
||||
matmul2(&x, &self.lm_head_t) // [T, vocab]
|
||||
}
|
||||
|
||||
/// MoE FFN over [T, hidden]: router top-k softmax, per-token weighted sum of
|
||||
/// its top-k experts' clamped-SwiGLU outputs. Correctness-first (per-token).
|
||||
fn moe_ffn(&self, x: &Tensor, layer: &Block, hidden: usize) -> Tensor {
|
||||
let t = x.shape()[0];
|
||||
let top_k = self.config.experts_per_tok();
|
||||
let n_experts = self.config.num_experts();
|
||||
let limit = self.config.swiglu_limit();
|
||||
|
||||
// router logits [T, n_experts] on host.
|
||||
let logits = add_bias(&matmul2(x, &layer.router_wt), &layer.router_bias);
|
||||
let logits_h = logits.to_device(Device::Cpu);
|
||||
let lg = logits_h.as_slice::<bf16>();
|
||||
|
||||
// Per-token top-k indices + softmax weights (over the chosen k).
|
||||
let mut out_rows: Vec<Tensor> = Vec::with_capacity(t);
|
||||
for ti in 0..t {
|
||||
let row = &lg[ti * n_experts..(ti + 1) * n_experts];
|
||||
let mut idx: Vec<usize> = (0..n_experts).collect();
|
||||
idx.sort_by(|&a, &b| row[b].to_f32().partial_cmp(&row[a].to_f32()).unwrap());
|
||||
let top = &idx[..top_k];
|
||||
let maxv = row[top[0]].to_f32();
|
||||
let exps: Vec<f32> = top.iter().map(|&e| (row[e].to_f32() - maxv).exp()).collect();
|
||||
let sum: f32 = exps.iter().sum();
|
||||
let weights: Vec<f32> = exps.iter().map(|w| w / sum).collect();
|
||||
|
||||
// x row as [1, hidden].
|
||||
let xr = row_view(x, ti);
|
||||
let mut acc: Option<Tensor> = None;
|
||||
for (j, &e) in top.iter().enumerate() {
|
||||
let y = expert_forward(&xr, &layer.gate_up_wt[e], &layer.gate_up_bias[e],
|
||||
&layer.down_wt[e], &layer.down_bias[e], limit); // [1, hidden]
|
||||
let yw = scale_tensor(&y, weights[j]);
|
||||
acc = Some(match acc { Some(a) => add(&a, &yw), None => yw });
|
||||
}
|
||||
out_rows.push(acc.unwrap_or_else(|| zeros_row(hidden)));
|
||||
}
|
||||
concat_rows(&out_rows) // [T, hidden]
|
||||
}
|
||||
}
|
||||
|
||||
// ---------- helpers ----------
|
||||
|
||||
/// Build a YaRN-scaled RoPE cos/sin cache (gpt-oss uses rope_type "yarn").
|
||||
/// Mirrors HF `_compute_yarn_parameters`: per-dim interpolation/extrapolation
|
||||
/// ramp between the scaled (theta*factor) and unscaled frequencies, plus a global
|
||||
/// attention scaling (mscale) folded into cos/sin. Cache layout matches xserv's
|
||||
/// rope kernel: f32 [max_seq, half_dim], cos[pos*half+i] = cos(pos*invfreq[i])*mscale.
|
||||
fn yarn_rope_cache(config: &ModelConfig) -> RopeCache {
|
||||
use std::f64::consts::PI;
|
||||
let head_dim = config.head_dim();
|
||||
let half = head_dim / 2;
|
||||
let max_seq = config.max_seq_len();
|
||||
let base = config.rope_theta.unwrap_or(150000.0);
|
||||
// gpt-oss rope_scaling: yarn, factor 32, beta_fast 32, beta_slow 1, orig 4096,
|
||||
// truncate false (keep correction range as floats).
|
||||
let factor = 32.0f64;
|
||||
let (beta_fast, beta_slow) = (32.0f64, 1.0f64);
|
||||
let orig_max = 4096.0f64;
|
||||
let dim = head_dim as f64;
|
||||
|
||||
let find_dim = |num_rot: f64| (dim * (orig_max / (num_rot * 2.0 * PI)).ln()) / (2.0 * base.ln());
|
||||
let low = find_dim(beta_fast).max(0.0);
|
||||
let high = find_dim(beta_slow).min(dim - 1.0);
|
||||
let denom = (high - low).max(1e-3);
|
||||
|
||||
let mut inv_freq = vec![0f64; half];
|
||||
for i in 0..half {
|
||||
let pos_freq = base.powf((2 * i) as f64 / dim);
|
||||
let extrap = 1.0 / pos_freq; // unscaled (extrapolation)
|
||||
let interp = 1.0 / (factor * pos_freq); // scaled (interpolation)
|
||||
let ramp = ((i as f64 - low) / denom).clamp(0.0, 1.0);
|
||||
let mask = 1.0 - ramp; // extrapolation factor
|
||||
inv_freq[i] = interp * (1.0 - mask) + extrap * mask;
|
||||
}
|
||||
// mscale: 0.1*ln(factor)+1 for factor>1.
|
||||
let mscale = (0.1 * factor.ln() + 1.0) as f64;
|
||||
|
||||
let mut cos = vec![0f32; max_seq * half];
|
||||
let mut sin = vec![0f32; max_seq * half];
|
||||
for p in 0..max_seq {
|
||||
for i in 0..half {
|
||||
let ang = p as f64 * inv_freq[i];
|
||||
cos[p * half + i] = (ang.cos() * mscale) as f32;
|
||||
sin[p * half + i] = (ang.sin() * mscale) as f32;
|
||||
}
|
||||
}
|
||||
let bytes = max_seq * half * std::mem::size_of::<f32>();
|
||||
let mut cos_buf = xserv_cuda::GpuBuffer::alloc(bytes).expect("alloc yarn cos");
|
||||
let mut sin_buf = xserv_cuda::GpuBuffer::alloc(bytes).expect("alloc yarn sin");
|
||||
let cb = unsafe { std::slice::from_raw_parts(cos.as_ptr() as *const u8, bytes) };
|
||||
let sb = unsafe { std::slice::from_raw_parts(sin.as_ptr() as *const u8, bytes) };
|
||||
cos_buf.copy_from_host(cb).unwrap();
|
||||
sin_buf.copy_from_host(sb).unwrap();
|
||||
RopeCache { cos: cos_buf, sin: sin_buf, max_seq_len: max_seq, half_dim: half }
|
||||
}
|
||||
|
||||
fn matmul2(a: &Tensor, b: &Tensor) -> Tensor {
|
||||
matmul(a, b, GemmBackend::CuBlas)
|
||||
}
|
||||
|
||||
/// One expert: clamped SwiGLU. x:[*,hidden] -> [*,hidden].
|
||||
/// gate_up = x@gate_up_wt + bias; gate=even cols, up=odd cols (interleaved);
|
||||
/// gate.clamp(max=limit); up.clamp(-limit,limit); h=(up+1)*gate*sigmoid(gate*1.702); h@down_wt+bias.
|
||||
fn expert_forward(x: &Tensor, gate_up_wt: &Tensor, gate_up_bias: &Tensor,
|
||||
down_wt: &Tensor, down_bias: &Tensor, limit: f32) -> Tensor {
|
||||
// Upload this expert's CPU-resident weights to x's device just for this call.
|
||||
let dev = x.device();
|
||||
let gate_up_wt = gate_up_wt.to_device(dev);
|
||||
let gate_up_bias = gate_up_bias.to_device(dev);
|
||||
let down_wt = down_wt.to_device(dev);
|
||||
let down_bias = down_bias.to_device(dev);
|
||||
let gate_up = add_bias(&matmul2(x, &gate_up_wt), &gate_up_bias); // [*, 2*inter]
|
||||
let h = clamped_swiglu(&gate_up, limit); // [*, inter]
|
||||
add_bias(&matmul2(&h, &down_wt), &down_bias) // [*, hidden]
|
||||
}
|
||||
|
||||
/// Clamped interleaved SwiGLU on host (correctness-first). [*, 2I] -> [*, I].
|
||||
fn clamped_swiglu(gate_up: &Tensor, limit: f32) -> Tensor {
|
||||
const ALPHA: f32 = 1.702;
|
||||
let rows = gate_up.shape()[0];
|
||||
let two_i = gate_up.shape()[1];
|
||||
let inter = two_i / 2;
|
||||
let h = gate_up.to_device(Device::Cpu);
|
||||
let s = h.as_slice::<bf16>();
|
||||
let mut out = vec![bf16::ZERO; rows * inter];
|
||||
for r in 0..rows {
|
||||
for i in 0..inter {
|
||||
let g = s[r * two_i + 2 * i].to_f32();
|
||||
let u = s[r * two_i + 2 * i + 1].to_f32();
|
||||
let g = g.min(limit);
|
||||
let u = u.clamp(-limit, limit);
|
||||
let glu = g * (1.0 / (1.0 + (-(g * ALPHA)).exp()));
|
||||
out[r * inter + i] = bf16::from_f32((u + 1.0) * glu);
|
||||
}
|
||||
}
|
||||
Tensor::from_slice(&out, &[rows, inter]).to_device(gate_up.device())
|
||||
}
|
||||
|
||||
/// Naive multi-head attention with per-head sink logits, on host (correctness-first).
|
||||
/// q:[1,n_heads,T,D] k,v:[1,n_kv,T,D] sinks:[n_heads] (host). Returns [T, n_heads*D].
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn attention_with_sinks(q: &Tensor, k: &Tensor, v: &Tensor, sinks: &Tensor,
|
||||
n_heads: usize, n_kv: usize, head_dim: usize, t: usize,
|
||||
window: Option<usize>) -> Tensor {
|
||||
let scale = (head_dim as f32).powf(-0.5);
|
||||
let n_rep = n_heads / n_kv;
|
||||
let qh = q.to_device(Device::Cpu); let qd = qh.as_slice::<bf16>();
|
||||
let kh = k.to_device(Device::Cpu); let kd = kh.as_slice::<bf16>();
|
||||
let vh = v.to_device(Device::Cpu); let vd = vh.as_slice::<bf16>();
|
||||
let sh = sinks.to_device(Device::Cpu); let sd = sh.as_slice::<bf16>();
|
||||
let hidden = n_heads * head_dim;
|
||||
let mut out = vec![bf16::ZERO; t * hidden];
|
||||
// index helpers: layout [H, T, D] within each (head) block.
|
||||
let qi = |h: usize, i: usize, d: usize| (h * t + i) * head_dim + d;
|
||||
let kvi = |h: usize, j: usize, d: usize| (h * t + j) * head_dim + d;
|
||||
for h in 0..n_heads {
|
||||
let kv = h / n_rep;
|
||||
for i in 0..t {
|
||||
// scores over valid keys j<=i (causal), and j>i-window (sliding).
|
||||
let lo = match window { Some(wn) if i + 1 > wn => i + 1 - wn, _ => 0 };
|
||||
let mut scores = vec![0f32; i - lo + 1];
|
||||
let mut maxv = sd[h].to_f32(); // sink participates in the max
|
||||
for j in lo..=i {
|
||||
let mut dot = 0f32;
|
||||
for d in 0..head_dim {
|
||||
dot += qd[qi(h, i, d)].to_f32() * kd[kvi(kv, j, d)].to_f32();
|
||||
}
|
||||
let s = dot * scale;
|
||||
scores[j - lo] = s;
|
||||
if s > maxv { maxv = s; }
|
||||
}
|
||||
let mut denom = (sd[h].to_f32() - maxv).exp(); // sink column
|
||||
for s in &scores { denom += (*s - maxv).exp(); }
|
||||
// weighted sum of v (sink contributes no value -> just inflates denom).
|
||||
for d in 0..head_dim {
|
||||
let mut acc = 0f32;
|
||||
for j in lo..=i {
|
||||
let p = (scores[j - lo] - maxv).exp() / denom;
|
||||
acc += p * vd[kvi(kv, j, d)].to_f32();
|
||||
}
|
||||
out[i * hidden + h * head_dim + d] = bf16::from_f32(acc);
|
||||
}
|
||||
}
|
||||
}
|
||||
Tensor::from_slice(&out, &[t, hidden]).to_device(q.device())
|
||||
}
|
||||
|
||||
/// Row-broadcast bias add: x:[T,N] + bias:[N] -> [T,N], via ones[T,1]@bias[1,N].
|
||||
fn add_bias(x: &Tensor, bias: &Tensor) -> Tensor {
|
||||
let t = x.shape()[0];
|
||||
let n = x.shape()[1];
|
||||
let ones = Tensor::from_slice(&vec![bf16::from_f32(1.0); t], &[t, 1]).to_device(x.device());
|
||||
let bias_row = bias.reshape(&[1, n]);
|
||||
let broadcast = matmul2(&ones, &bias_row); // [T, N]
|
||||
add(x, &broadcast)
|
||||
}
|
||||
|
||||
/// Slice expert `e` out of a fused [E, rows, cols] tensor -> [rows, cols].
|
||||
fn slice_expert(t: &Tensor, e: usize, rows: usize, cols: usize) -> Tensor {
|
||||
let host = t.to_device(Device::Cpu);
|
||||
let s = host.as_slice::<bf16>();
|
||||
let stride = rows * cols;
|
||||
Tensor::from_slice(&s[e * stride..(e + 1) * stride], &[rows, cols])
|
||||
}
|
||||
|
||||
/// Slice row `e` out of [E, n] -> [n].
|
||||
fn slice_row(t: &Tensor, e: usize, n: usize) -> Tensor {
|
||||
let host = t.to_device(Device::Cpu);
|
||||
let s = host.as_slice::<bf16>();
|
||||
Tensor::from_slice(&s[e * n..(e + 1) * n], &[n])
|
||||
}
|
||||
|
||||
fn row_view(t: &Tensor, row: usize) -> Tensor {
|
||||
let cols = t.shape()[1];
|
||||
let host = t.to_device(Device::Cpu);
|
||||
let s = host.as_slice::<bf16>();
|
||||
Tensor::from_slice(&s[row * cols..(row + 1) * cols], &[1, cols]).to_device(t.device())
|
||||
}
|
||||
|
||||
fn scale_tensor(t: &Tensor, s: f32) -> Tensor {
|
||||
let host = t.to_device(Device::Cpu);
|
||||
let data = host.as_slice::<bf16>();
|
||||
let out: Vec<bf16> = data.iter().map(|v| bf16::from_f32(v.to_f32() * s)).collect();
|
||||
Tensor::from_slice(&out, t.shape()).to_device(t.device())
|
||||
}
|
||||
|
||||
fn zeros_row(n: usize) -> Tensor {
|
||||
Tensor::from_slice(&vec![bf16::ZERO; n], &[1, n]).to_device(Device::Cuda(0))
|
||||
}
|
||||
|
||||
fn concat_rows(rows: &[Tensor]) -> Tensor {
|
||||
let n = rows[0].shape()[1];
|
||||
let mut out = Vec::with_capacity(rows.len() * n);
|
||||
for r in rows {
|
||||
let h = r.to_device(Device::Cpu);
|
||||
out.extend_from_slice(h.as_slice::<bf16>());
|
||||
}
|
||||
Tensor::from_slice(&out, &[rows.len(), n]).to_device(Device::Cuda(0))
|
||||
}
|
||||
@@ -13,6 +13,7 @@ pub use gpt2::{GPT2, KVCache};
|
||||
pub use kv_cache::GpuKVCache;
|
||||
pub use paged_kv_cache::{BlockAllocator, Location, PagedKVCache, BLOCK_SIZE};
|
||||
pub use qwen3::Qwen3;
|
||||
pub use gptoss::GptOss;
|
||||
pub use sampling::{SamplingParams, sample};
|
||||
|
||||
/// Initialize GPU kernel hooks. Called automatically by model constructors,
|
||||
|
||||
128
docs/19-moe-gpt-oss.md
Normal file
128
docs/19-moe-gpt-oss.md
Normal file
@@ -0,0 +1,128 @@
|
||||
# Phase 19: MoE — gpt-oss-20b
|
||||
|
||||
> 目标:在 xserv 支持 **MoE**,用 `openai/gpt-oss-20b` 端到端跑通,并与 llama.cpp 在
|
||||
> AIME 2025 / GSM8K 上对比正确性与性能。MXFP4 expert 权重加载时反量化为 BF16;整模型
|
||||
> ~40GB 单卡放不下 → 复用 Phase 18 的 **PP**(PP=2 ~20GB/卡,PP=4 ~10GB/卡)。
|
||||
>
|
||||
> 实时进度与重启续作指南见 `docs/MOE_PROGRESS.md`。
|
||||
|
||||
## 1. 架构(config.json,已核对)
|
||||
|
||||
num_hidden_layers=24, hidden=2880, **head_dim=64**(≠hidden/heads), n_heads=64,
|
||||
n_kv_heads=8(GQA n_rep=8), expert intermediate=2880, **num_local_experts=32**,
|
||||
**num_experts_per_tok=4**, vocab=201088, max_pos=131072, rope_theta=150000,
|
||||
sliding_window=128(交替层,见 `layer_types`), rms_norm_eps=1e-5, swiglu_limit=7.0,
|
||||
alpha=1.702, tie_embeddings=false。
|
||||
量化:**MXFP4**,仅 expert MLP(gate_up/down 的 `_blocks`+`_scales`);
|
||||
attn/router/embed/lm_head 为 BF16。
|
||||
|
||||
## 2. 参考数学(HF transformers `modeling_gpt_oss.py`,逐字核对)
|
||||
|
||||
### RMSNorm — 标准(fp32 算 variance,eps=1e-5)。
|
||||
|
||||
### Router(`GptOssTopKRouter`,softmax 在 topk **之后**,含 bias)
|
||||
```
|
||||
logits = x @ W_router^T + b_router # [T, 32]
|
||||
top_val, idx = topk(logits, k=4, dim=-1) # [T, 4]
|
||||
top_val = softmax(top_val, dim=-1) # 仅对选中的 4 个归一化
|
||||
scores = zeros[T,32].scatter(1, idx, top_val)
|
||||
```
|
||||
|
||||
### Experts(`GptOssExperts`,fused gate_up,**interleaved**;clamped;(up+1)·glu)
|
||||
```
|
||||
alpha=1.702; limit=7.0
|
||||
gate_up = x @ gate_up_proj[e] + gate_up_proj_bias[e] # [.., 2*dim]
|
||||
gate = gate_up[..., ::2]; up = gate_up[..., 1::2] # 偶/奇 交错
|
||||
gate = clamp(gate, max=limit) # 仅上界
|
||||
up = clamp(up, min=-limit, max=limit)
|
||||
glu = gate * sigmoid(gate * alpha)
|
||||
h = (up + 1) * glu # 注意 (up+1)
|
||||
y_e = h @ down_proj[e] + down_proj_bias[e]
|
||||
out = Σ_{e∈top4} scores[t,e] * y_e
|
||||
```
|
||||
|
||||
### Attention(`eager_attention_forward`,**带 sinks**)
|
||||
```
|
||||
scaling = head_dim**-0.5 = 64**-0.5;q/k/v/o 都有 bias
|
||||
RoPE(theta=150000) on q,k;repeat_kv(n_rep=8)
|
||||
attn = (q @ k^T) * scaling + causal_mask # 滑窗层叠加 banded(window=128)
|
||||
sinks = module.sinks[head] # 每 head 一个标量
|
||||
combined = cat([attn, sinks broadcast], dim=-1) # 多一列
|
||||
combined -= combined.max(-1, keepdim) # 数值稳定
|
||||
probs = softmax(combined, -1)
|
||||
scores = probs[..., :-1] # 丢掉 sink 列 => 概率不归一到 1
|
||||
o = (scores @ v) -> merge heads -> @Wo + bo
|
||||
```
|
||||
> sinks 等价于 softmax 分母多了 `exp(sink)`——可学习的"不注意"通道。
|
||||
> 交替 sliding window:config `layer_types` 标明哪些层 window=128,其余全注意力。
|
||||
|
||||
与 Qwen3 的新增点:MoE FFN、MXFP4 反量化、attention sinks(softmax 多一列再丢)、
|
||||
交替 sliding window、q/k/v/o bias、head_dim=64、clamped `(up+1)*glu`、rope_theta=150000。
|
||||
|
||||
### 实测张量布局(layer 0,已用 `tools/mxfp4_probe.py` 核对)
|
||||
```
|
||||
self_attn.q_proj.weight [4096,2880] +bias[4096] # 64 heads*64
|
||||
self_attn.k_proj.weight [512,2880] +bias[512] # 8 kv*64
|
||||
self_attn.v_proj.weight [512,2880] +bias[512]
|
||||
self_attn.o_proj.weight [2880,4096] +bias[2880]
|
||||
self_attn.sinks [64] # 每 q-head 一个标量(BF16)
|
||||
input_layernorm.weight [2880]; post_attention_layernorm.weight [2880]
|
||||
mlp.router.weight [32,2880] +bias[32]
|
||||
mlp.experts.gate_up_proj_blocks [32,5760,90,16] U8 + _scales [32,5760,90] U8 + _bias[32,5760] BF16
|
||||
mlp.experts.down_proj_blocks [32,2880,90,16] U8 + _scales [32,2880,90] U8 + _bias[32,2880] BF16
|
||||
# 全局: model.embed_tokens.weight, model.norm.weight, lm_head.weight (BF16)
|
||||
```
|
||||
MXFP4 打包:`[..., nblk=90, 16]` U8,每 16 字节 = 32 个 FP4 码(低 nibble=偶 idx,高 nibble=奇 idx),
|
||||
每 block 一个 E8M0 scale;`90*32 = 2880 = 输入(hidden)维`。即 gate_up 每 expert 权重逻辑 shape
|
||||
`[5760 out, 2880 in]`(**已转置存储**:行=out,列=in,与 HF `nn.Linear` 一致 `y=x·Wᵀ`)。
|
||||
|
||||
### RoPE(**rotate_half,非 interleave**)
|
||||
```
|
||||
dim = head_dim = 64; base = rope_theta = 150000
|
||||
inv_freq = 1 / base^(arange(0,64,2)/64) # 32 项
|
||||
freqs = pos ⊗ inv_freq # [S, 32];cos/sin = cos(freqs)/sin(freqs) (不 doubling)
|
||||
# 应用: x=[.., 64], first=x[:32], second=x[32:]
|
||||
# out_first = first*cos - second*sin
|
||||
# out_second = second*cos + first*sin
|
||||
```
|
||||
> ⚠️ 与 Qwen3 的 RoPE kernel(interleave)不同 —— gptoss 走 rotate_half。需单独处理。
|
||||
|
||||
### Decoder layer(pre-norm 残差,结构同 Qwen3)
|
||||
```
|
||||
h = x + attn(input_norm(x)) # attn 含 sinks/bias/滑窗
|
||||
out = h + moe(post_norm(h)) # moe = router + top4 experts 加权和
|
||||
```
|
||||
最终:`logits = lm_head(norm(h_last))`。无 q_norm/k_norm(与 Qwen3 不同,gptoss 没有)。
|
||||
|
||||
## 3. MXFP4 反量化(expert 权重)
|
||||
|
||||
expert 张量名:`model.layers.{i}.mlp.experts.gate_up_proj_blocks/_scales`、
|
||||
`...down_proj_blocks/_scales`(bias 为 BF16)。MXFP4:每 32 元素一 block 共享一个
|
||||
E8M0(8-bit 指数) scale,每元素 4-bit FP4(E2M1)。反量化
|
||||
`val = fp4_lut[code] * 2^(e8m0 - 127)`。**P19.1 先用 Python(numpy) 反量化并与 HF 一层
|
||||
数值对照**(block 方向 / LUT / gate_up interleave),再写进 Rust loader。
|
||||
|
||||
## 4. 路线(正确优先)
|
||||
|
||||
1. **P19.1** Python 侦查 + MXFP4 反量化验证(不依赖 GPU)。
|
||||
2. **P19.2** `config.rs` 加 MoE 字段(Qwen3 路径不变)。
|
||||
3. **P19.3** `gptoss.rs`:dense(attn+sinks+bias+滑窗 / norm / lm_head)+ MoE FFN
|
||||
(正确优先:逐 token top-4 gather→clamped SwiGLU→加权和);MXFP4 在 `from_weights`
|
||||
反量化为 BF16。验收:prefill logits 与 HF BF16 容差内一致(top-1 一致)。
|
||||
4. **P19.4** 接 PP(experts 随层切),`--pp` 端到端;PP=2/4 与 PP=1 等价。
|
||||
5. **P19.5** llama.cpp 对比(升级 submodule 到支持 gpt-oss 的版本 + 取/转 GGUF),
|
||||
跑 AIME 2025 + GSM8K,复用 `tools/bench` + `summarize_fullq.py`。
|
||||
|
||||
## 5. 风险
|
||||
|
||||
- MXFP4 格式细节必须逐字对 → Python 反量化兜底。
|
||||
- attention sinks + 交替滑窗:现有 flash/paged kernel 未必支持 → 正确优先版本先走朴素
|
||||
attention(显式 mask + sink 列)。
|
||||
- llama.cpp pinned b9371 早于 gpt-oss(约 2025-08)→ 需升级 submodule,有连锁影响。
|
||||
- 性能:MoE 正确优先版本(逐 expert gather/scatter)会慢;先对再快。
|
||||
- **环境**:huggingface.co 被墙,需经代理 + hf-mirror 下载(见 `MOE_PROGRESS.md` §2)。
|
||||
|
||||
## 6. 不在本阶段范围
|
||||
|
||||
GPU 原生 MXFP4 + 按需反量化 kernel(先全 BF16);高性能 grouped-GEMM / expert parallel;
|
||||
TP×MoE;单卡运行(需 MXFP4-native)。
|
||||
177
docs/MOE_PROGRESS.md
Normal file
177
docs/MOE_PROGRESS.md
Normal file
@@ -0,0 +1,177 @@
|
||||
# MoE (gpt-oss-20b) — 工作进度与续作指南
|
||||
|
||||
> **中断原因**:用户要重启 dash5 机器(IP 等可能变),让我先把当前 MoE 支持工作的状态
|
||||
> 完整记录到本文件,重启后据此继续。本文件是"重启后从这里接着干"的唯一入口。
|
||||
|
||||
最后更新:Phase 18 (PP) 已完成并 push;Phase 19 (MoE/gpt-oss-20b) 刚起步(下载受阻,
|
||||
架构与参考数学已侦查清楚)。
|
||||
|
||||
---
|
||||
|
||||
## 0. 一句话现状
|
||||
|
||||
- ✅ **Phase 18 流水线并行 (PP)** 全部完成、验证、benchmark,已 commit 并 push 到
|
||||
`origin/phase18-pipeline-parallelism`(gitea)。
|
||||
- 🚧 **Phase 19 MoE (gpt-oss-20b)** 刚开始:架构 + HF 参考数学已核对(见
|
||||
`docs/19-moe-gpt-oss.md`),**模型还没下载完**(HF 被墙,正在解决下载路径),代码未动。
|
||||
|
||||
---
|
||||
|
||||
## 1. 环境关键事实(重启后很可能变 / 需重新确认)
|
||||
|
||||
- **本机**(开发机,非 GPU):`/home/gahow/projects/xserv`,有公网(走代理)。
|
||||
- **dash5**(GPU 机,8×RTX 5090,无 NVLink,0-3/4-7 分组):通过 `ssh dash5` 访问。
|
||||
- 远端仓库目录:`/opt/wjh/projects/xserv`,模型目录:`/opt/wjh/models/`。
|
||||
- **dash5 无外网、无 rsync**;同步用 `./tools/sync-and-build.sh`(tar over ssh)。
|
||||
- cargo 在 `$HOME/.cargo/bin`;CUDA 12.9 在 `/usr/local/cuda-12.9`。
|
||||
- ⚠️ **重启后 `ssh dash5` 的 IP/可达性可能变** —— 先 `ssh dash5 hostname` 确认;
|
||||
若连不上,检查 `~/.ssh/config` 里 `dash*` 配置 / 让用户给新地址。
|
||||
- **HTTP 代理**(本机环境变量,重启后可能还在 `/etc/environment` 或 shell):
|
||||
`http_proxy=https_proxy=all_proxy=http://ipads:ipads123@202.120.40.82:11235`
|
||||
- **huggingface.co 被墙**(`SSL_ERROR_SYSCALL`,即使过代理)。pypi 可过代理。
|
||||
- **`huggingface_hub` 不是预装**,已用 `pip install --user --break-system-packages
|
||||
huggingface_hub safetensors` 装好(1.17.0);venv 不可用(无 ensurepip)。
|
||||
|
||||
---
|
||||
|
||||
## 2. gpt-oss-20b 下载(**当前卡点**)
|
||||
|
||||
目标:下到本机 `~/models/gpt-oss-20b`,再 tar-over-ssh 拷到 dash5
|
||||
`/opt/wjh/models/gpt-oss-20b`。
|
||||
|
||||
**已验证可行的下载路径**(重启后照此做):
|
||||
- huggingface.co 直连/经代理都失败。
|
||||
- hf-mirror.com 的 `/resolve/` 会 **308 跳回 huggingface.co**(也被墙)——所以不能用
|
||||
`curl -L` 跟跳转,`huggingface_hub` 设 `HF_ENDPOINT` 在新版(1.17)上 HEAD 也失败。
|
||||
- ✅ **能用的办法**:直接走 **hf-mirror 的 `/raw/`(小文件)和实际 CDN,经代理 curl**。
|
||||
已成功取到 `config.json`(200, 1799 bytes):
|
||||
```bash
|
||||
curl -s -x "http://ipads:ipads123@202.120.40.82:11235" \
|
||||
"https://hf-mirror.com/openai/gpt-oss-20b/raw/main/config.json" -o config.json
|
||||
```
|
||||
大文件(safetensors)要用 `/resolve/main/<file>` 且 **指定 `-x` 代理、不要 `-L`**,
|
||||
若仍 308 跳回 hf.co,则改用 hf-mirror 的 LFS 直链或 `huggingface_hub` 配合
|
||||
`HF_ENDPOINT=https://hf-mirror.com` + 代理(库内部不跟 308)。**下载脚本草稿在
|
||||
`/tmp/dl_shards.sh`(重启后 /tmp 会清空,需重建)。**
|
||||
|
||||
**待下载文件**(3 个分片 + 元数据,总 ~13.5GB MXFP4):
|
||||
- `model-00000-of-00002.safetensors`、`model-00001-of-00002.safetensors`、
|
||||
`model-00002-of-00002.safetensors`(注意是 0/1/2 三个,命名 of-00002)
|
||||
- `model.safetensors.index.json`、`config.json`、`tokenizer.json`、
|
||||
`tokenizer_config.json`、`special_tokens_map.json`、`generation_config.json`、
|
||||
`chat_template.jinja`
|
||||
|
||||
**重启后第一步**:`ls -la ~/models/gpt-oss-20b/` 看已下了哪些、`wc -c` 校验分片大小,
|
||||
断点续传用 `curl -C -`。
|
||||
|
||||
---
|
||||
|
||||
## 3. gpt-oss-20b 架构(config.json 已核对)
|
||||
|
||||
| 字段 | 值 |
|
||||
|------|----|
|
||||
| layers | 24;hidden 2880;**head_dim 64**(≠ hidden/heads!)|
|
||||
| heads | 64 q-heads / 8 kv-heads(GQA,n_rep=8)|
|
||||
| experts | num_local_experts **32**,num_experts_per_tok **4**(top-4)|
|
||||
| expert intermediate | 2880 |
|
||||
| vocab | 201088;max_pos 131072;tie_embeddings false |
|
||||
| rope_theta | 150000(核对是否有 rope_scaling/YaRN)|
|
||||
| sliding_window | 128(**交替层**,见 config `layer_types`)|
|
||||
| rms_norm_eps | 1e-5;swiglu_limit 7.0;alpha 1.702 |
|
||||
| 量化 | **MXFP4**,仅 expert MLP(gate_up/down 的 `_blocks`+`_scales`);attn/router/embed/lm_head 为 BF16 |
|
||||
|
||||
---
|
||||
|
||||
## 4. HF 参考数学(已从 transformers `modeling_gpt_oss.py` 逐字核对,务必照抄)
|
||||
|
||||
完整版见 `docs/19-moe-gpt-oss.md` §2。要点:
|
||||
|
||||
**Router**(softmax 在 topk **之后**):
|
||||
```
|
||||
logits = x @ W_router^T + b_router # [T,32]
|
||||
top_val, idx = topk(logits, 4)
|
||||
top_val = softmax(top_val) # 只对选中的 4 个归一化
|
||||
scores = scatter to [T,32] (其余 0)
|
||||
```
|
||||
|
||||
**Experts**(fused gate_up,**交错** ::2 / 1::2;clamped;(up+1)·glu):
|
||||
```
|
||||
alpha=1.702, limit=7.0
|
||||
gate_up = x @ gate_up_proj[e] + bias # [.., 2*2880]
|
||||
gate = gate_up[..., ::2]; up = gate_up[..., 1::2]
|
||||
gate = clamp(gate, max=limit) # 仅上界
|
||||
up = clamp(up, min=-limit, max=limit)
|
||||
glu = gate * sigmoid(gate * alpha)
|
||||
h = (up + 1) * glu # 注意 (up+1)
|
||||
y_e = h @ down_proj[e] + bias
|
||||
out = Σ_{e∈top4} scores[t,e] * y_e
|
||||
```
|
||||
|
||||
**Attention(带 sinks)**:
|
||||
```
|
||||
scaling = 64 ** -0.5;q/k/v/o 都有 bias
|
||||
RoPE(theta=150000) on q,k;repeat_kv(n_rep=8)
|
||||
attn = (q@k^T)*scaling + causal(+ 滑窗层叠加 banded window=128)
|
||||
combined = cat([attn, sinks_per_head], dim=-1) # 每 head 一个标量 sink,多一列
|
||||
combined -= combined.max(-1, keepdim) # 数值稳定
|
||||
probs = softmax(combined, -1)
|
||||
scores = probs[..., :-1] # 丢掉 sink 列(概率不归一到 1!)
|
||||
o = scores @ v -> merge heads -> @Wo + bo
|
||||
```
|
||||
|
||||
**RMSNorm**:标准(fp32 算 variance,eps=1e-5)。
|
||||
|
||||
参考源码已存(重启后 /tmp 清空需重取):`pip download transformers --no-deps`
|
||||
解 wheel 取 `transformers/models/gpt_oss/modeling_gpt_oss.py`(967 行)。
|
||||
|
||||
---
|
||||
|
||||
## 5. MXFP4 反量化(expert 权重)
|
||||
|
||||
- expert 张量名:`model.layers.{i}.mlp.experts.gate_up_proj_blocks` + `..._scales`,
|
||||
`...down_proj_blocks` + `..._scales`(bias 是 BF16 的 `gate_up_proj_bias`/`down_proj_bias`)。
|
||||
- MXFP4:每 **32** 元素一 block,共享一个 **E8M0**(8-bit 指数)scale,每元素 4-bit
|
||||
FP4(E2M1,16 码字)。反量化 `val = fp4_lut[code] * 2^(e8m0 - 127)`。
|
||||
- **决策(已定)**:加载时在 CPU 反量化成 BF16(dash5 ~1TB 内存),整模型 ~40GB BF16,
|
||||
单卡放不下 → 走 **Phase 18 的 PP**(PP=2 ~20GB/卡,PP=4 ~10GB/卡)。不写 GPU 原生
|
||||
MXFP4 kernel(风险高、慢),先正确跑通+对比,后续再优化。
|
||||
|
||||
---
|
||||
|
||||
## 6. 实施路线(Phase 19,逐步可验证)
|
||||
|
||||
1. **P19.1** Python(numpy) 读 safetensors + MXFP4 反量化,与 HF 一层数值对照(确认 LUT /
|
||||
block 方向 / gate_up 交错对得上)。**不依赖 GPU,重启后可先做。**
|
||||
2. **P19.2** `crates/xserv-model/src/config.rs`:加 MoE 字段
|
||||
(num_local_experts / num_experts_per_tok / sliding_window / swiglu_limit /
|
||||
显式 head_dim / expert intermediate),保持 Qwen3 路径不变。
|
||||
3. **P19.3** 新文件 `crates/xserv-model/src/gptoss.rs`:dense(attn+sinks+bias+滑窗 /
|
||||
RMSNorm / lm_head)+ MoE FFN(正确优先:逐 token top-4 gather→clamped SwiGLU→加权和)。
|
||||
MXFP4 在 `from_weights` 反量化为 BF16。验收:prefill logits 与 HF BF16 容差内一致。
|
||||
4. **P19.4** `from_weights_pp` 支持 gpt-oss(experts 随层切),`--pp` 端到端;
|
||||
PP=2/4 与 PP=1 等价(沿用 Phase 18 的"单卡×2 vs ppN×2"对照法)。注:~40GB 需 PP≥2。
|
||||
5. **P19.5** llama.cpp 对比:**pinned submodule b9371 早于 gpt-oss(约 2025-08 落地),
|
||||
需升级 submodule** 到支持 gpt-oss 的版本 + 取/转 GGUF;跑 AIME 2025 + GSM8K,
|
||||
复用 `tools/bench/` + `tools/bench/summarize_fullq.py`(已有,PP 阶段写的)。
|
||||
|
||||
---
|
||||
|
||||
## 7. 复用 Phase 18 的资产
|
||||
|
||||
- 多卡:`--pp N`(已验证),`crates/xserv-distributed`(NCCL P2P + AllReduce)。
|
||||
- bench:`tools/bench/runner.py`(支持 `--pp`/`--tp`)、`summarize_fullq.py`、
|
||||
`tools/pp_quality_full.sh`(xserv 0-3 ‖ llama 4-7 并行跑 AIME+GSM8K 的范式可直接改用)。
|
||||
- 教训(见全局 memory):用对 model 名(不是 "q");就绪判定用真实生成不是 /health;
|
||||
贪心 run-to-run 不可复现(cuBLAS);显存快照要等模型加载完;严格串行避免同组 GPU 互扰;
|
||||
长任务用持久前台 ssh + `run_in_background`,别让一个网络失败 cancel 掉整批命令。
|
||||
|
||||
---
|
||||
|
||||
## 8. 重启后立即要做(checklist)
|
||||
|
||||
1. `ssh dash5 hostname` 确认 GPU 机可达(不行就问用户新地址 / 改 ~/.ssh/config)。
|
||||
2. `git -C ~/projects/xserv log --oneline -6` 确认 PP 5 个 commit 还在
|
||||
(`859c0cc..` 那串,分支 `phase18-pipeline-parallelism`)。
|
||||
3. `ls -la ~/models/gpt-oss-20b/` 看下载进度,续传缺的分片(§2)。
|
||||
4. 重新 `pip download transformers` 取参考源码(/tmp 已清)。
|
||||
5. 从 §6 的 P19.1 接着干。
|
||||
79
tools/gptoss_dequant.py
Normal file
79
tools/gptoss_dequant.py
Normal file
@@ -0,0 +1,79 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Dequantize gpt-oss-20b MXFP4 expert weights -> a plain BF16 safetensors dir.
|
||||
|
||||
Only the expert MLPs are MXFP4 (`*_blocks` uint8 packed 4-bit + `*_scales` uint8
|
||||
E8M0, block=32); everything else is already BF16. We decode experts to BF16 and
|
||||
re-emit a standard HF-format dir so xserv's normal safetensors loader reads it
|
||||
(keeps the first MoE pass free of any MXFP4 code in Rust).
|
||||
|
||||
Fused expert outputs (per layer i), matching HF `GptOssExperts` param shapes:
|
||||
model.layers.{i}.mlp.experts.gate_up_proj [E, hidden, 2*inter] bf16
|
||||
model.layers.{i}.mlp.experts.down_proj [E, inter, hidden] bf16
|
||||
(*_bias tensors pass through unchanged)
|
||||
|
||||
NOTE on transpose: the MXFP4 `_blocks` decode to [E, OUT, IN] (out-major, the
|
||||
contraction dim = nblk*32 last). HF's nn.Parameter for these is [E, IN, OUT]
|
||||
(it does `x @ gate_up_proj`). We emit [E, IN, OUT] (transpose last two dims) so
|
||||
the names/shapes match HF exactly and xserv can treat them uniformly.
|
||||
|
||||
Run on the GPU host (torch + the model + disk):
|
||||
python3 tools/gptoss_dequant.py /opt/wjh/models/gpt-oss-20b /opt/wjh/models/gpt-oss-20b-bf16
|
||||
"""
|
||||
import sys, os, json, glob
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import save_file
|
||||
|
||||
# FP4 E2M1 code -> value (OCP MX). 16 entries.
|
||||
FP4 = torch.tensor(
|
||||
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
|
||||
-0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], dtype=torch.float32)
|
||||
|
||||
def dequant(blocks: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
|
||||
"""blocks uint8 [..., nblk, 16], scales uint8 [..., nblk] -> bf16 [..., nblk*32]."""
|
||||
blocks = blocks.to(torch.int64)
|
||||
lo = blocks & 0xF
|
||||
hi = (blocks >> 4) & 0xF
|
||||
codes = torch.stack([lo, hi], dim=-1).reshape(*blocks.shape[:-1], blocks.shape[-1] * 2)
|
||||
vals = FP4[codes] # [..., nblk, 32] f32
|
||||
scale = torch.exp2(scales.to(torch.float32) - 127.0) # [..., nblk]
|
||||
out = vals * scale[..., None]
|
||||
out = out.reshape(*out.shape[:-2], out.shape[-2] * out.shape[-1])
|
||||
return out.to(torch.bfloat16)
|
||||
|
||||
def main():
|
||||
src, dst = sys.argv[1], sys.argv[2]
|
||||
os.makedirs(dst, exist_ok=True)
|
||||
wm = json.load(open(os.path.join(src, "model.safetensors.index.json")))["weight_map"]
|
||||
|
||||
shards = {}
|
||||
for name, shard in wm.items():
|
||||
shards.setdefault(shard, []).append(name)
|
||||
|
||||
out = {}
|
||||
for shard in sorted(shards):
|
||||
h = safe_open(os.path.join(src, shard), framework="pt")
|
||||
keys = set(h.keys())
|
||||
for name in shards[shard]:
|
||||
if name.endswith("_scales"):
|
||||
continue
|
||||
if name.endswith("_blocks"):
|
||||
base = name[:-len("_blocks")] # ...gate_up_proj / ...down_proj
|
||||
sc = base + "_scales"
|
||||
sc_h = h if sc in keys else safe_open(os.path.join(src, wm[sc]), framework="pt")
|
||||
deq = dequant(h.get_tensor(name), sc_h.get_tensor(sc)) # [E, OUT, IN]
|
||||
out[base] = deq.transpose(1, 2).contiguous() # [E, IN, OUT] (HF param layout)
|
||||
print("dequant", base, tuple(out[base].shape), flush=True)
|
||||
else:
|
||||
out[name] = h.get_tensor(name) # already bf16/other; pass through
|
||||
|
||||
save_file(out, os.path.join(dst, "model.safetensors"), metadata={"format": "pt"})
|
||||
for f in glob.glob(os.path.join(src, "*.json")) + glob.glob(os.path.join(src, "*.jinja")):
|
||||
b = os.path.basename(f)
|
||||
if b == "model.safetensors.index.json":
|
||||
continue
|
||||
open(os.path.join(dst, b), "wb").write(open(f, "rb").read())
|
||||
print("DEQUANT_DONE ->", dst, flush=True)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
82
tools/mxfp4_probe.py
Normal file
82
tools/mxfp4_probe.py
Normal file
@@ -0,0 +1,82 @@
|
||||
#!/usr/bin/env python3
|
||||
"""P19.1 — inspect gpt-oss-20b safetensors + validate MXFP4 dequant (CPU only).
|
||||
|
||||
Run: python3 tools/mxfp4_probe.py /path/to/gpt-oss-20b
|
||||
|
||||
Does three things:
|
||||
1. List the layer-0 tensor names, shapes, dtypes (esp. expert _blocks/_scales).
|
||||
2. Dequantize one expert's gate_up_proj from MXFP4 -> fp32 with our own LUT/decode.
|
||||
3. Print stats so we can eyeball sanity (range, mean, a few values).
|
||||
|
||||
MXFP4 (OCP microscaling) as used by gpt-oss:
|
||||
- elements are FP4 E2M1 (4-bit), packed 2-per-byte.
|
||||
- every block of 32 consecutive elements shares one E8M0 scale (8-bit exponent).
|
||||
- value = fp4_e2m1_lut[code] * 2**(scale_e8m0 - 127)
|
||||
The `_blocks` tensor holds the packed 4-bit codes; `_scales` holds the per-block
|
||||
E8M0 exponents (uint8). We confirm shapes line up (last dim of blocks * 2 / ... )
|
||||
and that decoded values are in a sane range.
|
||||
"""
|
||||
import sys, json, os
|
||||
import numpy as np
|
||||
|
||||
# FP4 E2M1 code -> value lookup (OCP MX spec). 16 codes: sign(1) exp(2) mant(1).
|
||||
# Values: 0, 0.5, 1, 1.5, 2, 3, 4, 6 and their negatives.
|
||||
FP4_E2M1 = np.array(
|
||||
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
|
||||
-0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0],
|
||||
dtype=np.float32,
|
||||
)
|
||||
|
||||
def dequant_mxfp4(blocks: np.ndarray, scales: np.ndarray) -> np.ndarray:
|
||||
"""gpt-oss layout:
|
||||
blocks: uint8 [..., nblk, 16] -> each 16-byte row = 32 FP4 codes (2/byte)
|
||||
scales: uint8 [..., nblk] -> one E8M0 exponent per block
|
||||
Returns fp32 [..., nblk*32] (the input/contraction dim)."""
|
||||
lo = blocks & 0x0F
|
||||
hi = (blocks >> 4) & 0x0F
|
||||
# within a byte, low nibble is element 2i, high nibble is element 2i+1
|
||||
codes = np.empty(blocks.shape[:-1] + (blocks.shape[-1] * 2,), dtype=np.uint8)
|
||||
codes[..., 0::2] = lo
|
||||
codes[..., 1::2] = hi # [..., nblk, 32]
|
||||
vals = FP4_E2M1[codes] # [..., nblk, 32]
|
||||
scale_f = np.power(2.0, scales.astype(np.float32) - 127.0) # [..., nblk]
|
||||
out = vals * scale_f[..., None] # [..., nblk, 32]
|
||||
return out.reshape(out.shape[:-2] + (out.shape[-2] * out.shape[-1],)) # [..., nblk*32]
|
||||
|
||||
def main():
|
||||
d = sys.argv[1] if len(sys.argv) > 1 else "/home/gahow/models/gpt-oss-20b"
|
||||
from safetensors import safe_open
|
||||
idx = json.load(open(os.path.join(d, "model.safetensors.index.json")))
|
||||
wm = idx["weight_map"]
|
||||
l0 = {k: v for k, v in wm.items() if "layers.0." in k}
|
||||
print("=== layer 0 tensors ===")
|
||||
# open each shard lazily
|
||||
handles = {}
|
||||
def get(name):
|
||||
shard = wm[name]
|
||||
if shard not in handles:
|
||||
handles[shard] = safe_open(os.path.join(d, shard), framework="numpy")
|
||||
return handles[shard]
|
||||
for k in sorted(l0):
|
||||
h = get(k)
|
||||
t = h.get_slice(k)
|
||||
print(f" {k.replace('model.layers.0.','L0.')} shape={t.get_shape()} dtype={t.get_dtype()}")
|
||||
|
||||
# find expert gate_up blocks/scales
|
||||
gu_b = next((k for k in l0 if "gate_up_proj_blocks" in k), None)
|
||||
gu_s = next((k for k in l0 if "gate_up_proj_scales" in k), None)
|
||||
if gu_b and gu_s:
|
||||
print(f"\n=== dequant {gu_b.split('.')[-1]} (expert 0) ===")
|
||||
blocks = get(gu_b).get_tensor(gu_b)
|
||||
scales = get(gu_s).get_tensor(gu_s)
|
||||
print(" blocks", blocks.shape, blocks.dtype, " scales", scales.shape, scales.dtype)
|
||||
b0 = blocks[0]; s0 = scales[0] # expert 0: blocks [5760,90,16], scales [5760,90]
|
||||
deq = dequant_mxfp4(b0.astype(np.uint8), s0.astype(np.uint8)) # -> [5760, 2880]
|
||||
print(" expert0 dequant shape", deq.shape, "(expect [5760, 2880] = [2*inter, hidden])",
|
||||
"\n min %.4f max %.4f mean %.4f std %.4f" % (deq.min(), deq.max(), deq.mean(), deq.std()))
|
||||
print(" row0 first 8 vals:", np.round(deq[0, :8], 4))
|
||||
else:
|
||||
print("\n(no gate_up_proj_blocks/_scales found — check tensor names above)")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user