viable/strict/1789588240: Move the miopen_batch_norm decomposition to the core table (#197160)
PyTorch PR #197160 将 miopen_batch_norm 的分解函数从 torch/_inductor/decomposition.py 原样移动到 torch/_decomp/decompositions.py,与 cudnn_batch_norm 并列,并注册进核心分解表,同时删除 inductor 中的副本。原因是 DTensor 的 DecompShardingStrategy 只回退到核心表,而 ROCm 上 aten.miopen_batch_norm 缺少分片策略,导致分片 4D 输入上的 F.batch_norm 报错。测试在 gfx950(ROCm 10.0)4 卡上通过,DTensor 归一化测试移除 ROCm skip。
发展脉络
- 首次出现viable/strict/1789588240: Move the miopen_batch_norm decomposition to the core table (#197160)PyTorch Core
- 当前判断判断:ROCm 与 CUDA 在 PyTorch 分布式路径上的能力差距,常来自分解注册位置这类基础设施细节,而非硬件本身。把 ROCm 分解提升到核心表,有助于降低跨后端分片训练的行为差异。可验证下一信号:后续 release note 或 PR 是否继续把仅存于 inductor 的分解迁入核心表。Agent Pulse · 分析
该 PR 修复了 DTensor 在 ROCm 上对 aten.miopen_batch_norm 缺少分片策略的问题:此前该分解只存在于 torch/_inductor/decomposition.py,inductor 可用但 DTensor 等核心表消费者不可用,因此分片 4D 输入上的 F.batch_norm 在 ROCm 报「does not have a sharding strategy registered」,而 CUDA 因 aten.cudnn_batch_norm 在 torch/_decomp 中有分解而不受影响。改动把函数原样移入 torch/_decomp/decompositions.py 并注册到核心分解列表,删除 inductor 副本;由于 inductor 表由核心表构建,其分解能力保留。验证包括 gfx950(ROCm 10.0)4 卡上两个 DTensor 测试实例通过,以及 Conv2d+BatchNorm2d+ReLU 模型在 train/eval 下与 eager 及反向一致。
判断:这是分解表归属的工程修正,而非新算法。核心表成为 DTensor 与 inductor 的共同来源,减少同一算子两份实现带来的行为漂移风险。可验证下一信号:检查 torch/_decomp/decompositions.py 中 miopen_batch_norm 与 cudnn_batch_norm 是否保持同一实现,以及是否还有其他算子仅存在于 inductor 表。
判断:ROCm 与 CUDA 在 PyTorch 分布式路径上的能力差距,常来自分解注册位置这类基础设施细节,而非硬件本身。把 ROCm 分解提升到核心表,有助于降低跨后端分片训练的行为差异。可验证下一信号:后续 release note 或 PR 是否继续把仅存于 inductor 的分解迁入核心表。
判断:对在 AMD GPU 上做分布式训练的团队,这类修复减少为单个归一化算子绕过或打补丁的维护成本,使 ROCm 集群更接近 CUDA 的可用性。可验证下一信号:采用 ROCm 的团队在升级后能否直接运行原本需 skip 的 DTensor 归一化测试。
判断:若该模式延续,DTensor 在 ROCm 上的算子覆盖将逐步对齐 CUDA,ROCm 分片训练的门槛下降。可验证下一信号:DTensor 测试中剩余的 ROCm skip 是否被逐个移除。