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.
| 58 | """ |
| 59 | |
| 60 | class ParallelMLP(MegatronModule): |
| 61 | """MLP. |
| 62 | |
| 63 | MLP will take the input with h hidden state, project it to 4*h |
| 64 | hidden dimension, perform nonlinear transformation, and project the |
| 65 | state back into h hidden dimension. At the end, dropout is also |
| 66 | applied. |
| 67 | """ |
| 68 | |
| 69 | def __init__(self, init_method, output_layer_init_method): |
| 70 | super(ParallelMLP, self).__init__() |
| 71 | args = get_args() |
| 72 | |
| 73 | # Project to 4h. |
| 74 | if not args.memory_centric_tiled_linear: |
| 75 | self.dense_h_to_4h = mpu.ColumnParallelLinear( |
| 76 | args.hidden_size, |
| 77 | 4 * args.hidden_size, |
| 78 | gather_output=False, |
| 79 | init_method=init_method, |
| 80 | skip_bias_add=True) |
| 81 | else: |
| 82 | self.dense_h_to_4h = deepspeed.zero.TiledLinearReturnBias( |
| 83 | in_features=args.hidden_size, |
| 84 | out_features=4*args.hidden_size, |
| 85 | linear_cls=mpu.ColumnParallelLinear, |
| 86 | in_splits=args.tile_factor, |
| 87 | out_splits=4*args.tile_factor, |
| 88 | combine_out_splits=True, |
| 89 | gather_output=False, |
| 90 | init_method=init_method, |
| 91 | skip_bias_add=True) |
| 92 | |
| 93 | self.bias_gelu_fusion = args.bias_gelu_fusion |
| 94 | self.activation_func = F.gelu |
| 95 | if args.openai_gelu: |
| 96 | self.activation_func = openai_gelu |
| 97 | elif args.onnx_safe: |
| 98 | self.activation_func = erf_gelu |
| 99 | |
| 100 | # Project back to h. |
| 101 | if not args.memory_centric_tiled_linear: |
| 102 | self.dense_4h_to_h = mpu.RowParallelLinear( |
| 103 | 4 * args.hidden_size, |
| 104 | args.hidden_size, |
| 105 | input_is_parallel=True, |
| 106 | init_method=output_layer_init_method, |
| 107 | skip_bias_add=True) |
| 108 | else: |
| 109 | self.dense_4h_to_h = deepspeed.zero.TiledLinearReturnBias( |
| 110 | in_features=4*args.hidden_size, |
| 111 | out_features=args.hidden_size, |
| 112 | linear_cls=mpu.RowParallelLinear, |
| 113 | in_splits=4*args.tile_factor, |
| 114 | out_splits=args.tile_factor, |
| 115 | input_is_already_split=False, |
| 116 | combine_out_splits=True, |
| 117 | input_is_parallel=True, |