@PyTorch: At #PyTorchCon Europe 2026, @ezyang (@Meta) explains why many developers find tensor parallelism difficult to work with…

At #PyTorchCon Europe 2026, @ezyang (@Meta) explains why many developers find tensor parallelism difficult to work with and how PyTorch is exploring SPMD types to catch mistakes through type checking. Watch the full keynote: https://youtu.be/xvNh5F9t-d4?si=zf4P0RFhNB6tx6zh…

TL;DR: PyTorch 核心维护者 Edward 在 PyTorchCon Europe 2026 上介绍了 torch.compile 的预编译支持（解决分布式训练中 JIT 编译的慢和不一致问题）以及 SPMD 类型系统（帮助开发者编写正确的张量并行代码，确保梯度正确）。

PyTorch 生态现状

截至 2026 年 3 月初，PyTorch 在 GitHub 上的星标数和贡献量持续增长。过去 12 个月收到 15,000 个对核心库的贡献，新增 1,250 位贡献者。PyTorch 在顶级机器学习会议中被 90% 的研究项目和 90% 的开源项目使用。特别值得关注的是 dtensor 和 device mesh 的使用量持续上升，超过一千个开源项目采用了这些新功能，说明社区对分布式训练优化的需求旺盛。

预编译支持（Pre-compilation）

为什么需要预编译？

当前 torch.compile 采用即时编译（JIT）机制——每次运行前，每个节点都独立执行编译过程。这在分布式训练中带来两个问题：

1. 速度慢：每次任务重启都要重新编译，缓存预热耗时较长，且缓存查找时必须在 Dynamo 中完整跟踪一遍。
2. 结果不一致：不同节点可能执行不同路径。例如推荐系统中动态形状常见，一个节点因看到不同形状分布而重新编译，其他节点等待可能造成超时；即使编译器确定，也可能因图形不同而做出不同优化选择，若错误重排通信操作会导致死锁。

预编译方案

核心思路：提前编译一次，然后将编译产物分发给所有节点。运行时不再编译，保证所有节点执行相同代码。

需要解决两个问题：

• 完整模型捕获：要提前编译，必须知道所有需要编译的内容。如果提供完整计算图，事情会简单很多。
• 一次编译，处处可用：代码不能依赖某具体进程的 rank 信息，必须以通用方式编写。

当前进展

Torch Titan 中新增了图形训练器（Graph Trainer），专门针对完整模型捕获工作流。已在 Llama 3 上证明达到比特级等价，性能完全匹配。新接口 pre_compile_main：传入训练参数，将编译产物转储到目录。实际训练时只需要指定产物位置，直接使用，无需重新跟踪模型。

下一步计划：确保适用于 Torch Titan 中的所有模型，同时支持需要不同节点不同行为的情况（如 FSDP 参数非均匀切分，或日志记录仅在一个节点进行）。相关 PR 上周刚合并到 Torch Titan。

SPMD 类型（Single Program Multiple Data 类型系统）

解决什么问题？

传统 PyTorch 中，通信算子（如 allreduce）没有可微的自动求导函数。像 Megatron 这样框架虽有自定义自动求导，但编写起来晦涩难懂——前向不做任何事、反向做 allreduce 的 f 函数就是典型例子。如果忘了调用正确的通信函数，会得到静默的错误梯度。

虽然 dtensor 能保证正确梯度，但它的全局语义让一些高级用户不满意：“我想写本地语义的程序，用普通张量和通信算子，不要假装我不知道并行性的存在。” 另外 dtensor 在某些场景下有时模式开销较高等问题。

SPMD 类型的设计

SPMD 类型描述张量（数据）在节点间的分布方式，但不要求知道如何拼接回去（称为局部 SPMD）。主要包括以下几种：

• varying：张量在不同节点上是不同的。不说明如何拼接，只知道它们不同。
• partial：张量不仅是不同的，而且必须通过规约才能得到真实值。允许线性操作（因为可分配到求和），但不允许非线性操作。
• invariant：前向中所有节点上的值完全相同，反向中的值也完全相同。

价值主张

在显式使用通信算子的 Megatron 风格代码上，只需给输入和权重添加类型注解，并用特殊的、能传播类型的通信算子。类型检查通过的程序保证获得正确的梯度。如果忘记调用某些奇怪的函数，类型检查器会警告，并提示在哪里缺少通信算子。类型检查完全可选，无运行时开销，可以在单元测试中运行。所有现有通信模式都能原样工作。

该想法受 JAX 的 sharding in types（尤其是未文档化的 reduced/unreduced 类型）启发，已适配到 PyTorch 生态。

Source: @PyTorch: At #PyTorchCon Europe 2026, @ezyang (@Meta) explains why many developers find tensor parallelism difficult to work with…