如何为SFTTrainer配置跨多主机的GPU训练环境？

针对你的多主机GPU环境（Host1:1卡、Host2:1卡、Host3:2卡、Host4:2卡），要让SFTTrainer利用所有GPU，需按以下步骤配置：

一、前置环境准备 #

• 统一环境版本：所有主机必须安装完全一致的依赖包和系统组件：
  • Python版本、PyTorch、Transformers、TRL、Datasets版本完全匹配
  • GPU驱动、CUDA、NCCL版本统一（NCCL是多GPU通信核心，必须确保兼容性）
• 网络互通与免密登录：
  • 所有主机处于同一局域网，防火墙开放NCCL通信端口（默认无需手动指定，若有防火墙可临时关闭或放行相关端口）
  • 配置Host1到其他所有主机的SSH无密码登录，确保主节点能远程控制其他节点（或每个节点都能互相访问）
• 共享资源同步：
  • 训练脚本train.py需复制到所有主机的相同路径下
  • 模型权重（Llama-2-70b-hf）和数据集（IMDB）建议提前下载到共享存储（如NFS），或每个主机本地存储相同路径，避免重复下载浪费带宽

二、修改训练脚本 #

无需大幅改动原有代码，仅需补充模型和分词器的初始化细节，适配分布式训练：

from datasets import load_dataset
from trl import SFTTrainer
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载数据集（分布式训练会自动分片数据）
dataset = load_dataset("imdb", split="train")

# 加载模型与分词器
model = AutoModelForCausalLM.from_pretrained(
    "NousResearch/Llama-2-70b-hf",
    device_map="auto"  # 自动分配模型权重到可用GPU
)
tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-70b-hf")
tokenizer.pad_token = tokenizer.eos_token  # Llama默认无pad token，需手动指定

# 初始化SFTTrainer
trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    dataset_text_field="text",
    max_seq_length=512,
)

# 启动训练
trainer.train()

三、用Accelerate启动分布式训练（推荐） #

Accelerate是Hugging Face官方分布式训练工具，适配SFTTrainer的底层逻辑：

1. 生成配置文件：在Host1运行以下命令，按提示完成配置：

   accelerate config
   

   关键配置项：
   • 选择Distributed training
   • 输入节点总数4
   • 依次指定每个节点的GPU数量（Host1:1、Host2:1、Host3:2、Host4:2）
   • 指定主节点（Host1）的IP地址和通信端口（如29500）
2. 同步配置文件：将Host1生成的accelerate_config.yaml复制到其他所有主机的相同目录下
3. 启动各节点进程：
   • Host1（主节点，rank=0）：

     accelerate launch --config_file accelerate_config.yaml train.py
     

   • Host2（rank=1）：

     accelerate launch --config_file accelerate_config.yaml train.py
     

   • Host3（rank=2）：

     accelerate launch --config_file accelerate_config.yaml train.py
     

   • Host4（rank=3）：

     accelerate launch --config_file accelerate_config.yaml train.py
     

四、用Torchrun启动分布式训练（备选） #

若偏好原生PyTorch分布式工具，可直接用torchrun启动：

• Host1（主节点）：

  torchrun --nnodes=4 --nproc_per_node=1 --node_rank=0 --master_addr=<Host1的IP> --master_port=29500 train.py
  

• Host2：

  torchrun --nnodes=4 --nproc_per_node=1 --node_rank=1 --master_addr=<Host1的IP> --master_port=29500 train.py
  

• Host3：

  torchrun --nnodes=4 --nproc_per_node=2 --node_rank=2 --master_addr=<Host1的IP> --master_port=29500 train.py
  

• Host4：

  torchrun --nnodes=4 --nproc_per_node=2 --node_rank=3 --master_addr=<Host1的IP> --master_port=29500 train.py
  

关键注意事项 #

• 避免在Jupyter Notebook中直接启动分布式训练：多节点需独立启动进程，建议将代码保存为脚本，在各主机终端执行启动命令
• 调试通信问题：可设置环境变量export NCCL_DEBUG=INFO，查看NCCL通信日志定位网络或配置错误
• 模型权重加载：若模型过大，可开启load_in_8bit或load_in_4bit量化，减少单卡显存占用

内容的提问来源于stack exchange，提问作者user1564762