DeepSpeed多机多卡微调实验记录

1、由于本地的显卡是Tesla-T4系列，显存只有16G，这里选用了一个比较小的模型DistilBERTHuggingface地址，为了方便后续使用docker将其容器化训练，我这里写了一个下载脚本：

python 复制代码

import os
from datasets import load_dataset
from transformers import DistilBertForSequenceClassification, DistilBertTokenizer

def download_model_and_data():
    # 创建存储目录
    os.makedirs("./local_model", exist_ok=True)
    os.makedirs("./local_data", exist_ok=True) 

    # 下载 tokenizer 和模型
    tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
    model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)
    tokenizer.save_pretrained("./local_model")
    model.save_pretrained("./local_model")

    # 下载 IMDb 数据集
    dataset = load_dataset("imdb")
    dataset.save_to_disk("./local_data")   

if __name__ == "__main__":
    download_model_and_data()

2、本地配置如下：三台服务器，节点分别是192.168.2.186、192.168.2.185、192.168.2.184，每台服务器都有两张显卡，cuda是12.4，Nvidia驱动是550.54.15：

3、三台服务器都需要进行ssh免密操作，不会的请参考文章：ssh免密教学

4、接下来是脚本的编写，本脚本主要使用DeepSpeed，具体脚本如下：

python 复制代码

import torch
import os
from datasets import load_from_disk
from transformers import DistilBertForSequenceClassification, DistilBertTokenizer, Trainer, TrainingArguments

import deepspeed

 
def tokenize_function(examples):
    """Tokenize text dataset"""
    return tokenizer(examples["text"], padding="max_length", truncation=True)

def main():
    # DeepSpeed 多机训练的分布式初始化
    torch.distributed.init_process_group(backend="nccl", init_method="env://")
    # 加载本地 tokenizer
    global tokenizer
    tokenizer = DistilBertTokenizer.from_pretrained("./local_model")
    # 加载本地数据集
    dataset = load_from_disk("./local_data")

    # 预处理数据集（tokenize）
    train_dataset = dataset["train"].map(tokenize_function, batched=True)

    # 转换为 PyTorch 格式
    train_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])

    # 加载本地模型
    model = DistilBertForSequenceClassification.from_pretrained("./local_model", num_labels=2)
  
    # 训练参数

    training_args = TrainingArguments(
        output_dir="./results",
        per_device_train_batch_size=8,  # 适当调整 batch size
        num_train_epochs=3,
        save_strategy="epoch",  # 仅在每个 epoch 结束时保存
        logging_dir="./logs",
        logging_steps=10,
        report_to="none",  # 禁用日志上传到 wandb 等
        deepspeed="ds_config.json",  # 启用 DeepSpeed
        do_eval=False,  # 禁用 eval
    )

    # 初始化 Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        tokenizer=tokenizer,  # 添加 tokenizer
    )

    # 训练
    trainer.train()

if __name__ == "__main__":
    main()

5、需要安装的python第三方库如下（建议创建conda虚拟环境)，如果不会请参考conda创建虚拟环境教学：

requirements.txt 复制代码

transformers
datasets
deepspeed
torch
mpi4py
torchvision
torchaudio
importlib_metadata

6、运行可以通过hostfile的形式，或者pdsh命令进行运行。（1)第一种hostflie编写：

hostfile 复制代码

192.168.2.186 slot=2
192.168.2.185 slot=2
192.168.2.184 slot=2

其中slot2代表我需要使用两个GPU，请根据实际情况进行修改。然后在三台服务器上运行命令：

bash 复制代码

torchrun --nnodes=3 --nproc_per_node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.2.186:29500 train.py

（2）第二种通过pdsh命令，但是使用之前请先参考pdsh的配置文章：pdsh配置教学，然后运行命令（命令仅为参考，具体IP，以及目录需要修改成自己的：

bash 复制代码

pdsh -w ^hostfile "cd /path/to/your/project && torchrun --nnodes=3 --nproc_per_node=2 --rdzv_backend=c10d --rdzv_endpoint=192.168.2.186:29500 train.py"