| """ |
| MicroGLM: Tiny GLM-style model in PyTorch. |
| GLM (General Language Model) uses a prefix-LM architecture: |
| bidirectional attention on a prefix span + autoregressive generation on the rest. |
| This is a minimal implementation that fits in 6GB VRAM. |
| |
| Reference: "GLM: General Language Model Pretraining with Autoregressive Blank Infilling" (Du et al., 2021) |
| """ |
|
|
| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim, eps=1e-5): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(dim)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) * self.weight |
|
|
|
|
| class GLMAttention(nn.Module): |
| """ |
| Attention with support for 2D attention masks (for prefix-LM). |
| The mask has shape (B, 1, T, T) where: |
| - 0 means attend (bidirectional / causal allowed) |
| - -inf means blocked |
| """ |
| def __init__(self, n_embd, n_head, block_size, dropout): |
| super().__init__() |
| assert n_embd % n_head == 0 |
| self.n_head = n_head |
| self.head_dim = n_embd // n_head |
| self.qkv = nn.Linear(n_embd, n_embd * 3, bias=False) |
| self.proj = nn.Linear(n_embd, n_embd, bias=False) |
| self.attn_drop = nn.Dropout(dropout) |
| self.resid_drop = nn.Dropout(dropout) |
|
|
| def forward(self, x, attn_mask=None): |
| """ |
| x: (B, T, C) |
| attn_mask: (B, 1, T, T) or None (uses default causal mask) |
| """ |
| B, T, C = x.shape |
| qkv = self.qkv(x).reshape(B, T, 3, self.n_head, self.head_dim).permute(2, 0, 3, 1, 4) |
| q, k, v = qkv[0], qkv[1], qkv[2] |
| att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) |
|
|
| if attn_mask is not None: |
| |
| att = att + attn_mask |
| else: |
| |
| mask = torch.tril(torch.ones(T, T, device=x.device)).view(1, 1, T, T) |
| att = att.masked_fill(mask == 0, float('-inf')) |
|
|
| att = F.softmax(att, dim=-1) |
| att = self.attn_drop(att) |
| y = att @ v |
| y = y.transpose(1, 2).contiguous().view(B, T, C) |
| return self.resid_drop(self.proj(y)) |
|
|
|
|
| class GLMMLP(nn.Module): |
| def __init__(self, n_embd, dropout): |
| super().__init__() |
| self.fc = nn.Linear(n_embd, 4 * n_embd, bias=False) |
| self.proj = nn.Linear(4 * n_embd, n_embd, bias=False) |
| self.drop = nn.Dropout(dropout) |
|
|
| def forward(self, x): |
| x = F.gelu(self.fc(x)) |
| x = self.proj(x) |
| return self.drop(x) |
|
|
|
|
| class GLMBlock(nn.Module): |
| def __init__(self, n_embd, n_head, block_size, dropout): |
| super().__init__() |
| self.ln1 = RMSNorm(n_embd) |
| self.attn = GLMAttention(n_embd, n_head, block_size, dropout) |
| self.ln2 = RMSNorm(n_embd) |
| self.mlp = GLMMLP(n_embd, dropout) |
|
|
| def forward(self, x, attn_mask=None): |
| x = x + self.attn(self.ln1(x), attn_mask) |
| x = x + self.mlp(self.ln2(x)) |
| return x |
|
|
|
|
| class MicroGLM(nn.Module): |
| """ |
| Micro GLM: Prefix-LM decoder. |
| Supports two attention modes: |
| - Causal LM (default): standard autoregressive generation |
| - Prefix LM: bidirectional attention on first `prefix_len` tokens, causal on the rest |
| |
| This is controlled by passing a custom attention mask during forward(). |
| For training, we create a 2D mask where the prefix region is bidirectional |
| and the suffix region is causal. |
| """ |
|
|
| def __init__(self, vocab_size, block_size, n_layer=2, n_head=4, n_embd=128, dropout=0.1): |
| super().__init__() |
| self.block_size = block_size |
| self.wte = nn.Embedding(vocab_size, n_embd) |
| self.wpe = nn.Embedding(block_size, n_embd) |
| self.blocks = nn.ModuleList([GLMBlock(n_embd, n_head, block_size, dropout) for _ in range(n_layer)]) |
| self.ln_f = RMSNorm(n_embd) |
| self.lm_head = nn.Linear(n_embd, vocab_size, bias=False) |
| self.lm_head.weight = self.wte.weight |
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, module): |
| if isinstance(module, nn.Linear): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
| elif isinstance(module, nn.Embedding): |
| nn.init.normal_(module.weight, mean=0.0, std=0.02) |
|
|
| def _build_prefix_lm_mask(self, T, prefix_len, device): |
| """ |
| Build a 2D attention mask for prefix-LM. |
| - For positions i < prefix_len and j < prefix_len: bidirectional (0) |
| - For positions i >= prefix_len: causal (attend only to j <= i) |
| - All other positions: -inf |
| |
| Returns: (1, 1, T, T) mask where 0 = allowed, -inf = blocked |
| """ |
| |
| mask = torch.tril(torch.ones(T, T, device=device)) |
| |
| mask[:, :prefix_len] = 1.0 |
| |
| mask = mask.view(1, 1, T, T) |
| mask = mask.masked_fill(mask == 0, float('-inf')) |
| mask = mask.masked_fill(mask == 1.0, 0.0) |
| return mask |
|
|
| def forward(self, idx, targets=None, prefix_len=0): |
| """ |
| idx: (B, T) input tokens |
| targets: (B, T) target tokens (shifted for loss) |
| prefix_len: number of prefix tokens to use bidirectional attention |
| """ |
| B, T = idx.shape |
| if T > self.block_size: |
| raise ValueError(f'block size exceeded: {T} > {self.block_size}') |
| pos = torch.arange(T, device=idx.device) |
| x = self.wte(idx) + self.wpe(pos) |
|
|
| if prefix_len > 0: |
| attn_mask = self._build_prefix_lm_mask(T, prefix_len, idx.device) |
| else: |
| attn_mask = None |
|
|
| for block in self.blocks: |
| x = block(x, attn_mask) |
|
|
| x = self.ln_f(x) |
| logits = self.lm_head(x) |
| loss = None |
| if targets is not None: |
| loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1)) |
| return logits, loss |
|
|
| @torch.no_grad() |
| def generate(self, idx, max_new_tokens, temperature=1.0, top_k=40): |
| """Standard causal generation (no prefix).""" |
| self.eval() |
| for _ in range(max_new_tokens): |
| idx_cond = idx[:, -self.block_size:] |
| logits, _ = self(idx_cond, prefix_len=0) |
| logits = logits[:, -1, :] / temperature |
| if top_k is not None: |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| logits[logits < v[:, [-1]]] = -float('inf') |
| probs = F.softmax(logits, dim=-1) |
| idx_next = torch.multinomial(probs, num_samples=1) |
| idx = torch.cat([idx, idx_next], dim=1) |
| return idx |
|
|
| @torch.no_grad() |
| def generate_with_prefix(self, idx, prefix_len, max_new_tokens, temperature=1.0, top_k=40): |
| """ |
| Generate tokens where the first `prefix_len` tokens use bidirectional attention |
| and the rest are autoregressive. |
| """ |
| self.eval() |
| for _ in range(max_new_tokens): |
| idx_cond = idx[:, -self.block_size:] |
| T = idx_cond.shape[1] |
| prefix = min(prefix_len, T) |
| logits, _ = self(idx_cond, prefix_len=prefix) |
| logits = logits[:, -1, :] / temperature |
| if top_k is not None: |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| logits[logits < v[:, [-1]]] = -float('inf') |
| probs = F.softmax(logits, dim=-1) |
| idx_next = torch.multinomial(probs, num_samples=1) |
| idx = torch.cat([idx, idx_next], dim=1) |
| return idx |
