PyTorch 2.15：通信默认实现切换，平台性能改进

1. GPU 间通信切换至新实现

负责 NVIDIA GPU 间数据交换的通信层 NCCL 后端改用了新实现。指定 nccl 即会选择新实现，创建通信组时也会预先建立 GPU 间的通信通道。

还新增了用于排查通信停滞原因的记录功能。当某个训练进程超时时，所有进程都会将当时的任务状态写入磁盘，以便查看停滞时正在执行的操作。

以下是官方发行说明中的初始化示例。除默认实现外，还展示了使用独立对等通信通道的方式，以及选择旧实现的方式。

dist.init_process_group("nccl", device_id=torch.device("cuda", local_rank))  # new implementation
dist.init_process_group("nccl-lazy", device_id=torch.device("cuda", local_rank))  # independent P2P communicators
dist.init_process_group("nccl-legacy")  # legacy implementation, removed in the next release

2. 调整训练精度和内存使用时机

将模型状态分散存储在多块 GPU 上的 FSDP2 细化了精度控制。路由器门控或归一化权重等敏感参数可以保持 32 位，而模型其余部分则以较低精度训练。

当 reduce_dtype=None 时，梯度将按照各参数的梯度精度设置进行累积和归约。官方发行说明中，采用 32 位参数、bfloat16 模型计算的示例，其梯度处理也从 16 位变为 32 位，使梯度分片求和通信的数据量增至 2 倍。这一变化以更多通信量换取较高的梯度精度；如需维持原有行为，可显式指定归约精度。

将模型分为多个阶段执行的流水线并行，会在接收前才分配缓冲区，并释放传输完毕的输出。它不再提前堆积缓冲区，而是按需分配，以降低峰值内存占用。此外，还新增了让梯度归约与后续计算重叠执行的调节功能。

分布式通信新增了单边传输 API，无需通信组或集合通信，即可通过存储服务与对端建立连接。首个后端为 NIXL，可将传输加入 GPU 执行流，并纳入用于记录和重放一系列操作的 CUDA 图。

3. 加速 Apple 芯片上的文本生成运算

Apple 芯片的 GPU 执行层 MPS 将更多运算迁移至原生 Metal 内核。此前经过中间执行层的运算改由内核直接处理，以减少调用开销。

据官方公布的数据，将一行中 128,256 个词汇的分数转换为概率的 softmax，耗时降至原来的 1/8。放大尺寸的上采样运算，其反向传播速度最高提升至 21.7 倍。官方还给出了计算整个张量大小的范数运算提速超过 100 倍的结果。

据官方公布的数据，文本生成解码阶段的 F.linear 速度按几何平均值，即基于乘积计算的平均值，提升至 1.23 倍，小型投影运算最高提升至 5 倍。在混合专家模型等选择部分子网络运行的模型解码中，M3 Ultra 的新分组矩阵运算路径使吞吐量提升至 1.32~1.34 倍。

在另一项线性代数结果中，据官方公布的数据，512×512 矩阵的 QR 分解耗时从约 7 秒降至 1.5 毫秒。QR 分解是将矩阵分解为正交矩阵和三角矩阵的计算。

4. 配套选择安装包与运行环境

默认 wheel 安装包改为 CUDA 13.2，同时提供 CUDA 13.4 wheel。两者均要求 580 或更高版本的 NVIDIA 驱动，以及图灵架构或更新的 GPU。对于需要保留 CUDA 12.6 或 13.0 wheel 的环境，官方发行说明建议使用旧版本。

支持的 Python 版本范围为 3.11~3.15。从原有的 3.10 环境升级时，也需同步升级 Python。Apple 芯片的 MPS 加速要求 macOS 15 或更高版本。

AMD 平台新增了 ROCm 10.0 wheel。ROCm wheel 改为基于 TheRock 包，不再直接捆绑库，而是将其作为依赖项安装。

5. 准备迁移现有训练配置与存储数据

新通信后端会为未使用的组分配 GPU 内存，且每个组绑定单个设备，因此单进程使用多块 GPU 的配置需要旧实现；该实现及通过环境变量回退的路径将在下一版本移除。

2.15 保存的分布式检查点无法由 2.14 及更早版本读取，因此若计划回退至旧版本，还需考虑存储数据的兼容性。

PyTorch 2.15 的采用价值在于各项运算提速和训练控制能力扩展，而迁移规划还应涵盖通信配置与存储数据。