人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载本篇文章以 CANN ops-nn 仓库中的aclnnForeachNonFiniteCheckAndUnscale算子为讲解对象完整覆盖其产品支持情况、数学语义、两段式 aclnn 接口原型、参数与错误码规范、约束条件并给出可直接复制的完整调用示例。同时结合仓库内 op_graph 原型、op_host tiling 与 op_kernel 内核源码深入解释如何检测 Inf/NaN与如何完成反缩放的底层实现帮助读者在混合精度训练梯度裁剪GradScaler场景中正确、高效地使用该算子。一、算子定位与应用场景ForeachNonFiniteCheckAndUnscale是 CANN ops-nn 中 foreach 系列算子之一源码位于 foreach/foreach_non_finite_check_and_unscale对外通过 aclnn 接口 aclnnForeachNonFiniteCheckAndUnscale 提供能力。它在混合精度训练中承担两项关键职责Inf/NaN 巡检遍历梯度张量列表scaledGrads中的每一个 Tensor检测是否存在 Inf 或 NaN梯度反缩放无论是否检测到异常都会将scaledGrads中所有 Tensor 统一乘以invScale恢复真实的梯度数值。该算子与 PyTorch 混合精度 APItorch.cuda.amp.GradScaler中的unscale_语义一致训练过程中梯度先被scale放大以避免下溢在更新参数前需要检测梯度是否溢出Inf/NaN若溢出则跳过本次更新否则用invScale还原梯度。ForeachNonFiniteCheckAndUnscale正是把检查与还原合并为一次算子调用减少一次核下发开销。说明该算子为inplace 语义——scaledGrads既是输入也是输出反缩放结果直接写回原张量foundInf同样可被改写检测到异常时置 1.0。这与普通输入/输出分离的算子使用方式有明显区别使用时需注意张量内存可写。二、产品支持情况根据 aclnnForeachNonFiniteCheckAndUnscale.md 中产品支持情况一节的说明该算子在以下产品上受支持产品是否支持Ascend 950PR / Ascend 950DT支持Atlas A3 训练系列产品 / Atlas A3 推理系列产品支持Atlas A2 训练系列产品 / Atlas A2 推理系列产品支持Atlas 200I / 500 A2 推理产品不支持Atlas 推理系列产品支持Atlas 训练系列产品不支持此外从算子目录下的 README.md 与算子定义文件 foreach_non_finite_check_and_unscale_def.cpp 可以看到Kirin X90、Kirin 9030 处理器系列产品同样支持该算子且这两个平台上scaled_grads的数据类型不支持 BFLOAT16这一点在调用时需特别注意。三、功能说明与计算公式3.1 语义描述接口功能遍历scaledGrads中的所有 Tensor检查是否存在 Inf 或 NaN如果存在则将foundInf设置为 1.0否则foundInf的值保持不变同时对scaledGrads中的所有 Tensor 执行反缩放乘以invScale。3.2 计算公式foundInf的更新规则$$ foundInf \begin{cases}1.0, 当(Inf \in scaledGrads)或(NaN \in scaledGrads),\ foundInf, 其他. \end{cases} $$scaledGrads的反缩放规则作用于列表中的每一个张量$$ scaledGrads_i scaledGrads_i * invScale $$3.3 与源码语义的印证算子 IR 原型 foreach_non_finite_check_and_unscale_proto.h 中对该算子给出了同样精确的注释检测scaled_grads中是否存在 Inf 或 NaN存在则将found_inf置 1不存在则不对其操作最终将scaled_grads的所有值乘以inv_scale并写回。其输入定义如下REG_OP(ForeachNonFiniteCheckAndUnscale) .DYNAMIC_INPUT(scaled_grads, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16})) .INPUT(found_inf, TensorType({DT_FLOAT})) .INPUT(inv_scale, TensorType({DT_FLOAT})) .OP_END_FACTORY_REG(ForeachNonFiniteCheckAndUnscale)其中scaled_grads为动态输入张量列表found_inf与inv_scale为单元素 FLOAT32 张量与 aclnn 接口的参数一一对应。四、两段式接口与函数原型CANN 算子库的 aclnn 接口采用两段式调用设计必须先调用GetWorkspaceSize接口获取计算所需的 workspace 大小以及封装了算子计算流程的执行器再调用真正的执行接口完成计算。4.1 第一段GetWorkspaceSize申请空间与获取执行器aclnnStatus aclnnForeachNonFiniteCheckAndUnscaleGetWorkspaceSize( const aclTensorList *scaledGrads, const aclTensor *foundInf, const aclTensor *invScale, uint64_t *workspaceSize, aclOpExecutor **executor)4.2 第二段执行计算aclnnStatus aclnnForeachNonFiniteCheckAndUnscale( void *workspace, uint64_t workspaceSize, aclOpExecutor *executor, aclrtStream stream)五、第一段接口参数详解以下参数说明摘自 aclnnForeachNonFiniteCheckAndUnscale.md 并保留全部细节参数名输入/输出描述使用说明数据类型数据格式维度(shape)非连续TensorscaledGradsaclTensorList*输入/输出表示进行反缩放计算的输入和输出张量列表对应公式中的scaledGrads不支持空 Tensor该参数中所有 Tensor 的数据类型保持一致支持的最大长度为 256 个FLOAT32、FLOAT16、BFLOAT16ND0-8×foundInfaclTensor*输入/输出表示用来标记输入scaledGrads中是否存在 Inf 或 NaN 的张量对应公式中的foundInf不支持空 Tensor仅包含一个元素FLOAT32ND0-8×invScaleaclTensor*输入表示进行反缩放计算的张量对应公式中的invScale不支持空 Tensor仅包含一个元素FLOAT32ND0-8√workspaceSizeuint64_t*输出返回需要在 Device 侧申请的 workspace 大小-----executoraclOpExecutor**输出返回 op 执行器包含了算子计算流程-----关键约束解读scaledGrads 是动态张量列表接口通过aclTensorList承载任意数量的 Tensor但最大长度限制为 256 个该限制同样体现在 tiling 头文件 foreach_non_finite_check_and_unscale_tiling.h 的MAX_TENSOR_COUNT 256常量中列表中所有 Tensor 数据类型必须一致允许的类型为 FLOAT32、FLOAT16、BFLOAT16Kirin X90/Kirin 9030 上不支持 BFLOAT16foundInf 与 invScale 都是单元素 FLOAT32 张量shape 为{1}非连续 Tensor 支持情况inplace 操作的输入/输出scaledGrads、foundInf不支持非连续 TensorinvScale则支持非连续 Tensor表中√。README 中亦强调inplace 操作的输入/输出不支持非连续 Tensor。六、返回值与错误码两个接口的返回值均为aclnnStatus状态码具体可参见 aclnn 返回码。第一段接口完成入参校验出现以下场景时报错返回码错误码描述ACLNN_ERR_PARAM_NULLPTR161001传入的 scaledGrads、foundInf、invScale 是空指针ACLNN_ERR_PARAM_INVALID161002scaledGrads、foundInf、invScale 的数据类型不在支持的范围之内ACLNN_ERR_INNER_TILING_ERROR561002scaledGrads 长度超过限制ACLNN_ERR_INNER_TILING_ERROR561002scaledGrads 中的 Tensor 数据类型不相同ACLNN_ERR_INNER_TILING_ERROR561002foundInf 或 invScale 中的元素个数不为 1从 tiling 源码 foreach_non_finite_check_and_unscale_tiling.cpp 可以印证上述部分校验逻辑Init()阶段会逐一检查scaled_grads各 Tensor 的数据类型是否一致、shape size 是否大于 0CheckParams()阶段会校验scaled_grads的 dtype 必须为 float16/bfloat16/float且found_inf、inv_scale的 dtype 必须为 float。七、第二段接口参数详解参数名输入/输出描述workspace输入在 Device 侧申请的 workspace 内存地址workspaceSize输入在 Device 侧申请的 workspace 大小由第一段接口aclnnForeachNonFiniteCheckAndUnscaleGetWorkspaceSize获取executor输入op 执行器包含了算子计算流程stream输入指定执行任务的 Stream八、约束说明确定性计算aclnnForeachNonFiniteCheckAndUnscale默认确定性实现即相同输入在多次运行中产生确定一致的结果便于调试与结果比对。九、源码级原理剖析为帮助读者深入理解该算子的实现机制下面结合仓库源码解析其 tiling 与 kernel 两个层面的设计。9.1 Tiling数据如何切分到多核Tiling 逻辑位于 foreach_non_finite_check_and_unscale_tiling.cpp按 32 字节对齐切分每个 Tensor 的元素数先按 32 字节BYTE_BLOCK对齐tensorDataCountAlignedList[i] CeilDiv(count, elementsPerBlock) * elementsPerBlock再统计总数据量多核分配根据平台核数coreNumPlatform与总数据量计算实际需要的核数通过AssignDataToEachCore将各 Tensor 的区间起始 Tensor 下标、起始偏移、结束 Tensor 下标、结束偏移均衡分配给每个核写入tensorStartList / tensorEndList / tensorStartOffsetList / tensorEndOffsetList等 tiling 数据UB 内存划分DivideUbMemory根据平台 UB 大小、数据类型float 与 half/bf16 的系数不同估算scaledGradsUbSize与reduceTempValUbSize并开启双缓冲BUFFER_NUM 2提升流水效率workspace 大小固定申请 32 字节WORKSPACE_SIZE用于承载张量列表中各 Tensor 的地址偏移表内核通过该表定位每个 Tensor 的实际 GM 地址见GetTensorAddr中对指针偏移的解析Tiling Key 区分类型FLOAT→1、HALF→2、BFLOAT16→3内核据此选择模板实例化类型。9.2 KernelInf/NaN 检测与反缩放的计算路径内核实现位于 foreach_non_finite_check_and_unscale.cpp 与 foreach_non_finite_check_and_unscale_n_d.h整体流程为经典的CopyIn → Cast(非 float) → Compute → Cast(非 float) → CopyOut流水数据搬入 UBCopyIn按块每块不超过maxDataCount将 GM 数据搬入copyInQueue非 32 字节对齐的尾部使用DataCopyPad右填充类型提升仅 FLOAT16/BFLOAT16CastToFloat将 half/bf16 数据转换为 float保证后续计算与检测在 float 精度下进行Inf/NaN 检测Compute中利用向量归约指令ReduceMax求取本块数据的最大值随后在IsNonFinite中检查该浮点数的位模式——(tempValue 0x7FFFFFFF) 23 0xFF即指数位全为 1意味着该数是 Inf 或 NaN若最大值正常再通过ReduceMin检查最小值这是因为 NaN 在归约中可能导致最大值/最小值出现差异双端检查可提高检出可靠性。任一核检测到异常即通过foundInfGM.SetValue(0, 1.0)将foundInf置为 1.0并用haveFoundInf标志保证同一核内只做一次检测反缩放Muls(computeOutLT, computeInLT, invScaleVal, dataCount)将整块数据乘以从 GM 读入的invScale单元素值类型还原仅非 floatCastToOriginalType使用CAST_RINT舍入模式将结果转回原 half/bf16 类型结果写回 GMCopyOut将 UB 中结果搬回原scaledGrads地址inplace 写回。由于scaledGrads在 GM 侧存储的是一张地址表首个地址为偏移表基址其后依次存放各 Tensor 的实际地址GetTensorAddr通过两级指针解析获得第index个 Tensor 的真实 GM 地址这正是 aclnn 接口以aclTensorList传入多张量时的内核侧解包方式。9.3 平台差异与编译配置不同产品的算子配置dtype、format、动态 shape 支持等定义于 foreach_non_finite_check_and_unscale_def.cpp并为每个平台维护独立的 binary 配置例如 ascend910b 的 binary.json 中分别为 float16、float32、bfloat16 三种输入类型生成独立的算子二进制条目。Kirin 系列与 ascend310p 等平台的配置中scaled_grads仅支持 FLOAT16、FLOAT与 README 中Kirin 系列不支持 BFLOAT16的说明一致。十、完整调用示例示例代码位于 examples/test_aclnn_foreach_non_finite_check_and_unscale.cpp与 aclnnForeachNonFiniteCheckAndUnscale.md 中调用示例章节的代码等价以下为完整实现。具体编译与执行流程请参考编译与运行样例。#include iostream #include vector #include acl/acl.h #include aclnnop/aclnn_foreach_non_finite_check_and_unscale.h #define CHECK_RET(cond, return_expr) \ do { \ if (!(cond)) { \ return_expr; \ } \ } while (0) #define LOG_PRINT(message, ...) \ do { \ printf(message, ##__VA_ARGS__); \ } while (0) int64_t GetShapeSize(const std::vectorint64_t shape) { int64_t shapeSize 1; for (auto i : shape) { shapeSize * i; } return shapeSize; } int Init(int32_t deviceId, aclrtStream* stream) { // 固定写法资源初始化 auto ret aclInit(nullptr); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclInit failed. ERROR: %d\n, ret); return ret); ret aclrtSetDevice(deviceId); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSetDevice failed. ERROR: %d\n, ret); return ret); ret aclrtCreateStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtCreateStream failed. ERROR: %d\n, ret); return ret); return 0; } template typename T int CreateAclTensor(const std::vectorT hostData, const std::vectorint64_t shape, void** deviceAddr, aclDataType dataType, aclTensor** tensor) { auto size GetShapeSize(shape) * sizeof(T); // 调用aclrtMalloc申请device侧内存 auto ret aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMalloc failed. ERROR: %d\n, ret); return ret); // 调用aclrtMemcpy将host侧数据复制到device侧内存上 ret aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtMemcpy failed. ERROR: %d\n, ret); return ret); // 计算连续tensor的strides std::vectorint64_t strides(shape.size(), 1); for (int64_t i shape.size() - 2; i 0; i--) { strides[i] shape[i 1] * strides[i 1]; } // 调用aclCreateTensor接口创建aclTensor *tensor aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND, shape.data(), shape.size(), *deviceAddr); return 0; } int main() { // 1. 固定写法device/stream初始化参考acl API手册 // 根据自己的实际device填写deviceId int32_t deviceId 0; aclrtStream stream; auto ret Init(deviceId, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(Init acl failed. ERROR: %d\n, ret); return ret); // 2. 构造输入与输出需要根据API的接口自定义构造 std::vectorint64_t selfShape1 {2, 3}; std::vectorint64_t selfShape2 {1, 3}; std::vectorint64_t foundInfShape {1}; std::vectorint64_t invScaleShape {1}; void* input1DeviceAddr nullptr; void* input2DeviceAddr nullptr; void* foundInfDeviceAddr nullptr; void* invScaleDeviceAddr nullptr; aclTensor* input1 nullptr; aclTensor* input2 nullptr; aclTensor* foundInf nullptr; aclTensor* invScale nullptr; std::vectorfloat input1HostData {1, 2, 3, 4, 5, 6}; std::vectorfloat input2HostData {7, 8, 9}; std::vectorfloat foundInfHostData {-1.0f}; std::vectorfloat invScaleHostData {3.1f}; // 创建input1 aclTensor ret CreateAclTensor(input1HostData, selfShape1, input1DeviceAddr, aclDataType::ACL_FLOAT, input1); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建input2 aclTensor ret CreateAclTensor(input2HostData, selfShape2, input2DeviceAddr, aclDataType::ACL_FLOAT, input2); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建foundInf aclTensor ret CreateAclTensor(foundInfHostData, foundInfShape, foundInfDeviceAddr, aclDataType::ACL_FLOAT, foundInf); CHECK_RET(ret ACL_SUCCESS, return ret); // 创建invScale aclTensor ret CreateAclTensor(invScaleHostData, invScaleShape, invScaleDeviceAddr, aclDataType::ACL_FLOAT, invScale); CHECK_RET(ret ACL_SUCCESS, return ret); std::vectoraclTensor* tempInput{input1, input2}; aclTensorList* tensorListInput aclCreateTensorList(tempInput.data(), tempInput.size()); // 3. 调用CANN算子库API需要修改为具体的API名称 uint64_t workspaceSize 0; aclOpExecutor* executor; // 调用aclnnForeachNonFiniteCheckAndUnscale第一段接口 ret aclnnForeachNonFiniteCheckAndUnscaleGetWorkspaceSize(tensorListInput, foundInf, invScale, workspaceSize, executor); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnForeachNonFiniteCheckAndUnscaleGetWorkspaceSize failed. ERROR: %d\n, ret); return ret); // 根据第一段接口计算出的workspaceSize申请device内存 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(allocate workspace failed. ERROR: %d\n, ret); return ret); } // 调用aclnnForeachNonFiniteCheckAndUnscale第二段接口 ret aclnnForeachNonFiniteCheckAndUnscale(workspaceAddr, workspaceSize, executor, stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclnnForeachNonFiniteCheckAndUnscale failed. ERROR: %d\n, ret); return ret); // 4. 固定写法同步等待任务执行结束 ret aclrtSynchronizeStream(stream); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(aclrtSynchronizeStream failed. ERROR: %d\n, ret); return ret); // 5. 获取输出的值将device侧内存上的结果复制至host侧需要根据具体API的接口定义修改 auto size GetShapeSize(selfShape1); std::vectorfloat self1Data(size, 0); ret aclrtMemcpy(self1Data.data(), self1Data.size() * sizeof(self1Data[0]), input1DeviceAddr, size * sizeof(self1Data[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy self1 result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(self1 result[%ld] is: %f\n, i, self1Data[i]); } size GetShapeSize(selfShape2); std::vectorfloat self2Data(size, 0); ret aclrtMemcpy(self2Data.data(), self2Data.size() * sizeof(self2Data[0]), input2DeviceAddr, size * sizeof(self2Data[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy self2 result from device to host failed. ERROR: %d\n, ret); return ret); for (int64_t i 0; i size; i) { LOG_PRINT(self2 result[%ld] is: %f\n, i, self2Data[i]); } size GetShapeSize(foundInfShape); std::vectorfloat foundInfData(size, 0); ret aclrtMemcpy(foundInfData.data(), foundInfData.size() * sizeof(foundInfData[0]), foundInfDeviceAddr, size * sizeof(foundInfData[0]), ACL_MEMCPY_DEVICE_TO_HOST); CHECK_RET(ret ACL_SUCCESS, LOG_PRINT(copy foundInf result from device to host failed. ERROR: %d\n, ret); return ret); LOG_PRINT(foundInf result is: %f\n, foundInfData[0]); // 6. 释放aclTensor需要根据具体API的接口定义修改 aclDestroyTensorList(tensorListInput); aclDestroyTensor(foundInf); aclDestroyTensor(invScale); // 7.释放device资源需要根据具体API的接口定义修改 aclrtFree(input1DeviceAddr); aclrtFree(input2DeviceAddr); aclrtFree(foundInfDeviceAddr); aclrtFree(invScaleDeviceAddr); if (workspaceSize 0) { aclrtFree(workspaceAddr); } aclrtDestroyStream(stream); aclrtResetDevice(deviceId); aclFinalize(); return 0; }10.1 示例运行结果预期以示例数据为例input1 {1,2,3,4,5,6}、input2 {7,8,9}、foundInf -1.0、invScale 3.1由于输入中不含 Inf/NaNfoundInf保持原值-1.0不会被改写为 1.0input1每个元素乘以 3.1得到{3.1, 6.2, 9.3, 12.4, 15.5, 18.6}input2每个元素乘以 3.1得到{21.7, 24.8, 27.9}。若将input1HostData中任一元素改为std::numeric_limitsfloat::infinity()或std::nanf()则foundInf会被置为1.0同时所有张量仍完成反缩放。这一行为可由仓库测试脚本 tests/assets/golden.py 中的 golden 函数验证它逐张量检查np.isinf/np.isnan存在异常时置found_inf[0] 1.0并返回inp * inv_scale后的结果与公式和上述预期完全一致。十一、测试与验证仓库为算子提供了完整的测试支撑UTop_hosttest_foreach_non_finite_check_and_unscale_tiling.cpp 覆盖 tiling 数据生成逻辑UTop_kerneltest_foreach_non_finite_check_and_unscale.cpp 配套 gen_data.py 与 compare_data.py 完成数据生成与结果比对STarch35ttk_kernel_foreach_non_finite_check_and_unscale.csv 定义了全量场景的测试用例参数。这些测试共同保证了算子在多产品、多数据类型FLOAT32/FLOAT16/BFLOAT16、含 Inf/NaN 与不含异常两种路径下结果正确。十二、总结ForeachNonFiniteCheckAndUnscale是 CANN ops-nn 中面向混合精度训练的关键组合算子一次调用同时完成梯度溢出检测与梯度反缩放。使用时的要点可归纳为两段式调用先GetWorkspaceSize取 workspace 大小与执行器再执行计算workspace 需在 Device 侧申请并在结束后释放inplace 语义scaledGrads、foundInf会被原地改写且不支持非连续 Tensor参数约束列表长度 ≤ 256、列表内 dtype 一致、foundInf/invScale为单元素 FLOAT32产品差异Atlas 200I/500 A2 与 Atlas 训练系列不支持Kirin 系列不支持 BFLOAT16底层机制tiling 层按 32 字节对齐、多核均衡切分并固定分配 32 字节 workspacekernel 层通过 ReduceMax/ReduceMin 指数位全 1 判定 Inf/NaN以 Muls 完成反缩放非 float 类型先提升精度再还原兼顾正确性与确定性。如需进一步了解 aclnn 接口的通用调用框架、返回码语义或样例工程的编译运行方式可继续阅读 两段式接口说明、aclnn 返回码 与编译与运行样例。赞分享人工智能算子库深度学习CANNAscend【免费下载链接】ops-nn本项目是CANN提供的神经网络类计算算子库实现网络在NPU上加速计算。项目地址https://gitcode.com/cann/ops-nn点击查看免费下载相关推荐CANN ops-nn ForeachNonFiniteCheckAndUnscale 算子详解梯度 Inf/NaN 检测与反缩放实战CANN ops nn ForeachNonFiniteCheckAndUnscale 算子详解梯度 Inf/NaN 检测与反缩放实战 导读 本文基于 CAN人工智能算子库深度学习CANNAscend猫抓浏览器插件终极免费资源嗅探工具轻松下载网页媒体资源猫抓浏览器插件终极免费资源嗅探工具轻松下载网页媒体资源 你是否也曾遇到这样的情况在线观看精彩的视频教程想要收藏却无法下载浏览网页时发现精美的图片素材人工智能算子库深度学习CANNAscendCANN ops-nn 算子解析aclnnHardswishBackward 两段式 ACLNN 接口实现 HardSwish 反向梯度计算CANN ops nn 算子解析aclnnHardswishBackward 两段式 ACLNN 接口实现 HardSwish 反向梯度计算 导读 aclnn人工智能算子库深度学习CANNAscend创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考