[Note: On ``init_weights`` vs. ``reset_parameters``] Modules may define ``reset_parameters`` to initialize parameter values. ``reset_parameters`` is meant to only initialize directly owned parameters/buffers, not those of their child modules, and it can be us
(self)
| 393 | self.init_weights() |
| 394 | |
| 395 | def init_weights(self): |
| 396 | """ |
| 397 | [Note: On ``init_weights`` vs. ``reset_parameters``] |
| 398 | Modules may define ``reset_parameters`` to initialize parameter values. |
| 399 | ``reset_parameters`` is meant to only initialize directly owned |
| 400 | parameters/buffers, not those of their child modules, and it can be |
| 401 | used to give the initial values for these tensors. |
| 402 | Separately, users may want custom initialization for their modules, |
| 403 | different from that in ``reset_parameters``. For this, we define |
| 404 | ``init_weights``. We only call it in the constructor of this |
| 405 | ``Transformer`` root module to avoid reinitializing tensors. |
| 406 | """ |
| 407 | with torch.device(self.freqs_cis.device): |
| 408 | self.freqs_cis = precompute_freqs_cis( |
| 409 | self.model_args.dim // self.model_args.n_heads, |
| 410 | # Need to compute until at least the max token limit for generation |
| 411 | # (use 2x max sequence length to be safe) |
| 412 | self.model_args.max_seq_len * 2, |
| 413 | ) |
| 414 | nn.init.normal_(self.tok_embeddings.weight) |
| 415 | for layer in self.layers: |
| 416 | layer.init_weights() |
| 417 | self.norm.reset_parameters() |
| 418 | final_out_std = self.model_args.dim**-0.5 |
| 419 | cutoff_factor = 3 |
| 420 | nn.init.trunc_normal_( |
| 421 | self.output.weight, |
| 422 | mean=0.0, |
| 423 | std=final_out_std, |
| 424 | a=-cutoff_factor * final_out_std, |
| 425 | b=cutoff_factor * final_out_std, |
| 426 | ) |
| 427 | |
| 428 | def forward(self, tokens: torch.Tensor): |
| 429 | """ |
no test coverage detected