| |
| |
| |
| |
|
|
| from dataclasses import dataclass |
| import inspect |
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class CombinedOptimizer: |
| def __init__(self, *optimizers): |
| self.optimizers = [opt for opt in optimizers if opt is not None] |
| self.param_groups = [] |
| for opt in self.optimizers: |
| self.param_groups.extend(opt.param_groups) |
|
|
| def step(self, *args, **kwargs): |
| out = None |
| for opt in self.optimizers: |
| out = opt.step(*args, **kwargs) |
| return out |
|
|
| def zero_grad(self, *args, **kwargs): |
| for opt in self.optimizers: |
| opt.zero_grad(*args, **kwargs) |
|
|
| def state_dict(self): |
| return {"optimizers": [opt.state_dict() for opt in self.optimizers]} |
|
|
| def load_state_dict(self, state_dict): |
| states = state_dict["optimizers"] |
| if len(states) != len(self.optimizers): |
| raise ValueError(f"optimizer count mismatch: {len(states)} != {len(self.optimizers)}") |
| for opt, state in zip(self.optimizers, states): |
| opt.load_state_dict(state) |
|
|
|
|
| @dataclass |
| class LyraConfig: |
| block_size: int = 1024 |
| vocab_size: int = 50304 |
| n_layer: int = 12 |
| n_head: int = 10 |
| n_embd: int = 640 |
| rope_base: float = 10000.0 |
| use_qk_norm: bool = True |
| logit_softcap: float = 0.0 |
| mlp_hidden: int = 2048 |
| dropout: float = 0.0 |
| bias: bool = False |
|
|
|
|
| class ChannelRMSNorm(nn.Module): |
| def __init__(self, dim: int, eps: float = 1e-6): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dim)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| x_fp = x.float() |
| rms = x_fp.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt() |
| return (x_fp * rms).to(x.dtype) * self.weight |
|
|
|
|
| def build_rotary_cache(head_dim: int, max_seq_len: int, base: float = 10000.0): |
| inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) |
| t = torch.arange(max_seq_len, dtype=torch.float32) |
| freqs = torch.outer(t, inv_freq) |
| return freqs.cos(), freqs.sin() |
|
|
|
|
| def apply_rotary_embedding(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| x1, x2 = x.chunk(2, dim=-1) |
| cos = cos[None, None, :, :].to(x.dtype) |
| sin = sin[None, None, :, :].to(x.dtype) |
| return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) |
|
|
|
|
| class RotarySelfAttention(nn.Module): |
| def __init__(self, config: LyraConfig): |
| super().__init__() |
| assert config.n_embd % config.n_head == 0 |
| self.n_head = config.n_head |
| self.n_embd = config.n_embd |
| self.head_dim = config.n_embd // config.n_head |
| assert self.head_dim % 2 == 0, "head_dim must be even for RoPE" |
|
|
| self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False) |
| self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False) |
| self.c_proj.NANOGPT_SCALE_INIT = 1 |
|
|
| self.use_qk_norm = config.use_qk_norm |
| if self.use_qk_norm: |
| self.q_norm = ChannelRMSNorm(self.head_dim) |
| self.k_norm = ChannelRMSNorm(self.head_dim) |
|
|
| def forward(self, x, cos, sin): |
| B, T, C = x.size() |
| qkv = self.c_attn(x) |
| q, k, v = qkv.split(self.n_embd, dim=2) |
| q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
| v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2) |
|
|
| if self.use_qk_norm: |
| q = self.q_norm(q) |
| k = self.k_norm(k) |
|
|
| q = apply_rotary_embedding(q, cos[:T], sin[:T]) |
| k = apply_rotary_embedding(k, cos[:T], sin[:T]) |
|
|
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| y = y.transpose(1, 2).contiguous().view(B, T, C) |
| return self.c_proj(y) |
|
|
|
|
| class GatedFeedForward(nn.Module): |
| def __init__(self, config: LyraConfig): |
| super().__init__() |
| hidden = config.mlp_hidden |
| self.c_gate = nn.Linear(config.n_embd, hidden, bias=False) |
| self.c_up = nn.Linear(config.n_embd, hidden, bias=False) |
| self.c_proj = nn.Linear(hidden, config.n_embd, bias=False) |
| self.c_proj.NANOGPT_SCALE_INIT = 1 |
|
|
| def forward(self, x): |
| return self.c_proj(F.silu(self.c_gate(x)) * self.c_up(x)) |
|
|
|
|
| class DecoderLayer(nn.Module): |
| def __init__(self, config: LyraConfig): |
| super().__init__() |
| self.ln_1 = ChannelRMSNorm(config.n_embd) |
| self.attn = RotarySelfAttention(config) |
| self.ln_2 = ChannelRMSNorm(config.n_embd) |
| self.mlp = GatedFeedForward(config) |
|
|
| def forward(self, x, cos, sin): |
| x = x + self.attn(self.ln_1(x), cos, sin) |
| x = x + self.mlp(self.ln_2(x)) |
| return x |
|
|
|
|
| class LyraLM(nn.Module): |
| def __init__(self, config: LyraConfig): |
| super().__init__() |
| self.config = config |
|
|
| self.transformer = nn.ModuleDict(dict( |
| wte=nn.Embedding(config.vocab_size, config.n_embd), |
| h=nn.ModuleList([DecoderLayer(config) for _ in range(config.n_layer)]), |
| ln_f=ChannelRMSNorm(config.n_embd), |
| )) |
| self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) |
| self.transformer.wte.weight = self.lm_head.weight |
|
|
| head_dim = config.n_embd // config.n_head |
| cos, sin = build_rotary_cache(head_dim, config.block_size, config.rope_base) |
| self.register_buffer("rope_cos", cos, persistent=False) |
| self.register_buffer("rope_sin", sin, persistent=False) |
|
|
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, module): |
| if isinstance(module, nn.Linear): |
| std = 0.02 |
| if hasattr(module, "NANOGPT_SCALE_INIT"): |
| std *= (2 * self.config.n_layer) ** -0.5 |
| 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=0.02) |
|
|
| def forward(self, idx, targets=None): |
| B, T = idx.size() |
| assert T <= self.config.block_size, ( |
| f"sequence length {T} > block_size {self.config.block_size}" |
| ) |
| x = self.transformer.wte(idx) |
| for block in self.transformer.h: |
| x = block(x, self.rope_cos, self.rope_sin) |
| x = self.transformer.ln_f(x) |
| logits = self.lm_head(x) |
| if self.config.logit_softcap > 0: |
| cap = self.config.logit_softcap |
| logits = cap * torch.tanh(logits / cap) |
| logits = logits[:, :, :50257] |
| if targets is None: |
| return logits |
| loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1)) |
| return logits, loss |
|
|
| def configure_optimizers( |
| self, |
| weight_decay, |
| learning_rate, |
| betas, |
| device_type, |
| use_muon=False, |
| muon_lr=4e-4, |
| muon_momentum=0.95, |
| muon_ns_steps=5, |
| muon_nesterov=True, |
| muon_adjust_lr_fn="original", |
| embed_lr=0.0, |
| scalar_lr=0.0, |
| ): |
| param_dict = {pn: p for pn, p in self.named_parameters() if p.requires_grad} |
| if use_muon: |
| embed_names = ("transformer.wte", "lm_head") |
| muon_items = [ |
| (n, p) for n, p in param_dict.items() |
| if p.dim() == 2 and not any(name in n for name in embed_names) |
| ] |
| muon_ids = {id(p) for _, p in muon_items} |
| adamw_items = [(n, p) for n, p in param_dict.items() if id(p) not in muon_ids] |
| embed_params = [p for n, p in adamw_items if p.dim() >= 2] |
| scalar_params = [p for n, p in adamw_items if p.dim() < 2] |
|
|
| print( |
| f"Muon parameter tensors: {len(muon_items)}, " |
| f"with {sum(p.numel() for _, p in muon_items):,} parameters" |
| ) |
| print( |
| f"AdamW embed tensors: {len(embed_params)}, " |
| f"with {sum(p.numel() for p in embed_params):,} parameters" |
| ) |
| print( |
| f"AdamW scalar tensors: {len(scalar_params)}, " |
| f"with {sum(p.numel() for p in scalar_params):,} parameters" |
| ) |
|
|
| muon_optimizer = torch.optim.Muon( |
| [{"params": [p for _, p in muon_items], "lr": muon_lr, "initial_lr": muon_lr}], |
| lr=muon_lr, |
| weight_decay=weight_decay, |
| momentum=muon_momentum, |
| nesterov=muon_nesterov, |
| ns_steps=muon_ns_steps, |
| adjust_lr_fn=muon_adjust_lr_fn, |
| ) |
| adamw_groups = [] |
| if embed_params: |
| elr = embed_lr if embed_lr > 0 else learning_rate |
| adamw_groups.append({ |
| "params": embed_params, |
| "lr": elr, |
| "initial_lr": elr, |
| "weight_decay": 0.0, |
| }) |
| if scalar_params: |
| slr = scalar_lr if scalar_lr > 0 else learning_rate |
| adamw_groups.append({ |
| "params": scalar_params, |
| "lr": slr, |
| "initial_lr": slr, |
| "weight_decay": 0.0, |
| }) |
| fused_available = "fused" in inspect.signature(torch.optim.AdamW).parameters |
| use_fused = fused_available and device_type == "cuda" |
| adamw_optimizer = torch.optim.AdamW( |
| adamw_groups, |
| lr=learning_rate, |
| betas=betas, |
| **(dict(fused=True) if use_fused else {}), |
| ) |
| print( |
| "using Muon + AdamW: " |
| f"muon_lr={muon_lr}, adamw_lr={learning_rate}, fused AdamW={use_fused}" |
| ) |
| return CombinedOptimizer(muon_optimizer, adamw_optimizer) |
|
|
| decay_params = [p for p in param_dict.values() if p.dim() >= 2] |
| nodecay_params = [p for p in param_dict.values() if p.dim() < 2] |
| optim_groups = [ |
| {"params": decay_params, "weight_decay": weight_decay}, |
| {"params": nodecay_params, "weight_decay": 0.0}, |
| ] |
| num_decay_params = sum(p.numel() for p in decay_params) |
| num_nodecay_params = sum(p.numel() for p in nodecay_params) |
| print(f"num decayed parameter tensors: {len(decay_params)}, with {num_decay_params:,} parameters") |
| print(f"num non-decayed parameter tensors: {len(nodecay_params)}, with {num_nodecay_params:,} parameters") |
| fused_available = "fused" in inspect.signature(torch.optim.AdamW).parameters |
| use_fused = fused_available and device_type == "cuda" |
| extra_args = dict(fused=True) if use_fused else dict() |
| optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **extra_args) |
| print(f"using fused AdamW: {use_fused}") |
| return optimizer |
|
|
| def estimate_mfu(self, fwdbwd_per_iter, dt): |
| N = num_params(self) |
| cfg = self.config |
| L, H, Q, T = cfg.n_layer, cfg.n_head, cfg.n_embd // cfg.n_head, cfg.block_size |
| flops_per_token = 6 * N + 12 * L * H * Q * T |
| flops_per_fwdbwd = flops_per_token * T |
| flops_per_iter = flops_per_fwdbwd * fwdbwd_per_iter |
| flops_achieved = flops_per_iter * (1.0 / dt) |
| flops_promised = 312e12 |
| return flops_achieved / flops_promised |
|
|
|
|
| def num_params(model: LyraLM) -> int: |
| seen = set() |
| total = 0 |
| for p in model.parameters(): |
| if id(p) in seen: |
| continue |
| seen.add(id(p)) |
| total += p.numel() |
| return total |
|
|
|
|
| def strip_runtime_prefixes(state_dict): |
| clean = {} |
| for key, value in state_dict.items(): |
| for prefix in ("_orig_mod.", "module."): |
| if key.startswith(prefix): |
| key = key[len(prefix):] |
| clean[key] = value |
| return clean |
|
|
|
|
| def load_model(checkpoint_path: str, device: str = "cuda") -> nn.Module: |
| cfg = LyraConfig() |
| model = LyraLM(cfg) |
| checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False) |
| state_dict = checkpoint["model"] if isinstance(checkpoint, dict) and "model" in checkpoint else checkpoint |
| model.load_state_dict(strip_runtime_prefixes(state_dict), strict=True) |
| model.to(device) |
| model.eval() |
| return model |
|
|
|
|
| if __name__ == "__main__": |
| cfg = LyraConfig() |
| m = LyraLM(cfg) |
| n = num_params(m) |
| print(f"LyraConfig = {cfg}") |
| print(f"params: {n:,} ({n/1e6:.2f} M)") |
| assert n < 100_000_000, "OVER 100M PARAM CAP" |
| x = torch.randint(0, 50257, (2, 64)) |
| logits, loss = m(x, x) |
| print(f"in: {tuple(x.shape)} out: {tuple(logits.shape)} loss: {loss.item():.3f}") |
| print("ok") |
|
|
