| 74 | |
| 75 | |
| 76 | class TransformerBlock(nn.Module): |
| 77 | def __init__(self, args: ModelArgs): |
| 78 | super().__init__() |
| 79 | self.attention_norm = nn.LayerNorm(args.dim) |
| 80 | self.attention = Attention(args) |
| 81 | self.ffn_norm = nn.LayerNorm(args.dim) |
| 82 | self.feed_forward = FeedForward( |
| 83 | args.dim, hidden_dim=4 * args.dim, dropout_p=args.dropout_p |
| 84 | ) |
| 85 | |
| 86 | def forward(self, x): |
| 87 | h = x + self.attention(self.attention_norm(x)) |
| 88 | out = h + self.feed_forward(self.ffn_norm(h)) |
| 89 | return out |
| 90 | |
| 91 | def reset_parameters(self): |
| 92 | self.attention_norm.reset_parameters() |
| 93 | self.attention.reset_parameters() |
| 94 | self.ffn_norm.reset_parameters() |
| 95 | self.feed_forward.reset_parameters() |
| 96 | |
| 97 | |
| 98 | # A toy transformer model, partly inspired by the nanoGPT model: |