diff --git a/docs/supported-models.md b/docs/supported-models.md index a558ff696..9550bc9d6 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -53,6 +53,7 @@ Implemented model families include: - GPT-BigCode (`gpt_bigcode`, BigCode: the StarCoder / SantaCoder code models, a GPT-2 derivative whose defining change is Multi-Query Attention. With `multi_query: true` a single KV head is shared by every query head, so the fused `c_attn` projection produces `n_embd + 2 * kv_dim` features rather than `3 * n_embd` and is split at `[n_embd, n_embd + kv_dim]`, which is not an even three-way split; on `bigcode/gpt_bigcode-santacoder` that is 2048 + 2 * 128 = 2304 for 16 query heads of width 128 over one KV head. MQA maps onto the existing GQA attention path with `n_kv_heads = 1`, so no new attention code is involved. The rest of the block follows GPT-2: learned absolute position embeddings added to the token embeddings at the input boundary with no RoPE anywhere, `LayerNorm` with bias, and a tied output head. Unlike GPT-2 there is no `Conv1D` weight layout. HuggingFace `GPTBigCode` builds its projections with `nn.Linear`, so `c_attn` / `c_proj` / `c_fc` are already stored `[out, in]` (`transformer.h.0.attn.c_attn.weight` is `[2304, 2048]`, where the GPT-2 equivalent is `[768, 2304]`) and the loader must not transpose them; a weight that arrives in the transposed orientation is rejected by name rather than accepted. The MLP width comes from `n_inner`, not a hardcoded `4 * n_embd`, and there is no `h.N.attn.bias` causal-mask buffer to strip because HuggingFace registers this family's mask non-persistently. Loading rejects a `wpe` or `wte` table with fewer rows than `config.json`'s `n_positions` / `vocab_size` claims, a projection or norm whose shape disagrees with the config in any layer, and an untied config with no `lm_head` tensor, since the position clamp and the token lookup both trust config-derived bounds and an embedding gather does not range-check a positive index. A quantized `c_attn` is packed on the input axis only, so its output row count is checked against the config-derived width on the same path as an unquantized checkpoint; without that check, a 4-bit or 8-bit conversion whose `config.json` omits `multi_query` (which then defaults to `true`) against a packed weight sized for full multi-head attention would leave every slice in bounds while `Attention::forward`'s config-derived offsets read K and V from the wrong channels. Validated checkpoint: `bigcode/gpt_bigcode-santacoder` (f16, 1.1B). Code checkpoints ship no chat template, so the CLI renders the prompt verbatim.) - GPT-NeoX (`gpt_neox`, EleutherAI: the decoder behind the Pythia suite and the NeoX-derived checkpoints. Three things define it. First, the fused `query_key_value` projection is **interleaved per head**: it is `[3 * hidden_size, hidden_size]` in plain `nn.Linear` orientation, but its output is not three contiguous `hidden_size` blocks. The projection is reshaped to `(..., num_heads, 3 * head_dim)` and split on *that* axis, so the layout is head-major, `[q_i | k_i | v_i]` per head; on `EleutherAI/pythia-1b` that is 8 heads of width 256 contributing 768 contiguous channels each. A flat three-way split of the last axis (the GPT-2 and GPT-BigCode pattern) yields Q, K and V of exactly the right shape assembled from the wrong channels, so nothing throws and the model still emits fluent English out of a scrambled attention; the unit tests therefore pin the layout by channel value, not by shape. Second, **partial RoPE**: only `int(head_dim * rotary_pct)` channels of each head rotate and the rest pass through untouched (64 of 256 for Pythia's `rotary_pct` 0.25), expressed through the `dims` argument of `fast_rope` with `traditional = false`. Third, the **optional parallel residual**: with `use_parallel_residual: true` (every Pythia checkpoint) the attention and MLP sub-layers both read the same pre-norm input and their outputs are summed into the residual, `x + attn(input_layernorm(x)) + mlp(post_attention_layernorm(x))`; with `false` the block is the ordinary chained form where the MLP norm reads the post-attention residual. Both layouts are implemented, and running the wrong one is a silent quality regression rather than a crash. The rest of the block is `LayerNorm` with bias (not RMSNorm), no learned position embeddings at all, an MLP of `dense_4h_to_h(gelu(dense_h_to_4h(x)))` with no gate/up pattern, and an untied `embed_out` head. HuggingFace registers three per-layer PyTorch buffers that reach the checkpoint and are dropped at load: `attention.bias` (a `[1, 1, n, n]` causal mask whose name collides with the projection-bias namespace, exactly as GPT-2's `h.N.attn.bias` does), `attention.masked_bias`, and `attention.rotary_emb.inv_freq` (a precomputed RoPE table that `fast_rope` recomputes from `rotary_emb_base`). Both the raw HuggingFace key layout (`gpt_neox.layers.N`, top-level `embed_out`) and the layout produced by upstream mlx-lm's `sanitize` (`model.h.N`, `model.embed_out`) load. Loading rejects an `embed_in` table with fewer rows than `config.json`'s `vocab_size` claims, since an embedding gather wraps a negative index but does not range-check a positive one; a projection, bias or norm whose shape disagrees with the config in any layer, including on the quantized path, where packing compresses the input axis only and the output row count is still the width the fused reshape depends on; a `rotary_pct` whose rotary dimension count is zero, negative, odd or larger than `head_dim`, because MLX enforces exactly that contract on `rope`'s `dims` argument by throwing, and an MLX C++ exception crossing the cxx bridge is an uncatchable `std::terminate` at the first forward pass rather than a load error; a non-positive or non-finite `rotary_emb_base`, which would make every rotated channel NaN; a non-finite, negative or zero `layer_norm_eps`, since `fast::layer_norm` never looks at `eps` and computes `x * rsqrt(mean(x^2) + eps)`, so a bad value turns every hidden state into NaN without anything throwing and the checkpoint would otherwise load cleanly and generate uniform garbage; and a `quantization` block whose `bits` falls outside `1..=32` or whose `group_size` is non-positive, since both reach `quantized_matmul` unreconciled and MLX's own check on the derived unpacked width divides by zero or collapses to zero before throwing, the same uncatchable-`std::terminate` shape the rope guard exists to prevent. Validated checkpoint: `EleutherAI/pythia-1b` (f16, 244 tensors). Pythia ships no chat template, so the CLI renders the prompt verbatim and `--no-chat-template` is optional.) - GPT-OSS +- Kyutai Helium (`helium`, Kyutai: a dense Llama-shaped decoder. RMSNorm before attention and before the MLP, grouped-query attention, a SwiGLU MLP over `gate_proj` / `up_proj` / `down_proj`, no QK-norm, no MoE, no sliding window, and Llama's own weight key names, so the decoder block, attention and MLP are the existing dense path unchanged. **The one architectural difference is the RoPE convention**: upstream builds `nn.RoPE(head_dim, traditional=True, base=rope_theta)`, so Helium rotates interleaved channel pairs `(2i, 2i+1)` where every other Llama-family model in this tree rotates the split-half pairs `(i, i + dims/2)`. The two produce identically shaped tensors from identical weights, so running the wrong one is a silent quality regression rather than a crash, and no shape assertion can catch it. The flag is threaded from the loader through to every RoPE call the family can reach, including the batched decode path; it also disables the two fused quantized RoPE fast paths (`MLXCEL_ENABLE_FUSED_QKV_SPLIT_ROPE` and `MLXCEL_ENABLE_FUSED_CAUSAL_PREFILL_ATTENTION`), whose C++ launchers hardcode the split-half rotation and take no flag, so a quantized Helium would otherwise be mis-rotated under either of those opt-in gates. Tensor parallelism is deliberately refused: the TP runtime builds its per-rank model by parsing `config.json` straight into the shared Llama config, which cannot carry the flag. `head_dim` comes from `hidden_size / num_attention_heads` as upstream does, and a config whose declared `head_dim` disagrees is rejected rather than one being silently preferred. Loading also rejects an embedding table or output head with fewer rows than `config.json`'s `vocab_size` claims, since an embedding gather wraps a negative index but does not range-check a positive one; a projection, bias or norm whose shape disagrees with the config in any layer, on both axes and including on the quantized path, where packing compresses the input axis only so the output row count is still the width the fused QKV concatenation and the attention reshape depend on and the input width is recovered from the scales the way MLX recovers it, `scales.shape(-1) * group_size`; quantization `biases` whose shape disagrees with the `scales` they are the zero points for; an attention block whose `q_proj`, `k_proj` and `v_proj` disagree on whether they are quantized, on whether they carry affine `biases`, or on whether they carry a dense `bias`, since the fused loader decides from `q_proj` alone and silently drops a whole bias set when one is missing; an odd or non-positive head width, because MLX enforces exactly that contract on `rope`'s `dims` argument by throwing and an MLX C++ exception crossing the cxx bridge is an uncatchable `std::terminate` at the first forward pass rather than a load error; a non-positive or non-finite `rope_theta`, which would make every rotated channel NaN; a non-finite, negative or zero `rms_norm_eps` (Helium's own `1e-08` is unusually small for this family and stays accepted); and a `quantization` block whose `bits` falls outside `1..=32` or whose `group_size` is non-positive. Stop tokens come from `config.json`'s `eos_token_id`, since Helium's `tokenizer_config.json` declares neither an EOS token nor a chat template. Validated checkpoint: `kyutai/helium-1-preview-2b` (mlx-community 4-bit, 559 tensors). Helium 1 preview is an English/French base model and ships no chat template, so the CLI renders the prompt verbatim.) Many of these families have checkpoint-specific config or weight-layout requirements. If a checkpoint fails detection or loading, inspect its diff --git a/src/distributed/tensor_parallel/inference.rs b/src/distributed/tensor_parallel/inference.rs index 36f60c050..a397b24a6 100644 --- a/src/distributed/tensor_parallel/inference.rs +++ b/src/distributed/tensor_parallel/inference.rs @@ -240,6 +240,18 @@ fn fallback_architecture(model_type: ModelType) -> &'static str { // V rather than between whole heads. This arm keeps the dispatch table // total. ModelType::GptNeoX => "gpt_neox", + // Helium is structurally a shardable dense Llama, but TP is refused for + // it on purpose and this arm only keeps the dispatch table total. The + // tensor-parallel runtime builds its per-rank model by parsing + // `config.json` straight into `llama3::ModelArgs`, and + // `llama3::ModelArgs::rope_traditional` is deliberately not + // deserializable, so a sharded Helium would silently rotate split-half + // pairs while the single-process path rotates interleaved pairs. The + // refusal comes from `validate_supported_runtime`, whose + // `runtime_kind_for` match has no arm for `ModelType::Helium` and so + // returns `None`. Enabling TP here means threading the flag through the + // rank-local config first. + ModelType::Helium => "helium", ModelType::StarCoder2 => "starcoder2", ModelType::Mellum => "mellum", ModelType::MiniCPM | ModelType::MiniCPMOVLM => "minicpm", diff --git a/src/distributed/tensor_parallel/llama_runtime_tests.rs b/src/distributed/tensor_parallel/llama_runtime_tests.rs index d2ec5d27e..bb9e22d7f 100644 --- a/src/distributed/tensor_parallel/llama_runtime_tests.rs +++ b/src/distributed/tensor_parallel/llama_runtime_tests.rs @@ -52,6 +52,9 @@ fn make_test_model_args() -> LlamaModelArgs { rope_scaling: None, quantization: None, tie_word_embeddings: false, + // Every Llama-family checkpoint rotates split-half; only the Helium + // loader sets this. + rope_traditional: false, } } @@ -913,6 +916,7 @@ fn local_llama_args_preserves_computed_head_dim_when_config_omits_it() { rope_scaling: None, quantization: None, tie_word_embeddings: true, + rope_traditional: false, }; let plan = generate_shard_plan("llama", 1, &ShardConfig::with_tp_size(2)).unwrap(); diff --git a/src/loaded_model.rs b/src/loaded_model.rs index 9c6c67f4f..c2735f92b 100644 --- a/src/loaded_model.rs +++ b/src/loaded_model.rs @@ -164,6 +164,7 @@ pub enum LoadedModel { Gpt2(models::Gpt2Model), GptBigCode(models::GptBigCodeModel), GptNeoX(models::GptNeoxModel), + Helium(models::HeliumModel), StarCoder2(models::StarCoder2Model), // Mixed full/sliding caches use a wrapper Mellum(models::MellumWrapper), @@ -313,6 +314,7 @@ macro_rules! delegate_language_model { LoadedModel::Gpt2(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::GptBigCode(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::GptNeoX(inner) => LanguageModel::$method(inner, $($arg),*), + LoadedModel::Helium(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::StarCoder2(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::Mellum(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::MiniCPM(inner) => LanguageModel::$method(inner, $($arg),*), diff --git a/src/model_metadata.rs b/src/model_metadata.rs index 40d625572..c97041fab 100644 --- a/src/model_metadata.rs +++ b/src/model_metadata.rs @@ -203,6 +203,7 @@ macro_rules! for_each_model_registration { Gpt2 => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::Gpt2Model::load, args: models::gpt2::ModelArgs, weight_builder: models::Gpt2Model::from_weights, wrap: LoadedModel::Gpt2 } }; GptBigCode => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::GptBigCodeModel::load, args: models::gpt_bigcode::ModelArgs, weight_builder: models::GptBigCodeModel::from_weights, wrap: LoadedModel::GptBigCode } }; GptNeoX => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::GptNeoxModel::load, args: models::gpt_neox::ModelArgs, weight_builder: models::GptNeoxModel::from_weights, wrap: LoadedModel::GptNeoX } }; + Helium => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::HeliumModel::load, args: models::helium::ModelArgs, weight_builder: models::HeliumModel::from_weights, wrap: LoadedModel::Helium } }; StarCoder2 => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::StarCoder2Model::load, args: models::starcoder2::StarCoder2Config, weight_builder: models::StarCoder2Model::from_weights, wrap: LoadedModel::StarCoder2 } }; Mellum => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::MellumModel::load, args: models::mellum::ModelArgs, weight_builder: models::MellumModel::from_weights, wrap: |m| LoadedModel::Mellum(models::MellumWrapper::new(m)) } }; MiniCPM => { kind: Text, directory: ConfigBacked, weight: Some(WeightLoadRoute::ConfigBacked), adapter: None, config_backed: { dir_loader: models::MiniCPMModel::load, args: models::minicpm::ModelArgs, weight_builder: models::MiniCPMModel::from_weights, wrap: LoadedModel::MiniCPM } }; diff --git a/src/models/detection.rs b/src/models/detection.rs index 3e205ea9a..43f6843bc 100644 --- a/src/models/detection.rs +++ b/src/models/detection.rs @@ -200,6 +200,7 @@ pub fn get_model_type(model_path: &Path) -> Result { "gpt2" => Ok(ModelType::Gpt2), "gpt_bigcode" => Ok(ModelType::GptBigCode), "gpt_neox" => Ok(ModelType::GptNeoX), + "helium" => Ok(ModelType::Helium), "starcoder2" => Ok(ModelType::StarCoder2), "mellum" => Ok(ModelType::Mellum), "minicpm" => Ok(ModelType::MiniCPM), diff --git a/src/models/detection_tests.rs b/src/models/detection_tests.rs index 063aa8355..7985867c8 100644 --- a/src/models/detection_tests.rs +++ b/src/models/detection_tests.rs @@ -182,6 +182,42 @@ fn gpt_neox_model_type_is_detected() { fs::remove_dir_all(model_dir).unwrap(); } +#[test] +fn helium_model_type_is_detected() { + // Helium's config is field-for-field a Llama config apart from `model_type` + // and the smaller `rms_norm_eps`, so detection must key off `model_type` + // alone and must not fall through to the Llama arm. The two decode + // differently: Helium rotates interleaved RoPE pairs, Llama split-half ones. + let model_dir = temp_path("helium_text"); + fs::create_dir_all(&model_dir).unwrap(); + fs::write( + model_dir.join("config.json"), + r#"{ + "model_type": "helium", + "architectures": ["HeliumForCausalLM"], + "hidden_size": 2560, + "num_attention_heads": 20, + "num_key_value_heads": 20, + "num_hidden_layers": 24, + "intermediate_size": 7040, + "head_dim": 128, + "max_position_embeddings": 4096, + "rms_norm_eps": 1e-08, + "rope_theta": 100000.0, + "attention_bias": false, + "mlp_bias": false, + "tie_word_embeddings": false, + "vocab_size": 48000 + }"#, + ) + .unwrap(); + + let detected = super::detection::get_model_type(&model_dir).unwrap(); + assert_eq!(detected, ModelType::Helium); + + fs::remove_dir_all(model_dir).unwrap(); +} + #[test] fn fastvlm_model_type_is_detected() { for model_type in ["fastvlm", "llava_qwen2"] { diff --git a/src/models/helium.rs b/src/models/helium.rs new file mode 100644 index 000000000..20e9442c5 --- /dev/null +++ b/src/models/helium.rs @@ -0,0 +1,956 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Kyutai Helium (`helium`). +//! +//! Ported from mlx-lm's +//! [`helium.py`](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/models/helium.py). +//! +//! Helium is a dense Llama-shaped decoder: RMSNorm before attention and before +//! the MLP, grouped-query attention, a SwiGLU MLP over `gate_proj` / `up_proj` / +//! `down_proj`, no QK-norm, no MoE, no sliding window, and Llama's own weight +//! key names. The decoder block, attention and MLP therefore come from +//! [`crate::models::llama3`] unchanged rather than being copied. +//! +//! # The one architectural difference +//! +//! Upstream builds `nn.RoPE(head_dim, traditional=True, base=rope_theta)`. +//! Every other Llama-family model in this tree rotates split-half pairs +//! `(i, i + dims/2)`; Helium rotates interleaved pairs `(2i, 2i+1)`. The two +//! produce identically shaped tensors from identical weights, so running the +//! wrong one is a silent quality regression that no shape assertion can catch. +//! +//! [`ModelArgs::to_llama3_args`] therefore sets +//! [`crate::models::llama3::ModelArgs::rope_traditional`], which +//! [`crate::models::llama3::Attention::forward`] honors on the graph path and +//! which additionally disables the two fused RoPE fast paths, because their C++ +//! launchers hardcode `traditional = false` and take no flag. See that method's +//! doc comment for the full routing argument. +//! +//! # Untrusted config +//! +//! `config.json` arrives from a third-party HuggingFace repo in the common +//! `mlxcel generate -m /` flow, so [`ModelArgs::validate`] rejects +//! every scalar that could size an allocation, divide, truncate through an +//! `as i32` cast, or violate an undocumented precondition of an MLX C++ entry +//! point, and [`validate_weights`] rejects every tensor whose real shape +//! disagrees with the config, on both axes and on both the float and the +//! quantized path, before either can reach a kernel. An MLX C++ exception +//! crossing the cxx bridge is an uncatchable `std::terminate`, not a Rust error, +//! so a check that happens at the first forward pass is not a check. + +use mlxcel_core::generate::{DecodeBatchContext, LanguageModel}; +use mlxcel_core::layers::KVCache; +use mlxcel_core::weights::WeightMap; +use mlxcel_core::{MlxArray, UniquePtr}; +use serde::Deserialize; +use std::path::Path; + +use super::llama3::Llama3Model; + +// Configuration. + +/// Helium `config.json`. +/// +/// Field-for-field the upstream `ModelArgs`, plus `eos_token_id` (Helium ships +/// no chat template and no tokenizer-level EOS, so `config.json` is the only +/// source of the stop token) and the standard `quantization` block that an +/// mlx-community conversion adds. +#[derive(Debug, Clone, Deserialize)] +pub struct ModelArgs { + #[serde(default = "default_model_type")] + pub model_type: String, + + pub hidden_size: usize, + pub num_hidden_layers: usize, + pub intermediate_size: usize, + pub num_attention_heads: usize, + + /// Absent means multi-head attention. Upstream asserts this is set and then + /// uses it directly; the `Option` here only tolerates a config that omits + /// it, which resolves to `num_attention_heads`. + #[serde(default)] + pub num_key_value_heads: Option, + + /// Present in every published Helium `config.json`, but **not** what + /// upstream's attention uses. See [`ModelArgs::head_dim`]. + #[serde(default)] + pub head_dim: Option, + + #[serde(default = "default_rms_norm_eps")] + pub rms_norm_eps: f32, + + pub vocab_size: usize, + + #[serde(default)] + pub attention_bias: bool, + + #[serde(default)] + pub mlp_bias: bool, + + #[serde(default = "default_rope_theta")] + pub rope_theta: f32, + + #[serde(default = "default_max_position_embeddings")] + pub max_position_embeddings: usize, + + #[serde(default)] + pub tie_word_embeddings: bool, + + #[serde(default)] + pub eos_token_id: Option, + + #[serde(default)] + pub quantization: Option, +} + +/// A `config.json` token-id field, which may be a single int or a list of ints. +/// +/// Same shape as the per-family enums in [`crate::models::gpt2`] and +/// [`crate::models::gpt_neox`]; serde fails the whole config when one field +/// does not match its declared type, so the list form has to be accepted even +/// though every published Helium checkpoint writes a single int. +#[derive(Debug, Clone, Deserialize)] +#[serde(untagged)] +pub enum TokenIdField { + Single(i32), + Multiple(Vec), +} + +impl TokenIdField { + fn ids(&self) -> Vec { + match self { + Self::Single(id) => vec![*id], + Self::Multiple(ids) => ids.clone(), + } + } +} + +#[derive(Debug, Clone, Deserialize)] +pub struct Quantization { + pub group_size: i32, + pub bits: i32, +} + +fn default_model_type() -> String { + "helium".to_string() +} +fn default_rms_norm_eps() -> f32 { + 1e-8 +} +fn default_rope_theta() -> f32 { + 100_000.0 +} +fn default_max_position_embeddings() -> usize { + 4096 +} + +/// Upper bounds on the architecture scalars a Helium `config.json` may declare. +/// +/// `config.json` is untrusted input: `mlxcel generate -m /` downloads +/// a third-party HuggingFace repo and loads it in the same command, and the +/// download layer validates repo ids, filenames and transport but never parses +/// the file, so these fields arrive exactly as the checkpoint author wrote them. +/// +/// Each ceiling sits orders of magnitude above Helium 1 (24 layers, +/// `hidden_size` 2560, `intermediate_size` 7040, `vocab_size` 48000). They exist +/// so `num_hidden_layers` cannot size the `Vec::with_capacity` in +/// [`Llama3Model::from_weights`], and so the `as i32` casts these values feed +/// (`head_dim`, `num_heads * head_dim`, `num_kv_heads * head_dim`) stay inside +/// `i32` instead of truncating to a negative number. +const MAX_NUM_HIDDEN_LAYERS: usize = 1024; +/// See [`MAX_NUM_HIDDEN_LAYERS`]. +const MAX_HIDDEN_SIZE: usize = 65_536; +/// See [`MAX_NUM_HIDDEN_LAYERS`]. +const MAX_INTERMEDIATE_SIZE: usize = 1 << 22; +/// See [`MAX_NUM_HIDDEN_LAYERS`]. +const MAX_MAX_POSITION_EMBEDDINGS: usize = 1 << 22; +/// See [`MAX_NUM_HIDDEN_LAYERS`]. +const MAX_VOCAB_SIZE: usize = 1 << 24; + +impl ModelArgs { + /// Head width, and the number of channels RoPE rotates. + /// + /// Upstream `HeliumAttention` computes `args.hidden_size // n_heads` and + /// never reads the `head_dim` field, even though the dataclass requires it, + /// so `hidden_size / num_attention_heads` is authoritative here too. On + /// `kyutai/helium-1-preview-2b` the two agree (2560 / 20 = 128 = the + /// declared `head_dim`); [`ModelArgs::validate`] rejects a config where they + /// disagree rather than silently preferring one, because a checkpoint whose + /// projections were built for the other value would produce a reshape + /// mismatch inside MLX at the first forward pass. + /// + /// Only valid after [`ModelArgs::validate`], which rejects + /// `num_attention_heads == 0`. + pub fn head_dim(&self) -> usize { + self.hidden_size / self.num_attention_heads + } + + pub fn num_kv_heads(&self) -> usize { + self.num_key_value_heads.unwrap_or(self.num_attention_heads) + } + + pub fn group_size(&self) -> i32 { + self.quantization + .as_ref() + .map(|q| q.group_size) + .unwrap_or(64) + } + + pub fn bits(&self) -> i32 { + self.quantization.as_ref().map(|q| q.bits).unwrap_or(4) + } + + /// Stop tokens. + /// + /// Helium's `tokenizer_config.json` declares no `eos_token` and no chat + /// template, so `config.json`'s `eos_token_id` (2, ``) is the only + /// source. Delegating to [`Llama3Model`] instead would return Llama 3's + /// hardcoded `[128001, 128009]`, which are outside Helium's 48000-entry + /// vocabulary and would therefore never match, so generation would only ever + /// stop at the token limit. + pub fn eos_token_ids(&self) -> Vec { + self.eos_token_id + .as_ref() + .map(TokenIdField::ids) + .unwrap_or_default() + } + + /// Reject a `config.json` that cannot describe a real Helium, before any of + /// its fields sizes an allocation, divides, or reaches an MLX kernel. + pub fn validate(&self) -> Result<(), String> { + // The zero checks come first. `0.is_multiple_of(0)` is true, so a config + // with `hidden_size == num_attention_heads == 0` would pass the + // divisibility check below and then divide by zero in `head_dim()`. + if self.num_attention_heads == 0 { + return Err("Helium num_attention_heads must be at least 1".to_string()); + } + if self.hidden_size == 0 || self.hidden_size > MAX_HIDDEN_SIZE { + return Err(format!( + "Helium hidden_size ({}) must be between 1 and {MAX_HIDDEN_SIZE}", + self.hidden_size + )); + } + if self.num_hidden_layers == 0 || self.num_hidden_layers > MAX_NUM_HIDDEN_LAYERS { + return Err(format!( + "Helium num_hidden_layers ({}) must be between 1 and {MAX_NUM_HIDDEN_LAYERS}", + self.num_hidden_layers + )); + } + if self.intermediate_size == 0 || self.intermediate_size > MAX_INTERMEDIATE_SIZE { + return Err(format!( + "Helium intermediate_size ({}) must be between 1 and {MAX_INTERMEDIATE_SIZE}", + self.intermediate_size + )); + } + if self.vocab_size == 0 || self.vocab_size > MAX_VOCAB_SIZE { + return Err(format!( + "Helium vocab_size ({}) must be between 1 and {MAX_VOCAB_SIZE}", + self.vocab_size + )); + } + if self.max_position_embeddings == 0 + || self.max_position_embeddings > MAX_MAX_POSITION_EMBEDDINGS + { + return Err(format!( + "Helium max_position_embeddings ({}) must be between 1 and \ + {MAX_MAX_POSITION_EMBEDDINGS}", + self.max_position_embeddings + )); + } + if !self.hidden_size.is_multiple_of(self.num_attention_heads) { + return Err(format!( + "Helium hidden_size ({}) must be divisible by num_attention_heads ({})", + self.hidden_size, self.num_attention_heads + )); + } + + let num_kv_heads = self.num_kv_heads(); + if num_kv_heads == 0 { + return Err("Helium num_key_value_heads must be at least 1".to_string()); + } + if !self.num_attention_heads.is_multiple_of(num_kv_heads) { + return Err(format!( + "Helium num_attention_heads ({}) must be divisible by num_key_value_heads \ + ({num_kv_heads}) for grouped-query attention", + self.num_attention_heads + )); + } + + if let Some(declared) = self.head_dim { + let derived = self.head_dim(); + if declared != derived { + return Err(format!( + "Helium config declares head_dim {declared} but hidden_size ({}) / \ + num_attention_heads ({}) is {derived}. Upstream Helium builds its \ + projections from hidden_size // num_attention_heads and never reads the \ + head_dim field, so the two must agree; a checkpoint built for either value \ + would mis-shape the other's reshape.", + self.hidden_size, self.num_attention_heads + )); + } + } + + self.validate_rope()?; + self.validate_norm_eps()?; + self.validate_quantization()?; + Ok(()) + } + + /// Reject RoPE parameters that MLX would throw on, or that would silently + /// NaN every rotated channel. + /// + /// `mlx::core::fast::rope` requires `dims` to be positive, even, and no + /// larger than the input's last axis. Helium rotates the full head, so + /// `dims == head_dim` and the third condition holds by construction, but the + /// first two are config-controlled through `hidden_size` and + /// `num_attention_heads`. MLX enforces them by throwing + /// `std::invalid_argument`, and `fast_rope` crosses the cxx bridge as + /// `UniquePtr` rather than a `Result`, so that throw is an + /// uncatchable `std::terminate` at the **first forward pass**, long after + /// the checkpoint appeared to load cleanly. Both are therefore rejected + /// here, at load. + /// + /// `rope_theta` is the base of the frequency exponentiation. A zero, + /// negative or non-finite base makes every rotated channel NaN without + /// anything throwing at all, which is harder to diagnose than a crash. + fn validate_rope(&self) -> Result<(), String> { + let head_dim = self.head_dim(); + if head_dim < 2 || !head_dim.is_multiple_of(2) { + return Err(format!( + "Helium head_dim ({head_dim}, from hidden_size {} / num_attention_heads {}) must \ + be even and at least 2. MLX rotates channel pairs and rejects an odd or \ + non-positive rope dimension by throwing, which crosses the cxx bridge as an \ + uncatchable abort at the first forward pass rather than a load error.", + self.hidden_size, self.num_attention_heads + )); + } + if !self.rope_theta.is_finite() || self.rope_theta <= 0.0 { + return Err(format!( + "Helium rope_theta ({}) must be a finite positive number; RoPE exponentiates it \ + per channel, so a zero, negative or non-finite base makes every rotated channel \ + NaN and that NaN reaches the logits without anything throwing", + self.rope_theta + )); + } + Ok(()) + } + + /// Reject an `rms_norm_eps` that would NaN every hidden state. + /// + /// `fast::rms_norm` computes `x * weight * rsqrt(mean(x^2) + eps)` and never + /// inspects `eps`, so a non-finite, negative or zero value produces NaN + /// hidden states with no error at all: the checkpoint loads, generation + /// runs, and the output is uniform garbage. + /// + /// The bound is a range, not an allowlist. Helium 1's `rms_norm_eps` is + /// `1e-08`, which is unusually small for this family (Llama uses `1e-05`), + /// and must stay accepted. + fn validate_norm_eps(&self) -> Result<(), String> { + if !self.rms_norm_eps.is_finite() || self.rms_norm_eps <= 0.0 { + return Err(format!( + "Helium rms_norm_eps ({}) must be a finite positive number; it is added to the \ + mean square under an rsqrt, so a non-finite, negative or zero value makes every \ + normalized hidden state NaN and that NaN reaches the logits without anything \ + throwing", + self.rms_norm_eps + )); + } + Ok(()) + } + + /// Reject a `quantization` block that would abort an MLX kernel. + /// + /// `group_size` and `bits` reach `quantized_matmul` and `dequantize` + /// unreconciled when the tensor shapes carry too little information to + /// re-derive them. MLX's own `packed_in * 32 / bits` check divides by zero at + /// `bits == 0` and collapses to zero above 32, throwing the same uncatchable + /// `std::terminate` the rope guard exists to prevent. + /// + /// This is a range check rather than an allowlist on purpose: mlxcel + /// deliberately re-derives an effective bit width from the tensor shapes when + /// the declared one disagrees, and an allowlist would reject the + /// mixed-precision exports that behavior serves. + fn validate_quantization(&self) -> Result<(), String> { + let Some(quantization) = self.quantization.as_ref() else { + return Ok(()); + }; + if quantization.bits < 1 || quantization.bits > 32 { + return Err(format!( + "Helium quantization.bits ({}) must be between 1 and 32; MLX derives the unpacked \ + width as packed_in * 32 / bits, which divides by zero at 0 and collapses to zero \ + above 32, and the resulting C++ throw crosses the cxx bridge as an uncatchable \ + abort", + quantization.bits + )); + } + if quantization.group_size < 1 { + return Err(format!( + "Helium quantization.group_size ({}) must be positive; it divides the input axis \ + inside every quantized kernel", + quantization.group_size + )); + } + Ok(()) + } + + /// Build the `llama3::ModelArgs` that drives the shared dense decoder. + /// + /// The struct literal is deliberate: every field of the shared config is + /// named here, so a future field added to [`crate::models::llama3::ModelArgs`] + /// fails this file to compile rather than silently defaulting for Helium. + /// + /// `rope_traditional` is the whole reason this conversion exists. + pub fn to_llama3_args(&self) -> super::llama3::ModelArgs { + super::llama3::ModelArgs { + model_type: self.model_type.clone(), + hidden_size: self.hidden_size, + num_hidden_layers: self.num_hidden_layers, + intermediate_size: self.intermediate_size, + num_attention_heads: self.num_attention_heads, + rms_norm_eps: self.rms_norm_eps, + vocab_size: self.vocab_size, + // Pass the derived width explicitly rather than the declared field: + // `validate` has already rejected a config where they disagree, and + // this keeps `hidden_size / num_attention_heads` the single source. + head_dim: Some(self.head_dim()), + num_key_value_heads: Some(self.num_kv_heads()), + attention_bias: self.attention_bias, + mlp_bias: self.mlp_bias, + rope_theta: self.rope_theta, + // Helium has no RoPE scaling of any kind upstream. + rope_scaling: None, + quantization: self + .quantization + .as_ref() + .map(|q| super::llama3::Quantization { + group_size: q.group_size, + bits: q.bits, + }), + tie_word_embeddings: self.tie_word_embeddings, + rope_traditional: true, + } + } +} + +// Weight-shape validation. + +/// Whether an MLX dimension equals an expected extent. +/// +/// `array_shape` returns `i32`; a negative or oversized value can never match a +/// real extent, so the conversion failing is itself a mismatch. Mirrors the same +/// check in the GPT-2 lineage loaders. +fn dim_eq(dim: i32, expected: usize) -> bool { + usize::try_from(dim).is_ok_and(|d| d == expected) +} + +/// What [`validate_projection`] observed about one linear projection, so callers +/// can cross-check projections that get concatenated together. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct ProjectionShape { + quantized: bool, + /// Second axis of `.weight`. Equals `in_features` on the float path and the + /// packed width on the quantized path. + cols: i32, + /// Second axis of `.scales`, or `None` on the float path. + scale_cols: Option, + has_quant_biases: bool, + /// Whether a dense `.bias` is present, cross-checked across `q_proj` / + /// `k_proj` / `v_proj` for the same reason `has_quant_biases` is. + has_dense_bias: bool, +} + +/// Check a quantized projection's `.scales` (and `.biases`) against the input +/// width `config.json` claims. +/// +/// A quantized `[out_features, in_features]` matrix packs along the input axis +/// only, so the row check in [`validate_projection`] still applies but says +/// nothing about the input width. MLX reconstructs that width as +/// `scales.shape(-1) * group_size` and `extract_quantized_matmul_dims` throws +/// `std::invalid_argument` when it disagrees with the activation's last axis; +/// `validate_quantized_input` throws again when `.biases` and `.scales` differ +/// in shape. `quantized_matmul` crosses the cxx bridge as +/// `UniquePtr` rather than a `Result`, so either throw is an +/// uncatchable `std::terminate` at the **first forward pass**, long after the +/// checkpoint appeared to load. The q/k/v cross-check in [`validate_weights`] +/// cannot substitute: it only proves the three projections agree with each +/// other, not that any of them agrees with `hidden_size`. +/// +/// The declared `group_size` is the right one to check against, not a +/// shape-derived one. The affine loader trusts the declared group size and +/// re-derives `bits` from the shapes, and +/// `FusedQKVLinear::from_weights_separate` does so unconditionally, so the +/// declared value is what reaches MLX. That makes this stricter than the loader +/// for one layout the loader would repair (a declared `group_size` of 16 at +/// 4 bits, the NVFP4-fallback repack, where `reconcile_quantization_layout` +/// re-derives the group size instead); no Helium checkpoint is packed that way, +/// and rejecting it names the mismatch instead of aborting on it. +/// +/// Returns the scales column count, or `None` when the projection is not +/// quantized. +fn validate_quantized_scales( + weights: &WeightMap, + prefix: &str, + out_features: usize, + in_features: usize, + group_size: i32, +) -> Result, String> { + let Some(scales) = weights.get(&format!("{prefix}.scales")) else { + return Ok(None); + }; + let scale_shape = mlxcel_core::array_shape(scales); + if scale_shape.len() != 2 || !dim_eq(scale_shape[0], out_features) { + return Err(format!( + "unexpected {prefix}.scales shape {scale_shape:?}: expected {out_features} rows to \ + match {prefix}.weight" + )); + } + + let groups = usize::try_from(scale_shape[1]).unwrap_or(0); + let group_size = usize::try_from(group_size).unwrap_or(0); + let described = groups.checked_mul(group_size); + if described != Some(in_features) { + return Err(format!( + "unexpected {prefix}.scales shape {scale_shape:?}: {groups} quantization groups at \ + group_size {group_size} describe an input width of {}, but the config says \ + {in_features}. MLX reconstructs a quantized matrix's input width as \ + scales.shape(-1) * group_size and throws when it disagrees with the activation, and \ + that throw crosses the cxx bridge as an uncatchable abort at the first forward pass \ + rather than a load error. Packing compresses the input axis only, so the row count \ + is still correct and cannot catch this.", + described + .map(|d| d.to_string()) + .unwrap_or_else(|| "an overflowing number of".to_string()) + )); + } + + if let Some(biases) = weights.get(&format!("{prefix}.biases")) { + let bias_shape = mlxcel_core::array_shape(biases); + if bias_shape != scale_shape { + return Err(format!( + "unexpected {prefix}.biases shape {bias_shape:?}: the affine zero points must \ + have the same shape as {prefix}.scales ({scale_shape:?}). MLX rejects a mismatch \ + by throwing, which crosses the cxx bridge as an uncatchable abort at the first \ + forward pass." + )); + } + } + + Ok(Some(scale_shape[1])) +} + +/// Check one `[out_features, in_features]` projection against the config. +/// +/// A quantized weight is packed along the input axis only, so its row count is +/// still `out_features` and is checked on the same path as a float weight; +/// skipping the row check on the quantized path is exactly the carve-out that +/// let K and V be sliced from arbitrary interior channels in an earlier port in +/// this chain. The input axis is checked too, through the scales rather than +/// through the packed `.weight`, because the packed width alone does not fix a +/// width without a bit count: see [`validate_quantized_scales`]. Only the +/// declared-versus-derived bit width itself is left to `UnifiedLinear`, which +/// reconciles it from these two shapes and returns a load error, not an abort, +/// when they cannot agree. +fn validate_projection( + weights: &WeightMap, + prefix: &str, + out_features: usize, + in_features: usize, + group_size: i32, +) -> Result { + let weight_name = format!("{prefix}.weight"); + let weight = weights + .get(&weight_name) + .ok_or_else(|| format!("Weight not found: {weight_name}"))?; + let shape = mlxcel_core::array_shape(weight); + if shape.len() != 2 { + return Err(format!( + "unexpected {weight_name} shape {shape:?}: expected a 2-D [out, in] projection" + )); + } + let quantized = weights.contains_key(&format!("{prefix}.scales")); + + if !dim_eq(shape[0], out_features) { + let looks_transposed = !quantized && dim_eq(shape[0], in_features); + let hint = if looks_transposed { + " That is the [in, out] orientation; Llama-shaped checkpoints store [out, in] and \ + must not be transposed." + } else { + "" + }; + return Err(format!( + "unexpected {weight_name} shape {shape:?}: expected {out_features} rows.{hint}" + )); + } + if !quantized && !dim_eq(shape[1], in_features) { + return Err(format!( + "unexpected {weight_name} shape {shape:?}: expected [{out_features}, {in_features}]" + )); + } + + let scale_cols = + validate_quantized_scales(weights, prefix, out_features, in_features, group_size)?; + + let dense_bias = weights.get(&format!("{prefix}.bias")); + if let Some(bias) = dense_bias { + let bias_shape = mlxcel_core::array_shape(bias); + if bias_shape.len() != 1 || !dim_eq(bias_shape[0], out_features) { + return Err(format!( + "unexpected {prefix}.bias shape {bias_shape:?}: expected [{out_features}]" + )); + } + } + + Ok(ProjectionShape { + quantized, + cols: shape[1], + scale_cols, + has_quant_biases: weights.contains_key(&format!("{prefix}.biases")), + has_dense_bias: dense_bias.is_some(), + }) +} + +/// Check a 1-D RMSNorm weight. +fn validate_norm(weights: &WeightMap, key: &str, dim: usize) -> Result<(), String> { + let weight = weights + .get(key) + .ok_or_else(|| format!("Weight not found: {key}"))?; + let shape = mlxcel_core::array_shape(weight); + if shape.len() != 1 || !dim_eq(shape[0], dim) { + return Err(format!( + "unexpected {key} shape {shape:?}: expected [{dim}]" + )); + } + Ok(()) +} + +/// Check a 2-D embedding-shaped table and return its row count. +/// +/// The row count is the only place a config that overstates the vocabulary can +/// be caught: MLX's gather adds the axis size to a negative index but performs +/// no range check on a positive one, so an id past the last row reads whatever +/// follows the table in the buffer rather than faulting, and the result reaches +/// the logits. A table with *more* rows than the config claims is accepted, +/// which is how a vocabulary-padded head is stored. +fn validate_table( + weights: &WeightMap, + key: &str, + claimed_rows: usize, + expected_cols: usize, + group_size: i32, +) -> Result<(), String> { + let weight_name = format!("{key}.weight"); + let weight = weights + .get(&weight_name) + .ok_or_else(|| format!("Weight not found: {weight_name}"))?; + let shape = mlxcel_core::array_shape(weight); + if shape.len() != 2 { + return Err(format!( + "unexpected {weight_name} shape {shape:?}: expected a 2-D [rows, hidden_size] table" + )); + } + let rows = usize::try_from(shape[0]).unwrap_or(0); + if rows < claimed_rows { + return Err(format!( + "config vocab_size ({claimed_rows}) exceeds the {rows} rows present in {weight_name}. \ + Lookups into this table are bounded by vocab_size, so a config that overstates it \ + indexes past the end of the table, and the gather behind an embedding lookup does \ + not range-check a positive index." + )); + } + if weights.contains_key(&format!("{key}.scales")) { + // Quantized: the row count above is still the vocabulary bound, but the + // model width now lives in the scales, so it is checked there. Leaving + // it unchecked lets a table packed for a different width reach either + // `quantized_matmul` (an output head) or the first RMSNorm (an embedding + // table), both of which throw inside MLX and abort the process. + validate_quantized_scales(weights, key, rows, expected_cols, group_size)?; + } else if !dim_eq(shape[1], expected_cols) { + return Err(format!( + "{weight_name} is {shape:?} but hidden_size is {expected_cols}; an embedding table \ + must be the model width" + )); + } + Ok(()) +} + +/// Reject a checkpoint whose real tensor shapes disagree with `config.json`, +/// before any of them reaches MLX. +/// +/// This has to run **before** [`Llama3Model::from_weights`], not after. +/// `FusedQKVLinear::from_weights_separate` concatenates `q_proj`, `k_proj` and +/// `v_proj` (and their `scales` and `biases`) along axis 0 with no shape check +/// of its own, and `Attention::forward` then reshapes the result using +/// config-derived head counts. Both a mismatched concatenate and a mismatched +/// reshape throw inside MLX, and an MLX C++ exception crossing the cxx bridge is +/// an uncatchable `std::terminate` rather than a load error. +pub fn validate_weights(weights: &WeightMap, args: &ModelArgs) -> Result<(), String> { + let hidden = args.hidden_size; + let head_dim = args.head_dim(); + let q_out = args.num_attention_heads * head_dim; + let kv_out = args.num_kv_heads() * head_dim; + let group_size = args.group_size(); + + validate_table( + weights, + "model.embed_tokens", + args.vocab_size, + hidden, + group_size, + )?; + validate_norm(weights, "model.norm.weight", hidden)?; + if !args.tie_word_embeddings { + validate_table(weights, "lm_head", args.vocab_size, hidden, group_size)?; + } + + for layer in 0..args.num_hidden_layers { + let attn = format!("model.layers.{layer}.self_attn"); + let q = validate_projection( + weights, + &format!("{attn}.q_proj"), + q_out, + hidden, + group_size, + )?; + let k = validate_projection( + weights, + &format!("{attn}.k_proj"), + kv_out, + hidden, + group_size, + )?; + let v = validate_projection( + weights, + &format!("{attn}.v_proj"), + kv_out, + hidden, + group_size, + )?; + validate_projection( + weights, + &format!("{attn}.o_proj"), + hidden, + q_out, + group_size, + )?; + + // The fused QKV loader decides "is this quantized?" from `q_proj.scales` + // alone and then concatenates all three, and it keeps the affine + // `biases` only when all three carry one. A checkpoint where q, k and v + // disagree therefore either aborts inside `concatenate` or, worse for + // the `biases` case, silently dequantizes K and V without their zero + // points. Neither is recoverable downstream, so both are rejected here. + if q.quantized != k.quantized || q.quantized != v.quantized { + return Err(format!( + "{attn}: q_proj, k_proj and v_proj must all be quantized or all be float; they \ + are concatenated into one fused projection and the loader decides which path to \ + take from q_proj alone" + )); + } + if q.cols != k.cols + || q.cols != v.cols + || q.scale_cols != k.scale_cols + || q.scale_cols != v.scale_cols + { + return Err(format!( + "{attn}: q_proj, k_proj and v_proj must share an input width (got {}, {}, {}); \ + they are concatenated along the output axis", + q.cols, k.cols, v.cols + )); + } + if q.has_quant_biases != k.has_quant_biases || q.has_quant_biases != v.has_quant_biases { + return Err(format!( + "{attn}: q_proj, k_proj and v_proj must either all carry quantization biases or \ + none of them; the fused loader drops the whole set when one is missing, which \ + dequantizes the others without their zero points" + )); + } + // The dense `.bias` set has the same all-or-nothing rule as the affine + // `.biases` set above, and the same silent failure: the fused loader + // concatenates q/k/v biases only when all three are present and drops + // the whole set otherwise, so a checkpoint carrying a bias on some of + // them loads without complaint and runs those projections unbiased. + if q.has_dense_bias != k.has_dense_bias || q.has_dense_bias != v.has_dense_bias { + return Err(format!( + "{attn}: q_proj, k_proj and v_proj must either all carry a bias or none of them; \ + the fused loader concatenates the three biases into one and silently drops the \ + whole set when any is missing, which runs the projections that had one without \ + it" + )); + } + + let mlp = format!("model.layers.{layer}.mlp"); + validate_projection( + weights, + &format!("{mlp}.gate_proj"), + args.intermediate_size, + hidden, + group_size, + )?; + validate_projection( + weights, + &format!("{mlp}.up_proj"), + args.intermediate_size, + hidden, + group_size, + )?; + validate_projection( + weights, + &format!("{mlp}.down_proj"), + hidden, + args.intermediate_size, + group_size, + )?; + + validate_norm( + weights, + &format!("model.layers.{layer}.input_layernorm.weight"), + hidden, + )?; + validate_norm( + weights, + &format!("model.layers.{layer}.post_attention_layernorm.weight"), + hidden, + )?; + } + + Ok(()) +} + +// Model. + +/// Kyutai Helium. +/// +/// A [`Llama3Model`] built with traditional RoPE, plus Helium's own stop tokens. +/// Everything else, including batched decode, paged decode and the prompt cache, +/// is the shared dense path unchanged. +pub struct HeliumModel { + inner: Llama3Model, + eos_token_ids: Vec, +} + +impl HeliumModel { + /// Load from a checkpoint directory. + pub fn load>(model_dir: P) -> Result<(Self, ModelArgs), String> { + let model_dir = model_dir.as_ref(); + + let config_path = model_dir.join("config.json"); + let config_str = std::fs::read_to_string(&config_path) + .map_err(|e| format!("Failed to read config.json: {e}"))?; + let args: ModelArgs = serde_json::from_str(&config_str) + .map_err(|e| format!("Failed to parse config.json: {e}"))?; + args.validate()?; + + let weights = crate::models::load_text_weights(model_dir, None)?; + let model = Self::from_weights(&weights, &args)?; + + Ok((model, args)) + } + + pub fn from_weights(weights: &WeightMap, args: &ModelArgs) -> Result { + // `config.json` is untrusted, and `from_weights` is also reachable + // directly from the registration table, so revalidate here rather than + // trusting that `load` ran first. + args.validate()?; + validate_weights(weights, args)?; + + let inner = Llama3Model::from_weights(weights, &args.to_llama3_args())?; + + Ok(Self { + inner, + eos_token_ids: args.eos_token_ids(), + }) + } + + pub fn inner(&self) -> &Llama3Model { + &self.inner + } +} + +// LanguageModel trait implementation. + +impl LanguageModel for HeliumModel { + fn forward( + &self, + input_ids: &MlxArray, + caches: &mut [KVCache], + mask: Option<&MlxArray>, + ) -> UniquePtr { + Llama3Model::forward(&self.inner, input_ids, caches, mask) + } + + fn forward_with_embeddings( + &self, + input_ids: &MlxArray, + input_embeddings: Option<&MlxArray>, + caches: &mut [KVCache], + mask: Option<&MlxArray>, + ) -> UniquePtr { + self.inner + .forward_with_embeddings_impl(input_ids, input_embeddings, caches, mask) + } + + fn embed_tokens(&self, input_ids: &MlxArray) -> Option> { + Some(self.inner.get_embed_tokens(input_ids)) + } + + fn make_caches(&self) -> Vec { + Llama3Model::make_caches(&self.inner) + } + + fn num_layers(&self) -> usize { + self.inner.layers.len() + } + + fn eos_token_ids(&self) -> Vec { + self.eos_token_ids.clone() + } + + fn forward_batched( + &self, + input_ids: &MlxArray, + batch_caches: &mut [&mut [KVCache]], + mask: Option<&MlxArray>, + ) -> UniquePtr { + self.inner + .forward_batched_impl(input_ids, batch_caches, mask, None) + } + + fn forward_batched_with_context( + &self, + input_ids: &MlxArray, + batch_caches: &mut [&mut [KVCache]], + mask: Option<&MlxArray>, + context: Option<&DecodeBatchContext>, + ) -> UniquePtr { + self.inner + .forward_batched_impl(input_ids, batch_caches, mask, context) + } + + fn supports_batched_prefill(&self) -> bool { + true + } + + fn supports_maskless_padded_prefill(&self) -> bool { + true + } + + fn supports_paged_decode_backend(&self) -> bool { + true + } +} + +#[cfg(test)] +#[path = "helium_tests.rs"] +mod tests; diff --git a/src/models/helium_tests.rs b/src/models/helium_tests.rs new file mode 100644 index 000000000..6c5ae6497 --- /dev/null +++ b/src/models/helium_tests.rs @@ -0,0 +1,863 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Unit tests for the Helium config surface, its untrusted-config guards, its +//! weight-shape contract, and above all its RoPE convention. +//! +//! Everything here is checkpoint-free: the config tests parse the real +//! `kyutai/helium-1-preview-2b` field set, and the model tests build tiny +//! synthetic weight maps whose key names and shapes mirror the checkpoint. +//! +//! The tests that matter most are the RoPE ones. Traditional and split-half RoPE +//! produce identically shaped tensors from identical weights, so no shape +//! assertion, no cache assertion and no logits-shape assertion can tell them +//! apart; only the values can. Dropping the flag on any one of the three routes +//! that can rotate Q and K would leave the model loading, generating, and +//! quietly wrong. + +use super::{HeliumModel, ModelArgs, validate_weights}; +use crate::models::llama3::{Attention, Llama3Model}; +use mlxcel_core::generate::LanguageModel; +use mlxcel_core::layers::{FusedQKVLinear, KVCache}; +use mlxcel_core::weights::WeightMap; +use mlxcel_core::{MlxArray, UniquePtr}; + +// Config surface. + +/// The `mlx-community/helium-1-preview-2b-4bit` config, field-for-field. +const HELIUM_1_PREVIEW_2B_CONFIG: &str = r#"{ + "architectures": ["HeliumForCausalLM"], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 1, + "eos_token_id": 2, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 2560, + "initializer_range": 0.02, + "intermediate_size": 7040, + "max_position_embeddings": 4096, + "mlp_bias": false, + "model_type": "helium", + "num_attention_heads": 20, + "num_hidden_layers": 24, + "num_key_value_heads": 20, + "pretraining_tp": 1, + "quantization": { "group_size": 64, "bits": 4 }, + "quantization_config": { "group_size": 64, "bits": 4 }, + "rms_norm_eps": 1e-08, + "rope_theta": 100000.0, + "tie_word_embeddings": false, + "torch_dtype": "bfloat16", + "transformers_version": "4.45.0.dev0", + "use_cache": true, + "vocab_size": 48000 +}"#; + +fn helium_1_preview_2b() -> ModelArgs { + serde_json::from_str(HELIUM_1_PREVIEW_2B_CONFIG).expect("the real config must parse") +} + +#[test] +fn the_real_config_parses_and_validates() { + let args = helium_1_preview_2b(); + assert_eq!(args.model_type, "helium"); + assert_eq!(args.hidden_size, 2560); + assert_eq!(args.num_hidden_layers, 24); + assert_eq!(args.intermediate_size, 7040); + assert_eq!(args.num_attention_heads, 20); + assert_eq!(args.num_kv_heads(), 20); + assert_eq!(args.vocab_size, 48000); + assert_eq!(args.max_position_embeddings, 4096); + assert!(!args.attention_bias); + assert!(!args.mlp_bias); + assert!(!args.tie_word_embeddings); + assert_eq!(args.rope_theta, 100_000.0); + assert_eq!(args.group_size(), 64); + assert_eq!(args.bits(), 4); + + // 1e-08 is unusually small for this family (Llama uses 1e-05) and must not + // be mistaken for an invalid value by the eps guard. + assert_eq!(args.rms_norm_eps, 1e-8); + args.validate().expect("the real config must validate"); +} + +#[test] +fn head_dim_is_derived_from_hidden_size_and_head_count() { + let args = helium_1_preview_2b(); + // Upstream's attention computes hidden_size // n_heads and never reads the + // declared field. On this checkpoint the two agree. + assert_eq!(args.head_dim(), 128); + assert_eq!(args.head_dim, Some(128)); +} + +#[test] +fn the_real_checkpoint_shape_signature_matches_what_loading_expects() { + // The `mlx-community/helium-1-preview-2b-4bit` SafeTensors header, checked + // against the widths `validate_weights` derives from `config.json`. Pinning + // the arithmetic here costs nothing, where materializing 559 real tensors + // in a unit test would cost hundreds of megabytes. + let args = helium_1_preview_2b(); + let bits = args.bits(); + let group_size = args.group_size(); + + // self_attn.{q,k,v,o}_proj.weight rows, and lm_head / embed_tokens rows. + assert_eq!(args.num_attention_heads * args.head_dim(), 2560); + assert_eq!(args.num_kv_heads() * args.head_dim(), 2560); + assert_eq!(args.hidden_size, 2560); + assert_eq!(args.vocab_size, 48000); + // mlp.{gate,up}_proj.weight rows. + assert_eq!(args.intermediate_size, 7040); + + // 4-bit affine packs eight values into each u32 word, so the stored input + // axis is `in_features * bits / 32` and the scales axis is + // `in_features / group_size`. + assert_eq!(args.hidden_size as i32 * bits / 32, 320); + assert_eq!(args.hidden_size as i32 / group_size, 40); + assert_eq!(args.intermediate_size as i32 * bits / 32, 880); + assert_eq!(args.intermediate_size as i32 / group_size, 110); + + // The input-width guard reconstructs the input axis the way MLX does, + // `scales.shape(-1) * group_size`, so the real scales widths must land back + // on the config's own values. This is what proves the guard accepts this + // checkpoint rather than only rejecting broken ones. + assert_eq!(40 * group_size, args.hidden_size as i32); + assert_eq!(110 * group_size, args.intermediate_size as i32); +} + +#[test] +fn config_validation_rejects_a_head_dim_that_disagrees_with_the_head_split() { + let mut args = helium_1_preview_2b(); + args.head_dim = Some(64); + let err = args + .validate() + .expect_err("a head_dim that disagrees with hidden_size / num_attention_heads is rejected"); + assert!(err.contains("head_dim"), "{err}"); + + // Absent is fine: it resolves to the derived width. + args.head_dim = None; + args.validate().unwrap(); + assert_eq!(args.head_dim(), 128); +} + +#[test] +fn config_validation_rejects_impossible_architecture_scalars() { + // Zero heads must be rejected before the divisibility check, because + // `0.is_multiple_of(0)` is true and `head_dim()` would then divide by zero. + let mut zero_heads = helium_1_preview_2b(); + zero_heads.num_attention_heads = 0; + zero_heads.hidden_size = 0; + zero_heads.head_dim = None; + assert!( + zero_heads + .validate() + .expect_err("zero heads") + .contains("num_attention_heads") + ); + + for mutate in [ + (|a: &mut ModelArgs| a.hidden_size = 0) as fn(&mut ModelArgs), + |a: &mut ModelArgs| a.num_hidden_layers = 0, + |a: &mut ModelArgs| a.num_hidden_layers = 1 << 20, + |a: &mut ModelArgs| a.intermediate_size = 0, + |a: &mut ModelArgs| a.vocab_size = 0, + |a: &mut ModelArgs| a.vocab_size = 1 << 30, + |a: &mut ModelArgs| a.max_position_embeddings = 0, + // 2560 is not divisible by 7. + |a: &mut ModelArgs| a.num_attention_heads = 7, + // 20 query heads cannot be grouped into 3 KV heads. + |a: &mut ModelArgs| a.num_key_value_heads = Some(3), + |a: &mut ModelArgs| a.num_key_value_heads = Some(0), + ] { + let mut args = helium_1_preview_2b(); + args.head_dim = None; + mutate(&mut args); + assert!( + args.validate().is_err(), + "an impossible architecture scalar must be rejected: {args:?}" + ); + } +} + +#[test] +fn config_validation_rejects_a_rope_configuration_that_would_abort_an_mlx_kernel() { + // MLX's `fast::rope` requires an even, positive `dims`. It enforces that by + // throwing, and `fast_rope` crosses the cxx bridge as `UniquePtr` rather + // than `Result`, so the throw is an uncatchable `std::terminate` at the + // first forward pass rather than a load error. head_dim comes out odd here + // (2560 / 512 = 5). + let mut odd = helium_1_preview_2b(); + odd.num_attention_heads = 512; + odd.num_key_value_heads = Some(512); + odd.head_dim = None; + assert!( + odd.validate() + .expect_err("odd head_dim") + .contains("head_dim") + ); + + for theta in [0.0, -1.0, f32::NAN, f32::INFINITY] { + let mut args = helium_1_preview_2b(); + args.rope_theta = theta; + assert!( + args.validate().is_err(), + "rope_theta {theta} must be rejected: RoPE exponentiates it, and a bad base NaNs \ + every rotated channel without throwing" + ); + } + + // The real value, and the smallest legal head width, both stay accepted. + let mut minimal = helium_1_preview_2b(); + minimal.num_attention_heads = 1280; + minimal.num_key_value_heads = Some(1280); + minimal.head_dim = None; + minimal.validate().expect("a head_dim of 2 is legal"); +} + +#[test] +fn config_validation_rejects_an_rms_norm_eps_that_would_nan_every_hidden_state() { + for eps in [0.0, -1e-5, f32::NAN, f32::INFINITY] { + let mut args = helium_1_preview_2b(); + args.rms_norm_eps = eps; + assert!( + args.validate().is_err(), + "rms_norm_eps {eps} must be rejected: `fast::rms_norm` never inspects it, so a bad \ + value produces NaN hidden states with no error at all" + ); + } + // The checkpoint's own unusually small eps must stay accepted. + let mut small = helium_1_preview_2b(); + small.rms_norm_eps = 1e-8; + small.validate().unwrap(); +} + +#[test] +fn config_validation_rejects_a_quantization_block_that_would_abort_an_mlx_kernel() { + for (group_size, bits) in [(64, 0), (64, 33), (64, -4), (0, 4), (-1, 4)] { + let mut args = helium_1_preview_2b(); + args.quantization = Some(super::Quantization { group_size, bits }); + assert!( + args.validate().is_err(), + "quantization group_size {group_size} / bits {bits} must be rejected" + ); + } + // A range check, not an allowlist: mixed-precision exports re-derive an + // effective bit width from the tensor shapes, so unusual-but-legal values + // stay accepted. + for (group_size, bits) in [(64, 4), (32, 8), (128, 6), (64, 32), (1, 1)] { + let mut args = helium_1_preview_2b(); + args.quantization = Some(super::Quantization { group_size, bits }); + args.validate().unwrap(); + } +} + +#[test] +fn eos_token_ids_come_from_the_config_not_from_the_llama3_backbone() { + let args = helium_1_preview_2b(); + assert_eq!(args.eos_token_ids(), vec![2]); + // Llama 3's hardcoded stop tokens are outside Helium's 48000-entry + // vocabulary, so delegating would mean generation never stops early. + assert!(!args.eos_token_ids().contains(&128001)); + + let list: ModelArgs = serde_json::from_str( + &HELIUM_1_PREVIEW_2B_CONFIG.replace("\"eos_token_id\": 2", "\"eos_token_id\": [2, 5]"), + ) + .unwrap(); + assert_eq!(list.eos_token_ids(), vec![2, 5]); +} + +// The RoPE convention. + +#[test] +fn to_llama3_args_asks_the_shared_decoder_for_traditional_rope() { + let args = helium_1_preview_2b(); + let llama = args.to_llama3_args(); + assert!( + llama.rope_traditional, + "Helium's single architectural difference is `nn.RoPE(..., traditional=True)`" + ); + assert_eq!(llama.head_dim, Some(128)); + assert_eq!(llama.num_key_value_heads, Some(20)); + assert_eq!(llama.rope_theta, 100_000.0); + assert!(llama.rope_scaling.is_none()); +} + +#[test] +fn a_llama_config_cannot_turn_on_traditional_rope_through_json() { + // `rope_traditional` is `#[serde(skip)]` on purpose. Every family that + // already uses this decoder rotates split-half, and a checkpoint author must + // not be able to change how an existing checkpoint decodes by adding a key. + let args: crate::models::llama3::ModelArgs = serde_json::from_str( + r#"{ + "model_type": "llama", + "hidden_size": 64, + "num_hidden_layers": 1, + "intermediate_size": 128, + "num_attention_heads": 2, + "rms_norm_eps": 1e-5, + "vocab_size": 32, + "rope_traditional": true + }"#, + ) + .unwrap(); + assert!(!args.rope_traditional); +} + +#[test] +fn traditional_and_split_half_rope_are_different_rotations() { + // Why any of this matters: the two conventions rotate different channel + // pairs, so they disagree on values while agreeing on every shape. + let x = filled(&[1, 2, 4, 8]); + let traditional = mlxcel_core::fast_rope(&x, 8, true, 100_000.0, 1.0, 0); + let split_half = mlxcel_core::fast_rope(&x, 8, false, 100_000.0, 1.0, 0); + assert_eq!( + mlxcel_core::array_shape(&traditional), + mlxcel_core::array_shape(&split_half), + "the two conventions are shape-identical, which is exactly why a shape test cannot \ + distinguish them" + ); + assert!( + max_abs_diff(&traditional, &split_half) > 1e-2, + "the two conventions must actually differ, or every test below is vacuous" + ); +} + +#[test] +fn attention_carries_the_traditional_flag_only_for_helium() { + let args = tiny_args(); + let weights = tiny_weights(&args); + + let helium = + Attention::from_weights(&weights, &args.to_llama3_args(), "model.layers.0.self_attn") + .unwrap(); + assert!(helium.rope_traditional); + + let mut llama_args = args.to_llama3_args(); + llama_args.rope_traditional = false; + let llama = Attention::from_weights(&weights, &llama_args, "model.layers.0.self_attn").unwrap(); + assert!( + !llama.rope_traditional, + "every pre-Helium family must keep the split-half rotation" + ); +} + +#[test] +fn helium_logits_differ_from_the_same_weights_with_split_half_rope() { + // THE regression test. If the flag is ever dropped between the config and + // `fast_rope`, both models become the same graph and this fails. + let args = tiny_args(); + let weights = tiny_weights(&args); + + let traditional = Llama3Model::from_weights(&weights, &args.to_llama3_args()).unwrap(); + let mut split_half_args = args.to_llama3_args(); + split_half_args.rope_traditional = false; + let split_half = Llama3Model::from_weights(&weights, &split_half_args).unwrap(); + + let prompt = mlxcel_core::from_slice_i32(&[1, 2, 3, 4, 5, 6, 7, 8], &[1, 8]); + let mut a_caches = traditional.make_caches(); + let mut b_caches = split_half.make_caches(); + let a = traditional.forward(&prompt, &mut a_caches, None); + let b = split_half.forward(&prompt, &mut b_caches, None); + + assert_eq!(mlxcel_core::array_shape(&a), mlxcel_core::array_shape(&b)); + + // Two thresholds, because an absolute one alone can go blind. + // + // The gap this test measures is only meaningful next to the magnitude of + // the logits it is measured on. An earlier version of `filled` made every + // head nearly collinear and collapsed the gap to 2.4e-5 against a logits + // scale of 5.8, a relative separation of 4e-6: green, and useless. That + // particular collapse is now caught by the absolute floor, but its mirror + // image is not. A future change that widens `tiny_args` or rescales + // `filled` can grow the logits without growing the separation, and an + // unanchored absolute floor would keep passing while the test drifts back + // toward blindness. + // + // Measured on the pinned inputs below: gap 1.19e-2, logits scale 6.89, + // relative separation 1.7e-3. That is ~12x the absolute floor and ~17x the + // relative floor, and it is reproducible run to run because `filled` is a + // fixed-seed LCG over fixed shapes. Both floors must keep real margin; if a + // refactor pushes either one close, widen the model rather than lowering + // the floor. + let gap = max_abs_diff(&a, &b); + let logits_scale = max_abs(&a); + assert!( + gap > 1e-3, + "traditional and split-half RoPE must produce different logits from the same weights \ + (gap {gap}, logits scale {logits_scale})" + ); + assert!( + gap > logits_scale * 1e-4, + "the two rotations must be separated by a meaningful fraction of the logits, not just by \ + an absolute epsilon (gap {gap}, logits scale {logits_scale})" + ); +} + +#[test] +fn every_rope_route_honors_the_flag_consistently() { + // Three code paths can rotate Q and K on this decoder: the single-sequence + // graph path, the batched path (`fast_rope_batched`), and the two fused + // quantized launchers, which cannot express the convention and are therefore + // bypassed (see `the_fused_qkv_rope_launcher_cannot_express_traditional_rope` + // and `Attention::forward`). + // + // Honoring the flag on one route and dropping it on another would make a + // sequence decode differently depending on whether it was scheduled alone or + // in a batch. This pins the batched route against the single-sequence route, + // and scales the tolerance against the split-half difference so the + // assertion cannot pass by both being wrong in the same way. + let args = tiny_args(); + let weights = tiny_weights(&args); + let helium = HeliumModel::from_weights(&weights, &args).unwrap(); + + let single_ids = mlxcel_core::from_slice_i32(&[1, 2, 3, 4, 5, 6, 7, 8], &[1, 8]); + let mut single_caches = LanguageModel::make_caches(&helium); + let single = LanguageModel::forward(&helium, &single_ids, &mut single_caches, None); + + let batched_ids = + mlxcel_core::from_slice_i32(&[1, 2, 3, 4, 5, 6, 7, 8, 1, 2, 3, 4, 5, 6, 7, 8], &[2, 8]); + let mut row0 = LanguageModel::make_caches(&helium); + let mut row1 = LanguageModel::make_caches(&helium); + let mut batch: Vec<&mut [KVCache]> = vec![row0.as_mut_slice(), row1.as_mut_slice()]; + let batched = LanguageModel::forward_batched(&helium, &batched_ids, &mut batch, None); + + let batched_row0 = mlxcel_core::slice(&batched, &[0, 0, 0], &[1, i32::MAX, i32::MAX]); + let route_gap = max_abs_diff(&single, &batched_row0); + + let mut split_half_args = args.to_llama3_args(); + split_half_args.rope_traditional = false; + let split_half = Llama3Model::from_weights(&weights, &split_half_args).unwrap(); + let mut split_half_caches = split_half.make_caches(); + let split_half_logits = split_half.forward(&single_ids, &mut split_half_caches, None); + let convention_gap = max_abs_diff(&single, &split_half_logits); + + assert!( + route_gap < convention_gap / 100.0, + "the batched route must apply the same rotation as the single-sequence route \ + (route gap {route_gap}, convention gap {convention_gap})" + ); +} + +#[test] +fn the_fused_qkv_rope_launcher_cannot_express_traditional_rope() { + // This is the fact the bypass in `Attention::forward` rests on, so it is + // asserted rather than assumed. `fused_qkv_project_split_rope` applies RoPE + // inside C++ with `traditional` hardcoded to `false` and takes no flag, and + // it is reachable only for *quantized* weights, which is exactly what the + // Helium validation checkpoint is. + // + // If someone later teaches the launcher the flag, this test starts failing + // and points at the gate that should then be removed. + let hidden = 64; + let head_dim = 32; + let heads = 2; + let group_size = 32; + let bits = 4; + + let mut weights = WeightMap::new(); + for name in ["q_proj", "k_proj", "v_proj"] { + let dense = filled(&[hidden, hidden]); + let quantized = mlxcel_core::quantize_weights(&dense, group_size, bits); + weights.insert( + format!("attn.{name}.weight"), + mlxcel_core::quantized_weights_w(&quantized), + ); + weights.insert( + format!("attn.{name}.scales"), + mlxcel_core::quantized_weights_scales(&quantized), + ); + if mlxcel_core::quantized_weights_has_biases(&quantized) { + weights.insert( + format!("attn.{name}.biases"), + mlxcel_core::quantized_weights_biases(&quantized), + ); + } + } + + let fused = FusedQKVLinear::from_weights_separate( + &weights, "attn", group_size, bits, heads, heads, head_dim, + ) + .unwrap(); + + let x = filled(&[1, 3, hidden]); + let (fused_q, _, _) = fused + .forward_split_rope_quantized(&x, head_dim, 100_000.0, 0) + .expect("quantized weights take the fused launcher"); + + // The same projection, rotated in Rust with each convention. + let (raw_q, _, _) = fused.forward(&x); + let raw_q = mlxcel_core::reshape(&raw_q, &[1, 3, heads, head_dim]); + let raw_q = mlxcel_core::transpose_axes(&raw_q, &[0, 2, 1, 3]); + let split_half = mlxcel_core::fast_rope(&raw_q, head_dim, false, 100_000.0, 1.0, 0); + let traditional = mlxcel_core::fast_rope(&raw_q, head_dim, true, 100_000.0, 1.0, 0); + + assert!( + max_abs_diff(&fused_q, &split_half) < 1e-4, + "the fused launcher applies the split-half rotation" + ); + assert!( + max_abs_diff(&fused_q, &traditional) > 1e-3, + "the fused launcher cannot produce the traditional rotation, which is why Helium must \ + not route through it" + ); +} + +// Weight-shape contract. + +#[test] +fn a_well_formed_checkpoint_loads_and_generates() { + let args = tiny_args(); + let weights = tiny_weights(&args); + let model = HeliumModel::from_weights(&weights, &args).unwrap(); + + assert_eq!(LanguageModel::num_layers(&model), args.num_hidden_layers); + assert_eq!(LanguageModel::eos_token_ids(&model), vec![2]); + + let prompt = mlxcel_core::from_slice_i32(&[1, 2, 3], &[1, 3]); + let mut caches = LanguageModel::make_caches(&model); + let logits = LanguageModel::forward(&model, &prompt, &mut caches, None); + assert_eq!( + mlxcel_core::array_shape(&logits), + vec![1, 3, args.vocab_size as i32] + ); + + let next = mlxcel_core::from_slice_i32(&[4], &[1, 1]); + let step = LanguageModel::forward(&model, &next, &mut caches, None); + assert_eq!( + mlxcel_core::array_shape(&step), + vec![1, 1, args.vocab_size as i32] + ); +} + +#[test] +fn loading_rejects_an_embedding_table_the_config_overstates() { + let args = tiny_args(); + let mut weights = tiny_weights(&args); + weights.insert( + "model.embed_tokens.weight".into(), + filled(&[args.vocab_size as i32 - 1, args.hidden_size as i32]), + ); + let err = validate_weights(&weights, &args).expect_err("an undersized table is rejected"); + assert!(err.contains("vocab_size"), "{err}"); +} + +#[test] +fn loading_rejects_a_projection_whose_shape_disagrees_with_the_config() { + let args = tiny_args(); + let head_dim = args.head_dim() as i32; + + // A q_proj sized for a different head count. Every offset stays in bounds, + // so nothing downstream would fault; the reshape in `Attention::forward` + // would throw inside MLX instead, which is an uncatchable abort. + let mut wrong_rows = tiny_weights(&args); + wrong_rows.insert( + "model.layers.0.self_attn.q_proj.weight".into(), + filled(&[head_dim, args.hidden_size as i32]), + ); + assert!(validate_weights(&wrong_rows, &args).is_err()); + + // A norm of the wrong width would broadcast-fail inside MLX. + let mut wrong_norm = tiny_weights(&args); + wrong_norm.insert( + "model.layers.1.post_attention_layernorm.weight".into(), + filled(&[args.hidden_size as i32 + 1]), + ); + assert!( + validate_weights(&wrong_norm, &args).is_err(), + "every layer is checked, not just layer 0" + ); + + // A transposed projection is caught by name rather than silently loaded. + let mut transposed = tiny_weights(&args); + transposed.insert( + "model.layers.0.mlp.gate_proj.weight".into(), + filled(&[args.hidden_size as i32, args.intermediate_size as i32]), + ); + assert!(validate_weights(&transposed, &args).is_err()); +} + +#[test] +fn loading_rejects_a_partially_quantized_attention_block() { + // The fused QKV loader decides "is this quantized?" from `q_proj.scales` + // alone and then concatenates all three projections. A checkpoint where they + // disagree aborts inside `concatenate` rather than failing to load. + let args = tiny_args(); + let mut weights = tiny_weights(&args); + weights.insert( + "model.layers.0.self_attn.q_proj.scales".into(), + filled(&[ + (args.num_attention_heads * args.head_dim()) as i32, + scale_cols(&args, args.hidden_size), + ]), + ); + let err = validate_weights(&weights, &args).expect_err("mixed quantization is rejected"); + assert!(err.contains("quantized"), "{err}"); +} + +#[test] +fn loading_rejects_an_attention_block_with_partial_quantization_biases() { + // The fused loader keeps the affine `biases` only when all three + // projections carry one, and silently drops the whole set otherwise, which + // dequantizes the survivors without their zero points. Nothing throws. + let args = tiny_args(); + let mut weights = tiny_weights(&args); + let q_rows = (args.num_attention_heads * args.head_dim()) as i32; + let kv_rows = (args.num_kv_heads() * args.head_dim()) as i32; + let cols = scale_cols(&args, args.hidden_size); + for (name, rows) in [("q_proj", q_rows), ("k_proj", kv_rows), ("v_proj", kv_rows)] { + weights.insert( + format!("model.layers.0.self_attn.{name}.scales"), + filled(&[rows, cols]), + ); + } + weights.insert( + "model.layers.0.self_attn.q_proj.biases".into(), + filled(&[q_rows, cols]), + ); + let err = validate_weights(&weights, &args).expect_err("partial biases are rejected"); + assert!(err.contains("quantization biases"), "{err}"); +} + +#[test] +fn a_consistently_quantized_attention_block_is_accepted() { + // The positive control for the two rejections below: a quantized block whose + // scales describe exactly `hidden_size` must still load, so the width check + // cannot be passing by rejecting everything quantized. + let args = tiny_args(); + let mut weights = tiny_weights(&args); + quantize_attention_block(&args, &mut weights, 0, scale_cols(&args, args.hidden_size)); + validate_weights(&weights, &args).unwrap(); +} + +#[test] +fn loading_rejects_a_quantized_projection_packed_for_a_different_input_width() { + // Packing compresses the input axis only, so a projection built for a + // different `hidden_size` still has exactly the right number of rows and + // survives every row check, and q/k/v agreeing with each other proves + // nothing about whether any of them agrees with the config. MLX + // reconstructs the input width as `scales.shape(-1) * group_size` and + // throws from `extract_quantized_matmul_dims` when it disagrees with the + // activation, which crosses the cxx bridge as an uncatchable abort at the + // first forward pass rather than a load error. + let args = tiny_args(); + let mut weights = tiny_weights(&args); + let wrong = scale_cols(&args, args.hidden_size * 2); + quantize_attention_block(&args, &mut weights, 0, wrong); + let err = validate_weights(&weights, &args).expect_err("a mis-packed input width is rejected"); + assert!(err.contains("input width"), "{err}"); +} + +#[test] +fn loading_rejects_a_quantized_output_head_packed_for_a_different_input_width() { + // Same defect on the untied head, where the throw comes from the final + // `quantized_matmul` instead of an attention projection. + let args = tiny_args(); + let mut weights = tiny_weights(&args); + let vocab = args.vocab_size as i32; + weights.insert( + "lm_head.scales".into(), + filled(&[vocab, scale_cols(&args, args.hidden_size * 2)]), + ); + let err = validate_weights(&weights, &args).expect_err("a mis-packed head is rejected"); + assert!(err.contains("input width"), "{err}"); +} + +#[test] +fn loading_rejects_quantization_biases_that_disagree_with_the_scales() { + // MLX requires the affine zero points to have the same shape as the scales + // and throws otherwise. Presence alone is not enough to check. + let args = tiny_args(); + let mut weights = tiny_weights(&args); + let cols = scale_cols(&args, args.hidden_size); + quantize_attention_block(&args, &mut weights, 0, cols); + weights.insert( + "model.layers.0.self_attn.q_proj.biases".into(), + filled(&[ + (args.num_attention_heads * args.head_dim()) as i32, + cols + 1, + ]), + ); + let err = validate_weights(&weights, &args).expect_err("mis-shaped zero points are rejected"); + assert!(err.contains("same shape"), "{err}"); +} + +#[test] +fn loading_rejects_an_attention_block_with_a_partial_dense_bias() { + // The dense `.bias` set has the same all-or-nothing rule as the affine + // `.biases` set: the fused loader concatenates q/k/v biases only when all + // three are present and drops the whole set otherwise, so a checkpoint + // carrying a bias on only some of them loads and silently runs those + // projections unbiased. + let args = tiny_args(); + let mut weights = tiny_weights(&args); + weights.insert( + "model.layers.0.self_attn.q_proj.bias".into(), + filled(&[(args.num_attention_heads * args.head_dim()) as i32]), + ); + let err = validate_weights(&weights, &args).expect_err("a partial bias set is rejected"); + assert!(err.contains("all carry a bias"), "{err}"); +} + +#[test] +fn loading_rejects_a_missing_untied_head() { + let args = tiny_args(); + let mut weights = tiny_weights(&args); + weights.remove("lm_head.weight"); + assert!(validate_weights(&weights, &args).is_err()); + + // The same checkpoint is fine once the config says the head is tied. + let mut tied = tiny_args(); + tied.tie_word_embeddings = true; + validate_weights(&weights, &tied).unwrap(); +} + +// Helpers. + +/// A Helium small enough to run in a unit test, keeping the checkpoint's own +/// head split (multi-head, `num_key_value_heads == num_attention_heads`) and its +/// `rope_theta`. +fn tiny_args() -> ModelArgs { + serde_json::from_str( + r#"{ + "model_type": "helium", + "hidden_size": 64, + "num_hidden_layers": 2, + "intermediate_size": 128, + "num_attention_heads": 4, + "num_key_value_heads": 4, + "head_dim": 16, + "rms_norm_eps": 1e-08, + "rope_theta": 100000.0, + "attention_bias": false, + "mlp_bias": false, + "tie_word_embeddings": false, + "eos_token_id": 2, + "max_position_embeddings": 128, + "vocab_size": 32 + }"#, + ) + .unwrap() +} + +/// A float checkpoint with the Helium key layout: `model.embed_tokens`, +/// `model.layers.N.self_attn.{q,k,v,o}_proj`, +/// `model.layers.N.mlp.{gate,up,down}_proj`, the two RMSNorms, `model.norm`, and +/// an untied `lm_head`. +fn tiny_weights(args: &ModelArgs) -> WeightMap { + let hidden = args.hidden_size as i32; + let ff = args.intermediate_size as i32; + let vocab = args.vocab_size as i32; + let q_out = (args.num_attention_heads * args.head_dim()) as i32; + let kv_out = (args.num_kv_heads() * args.head_dim()) as i32; + + let mut w = WeightMap::new(); + w.insert("model.embed_tokens.weight".into(), filled(&[vocab, hidden])); + w.insert("model.norm.weight".into(), ones(&[hidden])); + w.insert("lm_head.weight".into(), filled(&[vocab, hidden])); + + for layer in 0..args.num_hidden_layers { + let p = format!("model.layers.{layer}"); + w.insert(format!("{p}.input_layernorm.weight"), ones(&[hidden])); + w.insert( + format!("{p}.post_attention_layernorm.weight"), + ones(&[hidden]), + ); + w.insert( + format!("{p}.self_attn.q_proj.weight"), + filled(&[q_out, hidden]), + ); + w.insert( + format!("{p}.self_attn.k_proj.weight"), + filled(&[kv_out, hidden]), + ); + w.insert( + format!("{p}.self_attn.v_proj.weight"), + filled(&[kv_out, hidden]), + ); + w.insert( + format!("{p}.self_attn.o_proj.weight"), + filled(&[hidden, q_out]), + ); + w.insert(format!("{p}.mlp.gate_proj.weight"), filled(&[ff, hidden])); + w.insert(format!("{p}.mlp.up_proj.weight"), filled(&[ff, hidden])); + w.insert(format!("{p}.mlp.down_proj.weight"), filled(&[hidden, ff])); + } + w +} + +/// Scales columns for a quantized `[out, in_features]` weight: one group per +/// `group_size` input channels, which is the width MLX reconstructs the input +/// axis from. +fn scale_cols(args: &ModelArgs, in_features: usize) -> i32 { + in_features as i32 / args.group_size() +} + +/// Turn one layer's `q_proj` / `k_proj` / `v_proj` / `o_proj` into a quantized +/// block by adding `.scales` and `.biases` with `cols` groups each, leaving the +/// `.weight` tensors from [`tiny_weights`] in place (only the shapes are read). +fn quantize_attention_block(args: &ModelArgs, weights: &mut WeightMap, layer: usize, cols: i32) { + let q_rows = (args.num_attention_heads * args.head_dim()) as i32; + let kv_rows = (args.num_kv_heads() * args.head_dim()) as i32; + let hidden = args.hidden_size as i32; + for (name, rows) in [ + ("q_proj", q_rows), + ("k_proj", kv_rows), + ("v_proj", kv_rows), + ("o_proj", hidden), + ] { + weights.insert( + format!("model.layers.{layer}.self_attn.{name}.scales"), + filled(&[rows, cols]), + ); + weights.insert( + format!("model.layers.{layer}.self_attn.{name}.biases"), + filled(&[rows, cols]), + ); + } +} + +/// Deterministic pseudo-random filler in `[-0.5, 0.5)`. +/// +/// A short repeating pattern (the obvious `(i % 7 - 3) * 0.1`) makes every head +/// nearly collinear, which flattens the attention scores and shrinks the +/// difference between two RoPE conventions to a few parts per million at the +/// logits. That would leave the convention tests technically green and +/// practically blind, so the filler is a fixed-seed LCG instead: still +/// bit-reproducible, but with enough spread for a rotation to matter. +fn filled(shape: &[i32]) -> UniquePtr { + let n: i32 = shape.iter().product(); + let mut state: u32 = 0x9E37_79B9; + let data: Vec = (0..n) + .map(|_| { + state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223); + (state >> 8) as f32 / (1u32 << 24) as f32 - 0.5 + }) + .collect(); + mlxcel_core::from_slice_f32(&data, shape) +} + +fn ones(shape: &[i32]) -> UniquePtr { + let n: i32 = shape.iter().product(); + mlxcel_core::from_slice_f32(&vec![1.0; n as usize], shape) +} + +fn max_abs_diff(a: &MlxArray, b: &MlxArray) -> f32 { + max_abs(&mlxcel_core::subtract(a, b)) +} + +/// Largest absolute value in `a`, used to scale a difference against the +/// magnitude it was measured on. +fn max_abs(a: &MlxArray) -> f32 { + mlxcel_core::item_f32(&mlxcel_core::max_all(&mlxcel_core::abs(a))) +} diff --git a/src/models/llama3.rs b/src/models/llama3.rs index 13c852e35..6a2795a24 100644 --- a/src/models/llama3.rs +++ b/src/models/llama3.rs @@ -62,6 +62,22 @@ pub struct ModelArgs { #[serde(default)] pub tie_word_embeddings: bool, + + /// Rotate interleaved channel pairs `(2i, 2i+1)` instead of the split-half + /// pairs `(i, i + dims/2)`. + /// + /// Deliberately **not** deserialized. Every family that reaches this loader + /// today rotates split-half, and a checkpoint author who writes + /// `"rope_traditional": true` into a Llama `config.json` must not be able to + /// silently change how an existing checkpoint decodes. The flag is set + /// programmatically by the loader of a family whose upstream definition + /// fixes it, which is currently only [`crate::models::helium`] (upstream + /// builds `nn.RoPE(head_dim, traditional=True, base=rope_theta)`). + /// + /// See [`Attention::forward`] for why setting this also disables the two + /// fused RoPE fast paths. + #[serde(skip)] + pub rope_traditional: bool, } #[derive(Debug, Clone, Deserialize)] @@ -122,9 +138,49 @@ pub struct Attention { pub scale: f32, pub rope_dims: i32, pub rope_base: f32, + /// See [`ModelArgs::rope_traditional`]. `false` for every family that used + /// this attention before Helium, so their graphs are unchanged. + pub rope_traditional: bool, } impl Attention { + /// Dense attention with RoPE. + /// + /// Used by: Llama (`llama` / `mistral` checkpoints), Qwen2 / Qwen2.5 (which + /// re-export this attention), Helium (which reuses it with + /// `rope_traditional` set), the Llama-3.2-Vision (`mllama`) text decoder's + /// self-attention layers, every VLM whose text backbone is `Llama3Model` or + /// `Qwen2Model` (Pixtral, LLaVA, SmolVLM / Idefics3, Idefics2, InternVL, + /// FastVLM, dots.ocr), the `llama` and `mistral` pipeline stage executors, + /// and the tensor-parallel Llama runtime. + /// + /// # Traditional RoPE and the fused fast paths + /// + /// Three code paths below can rotate Q and K, and all three must agree on + /// the rotation convention or the model decodes a different graph depending + /// on an environment variable and on whether the checkpoint is quantized: + /// + /// 1. `fused_causal_prefill_attention` (quantized prefill, opt-in through + /// `MLXCEL_ENABLE_FUSED_CAUSAL_PREFILL_ATTENTION`), + /// 2. [`FusedQKVLinear::forward_split_rope`] (quantized projection, opt-in + /// through `MLXCEL_ENABLE_FUSED_QKV_SPLIT_ROPE`), + /// 3. the graph fallback, which calls `fast_rope` directly. + /// + /// Only the third can express the convention. The first two apply RoPE + /// inside a C++ launcher that hardcodes `traditional = false` + /// (`mlx_cxx_bridge.cpp`, `fused_causal_prefill_attention` and + /// `fused_qkv_project_split_rope`), and neither takes a flag, so there is no + /// way to ask them for the interleaved rotation. Routing a traditional-RoPE + /// model through either one applies the wrong rotation to correctly shaped + /// tensors: nothing throws, the KV cache is the right shape, and the model + /// emits fluent text out of a mis-rotated attention. + /// + /// Both are therefore gated on `!self.rope_traditional`, so a + /// traditional-RoPE model always takes the graph fallback, which receives + /// the flag. Neither fast path is on by default, so this costs those models + /// nothing that is enabled today. Teaching the two launchers the flag would + /// be the better long-term fix, but it is a cxx bridge signature change, and + /// correctness must not wait on it. pub fn forward( &self, x: &MlxArray, @@ -139,6 +195,7 @@ impl Attention { if l > 1 && mask.is_none() && cache.is_empty() + && !self.rope_traditional && std::env::var("MLXCEL_ENABLE_FUSED_CAUSAL_PREFILL_ATTENTION").is_ok() && let (Some(qkv_weight), Some(o_weight)) = ( self.qkv_proj.qkv_proj.as_quantized_weight(), @@ -175,10 +232,16 @@ impl Attention { return output; } - let (q, k, v) = if let Some((q, k, v)) = + let fused_split_rope = if self.rope_traditional { + // The fused launcher hardcodes `traditional = false`; see the method + // doc comment. Skipping it is what keeps the rotation correct. + None + } else { self.qkv_proj .forward_split_rope(x, self.rope_dims, self.rope_base, offset) - { + }; + + let (q, k, v) = if let Some((q, k, v)) = fused_split_rope { (q, k, v) } else { // Fallback for non-quantized weights: preserve the existing Rust path. @@ -194,8 +257,22 @@ impl Attention { let k = mlxcel_core::transpose_axes(&k, &[0, 2, 1, 3]); let v = mlxcel_core::transpose_axes(&v, &[0, 2, 1, 3]); - let q = mlxcel_core::fast_rope(&q, self.rope_dims, false, self.rope_base, 1.0, offset); - let k = mlxcel_core::fast_rope(&k, self.rope_dims, false, self.rope_base, 1.0, offset); + let q = mlxcel_core::fast_rope( + &q, + self.rope_dims, + self.rope_traditional, + self.rope_base, + 1.0, + offset, + ); + let k = mlxcel_core::fast_rope( + &k, + self.rope_dims, + self.rope_traditional, + self.rope_base, + 1.0, + offset, + ); (q, k, v) }; @@ -268,7 +345,9 @@ impl Attention { /// cache updates and attention before concatenating the results back into /// `[B, T, hidden_dim]`. /// - /// Used by: Llama3 batched decode and full-sequence batched prefill + /// Used by: the batched decode and full-sequence batched prefill of every + /// family that reaches [`Attention::forward`]; see that method's `Used by` + /// list, which this path shares. pub fn forward_split_attention( &self, q_batched: &MlxArray, @@ -301,10 +380,14 @@ impl Attention { let k_batched = mlxcel_core::transpose_axes(&k_batched, &[0, 2, 1, 3]); let v_batched = mlxcel_core::transpose_axes(&v_batched, &[0, 2, 1, 3]); + // Batched decode / batched prefill share the same rotation convention as + // the single-sequence path above; dropping the flag on one route while + // honoring it on the other would make a sequence decode differently + // depending on whether it was scheduled alone or in a batch. let q_batched = mlxcel_core::fast_rope_batched( &q_batched, self.rope_dims, - false, + self.rope_traditional, self.rope_base, 1.0, &metadata.rope_offsets, @@ -312,7 +395,7 @@ impl Attention { let k_batched = mlxcel_core::fast_rope_batched( &k_batched, self.rope_dims, - false, + self.rope_traditional, self.rope_base, 1.0, &metadata.rope_offsets, @@ -453,6 +536,19 @@ impl Attention { result } + /// Build the dense attention block from a checkpoint. + /// + /// `rope_traditional` is carried over from [`ModelArgs`], where it is + /// `#[serde(skip)]`, so it is `false` for every family whose config is + /// parsed from JSON and `true` only when a loader sets it programmatically. + /// + /// Used by: Llama (`llama` / `mistral` checkpoints), Qwen2 / Qwen2.5 (which + /// re-export this attention), Helium (which reuses it with + /// `rope_traditional` set), the Llama-3.2-Vision (`mllama`) text decoder's + /// self-attention layers, every VLM whose text backbone is `Llama3Model` or + /// `Qwen2Model` (Pixtral, LLaVA, SmolVLM / Idefics3, Idefics2, InternVL, + /// FastVLM, dots.ocr), the `llama` and `mistral` pipeline stage executors, + /// and the tensor-parallel Llama runtime. pub fn from_weights( weights: &WeightMap, args: &ModelArgs, @@ -487,6 +583,7 @@ impl Attention { scale: 1.0 / (head_dim as f32).sqrt(), rope_dims: head_dim, rope_base: args.rope_theta, + rope_traditional: args.rope_traditional, }) } } diff --git a/src/models/mod.rs b/src/models/mod.rs index 7937ae6bc..12acd5957 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -68,6 +68,7 @@ pub mod gpt_neox; pub mod gpt_oss; pub mod granite; pub mod granitemoehybrid; +pub mod helium; pub mod hunyuan_moe; pub mod hunyuan_v1_dense; pub mod hunyuan_vl; @@ -172,6 +173,7 @@ pub use gpt_oss::{GptOssModel, GptOssWrapper}; pub use gpt2::Gpt2Model; pub use granite::GraniteModel; pub use granitemoehybrid::GraniteMoeHybridModel; +pub use helium::HeliumModel; pub use hunyuan_moe::HunyuanMoeModel; pub use hunyuan_v1_dense::HunyuanV1DenseModel; pub use internlm2::InternLM2Model; @@ -385,6 +387,7 @@ pub enum ModelType { Mellum, // Mellum 2 (JetBrains hybrid-attention MoE code model) // Other Transformer models + Helium, // Kyutai Helium (Llama-shaped dense decoder with traditional RoPE) MiniCPM, MiniCPM3, StableLM, @@ -572,6 +575,7 @@ pub const ALL_MODEL_TYPES: &[ModelType] = &[ ModelType::StarCoder2, ModelType::Mellum, // Other Transformer models + ModelType::Helium, ModelType::MiniCPM, ModelType::MiniCPM3, ModelType::StableLM, @@ -833,6 +837,10 @@ impl ModelType { "GPT-NeoX (EleutherAI Pythia, partial RoPE, parallel residual)", "Specialized", ), + ModelType::Helium => ( + "Kyutai Helium (dense Llama shape, traditional RoPE)", + "Specialized", + ), ModelType::StarCoder2 => ("StarCoder 2", "Specialized"), ModelType::Mellum => ("Mellum 2 (JetBrains code)", "Specialized"), ModelType::StableLM => ("StableLM", "Specialized"), @@ -1051,6 +1059,7 @@ mod metadata_tests { GptNeoX, StarCoder2, Mellum, + Helium, MiniCPM, MiniCPM3, StableLM,