| 42 | |
| 43 | class GPT(nn.Module): |
| 44 | def __init__(self, config: Config) -> None: |
| 45 | super().__init__() |
| 46 | assert config.padded_vocab_size is not None |
| 47 | self.config = config |
| 48 | |
| 49 | self.lm_head = nn.Linear(config.n_embd, config.padded_vocab_size, bias=False) |
| 50 | self.transformer = nn.ModuleDict( |
| 51 | dict( |
| 52 | wte=nn.Embedding(config.padded_vocab_size, config.n_embd), |
| 53 | h=nn.ModuleList(Block(config) for _ in range(config.n_layer)), |
| 54 | ln_f=config.norm_class(config.n_embd, eps=config.norm_eps), |
| 55 | ) |
| 56 | ) |
| 57 | self.max_seq_length = self.config.block_size |
| 58 | self.rope_cache: Optional[RoPECache] = None |
| 59 | self.mask_cache: Optional[torch.Tensor] = None |
| 60 | self.kv_caches: List[KVCache] = [] |
| 61 | |
| 62 | @property |
| 63 | def max_seq_length(self) -> int: |