AGAnchorGPU

工作负载指南

规划可复现的 PyTorch 运行

测量内存、选择并行策略、测试检查点恢复,并规划固定期限的 PyTorch 训练运行。

10 分钟阅读 · 更新于

在扩展前先验证一个完整步骤

锁定模型和数据集修订版本、依赖项版本、训练配置、随机种子和容器摘要。在单个加速器上运行模型加载、完整训练步骤、评估、检查点保存和检查点恢复。

使用与长期运行相同的序列长度和微批次设置。一个省略优化器状态或使用较短序列的小型试点可能会掩盖您试图测量的内存问题。

测量一个有代表性的优化步骤

在您现有的训练程序中,用一个峰值内存测量来包围一个代表性步骤。该示例假设模型、批次和优化器已在 GPU 上初始化。请包含优化器步骤,因为某些状态是延迟创建的。

torch.cuda.reset_peak_memory_stats()
optimizer.zero_grad(set_to_none=True)
loss = model(**batch).loss
loss.backward()
optimizer.step()
torch.cuda.synchronize()
peak_gib = torch.cuda.max_memory_allocated() / 2**30
print(f"Peak tensor memory: {peak_gib:.2f} GiB")

这报告的是峰值张量分配,而不是完整的设备占用。还应检查保留内存、运行时开销以及系统的设备内存报告。

出于正确的原因选择 DDP 或分片

DistributedDataParallel 在每个进程上保留模型副本并同步梯度。当完整训练状态可放入每个 GPU 且您希望并行处理更多数据时,它非常有用;它不会将 VRAM 合并为一个连续池。

完全分片数据并行可以分布参数、梯度和优化器状态。当内存压力证明分片合理时使用它,并考虑额外的通信和分布式检查点复杂性。

混合精度和激活检查点是独立的选择。请验证较低精度的数值行为;当检查点节省内存时,衡量重新计算激活的计算成本。

启动单节点试点

训练脚本必须初始化分布式执行,并使用本地排名信息为每块 GPU 分配一个进程。进程数应与可见 GPU 数量匹配。以下是一个支持分布式脚本的启动示例,并非完整的训练实现。

torchrun --standalone --nproc-per-node=4 \
  train.py --config configs/pilot.yaml

仅在已配置的节点上运行此操作。本地 AnchorGPU 演示不会执行训练,也不会分配远程 GPU。

根据试运行结果选择硬件

将 A100 用作价格较低的 80 GB CUDA 基准,如果工作负载有实用的 Hopper 专用路径,则测试 H100。H200 提供更大的 CUDA 内存配置。在本目录中,MI300X 提供更高的单加速器容量,前提是整个应用已在 ROCm 上完成验证。

对于 AMD,请将自定义扩展、内核、软件包版本以及经 ROCm 测试的 PyTorch 镜像作为一个完整集合进行检查。框架级别的兼容性并不能验证每一个可选操作。

将 H200 与 MI300X 进行比较

测试恢复,而不只是保存

在相关情况下记录模型、优化器、调度器、训练步骤、配置和混合精度缩放器状态。保存检查点,终止进程,然后恢复并继续,再承诺进行长时间运行。

使用适合分布式状态的检查点方法。测试已保存的工件能否独立于原始进程读取,并确保评估输出保持在预期容差范围内。

为整个流程制定预算

包括下载、预处理、数据加载、验证、检查点、导出以及将输出从节点复制出去。同时测量吞吐量与内存、数据加载停顿、通信时间、检查点持续时间和验证质量。

7 天期限适用于有边界的兼容性和性能分析。只有在您能够估算有效工作和运营缓冲之后,才选择 30 天预订。只有当额外计算超过通信和输入流水线限制时,更多 GPU 才有帮助。

了解固定期限计费

来源与延伸阅读

PyTorch 分布式概览PyTorch 峰值已分配内存PyTorch torchrunPyTorch 激活检查点PyTorch 分布式检查点PyTorch 在 ROCm 上
比较训练配置