一文搞懂大模型训练加速框架DeepSpeed的使用！

最新推荐文章于 2024-08-27 01:12:23 发布

雪碧没气阿

最新推荐文章于 2024-08-27 01:12:23 发布

阅读量2.5k

点赞数 9

文章标签：自然语言处理人工智能 ai大模型 DeepSheed 大模型训练 LLM 计算机技术

本文链接：https://blog.csdn.net/xxue345678/article/details/140145796

版权

在AI领域，常见的深度学习框架TensorFlow、PyTorch和Keras无疑是开发者们的得力工具，但随着模型规模的急剧膨胀，这些传统框架在应对大模型时往往会显得力不从心。

比如Pytorch的分布式并行计算框架DDP(Distributed Data Parallel)，尽管实现了数据并行，但是当模型大小超过单个GPU显存限制时显得捉襟见肘。此时，开发者往往只能手动进行复杂模型的参数拆分到各个GPU上，这无疑增加了研发的复杂性和门槛。

然而，微软推出的一款框架——DeepSpeed，可解决这一局限。它通过将模型参数拆散分布到各个GPU上，以实现大模型的计算。这也意味着，我们可以利用更少的硬件资源训练更大的模型，不再受限于单个GPU的显存限制。

安装DeepSpeed

pip install deepspeed

此外，还需要下载Pytorch，在官网选择自己对应的系统版本和环境，按照指示安装即可：

https://pytorch.org/get-started/locally/

使用DeepSpeed

载入数据集

在这里插入图片描述

创建模型

在这里插入图片描述

初始化Deepspeed

DeepSpeed通过输入参数来启动训练，因此需要使用argparse解析参数。

import argparse

def add_argument():
    # 创建一个ArgumentParser对象，设置描述为"CIFAR"    
    parser = argparse.ArgumentParser(description='CIFAR')    
    # 设置训练时的批大小，默认值为32    
    parser.add_argument('-b',
                        '--batch_size',                        
                        default=32,                        
                        type=int,                        
                        help='mini-batch size (default: 32)')    
    # 设置总的训练轮数，默认值为30    
    parser.add_argument('-e',       
                        '--epochs',                        
                        default=30,                        
                        type=int,                        
                        help='number of total epochs (default: 30)')    
    # 传递分布式训练中的排名，默认值为-1，表示未使用分布式训练    
    parser.add_argument('--local_rank',           
                        type=int,                        
                        default=-1,                        
                        help='local rank passed from distributed launcher')    
    # 设置输出日志信息的间隔，默认值为2000，即每2000次迭代打印一次日志    
    parser.add_argument('--log-interval',
                        type=int,                        
                        default=2000,                        
                        help="output logging information at a given interval")    
    # 添加与DeepSpeed相关的配置参数    
    parser = deepspeed.add_config_arguments(parser)   
    # 解析命令行参数，返回一个Namespace对象，其中包含了所有定义的参数及其对应的值    
    args = parser.parse_args()   
    # 返回解析后的参数对象args，供后续的训练脚本使用    
    return args

此外，模型初始化的时候除了参数，还需要model及其parameters，还有训练集：

# 启动DeepSpeed训练

# 调用之前定义的add_argument函数，解析命令行参数，并将结果存储在args变量中
args = add_argument()
# 创建Net类的实例
net = Net()
# 筛选出模型中需要梯度计算的参数
parameters = filter(lambda p: p.requires_grad, net.parameters())
# 使用deepspeed.initialize初始化模型引擎、优化器、数据加载器以及其他可能的组件
model_engine, optimizer, trainloader, __ = deepspeed.initialize(
    args=args, model=net, model_parameters=parameters, training_data=trainset)

训练

注意local_rank是不需要管的参数，在后面启动模型训练的时候，DeepSpeed会自动给这个参数赋值。

# 定义进行2个epoch的训练
for epoch in range(2):
    running_loss = 0.0    
    # 对于每个epoch，遍历训练数据加载器trainloader中的每一个小批量数据    
    # 同时提供索引i和数据data    
    for i, data in enumerate(trainloader):    
        # 将输入数据inputs和标签labels移动到当前GPU设备上，        
        # 具体是哪个GPU由model_engine.local_rank决定，        
        # 这对于分布式训练非常重要，确保数据被正确地分配到各个参与训练的GPU上        
        inputs, labels = data[0].to(model_engine.local_rank), data[1].to(  
            model_engine.local_rank)        
        # 通过model_engine执行前向传播，计算模型预测输出        
        outputs = model_engine(inputs)        
        # 计算预测输出outputs与真实标签labels之间的损失        
        loss = criterion(outputs, labels)        
        # 反向传播计算梯度        
        model_engine.backward(loss)        
        # 更新模型参数        
        model_engine.step()        
        
        # 计算并累加每个小批量的损失值        
        # 当达到args.log_interval指定的迭代次数时，打印平均损失值，        
        # 然后重置running_loss为0，以便计算下一个区间的平均损失        
        running_loss += loss.item()        
        if i % args.log_interval == (args.log_interval - 1):        
            print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / args.log_interval))            
            running_loss = 0.0

测试

模型测试和模型训练的逻辑类似：

在这里插入图片描述

编写模型参数

当前目录下新建一个config.json，写入调优器、训练batch等参数。

在这里插入图片描述

以上即为利用DeepSpeed开发模型的过程，由此可见，和Pytorch开发模型的过程大同小异，就是在初始化的时候使用了DeepSpeed，并以输入参数的形式初始化。

测试代码

首先，使用环境变量控制GPU，例如机器有10张GPU，但是只使用6, 7, 8, 9号GPU，输入命令：

export CUDA_VISIBLE_DEVICES="6,7,8,9"

然后开始运行代码：

deepspeed test.py --deepspeed_config config.json

如下图所示即为开始运行。

开始训练的时候DeepSpeed通常会打印更多的训练细节供用户监控，包括训练设置、性能统计和损失趋势，效果类似于：

这也说明第一个Deepspeed模型已完成，下来可以开始大规模训练之路了！

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述