स्टेप 11 · मॉडर्न आर्किटेक्चर अपग्रेड्स

class GQA(nn.Module):
    def __init__(self, d_model, n_head, n_kv_head):
        super().__init__()
        self.n_head, self.n_kv = n_head, n_kv_head
        self.d_head = d_model // n_head
        self.rep = n_head // n_kv_head          # query heads per kv head
        self.q  = nn.Linear(d_model, n_head * self.d_head, bias=False)
        self.kv = nn.Linear(d_model, 2 * n_kv_head * self.d_head, bias=False)
        self.proj = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        B, T, C = x.shape
        q = self.q(x).view(B, T, self.n_head, self.d_head).transpose(1, 2)
        k, v = self.kv(x).view(B, T, 2, self.n_kv, self.d_head).unbind(2)
        k = k.transpose(1, 2).repeat_interleave(self.rep, dim=1)
        v = v.transpose(1, 2).repeat_interleave(self.rep, dim=1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.proj(y.transpose(1, 2).reshape(B, T, C))

config = dict(
    norm="rmsnorm", norm_position="pre",
    position="rope", rope_theta=500000,     # large theta for long context
    ffn="swiglu", ffn_mult=8/3,
    attention="gqa", n_head=32, n_kv_head=8,
    bias=False,
    tie_embeddings=True,                    # for models under ~2B
    vocab_size=128256,
    dtype="bfloat16",
)