| """Transformer block for the dense SAGE model.""" |
|
|
| from __future__ import annotations |
|
|
| from typing import Optional |
|
|
| import torch |
| from torch import nn |
|
|
| from model.attention import GQAAttention |
| from model.config import ModelConfig |
| from model.mlp import SwiGLUMLP |
| from model.rmsnorm import RMSNorm |
|
|
|
|
| class TransformerBlock(nn.Module): |
| """Pre-norm transformer block with attention and SwiGLU.""" |
|
|
| def __init__(self, config: ModelConfig): |
| super().__init__() |
| self.norm1 = RMSNorm(config.d_model, eps=config.rms_norm_eps) |
| self.attn = GQAAttention(config) |
| self.norm2 = RMSNorm(config.d_model, eps=config.rms_norm_eps) |
| self.mlp = SwiGLUMLP(config) |
|
|
| def forward( |
| self, |
| hidden_states: torch.Tensor, |
| cos: torch.Tensor, |
| sin: torch.Tensor, |
| past_key_value: Optional[tuple[torch.Tensor, torch.Tensor]] = None, |
| ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: |
| """Forward pass with residual connections.""" |
| attn_output, present = self.attn(self.norm1(hidden_states), cos, sin, past_key_value=past_key_value) |
| hidden_states = hidden_states + attn_output |
| hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) |
| return hidden_states, present |
|
|
