| 232 | |
| 233 | |
| 234 | class MPTBlock(nn.Module): |
| 235 | def __init__(self, layer_id: int, args): |
| 236 | super().__init__() |
| 237 | self.n_heads = args.n_heads |
| 238 | self.dim = args.d_model |
| 239 | self.head_dim = args.d_model // args.n_heads |
| 240 | self.attn = MPTAttentionFused(args) |
| 241 | self.ffn = MPTMLP(d_model=args.d_model, expansion_ratio=4) |
| 242 | self.layer_id = layer_id |
| 243 | self.norm_1 = LPLayerNorm(args.d_model, eps=1e-6) |
| 244 | self.norm_2 = LPLayerNorm(args.d_model, eps=1e-6) |
| 245 | |
| 246 | def forward( |
| 247 | self, |
| 248 | x: torch.Tensor, |
| 249 | start_pos: int, |
| 250 | mask: Optional[torch.Tensor], |
| 251 | ): |
| 252 | h = x + self.attn.forward(self.norm_1(x), start_pos, mask) |
| 253 | out = h + self.ffn.forward(self.norm_2(h)) |
| 254 | return out |
| 255 | |
| 256 | |
| 257 | class Transformer(nn.Module): |