MarkBase Admin
ac75faa0cc
CI / build-and-test (push) Has been cancelled
- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully - All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance) - All stress tests passed (5/5 in 127.6s) - Concurrent inference - Memory stress (67.5 tok/s, 0 NaN) - Continuous generation - Batch processing - Long-running stability - Swift Metal inference engine with multimodal support
271 lines
10 KiB
Python
271 lines
10 KiB
Python
#!/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() |