分布式训练

使用 TRL 进行多 GPU 训练

TRL 中的训练器通过 PyTorch 的分布式启动器 torchrun 启动。它为每个 GPU 启动一个进程：

torchrun --nproc_per_node 8 train.py <SCRIPT_ARGS>

trl CLI 在底层执行相同的操作，因此 trl sft ... 会在机器的所有 GPU 上运行；传入 --nproc_per_node 可使用更少的 GPU。

在底层，🤗 Accelerate 会为每个 GPU 创建一个模型。每个进程：

• 处理自己的数据批次
• 计算该批次的损失和梯度
• 在所有 GPU 之间共享梯度更新

有效批量大小（effective batch size）的计算公式为： $$\text{Batch Size} = \text{per\_device\_train\_batch\_size} \times \text{num\_devices} \times \text{gradient\_accumulation\_steps}$$

在扩展到多 GPU 时，为保持批处理大小一致，请相应地更新 per_device_train_batch_size 和 gradient_accumulation_steps。

例如，以下配置是等效的，并且应产生相同的结果：

每个 GPU 上运行一个模型可能导致内存占用过高，对于大型模型或低内存 GPU 而言可能不可行。在这种情况下，您可以利用 DeepSpeed，它提供了模型分片、Zero Redundancy Optimizer、混合精度训练以及卸载到 CPU 或 NVMe 等优化。有关更多详细信息，请参阅我们的 DeepSpeed Integration 指南。

针对超长序列的训练有专门的指南：训练超过 100 万 Token。

多节点训练

当单台机器的 GPU 数量不足时，torchrun 可以将训练扩展到多台机器（节点）。每个节点都运行相同的命令，并指定节点总数、自身的 rank 以及主节点上的 rendezvous 端点。

选项 1：手动启动（非 HPC）

在每个节点上手动运行以下内容：

# Node 0 (main node, IP 10.0.0.1)
torchrun --nnodes 2 --nproc_per_node 8 --node_rank 0 --rdzv_backend c10d --rdzv_endpoint 10.0.0.1:29500 train.py

# Node 1
torchrun --nnodes 2 --nproc_per_node 8 --node_rank 1 --rdzv_backend c10d --rdzv_endpoint 10.0.0.1:29500 train.py

将 10.0.0.1 替换为 rank 0（主）节点的实际 IP 地址。

选项 2：SLURM 启动（HPC 集群）

对于使用 SLURM 作业调度器的集群，请创建一个作业脚本（例如 slurm_job.sh）：

#!/bin/bash
#SBATCH --nodes=2
#SBATCH --gpus-per-node=8
#SBATCH --job-name=trl_multi

MAIN_NODE=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
srun torchrun --nnodes 2 --nproc_per_node 8 --node_rank $SLURM_NODEID --rdzv_backend c10d --rdzv_endpoint $MAIN_NODE:29500 train.py

然后提交任务：

sbatch slurm_job.sh

srun 在每个节点上启动一个 torchrun，而每个 torchrun 又会在每个 GPU 上启动一个进程。

关键 SLURM 指令：

• --nodes=2：请求 2 个计算节点
• --gpus-per-node=8：为每个节点分配 8 个 GPU（共 16 个）
• --job-name：用于在作业队列中跟踪的标签

您可以通过在训练配置中设置 deepspeed 来将多节点与 DeepSpeed 结合使用。请参阅 DeepSpeed 集成指南。

延伸阅读

• torchrun（弹性启动）
• SLURM 工作负载管理器文档 - 用于集群作业调度