当前位置：首页 > news >正文

多机多卡分布式训练

news 2025/7/19 19:35:03

1. 环境搭建

分布式训练框架：accelerate+deepspeed+pdsh(可有可无)
基础环境：cuda、显卡驱动、pytorch

1.1 安装相关包

cuda安装：参考官网安装步骤

wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda-repo-rhel7-11-8-local-11.8.0_520.61.05-1.x86_64.rpm
sudo rpm -i cuda-repo-rhel7-11-8-local-11.8.0_520.61.05-1.x86_64.rpm
sudo yum clean all
sudo yum -y install nvidia-driver-latest-dkms
sudo yum -y install cuda

显卡驱动安装：下载官网驱动包并安装
pytorch安装：参考官网安装指令

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

accelerate安装：参考huggingface官网

pip install accelerate

deepspeed安装：参考deepspeed github

pdsh安装：官网说明

可以参考教程：并行分布式运维工具pdsh-阿里云开发者社区

tar jxvf pdsh-2.29.tar.bz2
cd pdsh-2.29
./configure --with-ssh --with-rsh --with-mrsh --with-dshgroups --with-machines=/etc/pdsh/machines
make
make install
pdsh -V

注意：所有机器均需要安装一模一样的环境：版本需要一致；conda安装路径一致；同时cuda和pytorch版本相对应，如下图所示。

2. 启动分布式训练

2.1 使用accelerate

# 1、生成accelerate配置文件，使用命令行生成
accelerate config

2.1 启动分布式训练脚本

方式一：使用pdsh，仅需要在主节点启动

# accelerate语法
accelerate launch --config_file <accelerate配置文件> python_script.py <python脚本环境变量参数># 示例如下
config_path=/data0/sdmt/mxm/kohya_ss/my_util/config/deepspeed_pdsh_config.yamlaccelerate launch --config_file $config_path \
train_text_to_image_sdxl.py --mixed_precision fp16 --enable_xformers_memory_efficient_attention --gradient_checkpointing  --noise_offset 0.05  --cache_dir "/data0/sdmt/mxm/datasets/" --num_train_epochs 20 --resolution 1024 --proportion_empty_prompts 0.2 --learning_rate 1e-06 --lr_scheduler "constant" --lr_warmup_steps 0 --validation_prompt "a pair of casual leather shoes" --validation_epochs 5  --pretrained_model_name_or_path "stabilityai/stable-diffusion-xl-base-1.0" --pretrained_vae_model_name_or_path "madebyollin/sdxl-vae-fp16-fix"  --train_data_dir "/data0/sdmt/train_img/000/10_train_1024_hug"

如下所示，启动2台服务器，服务器每台一张显卡。