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"
相关推荐
橘猫云计算机设计36 分钟前
基于Springboot的自习室预约系统的设计与实现(源码+lw+部署文档+讲解),源码可白嫖!
java·spring boot·后端·毕业设计
秋书一叶1 小时前
SpringBoot项目打包为window安装包
java·spring boot·后端
pwzs2 小时前
Spring MVC 执行流程全解析:从请求到响应的七步走
java·后端·spring·spring mvc
小兵张健2 小时前
互联网必备职场知识(4)—— 共情沟通能力
后端·产品经理·运营
AskHarries3 小时前
使用 acme.sh 自动更新 SSL 证书的指南
后端
Chandler243 小时前
Go:反射
开发语言·后端·golang
pwzs3 小时前
深入浅出 MVCC:MySQL 并发背后的多版本世界
数据库·后端·mysql
盒子69103 小时前
go for 闭环问题【踩坑记录】
开发语言·后端·golang
刘大猫264 小时前
Arthas monitor(方法执行监控)
人工智能·后端·监控
追逐时光者4 小时前
MongoDB从入门到实战之MongoDB简介
后端·mongodb