企业官网建设流程全解析

PyTorch torch.foreach 公开 API 完全指南：跨张量列表的融合算子应用与私有 API 迁移

【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch

torch.foreach是 PyTorch 中针对**张量列表（list/tuple of Tensor）**的批量算子接口：它把熟悉的逐张量算子（如torch.add、torch.tanh）"提升"到列表的每个位置上执行，并在输入满足条件时使用水平融合的多张量（multi-tensor）内核加速。本指南以 docs/source/foreach.md 为骨架，结合 torch/foreach/init.py 的源码实现与 test/test_foreach.py 等测试佐证，系统讲解其语义、参数形态组合、加速条件、从私有torch._foreach_*迁移的要点，以及当前完整支持的操作清单。

torch.foreach 是什么：把"循环"写成"一次调用"

torch.foreach的核心理念很简单：把一个作用于单个张量的普通 PyTorch 算子，等价地应用到一组张量（TensorList）的每一个位置上。文档中给出的示例清晰表达了这一语义：

inputs = [torch.ones(2), torch.ones(3)] result = torch.foreach.add(inputs, 2) # Equivalent to tuple(torch.add(tensor, 2) for tensor in inputs) torch.foreach.mul_(inputs, 3) # Mutates each tensor and returns `inputs`.

也就是说，torch.foreach.add(inputs, other)与一个逐位置调用torch.add的 Python 循环语义等价。两者的区别在于运行时路径：当输入满足特定条件时，foreach 算子会在加速器上采用**水平融合的多张量内核（horizontally fused multi-tensor kernel）**执行，从而减少内核启动开销、提升端到端运行时间；否则退化为逐张量的等价实现。

值得注意的是，模块源码在 torch/foreach/init.py 的模块级 docstring 中明确指出：调用该模块中的函数并不保证一定使用单个或融合内核。即"能融合就融合，不能融合就等价回退"是本 API 的默认行为契约。模块同时声明torch.foreach目前是beta API，其签名可能依据用户反馈调整；迁移期间原有的私有torch._foreach_*函数会保留可用。

返回约定

从源码中的文档生成逻辑（_result_doc）可以归纳出统一的返回契约：

• 就地（in-place）版本（函数名以下划线结尾）：就地修改inputs中的每个张量，并返回传入的原始容器对象本身；
• 非就地版本：不改动任何输入参数，返回一个元组，其中每个元素对应inputs中一个输入张量的结果。

例如在 abs_ 的实现中可以看到，作者没有直接返回torch._foreach_abs_(inputs)的返回值，而是执行后再return inputs。源码注释解释了原因：由于 dynamo（torch.compile）可能绕过 Python 绑定直接调用底层算子，显式返回原输入容器可以保证就地语义在编译路径下依然成立。

何时获得加速：水平融合内核的准入条件

文档以 CUDA 为例，给出了融合内核的常见准入门槛（common eligibility requirements）。只有当张量列表满足这些条件时，算子才可能被路由到加速实现：

• 张量为strided、互不重叠（non-overlapping）的稠密张量（dense tensors）；
• 所有张量位于同一设备（same device）；
• 各张量dtype 相互兼容（compatible dtypes）；
• 列表内对应位置的张量sizes 与 strides 一致（matching sizes and strides for corresponding tensors）。

这些约束本质上是为了让"多个独立的小张量"可以被当作"一个平坦的大缓冲区"来统一处理，从而让单个融合内核一次覆盖所有张量。一旦不满足（例如尺寸参差、设备不同），实现会回退到逐张量的语义等价执行。

源码在 docstring 里同样强调了这一"尽力而为"的性质——调用torch.foreach并不承诺一定产生单内核或融合内核，性能取决于输入形态与运行环境，因此在性能敏感场景下应自行验证。

API 覆盖：参数形态的组合空间与约束

foreach API 的本质是"把一个普通张量算子提升到列表上"，由此会产生组合爆炸般的签名空间（combinatorial space of possible signatures）。文档指出，依据算子的不同，一个参数可能存在多种提供方式：

• Tensor 参数：既可以作为被所有列表位置共享的单个Tensor，也可以按元素逐位置提供为TensorList；
• 标量参数：可能以Scalar、ScalarList、被共享的0 维（0-D）Tensor、或被打包的1 维 CPU Tensor（packed 1-D CPU Tensor，每个列表位置一个标量）等形式出现。

于是同一个算子可以衍生出诸如TensorList/TensorList、TensorList/Tensor、TensorList/ScalarList、TensorList/Scalar等形态；参数越多的算子组合越多。公开 foreach API 仅根据实际使用场景支持其中一部分组合。文档说明，如果读者需要某个缺失的组合，可以到 PyTorch 官方提交 issue 反馈。

在已支持的签名中，foreach API 维护以下跨签名约束：

