fuel-transformers 0.19.1

// SPDX-License-Identifier: MIT OR Apache-2.0
//! GLM-4 (new architecture, ChatGLM-4-style) ported to the
//! lazy-graph API.
//!
//! This is the HuggingFace-released GLM-4 architecture that
//! ships under `THUDM/glm-4-9b-chat` / `glm-4-9b-chat-hf` and
//! the upstream `glm4_new` module in fuel-transformers. It is
//! distinct from the older [`crate::models::lazy_glm4`] port (which
//! targeted the original CodeGeeX/ChatGLM3 lineage).
//!
//! # Differences from `lazy_glm4`
//!
//!   1. **Four RmsNorms per layer** instead of two. Sublayer
//!      structure is
//!      `input_norm → attn → post_self_attn_norm → +residual →
//!       post_attn_norm → mlp → post_mlp_norm → +residual`.
//!      Mirrors the Gemma-2 pattern: post-norms sit AFTER the
//!      sublayer's main op but BEFORE the residual add.
//!   2. **Fused `gate_up_proj`**: a single linear of width
//!      `2 * intermediate_size` whose output is sliced into
//!      gate (first half) and up (second half). MLP becomes
//!      `down(act(gate) * up)` with `act` configurable.
//!   3. **Interleaved partial RoPE**: same trick as
//!      `lazy_glm4` — `partial_rotary_factor` controls how
//!      many head-dim features get rotated; the remainder
//!      passes through unchanged. The lazy port reuses
//!      `lazy_glm4::apply_interleaved_partial_rope`.
//!   4. **Optional sliding-window mask** (eager carries it
//!      in the config but applies the standard causal mask
//!      by default — `sliding_window` is plumbed through but
//!      defaults to no window in the public release).
//!
//! Keeps standard GLM4 bits: separate Q/K/V (NOT fused),
//! optional QKV biases via `attention_bias`, GQA with
//! `num_key_value_heads`, tied embeddings.
//!
//! # Scope (v1)
//!
//! Forward-only, single sequence (`batch == 1`), no KV cache
//! (recomputes per call), F32. Returns vocab logits
//! `(1, seq, vocab_size)`. The standard causal mask is built
//! per-forward; sliding-window is plumbed through via the
//! config field.

use crate::models::lazy_glm4::apply_interleaved_partial_rope;
use fuel_ir::Shape;
use fuel_ir::error::Result;
use fuel_tensor::Device;
use fuel_tensor::{lazy::Tensor, lazy::WeightStorage};
use std::sync::Arc;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Glm4NewActivation {
    Silu,
    Gelu,
    GeluPytorchTanh,
}

#[derive(Debug, Clone, PartialEq)]
pub struct Glm4NewConfig {
    pub vocab_size: usize,
    pub hidden_size: usize,
    pub intermediate_size: usize,
    pub num_hidden_layers: usize,
    pub num_attention_heads: usize,
    pub num_key_value_heads: usize,
    /// `None` → derived from `hidden_size / num_attention_heads`.
    pub head_dim: Option<usize>,
    /// Fraction of `head_dim` that gets rotated. `None` → 1.0 (full).
    pub partial_rotary_factor: Option<f32>,
    pub attention_bias: bool,
    pub max_position_embeddings: usize,
    pub sliding_window: Option<usize>,
    pub tie_word_embeddings: bool,
    pub rope_theta: f64,
    pub rms_norm_eps: f64,
    pub hidden_act: Glm4NewActivation,
}

impl Glm4NewConfig {
    pub fn head_dim(&self) -> usize {
        self.head_dim
            .unwrap_or(self.hidden_size / self.num_attention_heads)
    }
    pub fn rotary_dim(&self) -> usize {
        let d = self.head_dim();
        match self.partial_rotary_factor {
            None => d,
            Some(f) => (f * d as f32) as usize,
        }
    }
    /// `THUDM/glm-4-9b-chat` preset (approximate; actual
    /// `config.json` from HuggingFace overrides).
    pub fn glm4_9b_chat() -> Self {
        Self {
            vocab_size: 151_552,
            hidden_size: 4_096,
            intermediate_size: 13_696,
            num_hidden_layers: 40,
            num_attention_heads: 32,
            num_key_value_heads: 2,
            head_dim: Some(128),
            partial_rotary_factor: Some(0.5),
            attention_bias: true,
            max_position_embeddings: 131_072,
            sliding_window: None,
            tie_word_embeddings: false,
            rope_theta: 5_000_000.0,
            rms_norm_eps: 1e-5,
            hidden_act: Glm4NewActivation::Silu,
        }
    }
}

/// Map a HuggingFace `hidden_act` string to [`Glm4NewActivation`]. Unknown values
/// ERROR rather than silently defaulting — a config naming an activation this port
/// does not implement is a fact worth surfacing, not swallowing.
fn glm4_new_activation_from_str(s: &str) -> fuel_ir::error::Result<Glm4NewActivation> {
    match s {
        "silu" | "swish" => Ok(Glm4NewActivation::Silu),
        "gelu" => Ok(Glm4NewActivation::Gelu),
        "gelu_pytorch_tanh" => Ok(Glm4NewActivation::GeluPytorchTanh),
        other => Err(fuel_ir::error::Error::Msg(format!(
            "unsupported GLM-4 hidden_act {other:?} (expected silu/gelu/gelu_pytorch_tanh)"
        ))),
    }
}

