| from dataclasses import dataclass |
| from typing import Literal |
|
|
| import torch |
| from torch.nn import functional as F |
|
|
|
|
| def gelu_approx(x): |
| return F.gelu(x, approximate="tanh") |
|
|
|
|
| @dataclass |
| class LinearWeights: |
| weight: torch.Tensor |
| bias: torch.Tensor |
|
|
|
|
| def linear(x: torch.Tensor, w: LinearWeights) -> torch.Tensor: |
| return F.linear(x, w.weight, w.bias) |
|
|
|
|
| @dataclass |
| class LayerNormWeights: |
| weight: torch.Tensor |
| bias: torch.Tensor |
|
|
|
|
| def layer_norm(x: torch.Tensor, w: LayerNormWeights) -> torch.Tensor: |
| return F.layer_norm(x, w.bias.shape, w.weight, w.bias) |
|
|
|
|
| @dataclass |
| class MLPWeights: |
| fc1: LinearWeights |
| fc2: LinearWeights |
| act: Literal["gelu_approx"] = "gelu_approx" |
|
|
|
|
| def mlp(x: torch.Tensor, w: MLPWeights) -> torch.Tensor: |
| x = w.fc1(x) |
| x = gelu_approx(x) |
| x = w.fc2(x) |
| return x |
|
|
|
|
| @dataclass |
| class AttentionWeights: |
| qkv: LinearWeights |
| proj: LinearWeights |
|
|
|
|
| def attn(x: torch.Tensor, w: AttentionWeights, n_heads: int) -> torch.Tensor: |
| bsz, q_len, d_model = x.shape |
| head_dim = d_model // n_heads |
|
|
| q, k, v = [ |
| t.view(bsz, q_len, n_heads, head_dim).transpose(1, 2) |
| for t in linear(x, w.qkv).chunk(3, dim=-1) |
| ] |
| out = F.scaled_dot_product_attention(q, k, v) |
| out = out.transpose(1, 2).reshape(bsz, q_len, d_model) |
| out = linear(out, w.proj) |
| return out |
|
|
