| 485 | |
| 486 | |
| 487 | class Mamba(nn.Module): |
| 488 | def __init__( |
| 489 | self, |
| 490 | num_tokens: int, |
| 491 | sequence_length: int, |
| 492 | config: MambaConfig, |
| 493 | return_embeddings: bool = True, |
| 494 | return_tokens: bool = True, |
| 495 | ): |
| 496 | super().__init__() |
| 497 | self.num_tokens = num_tokens |
| 498 | self.sequence_length = sequence_length |
| 499 | self.config = config |
| 500 | self.return_embeddings = return_embeddings |
| 501 | self.return_tokens = return_tokens |
| 502 | |
| 503 | # Embedding |
| 504 | self.token_embed = nn.Embedding(num_tokens, config.dim) |
| 505 | self.norm = nn.LayerNorm(config.dim) |
| 506 | |
| 507 | self.layers = nn.ModuleList( |
| 508 | [ResidualBlock(config) for _ in range(config.depth)] |
| 509 | ) |
| 510 | # self.norm_f = RMSNorm(config.dim) |
| 511 | |
| 512 | def forward(self, x): |
| 513 | # x : (B, L, D) |
| 514 | |
| 515 | # y : (B, L, D) |
| 516 | # Embedding |
| 517 | x = self.token_embed(x) |
| 518 | x = self.norm(x) |
| 519 | |
| 520 | for layer in self.layers: |
| 521 | x = layer(x) |
| 522 | |
| 523 | x = self.norm(x) |
| 524 | |
| 525 | # Return embeddings or Logits |
| 526 | # Return Tokens |
| 527 | if self.return_tokens: |
| 528 | x = OutputHead(self.config.dim, -1)(x) |
| 529 | return x |
| 530 | else: |
| 531 | return x |
| 532 | |
| 533 | def step(self, x, caches): |
| 534 | # x : (B, L, D) |
| 535 | # caches : [cache(layer) for all layers], cache : (h, inputs) |
| 536 | |
| 537 | # y : (B, L, D) |
| 538 | # caches : [cache(layer) for all layers], cache : (h, inputs) |
| 539 | |
| 540 | for i, layer in enumerate(self.layers): |
| 541 | x, caches[i] = layer.step(x, caches[i]) |
| 542 | |
| 543 | return x, caches |
| 544 | |