#!/usr/bin/env python3 """Debug E4B inference with MLX to compare with Swift Metal implementation.""" import mlx.core as mx import mlx.nn as nn import numpy as np import json from pathlib import Path MODEL_PATH = Path("/Users/accusys/MarkBase12B/models/E4B-MarkBase") def load_config(): with open(MODEL_PATH / "config.json") as f: return json.load(f) def debug_inference(): print("=" * 50) print("MLX E4B Debug Inference") print("=" * 50) # Load config config = load_config() text_config = config["text_config"] print(f"\nConfig: hidden={text_config['hidden_size']}, layers={text_config['num_hidden_layers']}") print(f" head_dim={text_config['head_dim']}, n_heads={text_config['num_attention_heads']}") print(f" n_kv_heads={text_config['num_key_value_heads']}") # We need to load the model - but MLX might not have Gemma4 support # Let's just load the safetensors directly and compute manually from safetensors import safe_open model_file = MODEL_PATH / "model.safetensors" def load_tensor(name): with safe_open(model_file, framework="mlx") as f: return f.get_tensor(name) print("\n--- Loading weights ---") # Load embedding embed_weight = load_tensor("language_model.model.embed_tokens.weight") embed_scales = load_tensor("language_model.model.embed_tokens.scales") embed_biases = load_tensor("language_model.model.embed_tokens.biases") print(f"embed_weight: shape={embed_weight.shape}, dtype={embed_weight.dtype}") print(f"embed_scales: shape={embed_scales.shape}, dtype={embed_scales.dtype}") print(f"embed_biases: shape={embed_biases.shape}, dtype={embed_biases.dtype}") # BOS token = 2 bos_token = 2 # Dequantize embedding for BOS token # Weight is U32 packed [vocab, hidden/8], scales/biases are BF16 [vocab, hidden/64] vocab_size, packed_dim = embed_weight.shape hidden_size = packed_dim * 8 num_groups = hidden_size // 64 print(f"\n--- Dequantizing BOS embedding ---") print(f"vocab_size={vocab_size}, hidden_size={hidden_size}, num_groups={num_groups}") # Get row for BOS token w_row = embed_weight[bos_token] # [320] s_row = embed_scales[bos_token].astype(mx.float32) # [40] b_row = embed_biases[bos_token].astype(mx.float32) # [40] # Dequantize embedding = [] for g in range(num_groups): scale = float(s_row[g]) bias = float(b_row[g]) for j in range(64): packed_idx = g * 8 + j // 8 shift = (j % 8) * 4 qval = int((w_row[packed_idx] >> shift) & 0xF) dq = qval * scale + bias embedding.append(dq) embedding = mx.array(embedding) print(f"embedding: shape={embedding.shape}") print(f" min={float(embedding.min()):.6f}, max={float(embedding.max()):.6f}") print(f" mean={float(embedding.mean()):.6f}") print(f" first 20: {[float(x) for x in embedding[:20]]}") # Apply embedding scale embed_scale = mx.sqrt(mx.array(float(hidden_size))) embedding_scaled = embedding * embed_scale print(f"\nAfter scaling (scale={float(embed_scale):.2f}):") print(f" min={float(embedding_scaled.min()):.6f}, max={float(embedding_scaled.max()):.6f}") print(f" first 20: {[float(x) for x in embedding_scaled[:20]]}") # Load layer 0 weights print("\n--- Loading Layer 0 weights ---") prefix = "language_model.model.layers.0" # Input norm input_norm_w = load_tensor(f"{prefix}.input_layernorm.weight").astype(mx.float32) print(f"input_norm: shape={input_norm_w.shape}") # Q, K, V projections q_proj_w = load_tensor(f"{prefix}.self_attn.q_proj.weight") q_proj_s = load_tensor(f"{prefix}.self_attn.q_proj.scales").astype(mx.float32) q_proj_b = load_tensor(f"{prefix}.self_attn.q_proj.biases").astype(mx.float32) k_proj_w = load_tensor(f"{prefix}.self_attn.k_proj.weight") k_proj_s = load_tensor(f"{prefix}.self_attn.k_proj.scales").astype(mx.float32) k_proj_b = load_tensor(f"{prefix}.self_attn.k_proj.biases").astype(mx.float32) v_proj_w = load_tensor(f"{prefix}.self_attn.v_proj.weight") v_proj_s = load_tensor(f"{prefix}.self_attn.v_proj.scales").astype(mx.float32) v_proj_b = load_tensor(f"{prefix}.self_attn.v_proj.biases").astype(mx.float32) print(f"q_proj: w={q_proj_w.shape}, s={q_proj_s.shape}, b={q_proj_b.shape}") print(f"k_proj: w={k_proj_w.shape}, s={k_proj_s.shape}, b={k_proj_b.shape}") print(f"v_proj: w={v_proj_w.shape}, s={v_proj_s.shape}, b={v_proj_b.shape}") # Q, K norms q_norm_w = load_tensor(f"{prefix}.self_attn.q_norm.weight").astype(mx.float32) k_norm_w = load_tensor(f"{prefix}.self_attn.k_norm.weight").astype(mx.float32) print(f"q_norm: shape={q_norm_w.shape}") print(f"k_norm: shape={k_norm_w.shape}") # O projection o_proj_w = load_tensor(f"{prefix}.self_attn.o_proj.weight") o_proj_s = load_tensor(f"{prefix}.self_attn.o_proj.scales").astype(mx.float32) o_proj_b = load_tensor(f"{prefix}.self_attn.o_proj.biases").astype(mx.float32) print(f"o_proj: w={o_proj_w.shape}, s={o_proj_s.shape}, b={o_proj_b.shape}") # --- Forward pass for Layer 0 --- print("\n--- Layer 0 Forward ---") hidden = embedding_scaled # 1. Input RMS norm eps = 1e-6 ss = mx.sum(hidden * hidden) / hidden_size rms = mx.sqrt(ss + eps) hidden_normed = hidden / rms * input_norm_w print(f"After input norm:") print(f" rms={float(rms):.6f}") print(f" min={float(hidden_normed.min()):.6f}, max={float(hidden_normed.max()):.6f}") print(f" first 20: {[float(x) for x in hidden_normed[:20]]}") # 2. Q projection def dequantize_matmul(x, w, s, b): """x: [inDim], w: [outDim, inDim/8], s/b: [outDim, inDim/64]""" out_dim, packed_dim = w.shape in_dim = packed_dim * 8 num_groups = in_dim // 64 result = [] for out_idx in range(out_dim): sum_val = 0.0 for g in range(num_groups): scale = float(s[out_idx, g]) bias_q = float(b[out_idx, g]) for j in range(64): packed_idx = g * 8 + j // 8 shift = (j % 8) * 4 qval = int((w[out_idx, packed_idx] >> shift) & 0xF) dq = qval * scale + bias_q sum_val += dq * float(x[j + g * 64]) result.append(sum_val) return mx.array(result) print("\nComputing Q projection...") q = dequantize_matmul(hidden_normed, q_proj_w, q_proj_s, q_proj_b) print(f"Q: shape={q.shape}") print(f" min={float(q.min()):.6f}, max={float(q.max()):.6f}") print(f" first 20: {[float(x) for x in q[:20]]}") print("\nComputing K projection...") k = dequantize_matmul(hidden_normed, k_proj_w, k_proj_s, k_proj_b) print(f"K: shape={k.shape}") print(f" min={float(k.min()):.6f}, max={float(k.max()):.6f}") print(f" first 20: {[float(x) for x in k[:20]]}") print("\nComputing V projection...") v = dequantize_matmul(hidden_normed, v_proj_w, v_proj_s, v_proj_b) print(f"V: shape={v.shape}") print(f" min={float(v.min()):.6f}, max={float(v.max()):.6f}") print(f" first 20: {[float(x) for x in v[:20]]}") # 3. Q, K norms (per-head RMS norm) n_heads = 8 n_kv_heads = 2 head_dim = 256 print("\nApplying Q norm (per-head)...") q_reshaped = q.reshape(n_heads, head_dim) q_normed = [] for h in range(n_heads): h_q = q_reshaped[h] ss = mx.sum(h_q * h_q) / head_dim rms = mx.sqrt(ss + eps) h_q_normed = h_q / rms * q_norm_w q_normed.append(h_q_normed) q_normed = mx.concatenate(q_normed) print(f"Q normed: shape={q_normed.shape}") print(f" min={float(q_normed.min()):.6f}, max={float(q_normed.max()):.6f}") print(f" first 20: {[float(x) for x in q_normed[:20]]}") print("\nApplying K norm (per-head)...") k_reshaped = k.reshape(n_kv_heads, head_dim) k_normed = [] for h in range(n_kv_heads): h_k = k_reshaped[h] ss = mx.sum(h_k * h_k) / head_dim rms = mx.sqrt(ss + eps) h_k_normed = h_k / rms * k_norm_w k_normed.append(h_k_normed) k_normed = mx.concatenate(k_normed) print(f"K normed: shape={k_normed.shape}") print(f" min={float(k_normed.min()):.6f}, max={float(k_normed.max()):.6f}") print(f" first 20: {[float(x) for x in k_normed[:20]]}") # 4. RoPE (sliding, theta=10000) # For position 0, RoPE doesn't rotate (cos(0)=1, sin(0)=0) # So q_rope = q_normed, k_rope = k_normed q_rope = q_normed k_rope = k_normed print(f"\nRoPE at position 0: no rotation") print(f"Q rope: same as Q normed") print(f"K rope: same as K normed") # 5. Attention at position 0 # Only attend to itself (position 0) # Q @ K.T for each head print("\n--- Attention at position 0 ---") attn_scale = 1.0 / mx.sqrt(mx.array(float(head_dim))) attn_out = [] for h in range(n_heads): kv_head = h % n_kv_heads h_q = q_rope[h * head_dim : (h+1) * head_dim] h_k = k_rope[kv_head * head_dim : (kv_head+1) * head_dim] h_v = v[kv_head * head_dim : (kv_head+1) * head_dim] # Score = Q @ K score = float(mx.sum(h_q * h_k) * attn_scale) print(f" Head {h}: score={score:.6f}") # Softmax: exp(score) / exp(score) = 1 # Output = 1 * V attn_out.append(h_v) attn_out = mx.concatenate(attn_out) print(f"\nAttention output: shape={attn_out.shape}") print(f" min={float(attn_out.min()):.6f}, max={float(attn_out.max()):.6f}") print(f" first 20: {[float(x) for x in attn_out[:20]]}") # 6. O projection print("\nComputing O projection...") o_out = dequantize_matmul(attn_out, o_proj_w, o_proj_s, o_proj_b) print(f"O output: shape={o_out.shape}") print(f" min={float(o_out.min()):.6f}, max={float(o_out.max()):.6f}") print(f" first 20: {[float(x) for x in o_out[:20]]}") # 7. Residual hidden_after_attn = hidden + o_out print(f"\nAfter attention residual:") print(f" min={float(hidden_after_attn.min()):.6f}, max={float(hidden_after_attn.max()):.6f}") print(f" first 20: {[float(x) for x in hidden_after_attn[:20]]}") print("\n" + "=" * 50) print("Debug complete. Compare with Swift output.") print("=" * 50) if __name__ == "__main__": debug_inference()