MLP. MLP will take the input with h hidden state, project it to 4*h hidden dimension, perform nonlinear transformation, and project the state back into h hidden dimension. At the end, dropout is also applied.
| 53 | |
| 54 | |
| 55 | class ParallelMLP(MegatronModule): |
| 56 | """MLP. |
| 57 | |
| 58 | MLP will take the input with h hidden state, project it to 4*h |
| 59 | hidden dimension, perform nonlinear transformation, and project the |
| 60 | state back into h hidden dimension. At the end, dropout is also |
| 61 | applied. |
| 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] |
| 95 | intermediate_parallel, _ = self.dense_h_to_4h(hidden_states) |
| 96 | intermediate_parallel = self.activation_func(intermediate_parallel) |
| 97 | # [s, b, h] |
| 98 | output, output_bias = self.dense_4h_to_h(intermediate_parallel) |
| 99 | |
| 100 | return output, output_bias |
| 101 | |
| 102 | |
| 103 | class ParallelSelfAttention(MegatronModule): |