(self, mpu_vocab_size, hidden_size, init_method,
layernorm_epsilon, parallel_output)
| 70 | """ |
| 71 | |
| 72 | def __init__(self, mpu_vocab_size, hidden_size, init_method, |
| 73 | layernorm_epsilon, parallel_output): |
| 74 | |
| 75 | super(BertLMHead, self).__init__() |
| 76 | |
| 77 | args = get_args() |
| 78 | |
| 79 | self.bias = torch.nn.Parameter(torch.zeros(mpu_vocab_size)) |
| 80 | self.bias.model_parallel = True |
| 81 | self.bias.partition_dim = 0 |
| 82 | self.bias.stride = 1 |
| 83 | self.parallel_output = parallel_output |
| 84 | |
| 85 | self.dense = get_linear_layer(hidden_size, hidden_size, init_method) |
| 86 | self.layernorm = LayerNorm(hidden_size, eps=layernorm_epsilon) |
| 87 | self.gelu = torch.nn.functional.gelu |
| 88 | if args.openai_gelu: |
| 89 | self.gelu = openai_gelu |
| 90 | elif args.onnx_safe: |
| 91 | self.gelu = erf_gelu |
| 92 | |
| 93 | def forward(self, hidden_states, word_embeddings_weight): |
| 94 | hidden_states = self.dense(hidden_states) |
nothing calls this directly
no test coverage detected