| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import math |
| from dataclasses import dataclass |
|
|
|
|
| @dataclass |
| class GLMConfig: |
| vocab_size: int = 151552 |
| hidden_size: int = 4096 |
| intermediate_size: int = 14336 |
| num_layers: int = 32 |
| num_heads: int = 32 |
| num_kv_heads: int = 8 |
| head_dim: int = 128 |
| eps: float = 1e-6 |
| rope_base: float = 10000.0 |
|
|
|
|
| |
| class RMSNorm(nn.Module): |
| def __init__(self, hidden_size, eps=1e-6): |
| super().__init__() |
| self.weight = nn.Parameter(torch.ones(hidden_size)) |
| self.eps = eps |
|
|
| def forward(self, x): |
| var = x.pow(2).mean(-1, keepdim=True) |
| x = x * torch.rsqrt(var + self.eps) |
| return self.weight * x |
|
|
|
|
| |
| class RotaryEmbedding(nn.Module): |
| def __init__(self, head_dim, base=10000.0): |
| super().__init__() |
| inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim)) |
| self.register_buffer("inv_freq", inv_freq, persistent=False) |
|
|
| def forward(self, seq_len, device, dtype): |
| t = torch.arange(seq_len, device=device).float() |
| freqs = torch.einsum("i,j->ij", t, self.inv_freq) |
| |
| emb = torch.cat([freqs, freqs], dim=-1) |
| return emb.cos().to(dtype), emb.sin().to(dtype) |
|
|
|
|
| def rotate_half(x): |
| x1 = x[..., : x.shape[-1] // 2] |
| x2 = x[..., x.shape[-1] // 2:] |
| return torch.cat([-x2, x1], dim=-1) |
|
|
|
|
| def apply_rotary(q, k, cos, sin): |
| cos = cos[None, None, :, :] |
| sin = sin[None, None, :, :] |
| q = (q * cos) + (rotate_half(q) * sin) |
| k = (k * cos) + (rotate_half(k) * sin) |
| return q, k |
|
|
|
|
| |
| class SelfAttention(nn.Module): |
| def __init__(self, cfg: GLMConfig): |
| super().__init__() |
| self.cfg = cfg |
| self.q_proj = nn.Linear(cfg.hidden_size, cfg.num_heads * cfg.head_dim, bias=False) |
| self.k_proj = nn.Linear(cfg.hidden_size, cfg.num_kv_heads * cfg.head_dim, bias=False) |
| self.v_proj = nn.Linear(cfg.hidden_size, cfg.num_kv_heads * cfg.head_dim, bias=False) |
| self.o_proj = nn.Linear(cfg.num_heads * cfg.head_dim, cfg.hidden_size, bias=False) |
|
|
| def forward(self, x, cos, sin): |
| B, T, _ = x.shape |
| q = self.q_proj(x).view(B, T, self.cfg.num_heads, self.cfg.head_dim).transpose(1, 2) |
| k = self.k_proj(x).view(B, T, self.cfg.num_kv_heads, self.cfg.head_dim).transpose(1, 2) |
| v = self.v_proj(x).view(B, T, self.cfg.num_kv_heads, self.cfg.head_dim).transpose(1, 2) |
|
|
| q, k = apply_rotary(q, k, cos, sin) |
|
|
| |
| if self.cfg.num_kv_heads != self.cfg.num_heads: |
| rep = self.cfg.num_heads // self.cfg.num_kv_heads |
| k = k.repeat_interleave(rep, dim=1) |
| v = v.repeat_interleave(rep, dim=1) |
|
|
| out = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| out = out.transpose(1, 2).contiguous().view(B, T, -1) |
| return self.o_proj(out) |
|
|
|
|
| |
| class MLP(nn.Module): |
| def __init__(self, cfg: GLMConfig): |
| super().__init__() |
| self.gate_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False) |
| self.up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False) |
| self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False) |
|
|
| def forward(self, x): |
| return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) |
|
|
|
|
| |
| class GLMBlock(nn.Module): |
| def __init__(self, cfg: GLMConfig): |
| super().__init__() |
| self.input_norm = RMSNorm(cfg.hidden_size, cfg.eps) |
| self.self_attn = SelfAttention(cfg) |
| self.post_norm = RMSNorm(cfg.hidden_size, cfg.eps) |
| self.mlp = MLP(cfg) |
|
|
| def forward(self, x, cos, sin): |
| |
| x = x + self.self_attn(self.input_norm(x), cos, sin) |
| |
| x = x + self.mlp(self.post_norm(x)) |
| return x |
|
|
|
|
| |
| class GLMForCausalLM(nn.Module): |
| def __init__(self, cfg: GLMConfig): |
| super().__init__() |
| self.cfg = cfg |
| self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size) |
| self.rotary_emb = RotaryEmbedding(cfg.head_dim, cfg.rope_base) |
| self.layers = nn.ModuleList([GLMBlock(cfg) for _ in range(cfg.num_layers)]) |
| self.final_norm = RMSNorm(cfg.hidden_size, cfg.eps) |
| self.lm_head = nn.Linear(cfg.hidden_size, cfg.vocab_size, bias=False) |
| |
| self.lm_head.weight = self.embed_tokens.weight |
|
|
| def forward(self, input_ids): |
| B, T = input_ids.shape |
| x = self.embed_tokens(input_ids) |
| cos, sin = self.rotary_emb(T, x.device, x.dtype) |
|
|
| for layer in self.layers: |
| x = layer(x, cos, sin) |
|
|
| x = self.final_norm(x) |
| return self.lm_head(x) |
|
|
|
|
| |
| if __name__ == "__main__": |
| cfg = GLMConfig() |
| model = GLMForCausalLM(cfg) |
| n_params = sum(p.numel() for p in model.parameters()) |
| print(f"Total parameters: {n_params/1e9:.2f}B") |
|
|
| ids = torch.randint(0, cfg.vocab_size, (2, 16)) |
| with torch.no_grad(): |
| logits = model(ids) |
| print(f"Output logits shape: {logits.shape}") |
|
|
