#!/usr/bin/env python3 """Debug E4B inference using torch/safetensors to compare with Swift Metal.""" import torch import numpy as np import json from pathlib import Path from safetensors.torch import safe_open 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 load_tensor(name): """Load tensor from safetensors.""" model_file = MODEL_PATH / "model.safetensors" with safe_open(model_file, framework="pt") as f: tensor = f.get_tensor(name) if tensor.dtype == torch.bfloat16: return tensor.float().numpy() if tensor.dtype == torch.uint32: return tensor.numpy() return tensor.float().numpy() def debug_inference(): print("=" * 50) print("Torch E4B Debug Inference") print("=" * 50) config = load_config() text_config = config["text_config"] hidden_size = text_config['hidden_size'] n_heads = text_config['num_attention_heads'] n_kv_heads = text_config['num_key_value_heads'] head_dim = text_config['head_dim'] print(f"\nConfig: hidden={hidden_size}, layers={text_config['num_hidden_layers']}") print(f" head_dim={head_dim}, n_heads={n_heads}, n_kv_heads={n_kv_heads}") # Load embedding print("\n--- Loading weights ---") 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}") # Dequantize BOS embedding bos_token = 2 vocab_size, packed_dim = embed_weight.shape num_groups = hidden_size // 64 print(f"\n--- Dequantizing BOS (token={bos_token}) ---") w_row = embed_weight[bos_token] s_row = embed_scales[bos_token] b_row = embed_biases[bos_token] embedding = np.zeros(hidden_size, dtype=np.float32) for g in range(num_groups): scale = s_row[g] bias = 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) embedding[g * 64 + j] = qval * scale + bias print(f"Embedding: min={embedding.min():.6f}, max={embedding.max():.6f}") print(f" first 20: {embedding[:20].tolist()}") # Scale embed_scale = np.sqrt(float(hidden_size)) embedding_scaled = embedding * embed_scale print(f"\nScaled (x{embed_scale:.2f}): min={embedding_scaled.min():.6f}, max={embedding_scaled.max():.6f}") print(f" first 20: {embedding_scaled[:20].tolist()}") # Layer 0 print("\n--- Layer 0 ---") prefix = "language_model.model.layers.0" input_norm_w = load_tensor(f"{prefix}.input_layernorm.weight") q_norm_w = load_tensor(f"{prefix}.self_attn.q_norm.weight") k_norm_w = load_tensor(f"{prefix}.self_attn.k_norm.weight") q_proj_w = load_tensor(f"{prefix}.self_attn.q_proj.weight") q_proj_s = load_tensor(f"{prefix}.self_attn.q_proj.scales") q_proj_b = load_tensor(f"{prefix}.self_attn.q_proj.biases") k_proj_w = load_tensor(f"{prefix}.self_attn.k_proj.weight") k_proj_s = load_tensor(f"{prefix}.self_attn.k_proj.scales") k_proj_b = load_tensor(f"{prefix}.self_attn.k_proj.biases") v_proj_w = load_tensor(f"{prefix}.self_attn.v_proj.weight") v_proj_s = load_tensor(f"{prefix}.self_attn.v_proj.scales") v_proj_b = load_tensor(f"{prefix}.self_attn.v_proj.biases") o_proj_w = load_tensor(f"{prefix}.self_attn.o_proj.weight") o_proj_s = load_tensor(f"{prefix}.self_attn.o_proj.scales") o_proj_b = load_tensor(f"{prefix}.self_attn.o_proj.biases") print(f"input_norm: {input_norm_w.shape}, q_norm: {q_norm_w.shape}, k_norm: {k_norm_w.shape}") print(f"q_proj: w={q_proj_w.shape}, s={q_proj_s.shape}") # Forward def dequantize_matmul(x, w, s, b): out_dim, packed_dim = w.shape num_groups = packed_dim * 8 // 64 result = np.zeros(out_dim, dtype=np.float32) 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 * x[g * 64 + j] result[out_idx] = sum_val return result hidden = embedding_scaled eps = 1e-6 # Input norm ss = np.sum(hidden * hidden) / hidden_size rms = np.sqrt(ss + eps) hidden_normed = hidden / rms * input_norm_w print(f"\nInput norm: rms={rms:.6f}") print(f" first 5: {hidden_normed[:5].tolist()}") # Q, K, V q = dequantize_matmul(hidden_normed, q_proj_w, q_proj_s, q_proj_b) k = dequantize_matmul(hidden_normed, k_proj_w, k_proj_s, k_proj_b) v = dequantize_matmul(hidden_normed, v_proj_w, v_proj_s, v_proj_b) print(f"\nQ: min={q.min():.6f}, max={q.max():.6f}") print(f" first 5: {q[:5].tolist()}") print(f"\nK: min={k.min():.6f}, max={k.max():.6f}") print(f" first 5: {k[:5].tolist()}") print(f"\nV: min={v.min():.6f}, max={v.max():.6f}") print(f" first 5: {v[:5].tolist()}") # Q, K norms q_reshaped = q.reshape(n_heads, head_dim) q_normed = np.zeros_like(q_reshaped) for h in range(n_heads): h_q = q_reshaped[h] ss = np.sum(h_q * h_q) / head_dim rms = np.sqrt(ss + eps) q_normed[h] = h_q / rms * q_norm_w q_normed = q_normed.flatten() k_reshaped = k.reshape(n_kv_heads, head_dim) k_normed = np.zeros_like(k_reshaped) for h in range(n_kv_heads): h_k = k_reshaped[h] ss = np.sum(h_k * h_k) / head_dim rms = np.sqrt(ss + eps) k_normed[h] = h_k / rms * k_norm_w k_normed = k_normed.flatten() print(f"\nQ normed: min={q_normed.min():.6f}, max={q_normed.max():.6f}") print(f" first 5: {q_normed[:5].tolist()}") print(f"\nK normed: min={k_normed.min():.6f}, max={k_normed.max():.6f}") print(f" first 5: {k_normed[:5].tolist()}") # Attention (pos 0) attn_scale = 1.0 / np.sqrt(float(head_dim)) attn_out = np.zeros(n_heads * head_dim, dtype=np.float32) for h in range(n_heads): kv_head = h % n_kv_heads h_q = q_normed[h * head_dim : (h+1) * head_dim] h_k = k_normed[kv_head * head_dim : (kv_head+1) * head_dim] h_v = v[kv_head * head_dim : (kv_head+1) * head_dim] score = float(np.sum(h_q * h_k) * attn_scale) print(f" Head {h}: score={score:.6f}") attn_out[h * head_dim : (h+1) * head_dim] = h_v print(f"\nAttn out: min={attn_out.min():.6f}, max={attn_out.max():.6f}") print(f" first 5: {attn_out[:5].tolist()}") # O proj o_out = dequantize_matmul(attn_out, o_proj_w, o_proj_s, o_proj_b) print(f"\nO out: min={o_out.min():.6f}, max={o_out.max():.6f}") print(f" first 5: {o_out[:5].tolist()}") # Residual hidden_after = hidden + o_out print(f"\nAfter attn: min={hidden_after.min():.6f}, max={hidden_after.max():.6f}") print(f" first 5: {hidden_after[:5].tolist()}") print("\n" + "=" * 50) print("KEY VALUES:") print("=" * 50) print(f"\nEmbedding scaled: {embedding_scaled[:5].tolist()}") print(f"After input norm: {hidden_normed[:5].tolist()}") print(f"Q projection: {q[:5].tolist()}") print(f"Q normed: {q_normed[:5].tolist()}") print(f"Attention out: {attn_out[:5].tolist()}") print(f"O projection: {o_out[:5].tolist()}") print(f"Hidden after attn: {hidden_after[:5].tolist()}") if __name__ == "__main__": debug_inference()