| import math |
| from dataclasses import dataclass |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
| from torch.utils.checkpoint import checkpoint |
|
|
|
|
| @dataclass |
| class Config: |
| vocab_size: int = 184704 |
| hidden_size: int = 2048 |
| intermediate_size: int = 6144 |
| num_hidden_layers: int = 28 |
| num_attention_heads: int = 16 |
| num_key_value_heads: int = 8 |
| head_dim: int = 128 |
| rope_theta: float = 1_000_000.0 |
| rms_norm_eps: float = 1e-6 |
| max_position_embeddings: int = 24576 |
| latent_dim: int = 64 |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim, eps): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dim)); self.eps = eps |
|
|
| def forward(self, x): |
| return F.rms_norm(x, self.weight.shape, self.weight.to(x.dtype), self.eps) |
|
|
|
|
| def rope(positions, head_dim, theta, device): |
| inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim)) |
| ang = positions.float().to(device)[:, None] * inv[None] |
| return ang.cos(), ang.sin() |
|
|
|
|
| def apply_rope(x, cos, sin): |
| f = x.float(); h = f.shape[-1] // 2 |
| lo, hi = f[..., :h], f[..., h:] |
| return torch.cat([lo * cos - hi * sin, hi * cos + lo * sin], dim=-1).to(x.dtype) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.h, self.kv, self.d = c.num_attention_heads, c.num_key_value_heads, c.head_dim |
| self.qkv_proj = nn.Linear(c.hidden_size, (self.h + 2 * self.kv) * self.d, bias=False) |
| self.o_proj = nn.Linear(self.h * self.d, c.hidden_size, bias=False) |
| self.q_norm, self.k_norm = RMSNorm(self.d, c.rms_norm_eps), RMSNorm(self.d, c.rms_norm_eps) |
|
|
| def qkv(self, x, cos, sin): |
| B, T, _ = x.shape |
| q, k, v = self.qkv_proj(x).split((self.h * self.d, self.kv * self.d, self.kv * self.d), dim=-1) |
| q = self.q_norm(q.view(B, T, self.h, self.d)).transpose(1, 2) |
| k = self.k_norm(k.view(B, T, self.kv, self.d)).transpose(1, 2) |
| v = v.view(B, T, self.kv, self.d).transpose(1, 2) |
| return apply_rope(q, cos, sin), apply_rope(k, cos, sin), v |
|
|
| def attend(self, q, k, v, mask=None): |
| o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, enable_gqa=self.h != self.kv) |
| return self.o_proj(o.transpose(1, 2).reshape(q.shape[0], q.shape[2], -1)) |
|
|
|
|
| class MLP(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.gate_up_proj = nn.Linear(c.hidden_size, 2 * c.intermediate_size, bias=False) |
| self.down_proj = nn.Linear(c.intermediate_size, c.hidden_size, bias=False) |
|
|
| def forward(self, x): |
| g, u = self.gate_up_proj(x).chunk(2, dim=-1) |
| return self.down_proj(F.silu(g) * u) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.input_layernorm = RMSNorm(c.hidden_size, c.rms_norm_eps) |
| self.self_attn = Attention(c) |
| self.post_attention_layernorm = RMSNorm(c.hidden_size, c.rms_norm_eps) |
| self.mlp = MLP(c) |
|
|
|
|
| def planner_block(layer, x, cos, sin, mask): |
| q, k, v = layer.self_attn.qkv(layer.input_layernorm(x), cos, sin) |
| x = x + layer.self_attn.attend(q, k, v, mask) |
| return x + layer.mlp(layer.post_attention_layernorm(x)) |
|
|
|
|
| def decoder_block(layer, s, cos, sin, prefix_k, prefix_v): |
| q, k, v = layer.self_attn.qkv(layer.input_layernorm(s), cos, sin) |
| k = torch.cat((prefix_k.to(k.dtype), k), dim=2); v = torch.cat((prefix_v.to(v.dtype), v), dim=2) |
| s = s + layer.self_attn.attend(q, k, v) |
| return s + layer.mlp(layer.post_attention_layernorm(s)) |
|
|
|
|
| class Planner(nn.Module): |
| def __init__(self, c: Config): |
| super().__init__() |
| self.c = c |
| self.embed_tokens = nn.Embedding(c.vocab_size, c.hidden_size) |
| self.layers = nn.ModuleList(Block(c) for _ in range(c.num_hidden_layers)) |
| self.norm = RMSNorm(c.hidden_size, c.rms_norm_eps) |
| self.lm_head = nn.Linear(c.hidden_size, c.vocab_size, bias=False) |
|
|
| def forward(self, ids, mask): |
| x = self.embed_tokens(ids) |
| cos, sin = rope(torch.arange(ids.shape[1]), self.c.head_dim, self.c.rope_theta, x.device) |
| mask = mask.to(x.dtype) |
| for layer in self.layers: |
| x = checkpoint(planner_block, layer, x, cos, sin, mask, use_reentrant=False) |
| return self.norm(x) |
|
|
| @torch.no_grad() |
| def prefill(self, ids, mask): |
| x = self.embed_tokens(ids[None]); T = ids.shape[0] |
| cos, sin = rope(torch.arange(T), self.c.head_dim, self.c.rope_theta, x.device) |
| ks, vs = [], [] |
| for layer in self.layers: |
| q, k, v = layer.self_attn.qkv(layer.input_layernorm(x), cos, sin) |
| x = x + layer.self_attn.attend(q, k, v, mask.to(x.dtype)) |
| x = x + layer.mlp(layer.post_attention_layernorm(x)) |
| ks.append(k); vs.append(v) |
| return ks, vs, self.norm(x) |
|
|
| @torch.no_grad() |
| def step(self, token, pos, kbuf, vbuf): |
| x = self.embed_tokens(torch.tensor([[token]], device=self.lm_head.weight.device)) |
| cos, sin = rope(torch.tensor([pos]), self.c.head_dim, self.c.rope_theta, x.device) |
| for i, layer in enumerate(self.layers): |
| q, k, v = layer.self_attn.qkv(layer.input_layernorm(x), cos, sin) |
| kbuf[i][:, :, pos] = k[:, :, 0]; vbuf[i][:, :, pos] = v[:, :, 0] |
| x = x + layer.self_attn.attend(q, kbuf[i][:, :, :pos + 1], vbuf[i][:, :, :pos + 1]) |
| x = x + layer.mlp(layer.post_attention_layernorm(x)) |
| return self.lm_head(self.norm(x)[:, -1]).float()[0] |
|
|
|
|
| class TimestepEmbedder(nn.Module): |
| def __init__(self, hidden): |
| super().__init__() |
| self.mlp = nn.Sequential(nn.Linear(256, hidden), nn.SiLU(), nn.Linear(hidden, hidden)) |
|
|
| def forward(self, t): |
| half = 128 |
| freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32, device=t.device) / half) |
| a = t.float()[:, None] * freqs[None] |
| return self.mlp(torch.cat([a.cos(), a.sin()], dim=-1).to(self.mlp[0].weight.dtype)) |
|
|
|
|
| def positions(n, dim, device): |
| pos = torch.arange(n, dtype=torch.float32, device=device)[:, None] |
| ang = pos * torch.exp(torch.arange(0, dim, 2, dtype=torch.float32, device=device) * (-math.log(10000.0) / dim)) |
| pe = torch.zeros(n, dim, device=device) |
| pe[:, 0::2], pe[:, 1::2] = ang.sin(), ang.cos() |
| return pe |
|
|
|
|
| class _Stack(nn.Module): |
| def __init__(self, c): |
| super().__init__() |
| self.layers = nn.ModuleList(Block(c) for _ in range(c.num_hidden_layers)) |
| self.norm = RMSNorm(c.hidden_size, c.rms_norm_eps) |
|
|
|
|
| class Decoder(nn.Module): |
| def __init__(self, c: Config): |
| super().__init__() |
| self.c = c |
| self.model = _Stack(c) |
| self.vae2llm = nn.Linear(c.latent_dim, c.hidden_size) |
| self.llm2vae = nn.Linear(c.hidden_size, c.latent_dim) |
| self.time_embedder = TimestepEmbedder(c.hidden_size) |
| self._pe = None |
|
|
| @torch.no_grad() |
| def velocity(self, x, sigma, prefix_k, prefix_v): |
| return self(x, sigma, prefix_k, prefix_v) |
|
|
| def forward(self, x, sigma, prefix_k, prefix_v, checkpointing=False): |
| dt = self.vae2llm.weight.dtype |
| A = prefix_k[0].shape[2]; n = x.shape[1] + 2 |
| s = F.pad(x.to(dt), (0, 0, 1, 1)) |
| if self._pe is None or self._pe.shape[0] < n or self._pe.device != x.device: |
| self._pe = positions(max(n, 1024), self.c.hidden_size, x.device).to(dt) |
| s = self.vae2llm(s) + self.time_embedder(sigma.to(dt))[:, None] + self._pe[:n][None] |
| cos, sin = rope(torch.arange(A, A + n), self.c.head_dim, self.c.rope_theta, x.device) |
| for i, layer in enumerate(self.model.layers): |
| if checkpointing: |
| s = checkpoint(decoder_block, layer, s, cos, sin, prefix_k[i], prefix_v[i], use_reentrant=False) |
| else: |
| s = decoder_block(layer, s, cos, sin, prefix_k[i], prefix_v[i]) |
| return self.llm2vae(self.model.norm(s))[:, 1:-1].float() |
|
|
