File size: 758 Bytes
33f617f 7812b29 33f617f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 | import torch.nn as nn
from .SwiGLU import SwiGLU
from .GroupedQueryAttention import GroupedQueryAttention
class DecoderBlock(nn.Module):
def __init__(self, d_model: int, num_heads: int, num_kv_heads: int, ff_dim: int, dropout: float):
super().__init__()
self.gqa = GroupedQueryAttention(d_model, num_heads, num_kv_heads, dropout)
self.ffn = SwiGLU(d_model, ff_dim)
self.norm1 = nn.RMSNorm(d_model, eps=1e-6)
self.norm2 = nn.RMSNorm(d_model, eps=1e-6)
self.drop = nn.Dropout(dropout)
def forward(self, x, cos, sin, attn_mask=None, causal: bool = True):
x = x + self.drop(self.gqa(self.norm1(x), cos, sin, attn_mask, causal))
x = x + self.drop(self.ffn(self.norm2(x)))
return x
|