Init method based on N(0, sigma/sqrt(2*num_layers).
(sigma, num_layers)
| 591 | |
| 592 | |
| 593 | def scaled_init_method(sigma, num_layers): |
| 594 | """Init method based on N(0, sigma/sqrt(2*num_layers).""" |
| 595 | std = sigma / math.sqrt(2.0 * num_layers) |
| 596 | |
| 597 | def init_(tensor): |
| 598 | return torch.nn.init.normal_(tensor, mean=0.0, std=std) |
| 599 | |
| 600 | return init_ |
| 601 | |
| 602 | |
| 603 | class GPT2ParallelTransformer(torch.nn.Module): |