स्टेप 4 · ट्रांसफॉर्मर ब्लॉक

class SwiGLU(nn.Module):
    def __init__(self, d_model, hidden=None):
        super().__init__()
        # 8/3 rather than 4x, because the gate adds a third matrix
        hidden = hidden or int(8 * d_model / 3 / 64) * 64
        self.up   = nn.Linear(d_model, hidden, bias=False)
        self.gate = nn.Linear(d_model, hidden, bias=False)
        self.down = nn.Linear(hidden, d_model, bias=False)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))

class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d))
    def forward(self, x):
        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return x * rms * self.weight

class Block(nn.Module):
    def __init__(self, d_model, n_head):
        super().__init__()
        self.n1 = RMSNorm(d_model)
        self.attn = CausalSelfAttention(d_model, n_head)
        self.n2 = RMSNorm(d_model)
        self.ff = SwiGLU(d_model)
    def forward(self, x):
        x = x + self.attn(self.n1(x))    # mix across positions
        x = x + self.ff(self.n2(x))      # think within each position
        return x