Initial commit: E4B-MarkBase model integration with passing tests
CI / build-and-test (push) Has been cancelled
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
This commit is contained in:
MarkBase Admin
@@ -0,0 +1,258 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Debug E4B inference using numpy/safetensors to compare with Swift Metal."""
|
||||
|
||||
import numpy as np
|
||||
import json
|
||||
from pathlib import Path
|
||||
from safetensors.numpy 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 bf16_to_f32(bf16_data):
|
||||
"""Convert BF16 uint16 array to float32."""
|
||||
bf16 = bf16_data.view(np.uint16)
|
||||
f32_bits = bf16.astype(np.uint32) << 16
|
||||
return f32_bits.view(np.float32)
|
||||
|
||||
def load_tensor(name):
|
||||
"""Load tensor from safetensors, converting BF16 to F32."""
|
||||
model_file = MODEL_PATH / "model.safetensors"
|
||||
with safe_open(model_file, framework="numpy") as f:
|
||||
data = f.get_tensor(name)
|
||||
# Check if dtype name contains 'bfloat' or is uint16 for scales/biases
|
||||
dtype_str = str(data.dtype)
|
||||
if 'bfloat' in dtype_str.lower() or dtype_str == 'e4m3bfloat16':
|
||||
return bf16_to_f32(data)
|
||||
if data.dtype == np.uint32:
|
||||
return data # Keep uint32 for weights
|
||||
return data.astype(np.float32)
|
||||
|
||||
def debug_inference():
|
||||
print("=" * 50)
|
||||
print("NumPy E4B Debug Inference")
|
||||
print("=" * 50)
|
||||
|
||||
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']}")
|
||||
|
||||
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("\n--- Loading weights ---")
|
||||
|
||||
# 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
|
||||
vocab_size, packed_dim = embed_weight.shape
|
||||
num_groups = hidden_size // 64
|
||||
|
||||
print(f"\n--- Dequantizing BOS embedding (token={bos_token}) ---")
|
||||
|
||||
w_row = embed_weight[bos_token].astype(np.uint32)
|
||||
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: shape={embedding.shape}")
|
||||
print(f" min={embedding.min():.6f}, max={embedding.max():.6f}")
|
||||
print(f" mean={embedding.mean():.6f}")
|
||||
print(f" first 20: {embedding[:20].tolist()}")
|
||||
|
||||
# Embedding scale
|
||||
embed_scale = np.sqrt(float(hidden_size))
|
||||
embedding_scaled = embedding * embed_scale
|
||||
print(f"\nAfter scaling (scale={embed_scale:.2f}):")
|
||||
print(f" min={embedding_scaled.min():.6f}, max={embedding_scaled.max():.6f}")
|
||||
print(f" first 20: {embedding_scaled[:20].tolist()}")
|
||||
|
||||
# Layer 0 weights
|
||||
print("\n--- Layer 0 weights ---")
|
||||
prefix = "language_model.model.layers.0"
|
||||
|
||||
input_norm_w = load_tensor(f"{prefix}.input_layernorm.weight")
|
||||
print(f"input_norm: shape={input_norm_w.shape}, first 5: {input_norm_w[:5].tolist()}")
|
||||
|
||||
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")
|
||||
|
||||
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_norm_w = load_tensor(f"{prefix}.self_attn.q_norm.weight")
|
||||
k_norm_w = load_tensor(f"{prefix}.self_attn.k_norm.weight")
|
||||
print(f"q_norm: shape={q_norm_w.shape}, first 5: {q_norm_w[:5].tolist()}")
|
||||
print(f"k_norm: shape={k_norm_w.shape}, first 5: {k_norm_w[:5].tolist()}")
|
||||
|
||||
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"o_proj: w={o_proj_w.shape}, s={o_proj_s.shape}, b={o_proj_b.shape}")
|
||||
|
||||
# --- Forward ---
|
||||
print("\n--- Forward pass ---")
|
||||
hidden = embedding_scaled
|
||||
|
||||
# 1. Input RMS norm
|
||||
eps = 1e-6
|
||||
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" min={hidden_normed.min():.6f}, max={hidden_normed.max():.6f}")
|
||||
print(f" first 20: {hidden_normed[:20].tolist()}")
|
||||
|
||||
# 2. Q, K, V projections
|
||||
def dequantize_matmul(x, w, s, b):
|
||||
out_dim, packed_dim = w.shape
|
||||
in_dim = packed_dim * 8
|
||||
num_groups = in_dim // 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 = s[out_idx, g]
|
||||
bias_q = 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
|
||||
|
||||
print("\nQ projection...")
|
||||
q = dequantize_matmul(hidden_normed, q_proj_w.astype(np.uint32), q_proj_s, q_proj_b)
|
||||
print(f"Q: shape={q.shape}")
|
||||
print(f" min={q.min():.6f}, max={q.max():.6f}")
|
||||
print(f" first 20: {q[:20].tolist()}")
|
||||
|
||||
print("\nK projection...")
|
||||
k = dequantize_matmul(hidden_normed, k_proj_w.astype(np.uint32), k_proj_s, k_proj_b)
|
||||
print(f"K: shape={k.shape}")
|
||||
print(f" min={k.min():.6f}, max={k.max():.6f}")
|
||||
print(f" first 20: {k[:20].tolist()}")
|
||||
|
||||
print("\nV projection...")
|
||||
v = dequantize_matmul(hidden_normed, v_proj_w.astype(np.uint32), v_proj_s, v_proj_b)
|
||||
print(f"V: shape={v.shape}")
|
||||
print(f" min={v.min():.6f}, max={v.max():.6f}")
|
||||
print(f" first 20: {v[:20].tolist()}")
|
||||
|
||||
# 3. Q, K norms
|
||||
print("\nQ norm (per-head)...")
|
||||
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()
|
||||
print(f"Q normed: shape={q_normed.shape}")
|
||||
print(f" min={q_normed.min():.6f}, max={q_normed.max():.6f}")
|
||||
print(f" first 20: {q_normed[:20].tolist()}")
|
||||
|
||||
print("\nK norm (per-head)...")
|
||||
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"K normed: shape={k_normed.shape}")
|
||||
print(f" min={k_normed.min():.6f}, max={k_normed.max():.6f}")
|
||||
print(f" first 20: {k_normed[:20].tolist()}")
|
||||
|
||||
# 4. Attention (position 0, only attend to self)
|
||||
print("\n--- Attention at position 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}")
|
||||
|
||||
# At position 0, softmax(score) = 1.0 for self-attention
|
||||
attn_out[h * head_dim : (h+1) * head_dim] = h_v
|
||||
|
||||
print(f"\nAttention output: shape={attn_out.shape}")
|
||||
print(f" min={attn_out.min():.6f}, max={attn_out.max():.6f}")
|
||||
print(f" first 20: {attn_out[:20].tolist()}")
|
||||
|
||||
# 5. O projection
|
||||
print("\nO projection...")
|
||||
o_out = dequantize_matmul(attn_out, o_proj_w.astype(np.uint32), o_proj_s, o_proj_b)
|
||||
print(f"O output: shape={o_out.shape}")
|
||||
print(f" min={o_out.min():.6f}, max={o_out.max():.6f}")
|
||||
print(f" first 20: {o_out[:20].tolist()}")
|
||||
|
||||
# 6. Residual
|
||||
hidden_after_attn = hidden + o_out
|
||||
print(f"\nAfter attention residual:")
|
||||
print(f" min={hidden_after_attn.min():.6f}, max={hidden_after_attn.max():.6f}")
|
||||
print(f" first 20: {hidden_after_attn[:20].tolist()}")
|
||||
|
||||
print("\n" + "=" * 50)
|
||||
print("KEY VALUES TO COMPARE WITH SWIFT:")
|
||||
print("=" * 50)
|
||||
print("\n1. Embedding (scaled):")
|
||||
print(f" first 5: {embedding_scaled[:5].tolist()}")
|
||||
print("\n2. After input norm:")
|
||||
print(f" first 5: {hidden_normed[:5].tolist()}")
|
||||
print("\n3. Q after projection:")
|
||||
print(f" first 5: {q[:5].tolist()}")
|
||||
print("\n4. Q after norm:")
|
||||
print(f" first 5: {q_normed[:5].tolist()}")
|
||||
print("\n5. Attention output:")
|
||||
print(f" first 5: {attn_out[:5].tolist()}")
|
||||
print("\n6. O projection output:")
|
||||
print(f" first 5: {o_out[:5].tolist()}")
|
||||
print("\n7. Hidden after attention:")
|
||||
print(f" first 5: {hidden_after_attn[:5].tolist()}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
debug_inference()
|
||||
Reference in New Issue
Block a user