【合入来源】

[Usage]: API一致性说明：torch.distributed.checkpoint.planner.WriteItem.tensor_storage_size 缺少专项验证用例 #1909
请勿携带内部流程信息（需求链接、问题单、内部issue等）

• 需求
• 问题单
• issue/工单
• 重构优化
• 资料更新

【修改方案】

请描述修改内容的具体实现，涉及哪些组件之间进行交互，可以用1、2、3、...进行罗列
如果是需求或者重构类的PR，需要补充详细设计文档（说明上下游组件关系、时序图、类图、DFX能力等内容）

本 PR 针对 torch.distributed.checkpoint.planner.WriteItem.tensor_storage_size 新增专项验证用例。

一、API 功能说明

torch.distributed.checkpoint.planner.WriteItem.tensor_storage_size 是 distributed checkpoint planner 中 WriteItem 提供的元数据辅助方法，用于返回当前写入项对应 tensor 的存储大小，单位为 bytes。

该 API 的核心行为如下：

1. 当 WriteItem.tensor_data 为 None 时，表示该写入项不是 tensor 写入项，例如 WriteItemType.BYTE_IO，此时返回 None。
2. 当 WriteItem.tensor_data 不为 None 时，根据 tensor_data.size 计算元素数量，并结合 tensor_data.properties.dtype 计算单个元素字节数，最终返回 tensor 存储字节数。
3. 该 API 不执行实际 tensor 数值计算，不涉及 NPU kernel 或底层算子调用，主要依赖 checkpoint planner 中已有的 tensor 元数据信息。

二、用例补齐说明

经检索，当前测试目录中未发现直接调用 WriteItem.tensor_storage_size() 并断言其返回值的专项测试用例。已有 checkpoint planner 相关测试主要覆盖 WriteItem、SavePlan、TensorWriteData 等结构和 planner 流程，但没有单独验证 tensor_storage_size 的返回值行为。

因此，本 PR 新增测试文件：
test/distributed/checkpoint/test_planner_api.py
该测试文件覆盖以下场景：
WriteItemType.TENSOR + tensor_data
构造普通 tensor 写入项，覆盖 torch.float32、torch.float16 和 torch.int8 三种 dtype，验证返回值等于：
tensor.numel() * tensor.element_size()
WriteItemType.SHARD + tensor_data
构造 shard 写入项，验证 tensor_storage_size 能够根据 shard 对应的 tensor metadata 返回正确的存储字节数。
WriteItemType.BYTE_IO 且 tensor_data 为 None
构造非 tensor 写入项，验证 tensor_storage_size 返回 None。

三、NPU 适配说明

torch.distributed.checkpoint.planner.WriteItem.tensor_storage_size 本身是元数据辅助方法，不直接依赖 NPU 算子。但测试中需要基于 tensor 构造 TensorWriteData 和 TensorProperties，因此测试 tensor 需要运行在 NPU 设备上。

本 PR 使用如下方式获取当前 accelerator：

device_type = acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu"
并通过如下方式将测试 tensor 迁移到当前设备：
tensor = torch.empty((2, 3), dtype=dtype).to(device_type)
因此，在 NPU 环境下，测试 tensor 会运行在 NPU 设备上，满足 NPU 适配要求。

四、涉及版本说明

本次用例补齐分别提交到以下目标分支：
v2.7.1
v2.9.0
v2.10.0
v2.11.0
v2.12.0
master
各版本均单独提交 PR，目标分支与源码分支一一对应。

【资料变更】

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

涉及资料支持情况检查。本 PR 不直接修改资料文件。

已检查 v2.7.1 分支 docs/zh/native_apis，当前文档中已有 torch.distributed.checkpoint.planner.WriteItem，但未检索到 torch.distributed.checkpoint.planner.WriteItem.tensor_storage_size 的方法级支持信息。

已单独提交资料补齐 PR 到 v2.7.1 分支，补充 docs/zh/native_apis 中该 API 的支持情况。资料 PR 刷新以下版本目录：

1. docs/zh/native_apis/pytorch_2-7-1/torch-distributed-checkpoint.md
2. docs/zh/native_apis/pytorch_2-9-0/torch-distributed-checkpoint.md
3. docs/zh/native_apis/pytorch_2-10-0/torch-distributed-checkpoint.md

pytorch_2-6-0 和 pytorch_2-8-0 按要求不刷新。

