diff --git a/docs/mindspore/features/alltoallvc.md b/docs/mindspore/features/alltoallvc.md index a0aaccef8a2141e0b61aae47d3115d0a00724a92..2145649dc24d60cc9c8fe533f4a6e224cd4b7eb3 100644 --- a/docs/mindspore/features/alltoallvc.md +++ b/docs/mindspore/features/alltoallvc.md @@ -21,8 +21,8 @@ MindSpore **AllToAllVC** 通信算子，支持不均匀通信，并在性能上 * **性能优化** - * AllToAllV 的实现中，通信过程需要进行 **两次 D2H（Device to Host）拷贝**。 - * AllToAllVC 算子优化为 **只进行一次拷贝**，显著降低了通信开销。 + * AllToAllV 的算子实现中，通信过程需要进行 **一次 D2H（Device to Host）拷贝**。 + * AllToAllVC 算子内优化为 **不进行 D2H 拷贝**，显著降低了通信开销。 * 在实际测试中，AllToAllVC 算子的性能优于基于 `torch.distributed.all_to_all_single` 的实现。 --- @@ -31,7 +31,11 @@ MindSpore **AllToAllVC** 通信算子，支持不均匀通信，并在性能上 （1）开启 AllToAllVC 特性，需要在运行时确保通信组已初始化，并加载 MindSpeed 中的通信优化模块。 -（2）在 msrun 启动bash 脚本中，启用参数 **--enable-a2avc** 即可开启使用 AllToAllVC 特性。 +（2）在 msrun 启动bash 脚本中，使用参数 **--enable-a2avc [int]** 即可控制是否使用 AllToAllVC 特性以及具体的模式。参数说明： + * **[int]** 为整数，代表AllToAllVC的模式，可选值为 0，1，或 2。 + * **[int]** = 0：默认值。不启用 AllToAllVC 功能。 + * **[int]** = 1：启用 AllToAllVC，并在运行过程中对传输参数和计算结果执行多层次一致性校验，以便及时发现并中止潜在的非法或不匹配传值。该模式在健壮性上更为严格，但性能开销较高，性能低于模式 2。 + * **[int]** = 2：启用 AllToAllVC，但运行时不进行多层次校验，仅执行必要的通信与计算，性能优于模式 1。 （3）需在 CANN 8.3.RC1 及以上版本使用 @@ -43,7 +47,7 @@ MindSpore **AllToAllVC** 通信算子，支持不均匀通信，并在性能上 ## 使用效果 -* **性能提升**：相对 AllToAllV，通信性能更优。 +* **性能提升**： 在使用模式 2 （**--enable-a2avc 2**）时,相对 AllToAllV，通信性能更优。在剪裁 deepseekv3 用例上性能比 AllToAllV整体优化 0.4%。 --- diff --git a/mindspeed_llm/mindspore/core/tensor_parallel/mappings.py b/mindspeed_llm/mindspore/core/tensor_parallel/mappings.py index b25c8e0518c7df6ef3988109a58110411fc09e6c..989bfc15bebaec30504290ad4f2fa4481e701105 100644 --- a/mindspeed_llm/mindspore/core/tensor_parallel/mappings.py +++ b/mindspeed_llm/mindspore/core/tensor_parallel/mappings.py @@ -6,10 +6,11 @@ import torch class _AllToAll(torch.autograd.Function): @staticmethod - def forward(ctx, group, input, send_count_matrix): + def forward(ctx, group, input, send_count_matrix, mode): """Forward function.""" ctx.group = group ctx.send_count_matrix = send_count_matrix + ctx.mode = mode world_size = torch.distributed.get_world_size(group=group) # Bypass the function if we are using only 1 GPU. @@ -41,11 +42,12 @@ class _AllToAll(torch.autograd.Function): """Backward function.""" return ( None, - _AllToAll.apply(ctx.group, *grad_output, ctx.send_count_matrix_T), + _AllToAll.apply(ctx.group, *grad_output, ctx.send_count_matrix_T, ctx.mode), + None, None, ) -def all_to_all(group, input_, send_count_matrix=None): +def all_to_all(group, input_, send_count_matrix=None, mode=None): """Wrapper for autograd function""" - return _AllToAll.apply(group, input_, send_count_matrix) \ No newline at end of file + return _AllToAll.apply(group, input_, send_count_matrix, mode) \ No newline at end of file diff --git a/mindspeed_llm/mindspore/mindspore_adaptor_v2.py b/mindspeed_llm/mindspore/mindspore_adaptor_v2.py index 638c252694335a46f677671f0a7496b4be42e406..18866523e00349f7c5d6b599d6e429884468c702 100644 --- a/mindspeed_llm/mindspore/mindspore_adaptor_v2.py +++ b/mindspeed_llm/mindspore/mindspore_adaptor_v2.py @@ -314,7 +314,9 @@ def patch_moe_fb_overlap(): def mindspore_register_args(group): - group.add_argument('--enable-a2avc', action='store_true', default=False, - help='enable a2avc') + group.add_argument('--enable-a2avc', type=int, choices=[0, 1, 2], default=0, + help='0: Disable a2avc,' + '1: Enable a2avc & Use mindspore comm_func.py & with verification (run slower than 2),' + '2: Enable a2avc & Use msadapter comm_func.py & without verification (run faster than 1)') group.add_argument('--enable-share-memory', action='store_true', default=False, help='Enable shared memory for passing actual_seq_len when reset-position-ids is enabled.')