[RFC]: [NN] DataParallel/DataLoader 行为补丁/[Tensor] .type() 重写补丁削减方案-pytorch-AtomGit

已关闭

[RFC]: [NN] DataParallel/DataLoader 行为补丁/[Tensor] .type() 重写补丁削减方案 #3686

wanglijun55创建于 8月5日关闭于 29 天前

8月5日创建

状态（Status）： Reviewing
作者（Authors）： 王立君
创建日期（Created）： 2026-07-21
更新日期（Updated）： 2026-07-21
相关 Issue/PR： N/A

1. 概述

1.1 简介

本评审报告针对 pytorch_npu 项目中两个 monkey-patch 补丁的处理策略变更进行评估。原计划将这些补丁贡献至 PyTorch 上游社区，经过深入分析后发现补丁函数与原函数之间的差异已可消除，因此建议将策略从【贡献上游社区】调整为【自行消除】。

涉及的两个补丁：

[NN] DataParallel/DataLoader 行为补丁：pytorch-npu/torch_npu/utils/_module.py
[Tensor] .type() 重写补丁：pytorch-npu/torch_npu/utils/tensor_methods.py

1.2 动机

当前痛点：monkey-patch 机制会增加代码维护复杂度，可能与上游 PyTorch 版本更新产生冲突
必要性：若差异可消除，自行移除补丁可简化代码架构，降低维护成本，提升与上游的兼容性
用户价值：提升 pytorch_npu 与原生 PyTorch 的兼容性，减少潜在的版本冲突风险

1.3 目标

目标：确认两个补丁的差异可消除，制定自行移除方案
非目标：本评审不涉及需要修改 PyTorch 上游代码的场景

2. 用例分析

2.1 [NN] DataParallel/DataLoader 行为补丁

文件位置：pytorch-npu/torch_npu/utils/_module.py

关键代码：

def _apply_module_patch():
    torch.nn.Module.npu = npu
    torch.nn.Module.to = to
    torch.nn.Module.cast_weight = cast_weight
    torch.nn.modules.rnn.LSTM.forward = _lstm_forward
    torch.nn.parallel.DataParallel.parallel_apply = npu_parallel_apply # 本次分析部分
    torch.nn.parallel.data_parallel = npu_data_parallel # 本次分析部分
    torch.utils.data.dataloader._MultiProcessingDataLoaderIter.__init__ = (
        _mpdl_iter_init
    )

补丁描述：

以 monkey-patch 形式重写了 torch.nn.parallel.DataParallel.parallel_apply 和 torch.nn.parallel.data_parallel 函数
原计划：修改 PyTorch 上游代码以适配 NPU 行为
现经分析：原函数和patch的差异部分已消除，部分可自行消除

已消除部分：https://github.com/pytorch/pytorch/pull/162573 已经将原函数的torch.cuda都改为torch.accelrator
可自行消除amp差异：torch_npu.npu.amp.autocast_mode和torch.amp.autocast在A5上验证输出一致
可自行消除set device差异：pta多一行setdevice的原因，是24年时cann没有默认device功能，新起一个线程如果直接调用get_device的话，会报错null context；目前已支持默认device功能

2.2 [Tensor] .type() 重写补丁

文件位置：pytorch-npu/torch_npu/utils/tensor_methods.py

关键代码：

def _add_tensor_methods():
    torch.Tensor.type_raw = torch.Tensor.type
    torch.Tensor.type = _npu_type

补丁描述：

通过 _npu_type 函数重写 PyTorch 原生 _type 函数，差异在于类型转换的方法
原计划：修改 PyTorch 上游代码
现经分析：原函数和patch的差异部分可自行消除

可自行消除amp差异：dtype(self.size()).copy_(self, non_blocking)和self.to(dtype=dtype.dtype, device='npu', non_blocking=non_blocking)在A5上验证输出一致

3. 方案设计

3.1 总体方案

根据本提案的场景用例及功能特点，阐述整体设计思路、技术方案、核心逻辑，可包括选择什么软/硬件平台、操作系统、编程模型、使用什么算法，系统架构如何布局，UI如何呈现等。并给出相应的限制和约束说明，如在什么场景或前置条件下，可以达成xx性能指标。
根据实现方案的复杂度，可选择采用自然语言，并结合架构图、时序图、活动图或状态机（算法）等适合的方式来辅助设计。

3.2 技术选型

列出考虑过但放弃的其他方案，给出优劣对比，说明不选择的理由。

3.3 功能与性能设计

结合场景用例分析的结果，对本提案的相关功能及性能指标影响展开设计，如功能的实现方案、运行的核心流程（文字描述或流程图）、数据模型定义或变更（若涉及）、影响范围等。

3.4 安全隐私与DFX设计

结合场景用例，对本提案所涉及的安全隐私及DFX（兼容性、可维护性、可测试性、可靠性...）等属性影响进行设计。

3.5 编程与调用设计

若本提案相关特性/功能组件/模块等支持被开发者集成调用（二次开发），则需要提供便捷易用的编程与调用能力。要站在开发者如何进行编程开发、接口调用及系统集成的使用方式上，给出相应的编程模型定义和设计，包括各要素的可获取方式和途径。

3.5.1 编程模型基本设计

开发环境设计：明确好开发者使用的软/硬件环境、开发&调试工具链、编程框架、要提供的加速库或算子等。
开发约束：开发者使用过程中的约束和限制说明，如硬件平台、编程语言限制等。
可验收设计：提供相应功能、性能指标等的验收环境、标准或用例设计，保证最终的实现可达成既定目标。

3.5.2 接口定义与设计

给出相关组件/模块被集成调用的API定义或变更、对接上下游主流生态技术栈的适配方案、提供功能被使用或集成的参考代码或方法等。

3.5.2.1 xxx（API1 Name）

接口描述：xxx
接口原型：xxx
输入/输出参数：

参数名称	输入/输出	类型	描述	取值范围

返回参数：

参数名称	类型	描述	取值范围

异常处理：xxx
约束说明：xxx
变更说明：xxx
调用参考代码：xxx

3.5.2.2 xxx（API2 Name）

...

3.5.3 编程手册设计

为了帮助开发者能快速上手开发，要设计好本提案相关特性/功能的《编程手册》，要包含哪些内容和章节，单独输出还是共用，在已有的手册中更新还是输出等。确保最后输出的《编程手册》中有相关变更内容。

4. 测试设计

4.1 单元测试

测试 DataParallel 在 NPU 设备上的基本功能
测试 Tensor.type() 类型转换的正确性

4.2 集成测试

验证分布式训练场景下的 DataParallel 行为
验证复杂网络中 Tensor 类型转换的正确性

4.3 回归测试

运行现有 pytorch_npu 测试套件，确保无功能回退介绍该功能的测试方法以及测试用例设计，包括单元测试（unit test），集成测试（integration test），端到端测试（e2e test）等。

5. 缺点和风险

5.1 潜在风险

风险	影响	应对措施
移除补丁后功能异常	高	充分的回归测试，灰度发布
与特定 PyTorch 版本不兼容	中	明确版本兼容性范围
未覆盖的边界场景	中	补充测试用例