viable/strict/1786787633: [CP] Fix FlexCP + load balancing enabled + aot_fx_trace issue (#188766)
PyTorch 发布修复，解决 TorchTitan graph_trainer 与 FlexCP 及负载均衡同时启用时出现的两个错误：未使用 squeeze(0) 导致 FunctionalTensor 使用错误，以及 BlockMask is_leaf 未修复导致 seq_dims 的 pytree 结构错误。
发展脉络
- 首次出现viable/strict/1786787633: [CP] Fix FlexCP + load balancing enabled + aot_fx_trace issue (#188766)PyTorch Core
- 行业反馈trunk/53126a784b3327da0a681a42d1613420967af2a2: [CP] Fix FlexCP + load balancing enabled + aot_fx_trace issue (#188766)PyTorch Core
- 当前判断PyTorch 持续优化大规模分布式训练框架，特别是针对长上下文和并行策略的组合。此修复表明，随着模型规模增大，训练框架的稳定性成为关键，需要不断修复边缘情况。Agent Pulse · 分析
PyTorch 在 2026 年 8 月 15 日发布了针对 TorchTitan graph_trainer 与 FlexCP 及负载均衡同时启用时的修复。该修复解决了两个具体错误：一是未使用 squeeze(0) 时出现 RuntimeError，提示尝试在 FunctionalTensor 上使用自身，而应使用对应的 FunctionalTensorMode；二是 BlockMask 的 is_leaf 属性未修复时出现 ValueError，提示 seq_dims 必须与 buffers 具有相同的 pytree 结构。此修复已合入 PyTorch 的 viable/strict 和 trunk 分支，由 SherlockNoMad 批准。
该修复针对 PyTorch 中 FlexCP（Flexible Context Parallelism）与负载均衡结合使用时的兼容性问题。错误涉及 FunctionalTensor 的使用和 BlockMask 的 pytree 结构，表明在自动微分和并行策略下，张量元数据（如 is_leaf）和维度处理需要特殊处理。修复可能涉及在特定条件下对张量进行 squeeze 操作，以及正确设置 BlockMask 的 is_leaf 属性，以确保 pytree 结构一致。
PyTorch 持续优化大规模分布式训练框架，特别是针对长上下文和并行策略的组合。此修复表明，随着模型规模增大，训练框架的稳定性成为关键，需要不断修复边缘情况。
此修复提升了 PyTorch 在分布式训练场景下的可靠性，降低了使用 TorchTitan 进行大规模训练时遇到错误的风险，对依赖 PyTorch 进行模型训练的企业具有积极影响。
未来，PyTorch 可能会继续改进 FlexCP 与负载均衡的集成，并可能推出更多针对分布式训练的稳定性修复。