MCPcopy Create free account
hub / github.com/pytorch/examples / init_weights

Method init_weights

distributed/tensor_parallelism/llama2_model.py:395–426  ·  view source on GitHub ↗

[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)

Source from the content-addressed store, hash-verified

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 """

Callers 1

__init__Method · 0.95

Calls 3

precompute_freqs_cisFunction · 0.85
init_weightsMethod · 0.45
reset_parametersMethod · 0.45

Tested by

no test coverage detected