torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_wrapper NPU 测试适配方案

一、API 功能说明

1. torch.distributed.algorithms.ddp_comm_hooks.default_hooks.fp16_compress_wrapper

该 API 接收一个 DDP 通信 hook，并返回包装后的 hook。包装函数执行时先将 GradBucket 缓冲区转换为 torch.float16，再把原 hook_state 和 bucket 传给内部 hook；内部 hook 返回的 Future 完成后，包装函数通过回调将结果复制回 bucket 缓冲区并返回 Future。PyTorch 文档将其描述为把通信输入压缩到 FP16，并在通信完成后恢复输入数据类型；在 DDP 梯度回写链路中，最终参数梯度保持原始 dtype。fp16_compress_wrapper(allreduce_hook) 与 fp16_compress_hook 的用途等价，也可以包装 PowerSGD 等其他返回 Future 的通信 hook。

该 API 的唯一显式入参 hook 必须是可调用对象，并应接受 (state, bucket)、返回 torch.futures.Future[torch.Tensor]。本次测试在 NPU 上验证 allreduce_hook、powerSGD_hook 和自定义 Future hook，以及 torch.float32、torch.float16、torch.bfloat16 输入。

二、PyTorch 社区测试用例情况

PyTorch 社区 test/distributed/test_c10d_nccl.py 中存在以下相关测试用例：

社区两个测试入口均直接验证目标 API，但依赖至少两张 CUDA 设备和 NCCL。现有覆盖未验证 HCCL/NPU、static_graph、process_group=None、FP16/BF16 输入、Future 与 state 传递契约、单 NPU执行、自定义子进程组、多 bucket、FP16 溢出边界及异常入参。

三、API 测试用例补齐及适配方案

1. 新增测试用例

测试内容：

在 test/distributed/algorithms/ddp_comm_hooks/test_fp16_compress_wrapper.py 中新增可扩展的独立测试文件和 TestFp16CompressWrapper 测试类。四个版本的测试文件内容一致，共包含 15 个测试方法：

测试通过 torch.accelerator.current_accelerator() 获取设备类型，参与运算的张量均迁移到 NPU。双进程场景使用 HCCL、动态空闲端口和显式进程组清理；Future 类型使用公共底层基类 torch._C.Future，兼容 HCCL/C++ Future。

新增原因及必要性：

PyTorch 社区用例已覆盖 CUDA/NCCL 下包装 all-reduce、PowerSGD 和 bucket view 的基本一致性，但没有独立的 HCCL/NPU验证。新增文件补齐 NPU直接验证，并覆盖内部 hook 类型、全部显式入参的正常与异常等价类、组合配置、dtype/Future/state 契约、边界值、子进程组、多 bucket 和单设备执行，为目标 API 的 NPU行为提供更完整的回归保护。

2. 单 NPU与多 NPU执行条件适配

适配原因：

13 个 DDP/HCCL 场景固定使用 WORLD_SIZE=2 并按 rank 设置设备；单 NPU服务器无法为 rank 1 设置设备。单 NPU环境仍可直接验证包装函数的压缩、回调和异常契约，因此不应跳过整个测试文件。

具体适配方案：

对统一启动双进程测试的 _spawn helper 应用 skipIfUnsupportMultiNPU(WORLD_SIZE)。服务器少于两张 NPU时，仅跳过依赖双进程的 13 个测试，继续执行 test_fp16_compress_wrapper_single_npu_contract 和 test_fp16_compress_wrapper_invalid_arguments；具备至少两张 NPU时执行全部 15 个测试。该处理避免非法设备索引，同时保留单 NPU环境可完成的有效功能验证。

四、资料变更

不涉及。

API本身已经适配，master/docs/zh/api/native_api目录文档已经记录，无需资料变更。

五、接口变更

不涉及。

本次只是纯新增测试用例，不涉及到API本身的调整。

六、功能验证

v2.7.1

执行命令：

python test/distributed/algorithms/ddp_comm_hooks/test_fp16_compress_wrapper.py -v

结果：