1. TensorList与ScalarList参数必须非空；
2. 对应位置上的张量列表与标量列表长度必须一致（源码文档生成函数_common_doc中的has_aligned_lists=True即为二元与逐点算子注入该长度要求）。

关于隐式标量转换与 Tensor 类型面的边界

文档特别澄清了**类型面（typed surface）**的边界问题：

• 只有当签名显式列出Tensor时，该签名才把对应参数纳入"支持 Tensor 类型"的类型面；
• 一个不要求梯度的 0 维 Tensor，有时可以通过**隐式标量转换（implicit scalar conversion）**被某个Scalar参数接受；
• 但把一个加速器上的 Tensor 做隐式标量转换意味着在主机端（host）读取它的值，这可能非常昂贵；在 CUDA 上这会同步 eager 执行，并且在CUDA graph 捕获（capture）期间不受支持。

因此文档明确建议：不要依赖隐式转换把 0 维 Tensor 当作"Tensor 重载"来使用。以 add 的实现 为例，源码用了一个内部哨兵_OMITTED_ALPHA来区分"未传alpha"与"显式传入alpha"：当other是共享 0 维张量时，必须显式传alpha（即使值就是1），以便路由到 Tensor 重载；省略alpha则可能把张量转换成主机标量（host scalar），从而触发上述同步开销。

已知限制

• 公开 foreach API不支持out=变体；
• 由于可能同时有多个中间结果存活，其内存占用可能高于逐个调用非 foreach 原始 API 的循环。

从私有 API 迁移到公开 beta API

许多读者早已熟悉 foreach 的私有拼写（private spelling）。文档以torch.add为例给出对照：

torch._foreach_add(inputs, other) # Private spelling torch.foreach.add(inputs, other) # Public beta spelling

迁移有两大保证：

1. 私有拼写保持可用：torch._foreach_*继续以不变签名存在，用于向后兼容，已有代码无需立即改动；
2. 公开函数调用的是同一批 ATen 算子：公开函数与私有算子共享底层实现，只是改善了 API 一致性（consistency）。

文档归纳了公开 API 在一致性上的两处改进：

1. 所有必需操作数均为位置限定（positional-only）参数，所有可选参数均为关键字限定（keyword-only）参数。这在源码的overload声明与实现签名中清晰可见，例如 add 的定义 中inputs、other之后带/，而可选的alpha只允许以关键字alpha=...传入；
2. 参数命名与对应的普通算子保持一致（parameter names align with the corresponding ordinary operation）。

同时，文档给出了命名约定：foreach API 所作用的主张量列表参数统一命名为inputs，其余参数保留普通算子的逻辑名称——即使当前支持的形态要求它以列表形式给出。这样当操作数形态未来演进（例如从标量演变为标量列表）时，签名依然保持可描述性。

torch.foreach模块下每个公开函数其实都是薄封装：源码中用_make_foreach_api装饰器（见 torch/foreach/init.py#L153-L214）完成三件事：

• 将内部实现中"省略即默认 1"的alpha哨兵替换为对外可见的默认值1；
• 为函数注入根据普通算子自动生成的 docstring（返回值/变更语义、非空与等长约束说明）；
• 支持列表/元组中张量类对象上的__torch_function__分发：当参数里嵌套存在支持__torch_function__的对象时，先展平参数并检查has_torch_function，再调用handle_torch_function路由到相应重载。

每个公开函数主体最终都转发到对应的私有算子，例如torch.foreach.abs(inputs)内部就是torch._foreach_abs(inputs)。据此推断，私有torch._foreach_*家族构成了真正的 ATen 落地实现，而torch.foreach是它之上的"签名规范化 + 文档化"公开层。

支持的操作清单（API 覆盖全集）

下表汇总了当前公开 foreach API 支持的算子全集，与 torch/foreach/init.py 中的__all__导出完全对应。除了文档按四类列出的清单，结合源码 docstring 还可以补充若干签名细节与不支持项。

一元操作（Unary Operations）

对inputs中每个张量施加对应一元函数，包括带_后缀的就地变体与zero_：

abs、abs_、acos、acos_、asin、asin_、atan、atan_、ceil、ceil_、cos、cos_、cosh、cosh_、erf、erf_、erfc、erfc_、exp、exp_、expm1、expm1_、floor、floor_、frac、frac_、lgamma、lgamma_、log、log_、log10、log10_、log1p、log1p_、log2、log2_、neg、neg_、reciprocal、reciprocal_、round、round_、rsqrt、rsqrt_、sigmoid、sigmoid_、sign、sign_、sin、sin_、sinh、sinh_、sqrt、sqrt_、tan、tan_、tanh、tanh_、trunc、trunc_、zero_

特别说明（来自源码 docstring）：

