auto: pre-test commit
This commit is contained in:
124
tests/unit/test_layer_comparison.py
Normal file
124
tests/unit/test_layer_comparison.py
Normal file
@@ -0,0 +1,124 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Layer-by-layer comparison: production kernel vs PyTorch reference.
|
||||
|
||||
This test loads both pipelines, runs the same input, and compares
|
||||
hidden states after each layer to find where the residual diverges.
|
||||
"""
|
||||
import os, sys, json, time, math, torch, torch.nn.functional as F
|
||||
from pathlib import Path
|
||||
|
||||
CHECKPOINT_DIR = os.environ.get("CHECKPOINT_DIR", "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4")
|
||||
DEVICE = "cuda:0"
|
||||
|
||||
def main():
|
||||
torch.manual_seed(42)
|
||||
|
||||
# Load config
|
||||
with open(os.path.join(CHECKPOINT_DIR, "config.json")) as f:
|
||||
cfg = json.load(f)
|
||||
n_layers = cfg["num_hidden_layers"]
|
||||
H = cfg["hidden_size"]
|
||||
hd = cfg["head_dim"]
|
||||
n_hc = cfg.get("n_hc", 4)
|
||||
print(f"Model: {n_layers} layers, {H} hidden, {hd} head_dim, {n_hc} mHC streams")
|
||||
|
||||
# --- Load production pipeline ---
|
||||
print("\nLoading production pipeline...")
|
||||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||||
from single_shot_inference import DSV4Model
|
||||
prod_model = DSV4Model(CHECKPOINT_DIR, device=DEVICE)
|
||||
print("Production pipeline loaded.")
|
||||
|
||||
# --- Load PyTorch reference pipeline ---
|
||||
print("\nLoading PyTorch reference pipeline...")
|
||||
from single_shot_PYTORCH_REFERENCE import mHCBlock, load_weights, forward_layer, rmsnorm
|
||||
all_w = load_weights(CHECKPOINT_DIR)
|
||||
print("Reference pipeline loaded.")
|
||||
|
||||
# --- Same input for both ---
|
||||
# Use the DeepSeek prompt
|
||||
from transformers import AutoTokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR, trust_remote_code=True)
|
||||
prompt = "The capital of France is"
|
||||
ids = tokenizer.encode(prompt, add_special_tokens=False)
|
||||
# Add chat template
|
||||
user_token = 128803
|
||||
asst_token = 128804
|
||||
chat_ids = [user_token] + ids + [asst_token]
|
||||
print(f"Input: {len(chat_ids)} tokens: {chat_ids}")
|
||||
|
||||
# --- Run production pipeline: prefill ---
|
||||
print("\n=== Production Pipeline: Prefill ===")
|
||||
prod_model.kv_cache.reset()
|
||||
prod_X = None
|
||||
prod_layer_states = [] # (X_l, X_mid, X_next) per layer
|
||||
|
||||
# Process tokens one at a time (decode style)
|
||||
for ti, tid in enumerate(chat_ids):
|
||||
token_id = torch.tensor([[tid]], dtype=torch.int32, device=DEVICE)
|
||||
if ti == len(chat_ids) - 1:
|
||||
# Save layer states for the last token
|
||||
# We need to modify the production pipeline to capture per-layer states
|
||||
# For now, just run and capture the final output
|
||||
pass
|
||||
prod_model.decode_step(token_id, position_offset=ti)
|
||||
|
||||
print("Production prefill done.")
|
||||
|
||||
# --- Run reference pipeline: prefill ---
|
||||
print("\n=== Reference Pipeline: Prefill ===")
|
||||
# Initialize mHC state
|
||||
emb_w = all_w.get("model.embed_tokens.weight")
|
||||
emb_ref = torch.nn.Embedding(emb_w.shape[0], emb_w.shape[1])
|
||||
emb_ref.weight.data = emb_w.bfloat16().to(DEVICE)
|
||||
|
||||
ref_X = mHCBlock.init_state(emb_ref(torch.tensor(chat_ids, device=DEVICE)), n_hc=n_hc)
|
||||
|
||||
# Build mHC blocks and norms for reference
|
||||
attn_mhcs, ffn_mhcs = [], []
|
||||
attn_norms, ffn_norms = [], []
|
||||
for li in range(n_layers):
|
||||
a_mhc = mHCBlock(H, n_hc, device=DEVICE)
|
||||
a_mhc.load(all_w[f"model.layers.{li}.attn_hc.fn"],
|
||||
all_w[f"model.layers.{li}.attn_hc.base"],
|
||||
all_w[f"model.layers.{li}.attn_hc.scale"])
|
||||
attn_mhcs.append(a_mhc)
|
||||
|
||||
f_mhc = mHCBlock(H, n_hc, device=DEVICE)
|
||||
f_mhc.load(all_w[f"model.layers.{li}.ffn_hc.fn"],
|
||||
all_w[f"model.layers.{li}.ffn_hc.base"],
|
||||
all_w[f"model.layers.{li}.ffn_hc.scale"])
|
||||
ffn_mhcs.append(f_mhc)
|
||||
|
||||
attn_norms.append(all_w[f"model.layers.{li}.input_layernorm.weight"].bfloat16().to(DEVICE))
|
||||
ffn_norms.append(all_w[f"model.layers.{li}.post_attention_layernorm.weight"].bfloat16().to(DEVICE))
|
||||
|
||||
# Run reference layer by layer
|
||||
print("Running reference layer by layer...")
|
||||
ref_kv_cache = {}
|
||||
for li in range(n_layers):
|
||||
w = all_w
|
||||
X_before = ref_X.clone()
|
||||
ref_X = forward_layer(ref_X, w, li, cfg, None, None,
|
||||
attn_mhcs[li], ffn_mhcs[li],
|
||||
attn_norms[li], ffn_norms[li],
|
||||
ref_kv_cache, torch.arange(len(chat_ids), device=DEVICE),
|
||||
0)
|
||||
x_max = ref_X.abs().max().item()
|
||||
if li % 10 == 0 or li >= 55:
|
||||
print(f" Ref L{li}: |X|={x_max:.1f}")
|
||||
|
||||
print("Reference prefill done.")
|
||||
print(f" Final |X|: {ref_X.abs().max().item():.1f}")
|
||||
|
||||
# Compare
|
||||
# We can't easily compare per-layer because the production pipeline
|
||||
# doesn't expose intermediate states. But we can compare the final
|
||||
# hidden state and the decoded token.
|
||||
|
||||
print("\n=== Summary ===")
|
||||
print(f"Production final |X|: N/A (need to instrument)")
|
||||
print(f"Reference final |X|: {ref_X.abs().max().item():.1f}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
169
tests/unit/test_mhc_comparison.py
Normal file
169
tests/unit/test_mhc_comparison.py
Normal file
@@ -0,0 +1,169 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Focused comparison: production MoE vs PyTorch reference MoE at specific layers.
|
||||
|
||||
This test:
|
||||
1. Loads both pipelines
|
||||
2. Processes the same input token through 1 layer
|
||||
3. Compares F_attn and F_ffn magnitudes between production and reference
|
||||
4. Identifies where the magnitude diverges
|
||||
"""
|
||||
import os, sys, json, time, math, torch, torch.nn.functional as F
|
||||
from pathlib import Path
|
||||
|
||||
CHECKPOINT_DIR = os.environ.get("CHECKPOINT_DIR", "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4")
|
||||
DEVICE = "cuda:0"
|
||||
HC_EPS = 1e-6
|
||||
|
||||
def sinkhorn_knopp(logits, t_max=20, eps=HC_EPS):
|
||||
M = torch.softmax(logits, -1) + eps
|
||||
M = M / (M.sum(-2, keepdim=True) + eps)
|
||||
for _ in range(t_max - 1):
|
||||
M = M / (M.sum(-1, keepdim=True) + eps)
|
||||
M = M / (M.sum(-2, keepdim=True) + eps)
|
||||
return M
|
||||
|
||||
def unweighted_rmsnorm(x, eps=1e-6):
|
||||
x_f = x.float()
|
||||
rms = x_f.pow(2).mean(-1, keepdim=True).add(eps).rsqrt()
|
||||
return (x_f * rms).to(x.dtype)
|
||||
|
||||
def rmsnorm(x, w, eps=1e-6):
|
||||
x_f = x.float()
|
||||
rms = x_f.pow(2).mean(-1, keepdim=True).add(eps).rsqrt()
|
||||
return (x_f * rms * w.float()).to(x.dtype)
|
||||
|
||||
FP4_LUT = torch.tensor([0., 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0])
|
||||
|
||||
def dequant_nvfp4(weight, weight_scale, weight_scale_2=None, input_scale=None):
|
||||
O, I2 = weight.shape; I = I2 * 2
|
||||
lo = (weight & 0x0F).to(torch.int8); hi = (weight >> 4).to(torch.int8)
|
||||
lut = FP4_LUT.to(device=weight.device, dtype=torch.float32)
|
||||
lo_f = lut[(lo & 0x07).long()] * torch.where((lo >> 3).bool(), -1., 1.)
|
||||
hi_f = lut[(hi & 0x07).long()] * torch.where((hi >> 3).bool(), -1., 1.)
|
||||
w = torch.stack([lo_f, hi_f], -1).reshape(O, I)
|
||||
s = weight_scale.float().repeat_interleave(16, 1)
|
||||
if weight_scale_2 is not None: s = s * weight_scale_2.float()
|
||||
return (w * s).bfloat16()
|
||||
|
||||
def main():
|
||||
torch.manual_seed(42)
|
||||
|
||||
with open(os.path.join(CHECKPOINT_DIR, "config.json")) as f:
|
||||
cfg = json.load(f)
|
||||
H = cfg["hidden_size"]
|
||||
n_hc = cfg.get("n_hc", 4)
|
||||
n_layers = cfg["num_hidden_layers"]
|
||||
n_experts = cfg["n_routed_experts"]
|
||||
top_k = cfg.get("num_experts_per_tok", 6)
|
||||
intermediate = cfg.get("intermediate_size", 18432)
|
||||
print(f"Model: {n_layers} layers, {H} hidden, {n_experts} experts, top-{top_k}")
|
||||
|
||||
# Load weights
|
||||
print("Loading weights...")
|
||||
from safetensors.torch import load_file
|
||||
cdir = Path(CHECKPOINT_DIR); wmap = {}
|
||||
idx = cdir / "model.safetensors.index.json"
|
||||
if idx.exists():
|
||||
with open(idx) as f: wmap = json.load(f).get("weight_map", {})
|
||||
shards = set(wmap.values()) if wmap else set(); all_w = {}
|
||||
for sn in sorted(shards):
|
||||
if (cdir / sn).exists(): all_w.update(load_file(str(cdir / sn)))
|
||||
print(f"Loaded {len(all_w)} tensors")
|
||||
|
||||
# Create a realistic hidden state (simulate running through a few layers)
|
||||
# Use token embedding + a few layers of mHC
|
||||
from single_shot_PYTORCH_REFERENCE import mHCBlock, load_weights as ref_load_weights, forward_layer
|
||||
ref_all_w = ref_load_weights(CHECKPOINT_DIR)
|
||||
|
||||
# Build mHC blocks for first 3 layers
|
||||
attn_mhcs, ffn_mhcs = [], []
|
||||
attn_norms, ffn_norms = [], []
|
||||
for li in range(min(5, n_layers)):
|
||||
a_mhc = mHCBlock(H, n_hc, device=DEVICE)
|
||||
a_mhc.load(ref_all_w[f"model.layers.{li}.attn_hc.fn"],
|
||||
ref_all_w[f"model.layers.{li}.attn_hc.base"],
|
||||
ref_all_w[f"model.layers.{li}.attn_hc.scale"])
|
||||
attn_mhcs.append(a_mhc)
|
||||
f_mhc = mHCBlock(H, n_hc, device=DEVICE)
|
||||
f_mhc.load(ref_all_w[f"model.layers.{li}.ffn_hc.fn"],
|
||||
ref_all_w[f"model.layers.{li}.ffn_hc.base"],
|
||||
ref_all_w[f"model.layers.{li}.ffn_hc.scale"])
|
||||
ffn_mhcs.append(f_mhc)
|
||||
attn_norms.append(ref_all_w[f"model.layers.{li}.input_layernorm.weight"].bfloat16().to(DEVICE))
|
||||
ffn_norms.append(ref_all_w[f"model.layers.{li}.post_attention_layernorm.weight"].bfloat16().to(DEVICE))
|
||||
|
||||
# Process one token through first 3 layers to get a realistic X state
|
||||
emb_w = ref_all_w["model.embed_tokens.weight"]
|
||||
emb = torch.nn.Embedding(emb_w.shape[0], emb_w.shape[1])
|
||||
emb.weight.data = emb_w.bfloat16().to(DEVICE)
|
||||
|
||||
# "The" token
|
||||
tid = 455
|
||||
X = mHCBlock.init_state(emb(torch.tensor([tid], device=DEVICE)), n_hc=n_hc)
|
||||
print(f"\nInitial |X| = {X.abs().max().item():.2f}")
|
||||
|
||||
# Run through first 3 layers using reference
|
||||
kv_cache = {}
|
||||
for li in range(3):
|
||||
X = forward_layer(X, ref_all_w, li, cfg, None, None,
|
||||
attn_mhcs[li], ffn_mhcs[li],
|
||||
attn_norms[li], ffn_norms[li],
|
||||
kv_cache, torch.tensor([3], device=DEVICE),
|
||||
tid)
|
||||
print(f" Ref L{li}: |X| = {X.abs().max().item():.2f}")
|
||||
|
||||
# Now X is a realistic hidden state after 3 layers
|
||||
# Save it for both production and reference comparison
|
||||
X_ref = X.clone()
|
||||
X_prod = X.clone()
|
||||
print(f"\nAfter 3 layers: |X| = {X_ref.abs().max().item():.2f}")
|
||||
|
||||
# --- Compare mHC at L3 ---
|
||||
li = 3
|
||||
print(f"\n=== Comparing mHC at L{li} ===")
|
||||
|
||||
# Reference mHC
|
||||
a_mhc = attn_mhcs[3] # Already loaded
|
||||
x_in_ref, ctx_ref = a_mhc.pre_block(X_ref)
|
||||
print(f" Ref x_in: |x| = {x_in_ref.abs().max().item():.4f}")
|
||||
print(f" Ref A: {ctx_ref['A'][0].tolist()}")
|
||||
print(f" Ref C: {ctx_ref['C'][0].tolist()}")
|
||||
print(f" Ref B row_sums: {ctx_ref['B'][0].sum(-1).tolist()}")
|
||||
|
||||
# Production mHC
|
||||
from dsv4.layers.mhc import mHCLayer
|
||||
prod_mhc = mHCLayer(hidden_dim=H, n_hc=n_hc, device=DEVICE)
|
||||
# Load weights
|
||||
fn = ref_all_w[f"model.layers.{li}.attn_hc.fn"].to(DEVICE, torch.float32)
|
||||
base = ref_all_w[f"model.layers.{li}.attn_hc.base"].to(DEVICE)
|
||||
scale = ref_all_w[f"model.layers.{li}.attn_hc.scale"].to(DEVICE)
|
||||
n = n_hc
|
||||
prod_mhc.load_weights(
|
||||
W_pre=fn[0:n], W_post=fn[n:2*n], W_comb=fn[2*n:],
|
||||
S_pre=base[0:n].reshape(1, n), S_post=base[n:2*n].reshape(n, 1),
|
||||
S_comb=base[2*n:].reshape(n, n),
|
||||
alpha_pre=scale[0].item(), alpha_post=scale[1].item(), alpha_comb=scale[2].item()
|
||||
)
|
||||
x_in_prod, ctx_prod = prod_mhc.pre_block(X_prod)
|
||||
print(f" Prod x_in: |x| = {x_in_prod.abs().max().item():.4f}")
|
||||
A_prod = ctx_prod.A_l
|
||||
C_prod = ctx_prod.C_l
|
||||
B_prod = ctx_prod.B_l
|
||||
print(f" Prod A: {A_prod[0].tolist()}")
|
||||
print(f" Prod C: {C_prod[0].tolist()}")
|
||||
print(f" Prod B row_sums: {B_prod[0].sum(-1).tolist()}")
|
||||
|
||||
# Compare
|
||||
cos_xin = F.cosine_similarity(x_in_ref.flatten().float(), x_in_prod.flatten().float(), dim=0).item()
|
||||
cos_A = F.cosine_similarity(ctx_ref['A'].flatten().float(), A_prod.flatten().float(), dim=0).item()
|
||||
cos_C = F.cosine_similarity(ctx_ref['C'].flatten().float(), C_prod.flatten().float(), dim=0).item()
|
||||
cos_B = F.cosine_similarity(ctx_ref['B'].flatten().float(), B_prod.flatten().float(), dim=0).item()
|
||||
print(f"\n cos(x_in): {cos_xin:.6f}")
|
||||
print(f" cos(A): {cos_A:.6f}")
|
||||
print(f" cos(C): {cos_C:.6f}")
|
||||
print(f" cos(B): {cos_B:.6f}")
|
||||
|
||||
print("\nDone.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user