Несколько строк кода научат вас легко выполнять рассуждения о сверхбольших моделях: LLaMA-30B+TITAN RTX*4+ускорение

Вы не можете с удовольствием играть с ChatGPT, потому что у вас нет ключа API или бесплатного токена? Хотите развернуть большие модели и поиграть с ними на ограниченных вычислительных ресурсах? ? Хотите разблокировать другие функции большой модели помимо диалога? ? ? Несколько строк кода научат вас развертывать очень большие модели и выполнять логические выводы в условиях ограниченных вычислительных ресурсов.

Подготовить

Очень большая языковая модель . OPT, GPT, LLaMA — все в порядке, если они с открытым исходным кодом. Зайдите в Hugging Face, чтобы найти понравившуюся модель, всегда найдется та, которая вам подходит. Я использую LLaMA-30B, нужно подготовить следующий пакет документов с официального сайта:

Соответствующая среда зависит от . Как стрелочник, само собой элементарные питорчи, трансформеры и т.д. В этот раз я представлю главного героя этого выпуска ** ускорение **! ! !
графические процессоры . TITAN RTX×4, то есть вычислительная мощность 4×24G.

# Наступи на яму , не наступи на яму ! ! ! Потому что я наступил на него.

начать прямо

Для начала кратко поясним принцип реализации акселератора: модель делится и хранится, GPU нельзя размещать на CPU, а CPU нельзя размещать в памяти. Затем используйте какую часть для загрузки какой части. Подробности смотрите во вступительном видеоролике на YouTube . Расчет того, какая часть должна быть загружена, где должен быть предварительно рассчитан. Но как рассчитать и выделить память без загрузки модели? Init_empty_weights новой версии pytorch поможет вам в этом. Что касается конкретного принципа , давайте посмотрим, что говорят другие: Я в основном отвечаю за то, что наступил на яму.
на коде.

from transformers import LlamaTokenizer, AutoConfig, AutoModelForCausalLM
import os
import torch
from accelerate import infer_auto_device_map, init_empty_weights, load_checkpoint_and_dispatch

os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'

# model_name = 'llama-7b'
model_name = 'llama-30b'

# 首先使用虚拟内存加载模型，这个内存被认为是无限大
# "/{}/".format(model_name)是你自己模型存放的路径，这个不会改的也别来部署大模型了
with init_empty_weights():
    model = AutoModelForCausalLM.from_pretrained("/{}/".format(model_name),
                                                 trust_remote_code=True, torch_dtype=torch.float16)

# The model weights are not tied.
# Please use the `tie_weights` method before using the `infer_auto_device` function.
model.tie_weights()
# 人为设置不同的device所分配的内存， devide:0分配得少一些方便后续推理，实际操作中不写也行
# max_memory={0: "16GiB", 1: "23GiB", 2: "23GiB", 3:"23GiB", "cpu":"80Gib"}  # cpu辛苦你啦
# device_map = infer_auto_device_map(model, max_memory=max_memory)    # 自动推断device_map
#
# print(device_map)  # 可以看看模型不同层被分在了哪里

model = load_checkpoint_and_dispatch(model, checkpoint="/{}/".format(model_name),
                    device_map='auto', offload_folder="offload", no_split_module_classes=["LlamaDecoderLayer"],
                    offload_state_dict=True, dtype=torch.float16).half()
        # half()强制数据转化，看名字应该是一个内存优化的数据类型

print("(*^_^*) model load finished!!!! please play with {}".format(model_name))


while True:
    text_in = input('please input your question: ')
    # 注意，使用Autokenizer会很慢，不知道为什么
    tokenizer = LlamaTokenizer.from_pretrained("/{}/".format(model_name))

    inputs = tokenizer(text_in, return_tensors="pt").to(0)

    outputs = model.generate(inputs["input_ids"], max_new_tokens=20)
    text_out = tokenizer.decode(outputs[0].tolist(), skip_special_tokens=True)

    print(text_out)

Примечание :

1. При использовании LLaMA не используйте Autokenizer, из-за которого токенизатор будет загружаться очень медленно.
2. При загрузке модели обязательно используйте: no_split_module_classes. Он указывает, какие части модели нельзя разделить между разными графическими процессорами. В частности, для части, содержащей остаток, не должно быть разделения! ! !

Наконец, видеопамять, занимаемая моделью 30B, выглядит следующим образом, что весьма немало: