| 478 | |
| 479 | |
| 480 | class T5Stack(torch.nn.Module): |
| 481 | def __init__(self, num_layers, model_dim, inner_dim, ff_dim, num_heads, vocab_size, dtype, device): |
| 482 | super().__init__() |
| 483 | self.embed_tokens = torch.nn.Embedding(vocab_size, model_dim, device=device) |
| 484 | self.block = torch.nn.ModuleList([T5Block(model_dim, inner_dim, ff_dim, num_heads, relative_attention_bias=(i == 0), dtype=dtype, device=device) for i in range(num_layers)]) |
| 485 | self.final_layer_norm = T5LayerNorm(model_dim, dtype=dtype, device=device) |
| 486 | |
| 487 | def forward(self, input_ids, intermediate_output=None, final_layer_norm_intermediate=True): |
| 488 | intermediate = None |
| 489 | x = self.embed_tokens(input_ids) |
| 490 | past_bias = None |
| 491 | for i, l in enumerate(self.block): |
| 492 | x, past_bias = l(x, past_bias) |
| 493 | if i == intermediate_output: |
| 494 | intermediate = x.clone() |
| 495 | x = self.final_layer_norm(x) |
| 496 | if intermediate is not None and final_layer_norm_intermediate: |
| 497 | intermediate = self.final_layer_norm(intermediate) |
| 498 | return x, intermediate |
| 499 | |
| 500 | |
| 501 | class T5(torch.nn.Module): |