// ROADMAP item 8 (II): config-from-path on the #57 template. A `serde` raw with
// HF field names + GLM-4's own constant defaults, then `resolve` routes kv heads
// through the shared `fuel_loaders::hf_config` rule. `head_dim` and
// `partial_rotary_factor` are already `Option` on Glm4NewConfig (derived at
// use-site by `head_dim()` / `rotary_dim()`), so they pass through verbatim —
// GLM-4-0414 ships an explicit `partial_rotary_factor: 0.5`, read faithfully.
// `hidden_act` is a non-serde enum, so it is parsed as a string and mapped.
#[derive(Debug, Clone, serde::Deserialize)]
struct Glm4NewConfigRaw {
    vocab_size: usize,
    hidden_size: usize,
    intermediate_size: usize,
    num_hidden_layers: usize,
    num_attention_heads: usize,
    #[serde(default)]
    num_key_value_heads: Option<usize>,
    #[serde(default)]
    head_dim: Option<usize>,
    #[serde(default)]
    partial_rotary_factor: Option<f32>,
    #[serde(default)]
    attention_bias: bool,
    max_position_embeddings: usize,
    #[serde(default)]
    sliding_window: Option<usize>,
    #[serde(default)]
    tie_word_embeddings: bool,
    #[serde(default = "default_glm4_new_rope_theta")]
    rope_theta: f64,
    #[serde(default = "default_glm4_new_rms_norm_eps")]
    rms_norm_eps: f64,
    #[serde(default)]
    hidden_act: Option<String>,
}

fn default_glm4_new_rope_theta() -> f64 {
    10_000.0
}
fn default_glm4_new_rms_norm_eps() -> f64 {
    1e-5
}

impl Glm4NewConfigRaw {
    fn from_json_str(json: &str) -> fuel_ir::error::Result<Self> {
        serde_json::from_str(json)
            .map_err(|e| fuel_ir::error::Error::Msg(format!("parsing GLM-4 config.json: {e}")))
    }

    fn resolve(self) -> fuel_ir::error::Result<Glm4NewConfig> {
        let hidden_act = match self.hidden_act.as_deref() {
            None => Glm4NewActivation::Silu,
            Some(s) => glm4_new_activation_from_str(s)?,
        };
        Ok(Glm4NewConfig {
            vocab_size: self.vocab_size,
            hidden_size: self.hidden_size,
            intermediate_size: self.intermediate_size,
            num_hidden_layers: self.num_hidden_layers,
            num_attention_heads: self.num_attention_heads,
            num_key_value_heads: fuel_loaders::hf_config::num_key_value_heads(
                self.num_key_value_heads,
                self.num_attention_heads,
            )?,
            // Option, derived at use-site by `head_dim()` — passed through.
            head_dim: self.head_dim,
            partial_rotary_factor: self.partial_rotary_factor,
            attention_bias: self.attention_bias,
            max_position_embeddings: self.max_position_embeddings,
            sliding_window: self.sliding_window,
            tie_word_embeddings: self.tie_word_embeddings,
            rope_theta: self.rope_theta,
            rms_norm_eps: self.rms_norm_eps,
            hidden_act,
        })
    }
}

impl Glm4NewConfig {
    /// Parse a HuggingFace `config.json` string into a [`Glm4NewConfig`].
    ///
    /// ROADMAP item 8 (II): reads the artifact rather than returning a preset —
    /// see the born-red `glm4_new_config_from_hf_json_parses_the_artifact_not_a_preset`.
    pub fn from_hf_json_str(json: &str) -> fuel_ir::error::Result<Self> {
        Glm4NewConfigRaw::from_json_str(json)?.resolve()
    }
}

#[derive(Debug, Clone)]
pub struct Glm4NewLayerWeights {
    pub input_norm_gain: Arc<[f32]>,
    pub post_self_attn_norm_gain: Arc<[f32]>,
    pub post_attn_norm_gain: Arc<[f32]>,
    pub post_mlp_norm_gain: Arc<[f32]>,
    pub q: WeightStorage,
    pub q_bias: Option<Arc<[f32]>>,
    pub k: WeightStorage,
    pub k_bias: Option<Arc<[f32]>>,
    pub v: WeightStorage,
    pub v_bias: Option<Arc<[f32]>>,
    pub o: WeightStorage,
    /// Fused `[hidden, 2 * intermediate]`: first half = gate,
    /// second half = up.
    pub gate_up: WeightStorage,
    pub down: WeightStorage,
}

#[derive(Debug, Clone)]
pub struct Glm4NewWeights {
    pub token_embedding: Arc<[f32]>,
    pub layers: Vec<Glm4NewLayerWeights>,
    pub final_norm_gain: Arc<[f32]>,
    /// `None` when `tie_word_embeddings == true`; the embedding
    /// table is reused for the lm_head.
    pub lm_head: Option<WeightStorage>,
}

#[derive(Debug, Clone)]
pub struct Glm4NewModel {
    pub config: Glm4NewConfig,
    pub weights: Glm4NewWeights,
}

