已合并
[ZBAL] 修复代码安全编码问题 #1189
[ZBAL] 修复代码安全编码问题 #1189
已合并
pengchenhui创建于 8月28日
pengchenhui
pengchenhui成员
8月28日

合入来源

问题/功能描述

[ZBAL] 修复代码安全编码问题

修改方案描述

[ZBAL] 修复代码安全编码问题

  1. 修复AICPU启动时空句柄未检查的问题
  2. 修复共享内存创建时的参数校验逻辑
  3. 修复进程组析构时的错误处理逻辑
  4. 修复Moe分发算子的偏移量类型溢出问题
  5. 修复NPU通信器的内存释放和析构问题

image.png

image.png

image.png

是否涉及UT/ST

开发自检

likedislike
Pull Request已成功合入, 合并人@victor7wang
(感谢 pengchenhui 的贡献)
pengchenhuipengchenhui成员
8月28日 创建了 pull request,commit fb87153d
pengchenhuipengchenhui成员
8月28日 关联了issue:[Bug]: [ZBAl] 代码安全编码问题需要整改
atomgit-bot
atomgit-bot
8月28日 评论:

变更摘要

该 PR(#1189)针对 ZBAL 的代码安全编码问题进行系统性修复,主要围绕空句柄/空指针校验、参数合法性校验、失败路径的资源回滚与析构清理以及整数溢出防护展开,具体覆盖 AICPU 启动空句柄未检查、共享内存创建参数校验、Python 绑定库加载异常处理与 allocator 切换时机、进程组析构错误处理、Moe 分发算子偏移量类型溢出,以及 NPU 通信器的内存释放与析构问题,目标是避免空指针使用、资源(内存/句柄/通信组)泄漏和整数截断等安全编码缺陷。

主要改动

  • Moe 分发算子偏移量溢出修复: zbal_moe_distribute_combine.h 将 epDataOffsetOnWin_ 与 tpDataOffsetOnWin_ 从 uint32_t 改为 uint64_t,并在 Init 中通过 static_cast<uint64_t> 提升乘法运算(如 epRankId_ * moeExpertPerRankNum_ * expertPerSizeOnWin_ 及 axisH_ * sizeof(...)),避免偏移量计算溢出截断。
  • NPU 通信器资源释放与析构修复: zbal_communicator.cpp 的 Create 在 AssignGatherGroupId 失败时调用 DestroyInner(commInner) 回滚已登记资源,CreateInner 的多个失败路径补充 comm->DecreaseRef() 释放创建者引用;zbal_npu_communicator_base.cpp 的 DestroyProfMemory 新增 AclrtFreeHost(perfHostMemory_) 释放 host 内存;zbal_npu_communicator_aicpu.h 为 NpuCommunicatorAICPU 新增析构函数重写并显式调用 UnInitialize(),避免基类析构的虚分派跳过 AICPU 资源清理。
  • 进程组析构错误处理: zbal_pytorch_process_group.cpp 的 ~ProcessGroupZBAL 在 zbal_comm_destroy 失败时不再提前 return,改为记录警告日志后继续销毁其余 comm 并清空 groupP2pComms_,防止部分失败导致剩余通信资源泄漏。
  • Python 绑定库加载与 allocator 切换: zbal_module.py 在目录下找不到 zbal.*.so 时抛出 ImportError 提示安装异常;switch_to_allocator 将 torch_npu.npu.memory.change_current_allocator(new_alloc) 的执行时机调整到各 allocator 回调(set_*_fn)配置完成之后。
  • AICPU 启动空句柄校验与共享内存参数校验: zbal_npu_aicpu_launcher.cpp 的 Launch 在 launchHandle 为空时记录错误并返回 Z_INVALID_PARAM;zbal_mem_mf_bootstrap.cpp 的 CreateSHMSpace 将断言改为显式校验,当 totalMemSize > spaceSize 时记录错误、销毁临时共享内存句柄并返回 Z_ERROR。
likedislike
不准确?
atomgit-bot
atomgit-bot
8月28日 评论:

代码审查

✅ 未发现问题

likedislike
不准确?
此处折叠了53条消息 查看更多
ascend-robotascend-robot成员
9月4日 添加了label:ci-pipeline-passed
Zzhangshuren成员
9月4日 将SoulMeister设为评审人
Zzhangshuren成员
9月4日 通过了评审
Vvictor7wang成员
9月4日 通过审查
Vvictor7wang成员
9月4日 合入了pull request