性能实测: Llama 3.1 8B/70B 在 H100 集群上的 micro-pipeline 提速基准)
torchtitan 异步张量并行(Async TP)性能实测: Llama 3.1 8B/70B 在 H100 集群上的 micro-pipeline 提速基准【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan本文解读 torchtitan 仓库中 2025 年 6 月 PyTorch 团队完成的 Async TP(异步张量并行)基准测试报告 benchmarks/asyncTP_llama3_h100_2025-06_torchtitan.md。读完后你将了解:Async TP 相对原生 TP 在 Llama 3.1 8B/70B 训练上的实际吞吐增益(bf16 与 float8 两种精度、两种量化策略)、测试所用的硬件平台与训练配置,以及如何通过 torchtitan 的--compile.enable_async_tensor_parallel开关在源码层面启用该优化。什么是 Async TP:把 TP 集合通信藏进矩阵乘法原生(即 Vanilla)张量并行中,每个 GEMM 之后都必须同步执行一次 all-gather/reduce-scatter 等 TP 集合通信,计算流与通信流是串行的:GEMM 做完、通信做完,再做下一段计算。Async TP 的思路(在 torchtitan 配置中的官方描述是 pipeline tensor-parallel collectives with matrix multiplications,见 CompileConfig)是把 TP 的集合通信与矩阵乘法做 micro-pipeline 重叠,用后续 GEMM 的计算时间掩盖当前通信的延迟,从而在不改变数值语义的前提下提升端到端吞吐。从源码结构看,torchtitan 中该功能的落地路径非常清晰:配置项定义在 torchtitan/config/configs.py 的CompileConfig中,字段为enable_async_tensor_parallel: bool False,默认关闭;其__post_init__中包含一条硬性校验:启用 Async TP 必须同时满足--compile.enable且--compile.components中包含model,否则会直接抛出ValueError。这与基准报告中所有结果均基于 torch.compile 的测试前提一致——Async TP 依赖编译期对计算/通信图形的重排,不能脱离 torch.compile 独立生效;实际启用逻辑在 torchtitan/distributed/compile.py 的_maybe_enable_async_tp中:当tp_mesh存在(即tensor_parallel_degree 1,取自ParallelDims.get_dense_tp_mesh())且配置开启时,先为该 TP process group 注册 symmetric memory(enable_symm_mem_for_group),再打开 Inductor 的torch._inductor.config._micro_pipeline_tp True开关。这说明 Async TP 依赖两个底层能力:symmetric memory 提供低延迟的跨卡显存原语,Inductor micro-pipeline pass 负责把通信算子插入到计算流水中。测试模型与硬件平台基准报告使用的模型为Llama 3.1 8B 与 Llama 3.1 70B,与仓库中 torchtitan/models/llama3/config_registry.py 注册的llama3_8b、llama3_70b等配置对应。硬件为 Meta 的 Grand Teton 平台,报告中给出的关键规格:每台主机 8 张 NVIDIA H100 GPU,卡间通过 NVLink 全互联;每块 H100 配备 96GB HBM2e,峰值显存带宽 2.4 TB/s;主机之间通过后端 RDMA 网络连接,每 GPU 400 Gb/s;使用默认 500W 功耗上限;报告同时指出,把 TDP 调至 700W 可能进一步提速。这一平台规格对理解结果很重要:Async TP 的收益本质来自以计算掩盖通信,因此它天然更受益于通信占比更高的场景——大模型(70B)与更慢的跨卡通信。基准测试结果以下为报告中的完整数据,单位均为 tokens/sec(全局吞吐)。Llama 3.1 70B:256 张 H100,FSDP32,TP8,torch.compile,full AC,local batch size 16量化方式Vanilla TP tokens/secAsync TP tokens/secAsync TP 加速比None (bfloat16)597.3652.41.09float8 tensorwise809.8942.41.16float8 rowwise599.6624.81.04Llama 3.1 8B:64 张 H100,FSDP8,TP8,torch.compile,per op SAC,local batch size 12量化方式Vanilla TP tokens/secAsync TP tokens/secAsync TP 加速比None (bfloat16)43784809.41.10float8 tensorwise5078.15570.11.10float8 rowwise3708.53914.91.06报告特别注明:Vanilla TP 下 float8 rowwise 训练基线偏低的性能问题,由上游 issue pytorch/torchtitan #1207 跟进处理——这是解读 rowwise 列数据时必须了解的背景,其基线数字本身可能低估了 rowwise 路线的真实水平。从数据中可以提取几个值得注意的结论:Async TP 在全部 6 个配置下均取得正收益,加速比介于 1.04 到 1.16 之间,没有负收益或持平的极端情况;float8 tensorwise 收益最大(1.16x / 1.10x)。tensorwise 量化的 GEMM 吞吐更高、计算段更短,意味着通信占比相对更高,micro-pipeline 掩盖通信的空间更大;70B 的 bf16 场景(1.09x)明显强于 8B(1.10x 中 rowwise 仅 1.06x)的整体格局:70B 在 TP8 下通信/计算比更高,重叠收益更充分;8B 模型本身计算密集,GEMM 单段时间较长,可重叠的通信比例受限;rowwise 两行收益普遍最小(1.04x / 1.06x),与其 vanilla 基线本身偏低(见 #1207)这一事实相呼应。在 torchtitan 中启用 Async TP:从 CLI 到测试配方基准测试的复现入口是仓库根目录的 run_train.sh,它通过torchrun启动torchtitan.train并透传 Tyro 配置参数。以单节点 8 卡为例,启用 Async TP 的最小组合为:NGPU8 ./run_train.sh \ --parallelism.tensor_parallel_degree 8 \ --compile.enable \ --compile.enable_async_tensor_parallel其中--compile.enable与--compile.enable_async_tensor_parallel的组合由 CompileConfig 的校验逻辑 约束;--parallelism.tensor_parallel_degree提供tp_mesh,二者缺一不可,否则_maybe_enable_async_tp会因tp_mesh is None或配置未开启而静默跳过。仓库内的 GPU 测试配方提供了可直接参考的真实用法。torchtitan_recipes/tests/h100.py 中定义了多个 H100 集成测试配置,例如:llama3_debugmodel_tp2_asynctp_compile:在llama3_debugmodel(seq_len2048)上设置tensor_parallel_degree2、compile.enableTrue、compile.enable_async_tensor_parallelTrue;llama3_debugmodel_float8_fsdp2x2_tp2_pp2_asynctp_compile:进一步叠加 FSDP2、TP2、PP2、8 个 microbatch 并禁用 CUDA graph,验证 Async TP 与 float8 量化、流水并行组合时的行为。此外,torchtitan/models/llama3/config_registry.py 中llama3_405b的注册配置也默认携带CompileConfig(enableTrue, enable_async_tensor_parallelTrue),说明该开关已被视为大模型(405B, TP8)训练配方的一部分。仓库中还存在一条与 Async TP 相邻的优化路线:配置llama3_debugmodel_dist_gemm(见 config_registry 中的定义)通过tp_gemm_backenddist_gemm把 attention 的 TP 集合通信直接折叠进 GEMM,其文档说明明确标注了 Async-TP: the attention TP collectives are folded into their GEMMs,读者可以对照理解这两条通信-计算融合路线的差异。测试版本环境报告记录了基准运行时的依赖版本,复现或对比结果时应以该表为准:仓库commit日期torch38410cf92025-06-14torchao62430402024-06-13torchtitan820504e2024-06-13需要说明:表中 torchao 与 torchtitan 两行的日期按原文档原样保留,但结合同期 commit 与 torch 行的日期(2025-06-14),这两处日期疑为原文笔误。无论日期如何,基准所对应的代码快照以上述三个 commit 为准;若使用当前仓库 HEAD 复现,由于 Inductor 的_micro_pipeline_tppass 与 symmetric memory 接口在上游持续演进,吞吐数字可能与该报告存在偏差。小结这份 2025 年 6 月的基准报告证明了 Async TP 在 torchtitan 训练栈中是一致且无副作用的正向优化:Llama 3.1 8B/70B 在 bf16 与 float8(tensorwise/rowwise)共 6 组配置下,TP8 场景下稳定获得 4%~16% 的全局吞吐提升,其中 float8 tensorwise 与 bf16 大模型场景收益最明显。落地侧只需--compile.enable、--parallelism.tensor_parallel_degree与--compile.enable_async_tensor_parallel三个参数配合,底层由 torchtitan/distributed/compile.py 中的 symmetric memory 注册与 Inductor micro-pipeline 开关完成;torchtitan_recipes/tests/h100.py 中随附的 debug 模型配置则为小集群上的功能验证提供了现成入口。【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 torchtitan 异步张量并行(Async TP)性能实测: Llama 3.1 8B/70B 在 H100 集群上的 micro-pipeline 提速基准
torchtitan 异步张量并行(Async TP)性能实测: Llama 3.1 8B/70B 在 H100 集群上的 micro-pipeline 提速基准【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan本文解读 torchtitan 仓库中 2025 年 6 月 PyTorch 团队完成的 Async TP(异步张量并行)基准测试报告 benchmarks/asyncTP_llama3_h100_2025-06_torchtitan.md。读完后你将了解:Async TP 相对原生 TP 在 Llama 3.1 8B/70B 训练上的实际吞吐增益(bf16 与 float8 两种精度、两种量化策略)、测试所用的硬件平台与训练配置,以及如何通过 torchtitan 的--compile.enable_async_tensor_parallel开关在源码层面启用该优化。什么是 Async TP:把 TP 集合通信藏进矩阵乘法原生(即 Vanilla)张量并行中,每个 GEMM 之后都必须同步执行一次 all-gather/reduce-scatter 等 TP 集合通信,计算流与通信流是串行的:GEMM 做完、通信做完,再做下一段计算。Async TP 的思路(在 torchtitan 配置中的官方描述是 pipeline tensor-parallel collectives with matrix multiplications,见 CompileConfig)是把 TP 的集合通信与矩阵乘法做 micro-pipeline 重叠,用后续 GEMM 的计算时间掩盖当前通信的延迟,从而在不改变数值语义的前提下提升端到端吞吐。从源码结构看,torchtitan 中该功能的落地路径非常清晰:配置项定义在 torchtitan/config/configs.py 的CompileConfig中,字段为enable_async_tensor_parallel: bool False,默认关闭;其__post_init__中包含一条硬性校验:启用 Async TP 必须同时满足--compile.enable且--compile.components中包含model,否则会直接抛出ValueError。这与基准报告中所有结果均基于 torch.compile 的测试前提一致——Async TP 依赖编译期对计算/通信图形的重排,不能脱离 torch.compile 独立生效;实际启用逻辑在 torchtitan/distributed/compile.py 的_maybe_enable_async_tp中:当tp_mesh存在(即tensor_parallel_degree 1,取自ParallelDims.get_dense_tp_mesh())且配置开启时,先为该 TP process group 注册 symmetric memory(enable_symm_mem_for_group),再打开 Inductor 的torch._inductor.config._micro_pipeline_tp True开关。这说明 Async TP 依赖两个底层能力:symmetric memory 提供低延迟的跨卡显存原语,Inductor micro-pipeline pass 负责把通信算子插入到计算流水中。测试模型与硬件平台基准报告使用的模型为Llama 3.1 8B 与 Llama 3.1 70B,与仓库中 torchtitan/models/llama3/config_registry.py 注册的llama3_8b、llama3_70b等配置对应。硬件为 Meta 的 Grand Teton 平台,报告中给出的关键规格:每台主机 8 张 NVIDIA H100 GPU,卡间通过 NVLink 全互联;每块 H100 配备 96GB HBM2e,峰值显存带宽 2.4 TB/s;主机之间通过后端 RDMA 网络连接,每 GPU 400 Gb/s;使用默认 500W 功耗上限;报告同时指出,把 TDP 调至 700W 可能进一步提速。这一平台规格对理解结果很重要:Async TP 的收益本质来自以计算掩盖通信,因此它天然更受益于通信占比更高的场景——大模型(70B)与更慢的跨卡通信。基准测试结果以下为报告中的完整数据,单位均为 tokens/sec(全局吞吐)。Llama 3.1 70B:256 张 H100,FSDP32,TP8,torch.compile,full AC,local batch size 16量化方式Vanilla TP tokens/secAsync TP tokens/secAsync TP 加速比None (bfloat16)597.3652.41.09float8 tensorwise809.8942.41.16float8 rowwise599.6624.81.04Llama 3.1 8B:64 张 H100,FSDP8,TP8,torch.compile,per op SAC,local batch size 12量化方式Vanilla TP tokens/secAsync TP tokens/secAsync TP 加速比None (bfloat16)43784809.41.10float8 tensorwise5078.15570.11.10float8 rowwise3708.53914.91.06报告特别注明:Vanilla TP 下 float8 rowwise 训练基线偏低的性能问题,由上游 issue pytorch/torchtitan #1207 跟进处理——这是解读 rowwise 列数据时必须了解的背景,其基线数字本身可能低估了 rowwise 路线的真实水平。从数据中可以提取几个值得注意的结论:Async TP 在全部 6 个配置下均取得正收益,加速比介于 1.04 到 1.16 之间,没有负收益或持平的极端情况;float8 tensorwise 收益最大(1.16x / 1.10x)。tensorwise 量化的 GEMM 吞吐更高、计算段更短,意味着通信占比相对更高,micro-pipeline 掩盖通信的空间更大;70B 的 bf16 场景(1.09x)明显强于 8B(1.10x 中 rowwise 仅 1.06x)的整体格局:70B 在 TP8 下通信/计算比更高,重叠收益更充分;8B 模型本身计算密集,GEMM 单段时间较长,可重叠的通信比例受限;rowwise 两行收益普遍最小(1.04x / 1.06x),与其 vanilla 基线本身偏低(见 #1207)这一事实相呼应。在 torchtitan 中启用 Async TP:从 CLI 到测试配方基准测试的复现入口是仓库根目录的 run_train.sh,它通过torchrun启动torchtitan.train并透传 Tyro 配置参数。以单节点 8 卡为例,启用 Async TP 的最小组合为:NGPU8 ./run_train.sh \ --parallelism.tensor_parallel_degree 8 \ --compile.enable \ --compile.enable_async_tensor_parallel其中--compile.enable与--compile.enable_async_tensor_parallel的组合由 CompileConfig 的校验逻辑 约束;--parallelism.tensor_parallel_degree提供tp_mesh,二者缺一不可,否则_maybe_enable_async_tp会因tp_mesh is None或配置未开启而静默跳过。仓库内的 GPU 测试配方提供了可直接参考的真实用法。torchtitan_recipes/tests/h100.py 中定义了多个 H100 集成测试配置,例如:llama3_debugmodel_tp2_asynctp_compile:在llama3_debugmodel(seq_len2048)上设置tensor_parallel_degree2、compile.enableTrue、compile.enable_async_tensor_parallelTrue;llama3_debugmodel_float8_fsdp2x2_tp2_pp2_asynctp_compile:进一步叠加 FSDP2、TP2、PP2、8 个 microbatch 并禁用 CUDA graph,验证 Async TP 与 float8 量化、流水并行组合时的行为。此外,torchtitan/models/llama3/config_registry.py 中llama3_405b的注册配置也默认携带CompileConfig(enableTrue, enable_async_tensor_parallelTrue),说明该开关已被视为大模型(405B, TP8)训练配方的一部分。仓库中还存在一条与 Async TP 相邻的优化路线:配置llama3_debugmodel_dist_gemm(见 config_registry 中的定义)通过tp_gemm_backenddist_gemm把 attention 的 TP 集合通信直接折叠进 GEMM,其文档说明明确标注了 Async-TP: the attention TP collectives are folded into their GEMMs,读者可以对照理解这两条通信-计算融合路线的差异。测试版本环境报告记录了基准运行时的依赖版本,复现或对比结果时应以该表为准:仓库commit日期torch38410cf92025-06-14torchao62430402024-06-13torchtitan820504e2024-06-13需要说明:表中 torchao 与 torchtitan 两行的日期按原文档原样保留,但结合同期 commit 与 torch 行的日期(2025-06-14),这两处日期疑为原文笔误。无论日期如何,基准所对应的代码快照以上述三个 commit 为准;若使用当前仓库 HEAD 复现,由于 Inductor 的_micro_pipeline_tppass 与 symmetric memory 接口在上游持续演进,吞吐数字可能与该报告存在偏差。小结这份 2025 年 6 月的基准报告证明了 Async TP 在 torchtitan 训练栈中是一致且无副作用的正向优化:Llama 3.1 8B/70B 在 bf16 与 float8(tensorwise/rowwise)共 6 组配置下,TP8 场景下稳定获得 4%~16% 的全局吞吐提升,其中 float8 tensorwise 与 bf16 大模型场景收益最明显。落地侧只需--compile.enable、--parallelism.tensor_parallel_degree与--compile.enable_async_tensor_parallel三个参数配合,底层由 torchtitan/distributed/compile.py 中的 symmetric memory 注册与 Inductor micro-pipeline 开关完成;torchtitan_recipes/tests/h100.py 中随附的 debug 模型配置则为小集群上的功能验证提供了现成入口。【免费下载链接】torchtitanA PyTorch native platform for training generative AI models项目地址: https://gitcode.com/GitHub_Trending/to/torchtitan创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考