了解张量并行(TP)和流水线并行(PP)的理论很重要,但将它们付诸实践需要专门的工具。Megatron-LM 是 NVIDIA 开发的一个著名库,专门设计用于实现这些复杂的并行策略,尤其适合训练大型 Transformer 模型。它提供并行层的优化实现,并管理所需的复杂通信模式。本节说明如何配置 Megatron-LM 以使用 TP 和 PP。Megatron-LM 中的配置通常通过命令行参数传递给训练脚本(如 pretrain_gpt.py 或类似脚本)来处理。我们来看具体如何设置张量并行和流水线并行。张量并行 (TP) 配置张量并行,有时也称为层内模型并行,涉及将单个大型层(如 MLP 块或注意力机制中的权重矩阵)的执行拆分到多个 GPU 上。Megatron-LM 提供专门的层实现(例如 ColumnParallelLinear、RowParallelLinear),它们自动处理此分区和必要的通信(如 AllReduce 或 AllGather)。启用和控制 TP 的主要参数是 --tensor-model-parallel-size(或根据 Megatron-LM 的特定版本或分支而异的类似变体)。这个值,我们称之为 $TP_{size}$,指定每个张量并行层将被拆分到多少个 GPU 上。例如,如果设置 --tensor-model-parallel-size 4,一个大型线性层的权重矩阵将被划分为 4 个列或行段,每个段驻留在张量并行组内的不同 GPU 上。在前向和反向传播过程中,Megatron-LM 协调这 4 个 GPU 之间所需的数据交换。# 在 Megatron-LM 环境中的示例 # (实际使用涉及 Megatron 的模型定义工具) # 假设 tp_size = 2 # GPU 0 存放 W_A,GPU 1 存放 W_B,并且 W = [W_A W_B] tp_group = get_tensor_model_parallel_group() # 获取 TP 通信组的函数 # ColumnParallelLinear: 拆分权重列, # 输入被广播,输出被收集 # 输入 X -> [X W_A, X W_B] -> AllGather -> 输出 Y linear_column_parallel = ColumnParallelLinear( input_size, output_size, tp_group=tp_group, ) # RowParallelLinear: 拆分权重行, # 输入被拆分,输出被归约 # 输入 X -> [X_A, X_B] -> [X_A W_A, X_B W_B] -> AllReduce -> 输出 Y linear_row_parallel = RowParallelLinear( input_size, output_size, tp_group=tp_group, ) # Transformer 块内的使用示例 (简化版) # mlp_output = linear_row_parallel(dropout( # linear_column_parallel(layer_input)))TP 的重要考虑因素:可除性: $TP_{size}$ 必须能整除模型的隐藏维度(--hidden-size)和注意力头的数量(--num-attention-heads)。Megatron-LM 的并行层依赖此进行分区。通信: TP 会在张量并行组内引入通信开销,通常是 AllReduce 操作。效率在很大程度上取决于这些 GPU 之间的互连速度(例如 NVLink)。内存: TP 减少了组内每个 GPU 上权重、梯度和优化器状态所需的内存,但由于通信需求,增加了激活内存使用量。流水线并行 (PP) 配置流水线并行涉及将模型的层按顺序划分到不同的 GPU 上,形成一个流水线。每个 GPU(或结合 TP/DP 时的 GPU 组)构成流水线中的一个“阶段”,负责只执行模型层的一个子集。配置 PP 的主要参数是 --pipeline-model-parallel-size(或类似参数)。这个值,$PP_{size}$,决定流水线的阶段数量。例如,--pipeline-model-parallel-size 4 创建一个 4 阶段的流水线。为了保持流水线阶段的利用率并最小化空闲时间(“气泡”),Megatron-LM 采用微批处理。整体训练批次被拆分为更小的微批,这些微批并发地流经流水线阶段。微批数量是一个重要的调整参数,通常由 --num-microbatches 控制,或根据全局批次大小、微批次大小和数据并行度计算得出。常见的调度方法是 1F1B(每个微批一次前向传播,一次反向传播),这有助于平衡计算和内存使用。digraph G { rankdir=LR; node [shape=box, style=filled, color="#ced4da", fontname="helvetica", fontsize=24]; edge [fontname="helvetica"]; subgraph cluster_0 { label = "阶段 0"; style=filled; color="#e9ecef"; fontsize=24; node [color="#74c0fc", fontsize=24]; s0_mb1 [label="前向 MB1", fontsize=24]; s0_mb2 [label="前向 MB2", fontsize=24]; s0_mb3 [label="前向 MB3", fontsize=24]; s0_mb1 -> s0_mb2 -> s0_mb3 [style=invis, fontsize=24]; // 阶段内顺序 } subgraph cluster_1 { label = "阶段 1"; style=filled; color="#e9ecef"; fontsize=24; node [color="#91a7ff", fontsize=24]; s1_mb1 [label="前向 MB1", fontsize=24]; s1_mb2 [label="前向 MB2", fontsize=24]; s1_mb3 [label="前向 MB3", fontsize=24]; s1_mb1 -> s1_mb2 -> s1_mb3 [style=invis, fontsize=24]; } subgraph cluster_2 { label = "阶段 2"; style=filled; color="#e9ecef"; fontsize=24; node [color="#e599f7", fontsize=24]; s2_mb1 [label="前向 MB1", fontsize=24]; s2_mb2 [label="前向 MB2", fontsize=24]; s2_mb3 [label="前向 MB3", fontsize=24]; s2_mb1 -> s2_mb2 -> s2_mb3 [style=invis, fontsize=24]; } s0_mb1 -> s1_mb1 [label=" 激活值", fontsize=24]; s0_mb2 -> s1_mb2 [label=" 激活值", fontsize=24]; s0_mb3 -> s1_mb3 [label=" 激活值", fontsize=24]; s1_mb1 -> s2_mb1 [label=" 激活值", fontsize=24]; s1_mb2 -> s2_mb2 [label=" 激活值", fontsize=24]; s1_mb3 -> s2_mb3 [label=" 激活值", fontsize=24]; // 可以添加反向传播箭头以完整展示 1F1B }微批(MB1、MB2、MB3)通过 3 阶段流水线($PP_{size}=3$)的前向传播(Fwd)流程。激活值(Actvs)在阶段之间传递。反向传播遵循类似的反向流程。PP 的重要考虑因素:层数: 变压器层的总数(--num-layers)理想情况下应能被 $PP_{size}$ 整除,以实现负载均衡,尽管 Megatron-LM 可以处理不均匀的分布。通信: PP 涉及相邻阶段之间的点对点通信,以向前传递激活值和向后传递梯度。流水线气泡: PP 的效率取决于最小化流水线气泡(处理批次开始和结束时的空闲时间)。增加微批数量有帮助,但也会增加激活内存。内存平衡: 内存使用在不同阶段可能无法完美平衡,因为第一和最后阶段通常包含嵌入层和输出层,这些层的大小可能与中间变压器块不同。结合张量并行与流水线并行Megatron-LM 擅长结合不同的并行策略。通常,TP 和 PP 会一起使用。在这种设置中,模型并行涉及的 GPU 总数是 $TP_{size} \times PP_{size}$。每个流水线阶段本身可能包含多个 GPU 通过张量并行一起工作。数据并行(DP)可以在此之上叠加,复制这种 TP/PP 结构。训练中使用的 GPU 总数将是 $N_{gpus} = DP_{size} \times TP_{size} \times PP_{size}$。digraph G { rankdir=LR; // 为所有元素设置全局字体大小 fontsize=24; node [shape=record, style=filled, fontname="helvetica", fontsize=24]; edge [fontname="helvetica", fontsize=24]; subgraph cluster_PP0 { label = "流水线阶段 0"; style=filled; color="#e9ecef"; fontsize=24; // 显式设置集群标签字体大小 node [color="#74c0fc", fontsize=24]; PP0 [label="{ GPU 0 | GPU 1 } | TP 组 0", fontsize=24]; } subgraph cluster_PP1 { label = "流水线阶段 1"; style=filled; color="#e9ecef"; fontsize=24; // 显式设置集群标签字体大小 node [color="#91a7ff", fontsize=24]; PP1 [label="{ GPU 2 | GPU 3 } | TP 组 1", fontsize=24]; } subgraph cluster_PP2 { label = "流水线阶段 2"; style=filled; color="#e9ecef"; fontsize=24; // 显式设置集群标签字体大小 node [color="#e599f7", fontsize=24]; PP2 [label="{ GPU 4 | GPU 5 } | TP 组 2", fontsize=24]; } subgraph cluster_PP3 { label = "流水线阶段 3"; style=filled; color="#e9ecef"; fontsize=24; // 显式设置集群标签字体大小 node [color="#ffc9c9", fontsize=24]; PP3 [label="{ GPU 6 | GPU 7 } | TP 组 3", fontsize=24]; } PP0 -> PP1 [label=" 激活值/梯度", fontsize=24]; PP1 -> PP2 [label=" 激活值/梯度", fontsize=24]; PP2 -> PP3 [label=" 激活值/梯度", fontsize=24]; }8 个 GPU 的配置示例:$PP_{size}=4$ 和 $TP_{size}=2$。每个流水线阶段使用 2 个 GPU 进行张量并行。通信发生在阶段之间(PP)的点对点方式,以及每个 TP 组内部(TP)的集体操作。配置示例这里有一个简化的示例,说明在使用 TP 和 PP 配置 Megatron-LM 进行 GPT 风格模型训练时,命令行参数可能是什么样子(假设总共有 8 个 GPU,其中 $TP_{size}=2$ 和 $PP_{size}=4$,为简单起见,数据并行 $DP_{size}=1$):# 使用 Megatron-LM 参数的训练命令示例 # (实际脚本名称和其他参数可能有所不同) python pretrain_gpt.py \ --num-layers 24 \ --hidden-size 2048 \ --num-attention-heads 32 \ --seq-length 2048 \ \ --tensor-model-parallel-size 2 \ --pipeline-model-parallel-size 4 \ \ --micro-batch-size 4 \ --global-batch-size 128 \ # 假设 DP 大小为 1,微批数量 = 全局批大小 / (微批大小 * dp_大小) # 需要根据 DP 大小仔细计算。这里假设 DP=1。 # 每个流水线的有效批次 = 微批大小 * 微批数量 # 全局批大小 = 微批大小 * 微批数量 * dp_大小 # 微批数量 = 全局批大小 / (微批大小 * dp_大小) = 128 / (4 * 1) = 32 --num-microbatches 32 \ \ --optimizer adam \ --learning-rate 1.0e-4 \ --weight-decay 0.01 \ --clip-grad 1.0 \ \ --train-data-path <path_to_train_data> \ --valid-data-path <path_to_valid_data> \ --tokenizer-type SentencePieceTokenizer \ --tokenizer-model <path_to_tokenizer> \ \ --distributed-backend nccl \ --save <path_to_save_checkpoints> \ --load <path_to_load_checkpoints> \ --log-interval 10 \ --save-interval 1000 \ --eval-interval 100 \ --num-workers 2在此示例中:--tensor-model-parallel-size 2:将线性层和注意力层等拆分到 2 个 GPU 上。--pipeline-model-parallel-size 4:将 24 个层拆分到 4 个顺序阶段(每个阶段大约 6 层)。--num-microbatches 32:用于高效地为 4 阶段流水线提供数据,根据全局批次大小、微批次大小和数据并行度计算得出。正确配置这些参数对于平衡计算、内存使用和可用硬件上的通信开销非常重要。Megatron-LM 提供底层的机制,但确定最佳的 $TP_{size}$、$PP_{size}$ 和微批数量通常需要根据具体的模型架构、硬件配置(GPU 类型、互连)和所需的批次大小进行实验。