From 86fb17366bc9a90181aa5ddee8a214a89193aeae Mon Sep 17 00:00:00 2001 From: juliankang4 <200967947+juliankang4@users.noreply.github.com> Date: Wed, 21 Oct 2026 13:39:36 +0900 Subject: [PATCH] prepare, levers: keep checkpoint matrices in their stored format by default --- Cargo.lock | 1 + .../gpu-gates/src/bin/gate_qwen4exp_nvfp4.rs | 275 +++++++++++-- crates/hf/src/modelopt.rs | 5 + crates/levers/src/lib.rs | 13 +- crates/levers/src/registry.rs | 25 ++ crates/levers/src/tests.rs | 12 +- crates/model-hf/Cargo.toml | 2 + crates/model-hf/src/bin/prepare.rs | 367 ++++++++++++++---- crates/model/src/arch/qwen35moe/mtp.rs | 2 +- crates/refset/src/arch/qwen4exp/nvfp4.rs | 15 +- justfile | 10 +- 11 files changed, 591 insertions(+), 136 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 36a9e18f4..0a50d4ad6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -341,6 +341,7 @@ version = "0.1.0" dependencies = [ "bloomery-gguf", "bloomery-hf", + "bloomery-levers", "bloomery-model", "bloomery-refset", "bloomery-safetensors", diff --git a/crates/gpu-gates/src/bin/gate_qwen4exp_nvfp4.rs b/crates/gpu-gates/src/bin/gate_qwen4exp_nvfp4.rs index b7da4dac1..5b799dd95 100644 --- a/crates/gpu-gates/src/bin/gate_qwen4exp_nvfp4.rs +++ b/crates/gpu-gates/src/bin/gate_qwen4exp_nvfp4.rs @@ -1,8 +1,12 @@ //! The Qwen3.8-Flash-Next NVFP4 gate: the split `prepare` writes from `nvidia/Qwen3.8-Flash-Next-NVFP4` -//! (routed experts NVFP4 with a tensor scale an expert, dense projections Q8_0, gains F32, the PLE -//! table in its F8 rows) against the checkpoint's own forward with the weights the split holds -//! (refset `hf-qwen4exp-nvfp4`, `tools/ref/qwen4exp-nvfp4-oracle.py`), by the e2e gate's bands -//! (`shared/q38_oracle.rs`). The plan is `--place a`'s on the one visible card, a DGX Spark's GB10, +//! (routed experts NVFP4 with a tensor scale an expert, dense projections BF16 as the checkpoint +//! stores them or, under `BLOOMERY_PREPARE_Q8=1`, rewritten to Q8_0, gains F32, the PLE table in +//! its F8 rows) against the checkpoint's own forward (refset +//! `hf-qwen4exp-nvfp4`, `tools/ref/qwen4exp-nvfp4-oracle.py`: the `_bf16` sets for the stored +//! projections, the others for the Q8_0 ones; the head's type says which), by the e2e gate's bands +//! (`shared/q38_oracle.rs`). The oracle keeps indexer projections in BF16; the model loader +//! refuses Q8_0 indexer projections. +//! The plan is `--place a`'s on the one visible card, a DGX Spark's GB10, //! whose budget holds every routed expert: no CPU kernel reads an NVFP4 stack. //! //! What is asserted: @@ -19,10 +23,11 @@ //! named tie (ours the oracle's runner-up, the oracle's margin at most twice our distance at the //! two ids, the logits within `FREE_BAND`). The layer outputs are printed: a flip in any of the //! 3,000 positions before it may lie on their path. -//! - Recorded, not judged: the batch run and the D3K step against the sets that keep the -//! checkpoint's BF16 projections, the Q8_0 rounding's share of the distance. +//! - Recorded, not judged: the batch run and the D3K step against the other pair of sets, the Q8_0 +//! rounding's share of the distance. //! //! Real-only: the fixture holds no NVFP4 stack. Loads the whole file on the card. +//! `BLOOMERY_REF_MODEL` selects a split outside the recipe's model folder. #[cfg(feature = "gpu")] fn main() -> std::process::ExitCode { @@ -66,27 +71,41 @@ mod gate { use bloomery_gpu::model::StepMode; use bloomery_gpu_gates::flip; use bloomery_gpu_gates::generate::Place; - use bloomery_gpu_gates::{GateError, RefManifest, checks_failed, data_dir, verdict}; + use bloomery_gpu_gates::{ + GateError, RefManifest, checks_failed, data_dir, ref_model_path, verdict, + }; use gguf::GgmlType; use model::arch::qwen35moe::names; use model::arch::qwen35moe::place::{Experts, PlanInputs}; - use model::placement::PlanLevers; + use model::placement::{ModelTensors, PlanLevers, Role}; use refset::arch::qwen4exp::nvfp4::{BATCH, BATCH_BF16, D3K, D3K_BF16, HF}; - /// The split `prepare` writes from the family's folder, first of its shards. - const MODEL: &str = - "/models/Qwen3.8-Flash-Next-NVFP4/Qwen3.8-Flash-Next-NVFP4-00001-of-00004.gguf"; - /// Cache rows: D3K's step at position 3,000, with room. const CTX: usize = 3072; /// What the clauses read of the file's header: the trunk's layers, a layer output's values - /// (every stream of the hidden width), and the routed experts and the slots a token takes. + /// (every stream of the hidden width), the routed experts and the slots a token takes, and + /// the sets its dense form is judged by. struct Dims { n_layer: usize, row: usize, n_expert: usize, n_used: usize, + dense: Dense, + } + + /// The form of a split's dense matrices ([`dense_of`]) and the oracle's sets it is judged by, + /// each beside the set of the other form, recorded. + #[derive(Clone, Copy, Debug, PartialEq, Eq)] + struct Dense { + /// The head's type: BF16 or Q8_0. + ty: GgmlType, + /// At least one indexer projection is Q8_0; the oracle's are all BF16. + q8_indexer: bool, + /// The batch set judged, then the one recorded. + batch: (&'static str, &'static str), + /// The D3K set judged, then the one recorded. + d3k: (&'static str, &'static str), } /// Steps' last logits, and with the taps armed each step's layer outputs and route taps. @@ -104,8 +123,8 @@ mod gate { } /// (l): the plan on the one visible card, every routed expert on it, and the load. - fn open(d: &Dims) -> Result<(Qwen38Model, bool), GateError> { - let file = open_split(Path::new(MODEL))?; + fn open(path: &Path, d: &Dims) -> Result<(Qwen38Model, bool), GateError> { + let file = open_split(path)?; let levers = bloomery_levers::at_main(&[])?; let specs = Place::A.on_host()?.card_specs()?; let [card] = specs[..] else { @@ -256,13 +275,13 @@ mod gate { } /// (t): D3K's prefill by ubatches from a reset, then its step by an eager step with the taps - /// armed, against `D3K` and, recorded, `D3K_BF16`. + /// armed, against the judged set `man` and, recorded, the other one. fn d3k_step( m: &mut Qwen38Model, d: &Dims, vocab: usize, - man: &RefManifest, - bf16: &RefManifest, + (name, man): (&str, &RefManifest), + (other, recorded): (&str, &RefManifest), ) -> Result { let (pos, step, prefill) = man.step()?; let tok = step[0]; @@ -273,7 +292,7 @@ mod gate { let r = run_steps(m, &[tok])?; let ours = r.logits.last().ok_or("no logits")?; let mut ok = true; - for (name, man, judged) in [(D3K, man, true), (D3K_BF16, bf16, false)] { + for (name, man, judged) in [(name, man, true), (other, recorded, false)] { let oracle = ik_last(man, vocab)?; let (top, oracle_top, oracle_2, margin, dist, logits_rel) = tie_numbers(ours, &oracle); let tie = top != oracle_top @@ -302,49 +321,127 @@ mod gate { Ok(ok) } - /// The header's dims, and the file's routed gate type, which must be NVFP4. - fn dims() -> Result { - let file = open_split(Path::new(MODEL))?; - let hp = PlanInputs::describe(&file)?.hp; + /// The dense form of the loaded dense projections of `model`: the head's + /// type, BF16 as the checkpoint stores it (what `prepare` writes by default, judged by the + /// `_bf16` sets) or Q8_0 (what `BLOOMERY_PREPARE_Q8=1` writes, or another converter, judged by + /// the others). Every other BF16 or Q8_0 tensor but the indexer's projections, which `prepare` + /// keeps in BF16 beside either form, must be of the head's type: the sets hold one form or the + /// other. This picks the reference sets only; the model loader refuses Q8_0 indexer projections. + /// A head of another type or none, and a mixed split, are refused by name. + fn dense_of(model: &ModelTensors) -> Result { + let tensors: Vec<(&str, GgmlType)> = model + .tensors + .iter() + .filter(|t| t.role != Role::Unused) + .filter(|t| (model.bf16_words)(t)) + .map(|t| (t.name.as_str(), t.ty)) + .collect(); + let head = names::output(); + let ty = match tensors.iter().find(|(n, _)| *n == head) { + Some(&(_, ty @ (GgmlType::BF16 | GgmlType::Q8_0))) => ty, + Some((_, ty)) => return Err(format!("{head} is {ty:?}, not BF16 or Q8_0")), + None => return Err(format!("the file has no {head}")), + }; + let indexer: Vec = (0..model.layers) + .flat_map(|l| [names::indexer_q_proj(l), names::indexer_k_proj(l)]) + .collect(); + if let Some((name, other)) = tensors.iter().find(|(n, t)| { + matches!(t, GgmlType::BF16 | GgmlType::Q8_0) + && *t != ty + && !indexer.iter().any(|i| i == n) + }) { + return Err(format!( + "{name} is {other:?} beside a {ty:?} {head}, a dense form no pair of sets holds" + )); + } + let q8_indexer = tensors + .iter() + .any(|(n, t)| *t == GgmlType::Q8_0 && indexer.iter().any(|i| i == n)); + Ok(if ty == GgmlType::BF16 { + Dense { + ty, + q8_indexer, + batch: (BATCH_BF16, BATCH), + d3k: (D3K_BF16, D3K), + } + } else { + Dense { + ty, + q8_indexer, + batch: (BATCH, BATCH_BF16), + d3k: (D3K, D3K_BF16), + } + }) + } + + /// The header's dims, the file's routed gate type, which must be NVFP4, and its dense form + /// ([`dense_of`]). + fn dims(path: &Path) -> Result { + let file = open_split(path)?; + let inputs = PlanInputs::describe(&file)?; + let hp = &inputs.hp; let gate = file .find(&names::ffn_gate_exps(0)) .map(|(_, t)| t.ty) .ok_or("the file has no routed gate stack on layer 0")?; if gate != GgmlType::NVFP4 { - return Err(format!("{MODEL}: layer 0's routed gate is {gate:?}, not NVFP4").into()); + return Err(format!( + "{}: layer 0's routed gate is {gate:?}, not NVFP4", + path.display() + ) + .into()); } + let dense = dense_of(&inputs.model).map_err(|e| format!("{}: {e}", path.display()))?; let streams = hp.exp.as_ref().map_or(1, |e| e.hc_streams); Ok(Dims { n_layer: hp.n_trunk, row: streams * hp.n_embd, n_expert: hp.n_expert, n_used: hp.n_used, + dense, }) } pub fn run() -> Result<(), GateError> { - let batch = set(BATCH)?; - let batch_bf16 = set(BATCH_BF16)?; - let d3k = set(D3K)?; - let d3k_bf16 = set(D3K_BF16)?; - for (q8, bf16) in [(&batch, &batch_bf16), (&d3k, &d3k_bf16)] { - if q8.header.tokens != bf16.header.tokens { + let path = ref_model_path()?; + println!("model {}", path.display()); + let d = dims(&path)?; + let Dense { + ty, + q8_indexer, + batch: (batch_name, other_batch), + d3k: (d3k_name, other_d3k), + } = d.dense; + println!( + "dense projections {}: judged by {batch_name} and {d3k_name}; indexer {}, the oracle's BF16", + if ty == GgmlType::BF16 { + "BF16 as stored" + } else { + "Q8_0" + }, + if q8_indexer { "includes Q8_0" } else { "BF16" } + ); + let batch = set(batch_name)?; + let batch_other = set(other_batch)?; + let d3k = set(d3k_name)?; + let d3k_other = set(other_d3k)?; + for (judged, other) in [(&batch, &batch_other), (&d3k, &d3k_other)] { + if judged.header.tokens != other.header.tokens { return Err(format!( "{}: oracle token ids differ from {}", - bf16.dir.display(), - q8.dir.display() + other.dir.display(), + judged.dir.display() ) .into()); } } let (_, toks, _) = batch.step()?; - let d = dims()?; - let (mut m, mut ok) = open(&d)?; + let (mut m, mut ok) = open(&path, &d)?; let vocab = m.body("run")?.vocab(); m.reset()?; let r = run_steps(&mut m, toks)?; ok &= free(&batch, &r, &d, vocab, true)?; - ok &= free(&batch_bf16, &r, &d, vocab, false)?; + ok &= free(&batch_other, &r, &d, vocab, false)?; ok &= batch_by( &mut m, toks, @@ -359,7 +456,113 @@ mod gate { ("g", Prompt38::Gemm, "one ubatch"), &batch, )?; - ok &= d3k_step(&mut m, &d, vocab, &d3k, &d3k_bf16)?; + ok &= d3k_step(&mut m, &d, vocab, (d3k_name, &d3k), (other_d3k, &d3k_other))?; if ok { Ok(()) } else { Err(checks_failed()) } } + + #[cfg(test)] + mod tests { + use super::*; + + /// The sets follow the split's dense form: `prepare`'s default split (every dense matrix + /// BF16) the `_bf16` pair; its `BLOOMERY_PREPARE_Q8=1` split (Q8_0, the indexer's + /// projections BF16) and a Q8_0 split another converter wrote (the indexer Q8_0 too) the + /// Q8_0 pair; a projection of the other form, a head of a third type and no head are + /// refused by name. Mutants: the form read as one type whatever the head, the pairs + /// swapped, the mixed check dropped, the indexer not exempt. + #[test] + fn a_split_is_judged_by_the_sets_of_its_dense_form() { + let split = |dense, indexer| -> Vec<(String, GgmlType)> { + vec![ + (names::output(), dense), + (names::token_embd(), dense), + (names::ssm_out(0), dense), + (names::ffn_down_shexp(0), dense), + (names::ffn_gate_exps(0), GgmlType::NVFP4), + (names::attn_q_norm(3), GgmlType::F32), + (names::attn_output(3), dense), + (names::indexer_q_proj(3), indexer), + (names::indexer_k_proj(3), indexer), + ] + }; + let of = |t: &[(String, GgmlType)]| { + use model::placement::ModelTensor; + dense_of(&ModelTensors { + tensors: t + .iter() + .map(|(n, ty)| ModelTensor { + name: n.clone(), + shard: 0, + layer: None, + role: if n.starts_with("blk.4.") { + Role::Unused + } else { + Role::Head + }, + ty: *ty, + dims: vec![32, 32], + file_bytes: 0, + gathered_rows: None, + }) + .collect(), + layers: 4, + experts: 1, + experts_used: 1, + bf16_words: |t| t.name != names::ffn_gate_inp(0), + }) + }; + let (bf16, q8) = (GgmlType::BF16, GgmlType::Q8_0); + let stored = Dense { + ty: bf16, + q8_indexer: false, + batch: (BATCH_BF16, BATCH), + d3k: (D3K_BF16, D3K), + }; + let rewritten = Dense { + ty: q8, + q8_indexer: false, + batch: (BATCH, BATCH_BF16), + d3k: (D3K, D3K_BF16), + }; + let mut unused = split(q8, bf16); + unused.push((names::indexer_q_proj(4), bf16)); + unused.push((names::attn_output(4), bf16)); + assert_eq!(of(&unused), Ok(rewritten)); + let mut values = split(q8, bf16); + values.push((names::ffn_gate_inp(0), bf16)); + assert_eq!(of(&values), Ok(rewritten)); + assert_eq!(of(&split(bf16, bf16)), Ok(stored)); + assert_eq!(of(&split(q8, bf16)), Ok(rewritten)); + for (ty, expected) in [(bf16, stored), (q8, rewritten)] { + assert_eq!( + of(&split(ty, q8)), + Ok(Dense { + q8_indexer: true, + ..expected + }) + ); + let mut mixed = split(ty, bf16); + mixed[7].1 = q8; + assert_eq!( + of(&mixed), + Ok(Dense { + q8_indexer: true, + ..expected + }) + ); + } + for (at, ty, named) in [ + (3, bf16, names::ffn_down_shexp(0)), + (6, q8, names::attn_output(3)), + (0, GgmlType::Q6_K, format!("{} is Q6_K", names::output())), + ] { + let mut t = split(if ty == q8 { bf16 } else { q8 }, bf16); + t[at].1 = ty; + let e = of(&t).unwrap_err(); + assert!(e.contains(&named), "{e}"); + } + let e = of(&split(q8, bf16)[1..]).unwrap_err(); + assert!(e.contains(&names::output()), "{e}"); + } + } } diff --git a/crates/hf/src/modelopt.rs b/crates/hf/src/modelopt.rs index 017a42295..69cf0cbb5 100644 --- a/crates/hf/src/modelopt.rs +++ b/crates/hf/src/modelopt.rs @@ -1093,6 +1093,11 @@ mod tests { ); let kept = normalize(get, &[("b.weight", &b)], Use::Matrix(Plain::Bf16)).unwrap(); assert_eq!((kept.type_id, &kept.bytes), (30, &bf16)); + let f32 = placed(Format::F32, &[]); + let kept = normalize(get, &[("f.weight", &f32)], Use::Matrix(Plain::F32)).unwrap(); + assert_eq!((kept.type_id, &kept.bytes), (0, &f32s)); + let q = normalize(get, &[("f.weight", &f32)], m).unwrap(); + assert_eq!((q.type_id, &q.bytes), (8, &want)); let rows = normalize( get, &[("t.weight", &placed(Format::Fp8Tensor, &["t.scale"]))], diff --git a/crates/levers/src/lib.rs b/crates/levers/src/lib.rs index da000adae..e5c138e34 100644 --- a/crates/levers/src/lib.rs +++ b/crates/levers/src/lib.rs @@ -39,7 +39,7 @@ pub use registry::{ PREFILL_GROUP_DEFAULT, PREFILL_GROUP_MAX, QWEN3_KV, QWEN38_EXPERTS, R8, RESIDENCY, ROUTE_TRACE, SCREEN, SPIN, STEP_STATS, THREADS, TIER_LINK, TIER_NODES, TP_NODES, WARM_PROFILE, }; -pub use registry::{SCHEMAJUMP, SPLIT_WALK, XSTREAM}; +pub use registry::{PREPARE_Q8, SCHEMAJUMP, SPLIT_WALK, XSTREAM}; #[cfg(test)] mod tests; @@ -1183,6 +1183,13 @@ impl Levers { self.flag(SCHEMAJUMP) } + /// `BLOOMERY_PREPARE_Q8`: `prepare` writes the BF16 or F32 matrices its map + /// reads as Q8_0 in Q8_0, not as stored. + #[must_use] + pub fn prepare_q8(&self) -> bool { + self.flag(PREPARE_Q8) + } + /// The host tier's load settings: [`HOST_POPULATE`], [`HOST_LOCK`], /// [`CARD_DONTNEED`], [`R8`], and whether [`ROUTE_TRACE`] is set. #[must_use] @@ -2025,7 +2032,7 @@ pub enum Draft38Off { /// Unset, and the last window needs `need` positions of `ctx`. Ctx { need: usize, ctx: usize }, /// Unset on the Qwen3.8 seat, and the target's `name`, a matrix the - /// draft borrows, is `ty` (or `absent`), not the Q8_0 the draft reads. + /// draft borrows, is `ty` (or `absent`), not a form the draft reads. Borrowed { name: String, ty: String }, /// Unset, and the draft yielded to the context /// ([`RungCost::of`]'s answer, the seats' plan-side rule). @@ -2072,7 +2079,7 @@ impl fmt::Display for Draft38Off { ), Draft38Off::Borrowed { name, ty } => write!( f, - "unset: the target's {name} is {ty}; the MTP draft reads it as Q8_0" + "unset: the target's {name} is {ty}, not a form the MTP draft reads" ), Draft38Off::Yield(y) => y.fmt(f), Draft38Off::Paged { arena } => write!( diff --git a/crates/levers/src/registry.rs b/crates/levers/src/registry.rs index e78d0490c..b904e7653 100644 --- a/crates/levers/src/registry.rs +++ b/crates/levers/src/registry.rs @@ -51,6 +51,7 @@ pub const TIER_NODES: &str = "BLOOMERY_TIER_NODES"; pub const TP_NODES: &str = "BLOOMERY_TP_NODES"; pub const SCREEN: &str = "BLOOMERY_SCREEN"; pub const SCHEMAJUMP: &str = "BLOOMERY_SCHEMAJUMP"; +pub const PREPARE_Q8: &str = "BLOOMERY_PREPARE_Q8"; /// The largest `BLOOMERY_PREFILL_GROUP`: the batches a V4.1 prompt group /// holds at most, which the body's buffers are sized for. @@ -1073,6 +1074,30 @@ pub(crate) static REGISTRY: &[LeverSpec] = &[ phase: Phase::Load, }, }, + LeverSpec { + name: PREPARE_Q8, + class: Class::M, + kind: Kind::Flag, + default: Unset::Is("0"), + doc: "`prepare`: `1` writes each BF16 or F32 matrix its checkpoint map reads as Q8_0 (the \ + attention and linear-attention projections, the shared expert, the PLE key and \ + value, the hyper-connections' down and up, the head and the token embedding) as \ + ggml's reference Q8_0; `0` keeps each single BF16 matrix as stored and refuses \ + F32 matrices read as Q8_0. BF16 and F32 expert stacks are refused, and the \ + quantized tensors (NVFP4, MXFP8, FP8) keep their stored blocks, under both \ + values. The rewrite is lossy: on \ + `nvidia/Qwen3.8-Flash-Next-NVFP4` at fc694b54, the checkpoint's own f32 forward \ + with its BF16 projections rounded to Q8_0 is at a KL divergence of 4.288e-02 from \ + the forward with them as stored, at the last of the five positions of the \ + batch set, and 4.181e-02 at position 3,000 of the D3K set (refset \ + `hf-qwen4exp-nvfp4`, the `_bf16` sets against the others), so it is asked for, \ + never a default. `--keep-bf16`, which asks for the stored form, is refused \ + beside it.", + site: Site::Parsed { + left: &[], + phase: Phase::Load, + }, + }, LeverSpec { name: "BLOOMERY_POISON", class: Class::D, diff --git a/crates/levers/src/tests.rs b/crates/levers/src/tests.rs index 78fa54365..b4bce1e82 100644 --- a/crates/levers/src/tests.rs +++ b/crates/levers/src/tests.rs @@ -404,6 +404,7 @@ fn accessors_read_their_rows() { assert_eq!(unset.host_lanes(), HOST_LANES_DEFAULT); assert_eq!(unset.nvtier_mem(), NvTierMem::Auto); assert_eq!(unset.gen_slots(), 1); + assert!(!unset.prepare_q8()); assert_eq!( unset.host(), HostCfg { @@ -448,6 +449,7 @@ fn accessors_read_their_rows() { (NVTIER_MEM, "pinned"), (PREFILL_GROUP, "8"), (GEN_SLOTS, "2"), + (PREPARE_Q8, "1"), ]), Scope::Every, ) @@ -498,6 +500,7 @@ fn accessors_read_their_rows() { ); } assert_eq!(set.gen_slots(), 2); + assert!(set.prepare_q8()); assert_eq!( set.host(), HostCfg { @@ -553,7 +556,8 @@ fn accessors_read_their_rows() { GEN_SLOTS, FORCE_IDS, SCREEN, - SCHEMAJUMP + SCHEMAJUMP, + PREPARE_Q8 ] ); } @@ -1303,10 +1307,10 @@ fn draft38_unset_follows_the_place_and_the_file() { assert_eq!( Draft38Off::Borrowed { name: "output.weight".to_string(), - ty: "q6_K".to_string() + ty: "f16".to_string() } .to_string(), - "unset: the target's output.weight is q6_K; the MTP draft reads it as Q8_0" + "unset: the target's output.weight is f16, not a form the MTP draft reads" ); } @@ -1449,7 +1453,7 @@ fn a_paged_plan_turns_only_an_unset_draft_off() { Draft38Off::Ctx { need: 5, ctx: 4 }, Draft38Off::Borrowed { name: "output.weight".to_string(), - ty: "q6_K".to_string(), + ty: "f16".to_string(), }, ] { assert_eq!( diff --git a/crates/model-hf/Cargo.toml b/crates/model-hf/Cargo.toml index 7482df569..bb109eb6a 100644 --- a/crates/model-hf/Cargo.toml +++ b/crates/model-hf/Cargo.toml @@ -15,6 +15,8 @@ workspace = true [dependencies] bloomery-gguf = { path = "../gguf" } bloomery-hf = { path = "../hf" } +# The `prepare` bin reads its one lever, `BLOOMERY_PREPARE_Q8`. +bloomery-levers = { path = "../levers" } bloomery-model = { path = "../model" } bloomery-safetensors = { path = "../safetensors" } # The `prepare` bin builds serving's tokenizer from the GGUF header it wrote and checks its split. diff --git a/crates/model-hf/src/bin/prepare.rs b/crates/model-hf/src/bin/prepare.rs index 7d1778630..9d6b3e7ed 100644 --- a/crates/model-hf/src/bin/prepare.rs +++ b/crates/model-hf/src/bin/prepare.rs @@ -8,11 +8,16 @@ //! //! writes `-00001-of-0000N.gguf` and on, each shard at most `n` GiB (48 by default) //! unless one tensor with its scale and file layout is larger, the metadata in the first. -//! `--keep-bf16` writes a BF16 matrix the map would make Q8_0 as it is stored, the head and the -//! token embedding among them; the server reads each through its BF16 site. Qwen3.8's decode then -//! reads about 9.27 GiB of weights a token where the Q8_0 split reads 5.73, the head alone 0.55 -//! GiB more (1,271,398,400 bytes of BF16 against 34 bytes a 32 values). Each shard is -//! written as `.part` and renamed when whole. The disk must hold the split and a tenth +//! A quantized tensor keeps its stored format (NVFP4, MXFP8, FP8), and a BF16 matrix stays BF16, +//! the head and the token embedding among them; the server reads each through its BF16 site. +//! BF16 and F32 expert stacks are refused by name. An F32 matrix the map reads as Q8_0 is +//! refused unless `BLOOMERY_PREPARE_Q8=1` writes each single BF16 or F32 matrix read as Q8_0 as +//! ggml's reference Q8_0 instead, a lossy rewrite whose cost its lever row states; `--keep-bf16` +//! asks for the stored form and is refused beside it. Qwen3.8's decode reads about 9.27 GiB +//! of weights a token from the stored split where the Q8_0 one reads 5.73, the head alone +//! 0.55 GiB more +//! (1,271,398,400 bytes of BF16 against 34 bytes a 32 values). Each shard is written as +//! `.part` and renamed when whole. The disk must hold the split and a tenth //! more. The written split is opened and its hyperparameters read as the server reads them. //! One line per tensor, a summary line, and exit status 1 with the error on any failure. Before //! any shard is written, the tokenizer serving builds from the header is checked to split text @@ -25,10 +30,11 @@ use std::path::{Path, PathBuf}; use std::process::ExitCode; use std::time::Instant; +use bloomery_levers::PREPARE_Q8; use gguf::write::{Layout, TensorDecl, Writer}; use gguf::{GgmlType, Value}; use hf::modelopt::{self, Normal, Plain, Use}; -use hf::quant_config::{self, Config, Format}; +use hf::quant_config::{self, Config, Format, Plan}; use model_hf::arch::qwen35moe::{Checkpoint, Mapped}; use safetensors::{Dtype, Folder}; @@ -37,8 +43,13 @@ const USAGE: &str = "usage: prepare [--shard-gib ] [--k type Res = Result>; fn main() -> ExitCode { - let args: Vec = std::env::args().skip(1).collect(); - match run(&args) { + let ran = bloomery_levers::at_main(&[PREPARE_Q8]) + .map_err(Into::into) + .and_then(|levers| { + let args: Vec = std::env::args().skip(1).collect(); + run(&args, levers.prepare_q8()) + }); + match ran { Ok(()) => ExitCode::SUCCESS, Err(e) => { eprintln!("prepare: error: {e}"); @@ -60,7 +71,9 @@ enum Make { Scale { of: String }, } -fn run(args: &[String]) -> Res<()> { +/// The split of the folder `args` name; `q8` (`BLOOMERY_PREPARE_Q8`) writes its BF16 matrices the +/// map reads as Q8_0, and its F32 matrices, in Q8_0. +fn run(args: &[String], q8: bool) -> Res<()> { let (mut shard_gib, mut keep_bf16, mut paths) = (48u64, false, vec![]); let mut it = args.iter(); while let Some(a) = it.next() { @@ -73,6 +86,12 @@ fn run(args: &[String]) -> Res<()> { let [dir, prefix] = paths[..] else { return Err(USAGE.into()); }; + if keep_bf16 && q8 { + return Err(format!( + "--keep-bf16 keeps the BF16 matrices {PREPARE_Q8}=1 asks to write as Q8_0: set one" + ) + .into()); + } let limit = u64::checked_mul(shard_gib, 1 << 30) .filter(|&l| l > 0) .ok_or_else(|| { @@ -91,77 +110,7 @@ fn run(args: &[String]) -> Res<()> { .map(|r| (n.to_string(), r.dtype, r.shape.to_vec())) }) .collect::>()?; - let plan = quant_config::plan( - &Config::read(dir)?, - entries - .iter() - .map(|(n, d, s)| (n.as_str(), *d, s.as_slice())), - ); - let logical: Vec<(&str, Vec)> = entries - .iter() - .filter_map( - |(n, _, s)| match (plan.placed.get(n), plan.refused.contains_key(n)) { - (Some(p), _) => Some((n.as_str(), p.format.logical(s))), - (None, true) => Some((n.as_str(), s.clone())), - (None, false) => None, - }, - ) - .collect(); - let mapped = checkpoint.map(logical.iter().map(|(n, s)| (*n, s.as_slice())))?; - - let mut planned = Vec::new(); - for (name, mut m) in mapped { - let source = &m.sources[0]; - let format = modelopt::format_of(&m.placed(&plan)?)?; - if keep_bf16 && format == Format::Bf16 && m.role == Use::Matrix(Plain::Q8_0) { - m.role = Use::Matrix(Plain::Bf16); - } - let type_id = modelopt::type_of(format, m.role).ok_or_else(|| { - format!( - "{source}: {format} read as {:?} has no GGUF type here", - m.role - ) - })?; - let nbytes = size_of(type_id, &m.dims).ok_or_else(|| { - format!( - "{name}: dims {:?} are not whole blocks of type {type_id}", - m.dims - ) - })?; - // One tensor scale an expert of a stack; one for a row table, whatever its parts. - let scales = GgmlType::from_u32(type_id) - .reads_scale() - .then_some(match m.role { - Use::Rows => 1, - Use::Matrix(_) => m.sources.len(), - }); - let decl = TensorDecl { - name: name.clone(), - dims: m.dims.clone(), - type_id, - nbytes, - }; - planned.push(Planned { - decl, - make: Make::Tensor(m), - }); - if let Some(n) = scales { - let n = u64::try_from(n).map_err(|_| format!("{name}: scale count exceeds u64"))?; - let nbytes = size_of(0, &[n]) - .ok_or_else(|| format!("{name}: scale byte count overflows u64"))?; - let stem = name.strip_suffix(".weight").unwrap_or(&name); - let decl = TensorDecl { - name: format!("{stem}.scale"), - dims: vec![n], - type_id: 0, - nbytes, - }; - planned.push(Planned { - decl, - make: Make::Scale { of: name }, - }); - } - } + let (plan, planned) = plan_tensors(&checkpoint, &Config::read(dir)?, &entries, q8)?; let mut kvs = checkpoint.header( planned @@ -286,6 +235,104 @@ fn run(args: &[String]) -> Res<()> { Ok(()) } +/// The folder's `entries` (name, dtype, shape) placed by its `config`, and each tensor the map +/// makes of them declared in the type its placed format and role give, its `.scale` beside it +/// where the type reads one. An unquantized matrix the map reads as Q8_0 stays in its stored +/// format unless `q8` asks for the rewrite; F32 matrices need that rewrite. Unquantized expert +/// stacks are unsupported. +fn plan_tensors( + checkpoint: &Checkpoint, + config: &Config, + entries: &[(String, Dtype, Vec)], + q8: bool, +) -> Res<(Plan, Vec)> { + let plan = quant_config::plan( + config, + entries + .iter() + .map(|(n, d, s)| (n.as_str(), *d, s.as_slice())), + ); + let logical: Vec<(&str, Vec)> = entries + .iter() + .filter_map( + |(n, _, s)| match (plan.placed.get(n), plan.refused.contains_key(n)) { + (Some(p), _) => Some((n.as_str(), p.format.logical(s))), + (None, true) => Some((n.as_str(), s.clone())), + (None, false) => None, + }, + ) + .collect(); + let mapped = checkpoint.map(logical.iter().map(|(n, s)| (*n, s.as_slice())))?; + let mut planned = Vec::new(); + for (name, mut m) in mapped { + let source = &m.sources[0]; + let format = modelopt::format_of(&m.placed(&plan)?)?; + if m.role == Use::Matrix(Plain::Q8_0) && matches!(format, Format::Bf16 | Format::F32) { + if m.dims.len() == 3 { + return Err(format!( + "{source}: {name}: {format} expert stacks are unsupported by the converter" + ) + .into()); + } + if !q8 { + if format == Format::F32 { + return Err(format!( + "{source}: {name} is stored {format}, a form its card site does not read; {PREPARE_Q8}=1 writes Q8_0" + ) + .into()); + } + m.role = Use::Matrix(Plain::Bf16); + } + } + let type_id = modelopt::type_of(format, m.role).ok_or_else(|| { + format!( + "{source}: {format} read as {:?} has no GGUF type here", + m.role + ) + })?; + let nbytes = size_of(type_id, &m.dims).ok_or_else(|| { + format!( + "{name}: dims {:?} are not whole blocks of type {type_id}", + m.dims + ) + })?; + // One tensor scale an expert of a stack; one for a row table, whatever its parts. + let scales = GgmlType::from_u32(type_id) + .reads_scale() + .then_some(match m.role { + Use::Rows => 1, + Use::Matrix(_) => m.sources.len(), + }); + let decl = TensorDecl { + name: name.clone(), + dims: m.dims.clone(), + type_id, + nbytes, + }; + planned.push(Planned { + decl, + make: Make::Tensor(m), + }); + if let Some(n) = scales { + let n = u64::try_from(n).map_err(|_| format!("{name}: scale count exceeds u64"))?; + let nbytes = size_of(0, &[n]) + .ok_or_else(|| format!("{name}: scale byte count overflows u64"))?; + let stem = name.strip_suffix(".weight").unwrap_or(&name); + let decl = TensorDecl { + name: format!("{stem}.scale"), + dims: vec![n], + type_id: 0, + nbytes, + }; + planned.push(Planned { + decl, + make: Make::Scale { of: name }, + }); + } + } + Ok((plan, planned)) +} + fn plan_shards( planned: Vec, kvs: &[(String, Value)], @@ -749,8 +796,166 @@ mod tests { fn a_shard_size_out_of_range_is_refused() { for n in ["0", "17179869184", "17179869185"] { let args = ["no-folder", "out", "--shard-gib", n].map(String::from); - let e = super::run(&args).unwrap_err().to_string(); + let e = super::run(&args, false).unwrap_err().to_string(); assert!(e.contains("--shard-gib"), "{n}: {e}"); } } + + /// Every tensor of layers 0, 1 and 3 and the unlayered text tensors of + /// nvidia/Qwen3.8-Flash-Next-NVFP4 is declared in its stored format: each one unsloth's + /// UD-Q4_K_XL holds in Q8_0, a BF16 matrix in the checkpoint (the head and the token + /// embedding among them), is BF16, and Q8_0 only when `BLOOMERY_PREPARE_Q8` asks; every other + /// tensor (the NVFP4 stacks, the PLE table's F8 rows, the gains, the scales) is the same + /// either way. Synthetic F32 matrices require the Q8_0 rewrite; BF16 and F32 routed stacks + /// are refused with either lever value. Dropping either refusal fails these assertions. + /// `--keep-bf16` asks for the stored form: taken alone, refused beside the lever. + #[test] + fn a_bf16_matrix_is_written_as_stored_unless_q8_is_asked() { + use super::*; + use model::arch::qwen35moe::names; + use std::collections::BTreeMap; + + let gz = |name: &str| -> serde_json::Value { + let path = format!("{}/../hf/tests/fixtures/{name}", env!("CARGO_MANIFEST_DIR")); + let file = File::open(&path).unwrap_or_else(|e| panic!("{path}: {e}")); + serde_json::from_reader(flate2::read::GzDecoder::new(file)).unwrap() + }; + let config = gz("quant-configs/nvidia__Qwen3.8-Flash-Next-NVFP4__config.json.gz"); + let hf_quant = + gz("quant-configs/nvidia__Qwen3.8-Flash-Next-NVFP4__hf_quant_config.json.gz"); + let quant = Config::parse(&config, Some(&hf_quant)).unwrap(); + let checkpoint = Checkpoint::parse(&config).unwrap(); + let header = gz("checkpoint-headers/qwen3.8-flash-next-nvfp4-layers-0-1-3.tensors.json.gz"); + let entries: Vec<(String, Dtype, Vec)> = header + .as_object() + .unwrap() + .iter() + .map(|(n, e)| { + let shape = e["shape"].as_array().unwrap(); + ( + n.clone(), + Dtype::of(e["dtype"].as_str().unwrap()).unwrap(), + shape.iter().map(|d| d.as_u64().unwrap() as usize).collect(), + ) + }) + .collect(); + let unsloth: Vec<(String, Vec, u32)> = serde_json::from_value(gz( + "checkpoint-headers/unsloth-qwen3.8-flash-next-ud-q4_k_xl-layers-0-1-3.tensors.json.gz", + )) + .unwrap(); + let q8_0: Vec<&str> = unsloth + .iter() + .filter(|(_, _, ty)| *ty == GgmlType::Q8_0.as_u32()) + .map(|(n, _, _)| n.as_str()) + .collect(); + for name in [names::output(), names::token_embd(), names::attn_output(3)] { + assert!(q8_0.contains(&name.as_str()), "{name}"); + } + let types = |q8| -> BTreeMap { + let (_, planned) = plan_tensors(&checkpoint, &quant, &entries, q8).unwrap(); + planned + .into_iter() + .map(|p| (p.decl.name, GgmlType::from_u32(p.decl.type_id))) + .collect() + }; + let (stored, rewritten) = (types(false), types(true)); + assert_eq!( + stored.keys().collect::>(), + rewritten.keys().collect::>() + ); + for (name, ty) in &stored { + let want = if q8_0.contains(&name.as_str()) { + (GgmlType::BF16, GgmlType::Q8_0) + } else { + (*ty, *ty) + }; + assert_eq!((*ty, rewritten[name]), want, "{name}"); + } + for (target, source) in [ + ( + names::attn_output(3), + "model.language_model.layers.3.self_attn.o_proj.weight", + ), + (names::output(), "lm_head.weight"), + ( + names::token_embd(), + "model.language_model.embed_tokens.weight", + ), + ] { + let mut f32 = entries.clone(); + f32.iter_mut().find(|(n, _, _)| n == source).unwrap().1 = Dtype::F32; + for q8 in [false, true] { + let result = plan_tensors(&checkpoint, &quant, &f32, q8); + if !q8 { + let Err(e) = result else { + panic!("F32 matrix {target} must be refused"); + }; + let e = e.to_string(); + assert!( + e.contains(source) && e.contains(&target) && e.contains(PREPARE_Q8), + "{e}" + ); + } else { + let (_, planned) = result.unwrap(); + let p = planned.iter().find(|p| p.decl.name == target).unwrap(); + assert_eq!(p.decl.type_id, 8, "{target}"); + } + } + } + let bf16: Vec<_> = entries + .iter() + .filter_map(|(n, d, s)| { + if n.contains("layers.0.mlp.experts.") && n.contains(".gate_proj.") { + n.ends_with(".weight") + .then(|| (n.clone(), Dtype::Bf16, Format::Nvfp4.logical(s))) + } else { + Some((n.clone(), *d, s.clone())) + } + }) + .collect(); + let mut single = config.clone(); + single["text_config"]["num_experts"] = 1.into(); + single["text_config"]["num_experts_per_tok"] = 1.into(); + let single = Checkpoint::parse(&single).unwrap(); + for dtype in [Dtype::Bf16, Dtype::F32] { + let mut stack = bf16.clone(); + for (n, d, _) in &mut stack { + if n.contains("layers.0.mlp.experts.") && n.contains(".gate_proj.") { + *d = dtype; + } + } + let source = "model.language_model.layers.0.mlp.experts.0.gate_proj.weight"; + let one: Vec<_> = stack + .iter() + .filter(|(n, _, _)| { + n == source + || n == "lm_head.weight" + || n == "model.language_model.embed_tokens.weight" + }) + .cloned() + .collect(); + for (checkpoint, stack) in [(&checkpoint, &stack), (&single, &one)] { + for q8 in [false, true] { + let Err(e) = plan_tensors(checkpoint, &quant, stack, q8) else { + panic!( + "{dtype:?} expert stacks must be refused for q8={q8}, entries={}", + stack.len() + ); + }; + let e = e.to_string(); + assert!( + e.contains(source) + && e.contains(&names::ffn_gate_exps(0)) + && e.contains("expert stacks are unsupported"), + "{e}" + ); + } + } + } + for q8 in [false, true] { + let args = ["no-folder", "out", "--keep-bf16"].map(String::from); + let e = run(&args, q8).unwrap_err().to_string(); + assert_eq!(e.contains("--keep-bf16"), q8, "q8={q8}: {e}"); + } + } } diff --git a/crates/model/src/arch/qwen35moe/mtp.rs b/crates/model/src/arch/qwen35moe/mtp.rs index b33bade73..ff8cbabd8 100644 --- a/crates/model/src/arch/qwen35moe/mtp.rs +++ b/crates/model/src/arch/qwen35moe/mtp.rs @@ -1367,7 +1367,7 @@ mod tests { } /// A target whose load keeps its BF16 head and embedding in their words - /// (`prepare --keep-bf16`) is admitted with both borrowed, and + /// (`prepare`'s default, the stored form) is admitted with both borrowed, and /// `head_kind` records the BF16 form the draft's head reads. #[test] fn a_bf16_target_head_and_embedding_are_admitted() { diff --git a/crates/refset/src/arch/qwen4exp/nvfp4.rs b/crates/refset/src/arch/qwen4exp/nvfp4.rs index 130f3bb05..36bc19657 100644 --- a/crates/refset/src/arch/qwen4exp/nvfp4.rs +++ b/crates/refset/src/arch/qwen4exp/nvfp4.rs @@ -1,11 +1,12 @@ //! The Qwen3.8-Flash-Next NVFP4 oracle: the forward of `nvidia/Qwen3.8-Flash-Next-NVFP4` in -//! transformers' qwen4_exp, f32 on the CPU, with the weights the split `prepare` writes from that +//! transformers' qwen4_exp, f32 on the CPU, with the weights a split `prepare` writes from that //! folder holds — the routed experts by ModelOpt's NVFP4 rule, the PLE rows by their F8 rule, the -//! BF16 projections `prepare` writes as Q8_0 rounded by ggml's reference Q8_0 rule -//! (`tools/ref/qwen4exp-nvfp4-oracle.py`). Its sets are in ik's dump layout and carry ik's tap -//! names, so a gate reads them as it reads [`super::IK`]'s; the ids are the batch set's and -//! [`super::D3K`]'s. The two `_bf16` sets widen the checkpoint's BF16 projections instead: the -//! Q8_0 rounding's own share of the difference, recorded beside the oracle. +//! BF16 projections `prepare` rewrites to Q8_0 under `BLOOMERY_PREPARE_Q8=1` rounded by ggml's +//! reference Q8_0 rule (`tools/ref/qwen4exp-nvfp4-oracle.py`). Its sets are in ik's dump layout +//! and carry ik's tap names, so a gate reads them as it reads [`super::IK`]'s; the ids are the +//! batch set's and [`super::D3K`]'s. The two `_bf16` sets widen the checkpoint's BF16 projections +//! instead, as the split `prepare` writes by default holds them; between the two pairs lies the +//! Q8_0 rounding's own share. use super::ARCH; use crate::RefError; @@ -28,7 +29,7 @@ pub const BATCH: &str = "hf-qwen4exp-nvfp4/hf_qwen4exp_nvfp4"; /// The 3,001 ids of [`super::D3K`] as a decode-step set: the step at position 3,000, every /// position's taps kept. pub const D3K: &str = "hf-qwen4exp-nvfp4/hf_qwen4exp_nvfp4_d3k"; -/// [`BATCH`] with the BF16 projections widened. +/// [`BATCH`] with the BF16 projections widened, as stored. pub const BATCH_BF16: &str = "hf-qwen4exp-nvfp4/hf_qwen4exp_nvfp4_bf16"; /// [`D3K`] with the BF16 projections widened: the layer outputs and logits of the last position, /// the routes of every position. diff --git a/justfile b/justfile index e03ab2df4..dd65c24a5 100644 --- a/justfile +++ b/justfile @@ -1405,11 +1405,13 @@ gate-gpu-qwen4exp-iqleg: # Qwen3.8-Flash-Next from the NVIDIA NVFP4 checkpoint (the split `prepare` writes from nvidia/Qwen3.8-Flash-Next-NVFP4) on # one card whose plan holds every routed expert (`--place a`; a DGX Spark's GB10): the load puts no expert on the host; the # batch set's eager steps hold every layer's streams within the e2e gate's band of the checkpoint's own forward (refset -# hf-qwen4exp-nvfp4), the router's flips named, and its last argmax; the batch set by one ubatch and D3K's step after its -# prefill by ubatches hold the oracle's argmax (D3K's a named tie allowed); the distances to the BF16-dense sets are -# recorded. Real-only: the fixture holds no NVFP4 stack. +# hf-qwen4exp-nvfp4, the sets of the split's dense form: BF16 as stored, or the Q8_0 of BLOOMERY_PREPARE_Q8=1), the +# router's flips named, and its last argmax; the batch set by one ubatch and D3K's step after its prefill by ubatches hold +# the oracle's argmax (D3K's a named tie allowed); the distances to the other form's sets are recorded. The bin's unit +# test, which holds the choice of sets to the split's dense form, runs first. BLOOMERY_REF_MODEL moves the split from +# the model folder below, including to prepare's stored-form output. Real-only: the fixture holds no NVFP4 stack. gate-gpu-qwen4exp-nvfp4: - BLOOMERY_MODEL=qwen4exp ./tools/box.sh 'bash tools/ref/real-only.sh gate-gpu-qwen4exp-nvfp4 && cargo oxide build --arch {{arch}} -- -p bloomery-gpu-gates --features gpu --release --bin gate_qwen4exp_nvfp4 && BLOOMERY_GATE_CARD=${BLOOMERY_GATE_CARD:-any} bash tools/gpu-gate.sh gate_qwen4exp_nvfp4' + BLOOMERY_REF_MODEL="${BLOOMERY_REF_MODEL:-/models/Qwen3.8-Flash-Next-NVFP4/Qwen3.8-Flash-Next-NVFP4-00001-of-00004.gguf}" BLOOMERY_MODEL=qwen4exp ./tools/box.sh 'bash tools/ref/real-only.sh gate-gpu-qwen4exp-nvfp4 && bash tools/gate.sh --oxide -p bloomery-gpu-gates --release --features gpu --bin gate_qwen4exp_nvfp4 && cargo oxide build --arch {{arch}} -- -p bloomery-gpu-gates --features gpu --release --bin gate_qwen4exp_nvfp4 && BLOOMERY_GATE_CARD=${BLOOMERY_GATE_CARD:-any} bash tools/gpu-gate.sh gate_qwen4exp_nvfp4' # Qwen3.8's MTP draft program on the 3090 beside the target: first, before any walk, the draft's load — its weights, # store and row map hold its plan's bytes, its token_embd and output are the target's own buffers at their addresses,