【合入来源】

• 问题单：https://gitcode.com/Ascend/pytorch/issues/1883

本 PR 覆盖以下 API：

• torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader
• torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.read_metadata
• torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.prepare_local_plan
• torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.prepare_global_plan
• torch.distributed.checkpoint.format_utils.BroadcastingTorchSaveReader.read_data

【修改方案】

一、解决方案

针对 issue 中提出的 BroadcastingTorchSaveReader API 专项测试缺失问题，本 PR 新增 test_broadcasting_reader_api.py，以“API 级行为验证 + NPU 分布式数据面验证”的方式补齐相关测试。

本 PR 做如下修改：

• 新增 test_broadcasting_reader_api.py，覆盖 BroadcastingTorchSaveReader 的 metadata、plan 透传、数据加载和错误边界行为。
• 使用 with_comms 和 skipIfUnsupportMultiNPU(2) 验证 read_data 在 NPU/HCCL 分布式环境下的 broadcast 加载路径。
• 使用 DynamicMetaLoadPlanner 覆盖普通 torch save 文件无 DCP metadata 时的推荐加载组合。
• 不修改 torch-npu 生产代码，不改变 API 签名、返回结构或运行时语义。

新增测试文件：

test/distributed/checkpoint/test_broadcasting_reader_api.py

二、用例完备性说明

本次新增 6 个测试用例，覆盖 BroadcastingTorchSaveReader 的生命周期入口、计划透传、数据加载主路径和错误边界。

1. test_read_metadata_returns_empty_metadata

覆盖 API：read_metadata

验证内容：

• 返回对象是 Metadata。
• state_dict_metadata 为空字典。
• 不需要真实 checkpoint 文件，也不依赖分布式环境。

该用例证明 BroadcastingTorchSaveReader 在 torch save 无 metadata 文件的情况下，能按设计返回空 metadata，并把 metadata 构造职责交给 planner。

2. test_prepare_local_plan_returns_input_plan

覆盖 API：prepare_local_plan

验证内容：

• 构造带 storage_data 和 planner_data 的 LoadPlan。
• 调用 prepare_local_plan 后返回值与输入对象是同一个对象。
• 输入计划中的附加数据不会被替换或丢失。

该用例证明 reader 不会修改本地 planner 结果，符合该 API 的 no-op 设计。

3. test_prepare_global_plan_returns_input_plans

覆盖 API：prepare_global_plan

验证内容：

• 构造多个 rank 的 LoadPlan 列表。
• 调用 prepare_global_plan 后返回值与输入列表是同一个列表。
• 各 rank 的 plan 不被重排、不被替换。

该用例证明 reader 在全局 planning 阶段不引入额外调度逻辑，保持 DCP 默认计划流转。

4. test_read_data_rejects_byte_io_items

覆盖 API：read_data

验证内容：

• 手工构造 LoadItemType.BYTE_IO 类型的 ReadItem。
• 使用真实 torch save 文件路径初始化 reader。
• 调用 read_data 时断言抛出 only supports loading Tensors。

该用例覆盖官方端到端成功路径没有覆盖的错误边界，证明 reader 对非 tensor 对象的限制是显式、可预期的。

5. test_read_data_loads_torch_save_tensor_to_npu_state_dict

覆盖 API：read_data

验证内容：

• rank 0 使用 torch.save 保存普通 tensor checkpoint。
• 使用 with_comms 启动 NPU/HCCL 分布式环境。
• 每个 rank 构造 NPU 目标 tensor。
• 通过 dcp.load + DynamicMetaLoadPlanner + BroadcastingTorchSaveReader 加载。
• 校验加载后的 NPU tensor 与源 tensor 一致。

该用例覆盖 read_data 的核心成功路径，证明 torch save 文件可以经 coordinator 读取、broadcast 分发并写回 NPU state_dict。

6. test_read_data_handles_nested_state_dict

覆盖 API：read_data

验证内容：

• 使用嵌套结构 {"model": {"weight": ..., "bias": ...}} 保存 torch save checkpoint。
• 目标 state_dict 使用同样嵌套结构，但 tensor 分配在 NPU。
• 使用 DynamicMetaLoadPlanner 默认 flatten 行为加载。
• 校验嵌套 key 下的 weight 和 bias 均正确写回。

该用例覆盖 read_data 中 flatten_state_dict 为 True 时的关键路径，证明嵌套 state_dict 在 NPU 分布式加载场景下可用。

覆盖关系总结

• metadata 阶段：test_read_metadata_returns_empty_metadata
• local plan 阶段：test_prepare_local_plan_returns_input_plan
• global plan 阶段：test_prepare_global_plan_returns_input_plans
• read_data 错误边界：test_read_data_rejects_byte_io_items
• read_data 普通 tensor 成功路径：test_read_data_loads_torch_save_tensor_to_npu_state_dict
• read_data 嵌套/flatten 成功路径：test_read_data_handles_nested_state_dict

上述用例覆盖了 reader 生命周期中本 PR 涉及的全部 API，并将纯 Python 计划行为与真实 NPU 通信加载行为分开验证，便于定位问题。

三、NPU 适配说明

本次不修改 torch-npu 生产代码，仅新增测试用例。

适配验证策略：

• 使用真实 NPU tensor 作为目标 state_dict，验证 reader 写回路径能落到 NPU 设备。
• 使用 with_comms 和 skipIfUnsupportMultiNPU(2) 启动 HCCL 分布式环境，覆盖 dist.broadcast 的 NPU 通信路径。
• torch save 文件由 rank 0 写入，其他 rank 通过 barrier 等待，贴近 BroadcastingTorchSaveReader 的 coordinator 读取模型。
• 使用 DynamicMetaLoadPlanner 复现官方推荐组合，覆盖无 DCP metadata 的 torch save 加载场景。

经检查，2.7.1、2.9.0、2.10.0、2.11.0、2.12.0 和 master 中本测试依赖的 BroadcastingTorchSaveReader 核心语义、LoadPlan / ReadItem 字段以及 dcp.load 调用方式保持兼容。高版本主要是类型标注和内部异常写法调整，不影响本用例验证的 metadata、plan 透传、BYTE_IO 错误边界和 tensor broadcast 加载路径，因此不同版本不需要拆分修改测试逻辑。

为什么不需要修改 API：

• read_metadata、prepare_local_plan、prepare_global_plan 是 DCP storage reader 生命周期的通用 Python 接口，不依赖硬件后端。
• read_data 中与硬件相关的部分是 tensor device 和 dist.broadcast，这部分通过真实 NPU tensor 和 HCCL 进程组验证即可。
• 该 reader 的行为约束来自 PyTorch DCP 通用设计，NPU 适配目标是验证现有语义在 NPU 下可用，而不是改变 API 签名或返回结构。

【资料变更】

不涉及

【接口变更】

不涉及

【功能验证】

在 2.7.1、2.9.0、2.10.0、2.11.0、2.12.0, master版本上执行新增测试，均通过。

----------------------------------------------------------------------
Ran 6 tests in 1.367s

OK

【CheckList】

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

• 代码注释完备，正确记录错误日志
• 代码实现进行了返回值、空指针等校验
• PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
• PR持续集成流水线（CI）执行通过，代码检查无异常