[deepspeed]deepspeed安装和测试代码

news2024/10/1 21:37:47

deepspeed官方对linux系统支持非常好，安装流程较为简单，推荐使用linux系统使用deepspeed.deepspeed由于要使用大模型进行训练和推理，建议显存>=24GB。windows上官方不直接支持，但是网上有安装whl文件，只能0.8.3这样老版本，由于deepspped需要aio等等，即使安装不能训练，只能推理部分模型。因此只能适合简单学习，本人使用deepspeed是0.12.7版本进行源码编译，测试发现deepspeed-mii无法正常使用，但是有些模型是可以用的，下面代码测试通过

import os
import deepspeed
import torch
from transformers import pipeline

local_rank = int(os.getenv('LOCAL_RANK', '0'))
world_size = int(os.getenv('WORLD_SIZE', '1'))
generator = pipeline('text-generation', model='EleutherAI/gpt-neo-2.7B',
                     device=local_rank)



generator.model = deepspeed.init_inference(generator.model,
                                           mp_size=world_size,
                                           dtype=torch.float,
                                           replace_with_kernel_inject=True)

string = generator("DeepSpeed is", do_sample=True, min_length=50)
if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
    print(string)

目测上面模型10GB，太大了，普通显卡根本无法玩起来