Spaces:
Runtime error
Runtime error
| """ | |
| One transformer block: a pre-norm residual around attention, then another | |
| around the SwiGLU MLP | |
| x = x + attn(norm1(x)) | |
| x = x + mlp(norm2(x)) | |
| Pre-norm (norm before each sublayer) keeps a clean residual highway, which | |
| lets deep stacks train stably. | |
| """ | |
| import torch.nn as nn | |
| from .norm import AbyssRMSNorm | |
| from .attention import AbyssCausalSelfAttention | |
| from .mlp import AbyssSwiGLU | |
| class AbyssBlock(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.norm1 = AbyssRMSNorm(config.n_embd, config.norm_eps) | |
| self.attn = AbyssCausalSelfAttention(config) | |
| self.norm2 = AbyssRMSNorm(config.n_embd, config.norm_eps) | |
| self.mlp = AbyssSwiGLU(config) | |
| def forward(self, x): | |
| x = x + self.attn(self.norm1(x)) | |
| x = x + self.mlp(self.norm2(x)) | |
| return x |