135 lines
4.2 KiB
Python
135 lines
4.2 KiB
Python
"""
|
|
Inference script for the 1B Transformer — Single GPU.
|
|
|
|
Usage:
|
|
python inference.py # auto-finds latest checkpoint
|
|
python inference.py /path/to/checkpoint.pt # specific checkpoint
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
import glob
|
|
import time
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from model.config import ModelConfig
|
|
from model.transformer import Transformer
|
|
from model.data import get_tokenizer
|
|
|
|
|
|
def find_latest_checkpoint(checkpoint_dir="/jfs/deepak-kumar/checkpoints"):
|
|
files = glob.glob(os.path.join(checkpoint_dir, "step_*.pt"))
|
|
if not files:
|
|
final = os.path.join(checkpoint_dir, "final.pt")
|
|
return final if os.path.exists(final) else None
|
|
return max(files, key=lambda f: int(os.path.basename(f).split("_")[1].split(".")[0]))
|
|
|
|
|
|
def load_model(checkpoint_path, device="cuda:0"):
|
|
config = ModelConfig()
|
|
model = Transformer(config)
|
|
|
|
print(f"Loading checkpoint: {checkpoint_path}")
|
|
ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
|
|
model.load_state_dict(ckpt["model"])
|
|
model = model.to(device).bfloat16().eval()
|
|
|
|
step = ckpt.get("step", "?")
|
|
loss = ckpt.get("loss", "?")
|
|
print(f" Step: {step} | Loss: {loss}")
|
|
print(f" Params: {sum(p.numel() for p in model.parameters()):,}")
|
|
print(f" Device: {device}")
|
|
del ckpt
|
|
torch.cuda.empty_cache()
|
|
return model, config
|
|
|
|
|
|
@torch.no_grad()
|
|
def generate(model, tokenizer, prompt, max_new_tokens=200,
|
|
temperature=0.8, top_k=50, top_p=0.9, device="cuda:0"):
|
|
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
|
|
t0 = time.time()
|
|
|
|
for i in range(max_new_tokens):
|
|
if input_ids.shape[1] >= model.config.max_seq_len:
|
|
break
|
|
|
|
with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
|
|
logits, _ = model(input_ids)
|
|
|
|
logits = logits[:, -1, :] / temperature
|
|
|
|
if top_k > 0:
|
|
topk_vals, _ = torch.topk(logits, top_k)
|
|
logits[logits < topk_vals[:, -1:]] = float("-inf")
|
|
|
|
if top_p < 1.0:
|
|
sorted_logits, sorted_idx = torch.sort(logits, descending=True)
|
|
cum_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
|
|
mask = cum_probs - F.softmax(sorted_logits, dim=-1) >= top_p
|
|
sorted_logits[mask] = float("-inf")
|
|
logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)
|
|
|
|
probs = F.softmax(logits, dim=-1)
|
|
next_token = torch.multinomial(probs, num_samples=1)
|
|
|
|
if next_token.item() == tokenizer.eos_token_id:
|
|
break
|
|
|
|
input_ids = torch.cat([input_ids, next_token], dim=1)
|
|
|
|
elapsed = time.time() - t0
|
|
gen_tokens = input_ids.shape[1] - len(tokenizer.encode(prompt))
|
|
tok_per_sec = gen_tokens / max(elapsed, 1e-9)
|
|
|
|
text = tokenizer.decode(input_ids[0], skip_special_tokens=True)
|
|
return text, gen_tokens, tok_per_sec
|
|
|
|
|
|
def main():
|
|
device = "cuda:0"
|
|
if len(sys.argv) > 1:
|
|
checkpoint = sys.argv[1]
|
|
else:
|
|
checkpoint = find_latest_checkpoint()
|
|
if checkpoint is None:
|
|
print("No checkpoint found!")
|
|
sys.exit(1)
|
|
|
|
model, config = load_model(checkpoint, device)
|
|
tokenizer = get_tokenizer()
|
|
|
|
prompts = [
|
|
"The meaning of life is",
|
|
"In machine learning, a neural network",
|
|
"The capital of France is",
|
|
"Once upon a time, there was a",
|
|
"To solve a quadratic equation, you need to",
|
|
"The theory of relativity explains that",
|
|
"Python is a programming language that",
|
|
"The sun rises in the east and",
|
|
]
|
|
|
|
print("\n" + "=" * 70)
|
|
print(" INFERENCE — 1B Transformer (Single GPU)")
|
|
print("=" * 70)
|
|
|
|
for prompt in prompts:
|
|
print(f"\n{'─' * 60}")
|
|
print(f"PROMPT: {prompt}")
|
|
print(f"{'─' * 60}")
|
|
text, n_tok, tps = generate(model, tokenizer, prompt,
|
|
max_new_tokens=150, temperature=0.8,
|
|
top_k=50, device=device)
|
|
generated = text[len(prompt):]
|
|
print(f"OUTPUT:{generated}")
|
|
print(f" [{n_tok} tokens, {tps:.1f} tok/s]")
|
|
|
|
print("\n" + "=" * 70)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|