[Usage]: torch.utils.data 相关 API 资料缺失及社区用例 NPU 适配-pytorch-AtomGit

已关闭

[Usage]: torch.utils.data 相关 API 资料缺失及社区用例 NPU 适配 #2202

yulin520创建于 6月1日关闭于 6月8日

yulin520

6月1日创建

在提交新问题之前，请确保您已经在社区中搜索过相关问题，并使用了社区中提供的资源/工具后，仍未找到满意的解决方式。

环境信息

初始验证环境

操作系统：Ubuntu
昇腾硬件信息：910B4
CANN 软件版本：8.5.0
安装的软件版本：torch 2.9.0，torch-npu 2.9.0

初始验证环境信息：

torch version: 2.9.0+cpu
torch_npu version: 2.9.0
current accelerator: npu
npu available: True
npu device count: 1

最终匹配版本复验环境

后续在重新搭建的在线开发环境中，对各目标分支补充完成匹配版本复验。

操作系统：Linux
昇腾硬件信息：910B3
CANN 软件版本：8.5.0
NPU Health：OK

各目标分支的 PyTorch 与 torch_npu 版本如下：

目标分支	PyTorch	torch_npu
`v2.7.1`	`2.7.1+cpu`	`2.7.1.post2`
`v2.9.0`	`2.9.0+cpu`	`2.9.0`
`v2.10.0`	`2.10.0+cpu`	`2.10.0rc2`
`v2.11.0`	`2.11.0+cu130`	`2.11.0rc1`
`v2.12.0`	`2.12.0+cu130`	`2.12.0rc1`

使用场景及问题

本 Issue 用于说明以下 10 个 torch.utils.data API 在 torch-npu 环境下的一致性验证结果：

torch.utils.data.Dataset
torch.utils.data.default_collate
torch.utils.data.default_convert
torch.utils.data.distributed.DistributedSampler
torch.utils.data.get_worker_info
torch.utils.data.graph_settings.apply_sharding
torch.utils.data.graph_settings.get_all_graph_pipes
torch.utils.data.IterableDataset
torch.utils.data.IterDataPipe
torch.utils.data.MapDataPipe

经检索、分析和复验：

上述 10 个 API 均已支持；
PyTorch 官方社区已提供相关测试用例；
上述 API 本体主要属于 Python 层数据集、数据加载、采样、DataPipe 或 DataPipe 图配置逻辑，无需新增 API 对齐代码；
相关社区测试文件中存在 Tensor 构造、比较、数据加载和 DataPipe 组合场景，需要补充 NPU 用例迁移；
已针对各目标版本分别生成 test_upstream patch；
其中 6 个 API 已有资料说明；
其中 4 个 API 原资料缺失，已通过文档 PR 补齐；
已在各版本匹配的 PyTorch 与 torch_npu 环境中完成复验。

资料补齐 PR：

https://gitcode.com/Ascend/pytorch/pull/37312

一、API 功能、官方用例及适配结论

