已合并
feat: mc2 frozen weight #1917
aijgnem1创建于 2025年2月17日
feat: mc2 frozen weight #1917
已合并
从refs/pull/1917/head合入到core_r0.8.0
共 5 个文件变更+251-14
| @@ -14,7 +14,8 @@ | |||
| 14 | # See the License for the specific language governing permissions and | 14 | # See the License for the specific language governing permissions and |
| 15 | # limitations under the License. | 15 | # limitations under the License. |
| 16 | from .ascend_turbo_cfg import ascend_turbo_cfg | 16 | from .ascend_turbo_cfg import ascend_turbo_cfg |
| 17 | -from .mc2_linears_seq_parallel import ColumnSeqParallelLinear, RowSeqParallelLinear | 17 | +from .mc2_linears_seq_parallel import (ColumnSeqParallelLinear, RowSeqParallelLinear, |
| 18 | + ColumnSeqParallelLinearWithFrozenWeight, RowSeqParallelLinearWithFrozenWeight) | ||
| 18 | 19 | ||
| 19 | 20 | ||
| 20 | def column_parallel_forward(self, input_, weight=None): | 21 | def column_parallel_forward(self, input_, weight=None): |
| @@ -36,18 +37,28 @@ def column_parallel_forward(self, input_, weight=None): | |||
| 36 | 37 | ||
| 37 | bias = self.bias if not self.skip_bias_add else None | 38 | bias = self.bias if not self.skip_bias_add else None |
| 38 | 39 | ||
| 39 | - output = ColumnSeqParallelLinear.apply( | 40 | + if not weight.requires_grad: |
| 40 | - input_, weight, bias, ascend_turbo_cfg.get_group() | 41 | + output = ColumnSeqParallelLinearWithFrozenWeight.apply( |
| 41 | - ) | 42 | + input_, weight, bias, ascend_turbo_cfg.get_group() |
| 43 | + ) | ||
| 44 | + else: | ||
| 45 | + output = ColumnSeqParallelLinear.apply( | ||
| 46 | + input_, weight, bias, ascend_turbo_cfg.get_group() | ||
| 47 | + ) | ||
| 42 | 48 | ||
| 43 | output_bias = self.bias if self.skip_bias_add else None | 49 | output_bias = self.bias if self.skip_bias_add else None |
| 44 | return output, output_bias | 50 | return output, output_bias |
| 45 | 51 | ||
| 46 | 52 | ||
| 47 | def row_parallel_forward(self, input_): | 53 | def row_parallel_forward(self, input_): |
| 48 | - output = RowSeqParallelLinear.apply( | 54 | + if not self.weight.requires_grad: |
| 49 | - input_, self.weight, None, ascend_turbo_cfg.get_group() | 55 | + output = RowSeqParallelLinearWithFrozenWeight.apply( |
| 50 | - ) | 56 | + input_, self.weight, None, ascend_turbo_cfg.get_group() |
| 57 | + ) | ||
| 58 | + else: | ||
| 59 | + output = RowSeqParallelLinear.apply( | ||
| 60 | + input_, self.weight, None, ascend_turbo_cfg.get_group() | ||
| 61 | + ) | ||
| 51 | 62 | ||
| 52 | if not self.skip_bias_add: | 63 | if not self.skip_bias_add: |
| 53 | output = output + self.bias if self.bias is not None else output | 64 | output = output + self.bias if self.bias is not None else output |
| @@ -231,3 +231,122 @@ class RowSeqParallelLinear(torch.autograd.Function): | |||
| 231 | grad_bias = None | 231 | grad_bias = None |
| 232 | 232 | ||
| 233 | return grad_input, grad_weight, grad_bias, None | 233 | return grad_input, grad_weight, grad_bias, None |
| 234 | + | ||
| 235 | + | ||
| 236 | +class ColumnSeqParallelLinearWithFrozenWeight(ColumnSeqParallelLinear): | ||
| 237 | + | ||
| 238 | + def forward(ctx, input_, weight, bias, group): | ||
| 239 | + ctx.input_shape = input_.shape | ||
| 240 | + ctx.use_bias = bias is not None | ||
| 241 | + ctx.weight = weight | ||
| 242 | + | ||
| 243 | + rank = torch.distributed.get_rank(group) | ||
| 244 | + hcomm_info = None | ||
| 245 | + if torch.__version__ > "2.0": | ||
| 246 | + global_rank = torch.distributed.get_global_rank(group, rank) | ||
| 247 | + hcomm_info = group._get_backend(torch.device("npu")).get_hccl_comm_name( | ||
| 248 | + global_rank | ||
| 249 | + ) | ||
| 250 | + | ||
| 251 | + else: | ||
| 252 | + hcomm_info = group.get_hccl_comm_name(rank) | ||
| 253 | + | ||
| 254 | + x = input_.reshape(input_.shape[0] * input_.shape[1], input_.shape[2]) | ||
| 255 | + | ||
| 256 | + world_size = ascend_turbo_cfg.get_world_size() | ||
| 257 | + # npu_all_gather_base_mm currently do not support bias | ||
| 258 | + output, all_gather_grad_output = torch_npu.npu_all_gather_base_mm( | ||
| 259 | + x, | ||
| 260 | + weight.t(), | ||
| 261 | + hcomm_info, | ||
| 262 | + world_size, | ||
| 263 | + bias=None, | ||
| 264 | + gather_index=0, | ||
| 265 | + gather_output=(not ascend_turbo_cfg.all_gather_recomputation), | ||
| 266 | + ) | ||
| 267 | + | ||
| 268 | + if bias is not None: | ||
| 269 | + output = output + bias | ||
| 270 | + | ||
| 271 | + output = output.view( | ||
| 272 | + int(output.shape[0] / input_.shape[1]), input_.shape[1], output.shape[1] | ||
| 273 | + ) | ||
| 274 | + ctx.hcomm_info = hcomm_info | ||
| 275 | + ctx.world_size = world_size | ||
| 276 | + ctx.group = group | ||
| 277 | + return output | ||
| 278 | + | ||
| 279 | + | ||
| 280 | + def backward(ctx, grad_output): | ||
| 281 | + input_shape = ctx.input_shape | ||
| 282 | + weight = ctx.weight | ||
| 283 | + | ||
| 284 | + hcomm_info = ctx.hcomm_info | ||
| 285 | + world_size = ctx.world_size | ||
| 286 | + grad_output_ = grad_output.reshape( | ||
| 287 | + grad_output.shape[0] * grad_output.shape[1], grad_output.shape[2] | ||
| 288 | + ) | ||
| 289 | + | ||
| 290 | + sub_grad_input = torch_npu.npu_mm_reduce_scatter_base( | ||
| 291 | + grad_output_, weight, hcomm_info, world_size, bias=None | ||
| 292 | + ) | ||
| 293 | + | ||
| 294 | + sub_grad_input = sub_grad_input.view(input_shape) | ||
| 295 | + | ||
| 296 | + return sub_grad_input, None, None, None | ||
| 297 | + | ||
| 298 | + | ||
| 299 | +class RowSeqParallelLinearWithFrozenWeight(RowSeqParallelLinear): | ||
| 300 | + | ||
| 301 | + def forward(ctx, input_, weight, bias, group): | ||
| 302 | + ctx.input_shape = input_.shape | ||
| 303 | + ctx.use_bias = bias is not None | ||
| 304 | + ctx.weight = weight | ||
| 305 | + | ||
| 306 | + rank = torch.distributed.get_rank(group) | ||
| 307 | + world_size = ascend_turbo_cfg.get_world_size() | ||
| 308 | + hcomm_info = None | ||
| 309 | + if torch.__version__ > "2.0": | ||
| 310 | + global_rank = torch.distributed.get_global_rank(group, rank) | ||
| 311 | + hcomm_info = group._get_backend(torch.device("npu")).get_hccl_comm_name( | ||
| 312 | + global_rank | ||
| 313 | + ) | ||
| 314 | + else: | ||
| 315 | + hcomm_info = group.get_hccl_comm_name(rank) | ||
| 316 | + | ||
| 317 | + x = input_.reshape(input_.shape[0] * input_.shape[1], input_.shape[2]) | ||
| 318 | + | ||
| 319 | + # npu_mm_reduce_scatter_base currently do not support bias | ||
| 320 | + output = torch_npu.npu_mm_reduce_scatter_base( | ||
| 321 | + x, weight.t(), hcomm_info, world_size, reduce_op="sum", bias=None | ||
| 322 | + ) | ||
| 323 | + | ||
| 324 | + if bias is not None: | ||
| 325 | + output = output + bias | ||
| 326 | + | ||
| 327 | + ctx.hcomm_info = hcomm_info | ||
| 328 | + ctx.world_size = world_size | ||
| 329 | + | ||
| 330 | + output = output.view( | ||
| 331 | + int(output.shape[0] / input_.shape[1]), input_.shape[1], output.shape[1] | ||
| 332 | + ) | ||
| 333 | + | ||
| 334 | + return output | ||
| 335 | + | ||
| 336 | + | ||
| 337 | + def backward(ctx, grad_output): | ||
| 338 | + input_shape = ctx.input_shape | ||
| 339 | + weight = ctx.weight | ||
| 340 | + hcomm_info = ctx.hcomm_info | ||
| 341 | + world_size = ctx.world_size | ||
| 342 | + grad_output_ = grad_output.reshape( | ||
| 343 | + grad_output.shape[0] * grad_output.shape[1], grad_output.shape[2] | ||
| 344 | + ) | ||
| 345 | + | ||
| 346 | + grad_input, _ = torch_npu.npu_all_gather_base_mm( | ||
| 347 | + grad_output_, weight, hcomm_info, world_size, bias=None, gather_index=0 | ||
| 348 | + ) | ||
| 349 | + | ||
| 350 | + grad_input = grad_input.view(input_shape) | ||
| 351 | + | ||
| 352 | + return grad_input, None, None, None | ||
| @@ -29,7 +29,7 @@ def set_mc2_args(args): | |||
| 29 | class TestMC2(DistributedTest): | 29 | class TestMC2(DistributedTest): |
| 30 | world_size = 8 | 30 | world_size = 8 |
| 31 | 31 | ||
| 32 | - def test_MC2ColumnParallelLinear(self): | 32 | + def test_mc2_column_parallel_linear(self): |
| 33 | args = parse_args(None, True) | 33 | args = parse_args(None, True) |
| 34 | args = set_mc2_args(args) | 34 | args = set_mc2_args(args) |
| 35 | set_args(args) | 35 | set_args(args) |
| @@ -65,4 +65,32 @@ class TestMC2(DistributedTest): | |||
| 65 | ones = torch.ones(args.seq_len, args.batch_size * self.world_size).half().npu() | 65 | ones = torch.ones(args.seq_len, args.batch_size * self.world_size).half().npu() |
| 66 | dLdb = torch.matmul(ones.t(), dLdY).sum(dim=0).view(-1) | 66 | dLdb = torch.matmul(ones.t(), dLdY).sum(dim=0).view(-1) |
| 67 | assert torch.allclose(dLdA, linear_layer.weight.grad, rtol=0.005, atol=0.005) | 67 | assert torch.allclose(dLdA, linear_layer.weight.grad, rtol=0.005, atol=0.005) |
| 68 | - assert torch.allclose(dLdb, linear_layer.bias.grad, rtol=0.005, atol=0.005) | 68 | + assert torch.allclose(dLdb, linear_layer.bias.grad, rtol=0.005, atol=0.005) |
| 69 | + | ||
| 70 | + def test_mc2_column_parallel_linear_frozen(self): | ||
| 71 | + args = parse_args(None, True) | ||
| 72 | + args = set_mc2_args(args) | ||
| 73 | + set_args(args) | ||
| 74 | + initialize_cfg_from_args(args) | ||
| 75 | + transformer_config = TransformerConfig(num_layers=1, | ||
| 76 | + hidden_size=12, | ||
| 77 | + num_attention_heads=4, | ||
| 78 | + use_cpu_initialization=True) | ||
| 79 | + transformer_config.sequence_parallel = args.sequence_parallel | ||
| 80 | + set_random_seed(args.seed) | ||
| 81 | + input_size = args.input_size_coeff * args.tensor_model_parallel_size | ||
| 82 | + output_size = args.output_size_coeff * args.tensor_model_parallel_size | ||
| 83 | + linear_layer = ColumnParallelLinear(input_size, | ||
| 84 | + output_size, | ||
| 85 | + keep_master_weight_for_test=True, | ||
| 86 | + init_method=transformer_config.init_method, | ||
| 87 | + config=transformer_config).half().npu() | ||
| 88 | + linear_layer.weight.requires_grad_(False) | ||
| 89 | + setattr(linear_layer.weight, 'main_grad', linear_layer.weight.clone()) | ||
| 90 | + input_ = torch.rand(args.batch_size, args.seq_len, input_size).half().npu() | ||
| 91 | + output = linear_layer(input_) | ||
| 92 | + gather_list = [torch.zeros(input_.shape).half().npu() for _ in range(self.world_size)] | ||
| 93 | + torch.distributed.all_gather(gather_list, input_) | ||
| 94 | + gather_res = torch.concat(gather_list, dim=0) | ||
| 95 | + output_naive = torch.matmul(gather_res, linear_layer.weight.t()) | ||
| 96 | + assert torch.allclose(output_naive, output[0], rtol=0.005, atol=0.005) | ||
| @@ -17,7 +17,7 @@ def set_mc2_args(args): | |||
| 17 | args.use_unpad = False | 17 | args.use_unpad = False |
| 18 | args.seed = 2024 | 18 | args.seed = 2024 |
| 19 | args.seq_len = 256 | 19 | args.seq_len = 256 |
| 20 | - args.input_size_coeff = 128 | 20 | + args.input_size_coeff = 256 |
| 21 | args.output_size_coeff = 256 | 21 | args.output_size_coeff = 256 |
| 22 | args.batch_size = 8 | 22 | args.batch_size = 8 |
| 23 | args.optimize_recomp_communication_level = True | 23 | args.optimize_recomp_communication_level = True |
| @@ -29,8 +29,7 @@ def set_mc2_args(args): | |||
| 29 | class TestMC2(DistributedTest): | 29 | class TestMC2(DistributedTest): |
| 30 | world_size = 8 | 30 | world_size = 8 |
| 31 | 31 | ||
| 32 | - @pytest.mark.skip(reason='this UT need updata for new megatron version.') | 32 | + def test_mc2_row_parallel_linear(self): |
| 33 | - def test_MC2RowParallelLinear(self): | ||
| 34 | args = parse_args(None, True) | 33 | args = parse_args(None, True) |
| 35 | args = set_mc2_args(args) | 34 | args = set_mc2_args(args) |
| 36 | set_args(args) | 35 | set_args(args) |
| @@ -73,4 +72,36 @@ class TestMC2(DistributedTest): | |||
| 73 | ones = torch.ones(args.seq_len, args.batch_size // self.world_size).half().npu() | 72 | ones = torch.ones(args.seq_len, args.batch_size // self.world_size).half().npu() |
| 74 | dLdb = torch.matmul(ones.t(), dLdY).sum(dim=0).view(-1) | 73 | dLdb = torch.matmul(ones.t(), dLdY).sum(dim=0).view(-1) |
| 75 | assert torch.allclose(dLdA, linear_layer.weight.grad, rtol=0.005, atol=0.005) | 74 | assert torch.allclose(dLdA, linear_layer.weight.grad, rtol=0.005, atol=0.005) |
| 76 | - assert torch.allclose(dLdb, linear_layer.bias.grad, rtol=0.005, atol=0.005) | 75 | + assert torch.allclose(dLdb, linear_layer.bias.grad, rtol=0.005, atol=0.005) |
| 76 | + | ||
| 77 | + def test_mc2_row_parallel_linear_frozen(self): | ||
| 78 | + args = parse_args(None, True) | ||
| 79 | + args = set_mc2_args(args) | ||
| 80 | + set_args(args) | ||
| 81 | + initialize_cfg_from_args(args) | ||
| 82 | + transformer_config = TransformerConfig(num_layers=1, | ||
| 83 | + hidden_size=12, | ||
| 84 | + num_attention_heads=4, | ||
| 85 | + use_cpu_initialization=True) | ||
| 86 | + transformer_config.sequence_parallel = args.sequence_parallel | ||
| 87 | + set_random_seed(args.seed) | ||
| 88 | + input_size = args.input_size_coeff * args.tensor_model_parallel_size | ||
| 89 | + output_size = args.output_size_coeff * args.tensor_model_parallel_size | ||
| 90 | + linear_layer = RowParallelLinear(input_size, | ||
| 91 | + output_size, | ||
| 92 | + keep_master_weight_for_test=True, | ||
| 93 | + bias=True, input_is_parallel=True, | ||
| 94 | + skip_bias_add=False, | ||
| 95 | + init_method=transformer_config.init_method, | ||
| 96 | + config=transformer_config).half().npu() | ||
| 97 | + linear_layer.weight.requires_grad_(False) | ||
| 98 | + setattr(linear_layer.weight, 'main_grad', linear_layer.weight.clone()) | ||
| 99 | + input_ = torch.rand(args.batch_size, args.seq_len, args.input_size_coeff) | ||
| 100 | + input_ = input_.half().npu() | ||
| 101 | + output = linear_layer(input_) | ||
| 102 | + res = torch.matmul(input_, linear_layer.weight.npu().T) | ||
| 103 | + tensor = torch.empty(args.batch_size // self.world_size, args.seq_len, output_size) | ||
| 104 | + tensor = tensor.half().npu() | ||
| 105 | + scatter_list = list(torch.chunk(res, chunks=self.world_size, dim=0)) | ||
| 106 | + torch.distributed.reduce_scatter(tensor, scatter_list) | ||
| 107 | + assert torch.allclose(tensor, output[0], rtol=0.005, atol=0.005) | ||
| @@ -29,7 +29,7 @@ def set_mc2_args(args): | |||
| 29 | class TestMC2(DistributedTest): | 29 | class TestMC2(DistributedTest): |
| 30 | world_size = 8 | 30 | world_size = 8 |
| 31 | 31 | ||
| 32 | - def test_Mcore_MC2ColumnParallelLinear(self): | 32 | + def test_mcore_mc2_column_parallel_linear(self): |
| 33 | args = parse_args(None, True) | 33 | args = parse_args(None, True) |
| 34 | args = set_mc2_args(args) | 34 | args = set_mc2_args(args) |
| 35 | set_args(args) | 35 | set_args(args) |
| @@ -84,3 +84,51 @@ class TestMC2(DistributedTest): | |||
| 84 | assert torch.allclose(output_mc2_close, output_mc2_open, rtol=0.005, atol=0.005) | 84 | assert torch.allclose(output_mc2_close, output_mc2_open, rtol=0.005, atol=0.005) |
| 85 | assert torch.allclose(output_weight_grad_mc2_close, output_weight_grad_mc2_open, rtol=0.005, atol=0.005) | 85 | assert torch.allclose(output_weight_grad_mc2_close, output_weight_grad_mc2_open, rtol=0.005, atol=0.005) |
| 86 | assert torch.allclose(linear_layer_mc2_close.bias.grad, linear_layer_mc2_open.bias.grad, rtol=0.005, atol=0.005) | 86 | assert torch.allclose(linear_layer_mc2_close.bias.grad, linear_layer_mc2_open.bias.grad, rtol=0.005, atol=0.005) |
| 87 | + | ||
| 88 | + def test_mcore_mc2_column_parallel_linear_frozen(self): | ||
| 89 | + args = parse_args(None, True) | ||
| 90 | + args = set_mc2_args(args) | ||
| 91 | + set_args(args) | ||
| 92 | + | ||
| 93 | + transformer_config = TransformerConfig(num_layers=1, | ||
| 94 | + hidden_size=12, | ||
| 95 | + num_attention_heads=4, | ||
| 96 | + use_cpu_initialization=True) | ||
| 97 | + transformer_config.sequence_parallel = args.sequence_parallel | ||
| 98 | + input_size = args.input_size_coeff * args.tensor_model_parallel_size | ||
| 99 | + output_size = args.vocab_size | ||
| 100 | + input_ = torch.rand(args.seq_len, args.batch_size, input_size).half().npu() | ||
| 101 | + | ||
| 102 | + # get output_weight | ||
| 103 | + linear_layer = ColumnParallelLinear(input_size, | ||
| 104 | + output_size, | ||
| 105 | + keep_master_weight_for_test=True, | ||
| 106 | + init_method=transformer_config.init_method, | ||
| 107 | + config=transformer_config).half().npu() | ||
| 108 | + linear_layer.weight.requires_grad_(False) | ||
| 109 | + setattr(linear_layer.weight, 'main_grad', linear_layer.weight.clone()) | ||
| 110 | + output_weight_mc2_close = linear_layer.weight | ||
| 111 | + output_weight_mc2_open = linear_layer.weight | ||
| 112 | + | ||
| 113 | + # close mc2 forward and backward | ||
| 114 | + set_random_seed(args.seed) | ||
| 115 | + linear_layer_mc2_close = ColumnParallelLinear(input_size, | ||
| 116 | + output_size, | ||
| 117 | + keep_master_weight_for_test=True, | ||
| 118 | + init_method=transformer_config.init_method, | ||
| 119 | + config=transformer_config).half().npu() | ||
| 120 | + linear_layer_mc2_close.weight.requires_grad_(False) | ||
| 121 | + output_mc2_close, _ = linear_layer_mc2_close(input_, output_weight_mc2_close) | ||
| 122 | + | ||
| 123 | + # open mc2 forward and backward | ||
| 124 | + initialize_cfg_from_args(args) | ||
| 125 | + set_random_seed(args.seed) | ||
| 126 | + linear_layer_mc2_open = ColumnParallelLinear(input_size, | ||
| 127 | + output_size, | ||
| 128 | + keep_master_weight_for_test=True, | ||
| 129 | + init_method=transformer_config.init_method, | ||
| 130 | + config=transformer_config).half().npu() | ||
| 131 | + output_mc2_open, _ = linear_layer_mc2_open(input_, output_weight_mc2_open) | ||
| 132 | + | ||
| 133 | + # result compare | ||
| 134 | + assert torch.allclose(output_mc2_close, output_mc2_open, rtol=0.005, atol=0.005) | ||