SuperAPI AI 工具集

fake_distributed_backend

部署与运维

查看项目

更新于 10-22 04:51

技能介绍

fake_distributed_backend 是 TorchTitan 中的分布式仿真测试模式,用于在不进行完整真实通信的情况下检查分布式配置与内存分配行为。

用途与两种模式:

  1. fake:将所有分布式轴设为假,适用于配置、形状、归属关系和 rank 本地内存分配分析,不涉及真实传输。
  2. real_pp_fake_spmd:每个流水线 rank 使用一个物理进程;PP 通信为真实,数据、张量、上下文、专家并行轴为假,适用于需要检验流水线通信与缓冲区生命周期的场景。 两种模式均不能替代真实的分布式数值测试或性能测试。完整的面向用户启动约定(含命令示例和环境变量语义)见 docs/debugging.md。

逻辑 Rank 坐标:

  • 将 NGPU 设为逻辑 world size。纯 fake PP 还需设置:export FAKE_PP_RANK=<logical PP coordinate>
  • 对 PP degree P,非 PP 逻辑 world size 为 NGPU / P,全局逻辑 rank 为 FAKE_PP_RANK * (NGPU / P)。
  • 两种模式都表示 SPMD 坐标零。纯 fake 在 PP 关闭时默认 PP 坐标为零。混合模式下,物理 RANK 即为 PP 坐标,且 WORLD_SIZE 必须等于 PP degree;设置 FAKE_PP_RANK 无效。
  • COMM_BACKEND=fake 是 run_train.sh 识别的便捷方式,用于单进程纯 fake 运行。混合模式需用 torchrun 启动,将 NGPU 设为逻辑 world size,并传入 --comm-backend real_pp_fake_spmd;torchrun 会提供物理 RANK、WORLD_SIZE、LOCAL_RANK 和 rendezvous 变量。

内存调试流程:

  1. 使用生产环境的模型、recipe、dtype、并行度、microbatch 数量、FSDP 策略、激活检查点策略和 CUDA-graph 设置。
  2. 选择要调查其 SPMD-zero 归属的 PP 坐标。
  3. 在初始化后、完整优化器预热后、稳态前向/反向期间以及优化器完成后捕获分配器状态。
  4. 在相同逻辑 rank 和捕获点之间比较不同配置。
  5. 按以下顺序逐步升级:纯 fake -> 真实 PP/假 SPMD -> 真实分布式运行。

fake 执行所代表的内容包括:PyTorch 管理的模型参数、优化器状态、已准备的量化权重、激活、梯度、流水线缓冲区和 DistMoE scratch 或激活 arena。它不能忠实代表 NCCL 通信器分配、网络注册、传输 scratch、集合通信延迟或通信重叠。

可使用 torch.cuda.memory_snapshot() 查看分配来源,使用 torch.cuda.max_memory_allocated() / torch.cuda.max_memory_reserved() 获取可比较的总量。fake 运行可以确定归属和生命周期行为,但不能确定分布式数值一致性和吞吐量。

失败解读:

  • 整除或坐标错误属于启动约定失败,应修正逻辑拓扑,而非修改模型配置。
  • 混合模式 world-size 错误表示物理分配未做到每个 PP rank 恰好一个进程。
  • 纯 fake 成功但混合模式失败,通常可定位到 PP 传输、发送/接收归属或通信器初始化问题。
  • 混合模式成功但真实运行失败,通常可定位到 SPMD 集合通信、真实通信内存或规模相关调度问题。
  • 不要将不同 PP 坐标之间的比较作为内存回退的证据,因为它们对模型和激活的归属可能有意不同。