| 360 | """ |
| 361 | |
| 362 | def __init__(self, hidden_size, output_dropout_prob, init_method, |
| 363 | output_layer_init_method=None): |
| 364 | super(ParallelMLP, self).__init__() |
| 365 | # Set output layer initialization if not provided. |
| 366 | if output_layer_init_method is None: |
| 367 | output_layer_init_method = init_method |
| 368 | # Project to 4h. |
| 369 | self.dense_h_to_4h = ColumnParallelLinear(hidden_size, 4 * hidden_size, |
| 370 | gather_output=False, |
| 371 | init_method=init_method) |
| 372 | # Project back to h. |
| 373 | self.dense_4h_to_h = RowParallelLinear( |
| 374 | 4 * hidden_size, |
| 375 | hidden_size, |
| 376 | input_is_parallel=True, |
| 377 | init_method=output_layer_init_method) |
| 378 | self.dropout = torch.nn.Dropout(output_dropout_prob) |
| 379 | |
| 380 | def forward(self, hidden_states): |
| 381 | # [b, s, 4hp] |