快捷方式

多节点微调

恭喜!您终于摆脱了“GPU 贫乏”的困境,现在可以使用多节点设置了。您可以告别绞尽脑汁进行内存优化的日子,但请为迎接分布式计算复杂性带来的新挑战做好准备。

您将学习到
  • 为什么多节点训练很有用

  • 如何在 SLURM 集群上设置 torchtune 包

  • 如何使用全参数更新(而非 LoRA)微调 Llama3.3 70B 模型

前提条件
  • 熟悉 torchtune 中的分布式训练

  • 已经了解基本的 SLURM 命令

多节点训练的优势

更多机器意味着更多内存!这在几个方面都很棒

  1. 更大模型:内存更多,可以训练更大的模型,例如 Llama3.1 405BDeepseek-V3 等。

  2. 更长数据:对于许多微调任务(例如编写代码),长上下文长度很有帮助;但是,更长的上下文长度意味着激活需要更多内存。

  3. 更高质量:内存更多,可以进行全参数更新(而非 LoRA),并使用像 AdamW 这样的优化器(而非低精度优化器),这两者都可能提高训练质量。

  4. 更快训练:能够将更多数据放入内存,可以使用更高的批处理大小,并且关闭内存优化,例如激活检查点,从而减少训练完成所需的时间。

注意

低节点间带宽与 FSDP 我们利用 PyTorch 的 Fully Sharded Data Parallel (完全分片数据并行) 在多个设备上分发模型。为了分发训练,FSDP 在每次前向传播时运行一个 all-gather 操作,并在每次反向传播时运行一个 all-gather(通常)加上一个 reduce-scatter 操作。这些操作(通常)会阻塞训练直到完成,并且如果节点间连接速度慢,训练速度可能会降低。有关更多信息,请参阅 此 Github Issue

在 2 个节点上训练 Llama3.3 70B

开始训练吧!我们将使用一种常见的集群工作流管理器,名为 SLURM,并且假设您在本教程中对 SLURM 有一定的了解。首先,我们需要安装 torchtune。虽然这与常规安装说明一样简单直接,但建议您将软件包安装到集群中所有节点都可以访问的虚拟环境中,例如共享文件系统。

接下来,我们需要将 Llama3.3 70B 模型下载到您的共享文件系统。您需要按照此处概述的步骤确保拥有正确的凭据。

$ tune download meta-llama/Llama-3.3-70B-Instruct --ignore-patterns "consolidated/*.pth" --output-dir SHARED_FS/Llama-3.3-70B-Instruct

现在我们已经下载了模型,接下来看看我们的 SLURM bash 脚本示例。

#!/bin/bash
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# ---------- SBATCH commands ---------- #
#SBATCH --job-name=torchtune-multi-node
#SBATCH --ntasks=2
#SBATCH --nodes=2
#SBATCH --gpus-per-task=8
#SBATCH --cpus-per-task=96
#SBATCH --partition=train

# ---------- Set env variables ---------- #
# Grab the IP for head node:
# You may need to set this to the fully qualified domain name of your head node
nodes=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) )
nodes_array=($nodes)
head_node=${nodes_array[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)
echo Node IP: $head_node_ip

# You might need to explicitly set the network interface for distributed backends:
# export NCCL_SOCKET_IFNAME=...
# export GLOO_SOCKET_IFNAME=...

export TORCH_DIST_INIT_BARRIER=1
export LOGLEVEL=INFO

# ---------- Launch training ---------- #
# You probably want to load in a virtual env w/ conda...
# module load conda
# conda activate torchtune
# ...or venv
# source torchtune/bin/activate

SHARED_FS=/mnt/slurm # <-- Replace w/ your filesystem
CHECKPOINT_DIR="$SHARED_FS/Llama-3.3-70B-Instruct"
OUTPUT_DIR="$SHARED_FS/Llama3.3-70B-fft-output"

# Adjust sbatch --ntasks and sbatch --nodes above and --nnodes below to your specific node count
srun tune run --nnodes 2 --nproc_per_node 8 --rdzv_id 101 --rdzv_backend c10d --rdzv_endpoint "$head_node_ip:29500" \
    full_finetune_distributed --config llama3_3/70B_full_multinode checkpoint_dir=$CHECKPOINT_DIR output_dir=$OUTPUT_DIR

这个脚本中有很多信息,但以下是核心部分

  • 我们使用了 SLURM 特定命令,例如节点数、任务数、可用 CPU 等。

  • 我们正在使用 torchrunfull_finetune_distributed 代码示例集进行训练,就像在单节点上一样。

  • 您可以考虑一些集群特定的环境变量(NCCL_BUFFSIZE, NCCL_DEBUG, FI_PROVIDER 等),以最大化 GPU 利用率、进行调试等。

注意

我们可能需要显式设置分布式后端的网络接口。您可以在此处阅读更多关于 PyTorch 分布式后端的信息,同时了解通过在特定节点上运行 ipconfig 来查找您的网络接口也很有帮助。

在 bash 脚本中更新共享文件系统后,我们可以使用 sbatch 启动。

sbatch full_finetune_multinode.slurm

并且 squeue 的输出应该显示我们的作业正在运行。

$ squeue
JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
1     train         torchtun slurm R       0:03      2 slurm-worker-[1-2]

训练完成后(使用默认配置,总共大约需要七分钟,速度为 880 tok/s),我们可以按照此处的说明将我们漂亮的新模型上传到 Hugging Face Hub!

未来发展

我们已经介绍了如何在 SLURM 上使用 FSDP 在两个节点上启动微调作业的基础知识。我们还在准备更多内容,包括……

二维并行:利用 FSDP 张量并行(通常称为二维并行)将进一步降低内存需求,使我们能够更充分地利用上述列出的优势。

更长上下文(环形注意力等):更多内存和更多机器意味着我们可以在更长的序列上训练,并利用环形注意力等巧妙技巧(其中 token 在 GPU 之间分割)。您可以在此 Github RFC 中阅读有关 torchtune 未来计划的更多信息。

想要其他优化?随时可以通过在我们的仓库上提交 Github Issue 或在 Discord 中联系我们来告诉我们!

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深度教程

查看教程

资源

查找开发资源并获得问题解答

查看资源