File size: 1,185 Bytes
353cf0a | 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 29 30 31 32 33 | import torch
import torch.nn as nn
from .attention import GQA
from .SwiGLU import SwiGLU
class TransformerBlock(nn.Module):
def __init__(self, hidden_size: int = 1024, intermediate_size: int = 2816,
num_heads: int = 16, num_kv_heads: int = 4, max_seq_len: int = 2048,
rope_base: float = 10000.0, norm_eps: float = 1e-6,
layer_idx: int = 0):
super().__init__()
self.attn_norm = nn.RMSNorm(hidden_size, eps=norm_eps)
self.attn = GQA(
hidden_size=hidden_size,
num_heads=num_heads,
num_kv_heads=num_kv_heads,
max_seq_len=max_seq_len,
rope_base=rope_base,
layer_idx=layer_idx,
)
self.mlp_norm = nn.RMSNorm(hidden_size, eps=norm_eps)
self.mlp = SwiGLU(
hidden_size=hidden_size,
intermediate_size=intermediate_size,
)
def forward(self, x: torch.Tensor, past_key_values=None, past_len: int = 0):
x = x + self.attn(self.attn_norm(x),
past_key_values=past_key_values, past_len=past_len)
x = x + self.mlp(self.mlp_norm(x))
return x |