资料补齐 PR：https://gitcode.com/Ascend/pytorch/pull/35673

【接口变更】

请确认是否涉及跨代码仓或者客户面可见的接口变更。如涉及，需要详细说明接口以及对应的变更内容，同时需要在资料中体现。如不涉及，需填写“不涉及”

不涉及

【功能验证】

说明测试场景，测试方法。如果本次测试方式与常规单元测试不同，请详细说明您的测试步骤
新增/变更内容是否已新增/适配UT测试用例看护，并补充测试自验证截图

新增测试用例：
test/distributed/checkpoint/test_planner_api.py

测试命令：
python test/distributed/checkpoint/test_planner_api.py -v

测试场景：
WriteItemType.TENSOR + tensor_data，验证返回 tensor 存储字节数。
WriteItemType.SHARD + tensor_data，验证返回 shard tensor 存储字节数。
WriteItemType.BYTE_IO 且 tensor_data 为 None，验证返回 None。

已在以下版本执行该用例：
v2.7.1：通过
v2.9.0：通过
v2.10.0：通过
v2.11.0：通过
v2.12.0：通过
master：通过

不同版本运行用例结果示例：
V2.7.1

[root@734774725e04 workspace]# python /workspace/ascend-pytorch-v2.7.1/test/distributed/checkpoint/test_planner_api.py -v
test_write_item_tensor_storage_size_for_non_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_non_tensor) ... ok
test_write_item_tensor_storage_size_for_shard (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_shard) ... ok
test_write_item_tensor_storage_size_for_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_tensor) ... ok

----------------------------------------------------------------------
Ran 3 tests in 4.995s

OK
[root@734774725e04 workspace]# 

V2.9.0

[root@734774725e04 workspace]# python /workspace/ascend-pytorch-v2.9.0/test/distributed/checkpoint/test_planner_api.py -v

test_write_item_tensor_storage_size_for_non_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_non_tensor) ... ok
test_write_item_tensor_storage_size_for_shard (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_shard) ... ok
test_write_item_tensor_storage_size_for_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_tensor) ... ok

----------------------------------------------------------------------
Ran 3 tests in 1.023s

OK
[root@734774725e04 workspace]# 

V2.10.0

[root@734774725e04 workspace]# python /workspace/ascend-pytorch-v2.10.0/test/distributed/checkpoint/test_planner_api.py -v
test_write_item_tensor_storage_size_for_non_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_non_tensor) ... ok
test_write_item_tensor_storage_size_for_shard (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_shard) ... ok
test_write_item_tensor_storage_size_for_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_tensor) ... ok

----------------------------------------------------------------------
Ran 3 tests in 1.036s

OK
[root@734774725e04 workspace]# 

V2.11.0

[root@734774725e04 workspace]# python /workspace/ascend-pytorch-v2.11.0/test/distributed/checkpoint/test_planner_api.py -v
test_write_item_tensor_storage_size_for_non_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_non_tensor) ... ok
test_write_item_tensor_storage_size_for_shard (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_shard) ... ok
test_write_item_tensor_storage_size_for_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_tensor) ... ok

----------------------------------------------------------------------
Ran 3 tests in 1.021s

OK
[root@734774725e04 workspace]# 

V2.12.0

[root@734774725e04 workspace]# python /workspace/ascend-pytorch-v2.12.0/test/distributed/checkpoint/test_planner_api.py -v

test_write_item_tensor_storage_size_for_non_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_non_tensor) ... ok
test_write_item_tensor_storage_size_for_shard (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_shard) ... ok
test_write_item_tensor_storage_size_for_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_tensor) ... ok

----------------------------------------------------------------------
Ran 3 tests in 1.010s

OK
[root@734774725e04 workspace]# 

master

[root@734774725e04 workspace]# python /workspace/ascend-pytorch-master/test/distributed/checkpoint/test_planner_api.py -v

test_write_item_tensor_storage_size_for_non_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_non_tensor) ... ok
test_write_item_tensor_storage_size_for_shard (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_shard) ... ok
test_write_item_tensor_storage_size_for_tensor (__main__.TestPlannerAPI.test_write_item_tensor_storage_size_for_tensor) ... ok

----------------------------------------------------------------------
Ran 3 tests in 1.026s

OK
[root@734774725e04 workspace]# 

【CheckList】

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

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