Add system prompt, CLI args, inverse RoPE flag, minimal e2e test
- System prompt added via chat template (reasoning model needs instructions) - MAX_NEW_TOKENS=512 (reasoning chain-of-thought needs more tokens) - --no-inverse-rope flag to test without inverse RoPE on attn output - --skip-moe flag to debug with shared expert only - --max-tokens and --prompt CLI overrides - minimal_e2e_test(): processes 'The' through full model, checks logits, tracks per-layer residual stream evolution, reports NaN/Inf/spread - INVERSE_ROPE doc: explains partial RoPE only affects last 64/512 dims, first 448 always un-RoPE'd, relative encoding may be intentional
This commit is contained in:
@@ -36,7 +36,7 @@ Usage (on B200):
|
||||
cd /root/dsv4-nvfp4-workspace/kernel
|
||||
python3 single_shot_inference.py
|
||||
"""
|
||||
import os, sys, time, json, math
|
||||
import os, sys, time, json, math, argparse
|
||||
import torch
|
||||
from pathlib import Path
|
||||
|
||||
@@ -44,11 +44,27 @@ from pathlib import Path
|
||||
# Configuration
|
||||
# =====================================================================
|
||||
|
||||
def parse_args():
|
||||
p = argparse.ArgumentParser(description='DSV4 Single-Shot Inference')
|
||||
p.add_argument('--no-inverse-rope', action='store_true', help='Skip inverse RoPE on attention output')
|
||||
p.add_argument('--skip-moe', action='store_true', help='Only use shared expert (skip routed)')
|
||||
p.add_argument('--max-tokens', type=int, default=512, help='Max new tokens to generate')
|
||||
p.add_argument('--prompt', type=str, default=None, help='Override prompt')
|
||||
return p.parse_args()
|
||||
|
||||
_args = parse_args()
|
||||
|
||||
CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
|
||||
MAX_NEW_TOKENS = 10
|
||||
PROMPT = "The capital of France is"
|
||||
MAX_NEW_TOKENS = _args.max_tokens
|
||||
SYSTEM_PROMPT = "You are a helpful, harmless, and honest AI assistant. Answer the user's questions accurately and concisely. If you're unsure about something, say so rather than guessing. Follow the user's instructions carefully and ask for clarification when needed. Always respond in the same language the user is writing in."
|
||||
PROMPT = _args.prompt or "The capital of France is"
|
||||
NUM_GPUS = 8
|
||||
SKIP_ROUTED_MOE = False # If True, only use shared expert (debug)
|
||||
SKIP_ROUTED_MOE = _args.skip_moe # If True, only use shared expert (debug)
|
||||
INVERSE_ROPE = not _args.no_inverse_rope # If False, skip inverse RoPE on attention output (diagnostic)
|
||||
# When True: applies inverse RoPE at query position → converts absolute→relative
|
||||
# When False: leaves relative position encoding intact for output projection
|
||||
# DSV4 partial RoPE only affects last 64/512 dims; first 448 are always un-RoPE'd
|
||||
print(f"Config: INVERSE_ROPE={INVERSE_ROPE}, SKIP_ROUTED_MOE={SKIP_ROUTED_MOE}, MAX_NEW_TOKENS={MAX_NEW_TOKENS}")
|
||||
|
||||
# =====================================================================
|
||||
# NVFP4 dequantization — matches checkpoint format exactly
|
||||
@@ -435,7 +451,15 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
|
||||
|
||||
|
||||
# -- Inverse RoPE on attention output (paper §2.3.3) --
|
||||
attn_out = apply_inverse_rope(attn_out, positions_dev, rope_cos, rope_sin, hd, rd)
|
||||
# DSV4 uses K=V in MQA; both get RoPE'd. Inverse RoPE on the output
|
||||
# at query position q converts: R(q)⁻¹ Σ softmax(R(q)Q·R(p)K) R(p)V
|
||||
# For single KV entry at p: R(p-q)V (relative position encoding)
|
||||
# This only affects the last 64 dims (partial RoPE); first 448 unchanged.
|
||||
# The relative encoding in those 64 dims may be INTENTIONAL — the
|
||||
# output projection can use it for position-dependent computation.
|
||||
# Test both modes via INVERSE_ROPE flag.
|
||||
if INVERSE_ROPE:
|
||||
attn_out = apply_inverse_rope(attn_out, positions_dev, rope_cos, rope_sin, hd, rd)
|
||||
|
||||
# -- Output projection: wo_a (grouped BMM) + wo_b (NVFP4) --
|
||||
# wo_a: grouped linear, (n_h, hd) → (n_groups, o_rank) via BMM
|
||||
@@ -711,12 +735,39 @@ def main():
|
||||
t_compiled = time.time()
|
||||
print(f"Compile: {t_compiled - t_loaded:.1f}s")
|
||||
|
||||
# ==== Phase 3: Inference ====
|
||||
print(f"\n{'='*70}\nPhase 3: Inference\n{'='*70}")
|
||||
# ==== Phase 2.5: Minimal E2E test ====
|
||||
print(f"\n{'='*70}\nPhase 2.5: Minimal E2E Test (single token 'The')\n{'='*70}")
|
||||
from transformers import AutoTokenizer
|
||||
tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR)
|
||||
input_ids = tokenizer.encode(PROMPT, return_tensors="pt").cuda()
|
||||
print(f"Prompt: '{PROMPT}' → {input_ids.tolist()}")
|
||||
minimal_e2e_test(all_weights, cfg, rope_caches, attn_mhc_blocks,
|
||||
ffn_mhc_blocks, attn_norms, ffn_norms, embed, lm_w,
|
||||
final_norm_w, tokenizer)
|
||||
|
||||
# ==== Phase 3: Inference ====
|
||||
print(f"\n{'='*70}\nPhase 3: Inference\n{'='*70}")
|
||||
# Apply chat template with system prompt
|
||||
messages = [
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{"role": "user", "content": PROMPT},
|
||||
]
|
||||
if hasattr(tokenizer, 'apply_chat_template') and tokenizer.chat_template is not None:
|
||||
input_ids = tokenizer.apply_chat_template(
|
||||
messages, return_tensors="pt", add_generation_prompt=True
|
||||
).cuda()
|
||||
# Find where the user prompt starts for display
|
||||
user_only_ids = tokenizer.encode(PROMPT, return_tensors="pt")
|
||||
print(f"Chat template applied. Input: {input_ids.shape[1]} tokens")
|
||||
else:
|
||||
# Fallback: prepend system prompt manually
|
||||
sys_ids = tokenizer.encode(SYSTEM_PROMPT, return_tensors="pt")[0]
|
||||
user_ids = tokenizer.encode(PROMPT, return_tensors="pt")[0]
|
||||
# Add BOS + system + newline + user
|
||||
all_ids = [tokenizer.bos_token_id] if tokenizer.bos_token_id else []
|
||||
all_ids += sys_ids.tolist() + user_ids.tolist()
|
||||
input_ids = torch.tensor([all_ids], dtype=torch.long).cuda()
|
||||
print(f"Manual prompt assembly. Input: {input_ids.shape[1]} tokens")
|
||||
print(f"Prompt: '{PROMPT}' → {input_ids.tolist()[:20]}...")
|
||||
print(f"Decoded: '{tokenizer.decode(input_ids[0][:50])}'")
|
||||
|
||||
generated = input_ids[0].tolist()
|
||||
|
||||
@@ -832,5 +883,131 @@ def main():
|
||||
print(f"{'='*70}")
|
||||
|
||||
|
||||
# =====================================================================
|
||||
# Minimal end-to-end test — single token "The" through the model
|
||||
# =====================================================================
|
||||
|
||||
def minimal_e2e_test(all_weights, cfg, rope_caches, attn_mhc_blocks,
|
||||
ffn_mhc_blocks, attn_norms, ffn_norms, embed, lm_w,
|
||||
final_norm_w, tokenizer):
|
||||
"""Process a single token 'The' through the model and check output logits.
|
||||
|
||||
This is a focused diagnostic: if the model can't even produce reasonable
|
||||
logits for a single token, something is fundamentally wrong in the
|
||||
pipeline. We check:
|
||||
1. No NaN/Inf in any layer output
|
||||
2. Residual stream magnitude stays bounded
|
||||
3. Top-5 logits are sensible (not all Chinese tokens for English)
|
||||
4. Logit spread (max - min) is > 1.0 (not uniform)
|
||||
"""
|
||||
n_layers = cfg["num_hidden_layers"]
|
||||
H = cfg["hidden_size"]
|
||||
n_h = cfg["num_attention_heads"]
|
||||
hd = cfg["head_dim"]
|
||||
rd = cfg.get("qk_rope_head_dim", cfg.get("rope_dim", 64))
|
||||
n_hc = 4
|
||||
|
||||
# Tokenize just "The"
|
||||
tid = torch.tensor(tokenizer.encode("The"), dtype=torch.long, device='cuda:0')
|
||||
if tid.numel() > 1:
|
||||
# If tokenizer adds BOS, take last token
|
||||
print(f" Note: 'The' tokenized to {tid.numel()} tokens, using last one")
|
||||
tid = tid[-1:]
|
||||
print(f" Token ID: {tid.item()} = '{tokenizer.decode(tid.tolist())}'")
|
||||
|
||||
# Setup
|
||||
positions = torch.tensor([0], dtype=torch.long, device='cuda:0')
|
||||
emb = embed(tid) # (1, H)
|
||||
X = mHCBlock.init_state(emb, n_hc) # (1, n_hc, H)
|
||||
|
||||
# Track per-layer diagnostics
|
||||
layer_diags = []
|
||||
|
||||
for li in range(n_layers):
|
||||
gpu = li % NUM_GPUS
|
||||
dev = f"cuda:{gpu}"
|
||||
if X.device != torch.device(dev):
|
||||
X = X.to(dev)
|
||||
torch.cuda.set_device(gpu)
|
||||
|
||||
w = get_layer_weights(all_weights, li, dev)
|
||||
attn_mhc = attn_mhc_blocks.get(li)
|
||||
ffn_mhc = ffn_mhc_blocks.get(li)
|
||||
a_norm = attn_norms[li]
|
||||
f_norm = ffn_norms[li]
|
||||
rc, rs = rope_caches[gpu]
|
||||
kv_cache = SimpleKVCache(head_dim=hd, max_seq=8192, device=dev)
|
||||
|
||||
X = forward_layer(X, w, li, cfg, rc, rs,
|
||||
attn_mhc, ffn_mhc, a_norm, f_norm,
|
||||
kv_cache, tid, positions)
|
||||
|
||||
# Per-layer diagnostic
|
||||
x_max = X.abs().max().item()
|
||||
has_nan = torch.isnan(X.float()).any().item()
|
||||
has_inf = torch.isinf(X.float()).any().item()
|
||||
# Stream 0 (primary)
|
||||
x0 = X[:, 0, :]
|
||||
x0_mean = x0.float().abs().mean().item()
|
||||
x0_std = x0.float().std().item()
|
||||
layer_diags.append({
|
||||
'layer': li, 'gpu': gpu, 'x_max': x_max,
|
||||
'x0_mean': x0_mean, 'x0_std': x0_std,
|
||||
'nan': has_nan, 'inf': has_inf
|
||||
})
|
||||
|
||||
del w
|
||||
|
||||
if has_nan or has_inf:
|
||||
print(f" ❌ Layer {li}: NaN={has_nan} Inf={has_inf} — STOPPING")
|
||||
break
|
||||
|
||||
X = X.to('cuda:0')
|
||||
torch.cuda.set_device(0)
|
||||
|
||||
# Final norm + lm_head
|
||||
x_out = X[:, 0, :]
|
||||
if final_norm_w is not None:
|
||||
xf = x_out.float()
|
||||
rms = xf.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt()
|
||||
x_out = (xf * rms * final_norm_w.float()).bfloat16()
|
||||
logits = torch.nn.functional.linear(x_out, lm_w)
|
||||
|
||||
# Results
|
||||
print(f"\n === Minimal E2E Test Results ===")
|
||||
print(f" Logits: min={logits.float().min().item():.2f} max={logits.float().max().item():.2f} "
|
||||
f"spread={logits.float().max().item() - logits.float().min().item():.2f}")
|
||||
print(f" NaN={torch.isnan(logits.float()).any().item()} "
|
||||
f"Inf={torch.isinf(logits.float()).any().item()}")
|
||||
|
||||
top10_vals, top10_ids = torch.topk(logits[0], 10)
|
||||
print(f" Top-10 predictions:")
|
||||
for i, (tid_v, val) in enumerate(zip(top10_ids, top10_vals)):
|
||||
tok_str = tokenizer.decode([tid_v.item()])
|
||||
print(f" {i+1}. '{tok_str}' (id={tid_v.item()}, logit={val.item():.3f})")
|
||||
|
||||
# Print residual stream evolution
|
||||
print(f"\n Residual stream evolution (stream 0):")
|
||||
for d in layer_diags[::5]: # Every 5th layer
|
||||
print(f" L{d['layer']:2d}: |X|={d['x_max']:.1f} "
|
||||
f"mean={d['x0_mean']:.1f} std={d['x0_std']:.1f} "
|
||||
f"nan={d['nan']} inf={d['inf']}")
|
||||
# Always print last
|
||||
if layer_diags:
|
||||
d = layer_diags[-1]
|
||||
print(f" L{d['layer']:2d}: |X|={d['x_max']:.1f} "
|
||||
f"mean={d['x0_mean']:.1f} std={d['x0_std']:.1f} "
|
||||
f"nan={d['nan']} inf={d['inf']}")
|
||||
|
||||
# Check for reasonable output
|
||||
spread = logits.float().max().item() - logits.float().min().item()
|
||||
if spread < 1.0:
|
||||
print(f" ⚠️ Logit spread {spread:.2f} is very low — model is essentially uniform")
|
||||
else:
|
||||
print(f" ✓ Logit spread {spread:.2f} looks reasonable")
|
||||
|
||||
return logits, layer_diags
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user