已合并
feat: add get_p2p_stream_id interface for HCCL backend #32201
pengjingyou创建于 3月23日
feat: add get_p2p_stream_id interface for HCCL backend #32201
已合并
pengjingyou创建于 3月23日
pengjingyou成员
3月23日

【合入来源】

如有社区issue,请关联issue链接
请勿携带内部流程信息(需求链接、问题单、内部issue等)

为 HCCL 后端新增 get_p2p_stream_id 接口,用于获取 P2P 通信的 NPU stream ID,方便用户在分布式训练场景下进行 stream 级别的同步和管理。

【修改方案】

  1. 新增 C++ 方法 getP2PStreamId
    在 ProcessGroupHCCL 类中新增 getP2PStreamId 方法,支持两种模式:
  • 批量模式 (is_batched=1) :使用 getKeyFromDevices 生成 key,适用于 batch_isend_irecv 场景
  • 非批量模式 (is_batched=0) :使用 getKeySendRecv 生成 key,适用于单独的 send/recv 场景
  1. 新增 Python 绑定
    在 Init.cpp 中添加 Python 绑定,暴露 get_p2p_stream_id 接口到 Python 层

【资料变更】

请确认是否涉及资料变更。如涉及,需要在PR中体现,并简要说明修改内容。如不涉及,需填写“不涉及”

【接口变更】

接口名称 参数 返回值 说明
ProcessGroupHCCL.get_p2p_stream_id device: torch.device - NPU
设备 peer: int - 对端 rank
is_batched: int - 是否批量模式 (1=批量, 0=非批量)
int - stream ID,不存在时返回 -1 获取 P2P 通信的 stream ID

【功能验证】

  1. ut验证
    新增测试文件 test/distributed/test_get_p2p_stream_id.py ,包含以下测试场景:
测试用例 说明
test_p2p_stream_id_batched 测试批量模式 (is_batched=1) 下获取 stream ID
test_p2p_stream_id_non_batched 测试非批量模式 (is_batched=0) 下获取 stream ID
test_p2p_stream_id_invalid 测试未构造p2p stream场景下获取 stream ID 非法值-1

python -m pytest test/distributed/test_get_p2p_stream_id.py
image.png

  1. 业务场景验证
    调用dist.send/recv方法进行p2p通信后,通过该新增接口get_p2p_stream_id()获取p2p stream id,拿到id对应的Stream实例,指定在该Stream上进行计算操作。
    通过profiling文件可以看到计算在p2p stream上进行
    image.png

【CheckList】

PR提交人对以下CheckList自检项进行全量自检,自检通过或不涉及,均修改 [ ] 为 [x]

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 pengjingyou 的贡献)
Ppengjingyou成员
3月23日 创建了 pull request,commit ee23c030
ascend-robot
ascend-robot成员
3月23日 评论:

Thanks for your pull-request.
The full list of commands accepted by me can be found at here
You can get sig-info at here


PR Approval Progress

Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
test renyujin, wjlflyer (2/2) renyujin (1/1)
torch_npu/csrc/distributed renyujin, wjlflyer (2/2) renyujin (1/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

pengjingyou, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
ascend-robotascend-robot成员
3月23日 添加了label:ascend-cla/yes
pengjingyou成员
3月23日 评论:

compile

likedislike
ascend-robotascend-robot成员
3月23日 添加了label:ci-pipeline-running
ascend-robotascend-robot成员
3月23日 删除了label:ci-pipeline-running
ascend-robotascend-robot成员
3月23日 添加了label:ci-pipeline-failed
ascend-robot
ascend-robot成员
3月23日 评论:
流水线 PR-pipeline_pytorch#12038 运行失败
阶段 任务名 状态 详情
编译构建 Build_X86 >>>
Build_ARM >>>
Build_LibTorch_x86 >>>
Build_X86_torchair 🛑 >>>
Build_ARM_torchair 🛑 >>>
Build_ARM_inductor 🛑 >>>
恶意代码检查 Antipoison >>>
编码安全与规范检查 CodeCheck >>>
check_error >>>
开源片段检查 SCA >>>
开发者测试 UT_X86_Part_01 🛑 >>>
UT_X86_Part_02 🛑 >>>
UT_ARM_A3_Part_01 🛑 >>>
UT_ARM_A3_Part_02 🛑 >>>
UT_DIST_X86_Part_01 >>>
UT_DIST_X86_Part_02 >>>
UT_DIST_X86_Part_03 >>>
UT_DIST_X86_Part_04 >>>
UT_inductor_Part_01 🛑 >>>
UT_inductor_Part_02 🛑 >>>
UT_inductor_Part_03 🛑 >>>
UT_inductor_Part_04 🛑 >>>
UT_ARM_A2_Part_01 >>>
UT_ARM_A2_Part_02 >>>
流水线 PR-pipeline_pytorch >>>
likedislike
pengjingyou成员
3月24日 评论:

compile

likedislike
ascend-robotascend-robot成员
3月24日 删除了label:ci-pipeline-failed
此处折叠了7条事件消息 查看更多
ascend-robotascend-robot成员
3月24日 添加了label:ci-pipeline-passed
ascend-robot
ascend-robot成员
3月24日 评论:
流水线 PR-pipeline_pytorch#12072 (重试第2次)已完成
阶段 任务名 状态 详情
编译构建 Build_X86 >>>
Build_ARM >>>
Build_LibTorch_x86 >>>
Build_X86_torchair 🛑 >>>
Build_ARM_torchair 🛑 >>>
Build_ARM_inductor 🛑 >>>
恶意代码检查 Antipoison >>>
编码安全与规范检查 CodeCheck >>>
check_error >>>
开源片段检查 SCA >>>
开发者测试 UT_X86_Part_01 🛑 >>>
UT_X86_Part_02 🛑 >>>
UT_ARM_A3_Part_01 🛑 >>>
UT_ARM_A3_Part_02 🛑 >>>
UT_DIST_X86_Part_01 >>>
UT_DIST_X86_Part_02 >>>
UT_DIST_X86_Part_03 >>>
UT_DIST_X86_Part_04 >>>
UT_inductor_Part_01 🛑 >>>
UT_inductor_Part_02 🛑 >>>
UT_inductor_Part_03 🛑 >>>
UT_inductor_Part_04 🛑 >>>
UT_ARM_A2_Part_01 >>>
UT_ARM_A2_Part_02 >>>
流水线 PR-pipeline_pytorch >>>
likedislike
wjlflyer成员
3月24日 评论:

/lgtm

likedislike
renyujin成员
3月25日 评论:

/lgtm
/approve

likedislike
ascend-robotascend-robot成员
3月25日 添加了label:approvedlgtm
ascend-robot
ascend-robot成员
3月25日 评论:

Review Guide

This pull-request passes review.
Committers who wrote a comment of /approve are: renyujin.
Reviewers who wrote a comment of /lgtm are: renyujin, wjlflyer.

likedislike
ascend-robotascend-robot成员
3月25日 合入了pull request