moe(wip): gpt-oss-20b groundwork — config fields, arch doc, MXFP4 tools
Phase 19 start. config.rs: explicit head_dim (gpt-oss=64) + MoE fields (num_local_experts, num_experts_per_tok, swiglu_limit, sliding_window, layer_types) with accessors; Qwen3/GPT-2 paths unchanged (fall back to hidden/num_heads when head_dim absent). docs/19-moe-gpt-oss.md: architecture + exact HF reference math (router softmax-after-topk, interleaved clamped (up+1)*glu experts, attention sinks, alternating sliding window, rotate_half RoPE theta=150000, head_dim 64), verified tensor layout, MXFP4 dequant plan. docs/MOE_PROGRESS.md: resume/handoff snapshot. tools/mxfp4_probe.py: inspect safetensors + validate MXFP4 decode (done). tools/gptoss_dequant.py: MXFP4 experts -> plain BF16 safetensors dir so the existing loader reads it (no MXFP4 in Rust for the first pass). Verified: llama.cpp (dash5, LLM_ARCH_OPENAI_MOE) runs the gpt-oss-20b MXFP4 GGUF correctly (17*24 -> 408) = the correctness oracle. MXFP4 decode validated in numpy. Model + GGUF staged on dash5. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user