• clone也归入此类（位于 源码一元区附近，支持可选的memory_format关键字参数）；
• round/round_不支持decimals参数（The ``decimals`` argument is not supported.）。

二元操作（Binary Operations）

对inputs与逐位置或共享的other施加二元函数：

add、add_、sub、sub_、mul、mul_、div、div_、clamp_min、clamp_min_、clamp_max、clamp_max_、minimum、minimum_、maximum、maximum_、pow、pow_、copy_

补充签名细节：

• add/add_：other可为共享标量、标量列表、张量列表或共享 0 维标量张量；alpha仅当other为张量列表或共享 0 维标量张量时受支持，默认1；
• sub/sub_：alpha仅当other为张量列表时受支持；
• div/div_：不支持rounding_mode参数；
• clamp_min/clamp_max：仅指定单一上界或下界（普通torch.clamp的min/max二选一形态）；
• copy_：签名形如copy_(inputs, src, *, non_blocking=False)，把src列表中对应位置张量复制到inputs对应张量；
• pow的形态较为特殊，除pow(inputs, exponent_*)外还支持pow(input, exponent_list)这种标量底数 + 张量列表指数的反向形态。

逐点操作（Pointwise Operations）

对来自三张输入列表的对应位置张量组合计算：

addcmul、addcmul_、addcdiv、addcdiv_、lerp、lerp_

以 addcmul 实现 所在的逐点系列为例，签名统一为(inputs, tensor1, tensor2, *, value=...)；其中第三个标量参数value可以是一个共享标量、一个标量列表/元组，或一个打包的 1 维 CPU Tensor（每个列表位置一个标量），默认值为1；lerp一族则形如lerp(inputs, end, weight, /)。

归约与矩阵操作（Reductions and Matrix Operations）

max、norm、mm

• max(inputs)：对每个输入张量施加归约语义的torch.max（等价于逐位置应用）；
• norm(inputs, *, ord=2, dtype=None)：逐位置计算范数，默认二阶范数，可选 dtype；
• mm(inputs, mat2)：对inputs与mat2中对应位置的矩阵执行矩阵乘法（matmul 的批量应用）。

在框架内部的使用与测试验证

torch.foreach与其底层torch._foreach_*算子并非只服务于终端用户，它们已是 PyTorch 框架内部多个高频路径的组成部分（下述使用场景可以佐证 foreach API 的实战价值）：

• 优化器实现：从源码结构看，torch/optim下的 adam.py、adamw.py、sgd.py、adagrad.py、rmsprop.py、_adafactor.py 等多个优化器文件均引用了 foreach 相关算子。foreach 优化器把"参数 + 动量/二阶矩 + 梯度"打包成列表统一更新，从而在参数众多时显著减少内核启动次数；
• 分布式梯度同步：在 torch/nn/parallel/distributed.py 中可以看到将梯度拷贝"flushed as a single_foreach_copy_plus one flat"的实现模式，即 DDP 的参数梯度在转成扁平缓冲区后，借助 foreach 拷贝算子把多次小拷贝合并为一次批量操作；
• 编译路径：torch.foreach的就地实现特意"返回原输入容器"以保证 dynamo/torch.compile绕过 Python 绑定时就地不变量依然成立（见前述abs_的注释），说明该 API 与编译路径有专门设计的兼容行为；
• 测试覆盖：功能正确性由 test/test_foreach.py 覆盖（包括与逐张量循环等价性的对比测试等），而编译/融合相关路径则由 test/inductor/test_foreach.py 覆盖。

小结与使用建议

• 语义等价：torch.foreach.*永远只是把普通算子提升到列表，返回值/就地语义与逐张量循环一致，可放心作为优化手段替换手写循环；
• 加速是有条件的：是否走水平融合多张量内核取决于设备、dtype、strides、sizes 等条件，命中不了就自动回退，因此不应假设"调用即融合"；
• 签名注意点：主列表参数一律叫inputs，必需参数位置限定、可选参数关键字限定；共享 0 维张量走 Tensor 形态时需显式传alpha，谨防隐式标量转换带来的主机同步与 CUDA graph 不兼容问题；
• 迁移友好：私有torch._foreach_*签名保持不变、随时可继续使用，公开torch.foreach是推荐的长期方向（beta），二者调用同一批 ATen 算子；
• 边界意识：不支持out=；round不支持decimals、div不支持rounding_mode；逐位置与共享参数形态的组合仅支持常用子集，遇到缺失组合应向项目提交 issue。

无论你是在手写优化循环、改造自定义优化器，还是在阅读 DDP 与优化器的源码，把torch.foreach的这套"列表提升 + 条件融合"心智模型建立起来，都能帮助你写出更贴近框架原生加速路径的高效代码。

【免费下载链接】pytorchTensors and Dynamic neural networks in Python with strong GPU acceleration项目地址: https://gitcode.com/GitHub_Trending/py/pytorch