| 62 | """ |
| 63 | |
| 64 | def __init__( |
| 65 | self, |
| 66 | init_method, |
| 67 | output_layer_init_method, |
| 68 | scale: int = 4, |
| 69 | ): |
| 70 | super(ParallelMLP, self).__init__() |
| 71 | args = get_args() |
| 72 | |
| 73 | # Project to 4h. |
| 74 | self.dense_h_to_4h = mpu.ColumnParallelLinear( |
| 75 | args.hidden_size, |
| 76 | scale * args.hidden_size, |
| 77 | gather_output=False, |
| 78 | init_method=init_method, |
| 79 | # skip_bias_add=True, |
| 80 | ) |
| 81 | |
| 82 | self.activation_func = fast_gelu |
| 83 | |
| 84 | # Project back to h. |
| 85 | self.dense_4h_to_h = mpu.RowParallelLinear( |
| 86 | scale * args.hidden_size, |
| 87 | args.hidden_size, |
| 88 | input_is_parallel=True if args.tensor_model_parallel_size > 1 else False, |
| 89 | init_method=output_layer_init_method, |
| 90 | # skip_bias_add=True, |
| 91 | ) |
| 92 | |
| 93 | def forward(self, hidden_states): |
| 94 | # [s, b, 4hp] |