test_fp16_compress_wrapper_allreduce ... ok
test_fp16_compress_wrapper_allreduce_grad_is_view ... ok
test_fp16_compress_wrapper_allreduce_grad_is_view_static_graph ... ok
test_fp16_compress_wrapper_allreduce_none_pg ... ok
test_fp16_compress_wrapper_allreduce_static_graph ... ok
test_fp16_compress_wrapper_custom_subgroup ... ok
test_fp16_compress_wrapper_future_dtype_and_state ... ok
test_fp16_compress_wrapper_invalid_arguments ... ok
test_fp16_compress_wrapper_multiple_buckets ... ok
test_fp16_compress_wrapper_predivide_overflow_boundary ... ok
test_fp16_compress_wrapper_powersgd ... ok
test_fp16_compress_wrapper_powersgd_grad_is_view ... ok
test_fp16_compress_wrapper_powersgd_grad_is_view_static_graph ... ok
test_fp16_compress_wrapper_powersgd_static_graph ... ok
test_fp16_compress_wrapper_single_npu_contract ... ok

----------------------------------------------------------------------
Ran 15 tests in 215.079s

OK

v2.11.0

执行命令：

python test/distributed/algorithms/ddp_comm_hooks/test_fp16_compress_wrapper.py -v

结果：

test_fp16_compress_wrapper_allreduce ... ok
test_fp16_compress_wrapper_allreduce_grad_is_view ... ok
test_fp16_compress_wrapper_allreduce_grad_is_view_static_graph ... ok
test_fp16_compress_wrapper_allreduce_none_pg ... ok
test_fp16_compress_wrapper_allreduce_static_graph ... ok
test_fp16_compress_wrapper_custom_subgroup ... ok
test_fp16_compress_wrapper_future_dtype_and_state ... ok
test_fp16_compress_wrapper_invalid_arguments ... ok
test_fp16_compress_wrapper_multiple_buckets ... ok
test_fp16_compress_wrapper_predivide_overflow_boundary ... ok
test_fp16_compress_wrapper_powersgd ... ok
test_fp16_compress_wrapper_powersgd_grad_is_view ... ok
test_fp16_compress_wrapper_powersgd_grad_is_view_static_graph ... ok
test_fp16_compress_wrapper_powersgd_static_graph ... ok
test_fp16_compress_wrapper_single_npu_contract ... ok

----------------------------------------------------------------------
Ran 15 tests in 209.623s

OK

v2.12.0

执行命令：

python test/distributed/algorithms/ddp_comm_hooks/test_fp16_compress_wrapper.py -v

结果：

test_fp16_compress_wrapper_allreduce ... ok
test_fp16_compress_wrapper_allreduce_grad_is_view ... ok
test_fp16_compress_wrapper_allreduce_grad_is_view_static_graph ... ok
test_fp16_compress_wrapper_allreduce_none_pg ... ok
test_fp16_compress_wrapper_allreduce_static_graph ... ok
test_fp16_compress_wrapper_custom_subgroup ... ok
test_fp16_compress_wrapper_future_dtype_and_state ... ok
test_fp16_compress_wrapper_invalid_arguments ... ok
test_fp16_compress_wrapper_multiple_buckets ... ok
test_fp16_compress_wrapper_predivide_overflow_boundary ... ok
test_fp16_compress_wrapper_powersgd ... ok
test_fp16_compress_wrapper_powersgd_grad_is_view ... ok
test_fp16_compress_wrapper_powersgd_grad_is_view_static_graph ... ok
test_fp16_compress_wrapper_powersgd_static_graph ... ok
test_fp16_compress_wrapper_single_npu_contract ... ok

----------------------------------------------------------------------
Ran 15 tests in 214.918s

OK

master

执行命令：

python test/distributed/algorithms/ddp_comm_hooks/test_fp16_compress_wrapper.py -v

结果：

test_fp16_compress_wrapper_allreduce ... ok
test_fp16_compress_wrapper_allreduce_grad_is_view ... ok
test_fp16_compress_wrapper_allreduce_grad_is_view_static_graph ... ok
test_fp16_compress_wrapper_allreduce_none_pg ... ok
test_fp16_compress_wrapper_allreduce_static_graph ... ok
test_fp16_compress_wrapper_custom_subgroup ... ok
test_fp16_compress_wrapper_future_dtype_and_state ... ok
test_fp16_compress_wrapper_invalid_arguments ... ok
test_fp16_compress_wrapper_multiple_buckets ... ok
test_fp16_compress_wrapper_predivide_overflow_boundary ... ok
test_fp16_compress_wrapper_powersgd ... ok
test_fp16_compress_wrapper_powersgd_grad_is_view ... ok
test_fp16_compress_wrapper_powersgd_grad_is_view_static_graph ... ok
test_fp16_compress_wrapper_powersgd_static_graph ... ok
test_fp16_compress_wrapper_single_npu_contract ... ok

----------------------------------------------------------------------
Ran 15 tests in 189.099s

OK