| 435 | |
| 436 | |
| 437 | class BertAttention(nn.Module): |
| 438 | def __init__(self, config): |
| 439 | super(BertAttention, self).__init__() |
| 440 | self.self = BertSelfAttention(config) |
| 441 | # self.self = mpu.BertParallelSelfAttention( |
| 442 | # hidden_size=config.hidden_size, |
| 443 | # num_attention_heads=config.num_attention_heads, |
| 444 | # dropout_prob=config.attention_probs_dropout_prob, |
| 445 | # output_parallel=True, |
| 446 | # init_method=normal_init_method(mean=0.0, |
| 447 | # std=config.initializer_range)) |
| 448 | self.output = BertSelfOutput(config) |
| 449 | |
| 450 | def forward(self, input_tensor, attention_mask): |
| 451 | self_output = self.self(input_tensor, attention_mask) |
| 452 | attention_output = self.output(self_output, input_tensor) |
| 453 | return attention_output |
| 454 | |
| 455 | |
| 456 | class BertIntermediate(nn.Module): |