|
|
"""Nanochat model implementation and inference utilities.""" |
|
|
|
|
|
from __future__ import annotations |
|
|
|
|
|
import json |
|
|
import math |
|
|
import pickle |
|
|
from dataclasses import dataclass |
|
|
from pathlib import Path |
|
|
from typing import TYPE_CHECKING |
|
|
|
|
|
import torch |
|
|
import torch.nn.functional as F |
|
|
from torch import nn |
|
|
|
|
|
if TYPE_CHECKING: |
|
|
from collections.abc import Generator |
|
|
|
|
|
|
|
|
@dataclass |
|
|
class GPTConfig: |
|
|
"""Configuration for GPT model architecture. |
|
|
|
|
|
Attributes: |
|
|
sequence_len: Maximum sequence length |
|
|
vocab_size: Size of vocabulary |
|
|
n_layer: Number of transformer layers |
|
|
n_head: Number of attention heads |
|
|
n_kv_head: Number of key-value heads |
|
|
n_embd: Embedding dimension |
|
|
|
|
|
""" |
|
|
|
|
|
sequence_len: int = 1024 |
|
|
vocab_size: int = 50304 |
|
|
n_layer: int = 12 |
|
|
n_head: int = 6 |
|
|
n_kv_head: int = 6 |
|
|
n_embd: int = 768 |
|
|
|
|
|
|
|
|
def norm(x: torch.Tensor) -> torch.Tensor: |
|
|
"""Apply RMS normalization to input tensor.""" |
|
|
return F.rms_norm(x, (x.size(-1),)) |
|
|
|
|
|
|
|
|
_EXPECTED_NDIM = 4 |
|
|
|
|
|
|
|
|
def apply_rotary_emb( |
|
|
x: torch.Tensor, |
|
|
cos: torch.Tensor, |
|
|
sin: torch.Tensor, |
|
|
) -> torch.Tensor: |
|
|
"""Apply rotary positional embeddings to input tensor. |
|
|
|
|
|
Args: |
|
|
x: Input tensor of shape (batch, seq_len, n_heads, head_dim) |
|
|
cos: Cosine component of rotary embeddings |
|
|
sin: Sine component of rotary embeddings |
|
|
|
|
|
Returns: |
|
|
Tensor with rotary embeddings applied |
|
|
|
|
|
""" |
|
|
assert x.ndim == _EXPECTED_NDIM |
|
|
d = x.shape[3] // 2 |
|
|
x1, x2 = x[..., :d], x[..., d:] |
|
|
y1 = x1 * cos + x2 * sin |
|
|
y2 = x1 * (-sin) + x2 * cos |
|
|
return torch.cat([y1, y2], 3).to(x.dtype) |
|
|
|
|
|
|
|
|
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: |
|
|
"""Repeat key/value tensors for multi-head attention. |
|
|
|
|
|
Args: |
|
|
x: Input tensor of shape (batch, n_kv_heads, seq_len, head_dim) |
|
|
n_rep: Number of times to repeat |
|
|
|
|
|
Returns: |
|
|
Tensor with repeated key/value heads |
|
|
|
|
|
""" |
|
|
if n_rep == 1: |
|
|
return x |
|
|
bs, n_kv_heads, slen, head_dim = x.shape |
|
|
return ( |
|
|
x[:, :, None, :, :] |
|
|
.expand(bs, n_kv_heads, n_rep, slen, head_dim) |
|
|
.reshape(bs, n_kv_heads * n_rep, slen, head_dim) |
|
|
) |
|
|
|
|
|
|
|
|
class CausalSelfAttention(nn.Module): |
|
|
"""Causal self-attention with rotary position embeddings.""" |
|
|
|
|
|
def __init__(self, config: GPTConfig, layer_idx: int) -> None: |
|
|
"""Initialize attention layer. |
|
|
|
|
|
Args: |
|
|
config: Model configuration |
|
|
layer_idx: Layer index for KV cache |
|
|
|
|
|
""" |
|
|
super().__init__() |
|
|
self.layer_idx = layer_idx |
|
|
self.n_head = config.n_head |
|
|
self.n_kv_head = config.n_kv_head |
|
|
self.n_embd = config.n_embd |
|
|
self.head_dim = self.n_embd // self.n_head |
|
|
assert self.n_embd % self.n_head == 0 |
|
|
assert self.n_kv_head <= self.n_head |
|
|
assert self.n_head % self.n_kv_head == 0 |
|
|
self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False) |
|
|
self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False) |
|
|
self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False) |
|
|
self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False) |
|
|
|
|
|
def forward( |
|
|
self, |
|
|
x: torch.Tensor, |
|
|
cos_sin: tuple[torch.Tensor, torch.Tensor], |
|
|
kv_cache: object | None, |
|
|
) -> torch.Tensor: |
|
|
"""Forward pass of attention layer. |
|
|
|
|
|
Args: |
|
|
x: Input tensor |
|
|
cos_sin: Tuple of (cos, sin) rotary embeddings |
|
|
kv_cache: Optional KV cache for generation |
|
|
|
|
|
Returns: |
|
|
Output tensor after attention |
|
|
|
|
|
""" |
|
|
b, t, _c = x.size() |
|
|
q = self.c_q(x).view(b, t, self.n_head, self.head_dim) |
|
|
k = self.c_k(x).view(b, t, self.n_kv_head, self.head_dim) |
|
|
v = self.c_v(x).view(b, t, self.n_kv_head, self.head_dim) |
|
|
cos, sin = cos_sin |
|
|
q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin) |
|
|
q, k = norm(q), norm(k) |
|
|
q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) |
|
|
if kv_cache is not None: |
|
|
k, v = kv_cache.insert_kv(self.layer_idx, k, v) |
|
|
tq = q.size(2) |
|
|
tk = k.size(2) |
|
|
nrep = self.n_head // self.n_kv_head |
|
|
k, v = repeat_kv(k, nrep), repeat_kv(v, nrep) |
|
|
if kv_cache is None or tq == tk: |
|
|
y = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
|
|
elif tq == 1: |
|
|
y = F.scaled_dot_product_attention(q, k, v, is_causal=False) |
|
|
else: |
|
|
attn_mask = torch.zeros((tq, tk), dtype=torch.bool, device=q.device) |
|
|
prefix_len = tk - tq |
|
|
if prefix_len > 0: |
|
|
attn_mask[:, :prefix_len] = True |
|
|
attn_mask[:, prefix_len:] = torch.tril( |
|
|
torch.ones((tq, tq), dtype=torch.bool, device=q.device), |
|
|
) |
|
|
y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) |
|
|
y = y.transpose(1, 2).contiguous().view(b, t, -1) |
|
|
return self.c_proj(y) |
|
|
|
|
|
|
|
|
class MLP(nn.Module): |
|
|
"""Multi-layer perceptron with squared ReLU activation.""" |
|
|
|
|
|
def __init__(self, config: GPTConfig) -> None: |
|
|
"""Initialize MLP layer. |
|
|
|
|
|
Args: |
|
|
config: Model configuration |
|
|
|
|
|
""" |
|
|
super().__init__() |
|
|
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False) |
|
|
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False) |
|
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
|
"""Forward pass of MLP. |
|
|
|
|
|
Args: |
|
|
x: Input tensor |
|
|
|
|
|
Returns: |
|
|
Output tensor after MLP transformation |
|
|
|
|
|
""" |
|
|
x = self.c_fc(x) |
|
|
x = F.relu(x).square() |
|
|
return self.c_proj(x) |
|
|
|
|
|
|
|
|
class Block(nn.Module): |
|
|
"""Transformer block with attention and MLP.""" |
|
|
|
|
|
def __init__(self, config: GPTConfig, layer_idx: int) -> None: |
|
|
"""Initialize transformer block. |
|
|
|
|
|
Args: |
|
|
config: Model configuration |
|
|
layer_idx: Layer index |
|
|
|
|
|
""" |
|
|
super().__init__() |
|
|
self.attn = CausalSelfAttention(config, layer_idx) |
|
|
self.mlp = MLP(config) |
|
|
|
|
|
def forward( |
|
|
self, |
|
|
x: torch.Tensor, |
|
|
cos_sin: tuple[torch.Tensor, torch.Tensor], |
|
|
kv_cache: object | None, |
|
|
) -> torch.Tensor: |
|
|
"""Forward pass of transformer block. |
|
|
|
|
|
Args: |
|
|
x: Input tensor |
|
|
cos_sin: Tuple of (cos, sin) rotary embeddings |
|
|
kv_cache: Optional KV cache for generation |
|
|
|
|
|
Returns: |
|
|
Output tensor after block transformation |
|
|
|
|
|
""" |
|
|
x = x + self.attn(norm(x), cos_sin, kv_cache) |
|
|
return x + self.mlp(norm(x)) |
|
|
|
|
|
|
|
|
class GPT(nn.Module): |
|
|
"""GPT model with rotary position embeddings.""" |
|
|
|
|
|
def __init__(self, config: GPTConfig) -> None: |
|
|
"""Initialize GPT model. |
|
|
|
|
|
Args: |
|
|
config: Model configuration |
|
|
|
|
|
""" |
|
|
super().__init__() |
|
|
self.config = config |
|
|
self.transformer = nn.ModuleDict( |
|
|
{ |
|
|
"wte": nn.Embedding(config.vocab_size, config.n_embd), |
|
|
"h": nn.ModuleList( |
|
|
[Block(config, layer_idx) for layer_idx in range(config.n_layer)], |
|
|
), |
|
|
}, |
|
|
) |
|
|
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) |
|
|
self.rotary_seq_len = config.sequence_len * 10 |
|
|
head_dim = config.n_embd // config.n_head |
|
|
cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim) |
|
|
self.register_buffer("cos", cos, persistent=False) |
|
|
self.register_buffer("sin", sin, persistent=False) |
|
|
self.transformer.wte.to(dtype=torch.bfloat16) |
|
|
|
|
|
def init_weights(self) -> None: |
|
|
"""Initialize model weights.""" |
|
|
self.apply(self._init_weights) |
|
|
torch.nn.init.zeros_(self.lm_head.weight) |
|
|
for block in self.transformer.h: |
|
|
torch.nn.init.zeros_(block.mlp.c_proj.weight) |
|
|
torch.nn.init.zeros_(block.attn.c_proj.weight) |
|
|
head_dim = self.config.n_embd // self.config.n_head |
|
|
cos, sin = self._precompute_rotary_embeddings(self.rotary_seq_len, head_dim) |
|
|
self.cos, self.sin = cos, sin |
|
|
|
|
|
def _init_weights(self, module: nn.Module) -> None: |
|
|
"""Initialize weights for a single module. |
|
|
|
|
|
Args: |
|
|
module: Module to initialize |
|
|
|
|
|
""" |
|
|
if isinstance(module, nn.Linear): |
|
|
fan_out = module.weight.size(0) |
|
|
fan_in = module.weight.size(1) |
|
|
std = 1.0 / math.sqrt(fan_in) * min(1.0, math.sqrt(fan_out / fan_in)) |
|
|
torch.nn.init.normal_(module.weight, mean=0.0, std=std) |
|
|
if module.bias is not None: |
|
|
torch.nn.init.zeros_(module.bias) |
|
|
elif isinstance(module, nn.Embedding): |
|
|
torch.nn.init.normal_(module.weight, mean=0.0, std=1.0) |
|
|
|
|
|
def _precompute_rotary_embeddings( |
|
|
self, |
|
|
seq_len: int, |
|
|
head_dim: int, |
|
|
base: int = 10000, |
|
|
device: torch.device | str | None = None, |
|
|
) -> tuple[torch.Tensor, torch.Tensor]: |
|
|
"""Precompute rotary position embeddings. |
|
|
|
|
|
Args: |
|
|
seq_len: Maximum sequence length |
|
|
head_dim: Dimension of attention heads |
|
|
base: Base for frequency calculation |
|
|
device: Device to place tensors on |
|
|
|
|
|
Returns: |
|
|
Tuple of (cos, sin) tensors for rotary embeddings |
|
|
|
|
|
""" |
|
|
if device is None: |
|
|
device = self.transformer.wte.weight.device |
|
|
channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) |
|
|
inv_freq = 1.0 / (base ** (channel_range / head_dim)) |
|
|
t = torch.arange(seq_len, dtype=torch.float32, device=device) |
|
|
freqs = torch.outer(t, inv_freq) |
|
|
cos, sin = freqs.cos(), freqs.sin() |
|
|
cos, sin = cos.bfloat16(), sin.bfloat16() |
|
|
return cos[None, :, None, :], sin[None, :, None, :] |
|
|
|
|
|
def forward( |
|
|
self, |
|
|
idx: torch.Tensor, |
|
|
targets: torch.Tensor | None = None, |
|
|
kv_cache: object | None = None, |
|
|
) -> torch.Tensor: |
|
|
"""Forward pass of GPT model. |
|
|
|
|
|
Args: |
|
|
idx: Input token indices |
|
|
targets: Target token indices (unused in this implementation) |
|
|
kv_cache: Optional KV cache for generation |
|
|
|
|
|
Returns: |
|
|
Logits for next token prediction |
|
|
|
|
|
""" |
|
|
_b, t = idx.size() |
|
|
assert self.cos.size(1) >= t |
|
|
t0 = 0 if kv_cache is None else kv_cache.get_pos() |
|
|
cos_sin = self.cos[:, t0 : t0 + t], self.sin[:, t0 : t0 + t] |
|
|
x = self.transformer.wte(idx) |
|
|
x = norm(x) |
|
|
for block in self.transformer.h: |
|
|
x = block(x, cos_sin, kv_cache) |
|
|
x = norm(x) |
|
|
softcap = 15 |
|
|
logits = self.lm_head(x) |
|
|
return softcap * torch.tanh(logits / softcap) |
|
|
|
|
|
|
|
|
class NanochatModel: |
|
|
"""Wrapper class for loading and running inference with the nanochat model.""" |
|
|
|
|
|
def __init__(self, model_dir: str, device: str = "cpu") -> None: |
|
|
"""Initialize the NanochatModel. |
|
|
|
|
|
Args: |
|
|
model_dir: Directory containing model files |
|
|
device: Device to run inference on (default: "cpu") |
|
|
|
|
|
""" |
|
|
self.device = torch.device(device) |
|
|
self.model_dir = model_dir |
|
|
|
|
|
self.model = self._load_model() |
|
|
self.enc = self._load_tokenizer() |
|
|
self._setup_special_tokens() |
|
|
|
|
|
def _load_model(self) -> GPT: |
|
|
"""Load the model from the model directory.""" |
|
|
model_dir_path = Path(self.model_dir) |
|
|
model_files = list(model_dir_path.glob("model_*.pt")) |
|
|
if not model_files: |
|
|
msg = f"No model files found in {self.model_dir}" |
|
|
raise FileNotFoundError(msg) |
|
|
model_file = model_files[0] |
|
|
|
|
|
meta_files = list(model_dir_path.glob("meta_*.json")) |
|
|
if not meta_files: |
|
|
msg = f"No meta files found in {self.model_dir}" |
|
|
raise FileNotFoundError(msg) |
|
|
meta_file = meta_files[0] |
|
|
|
|
|
with meta_file.open() as f: |
|
|
meta = json.load(f) |
|
|
|
|
|
model_config_kwargs = meta["model_config"] |
|
|
|
|
|
model_config = GPTConfig(**model_config_kwargs) |
|
|
with torch.device("meta"): |
|
|
model = GPT(model_config) |
|
|
|
|
|
model_data = torch.load( |
|
|
model_file, |
|
|
map_location=self.device, |
|
|
weights_only=True, |
|
|
) |
|
|
model_data = {k.removeprefix("_orig_mod."): v for k, v in model_data.items()} |
|
|
|
|
|
model_data = { |
|
|
k: v.float() if v.dtype == torch.bfloat16 else v |
|
|
for k, v in model_data.items() |
|
|
} |
|
|
|
|
|
model.to_empty(device=self.device) |
|
|
model.init_weights() |
|
|
model.load_state_dict(model_data, strict=True, assign=True) |
|
|
model.eval() |
|
|
|
|
|
return model |
|
|
|
|
|
def _load_tokenizer(self) -> object: |
|
|
"""Load the tokenizer from the model directory. |
|
|
|
|
|
Returns: |
|
|
Loaded tokenizer object |
|
|
|
|
|
""" |
|
|
tokenizer_path = Path(self.model_dir) / "tokenizer.pkl" |
|
|
if not tokenizer_path.exists(): |
|
|
msg = f"Tokenizer not found at {tokenizer_path}" |
|
|
raise FileNotFoundError(msg) |
|
|
|
|
|
with tokenizer_path.open("rb") as f: |
|
|
return pickle.load(f) |
|
|
|
|
|
def _setup_special_tokens(self) -> None: |
|
|
"""Set up special token IDs for chat formatting.""" |
|
|
try: |
|
|
try: |
|
|
self.bos_token_id = self.enc.encode_single_token("<|bos|>") |
|
|
except KeyError: |
|
|
self.bos_token_id = self.enc.encode_single_token("<|endoftext|>") |
|
|
|
|
|
self.user_start_id = self.enc.encode_single_token("<|user_start|>") |
|
|
self.user_end_id = self.enc.encode_single_token("<|user_end|>") |
|
|
self.assistant_start_id = self.enc.encode_single_token( |
|
|
"<|assistant_start|>", |
|
|
) |
|
|
self.assistant_end_id = self.enc.encode_single_token("<|assistant_end|>") |
|
|
self.stop_tokens = {self.bos_token_id, self.assistant_end_id} |
|
|
except KeyError as e: |
|
|
msg = f"Required special token missing from tokenizer: {e}" |
|
|
raise ValueError(msg) from e |
|
|
|
|
|
def format_prompt(self, message: str) -> list[int]: |
|
|
"""Format a user message using chat format. |
|
|
|
|
|
Args: |
|
|
message: User's input message |
|
|
|
|
|
Returns: |
|
|
List of token IDs formatted for chat |
|
|
|
|
|
""" |
|
|
prompt_tokens = self.enc.encode_ordinary(message) |
|
|
return [ |
|
|
self.bos_token_id, |
|
|
self.user_start_id, |
|
|
*prompt_tokens, |
|
|
self.user_end_id, |
|
|
self.assistant_start_id, |
|
|
] |
|
|
|
|
|
def format_conversation(self, history: list[dict[str, str]]) -> list[int]: |
|
|
"""Format a multi-turn conversation using chat format. |
|
|
|
|
|
Args: |
|
|
history: List of message dictionaries with 'role' and 'content' keys |
|
|
role can be 'user' or 'assistant' |
|
|
|
|
|
Returns: |
|
|
List of token IDs formatted for multi-turn chat |
|
|
|
|
|
""" |
|
|
tokens = [self.bos_token_id] |
|
|
|
|
|
for message in history: |
|
|
role = message.get("role") |
|
|
content = message.get("content", "") |
|
|
content_tokens = self.enc.encode_ordinary(content) |
|
|
|
|
|
if role == "user": |
|
|
tokens.extend([ |
|
|
self.user_start_id, |
|
|
*content_tokens, |
|
|
self.user_end_id, |
|
|
]) |
|
|
elif role == "assistant": |
|
|
tokens.extend([ |
|
|
self.assistant_start_id, |
|
|
*content_tokens, |
|
|
self.assistant_end_id, |
|
|
]) |
|
|
|
|
|
tokens.append(self.assistant_start_id) |
|
|
|
|
|
return tokens |
|
|
|
|
|
def generate( |
|
|
self, |
|
|
prompt: str | None = None, |
|
|
history: list[dict[str, str]] | None = None, |
|
|
max_tokens: int = 512, |
|
|
temperature: float = 0.8, |
|
|
top_k: int = 50, |
|
|
) -> Generator[str, None, None]: |
|
|
"""Generate text from a prompt or conversation history. |
|
|
|
|
|
Args: |
|
|
prompt: The input text prompt (for single-turn) |
|
|
history: List of message dicts with 'role' and 'content' (for multi-turn) |
|
|
max_tokens: Maximum number of tokens to generate |
|
|
temperature: Sampling temperature |
|
|
top_k: Top-k sampling parameter |
|
|
|
|
|
Yields: |
|
|
Decoded token strings |
|
|
|
|
|
""" |
|
|
if history is not None: |
|
|
input_ids = self.format_conversation(history) |
|
|
elif prompt is not None: |
|
|
input_ids = self.format_prompt(prompt) |
|
|
else: |
|
|
msg = "Either prompt or history must be provided" |
|
|
raise ValueError(msg) |
|
|
|
|
|
x = torch.tensor([input_ids], dtype=torch.long, device=self.device) |
|
|
|
|
|
with torch.inference_mode(): |
|
|
for _ in range(max_tokens): |
|
|
logits = self.model(x) |
|
|
|
|
|
logits = logits[:, -1, :] |
|
|
|
|
|
logits = logits / temperature |
|
|
|
|
|
if top_k > 0: |
|
|
v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
|
|
logits[logits < v[:, [-1]]] = -float("inf") |
|
|
|
|
|
probs = F.softmax(logits, dim=-1) |
|
|
next_token = torch.multinomial(probs, num_samples=1) |
|
|
|
|
|
if next_token.item() in self.stop_tokens: |
|
|
break |
|
|
|
|
|
token_str = self.enc.decode([next_token.item()]) |
|
|
yield token_str |
|
|
|
|
|
x = torch.cat([x, next_token], dim=1) |
