| 68 | |
| 69 | class GPT(BaseModel): |
| 70 | def __init__(self, config: Config) -> None: |
| 71 | # Skip the parent class __init__ altogether and replace it to avoid useless allocations |
| 72 | nn.Module.__init__(self) |
| 73 | assert config.padded_vocab_size is not None |
| 74 | self.config = config |
| 75 | |
| 76 | self.lm_head = AdapterV2Linear( |
| 77 | config.n_embd, config.padded_vocab_size, bias=False |
| 78 | ) |
| 79 | self.transformer = nn.ModuleDict( |
| 80 | dict( |
| 81 | wte=nn.Embedding(config.padded_vocab_size, config.n_embd), |
| 82 | h=nn.ModuleList(Block(config, i) for i in range(config.n_layer)), |
| 83 | ln_f=config.norm_class(config.n_embd, eps=config.norm_eps), |
| 84 | ) |
| 85 | ) |
| 86 | |
| 87 | self.rope_cache: Optional[RoPECache] = None |
| 88 | self.mask_cache: Optional[torch.Tensor] = None |
| 89 | self.kv_caches: List[KVCache] = [] |
| 90 | self.adapter_kv_caches: List[KVCache] = [] |
| 91 | |
| 92 | @classmethod |
| 93 | def from_name(cls, name: str, **kwargs: Any) -> Self: |