决定模型分片策略#

大型模型通常无法容纳在单个 GPU 上，需要跨多个 GPU 进行分片。用于实现此目的的分片策略会对性能产生重大影响。本指南将逐步介绍如何确定张量并行、流水线并行或两者混合使用是否最适合您。如果您不熟悉张量并行和流水线并行，请参阅掌握 LLM 技术 - 推理优化

如何考虑模型分片：通信是关键#

在多个 GPU 之间拆分模型权重需要它们彼此之间通信激活，从而增加额外的开销。此开销在您的系统上的成本高低是确定最佳策略的关键因素。

在流水线并行中，模型被分成连续层的集合，每个 GPU 容纳这些集合中的一个。在这种设置中，唯一需要的通信是每个 GPU 将其集合的输出发送到具有下一个集合的 GPU。

Pipeline Parallel Visualization

另一方面，张量并行采用模型的每一层并将其拆分到 GPU 之间。这意味着每个 GPU 都容纳每一层的一部分。但是，由于每一层都需要前一层的完整输出作为输入，因此每个 GPU 必须执行更重的 All-Reduce 通信操作，以与其他所有 GPU 共享其结果，然后才能开始处理下一层。虽然这看起来不利，但由于每个 GPU 仅保存部分层，因此它也执行较小的矩阵乘法，从而使其能够更快地计算其输出。

Tensor Parallel Visualization

最终，决定最佳策略取决于 All-Reduce 操作的额外开销是否掩盖了较小矩阵乘法带来的收益。如果 GPU 之间的互连足够快，则每层计算负担减少带来的收益可能会超过额外的通信成本。因此，一般的经验法则是，如果您的 GPU 之间具有像 NVLink 这样的快速连接，那么张量并行可能是一个不错的选择。但是，如果通信将通过慢速连接（例如跨节点）进行，则流水线并行可能更好。总的来说，我们提供以下指南

如果您的模型适合在单个 GPU 上运行： 除非您有非常具体的原因，否则不要对模型进行分片。最佳通信开销是没有通信开销。

如果您的模型适合在单个节点上运行： 张量并行可能是这里的最佳选择，尤其是在 GPU 之间具有像 NVLink 这样的快速连接的情况下。如果没有，则可能需要流水线并行。从张量并行开始，并进行健全性检查以确定流水线并行是否更好。

如果您的模型跨多个节点分片： 节点间连接通常比节点内连接慢得多，因此，如果您在节点之间进行张量并行，则它会受到慢速互连的瓶颈影响。因此，一个好的起点是在节点内进行张量并行，在节点之间进行流水线并行。例外情况是您在 NVL36 或 NVL72 Blackwell 系统上运行。这些系统具有多节点 NVLink，因此只要您保持在 36 或 72 个 GPU 范围内，张量并行就不会受到节点间连接的瓶颈影响。

如何设置张量并行和流水线并行#

LLM 类采用 tensor_parallel_size 和 pipeline_parallel_size 作为参数。tensor_parallel_size * pipeline_parallel_size 应等于您在模型上进行分片的 GPU 总数，称为世界大小。例如，如果您要跨 2 个节点对模型进行分片，每个节点有 16 个 GPU，您可能会将张量并行设置为 8（用于节点内的张量并行），并将流水线并行设置为 2（节点之间的流水线并行），如下所示

    llm = LLM(
        model="/scratch/Llama-3.1-405B-Instruct",
        tensor_parallel_size=8,
        pipeline_parallel_size=2
    )

如果您使用的是用于构建引擎的 CLI 流程，您可以通过将 --tp_size 和 --tp_size 参数提供给 convert_checkpoint.py 来指定张量并行和流水线并行

python examples/llama/convert_checkpoint.py --model_dir ./tmp/llama/405B/ \
                            --output_dir ./tllm_checkpoint_16gpu_tp8_pp2 \
                            --dtype float16 \
                            --tp_size 8
                            --pp_size 2