mirror of
https://github.com/karpathy/nanochat.git
synced 2026-01-26 05:14:13 +00:00
348 lines
15 KiB
Python
348 lines
15 KiB
Python
"""
|
|
Engine for efficient inference of our models.
|
|
|
|
Everything works around token sequences:
|
|
- The user can send token sequences to the engine
|
|
- The engine returns the next token
|
|
|
|
Notes:
|
|
- The engine knows nothing about tokenization, it's purely token id sequences.
|
|
|
|
The whole thing is made as efficient as possible.
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import signal
|
|
import warnings
|
|
from contextlib import contextmanager
|
|
from collections import deque
|
|
from nanochat.common import compute_init, autodetect_device_type
|
|
from nanochat.checkpoint_manager import load_model
|
|
from contextlib import nullcontext
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Calculator tool helpers
|
|
@contextmanager
|
|
def timeout(duration, formula):
|
|
def timeout_handler(signum, frame):
|
|
raise Exception(f"'{formula}': timed out after {duration} seconds")
|
|
|
|
signal.signal(signal.SIGALRM, timeout_handler)
|
|
signal.alarm(duration)
|
|
yield
|
|
signal.alarm(0)
|
|
|
|
def eval_with_timeout(formula, max_time=3):
|
|
try:
|
|
with timeout(max_time, formula):
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore", SyntaxWarning)
|
|
return eval(formula, {"__builtins__": {}}, {})
|
|
except Exception as e:
|
|
signal.alarm(0)
|
|
# print(f"Warning: Failed to eval {formula}, exception: {e}") # it's ok ignore wrong calculator usage
|
|
return None
|
|
|
|
def use_calculator(expr):
|
|
"""
|
|
Evaluate a Python expression safely.
|
|
Supports both math expressions and string operations like .count()
|
|
"""
|
|
# Remove commas from numbers
|
|
expr = expr.replace(",", "")
|
|
|
|
# Check if it's a pure math expression (old behavior)
|
|
if all([x in "0123456789*+-/.() " for x in expr]):
|
|
if "**" in expr: # disallow power operator
|
|
return None
|
|
return eval_with_timeout(expr)
|
|
|
|
# Check if it's a string operation we support
|
|
# Allow: strings (single/double quotes), .count(), letters, numbers, spaces, parens
|
|
allowed_chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'\"()._ "
|
|
if not all([x in allowed_chars for x in expr]):
|
|
return None
|
|
|
|
# Disallow dangerous patterns
|
|
dangerous_patterns = ['__', 'import', 'exec', 'eval', 'compile', 'open', 'file',
|
|
'input', 'raw_input', 'globals', 'locals', 'vars', 'dir',
|
|
'getattr', 'setattr', 'delattr', 'hasattr']
|
|
expr_lower = expr.lower()
|
|
if any(pattern in expr_lower for pattern in dangerous_patterns):
|
|
return None
|
|
|
|
# Only allow .count() method for now (can expand later)
|
|
if '.count(' not in expr:
|
|
return None
|
|
|
|
# Evaluate with timeout
|
|
return eval_with_timeout(expr)
|
|
|
|
# -----------------------------------------------------------------------------
|
|
class KVCache:
|
|
"""
|
|
KV Cache designed for Flash Attention 3's flash_attn_with_kvcache API.
|
|
|
|
Key differences from FA2-style cache:
|
|
- Tensors are (B, T, H, D) not (B, H, T, D)
|
|
- FA3 updates the cache in-place during flash_attn_with_kvcache
|
|
- Position tracked per batch element via cache_seqlens tensor
|
|
"""
|
|
|
|
def __init__(self, batch_size, num_heads, seq_len, head_dim, num_layers, device, dtype=torch.bfloat16):
|
|
self.batch_size = batch_size
|
|
self.max_seq_len = seq_len
|
|
self.n_layers = num_layers
|
|
self.n_heads = num_heads
|
|
self.head_dim = head_dim
|
|
# Pre-allocate cache tensors: (n_layers, B, T, H, D)
|
|
self.k_cache = torch.zeros(num_layers, batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
|
|
self.v_cache = torch.zeros(num_layers, batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
|
|
# Current sequence length per batch element (FA3 needs int32)
|
|
self.cache_seqlens = torch.zeros(batch_size, dtype=torch.int32, device=device)
|
|
|
|
def reset(self):
|
|
"""Reset cache to empty state."""
|
|
self.cache_seqlens.zero_()
|
|
|
|
def get_pos(self):
|
|
"""Get current position (assumes all batch elements at same position)."""
|
|
return self.cache_seqlens[0].item()
|
|
|
|
def get_layer_cache(self, layer_idx):
|
|
"""Return (k_cache, v_cache) views for a specific layer."""
|
|
return self.k_cache[layer_idx], self.v_cache[layer_idx]
|
|
|
|
def advance(self, num_tokens):
|
|
"""Advance the cache position by num_tokens."""
|
|
self.cache_seqlens += num_tokens
|
|
|
|
def prefill(self, other):
|
|
"""
|
|
Copy cached KV from another cache into this one.
|
|
Used when we do batch=1 prefill and then want to generate multiple samples in parallel.
|
|
"""
|
|
assert self.get_pos() == 0, "Cannot prefill a non-empty KV cache"
|
|
assert self.n_layers == other.n_layers and self.n_heads == other.n_heads and self.head_dim == other.head_dim
|
|
assert self.max_seq_len >= other.max_seq_len
|
|
other_pos = other.get_pos()
|
|
self.k_cache[:, :, :other_pos, :, :] = other.k_cache[:, :, :other_pos, :, :]
|
|
self.v_cache[:, :, :other_pos, :, :] = other.v_cache[:, :, :other_pos, :, :]
|
|
self.cache_seqlens.fill_(other_pos)
|
|
|
|
# -----------------------------------------------------------------------------
|
|
@torch.inference_mode()
|
|
def sample_next_token(logits, rng, temperature=1.0, top_k=None):
|
|
"""Sample a single next token from given logits of shape (B, vocab_size). Returns (B, 1)."""
|
|
assert temperature >= 0.0, "temperature must be non-negative"
|
|
if temperature == 0.0:
|
|
return torch.argmax(logits, dim=-1, keepdim=True)
|
|
if top_k is not None and top_k > 0:
|
|
k = min(top_k, logits.size(-1))
|
|
vals, idx = torch.topk(logits, k, dim=-1)
|
|
vals = vals / temperature
|
|
probs = F.softmax(vals, dim=-1)
|
|
choice = torch.multinomial(probs, num_samples=1, generator=rng)
|
|
return idx.gather(1, choice)
|
|
else:
|
|
logits = logits / temperature
|
|
probs = F.softmax(logits, dim=-1)
|
|
return torch.multinomial(probs, num_samples=1, generator=rng)
|
|
|
|
# -----------------------------------------------------------------------------
|
|
|
|
class RowState:
|
|
# Per-row state tracking during generation
|
|
def __init__(self, current_tokens=None):
|
|
self.current_tokens = current_tokens or [] # Current token sequence for this row
|
|
self.forced_tokens = deque() # Queue of tokens to force inject
|
|
self.in_python_block = False # Whether we are inside a python block
|
|
self.python_expr_tokens = [] # Tokens of the current python expression
|
|
self.completed = False # Whether this row has completed generation
|
|
|
|
class Engine:
|
|
|
|
def __init__(self, model, tokenizer):
|
|
self.model = model
|
|
self.tokenizer = tokenizer # needed for tool use
|
|
|
|
@torch.inference_mode()
|
|
def generate(self, tokens, num_samples=1, max_tokens=None, temperature=1.0, top_k=None, seed=42):
|
|
"""Same as generate, but does single prefill and then clones the KV cache."""
|
|
assert isinstance(tokens, list) and isinstance(tokens[0], int), "expecting list of ints"
|
|
device = self.model.get_device()
|
|
rng = torch.Generator(device=device)
|
|
rng.manual_seed(seed)
|
|
|
|
# Get the special tokens we need to coordinate the tool use state machine
|
|
get_special = lambda s: self.tokenizer.encode_special(s)
|
|
python_start = get_special("<|python_start|>")
|
|
python_end = get_special("<|python_end|>")
|
|
output_start = get_special("<|output_start|>")
|
|
output_end = get_special("<|output_end|>")
|
|
assistant_end = get_special("<|assistant_end|>") # if sampled, ends row
|
|
bos = self.tokenizer.get_bos_token_id() # if sampled, ends row
|
|
|
|
# 1) Run a batch 1 prefill of the prompt tokens
|
|
m = self.model.config
|
|
kv_model_kwargs = {"num_heads": m.n_kv_head, "head_dim": m.n_embd // m.n_head, "num_layers": m.n_layer}
|
|
kv_cache_prefill = KVCache(
|
|
batch_size=1,
|
|
seq_len=len(tokens),
|
|
device=device,
|
|
**kv_model_kwargs,
|
|
)
|
|
ids = torch.tensor([tokens], dtype=torch.long, device=device)
|
|
logits = self.model.forward(ids, kv_cache=kv_cache_prefill)
|
|
logits = logits[:, -1, :].expand(num_samples, -1) # (num_samples, vocab_size)
|
|
|
|
# 2) Replicate the KV cache for each sample/row
|
|
kv_length_hint = (len(tokens) + max_tokens) if max_tokens is not None else self.model.config.sequence_len
|
|
kv_cache_decode = KVCache(
|
|
batch_size=num_samples,
|
|
seq_len=kv_length_hint,
|
|
device=device,
|
|
**kv_model_kwargs,
|
|
)
|
|
kv_cache_decode.prefill(kv_cache_prefill)
|
|
del kv_cache_prefill # no need to keep this memory around
|
|
|
|
# 3) Initialize states for each sample
|
|
row_states = [RowState(tokens.copy()) for _ in range(num_samples)]
|
|
|
|
# 4) Main generation loop
|
|
num_generated = 0
|
|
while True:
|
|
# Stop condition: we've reached max tokens
|
|
if max_tokens is not None and num_generated >= max_tokens:
|
|
break
|
|
# Stop condition: all rows are completed
|
|
if all(state.completed for state in row_states):
|
|
break
|
|
|
|
# Sample the next token for each row
|
|
next_ids = sample_next_token(logits, rng, temperature, top_k) # (B, 1)
|
|
sampled_tokens = next_ids[:, 0].tolist()
|
|
|
|
# Process each row: choose the next token, update state, optional tool use
|
|
token_column = [] # contains the next token id along each row
|
|
token_masks = [] # contains the mask (was it sampled (1) or forced (0)?) along each row
|
|
for i, state in enumerate(row_states):
|
|
# Select the next token in this row
|
|
is_forced = len(state.forced_tokens) > 0 # are there tokens waiting to be forced in deque?
|
|
token_masks.append(0 if is_forced else 1) # mask is 0 if forced, 1 if sampled
|
|
next_token = state.forced_tokens.popleft() if is_forced else sampled_tokens[i]
|
|
token_column.append(next_token)
|
|
# Update the state of this row to include the next token
|
|
state.current_tokens.append(next_token)
|
|
# On <|assistant_end|> or <|bos|>, mark the row as completed
|
|
if next_token == assistant_end or next_token == bos:
|
|
state.completed = True
|
|
# Handle tool logic
|
|
if next_token == python_start:
|
|
state.in_python_block = True
|
|
state.python_expr_tokens = []
|
|
elif next_token == python_end and state.in_python_block:
|
|
state.in_python_block = False
|
|
if state.python_expr_tokens:
|
|
expr = self.tokenizer.decode(state.python_expr_tokens)
|
|
result = use_calculator(expr)
|
|
if result is not None:
|
|
result_tokens = self.tokenizer.encode(str(result))
|
|
state.forced_tokens.append(output_start)
|
|
state.forced_tokens.extend(result_tokens)
|
|
state.forced_tokens.append(output_end)
|
|
state.python_expr_tokens = []
|
|
elif state.in_python_block:
|
|
state.python_expr_tokens.append(next_token)
|
|
|
|
# Yield the token column
|
|
yield token_column, token_masks
|
|
num_generated += 1
|
|
|
|
# Prepare logits for next iteration
|
|
ids = torch.tensor(token_column, dtype=torch.long, device=device).unsqueeze(1)
|
|
logits = self.model.forward(ids, kv_cache=kv_cache_decode)[:, -1, :] # (B, vocab_size)
|
|
|
|
def generate_batch(self, tokens, num_samples=1, **kwargs):
|
|
"""
|
|
Non-streaming batch generation that just returns the final token sequences.
|
|
Returns a list of token sequences (list of lists of ints).
|
|
Terminal tokens (assistant_end, bos) are not included in the results.
|
|
"""
|
|
assistant_end = self.tokenizer.encode_special("<|assistant_end|>")
|
|
bos = self.tokenizer.get_bos_token_id()
|
|
results = [tokens.copy() for _ in range(num_samples)]
|
|
masks = [[0] * len(tokens) for _ in range(num_samples)]
|
|
completed = [False] * num_samples
|
|
for token_column, token_masks in self.generate(tokens, num_samples, **kwargs):
|
|
for i, (token, mask) in enumerate(zip(token_column, token_masks)):
|
|
if not completed[i]:
|
|
if token == assistant_end or token == bos:
|
|
completed[i] = True
|
|
else:
|
|
results[i].append(token)
|
|
masks[i].append(mask)
|
|
# Stop if all rows are completed
|
|
if all(completed):
|
|
break
|
|
return results, masks
|
|
|
|
|
|
if __name__ == "__main__":
|
|
"""
|
|
Quick inline test to make sure that the naive/slow model.generate function
|
|
is equivalent to the faster Engine.generate function here.
|
|
"""
|
|
import time
|
|
# init compute
|
|
ddp, ddp_rank, ddp_local_rank, ddp_world_size, device = compute_init()
|
|
device_type = autodetect_device_type()
|
|
autocast_ctx = torch.amp.autocast(device_type=device_type, dtype=torch.bfloat16) if device_type == "cuda" else nullcontext()
|
|
|
|
# load the model and tokenizer
|
|
model, tokenizer, meta = load_model("base", device, phase="eval")
|
|
bos_token_id = tokenizer.get_bos_token_id()
|
|
# common hyperparameters
|
|
kwargs = dict(max_tokens=64, temperature=0.0)
|
|
# set the starting prompt
|
|
prompt_tokens = tokenizer.encode("The chemical formula of water is", prepend=bos_token_id)
|
|
# generate the reference sequence using the model.generate() function
|
|
generated_tokens = []
|
|
torch.cuda.synchronize()
|
|
t0 = time.time()
|
|
stream = model.generate(prompt_tokens, **kwargs)
|
|
with autocast_ctx:
|
|
for token in stream:
|
|
generated_tokens.append(token)
|
|
chunk = tokenizer.decode([token])
|
|
print(chunk, end="", flush=True)
|
|
print()
|
|
torch.cuda.synchronize()
|
|
t1 = time.time()
|
|
print(f"Reference time: {t1 - t0:.2f}s")
|
|
reference_ids = generated_tokens
|
|
# generate tokens with Engine
|
|
generated_tokens = []
|
|
engine = Engine(model, tokenizer)
|
|
stream = engine.generate(prompt_tokens, num_samples=1, **kwargs) # note: runs in fp32
|
|
torch.cuda.synchronize()
|
|
t0 = time.time()
|
|
with autocast_ctx:
|
|
for token_column, token_masks in stream:
|
|
token = token_column[0] # only print out the first row
|
|
generated_tokens.append(token)
|
|
chunk = tokenizer.decode([token])
|
|
print(chunk, end="", flush=True)
|
|
print()
|
|
torch.cuda.synchronize()
|
|
t1 = time.time()
|
|
print(f"Engine time: {t1 - t0:.2f}s")
|
|
# compare the two sequences
|
|
for i in range(len(reference_ids)):
|
|
if reference_ids[i] != generated_tokens[i]:
|
|
print(f"Mismatch at {i}: {reference_ids[i]} != {generated_tokens[i]}")
|
|
break
|
|
print(f"Match: {reference_ids == generated_tokens}")
|