| 475 | |
| 476 | |
| 477 | class BertOutput(nn.Module): |
| 478 | def __init__(self, config): |
| 479 | super(BertOutput, self).__init__() |
| 480 | if hasattr(config, 'deep_init') and config.deep_init: |
| 481 | init_method = scaled_init_method(mean=0.0, |
| 482 | std=config.initializer_range, |
| 483 | num_layers=config.num_hidden_layers) |
| 484 | else: |
| 485 | init_method = normal_init_method(mean=0.0, |
| 486 | std=config.initializer_range) |
| 487 | self.dense = nn.Linear(config.intermediate_size, config.hidden_size, bias=True) |
| 488 | # self.dense = mpu.RowParallelLinear( |
| 489 | # input_size=config.intermediate_size, |
| 490 | # output_size=config.hidden_size, |
| 491 | # bias=True, |
| 492 | # input_is_parallel=True, |
| 493 | # stride=1, |
| 494 | # init_method=init_method) |
| 495 | self.fp32_layernorm = config.fp32_layernorm |
| 496 | self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layernorm_epsilon) |
| 497 | self.dropout = nn.Dropout(config.hidden_dropout_prob) |
| 498 | |
| 499 | def forward(self, hidden_states, input_tensor): |
| 500 | hidden_states = self.dense(hidden_states) |
| 501 | hidden_states = self.dropout(hidden_states) |
| 502 | ln_input = hidden_states + input_tensor |
| 503 | previous_type = ln_input.type() |
| 504 | if self.fp32_layernorm: |
| 505 | ln_input = ln_input.float() |
| 506 | hidden_states = self.LayerNorm(ln_input) |
| 507 | if self.fp32_layernorm: |
| 508 | hidden_states = hidden_states.type(previous_type) |
| 509 | return hidden_states |
| 510 | |
| 511 | |
| 512 | class BertLayer(nn.Module): |