API	功能说明	PyTorch 官方测试用例	API 本体是否需要新增对齐代码	相关用例是否需要 NPU 场景检查或迁移	当前处理结果	资料情况
`torch.utils.data.Dataset`	Map-style 数据集基础类，用于通过索引读取样本	`test/test_dataloader.py` 中 `TensorDataset`、`ConcatDataset`、`StackDataset` 等相关用例	否	是	已完成相关 Tensor 场景迁移与复验	文档已有
`torch.utils.data.default_collate`	将多个样本整理为 batch，支持 Tensor、NumPy、映射和序列等类型	`test/test_dataloader.py` 中 `TestDataLoader.test_default_collate_*` 相关用例	否	是	已检查 NPU Tensor 场景，并保留原始语义用例复验	文档已有
`torch.utils.data.default_convert`	将 NumPy 数据转换为 Tensor，并递归处理映射、序列等结构	`test/test_dataloader.py` 中 `TestDataLoader.test_default_convert` 等相关用例	否	是	已完成相关 NPU 场景检查与复验	文档已有
`torch.utils.data.distributed.DistributedSampler`	在分布式训练场景下对数据索引进行分片采样	`test/test_dataloader.py` 中 DistributedSampler 相关用例	否	是	已完成 NPU 环境复验	文档已有
`torch.utils.data.get_worker_info`	获取 DataLoader worker 的编号、数量和数据集信息	`test/test_dataloader.py` 中 `TestDataLoader.test_get_worker_info`	否	是	已完成相关语义复验	文档已有
`torch.utils.data.graph_settings.apply_sharding`	对 DataPipe 图中的 `sharding_filter` 动态应用分片配置	`test/test_datapipe.py` 中 `TestSharding.test_simple_sharding`、`test_sharding_groups`、`test_sharding_length`、`test_multi_sharding`	否	是	已完成相关 NPU 场景适配与复验	原资料缺失，已通过 PR 补齐
`torch.utils.data.graph_settings.get_all_graph_pipes`	遍历 DataPipe 图并返回图中的全部 DataPipe 节点	`test/test_datapipe.py` 中 `TestGraph.test_simple_traverse`、`test_traverse_forked`、`test_traverse_mapdatapipe`、`test_traverse_circular_datapipe`	否	是	已完成相关 NPU 场景适配与复验	原资料缺失，已通过 PR 补齐
`torch.utils.data.IterableDataset`	可迭代式数据集基础类，适用于流式或按迭代方式产生样本的数据源	`test/test_dataloader.py` 中 `TestDataLoader.test_iterable_style_dataset`、`test_chain_iterable_style_dataset`、`test_iterabledataset_len` 等用例	否	是	已完成相关 NPU Tensor 场景迁移与复验	文档已有
`torch.utils.data.IterDataPipe`	可迭代式 DataPipe 基础类，支持 `map`、`filter`、`batch`、`zip`、`concat`、`fork` 等操作	`test/test_datapipe.py` 中 `TestFunctionalIterDataPipe`、`TestIterDataPipeSingletonConstraint`、`TestIterDataPipeCountSampleYielded` 相关用例	否	是	已完成相关 NPU Tensor 场景迁移与复验	原资料缺失，已通过 PR 补齐
`torch.utils.data.MapDataPipe`	Map-style DataPipe 基础类，通过索引访问数据，支持组合、映射、打乱和批处理等操作	`test/test_datapipe.py` 中 `TestFunctionalMapDataPipe`、`TestGraph.test_traverse_mapdatapipe`	否	是	已完成相关 NPU Tensor 场景迁移与复验	原资料缺失，已通过 PR 补齐

二、测试用例适配说明

上述 10 个 API 的 PyTorch 官方社区测试用例已经存在。

API 本体不需要新增对齐代码，但是相关社区测试文件中存在 Tensor 使用场景。为保证测试在 NPU 环境中执行，本次对测试文件进行 NPU 适配。

涉及的社区测试文件：

test/test_dataloader.py
test/test_datapipe.py

在 patch 中引入：

import torch_npu
from torch_npu.contrib import transfer_to_npu

并对需要迁移的 Tensor 显式迁移至 NPU。

各版本分别提交以下 patch 文件：

test_upstream/test/test_dataloader.py.patch
test_upstream/test/test_datapipe.py.patch

三、DataLoader 多进程场景说明

TestDataLoader.test_multiprocessing_contexts 涉及 DataLoader worker 间共享 NPU Tensor。

该场景需要启用：

PYTORCH_NPU_ALLOC_CONF=expandable_segments:True

因此，本次 patch 增加条件判断：

当启用 expandable_segments:True 时，可以执行对应 NPU IPC 场景；
当未启用该配置时，跳过该用例，避免 DataLoader worker 间共享 NPU Tensor 导致异常。

该条件跳过逻辑不影响其他 DataLoader 和 DataPipe 用例验证。

四、API 适配方案

4.1 API 对齐代码

不需要新增 API 对齐代码。

上述 API 主要属于 Python 层数据集、数据加载、采样、DataPipe 或 DataPipe 图配置逻辑，不涉及新增 NPU 算子实现。

4.2 社区测试用例 patch

需要提交测试适配 patch。

修改文件：

test_upstream/test/test_dataloader.py.patch
test_upstream/test/test_datapipe.py.patch

主要修改内容：

引入 torch_npu；
引入 transfer_to_npu；
将需要迁移的 Tensor 显式迁移至 NPU；
对 TestDataLoader.test_multiprocessing_contexts 增加 NPU IPC 条件跳过逻辑；
保留并验证社区用例原始语义，避免 NPU 迁移破坏原有行为。

五、验证命令

5.1 DataLoader 相关用例

python test_dataloader.py -v \
  TestTensorDataset \
  TestStackDataset \
  TestConcatDataset \
  TestDataLoader.test_distributed_sampler_invalid_rank \
  TestDataLoader.test_duplicating_data_with_drop_last \
  TestDataLoader.test_multiprocessing_contexts

5.2 DataPipe 相关用例

python test_datapipe.py -v \
  TestFunctionalIterDataPipe.test_map_iterdatapipe \
  TestFunctionalIterDataPipe.test_collate_iterdatapipe \
  TestFunctionalMapDataPipe.test_map_mapdatapipe \
  TestSharding.test_simple_sharding \
  TestSharding.test_sharding_groups \
  TestSharding.test_sharding_length \
  TestSharding.test_multi_sharding \
  TestGraph.test_simple_traverse \
  TestGraph.test_traverse_forked \
  TestGraph.test_traverse_mapdatapipe \
  TestGraph.test_traverse_circular_datapipe