impl Glm4NewModel {
    pub fn forward(&self, tokens: &[u32], start_pos: usize) -> Result<Tensor> {
        let h_norm = self.run_backbone(tokens, start_pos)?;
        self.apply_lm_head(&h_norm)
    }

    /// Run the decoder forward up to the final RmsNorm and
    /// return per-token hidden states `(1, seq, hidden_size)`.
    /// Skips tied/untied `lm_head` projection.
    pub fn forward_hidden(&self, tokens: &[u32], start_pos: usize) -> Result<Tensor> {
        self.run_backbone(tokens, start_pos)
    }

    /// Multimodal entry point. Skips token embedding; runs the decoder
    /// over pre-embedded inputs. GLM4-new does NOT scale embeddings.
    pub fn forward_embeds(&self, embeds: &Tensor, start_pos: usize) -> Result<Tensor> {
        let h_norm = self.run_backbone_embeds(embeds, start_pos)?;
        self.apply_lm_head(&h_norm)
    }

    /// Hidden-state variant of [`Self::forward_embeds`].
    pub fn forward_hidden_embeds(&self, embeds: &Tensor, start_pos: usize) -> Result<Tensor> {
        self.run_backbone_embeds(embeds, start_pos)
    }

    /// Build per-token embeddings without running the decoder.
    pub fn embed_tokens_anchored(&self, anchor: &Tensor, tokens: &[u32]) -> Result<Tensor> {
        let cfg = &self.config;
        anchor.embed_tokens_anchored(
            self.weights.token_embedding.clone(),
            cfg.vocab_size,
            cfg.hidden_size,
            tokens,
        )
    }

    fn apply_lm_head(&self, h_norm: &Tensor) -> Result<Tensor> {
        let cfg = &self.config;
        let weights = &self.weights;
        let logits = match &weights.lm_head {
            Some(w) => w.apply_linear(h_norm, cfg.hidden_size, cfg.vocab_size)?,
            None => {
                let lm_w = h_norm.const_f32_like(
                    Arc::clone(&weights.token_embedding),
                    Shape::from_dims(&[cfg.vocab_size, cfg.hidden_size]),
                )?;
                h_norm.matmul(&lm_w.transpose()?)?
            }
        };
        Ok(logits)
    }

    fn run_backbone(&self, tokens: &[u32], start_pos: usize) -> Result<Tensor> {
        let cfg = &self.config;
        let weights = &self.weights;
        let seq = tokens.len();
        assert!(seq > 0);

        let h = Tensor::embed_tokens(
            weights.token_embedding.clone(),
            cfg.vocab_size,
            cfg.hidden_size,
            tokens,
            &Device::cpu(),
        )?;
        self.run_backbone_embeds(&h, start_pos)
    }

    fn run_backbone_embeds(&self, embeds: &Tensor, start_pos: usize) -> Result<Tensor> {
        let cfg = &self.config;
        let weights = &self.weights;
        let dims = embeds.shape();
        let dims = dims.dims();
        if dims.len() != 3 || dims[0] != 1 || dims[2] != cfg.hidden_size {
            return Err(fuel_ir::error::Error::Msg(format!(
                "Glm4NewModel::forward_embeds: expected embeds shape (1, seq, hidden_size={}), got {:?}",
                cfg.hidden_size, dims,
            )).bt());
        }
        let seq = dims[1];
        if seq == 0 {
            return Err(fuel_ir::error::Error::Msg(
                "Glm4NewModel::forward_embeds: seq must be > 0".into(),
            )
            .bt());
        }
        let head_dim = cfg.head_dim();
        let q_dim = cfg.num_attention_heads * head_dim;
        let kv_dim = cfg.num_key_value_heads * head_dim;
        if !cfg
            .num_attention_heads
            .is_multiple_of(cfg.num_key_value_heads)
        {
            return Err(fuel_ir::error::Error::Msg(
                "Glm4NewConfig: num_attention_heads must be a multiple of num_key_value_heads"
                    .into(),
            )
            .bt());
        }
        let mut h = embeds.clone();

        let rope_dim = cfg.rotary_dim();
        let (rope_cos, rope_sin) = h.rope_tables_const(cfg.rope_theta, start_pos, seq, rope_dim);

        let mask = self.build_mask(&h, seq);

        for layer in &weights.layers {
            h = self.apply_layer(&h, layer, &rope_cos, &rope_sin, &mask, q_dim, kv_dim)?;
        }
        h.rms_norm_affine(
            std::sync::Arc::clone(&weights.final_norm_gain),
            cfg.rms_norm_eps,
        )
    }

    fn build_mask(&self, anchor: &Tensor, seq: usize) -> Tensor {
        let cfg = &self.config;
        let mut mask_data = vec![0.0_f32; seq * seq];
        for i in 0..seq {
            for j in 0..seq {
                let masked = if j > i {
                    true
                } else if let Some(w) = cfg.sliding_window {
                    j + w < i
                } else {
                    false
                };
                if masked {
                    mask_data[i * seq + j] = f32::NEG_INFINITY;
                }
            }
        }
        anchor.const_f32_like(mask_data, Shape::from_dims(&[1, 1, seq, seq]))
        .expect(
                "build_mask: buffer is vec![_; seq*seq] and the shape's elem_count is seq*seq -- \n             both derived from `seq` in this function; the loop writes in place",
        )
    }

