关联 Issue

• https://gitcode.com/Ascend/pytorch/issues/2202

修改内容

本 PR 对 torch.utils.data 相关社区用例进行 NPU 适配。

修改文件：

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

主要修改：

1. 对相关测试用例中的 Tensor 显式迁移至 NPU；

2. 在测试文件中引入：

   import torch_npu
   from torch_npu.contrib import transfer_to_npu
   

3. 对 TestDataLoader.test_multiprocessing_contexts 增加 NPU IPC 条件判断。当未启用 expandable_segments:True 时跳过该用例，避免 DataLoader worker 间 NPU Tensor 共享场景异常。

验证环境

• 操作系统：Linux
• 昇腾硬件：Ascend 910B3
• CANN：8.5.0
• 目标分支：v2.9.0
• PyTorch：2.9.0+cpu
• torch_npu：2.9.0

验证命令

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

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

验证结果

• DataLoader 相关用例：23个测试，OK；其中 test_multiprocessing_contexts 在未启用 expandable_segments:True 时按预期跳过；
• DataPipe 相关用例：11 个测试，OK；
• 保留原始语义测试：13 个测试，OK。

保留原始语义用例

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

Patch 生成与检查说明

本 PR 中的 patch 基于对应版本的 PyTorch 社区源码生成。已有 patch 时，先应用原 patch，再修改社区源码，最后通过 git diff 重新生成 patch 文件。未手动编辑 patch 文件。

已执行：

git apply --check
git diff --check

检查结果通过。

补充对应版本的用例执行结果和耗时截图

目标分支：v2.9.0

已在匹配版本的PyTorch和torch_npu运行时环境中完成以下验证：

DataLoader相关用例
DataPipe相关用例
原始语义回归用例

以上用例均已执行通过。

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

详细测试数量和执行耗时见附件截图。