File size: 1,004 Bytes
2ae6691 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 | import torch
import torch.nn as nn
from .rmsnorm import RMSNorm
from .attention import GroupedQueryAttention
from .mlp import SwiGLU
class DecoderBlock(nn.Module):
"""One transformer block: attention sub-layer, then MLP sub-layer.
Each sub-layer is pre-normed and wrapped in a residual (add-back) connection.
"""
def __init__(self, dim, n_heads, n_kv_heads, head_dim, hidden_dim, eps=1e-6):
super().__init__()
# a norm before each sub-layer
self.attn_norm = RMSNorm(dim, eps)
self.mlp_norm = RMSNorm(dim, eps)
self.attn = GroupedQueryAttention(dim, n_heads, n_kv_heads, head_dim)
self.mlp = SwiGLU(dim, hidden_dim)
def forward(self, x, cos, sin, block_mask):
# attention sub-layer: normalize, attend, add the result back
x = x + self.attn(self.attn_norm(x), cos, sin, block_mask)
# mlp sub-layer: normalize, process per-token, add the result back
x = x + self.mlp(self.mlp_norm(x))
return x |