已合并
【bugfix!!!】fbov COC&share_expert_sync fix #3005
EX_mitsuX创建于 2025年11月19日
【bugfix!!!】fbov COC&share_expert_sync fix #3005
已合并
共 7 个文件变更+45-24
| @@ -315,3 +315,13 @@ def alltoall_seq_token_unpermutation( | |||
| 315 | # Reshape the output tensor | 315 | # Reshape the output tensor |
| 316 | output = output.view(self.hidden_shape) | 316 | output = output.view(self.hidden_shape) |
| 317 | return output, None | 317 | return output, None |
| 318 | + | ||
| 319 | + | ||
| 320 | +def preprocess_sync_wrapper(fn): | ||
| 321 | + | ||
| 322 | + def wrapper(self, routing_map): | ||
| 323 | + fn(self, routing_map) | ||
| 324 | + if self.num_local_experts > 1: | ||
| 325 | + if self.config.moe_permute_fusion: | ||
| 326 | + self._maybe_update_cuda_sync_point("before_permutation_2") | ||
| 327 | + return wrapper | ||
E | |||
| @@ -95,10 +95,9 @@ def transformer_layer_forward_moe( | |||
| 95 | perm1_out, perm1_probs, tokens_per_expert = dispatcher.token_permute1(detached_mlp_input, probs_detached, routing_map) | 95 | perm1_out, perm1_probs, tokens_per_expert = dispatcher.token_permute1(detached_mlp_input, probs_detached, routing_map) |
| 96 | 96 | ||
| 97 | if use_shared_experts: | 97 | if use_shared_experts: |
| 98 | - with torch.npu.stream(dispatcher.overlap_stream): | 98 | + # Shared Experts Forward. |
| 99 | - # Shared Experts Forward. | 99 | + self.mlp.shared_experts.linear_fc1_forward_and_act() |
| 100 | - self.mlp.shared_experts.linear_fc1_forward_and_act() | 100 | + self.mlp.shared_experts.linear_fc2_forward() |
| 101 | - self.mlp.shared_experts.linear_fc2_forward() | ||
| 102 | if dispatcher.num_local_experts > 1: | 101 | if dispatcher.num_local_experts > 1: |
| 103 | # launch synchronization here to wait for non-blocking mem copy in preprocess func. | 102 | # launch synchronization here to wait for non-blocking mem copy in preprocess func. |
| 104 | dispatcher.cuda_sync_point = "no_sync" | 103 | dispatcher.cuda_sync_point = "no_sync" |
| @@ -217,10 +217,9 @@ def transformer_layer_forward_moe_backward_moe_overlaping( | |||
| 217 | perm1_out, perm1_probs, tokens_per_expert = fwd_dispatcher.token_permute1(detached_mlp_input, probs_detached, routing_map) | 217 | perm1_out, perm1_probs, tokens_per_expert = fwd_dispatcher.token_permute1(detached_mlp_input, probs_detached, routing_map) |
| 218 | 218 | ||
| 219 | if use_shared_experts: | 219 | if use_shared_experts: |
| 220 | - with torch.npu.stream(fwd_dispatcher.overlap_stream): | 220 | + # Shared Experts Forward. |
| 221 | - # Shared Experts Forward. | 221 | + fwd_shared_experts.linear_fc1_forward_and_act() |
| 222 | - fwd_shared_experts.linear_fc1_forward_and_act() | 222 | + fwd_shared_experts.linear_fc2_forward() |
| 223 | - fwd_shared_experts.linear_fc2_forward() | ||
| 224 | 223 | ||
| 225 | if args.moe_zero_memory != 'disable': | 224 | if args.moe_zero_memory != 'disable': |
| 226 | (bwd_perm_a2a_out, bwd_recomp_perm_a2a_handle), _ = bwd_dispatcher.async_dispatch_comm( | 225 | (bwd_perm_a2a_out, bwd_recomp_perm_a2a_handle), _ = bwd_dispatcher.async_dispatch_comm( |
| @@ -247,11 +246,17 @@ def transformer_layer_forward_moe_backward_moe_overlaping( | |||
| 247 | 246 | ||
| 248 | # Shared Experts Backward | 247 | # Shared Experts Backward |
| 249 | if use_shared_experts: | 248 | if use_shared_experts: |
| 250 | - WeightGradStore.start_decouple() | 249 | + with torch.npu.stream(bwd_dispatcher.overlap_stream): |
| 251 | - bwd_shared_experts.linear_fc2_act_fc1_backward(bwd_layer_graph.shared_experts_graph, keep_grad=True) | 250 | + WeightGradStore.start_decouple() |
| 252 | - WeightGradStore.end_decouple() | 251 | + bwd_shared_experts.linear_fc2_act_fc1_backward(bwd_layer_graph.shared_experts_graph, keep_grad=True) |
| 252 | + WeightGradStore.end_decouple() | ||
| 253 | 253 | ||
| 254 | with checkpoint_context: | 254 | with checkpoint_context: |
| 255 | + # Async Perm A2A. | ||
| 256 | + from ..modules.token_dispatcher import PREMUTE_FINISH_EVENT | ||
| 257 | + if PREMUTE_FINISH_EVENT is not None: | ||
| 258 | + #Wait for permute1 finish. | ||
| 259 | + torch.npu.current_stream().wait_event(PREMUTE_FINISH_EVENT) | ||
| 255 | (perm_a2a_out, perm_a2a_handle), (perm_prob_a2a_out, perm_prob_a2a_handle) = fwd_dispatcher.async_dispatch_comm( | 260 | (perm_a2a_out, perm_a2a_handle), (perm_prob_a2a_out, perm_prob_a2a_handle) = fwd_dispatcher.async_dispatch_comm( |
| 256 | perm1_out, perm1_probs, wait_event=last_comm_handle | 261 | perm1_out, perm1_probs, wait_event=last_comm_handle |
| 257 | ) | 262 | ) |
| @@ -263,8 +268,9 @@ def transformer_layer_forward_moe_backward_moe_overlaping( | |||
| 263 | 268 | ||
| 264 | with checkpoint_context: | 269 | with checkpoint_context: |
| 265 | if use_shared_experts: | 270 | if use_shared_experts: |
| 266 | - fwd_shared_experts.post_forward_comm(wait_event=last_comm_handle) | 271 | + with torch.npu.stream(fwd_dispatcher.overlap_stream): |
| 267 | - last_comm_handle = fwd_shared_experts.fc2_output_comm_handle | 272 | + fwd_shared_experts.post_forward_comm(wait_event=last_comm_handle) |
| 273 | + last_comm_handle = fwd_shared_experts.fc2_output_comm_handle | ||
| 268 | 274 | ||
| 269 | if recomp_norm: | 275 | if recomp_norm: |
| 270 | fwd_layer.norm_ckpt2.discard_output() | 276 | fwd_layer.norm_ckpt2.discard_output() |
| @@ -283,7 +289,8 @@ def transformer_layer_forward_moe_backward_moe_overlaping( | |||
| 283 | 289 | ||
| 284 | # launch shared experts post backward comm | 290 | # launch shared experts post backward comm |
| 285 | if use_shared_experts: | 291 | if use_shared_experts: |
| 286 | - bwd_shared_experts.post_backward_comm(wait_event=last_comm_handle) | 292 | + with torch.npu.stream(bwd_dispatcher.overlap_stream): |
| 293 | + bwd_shared_experts.post_backward_comm(wait_event=last_comm_handle) | ||
| 287 | 294 | ||
| 288 | # Grouped MLP dw computation | 295 | # Grouped MLP dw computation |
| 289 | if args.moe_zero_memory == 'level0': | 296 | if args.moe_zero_memory == 'level0': |
| @@ -327,7 +334,8 @@ def transformer_layer_forward_moe_backward_moe_overlaping( | |||
| 327 | unperm1_out = fwd_dispatcher.token_unpermute1(detached_expert_output, None) | 334 | unperm1_out = fwd_dispatcher.token_unpermute1(detached_expert_output, None) |
| 328 | expert_output.untyped_storage().resize_(0) | 335 | expert_output.untyped_storage().resize_(0) |
| 329 | if use_shared_experts: | 336 | if use_shared_experts: |
| 330 | - shared_expert_output, share_experts_graph = fwd_shared_experts.get_output() | 337 | + with torch.npu.stream(fwd_dispatcher.overlap_stream): |
| 338 | + shared_expert_output, share_experts_graph = fwd_shared_experts.get_output() | ||
| 331 | 339 | ||
| 332 | bwd_perm_a2a_handle.wait() | 340 | bwd_perm_a2a_handle.wait() |
| 333 | bwd_perm_a2a_handle = None | 341 | bwd_perm_a2a_handle = None |
| @@ -342,10 +350,12 @@ def transformer_layer_forward_moe_backward_moe_overlaping( | |||
| 342 | if bwd_prob_handle: | 350 | if bwd_prob_handle: |
| 343 | bwd_prob_handle.wait() | 351 | bwd_prob_handle.wait() |
| 344 | if use_shared_experts: | 352 | if use_shared_experts: |
| 345 | - shared_experts_grad = bwd_shared_experts.get_backward_grad() | 353 | + with torch.npu.stream(bwd_dispatcher.overlap_stream): |
| 346 | - if shared_experts_grad is not None: | 354 | + shared_experts_grad = bwd_shared_experts.get_backward_grad() |
| 347 | - bwd_layer_graph.pre_mlp_layernorm_graph[1].grad = shared_experts_grad | 355 | + if shared_experts_grad is not None: |
| 356 | + bwd_layer_graph.pre_mlp_layernorm_graph[1].grad = shared_experts_grad | ||
| 348 | 357 | ||
| 358 | + torch.npu.current_stream().wait_stream(bwd_dispatcher.overlap_stream) | ||
| 349 | run_graph_backward(bwd_layer_graph.perm1_graph, [perm1_out_grad, perm1_prob_out_grad]) | 359 | run_graph_backward(bwd_layer_graph.perm1_graph, [perm1_out_grad, perm1_prob_out_grad]) |
| 350 | perm1_out_grad.untyped_storage().resize_(0) | 360 | perm1_out_grad.untyped_storage().resize_(0) |
| 351 | 361 | ||
| @@ -399,6 +409,7 @@ def transformer_layer_forward_moe_backward_moe_overlaping( | |||
| 399 | unperm_a2a_out.untyped_storage().resize_(0) | 409 | unperm_a2a_out.untyped_storage().resize_(0) |
| 400 | 410 | ||
| 401 | if use_shared_experts: | 411 | if use_shared_experts: |
| 412 | + torch.npu.current_stream().wait_stream(fwd_dispatcher.overlap_stream) | ||
| 402 | detached_shared_expert_output = detach_tensor(shared_expert_output) | 413 | detached_shared_expert_output = detach_tensor(shared_expert_output) |
| 403 | mlp_output = route_expert_output + detached_shared_expert_output | 414 | mlp_output = route_expert_output + detached_shared_expert_output |
| 404 | shared_expert_output.untyped_storage().resize_(0) | 415 | shared_expert_output.untyped_storage().resize_(0) |
| @@ -87,13 +87,16 @@ class FusedMoEPermuteFeature(MindSpeedFeature): | |||
| 87 | 87 | ||
| 88 | if getattr(args, "moe_token_dispatcher_type", None) == "alltoall": | 88 | if getattr(args, "moe_token_dispatcher_type", None) == "alltoall": |
| 89 | from mindspeed.core.fusions.fused_moe_permute import ( | 89 | from mindspeed.core.fusions.fused_moe_permute import ( |
| 90 | - moe_alltoall_token_dispatcher_init_wrapper, maybe_dtoh_and_synchronize) | 90 | + moe_alltoall_token_dispatcher_init_wrapper, maybe_dtoh_and_synchronize, preprocess_sync_wrapper) |
| 91 | 91 | ||
| 92 | # Since fused_sort_chunks_by_index is not currently supported, when self.num_local_experts > 1, | 92 | # Since fused_sort_chunks_by_index is not currently supported, when self.num_local_experts > 1, |
| 93 | # move self.num_global_tokens_per_local_expert to cpu | 93 | # move self.num_global_tokens_per_local_expert to cpu |
| 94 | patch_manager.register_patch( | 94 | patch_manager.register_patch( |
| 95 | 'megatron.core.transformer.moe.token_dispatcher.MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize', | 95 | 'megatron.core.transformer.moe.token_dispatcher.MoEAlltoAllTokenDispatcher._maybe_dtoh_and_synchronize', |
| 96 | maybe_dtoh_and_synchronize) | 96 | maybe_dtoh_and_synchronize) |
| 97 | + patch_manager.register_patch( | ||
| 98 | + 'megatron.core.transformer.moe.token_dispatcher.MoEAlltoAllTokenDispatcher.preprocess', | ||
| 99 | + preprocess_sync_wrapper) | ||
| 97 | 100 | ||
| 98 | # Since fused_sort_chunks_by_index is not currently supported, set self.permute_idx_device to None | 101 | # Since fused_sort_chunks_by_index is not currently supported, set self.permute_idx_device to None |
| 99 | patch_manager.register_patch( | 102 | patch_manager.register_patch( |
| @@ -41,10 +41,6 @@ def patch_features(): | |||
| 41 | # apply megatron patches | 41 | # apply megatron patches |
| 42 | MindSpeedFeaturesManager.apply_features_patches(mindspeed_args) | 42 | MindSpeedFeaturesManager.apply_features_patches(mindspeed_args) |
| 43 | 43 | ||
| 44 | - # accelerate package will check TE on sys.modules, so we need remove this patch | ||
| 45 | - if 'transformer_engine' in sys.modules: | ||
| 46 | - del sys.modules["transformer_engine"] | ||
| 47 | - | ||
| 48 | 44 | ||
| 49 | def delete_lock_file(): | 45 | def delete_lock_file(): |
| 50 | """Delete lock file in multiprocess for JIT build..""" | 46 | """Delete lock file in multiprocess for JIT build..""" |
| @@ -82,7 +82,9 @@ class MindSpeedTELayerNormColumnParallelLinear(torch.nn.Module): | |||
| 82 | self.is_expert = is_expert | 82 | self.is_expert = is_expert |
| 83 | self.sequence_parallel = self.config.sequence_parallel | 83 | self.sequence_parallel = self.config.sequence_parallel |
| 84 | self.gradient_accumulation_fusion = self.config.gradient_accumulation_fusion | 84 | self.gradient_accumulation_fusion = self.config.gradient_accumulation_fusion |
| 85 | + self.parallel_mode = 'column' | ||
| 85 | self.fp8_meta = FP8Metadata(['inputs', 'weight', 'grads']) | 86 | self.fp8_meta = FP8Metadata(['inputs', 'weight', 'grads']) |
| 87 | + | ||
| 86 | # MindSpeedTELayerNormColumnParallelLinear check. | 88 | # MindSpeedTELayerNormColumnParallelLinear check. |
| 87 | if gather_output: | 89 | if gather_output: |
| 88 | raise ValueError('Transformer Engine linear layers do not support gather_output = True') | 90 | raise ValueError('Transformer Engine linear layers do not support gather_output = True') |
修复开启permute融合算子后可能的同步问题。