多机多卡(Distributed Training)微调大模型是一项复杂但非常高效的任务。它允许你利用多台机器的计算资源来训练一个模型,从而显著缩短训练时间。
多机多卡微调核心流程
整个流程可以概括为以下几个核心步骤:
- 环境准备与硬件配置
• 硬件:确保拥有多台(例如2-4台)服务器/工作站,每台机器配备多张GPU(通常是4卡或8卡)。机器之间需要通过高速网络互联(如InfiniBand或高速万兆以太网),低延迟和高带宽对通信效率至关重要。
• 软件与环境:
◦ 操作系统:通常使用Linux(如Ubuntu)。◦ 驱动:在所有机器上安装相同版本的NVIDIA GPU驱动。◦ CUDA & cuDNN:在所有机器上安装相同版本的CUDA工具包和cuDNN。◦ Python环境:使用conda或venv创建一致的Python环境,确保所有机器上的Python版本、PyTorch/TensorFlow版本、以及其他依赖库版本完全一致。这是避免莫名错误的关键。◦ 通信库:安装并配置好NCCL(NVIDIA Collective Communication Library),它是多机多卡通信的基础,通常随PyTorch等框架一起安装。
- 代码修改与分布式策略选择
这是最核心的技术部分。你需要修改你的训练代码,使其支持分布式训练。
主流方案选择:
• A. 基于数据并行(Data Parallelism)
◦ 核心思想:将一份完整的模型复制到每台机器的每个GPU上,然后将训练数据分成多个小块(batch),分配到各个GPU上并行计算。每个GPU计算完梯度后,再进行梯度同步(所有GPU的梯度求平均),最后每个GPU用同步后的梯度更新自己的模型参数。这样,所有GPU上的模型始终保持一致。◦ 实现方式:▪ DP (DataParallel):PyTorch早期的单机多卡方案,不适用于多机,且效率较低(主卡负载重),不推荐。▪ DDP (DistributedDataParallel):当前PyTorch的多机多卡首选方案。每个GPU都有一个独立的进程,通过NCCL在后台进行高效的All-Reduce通信来完成梯度同步,效率远高于DP。◦ 适用场景:模型本身可以完整地放入单个GPU的显存中。这是最常见和最简单的场景。
• B. 基于模型并行(Model Parallelism)或混合并行
◦ 核心思想:当模型非常大,单个GPU的显存放不下整个模型时,就需要将模型的不同层拆分到不同的GPU甚至不同的机器上。◦ 实现方式:▪ 流水线并行 (Pipeline Parallelism):将模型按层切分到多个GPU上。一个batch的数据会像在流水线上一样,依次经过这些GPU进行计算。需要解决“流水线气泡”问题。▪ 张量并行 (Tensor Parallelism):将模型内部的单个大权重矩阵运算(如GEMM)拆分到多个GPU上并行计算。例如,Megatron-LM和ColossalAI主要采用此方式。▪ FSDP (Fully Sharded Data Parallel):强烈推荐的现代方案。它结合了数据并行和模型并行。它会将模型参数、梯度、优化器状态分片到所有GPU上,从而极大地节省了每张GPU的显存占用。在计算时,需要哪个分片的参数就临时通过通信获取,计算完后再丢弃。这使得我们可以用更少的GPU来微调极大的模型(如LLaMA2-70B)。◦ 适用场景:微调超大规模模型(通常 > 10B 参数)。
一个简单的比喻:
•
DDP 就像是一个交响乐团。每个乐手(GPU)面前都有一份完整的乐谱(完整模型),大家各自演奏自己的部分(处理数据分片),但会听着指挥(梯度同步)来保持节奏一致。
•
FSDP 就像是一个剧团在排练一部厚厚的剧本。剧本(完整模型)被拆成很多页,分发给不同的演员(GPU)保管。当排练到某一场戏时,需要这场戏台词的那些演员就把自己保管的那几页纸拿出来念,念完就收回去。这样不需要每个演员都拿着整本厚厚的剧本
代码修改关键点(以PyTorch DDP为例):
-
初始化进程组:使用 torch.distributed.init_process_group 初始化分布式环境,指定后端(通常是nccl)、当前进程的rank、世界大小(总进程数=机器数*每台机器的GPU数)等。
-
包装模型:使用 torch.nn.parallel.DistributedDataParallel 包装你的模型。
-
分配数据:使用 torch.utils.data.distributed.DistributedSampler 确保每个进程/GPU只加载和数据整体的一部分,避免数据重复。
-
调整日志与评估:通常只在rank 0进程(主进程)上打印日志、保存模型、做验证评估,避免重复输出。
-
启动训练任务
使用PyTorch内置的 torch.distributed.launch 或更推荐的 torchrun 来启动训练脚本。
示例命令:
假设有2台机器(main_node: 192.168.1.1, node2: 192.168.1.2),每台机器有4张GPU。在每台机器上执行(或在主节点上通过ssh触发):
# 在主节点 (rank 0) 上执行
torchrun \--nnodes=2 \ # 总机器数--nproc_per_node=4 \ # 每台机器的GPU数--rdzv_id=456 \ # 任务唯一ID--rdzv_backend=c10d \ # 后端--rdzv_endpoint=192.168.1.1:1234 \ # 主节点地址和端口--node_rank=0 \ # 当前机器的rank(主节点为0)your_training_script.py \--your_script_args# 在第二台机器 (node2, rank 1) 上执行
torchrun \--nnodes=2 \--nproc_per_node=4 \--rdzv_id=456 \--rdzv_backend=c10d \--rdzv_endpoint=192.168.1.1:1234 \ # 注意这里依然是主节点地址--node_rank=1 \ # 当前机器的rank(第二台为1)your_training_script.py \--your_script_args
- 监控与调试
• 监控:使用nvidia-smi监控GPU利用率、显存占用。使用htop监控CPU和内存。使用NCCL调试日志(export NCCL_DEBUG=INFO)来观察通信是否正常。
• 日志:确保只有主进程(rank 0)记录日志和保存TensorBoard文件,避免混乱。
• 保存与加载 checkpoint:保存 checkpoint 时,通常也只在 rank 0 进程上进行。加载时,需要先将 checkpoint 广播到所有进程,或者使用 model.module.load_state_dict() 来加载。
关键注意事项
-
环境一致性:这是最大的坑!所有机器上的软件环境(Python, PyTorch, CUDA, 库版本)必须高度一致,否则会出现难以排查的错误。
-
网络配置:
◦ 防火墙:确保所有机器之间在指定端口(如上面的1234)上可以互相通信。关闭防火墙或配置规则允许分布式训练使用的端口范围。◦ SSH互信(可选但推荐):方便在主节点上一键启动所有机器上的任务。
-
数据准备:
◦ 数据集需要放在所有机器都能访问的位置(如共享文件系统NFS,或每台机器本地都有一份完整的数据副本)。◦ DistributedSampler 会自动为每个进程分配不重复的数据,确保每个batch的数据只被计算一次。
-
性能调优:
◦ 通信瓶颈:如果网络慢,梯度同步会成为瓶颈。确保使用高速网络(InfiniBand > 10GbE > 1GbE)。◦ Batch Size:全局batch size = 每卡batch_size * GPU总数。增大全局batch size时,可能需要调整学习率(线性缩放或使用LARS/LAMB等自适应优化器)。
◦ 梯度累积:如果每张GPU能承受的batch size很小,可以通过梯度累积来模拟更大的全局batch size,即多次前向后累积梯度再同步和更新。
-
容错性:
◦ 多机任务运行时间长,一台机器出错可能导致整个任务失败。需要思考如何做断点续训。◦ 使用 torchrun 时,它具备一定的弹性功能,可以处理部分节点失败的情况。
-
选择正确的并行策略:
◦ 如果显存够用,优先选择 DDP,最简单高效。◦ 如果显存不足,尝试使用 FSDP,这是目前社区微调大模型的主流和推荐方式。
◦ 除非模型极大(如万亿参数),否则一般不需要手动配置流水线并行或张量并行。
总结
多机多卡微调的核心流程是:环境配置 -> 选择并行策略并修改代码 -> 使用 torchrun 启动任务 -> 监控和调试。
对于初学者,建议从 单机多卡(DDP) 开始,彻底掌握其原理和流程后,再扩展到多机。对于微调超大规模模型(如70B),直接学习并使用 FSDP。
工具选择上,可以优先使用PyTorch原生的DDP/FSDP,如果追求更高级的功能和优化,也可以考虑基于PyTorch的第三方框架如 Deepspeed 或 ColossalAI。