    #[allow(clippy::too_many_arguments)]
    fn apply_layer(
        &self,
        x: &Tensor,
        layer: &Glm4NewLayerWeights,
        rope_cos: &Tensor,
        rope_sin: &Tensor,
        mask: &Tensor,
        q_dim: usize,
        kv_dim: usize,
    ) -> Result<Tensor> {
        let cfg = &self.config;
        let head_dim = cfg.head_dim();
        let rope_dim = cfg.rotary_dim();
        let dims = x.shape();
        let dims = dims.dims();
        let batch = dims[0];
        let seq = dims[1];

        // Attention sublayer: post_self_attn_norm(attn(input_norm(x))) + x.
        let x_in = x.rms_norm_affine(
            std::sync::Arc::clone(&layer.input_norm_gain),
            cfg.rms_norm_eps,
        )?;
        let q = layer
            .q
            .apply_linear(&x_in, cfg.hidden_size, q_dim)?
            .add_optional_trailing_bias(layer.q_bias.as_ref())?;
        let k = layer
            .k
            .apply_linear(&x_in, cfg.hidden_size, kv_dim)?
            .add_optional_trailing_bias(layer.k_bias.as_ref())?;
        let v = layer
            .v
            .apply_linear(&x_in, cfg.hidden_size, kv_dim)?
            .add_optional_trailing_bias(layer.v_bias.as_ref())?;

        let _ = (batch, seq);
        let q = q.split_heads(cfg.num_attention_heads, head_dim)?;
        let k = k.split_heads(cfg.num_key_value_heads, head_dim)?;
        let v = v.split_heads(cfg.num_key_value_heads, head_dim)?;

        let q_r = apply_interleaved_partial_rope(&q, rope_cos, rope_sin, head_dim, rope_dim)?;
        let k_r = apply_interleaved_partial_rope(&k, rope_cos, rope_sin, head_dim, rope_dim)?;

        let n_rep = cfg.num_attention_heads / cfg.num_key_value_heads;
        let k_full = k_r.repeat_interleave(1_usize, n_rep)?;
        let v_full = v.repeat_interleave(1_usize, n_rep)?;

        let scale = 1.0_f64 / (head_dim as f64).sqrt();
        let scores = q_r.matmul(&k_full.transpose()?)?.mul_scalar(scale);
        let scores = scores.broadcast_add(mask)?;
        let probs = scores.softmax_last_dim()?;
        let ctx = probs.matmul(&v_full)?;
        let merged = ctx.merge_heads()?;
        let attn_out = layer.o.apply_linear(&merged, q_dim, cfg.hidden_size)?;
        let attn_post = attn_out.rms_norm_affine(
            std::sync::Arc::clone(&layer.post_self_attn_norm_gain),
            cfg.rms_norm_eps,
        )?;
        let h1 = x.add(&attn_post)?;

        // MLP sublayer: post_mlp_norm(mlp(post_attn_norm(h1))) + h1.
        let h1_in = h1.rms_norm_affine(
            std::sync::Arc::clone(&layer.post_attn_norm_gain),
            cfg.rms_norm_eps,
        )?;
        let gate_up =
            layer
                .gate_up
                .apply_linear(&h1_in, cfg.hidden_size, 2 * cfg.intermediate_size)?;
        let gate = gate_up.slice(2_usize, 0, cfg.intermediate_size)?;
        let up = gate_up.slice(2_usize, cfg.intermediate_size, cfg.intermediate_size)?;
        let act = match cfg.hidden_act {
            Glm4NewActivation::Silu => gate.silu(),
            Glm4NewActivation::Gelu => gate.gelu(),
            Glm4NewActivation::GeluPytorchTanh => gate.gelu_tanh(),
        };
        let inner = act.mul(&up)?;
        let mlp_out = layer
            .down
            .apply_linear(&inner, cfg.intermediate_size, cfg.hidden_size)?;
        let mlp_post = mlp_out.rms_norm_affine(
            std::sync::Arc::clone(&layer.post_mlp_norm_gain),
            cfg.rms_norm_eps,
        )?;
        h1.add(&mlp_post)
    }
}

// ---- HuggingFace safetensors loader ----------------------------------------

