| 609 | |
| 610 | |
| 611 | class BertLMPredictionHead(nn.Module): |
| 612 | def __init__(self, config, bert_model_embedding_weights): |
| 613 | super(BertLMPredictionHead, self).__init__() |
| 614 | self.transform = BertPredictionHeadTransform(config) |
| 615 | |
| 616 | # The output weights are the same as the input embeddings, but there is |
| 617 | # an output-only bias for each token. |
| 618 | self.decoder = nn.Linear(bert_model_embedding_weights.size(1), |
| 619 | bert_model_embedding_weights.size(0), |
| 620 | bias=False) |
| 621 | # self.decoder_weight = bert_model_embedding_weights |
| 622 | # self.bias = nn.Parameter(torch.zeros(bert_model_embedding_weights.size(0))) |
| 623 | # self.bias.model_parallel = True |
| 624 | self.fp32_embedding = config.fp32_embedding |
| 625 | self.fp32_layernorm = config.fp32_layernorm |
| 626 | |
| 627 | def convert_to_type(tensor): |
| 628 | if self.fp32_embedding: |
| 629 | return tensor.half() |
| 630 | else: |
| 631 | return tensor |
| 632 | |
| 633 | self.type_converter = convert_to_type |
| 634 | self.converted = False |
| 635 | |
| 636 | def forward(self, hidden_states): |
| 637 | if not self.converted: |
| 638 | self.converted = True |
| 639 | if self.fp32_embedding: |
| 640 | self.transform.half() |
| 641 | if self.fp32_layernorm: |
| 642 | self.transform.LayerNorm.float() |
| 643 | hidden_states = self.transform(self.type_converter(hidden_states)) |
| 644 | hidden_states = self.decoder(hidden_states) + self.bias |
| 645 | # hidden_states = mpu.copy_to_model_parallel_region(hidden_states) |
| 646 | # hidden_states = F.linear(self.type_converter(hidden_states), |
| 647 | # self.type_converter(self.decoder_weight), |
| 648 | # self.type_converter(self.bias)) |
| 649 | return hidden_states |
| 650 | |
| 651 | |
| 652 | class BertOnlyMLMHead(nn.Module): |