5.3 保留原始语义用例

python test_dataloader.py -v \
  TestDataLoader.test_default_collate_dtype \
  TestDataLoader.test_default_collate_mapping_keep_type \
  TestDataLoader.test_default_collate_sequence_keep_type \
  TestDataLoader.test_default_collate_sequence_dont_keep_type \
  TestDataLoader.test_default_collate_bad_numpy_types \
  TestDataLoader.test_default_collate_numpy_memmap \
  TestDataLoader.test_default_collate_bad_sequence_type \
  TestDataLoader.test_default_collate_shared_tensor \
  TestDataLoader.test_get_worker_info \
  TestDataLoader.test_invalid_ctor_args_combinations \
  TestDataLoader.test_iterable_style_dataset \
  TestDataLoader.test_chain_iterable_style_dataset \
  TestDataLoader.test_iterabledataset_len

六、匹配版本复验结果

目标分支	PyTorch	torch_npu	DataLoader 用例	DataPipe 用例	保留原始语义用例	验证结果
`v2.7.1`	`2.7.1+cpu`	`2.7.1.post2`	23	11	13	通过
`v2.9.0`	`2.9.0+cpu`	`2.9.0`	23	11	13	通过
`v2.10.0`	`2.10.0+cpu`	`2.10.0rc2`	23	11	13	通过
`v2.11.0`	`2.11.0+cu130`	`2.11.0rc1`	24	11	13	通过
`v2.12.0`	`2.12.0+cu130`	`2.12.0rc1`	24	11	13	通过

其中，TestDataLoader.test_multiprocessing_contexts 在未启用 expandable_segments:True 时按预期跳过。

七、`v2.12.0` 验证环境补充说明

在使用以下环境搭建 v2.12.0 验证环境时：

torch==2.12.0
torch_npu==2.12.0rc1
CANN==8.5.0

直接导入：

from torch_npu.contrib import transfer_to_npu

会触发：

TypeError: _apply_patches() takes 0 positional arguments but 1 was given

经检查，本次通过 PyPI 安装的 torch_npu==2.12.0rc1 wheel 中：

torch_npu/contrib/transfer_to_npu.py

的接口调用方式与官方 v2.12.0 分支中的实现存在差异。

为保证验证环境与官方 v2.12.0 分支保持一致，本地验证时使用官方 v2.12.0 分支中的同名文件替换 wheel 内文件。完成替换后，重新执行 NPU 冒烟测试和相关用例，结果均通过。

该调整仅用于本地验证环境，不属于本次 PR 提交内容。

八、资料检查

已检查：

docs/zh/native_apis

以下 6 个 API 已有资料说明，无需修改：

torch.utils.data.Dataset
torch.utils.data.default_collate
torch.utils.data.default_convert
torch.utils.data.distributed.DistributedSampler
torch.utils.data.get_worker_info
torch.utils.data.IterableDataset

以下 4 个 API 原资料缺失，已通过文档 PR 补齐：

torch.utils.data.IterDataPipe
torch.utils.data.MapDataPipe
torch.utils.data.graph_settings.apply_sharding
torch.utils.data.graph_settings.get_all_graph_pipes

资料补齐 PR：

https://gitcode.com/Ascend/pytorch/pull/37312

九、相关 PR

9.1 资料 PR

v2.7.1：https://gitcode.com/Ascend/pytorch/pull/37312

9.2 测试 PR

v2.7.1：https://gitcode.com/Ascend/pytorch/pull/37376
v2.9.0：https://gitcode.com/Ascend/pytorch/pull/37377
v2.10.0：https://gitcode.com/Ascend/pytorch/pull/37378
v2.11.0：https://gitcode.com/Ascend/pytorch/pull/37379
v2.12.0：https://gitcode.com/Ascend/pytorch/pull/37380

当前 6 个 PR 均已通过：

ascend-cla/yes
docs-ci-pipeline-success

十、提交前审计结果

已完成提交前冗余文件审计。

测试 PR 中仅包含以下两个目标 patch 文件：

test_upstream/test/test_dataloader.py.patch
test_upstream/test/test_datapipe.py.patch

文档 PR 中仅包含以下 5 个目标文档文件：

docs/zh/native_apis/pytorch_2-7-1/torch-utils-data.md
docs/zh/native_apis/pytorch_2-9-0/torch-utils-data.md
docs/zh/native_apis/pytorch_2-10-0/torch-utils-data.md
docs/zh/native_apis/pytorch_2-11-0/torch-utils-data.md
docs/zh/native_apis/pytorch_2-12-0/torch-utils-data.md

审计结果：