impl Glm4NewWeights {
    /// Load GLM-4-new weights from HF safetensors (e.g.
    /// `THUDM/glm-4-9b-chat`). HF naming uses the standard
    /// `model.embed_tokens.weight`, `model.layers.{i}.*`, `model.norm.weight`,
    /// and `lm_head.weight` (when not tied). Per-layer:
    ///   - `model.layers.{i}.input_layernorm.weight`
    ///   - `model.layers.{i}.self_attn.{q,k,v,o}_proj.{weight,bias}`
    ///   - `model.layers.{i}.post_self_attn_layernorm.weight`
    ///   - `model.layers.{i}.post_attention_layernorm.weight`
    ///   - `model.layers.{i}.mlp.gate_up_proj.weight` (fused gate+up,
    ///     out_dim = 2 * intermediate_size)
    ///   - `model.layers.{i}.mlp.down_proj.weight`
    ///   - `model.layers.{i}.post_mlp_layernorm.weight`
    pub fn load_from_mmapped(
        st: &fuel_loaders::safetensors::MmapedSafetensors,
        cfg: &Glm4NewConfig,
    ) -> Result<Self> {
        use fuel_tensor::{lazy::load_tensor_as_f32, lazy::load_transposed_matrix_preserve_dtype};

        let h = cfg.hidden_size;
        let inter = cfg.intermediate_size;
        let head_dim = cfg.head_dim();
        let q_dim = cfg.num_attention_heads * head_dim;
        let kv_dim = cfg.num_key_value_heads * head_dim;

        let token_embedding = Arc::from(load_tensor_as_f32(st, "model.embed_tokens.weight")?);

        let mut layers: Vec<Glm4NewLayerWeights> = Vec::with_capacity(cfg.num_hidden_layers);
        for i in 0..cfg.num_hidden_layers {
            let p = format!("model.layers.{i}");
            let input_norm_gain = Arc::from(load_tensor_as_f32(
                st,
                &format!("{p}.input_layernorm.weight"),
            )?);
            let post_self_attn_norm_gain = Arc::from(load_tensor_as_f32(
                st,
                &format!("{p}.post_self_attn_layernorm.weight"),
            )?);
            let post_attn_norm_gain = Arc::from(load_tensor_as_f32(
                st,
                &format!("{p}.post_attention_layernorm.weight"),
            )?);
            let post_mlp_norm_gain = Arc::from(load_tensor_as_f32(
                st,
                &format!("{p}.post_mlp_layernorm.weight"),
            )?);

            let q = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.self_attn.q_proj.weight"),
                q_dim,
                h,
            )?;
            let k = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.self_attn.k_proj.weight"),
                kv_dim,
                h,
            )?;
            let v = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.self_attn.v_proj.weight"),
                kv_dim,
                h,
            )?;
            let o = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.self_attn.o_proj.weight"),
                h,
                q_dim,
            )?;

            let (q_bias, k_bias, v_bias) = if cfg.attention_bias {
                (
                    Some(Arc::from(load_tensor_as_f32(
                        st,
                        &format!("{p}.self_attn.q_proj.bias"),
                    )?)),
                    Some(Arc::from(load_tensor_as_f32(
                        st,
                        &format!("{p}.self_attn.k_proj.bias"),
                    )?)),
                    Some(Arc::from(load_tensor_as_f32(
                        st,
                        &format!("{p}.self_attn.v_proj.bias"),
                    )?)),
                )
            } else {
                (None, None, None)
            };

            let gate_up = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.mlp.gate_up_proj.weight"),
                2 * inter,
                h,
            )?;
            let down = load_transposed_matrix_preserve_dtype(
                st,
                &format!("{p}.mlp.down_proj.weight"),
                h,
                inter,
            )?;

            layers.push(Glm4NewLayerWeights {
                input_norm_gain,
                post_self_attn_norm_gain,
                post_attn_norm_gain,
                post_mlp_norm_gain,
                q,
                q_bias,
                k,
                k_bias,
                v,
                v_bias,
                o,
                gate_up,
                down,
            });
        }

        let final_norm_gain = Arc::from(load_tensor_as_f32(st, "model.norm.weight")?);
        let lm_head = if cfg.tie_word_embeddings {
            None
        } else {
            Some(load_transposed_matrix_preserve_dtype(
                st,
                "lm_head.weight",
                cfg.vocab_size,
                h,
            )?)
        };

        Ok(Self {
            token_embedding,
            layers,
            final_norm_gain,
            lm_head,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    // ROADMAP item 8 (II). Golden values from zai-org/GLM-4-9B-0414's real
    // config.json — its rope_theta (10000) / max_position_embeddings (32768)
    // differ from the `glm4_9b_chat()` preset (5e6 / 131072), so parse != preset.
    const GLM4_9B_0414_CONFIG_JSON: &str = r#"{
        "architectures": ["Glm4ForCausalLM"],
        "model_type": "glm4",
        "vocab_size": 151552,
        "hidden_size": 4096,
        "intermediate_size": 13696,
        "num_hidden_layers": 40,
        "num_attention_heads": 32,
        "num_key_value_heads": 2,
        "head_dim": 128,
        "partial_rotary_factor": 0.5,
        "attention_bias": true,
        "max_position_embeddings": 32768,
        "rms_norm_eps": 1e-05,
        "rope_theta": 10000.0,
        "tie_word_embeddings": false,
        "hidden_act": "silu"
    }"#;

    #[test]
    fn glm4_new_config_from_hf_json_parses_the_artifact_not_a_preset() {
        let cfg = Glm4NewConfig::from_hf_json_str(GLM4_9B_0414_CONFIG_JSON).unwrap();
        assert_eq!(cfg.hidden_size, 4096);
        assert_eq!(cfg.num_hidden_layers, 40);
        assert_eq!(cfg.num_attention_heads, 32);
        assert_eq!(cfg.vocab_size, 151_552);
        assert_eq!(cfg.intermediate_size, 13_696);
        // GQA: default would be num_attention_heads (32); 2 proves the key was READ.
        assert_eq!(cfg.num_key_value_heads, 2);
        // explicit head_dim (Option) read as Some(128); head_dim() returns 128.
        assert_eq!(cfg.head_dim, Some(128));
        assert_eq!(cfg.head_dim(), 128);
        // GLM-4 ships an explicit partial rotary factor — read faithfully.
        assert_eq!(cfg.partial_rotary_factor, Some(0.5));
        assert_eq!(cfg.rope_theta, 10_000.0);
        assert_eq!(cfg.max_position_embeddings, 32_768);
        assert!(cfg.attention_bias);
        assert_eq!(cfg.hidden_act, Glm4NewActivation::Silu);
        // Sabotage sibling (WEAKER): distinct from the preset (rope_theta/max_pos).
        assert_ne!(cfg, Glm4NewConfig::glm4_9b_chat());
    }

    /// A SECOND distinct config exercising the default path (partial_rotary_factor/
    /// rope_theta/hidden_act omitted → None / 10000 / Silu) AND an explicit,
    /// decoupled head_dim (80 ≠ 1024/16 = 64), which Glm4NewConfig keeps as Some.
    #[test]
    fn glm4_new_config_reads_a_second_distinct_config() {
        let json = r#"{
            "model_type": "glm4",
            "vocab_size": 60000,
            "hidden_size": 1024,
            "intermediate_size": 4096,
            "num_hidden_layers": 8,
            "num_attention_heads": 16,
            "num_key_value_heads": 4,
            "head_dim": 80,
            "attention_bias": true,
            "max_position_embeddings": 8192
        }"#;
        let cfg = Glm4NewConfig::from_hf_json_str(json).unwrap();
        assert_eq!(cfg.num_key_value_heads, 4);
        // explicit head_dim kept as Some, NOT collapsed to the quotient (1024/16=64).
        assert_eq!(cfg.head_dim, Some(80));
        assert_eq!(cfg.head_dim(), 80);
        // omitted → resolve defaults
        assert_eq!(cfg.partial_rotary_factor, None);
        assert_eq!(cfg.rope_theta, 10_000.0);
        assert_eq!(cfg.hidden_act, Glm4NewActivation::Silu);
        assert_ne!(cfg.hidden_size, 4096);
    }

    /// `num_key_value_heads` ABSENT → defaults to `num_attention_heads`; and true
    /// MQA (`1`) survives — both via the shared `hf_config::num_key_value_heads`.
    #[test]
    fn glm4_new_config_gqa_default_and_true_mqa() {
        let base = |kv: &str| {
            format!(
                r#"{{
                "model_type": "glm4",
                "vocab_size": 1000, "hidden_size": 64, "intermediate_size": 128,
                "num_hidden_layers": 2, "num_attention_heads": 8,
                "max_position_embeddings": 128 {kv}
            }}"#
            )
        };
        let absent = Glm4NewConfig::from_hf_json_str(&base("")).unwrap();
        assert_eq!(absent.num_key_value_heads, 8); // absent → num_attention_heads
        let mqa = Glm4NewConfig::from_hf_json_str(&base(", \"num_key_value_heads\": 1")).unwrap();
        assert_eq!(mqa.num_key_value_heads, 1); // TRUE MQA survives
    }

    /// The non-serde activation enum is parsed from the `hidden_act` string:
    /// known values map, an UNKNOWN one ERRORS rather than silently defaulting.
    #[test]
    fn glm4_new_config_maps_and_rejects_hidden_act() {
        let with_act = |act: &str| {
            format!(
                r#"{{
                "model_type": "glm4",
                "vocab_size": 1000, "hidden_size": 64, "intermediate_size": 128,
                "num_hidden_layers": 2, "num_attention_heads": 8,
                "max_position_embeddings": 128, "hidden_act": "{act}"
            }}"#
            )
        };
        assert_eq!(
            Glm4NewConfig::from_hf_json_str(&with_act("gelu"))
                .unwrap()
                .hidden_act,
            Glm4NewActivation::Gelu
        );
        assert_eq!(
            Glm4NewConfig::from_hf_json_str(&with_act("gelu_pytorch_tanh"))
                .unwrap()
                .hidden_act,
            Glm4NewActivation::GeluPytorchTanh
        );
        // Unknown activation is surfaced as an error, not swallowed.
        assert!(Glm4NewConfig::from_hf_json_str(&with_act("mish")).is_err());
    }

    fn vec_of(n: usize, next: &mut dyn FnMut() -> f32) -> Arc<[f32]> {
        Arc::from((0..n).map(|_| next()).collect::<Vec<_>>())
    }

    fn rng_seed(seed: u32) -> impl FnMut() -> f32 {
        let mut s = seed;
        move || {
            s = s.wrapping_mul(1103515245).wrapping_add(12345);
            ((s >> 16) as u16 as f32 / 65535.0 - 0.5) * 0.05
        }
    }

    fn tiny_cfg() -> Glm4NewConfig {
        Glm4NewConfig {
            vocab_size: 32,
            hidden_size: 16,
            intermediate_size: 24,
            num_hidden_layers: 2,
            num_attention_heads: 4,
            num_key_value_heads: 2,
            head_dim: Some(4),
            partial_rotary_factor: Some(0.5),
            attention_bias: true,
            max_position_embeddings: 32,
            sliding_window: None,
            tie_word_embeddings: true,
            rope_theta: 10_000.0,
            rms_norm_eps: 1e-5,
            hidden_act: Glm4NewActivation::Silu,
        }
    }

    fn tiny_weights(cfg: &Glm4NewConfig, seed: u32) -> Glm4NewWeights {
        let mut nb = rng_seed(seed);
        let h = cfg.hidden_size;
        let i = cfg.intermediate_size;
        let head_dim = cfg.head_dim();
        let q_dim = cfg.num_attention_heads * head_dim;
        let kv_dim = cfg.num_key_value_heads * head_dim;
        let token_embedding = vec_of(cfg.vocab_size * h, &mut nb);
        let layers: Vec<Glm4NewLayerWeights> = (0..cfg.num_hidden_layers)
            .map(|_| Glm4NewLayerWeights {
                input_norm_gain: Arc::from(vec![1.0_f32; h]),
                post_self_attn_norm_gain: Arc::from(vec![1.0_f32; h]),
                post_attn_norm_gain: Arc::from(vec![1.0_f32; h]),
                post_mlp_norm_gain: Arc::from(vec![1.0_f32; h]),
                q: WeightStorage::F32(vec_of(h * q_dim, &mut nb)),
                q_bias: if cfg.attention_bias {
                    Some(vec_of(q_dim, &mut nb))
                } else {
                    None
                },
                k: WeightStorage::F32(vec_of(h * kv_dim, &mut nb)),
                k_bias: if cfg.attention_bias {
                    Some(vec_of(kv_dim, &mut nb))
                } else {
                    None
                },
                v: WeightStorage::F32(vec_of(h * kv_dim, &mut nb)),
                v_bias: if cfg.attention_bias {
                    Some(vec_of(kv_dim, &mut nb))
                } else {
                    None
                },
                o: WeightStorage::F32(vec_of(q_dim * h, &mut nb)),
                gate_up: WeightStorage::F32(vec_of(h * 2 * i, &mut nb)),
                down: WeightStorage::F32(vec_of(i * h, &mut nb)),
            })
            .collect();
        let lm_head = if cfg.tie_word_embeddings {
            None
        } else {
            Some(WeightStorage::F32(vec_of(h * cfg.vocab_size, &mut nb)))
        };
        Glm4NewWeights {
            token_embedding,
            layers,
            final_norm_gain: Arc::from(vec![1.0_f32; h]),
            lm_head,
        }
    }

    #[test]
    fn forward_shape_and_finite() {
        let cfg = tiny_cfg();
        let model = Glm4NewModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg, 11),
        };
        let tokens = [1_u32, 2, 3, 4];
        let out = model.forward(&tokens, 0).unwrap();
        assert_eq!(out.shape().dims(), &[1, tokens.len(), cfg.vocab_size]);
        for &v in &out.realize_f32() {
            assert!(v.is_finite(), "non-finite logit: {v}");
        }
    }

    /// 4-norm-per-layer sublayer structure is wired. Zeroing
    /// the post_self_attn_norm gain must alter the output —
    /// confirms that path runs.
    #[test]
    fn post_self_attn_norm_is_wired() {
        let cfg = tiny_cfg();
        let mut base = tiny_weights(&cfg, 22);
        let mut modified = base.clone();
        let h = cfg.hidden_size;
        modified.layers[0].post_self_attn_norm_gain = Arc::from(vec![0.0_f32; h]);
        // Re-Arc base to break any potential aliasing.
        base.layers[0].post_self_attn_norm_gain =
            Arc::from(base.layers[0].post_self_attn_norm_gain.to_vec());

        let m_a = Glm4NewModel {
            config: cfg.clone(),
            weights: base,
        };
        let m_b = Glm4NewModel {
            config: cfg,
            weights: modified,
        };
        let tokens = [1_u32, 2, 3, 4];
        let a = m_a.forward(&tokens, 0).unwrap().realize_f32();
        let b = m_b.forward(&tokens, 0).unwrap().realize_f32();
        let mut max_diff = 0.0_f32;
        for (x, y) in a.iter().zip(b.iter()) {
            max_diff = max_diff.max((x - y).abs());
        }
        assert!(
            max_diff > 1e-7,
            "zeroing post_self_attn_norm_gain must alter output, max_diff = {max_diff}"
        );
    }

    /// Fused gate_up split is wired correctly: zeroing the
    /// gate half (first `intermediate_size` columns) zeroes the
    /// MLP contribution, must change output. If the slice
    /// boundaries were inverted (gate becomes up), the test
    /// would still pass because the MLP would also collapse.
    /// What we're really testing is "fused gate_up is split
    /// at the right boundary and the gate path is active".
    #[test]
    fn fused_gate_up_split() {
        let cfg = tiny_cfg();
        let base = tiny_weights(&cfg, 33);
        let mut modified = base.clone();
        let h = cfg.hidden_size;
        let i = cfg.intermediate_size;
        // Replace the FIRST half of gate_up with zeros (gate path).
        let orig = match &base.layers[0].gate_up {
            WeightStorage::F32(v) => v.to_vec(),
            _ => unreachable!(),
        };
        let mut zeroed = orig.clone();
        // Layout: (hidden, 2 * intermediate) row-major. For each row,
        // zero the first `intermediate` columns.
        for row in 0..h {
            for col in 0..i {
                zeroed[row * (2 * i) + col] = 0.0;
            }
        }
        modified.layers[0].gate_up = WeightStorage::F32(Arc::from(zeroed));

        let m_a = Glm4NewModel {
            config: cfg.clone(),
            weights: base,
        };
        let m_b = Glm4NewModel {
            config: cfg,
            weights: modified,
        };
        let tokens = [1_u32, 2, 3, 4];
        let a = m_a.forward(&tokens, 0).unwrap().realize_f32();
        let b = m_b.forward(&tokens, 0).unwrap().realize_f32();
        let mut max_diff = 0.0_f32;
        for (x, y) in a.iter().zip(b.iter()) {
            max_diff = max_diff.max((x - y).abs());
        }
        assert!(
            max_diff > 1e-7,
            "zeroing gate half of gate_up must alter output, max_diff = {max_diff}"
        );
    }

    /// rotary_dim is `partial_rotary_factor * head_dim`.
    #[test]
    fn partial_rotary_factor() {
        let mut cfg = tiny_cfg();
        cfg.head_dim = Some(8);
        cfg.partial_rotary_factor = Some(0.5);
        assert_eq!(cfg.rotary_dim(), 4);
        cfg.partial_rotary_factor = Some(1.0);
        assert_eq!(cfg.rotary_dim(), 8);
        cfg.partial_rotary_factor = None;
        assert_eq!(cfg.rotary_dim(), 8);
    }

    /// Untied embeddings: lm_head is a distinct WeightStorage.
    #[test]
    fn untied_embeddings_runs() {
        let mut cfg = tiny_cfg();
        cfg.tie_word_embeddings = false;
        let model = Glm4NewModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg, 44),
        };
        let tokens = [1_u32, 2, 3];
        let out = model.forward(&tokens, 0).unwrap();
        assert_eq!(out.shape().dims(), &[1, tokens.len(), cfg.vocab_size]);
        for &v in &out.realize_f32() {
            assert!(v.is_finite());
        }
    }

    #[test]
    fn forward_hidden_shape_and_finite() {
        let cfg = tiny_cfg();
        let model = Glm4NewModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg, 55),
        };
        let tokens = [1_u32, 2, 3, 4];
        let hidden = model.forward_hidden(&tokens, 0).unwrap();
        assert_eq!(hidden.shape().dims(), &[1, tokens.len(), cfg.hidden_size]);
        for &v in &hidden.realize_f32() {
            assert!(v.is_finite(), "non-finite hidden: {v}");
        }
    }

    #[test]
    fn forward_embeds_matches_forward_after_token_lookup() {
        let cfg = tiny_cfg();
        let model = Glm4NewModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg, 55),
        };
        let tokens = [1_u32, 2, 3];
        let logits_ref = model.forward(&tokens, 0).unwrap().realize_f32();
        let anchor =
            Tensor::from_f32(vec![0.0_f32], Shape::from_dims(&[1]), &Device::cpu()).unwrap();
        let embeds = model.embed_tokens_anchored(&anchor, &tokens).unwrap();
        let logits_via_embeds = model.forward_embeds(&embeds, 0).unwrap().realize_f32();
        let max_diff = logits_ref
            .iter()
            .zip(logits_via_embeds.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0_f32, f32::max);
        assert!(
            max_diff < 1e-5,
            "GLM4-new forward vs forward_embeds must agree (max diff {max_diff})"
        );
    }

    #[test]
    fn forward_embeds_rejects_bad_shape() {
        let cfg = tiny_cfg();
        let model = Glm4NewModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg, 55),
        };
        let bad = Tensor::from_f32(
            vec![0.0_f32; 3 * (cfg.hidden_size + 1)],
            Shape::from_dims(&[1, 3, cfg.hidden_size + 1]),
            &Device::cpu(),
        )
        .unwrap();
        assert!(model.forward_embeds(&bad, 0).is_err());
    }

    #[test]
    fn forward_hidden_embeds_matches_forward_hidden() {
        let cfg = tiny_cfg();
        let model = Glm4NewModel {
            config: cfg.clone(),
            weights: tiny_weights(&cfg, 55),
        };
        let tokens = [5_u32, 7];
        let h_ref = model.forward_hidden(&tokens, 0).unwrap().realize_f32();
        let anchor =
            Tensor::from_f32(vec![0.0_f32], Shape::from_dims(&[1]), &Device::cpu()).unwrap();
        let embeds = model.embed_tokens_anchored(&anchor, &tokens).unwrap();
        let h_via_embeds = model
            .forward_hidden_embeds(&embeds, 0)
            .unwrap()
            .realize_f32();
        let max_diff = h_ref
            .iter()
            .zip(h_via_embeds.iter())
            .map(|(a, b)| (a - b).abs())
            .fold(0.0_f32, f32::max);
        assert!(
            max_diff < 1e-5,
            "GLM4-new forward_hidden vs forward_hidden_embeds must agree (max diff {max_diff})"
        );
    }
}