biondizzle
49282fe206
Key corrections: - RMSNorm applied to projection output (mixes *= rsqrt(sqrsum/K + eps)) not to the input before projection - comb_mix uses softmax + Sinkhorn, NOT exp + Sinkhorn - pre_mix = sigmoid(logits) + eps (not matmul with X_l) - layer_input = sum(pre_mix * residual) — weighted sum, not bmm - post_mix = sigmoid * hc_post_mult_value (2.0) - bias split: [pre(4), post(4), comb(16)] not [pre(4), comb(16), post(4)]
469 lines
18 KiB
Python
469 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""Single-shot DSV4 inference — 8-GPU pipeline parallel with mHC.
|
|
|
|
Loads the full NVFP4 checkpoint across 8 B200 GPUs. Includes:
|
|
- mHC (Manifold-Constrained Hyper-Connections) — load-bearing residual
|
|
- Q low-rank projection + KV projection
|
|
- RoPE (partial, last 64 dims)
|
|
- Production FMHA kernel (tcgen05 MMA, hd=512, 128 heads)
|
|
- Output projection: wo_a (grouped BMM) → wo_b (NVFP4)
|
|
- Shared expert FFN (SwiGLU)
|
|
- NVFP4 dequant → BF16 matmul baseline for linear layers
|
|
|
|
Missing (causing incorrect output):
|
|
- Routed MoE experts (384 experts, top-6)
|
|
- KV cache across decode steps
|
|
- Compressor + indexer (CSA/HCA compressed KV)
|
|
|
|
Usage (on B200):
|
|
source /root/dsv4-nvfp4-workspace/venv/bin/activate
|
|
cd /root/dsv4-nvfp4-workspace/kernel
|
|
python3 single_shot_inference.py
|
|
"""
|
|
import os, sys, time, json, math
|
|
import torch
|
|
from pathlib import Path
|
|
|
|
CHECKPOINT_DIR = "/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4"
|
|
MAX_NEW_TOKENS = 10
|
|
PROMPT = "The capital of France is"
|
|
NUM_GPUS = 8
|
|
|
|
# =====================================================================
|
|
# NVFP4 dequantization
|
|
# =====================================================================
|
|
|
|
FP4_LUT = torch.tensor([0., 2., 3., 4., 6., 8., 12., 24.])
|
|
|
|
def dequant_nvfp4_weight(weight, weight_scale, weight_scale_2):
|
|
out_dim = weight.shape[0]
|
|
in_packed = weight.shape[1]
|
|
in_features = in_packed * 2
|
|
low = (weight & 0x0F).to(torch.int8)
|
|
high = (weight >> 4).to(torch.int8)
|
|
low_sign, low_idx = (low >> 3).bool(), (low & 0x07).long()
|
|
high_sign, high_idx = (high >> 3).bool(), (high & 0x07).long()
|
|
lut = FP4_LUT.to(device=weight.device, dtype=torch.float32)
|
|
low_f = lut[low_idx] * torch.where(low_sign, -1.0, 1.0)
|
|
high_f = lut[high_idx] * torch.where(high_sign, -1.0, 1.0)
|
|
w_f = torch.stack([low_f, high_f], dim=-1).reshape(out_dim, in_features)
|
|
scale_f = weight_scale.float() * weight_scale_2.float()
|
|
scale_expanded = scale_f.repeat_interleave(16, dim=1)
|
|
return (w_f * scale_expanded).bfloat16()
|
|
|
|
def nvfp4_linear(x, weight, weight_scale, weight_scale_2):
|
|
w = dequant_nvfp4_weight(weight, weight_scale, weight_scale_2)
|
|
return torch.nn.functional.linear(x, w)
|
|
|
|
def bf16_linear(x, weight):
|
|
return torch.nn.functional.linear(x, weight.bfloat16())
|
|
|
|
|
|
# =====================================================================
|
|
# mHC — Manifold-Constrained Hyper-Connections
|
|
# =====================================================================
|
|
|
|
def sinkhorn_knopp(M, t_max=20, eps=1e-6):
|
|
"""Project (T, n, n) positive matrices onto Birkhoff polytope."""
|
|
for _ in range(t_max):
|
|
M = M / (M.sum(dim=-1, keepdim=True) + eps)
|
|
M = M / (M.sum(dim=-2, keepdim=True) + eps)
|
|
return M
|
|
|
|
|
|
class mHCBlock:
|
|
"""One mHC block — matches vLLM torch reference exactly.
|
|
|
|
Checkpoint: fn (24,28672), base (24,), scale (3,)
|
|
Split: pre_mix (n_hc=4), post_mix (n_hc=4), comb_mix (n_hc*n_hc=16)
|
|
"""
|
|
def __init__(self, hidden_dim=7168, n_hc=4, sinkhorn_repeat=20, device='cuda'):
|
|
self.n_hc = n_hc
|
|
self.K = n_hc * hidden_dim
|
|
self.sinkhorn_repeat = sinkhorn_repeat
|
|
self.device = device
|
|
self.fn = None
|
|
self.hc_scale = None
|
|
self.hc_base = None
|
|
self.rms_eps = 1e-6
|
|
self.hc_pre_eps = 0.0
|
|
self.hc_sinkhorn_eps = 1e-6
|
|
self.hc_post_mult_value = 2.0
|
|
|
|
def load_from_checkpoint(self, fn, base, scale):
|
|
self.fn = fn.to(device=self.device, dtype=torch.float32).contiguous()
|
|
self.hc_base = base.to(device=self.device, dtype=torch.float32).contiguous()
|
|
self.hc_scale = scale.to(device=self.device, dtype=torch.float32).contiguous()
|
|
|
|
def pre_block(self, residual):
|
|
"""residual: (T, n_hc, d) BF16 → layer_input (T, d), ctx"""
|
|
n = self.n_hc
|
|
K = self.K
|
|
eps = self.rms_eps
|
|
T = residual.shape[0]
|
|
|
|
res_flat = residual.reshape(T, K).float() # (T, K)
|
|
|
|
# Project with RMSNorm on input
|
|
mixes = torch.matmul(res_flat, self.fn.t()) # (T, 24)
|
|
sqrsum = res_flat.square().sum(dim=-1, keepdim=True)
|
|
mixes = mixes * torch.rsqrt(sqrsum / K + eps)
|
|
|
|
# Split
|
|
pre_logits = mixes[:, :n] * self.hc_scale[0] + self.hc_base[:n]
|
|
pre_mix = torch.sigmoid(pre_logits) + self.hc_pre_eps # (T, 4)
|
|
|
|
post_logits = mixes[:, n:2*n] * self.hc_scale[1] + self.hc_base[n:2*n]
|
|
post_mix = torch.sigmoid(post_logits) * self.hc_post_mult_value # (T, 4)
|
|
|
|
comb_logits = (mixes[:, 2*n:].reshape(T, n, n) * self.hc_scale[2]
|
|
+ self.hc_base[2*n:].reshape(1, n, n))
|
|
comb_mix = torch.softmax(comb_logits, dim=-1) + self.hc_sinkhorn_eps
|
|
comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + self.hc_sinkhorn_eps)
|
|
for _ in range(self.sinkhorn_repeat - 1):
|
|
comb_mix = comb_mix / (comb_mix.sum(dim=-1, keepdim=True) + self.hc_sinkhorn_eps)
|
|
comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + self.hc_sinkhorn_eps)
|
|
|
|
# layer_input = sum(pre_mix * residual)
|
|
layer_input = (pre_mix.unsqueeze(-1) * residual.float()).sum(dim=1).bfloat16()
|
|
|
|
return layer_input, (post_mix, comb_mix)
|
|
|
|
def post_block(self, residual, F_out, ctx):
|
|
"""residual: (T, n_hc, d), F_out: (T, d) → (T, n_hc, d)"""
|
|
post_mix, comb_mix = ctx
|
|
mixed_residual = torch.einsum('tij,tjh->tjh', comb_mix, residual.float())
|
|
post_term = post_mix.unsqueeze(-1) * F_out.unsqueeze(1).float()
|
|
return (mixed_residual + post_term).bfloat16()
|
|
# RoPE
|
|
# =====================================================================
|
|
|
|
def build_rope_cache(max_pos, head_dim, rope_dim, device, theta=10000.0):
|
|
half = rope_dim // 2
|
|
freqs = 1.0 / (theta ** (torch.arange(0, rope_dim, 2, dtype=torch.float32) / rope_dim))
|
|
angles = torch.outer(torch.arange(max_pos, dtype=torch.float32), freqs)
|
|
return torch.cos(angles).to(device), torch.sin(angles).to(device)
|
|
|
|
def apply_rope(x, positions, cos_cache, sin_cache, rope_dim):
|
|
T, n_h, hd = x.shape
|
|
nope = hd - rope_dim
|
|
cos = cos_cache[positions].unsqueeze(1).to(x.dtype)
|
|
sin = sin_cache[positions].unsqueeze(1).to(x.dtype)
|
|
out = x.clone()
|
|
out[:, :, nope:][..., 0::2] = x[:, :, nope:][..., 0::2] * cos - x[:, :, nope:][..., 1::2] * sin
|
|
out[:, :, nope:][..., 1::2] = x[:, :, nope:][..., 0::2] * sin + x[:, :, nope:][..., 1::2] * cos
|
|
return out
|
|
|
|
|
|
# =====================================================================
|
|
# Weight loading
|
|
# =====================================================================
|
|
|
|
def load_all_weights(checkpoint_dir, num_layers):
|
|
from safetensors.torch import load_file
|
|
|
|
cdir = Path(checkpoint_dir)
|
|
index_path = cdir / "model.safetensors.index.json"
|
|
weight_map = {}
|
|
if index_path.exists():
|
|
with open(index_path) as f:
|
|
weight_map = json.load(f).get("weight_map", {})
|
|
|
|
shard_names = set(weight_map.values()) if weight_map else {
|
|
f"model-{i:05d}-of-00095.safetensors" for i in range(1, 96)
|
|
}
|
|
|
|
print(f"Loading {len(shard_names)} shards...")
|
|
all_weights = {}
|
|
loaded = 0
|
|
for shard_name in sorted(shard_names):
|
|
if not (cdir / shard_name).exists():
|
|
continue
|
|
data = load_file(str(cdir / shard_name))
|
|
all_weights.update(data)
|
|
loaded += 1
|
|
if loaded % 20 == 0:
|
|
print(f" {loaded}/{len(shard_names)} shards, {len(all_weights)} tensors")
|
|
print(f" Done: {len(all_weights)} tensors")
|
|
|
|
layer_weights = {}
|
|
global_weights = {}
|
|
|
|
print("Assigning to GPUs...")
|
|
for key, tensor in all_weights.items():
|
|
if key.startswith("model.layers."):
|
|
li = int(key.split(".")[2])
|
|
target_gpu = li % NUM_GPUS
|
|
target_device = f"cuda:{target_gpu}"
|
|
if li not in layer_weights:
|
|
layer_weights[li] = {"_device": target_device, "_gpu": target_gpu}
|
|
layer_weights[li][key] = tensor.to(target_device)
|
|
elif key.startswith("model.embed_tokens"):
|
|
global_weights[key] = tensor.to("cuda:0")
|
|
elif key.startswith("model.norm"):
|
|
global_weights[key] = tensor.to("cuda:0")
|
|
elif key.startswith("lm_head"):
|
|
global_weights[key] = tensor.to("cuda:0")
|
|
|
|
for gpu in range(NUM_GPUS):
|
|
alloc = torch.cuda.memory_allocated(gpu) / 1e9
|
|
print(f" GPU {gpu}: {alloc:.1f}GB")
|
|
|
|
return layer_weights, global_weights
|
|
|
|
|
|
# =====================================================================
|
|
# Single layer forward
|
|
# =====================================================================
|
|
|
|
def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin, attn_mhc, ffn_mhc):
|
|
"""Forward one layer with mHC.
|
|
|
|
X_l: (1, n_hc, H) BF16 → (1, n_hc, H) BF16
|
|
"""
|
|
device = X_l.device
|
|
H = cfg["hidden_size"]
|
|
n_h = cfg["num_attention_heads"]
|
|
hd = cfg["head_dim"]
|
|
rd = cfg["qk_rope_head_dim"]
|
|
o_rank = cfg["o_lora_rank"]
|
|
o_groups = cfg["o_groups"]
|
|
n_hc = 4
|
|
|
|
pre = f"model.layers.{li}.self_attn"
|
|
T = X_l.shape[0]
|
|
heads_per_group = n_h // o_groups
|
|
group_input_dim = heads_per_group * hd
|
|
|
|
# ==== mHC pre_block (attention) ====
|
|
x_in, attn_ctx = attn_mhc.pre_block(X_l) # x_in: (T, H) BF16
|
|
|
|
# ==== Q projection ====
|
|
c_Q = nvfp4_linear(x_in, w[f"{pre}.q_a_proj.weight"],
|
|
w[f"{pre}.q_a_proj.weight_scale"],
|
|
w[f"{pre}.q_a_proj.weight_scale_2"])
|
|
q = nvfp4_linear(c_Q, w[f"{pre}.q_b_proj.weight"],
|
|
w[f"{pre}.q_b_proj.weight_scale"],
|
|
w[f"{pre}.q_b_proj.weight_scale_2"])
|
|
|
|
# ==== KV projection ====
|
|
kv = nvfp4_linear(x_in, w[f"{pre}.kv_proj.weight"],
|
|
w[f"{pre}.kv_proj.weight_scale"],
|
|
w[f"{pre}.kv_proj.weight_scale_2"])
|
|
|
|
# ==== Reshape for attention ====
|
|
q_heads = q.reshape(T, n_h, hd).permute(1, 0, 2)
|
|
k = kv.reshape(T, 1, hd).permute(1, 0, 2)
|
|
v = k.clone()
|
|
|
|
# ==== RoPE ====
|
|
pos = torch.tensor([0], dtype=torch.long, device=device)
|
|
q_heads = apply_rope(q_heads, pos, rope_cos, rope_sin, rd)
|
|
k = apply_rope(k, pos, rope_cos, rope_sin, rd)
|
|
|
|
# ==== FMHA ====
|
|
from dsv4.kernels.attention.production import dsv4_attention
|
|
attn_out = dsv4_attention(q_heads, k, v)
|
|
attn_out = attn_out.permute(1, 0, 2).reshape(T, n_h * hd)
|
|
|
|
# ==== Output projection ====
|
|
attn_grouped = attn_out.reshape(T, o_groups, heads_per_group, hd).reshape(T, o_groups, group_input_dim)
|
|
oa_w = w[f"{pre}.o_a_proj.weight"].bfloat16()
|
|
oa_3d = oa_w.reshape(o_groups, o_rank, group_input_dim)
|
|
attn_for_bmm = attn_grouped.permute(1, 0, 2)
|
|
grouped_out = torch.bmm(attn_for_bmm, oa_3d.transpose(1, 2))
|
|
grouped_flat = grouped_out.permute(1, 0, 2).reshape(T, o_groups * o_rank)
|
|
F_attn = nvfp4_linear(grouped_flat, w[f"{pre}.o_b_proj.weight"],
|
|
w[f"{pre}.o_b_proj.weight_scale"],
|
|
w[f"{pre}.o_b_proj.weight_scale_2"])
|
|
|
|
# ==== mHC post_block (attention) ====
|
|
X_l = attn_mhc.post_block(X_l, F_attn, attn_ctx) # (T, n_hc, H)
|
|
|
|
# ==== mHC pre_block (FFN) ====
|
|
x_ffn, ffn_ctx = ffn_mhc.pre_block(X_l)
|
|
|
|
# ==== FFN: shared expert ====
|
|
se_pre = f"model.layers.{li}.mlp.shared_experts"
|
|
se_gate_w = w.get(f"{se_pre}.gate_proj.weight")
|
|
F_ffn = torch.zeros_like(x_ffn)
|
|
if se_gate_w is not None:
|
|
gate = nvfp4_linear(x_ffn, se_gate_w,
|
|
w[f"{se_pre}.gate_proj.weight_scale"],
|
|
w[f"{se_pre}.gate_proj.weight_scale_2"])
|
|
up = nvfp4_linear(x_ffn, w[f"{se_pre}.up_proj.weight"],
|
|
w[f"{se_pre}.up_proj.weight_scale"],
|
|
w[f"{se_pre}.up_proj.weight_scale_2"])
|
|
F_ffn = nvfp4_linear(
|
|
torch.nn.functional.silu(gate) * up,
|
|
w[f"{se_pre}.down_proj.weight"],
|
|
w[f"{se_pre}.down_proj.weight_scale"],
|
|
w[f"{se_pre}.down_proj.weight_scale_2"],
|
|
)
|
|
|
|
# ==== mHC post_block (FFN) ====
|
|
X_l = ffn_mhc.post_block(X_l, F_ffn, ffn_ctx)
|
|
|
|
return X_l
|
|
|
|
|
|
# =====================================================================
|
|
# Main
|
|
# =====================================================================
|
|
|
|
def main():
|
|
t_start = time.time()
|
|
print("=" * 70)
|
|
print("DSV4 Single-Shot Inference — 8-GPU with mHC")
|
|
print("=" * 70)
|
|
|
|
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"]
|
|
n_h = cfg["num_attention_heads"]
|
|
hd = cfg["head_dim"]
|
|
rd = cfg["qk_rope_head_dim"]
|
|
n_hc = 4
|
|
print(f"Model: {n_layers} layers, {n_h} heads, hd={hd}, rope_dim={rd}")
|
|
|
|
# ==== Phase 1: Load weights ====
|
|
print(f"\n{'='*70}\nPhase 1: Loading weights\n{'='*70}")
|
|
layer_weights, global_weights = load_all_weights(CHECKPOINT_DIR, n_layers)
|
|
t_loaded = time.time()
|
|
print(f"Weight loading: {t_loaded - t_start:.1f}s")
|
|
|
|
# ==== Build mHC blocks per layer ====
|
|
print("Building mHC blocks...")
|
|
attn_mhc_blocks = {}
|
|
ffn_mhc_blocks = {}
|
|
for li in range(n_layers):
|
|
gpu = li % NUM_GPUS
|
|
dev = f"cuda:{gpu}"
|
|
|
|
# Attention mHC
|
|
attn_fn = layer_weights[li].get(f"model.layers.{li}.attn_hc.fn")
|
|
attn_base = layer_weights[li].get(f"model.layers.{li}.attn_hc.base")
|
|
attn_scale = layer_weights[li].get(f"model.layers.{li}.attn_hc.scale")
|
|
if attn_fn is not None and attn_base is not None and attn_scale is not None:
|
|
attn_mhc = mHCBlock(hidden_dim=H, n_hc=n_hc, device=dev)
|
|
attn_mhc.load_from_checkpoint(attn_fn, attn_base, attn_scale)
|
|
attn_mhc_blocks[li] = attn_mhc
|
|
|
|
# FFN mHC
|
|
ffn_fn = layer_weights[li].get(f"model.layers.{li}.ffn_hc.fn")
|
|
ffn_base = layer_weights[li].get(f"model.layers.{li}.ffn_hc.base")
|
|
ffn_scale = layer_weights[li].get(f"model.layers.{li}.ffn_hc.scale")
|
|
if ffn_fn is not None and ffn_base is not None and ffn_scale is not None:
|
|
ffn_mhc = mHCBlock(hidden_dim=H, n_hc=n_hc, device=dev)
|
|
ffn_mhc.load_from_checkpoint(ffn_fn, ffn_base, ffn_scale)
|
|
ffn_mhc_blocks[li] = ffn_mhc
|
|
|
|
print(f" attn mHC: {len(attn_mhc_blocks)} layers")
|
|
print(f" ffn mHC: {len(ffn_mhc_blocks)} layers")
|
|
|
|
# ==== Global weights (gpu0) ====
|
|
torch.cuda.set_device(0)
|
|
embed_w = global_weights.get("model.embed_tokens.weight")
|
|
embed = torch.nn.Embedding.from_pretrained(embed_w.bfloat16())
|
|
lm_w = global_weights.get("lm_head.weight", embed_w).bfloat16()
|
|
final_norm_w = global_weights.get("model.norm.weight")
|
|
|
|
# RoPE caches per GPU
|
|
rope_caches = {g: build_rope_cache(8192, hd, rd, f"cuda:{g}") for g in range(NUM_GPUS)}
|
|
|
|
# ==== Phase 2: Compile ====
|
|
print(f"\n{'='*70}\nPhase 2: JIT compiling\n{'='*70}")
|
|
from dsv4.kernels.attention.production import dsv4_attention
|
|
dummy_q = torch.randn(n_h, 1, hd, dtype=torch.bfloat16, device='cuda:0')
|
|
dummy_k = torch.randn(1, 1, hd, dtype=torch.bfloat16, device='cuda:0')
|
|
try:
|
|
_ = dsv4_attention(dummy_q, dummy_k, dummy_k.clone())
|
|
print(" FMHA: compiled OK")
|
|
except Exception as e:
|
|
print(f" FMHA error: {e}")
|
|
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}")
|
|
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()}")
|
|
|
|
generated = input_ids[0].tolist()
|
|
|
|
for step in range(MAX_NEW_TOKENS):
|
|
t0 = time.time()
|
|
tid = torch.tensor([generated[-1]], dtype=torch.long, device='cuda:0')
|
|
|
|
# Embed → mHC init state
|
|
emb = embed(tid) # (1, H) on gpu0
|
|
X = emb.unsqueeze(1).expand(-1, n_hc, -1).clone() # (1, n_hc, H)
|
|
|
|
# Process layers
|
|
for li in range(n_layers):
|
|
gpu = li % NUM_GPUS
|
|
target_device = f"cuda:{gpu}"
|
|
if X.device != torch.device(target_device):
|
|
X = X.to(target_device)
|
|
torch.cuda.set_device(gpu)
|
|
|
|
attn_mhc = attn_mhc_blocks.get(li)
|
|
ffn_mhc = ffn_mhc_blocks.get(li)
|
|
rc, rs = rope_caches[gpu]
|
|
X = forward_layer(X, layer_weights[li], li, cfg, rc, rs, attn_mhc, ffn_mhc)
|
|
|
|
# NaN check every 5 layers or on first NaN
|
|
if li % 10 == 0 or (X.float().abs().max().item() > 1e4 or torch.isnan(X.float()).any().item()):
|
|
has_nan = torch.isnan(X.float()).any().item()
|
|
xmax = X.float().abs().max().item()
|
|
print(f" L{li}: nan={has_nan}, max_abs={xmax:.2f}")
|
|
if has_nan:
|
|
# Debug: check mHC outputs
|
|
x_out = X[:, 0, :]
|
|
print(f" L{li} stream0: nan={torch.isnan(x_out.float()).any().item()}, max={x_out.float().abs().max().item():.2f}")
|
|
break
|
|
|
|
# Back to gpu0
|
|
X = X.to('cuda:0')
|
|
torch.cuda.set_device(0)
|
|
|
|
# Read out stream 0 → RMSNorm → lm_head
|
|
x_out = X[:, 0, :] # (1, H)
|
|
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)
|
|
next_id = torch.argmax(logits, dim=-1).item()
|
|
generated.append(next_id)
|
|
|
|
tok_str = tokenizer.decode([next_id])
|
|
dt = time.time() - t0
|
|
|
|
# Check for NaN
|
|
has_nan = torch.isnan(logits.float()).any().item()
|
|
logit_range = f"[{logits.float().min().item():.1f}, {logits.float().max().item():.1f}]"
|
|
print(f" Step {step}: {next_id} '{tok_str}' ({dt:.2f}s) logits={logit_range} nan={has_nan}")
|
|
|
|
if has_nan:
|
|
print(" NaN detected, stopping")
|
|
break
|
|
if next_id == tokenizer.eos_token_id:
|
|
break
|
|
|
|
out = tokenizer.decode(generated, skip_special_tokens=True)
|
|
total = time.time() - t_start
|
|
print(f"\n{'='*70}")
|
|
print(f"Input: '{PROMPT}'")
|
|
print(f"Output: '{out}'")
|
|
print(f"Total: {total:.1f}s (load: {t_loaded-t_start:.1f}s, compile: {t_compiled-t_loaded:.1f}s, infer: {time.time()-t_compiled:.1f}s)")
|
|
print(f"{'='*70}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|