| 41 | |
| 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: |
| 64 | return self._max_seq_length |
| 65 | |
| 66 | @max_seq_length.setter |
| 67 | def max_seq_length(self, value: int) -> None: |
| 68 | """ |
| 69 | When doing inference, the sequences used might be shorter than the model's context length. |
| 70 | This allows setting a smaller number to avoid allocating unused memory |
| 71 | """ |
| 72 | if value > self.config.block_size: |
| 73 | raise ValueError(f"Cannot attend to {value}, block size is only {self.config.block_size}") |
| 74 | self._max_seq_length = value |
| 75 | if not hasattr(self, "cos"): |
| 76 | # first call |
| 77 | cos, sin = self.get_rope_cache() |
| 78 | self.register_buffer("cos", cos, persistent=False) |
| 79 | self.register_buffer("sin", sin, persistent=False) |
| 80 | # override |
| 81 | elif value != self.cos.size(0): |
| 82 | self.cos, self.sin = self.get_rope_cache(device=self.cos.device) |
| 83 | # the mask and kv cache size will get updated on `set_kv_cache`. we cannot update it here because we don't know |
| 84 | # if the kv cache is expected |
| 85 | |
| 86 | def _init_weights(self, module: nn.Module, n_layer) -> None: |
| 87 | """Meant to be used with `gpt.apply(gpt._init_weights)`.""" |
| 88 | # GPT-NeoX https://arxiv.org/pdf/2204.06745.pdf |
| 89 | if isinstance(module, nn.Embedding): |
| 90 | torch.nn.init.normal_( |
| 91 | module.weight, mean=0.0, std=math.sqrt(2.0 / 5 / self.config.n_embd) |
| 92 | ) |
| 93 | # RWKV: set it to 1e-4 |
| 94 | # torch.nn.init.uniform_(module.weight, -1e-4, 1e-4) |
| 95 | elif isinstance(module, nn.Linear): |
| 96 | torch.nn.init.normal_( |
| 97 | module.weight, mean=0.0, std=math.sqrt(2.0 / 5 / self.config.n_embd) |
| 98 | ) |
| 99 | if module.bias is not None: |
| 100 | torch.nn.init.zeros_(module.bias) |
no outgoing calls
no test coverage detected