| 453 | |
| 454 | |
| 455 | class Transformer(nn.Module): |
| 456 | def __init__( |
| 457 | self, |
| 458 | config, |
| 459 | max_seq_size=8192, |
| 460 | ): |
| 461 | super().__init__() |
| 462 | self.num_layers = config.num_layers |
| 463 | self.layers = self._build_layers(config) |
| 464 | |
| 465 | def _build_layers(self, config): |
| 466 | layers = torch.nn.ModuleList() |
| 467 | for layer_id in range(self.num_layers): |
| 468 | layers.append( |
| 469 | TransformerBlock( |
| 470 | config, |
| 471 | layer_id=layer_id + 1 , |
| 472 | ) |
| 473 | ) |
| 474 | return layers |
| 475 | |
| 476 | def forward( |
| 477 | self, |
| 478 | hidden_states, |
| 479 | attention_mask, |
| 480 | cu_seqlens=None, |
| 481 | max_seq_len=None, |
| 482 | ): |
| 483 | |
| 484 | if max_seq_len is not None and not isinstance(max_seq_len, torch.Tensor): |
| 485 | max_seq_len = torch.tensor(max_seq_len, dtype=torch.int32, device="cpu") |
| 486 | |
| 487 | for lid, layer in enumerate(self.layers): |
| 488 | hidden_states = layer( |
| 489 | hidden_states, |
| 490 | attention_mask, |
| 491 | cu_seqlens, |
| 492 | max_seq_len, |
| 493 | ) |
| 494 | return hidden_states |
| 495 | |
| 496 | |
| 497 | class Step1Model(PreTrainedModel): |