viable/strict/1791021071: [fsdp2][mp] per-param mixed_precision policy (#196626)
PyTorch 发布 viable/strict/1791021071 版本说明，介绍 FSDP2 的 per-param mixed_precision policy（PR #196626）。All-gather 阶段：同构路径不变；当参数需要不同 compute dtype 时，FSDP 按 source/target dtype 分组 staging 拷贝，复用现有 uint8 通信缓冲，每个 fully_shard 仍只发一次 AG。Reduce-scatter 阶段：不同 compute_dtype 会产生不同 dtype 的梯度，而 reduce-scatter 输入缓冲必须单一 dtype；当所有参数共享 reduce_dtype 时保留原有单 collective 行为，为此扩展 _chunk_cat.out CUDA 路径以读取异构 BF16/FP32 输入并在分块与 padding 时直接转换进 reduction buffer。
Development
- First Reportviable/strict/1791021071: [fsdp2][mp] per-param mixed_precision policy (#196626)PyTorch Core
- Current Assessment判断：混合精度策略从全局开关走向 per-param 粒度，反映大模型训练中不同层对数值稳定性的需求分化（如 KDA 等场景）。这属于框架层能力演进，短期不改变硬件供给格局，但会降低训练侧为精度妥协而做的模型结构折衷。Agent Pulse · analysis
PyTorch 在 viable/strict/1791021071 发布说明中给出 FSDP2 的 per-param mixed_precision policy（PR #196626），目标是让不同参数使用不同计算精度。实现上，all-gather 保留同构路径，异构时按 source/target dtype 分组 staging 拷贝并复用 uint8 通信缓冲，每个 fully_shard 仍只有一次 AG；benchmark 显示 copy-in 有轻微性能差异，原因被归为 foreach_copy 固定每线程四元素，uint8 相比 bf16 需要约两倍的元素级调度与指令工作来搬运相同字节数。reduce-scatter 方面，异构 compute_dtype 导致梯度 dtype 不同，而输入缓冲必须单一 dtype；当参数共享 reduce_dtype 时保留单 collective 行为，并扩展 _chunk_cat.out CUDA 路径以读取异构 BF16/FP32 输入、在分块与 padding 时直接转换进 reduction buffer。
这是通信与精度解耦的工程改动：AG 侧用 dtype 分组加 uint8 缓冲避免增加 collective 次数，RS 侧用融合 cast 避免为异构梯度引入额外 kernel 或多次 collective。可验证的下一信号是 PR 中 benchmark 的端到端吞吐与显存数据，以及 _chunk_cat.out 异构输入路径在非 BF16/FP32 组合下的覆盖情况。
判断：混合精度策略从全局开关走向 per-param 粒度，反映大模型训练中不同层对数值稳定性的需求分化（如 KDA 等场景）。这属于框架层能力演进，短期不改变硬件供给格局，但会降低训练侧为精度妥协而做的模型结构折衷。
对训练基础设施团队而言，per-param 精度策略可能减少为稳定性而牺牲吞吐的取舍，但收益取决于实际 benchmark 而非设计描述。评估时应要求端到端吞吐、显存与收敛对比数据，而非仅看 collective 次数不变这一结构性论据。
若该策略在后续版本默认可用，可观察是否出现按参数组配置精度的标准接口，以及上游训练框架是否跟进暴露同类配置；反之若长期停留在 PR 阶段，说明收益不足以抵消实现复杂度。