如何在单个 4GB 的 GPU 上运行 70B LLMs

低代码演说家

已于 2024-08-06 11:38:09 修改

阅读量609

点赞数 13

文章标签： llama python

于 2024-08-06 11:00:00 首次发布

本文链接：https://blog.csdn.net/2301_80441165/article/details/140906429

版权

你是否曾经梦想过使用最先进的大型语言模型 (LLMs) 来执行自然语言处理 (NLP) 任务，但因高内存要求而感到沮丧？如果是这样，你可能对 AirLLM 感兴趣，这是一个优化推理内存使用的 Python 包，允许 70B LLMs 在单个 4GB GPU 上运行推理。不需要会导致模型性能下降的量化、蒸馏、修剪或其他模型压缩技术。

什么是 AirLLM？它是如何运作的？

大型语言模型 (LLMs) 的计算成本很高，并且需要大量内存来训练和运行。原因是 LLMs 具有大量层 - 70B 模型可以具有超过 80 层。然而，在推理过程中，语言模型中的每一层都是独立的，仅依赖于前一层的输出。因此，没有必要将所有层保留在 GPU 内存中。相反，我们可以在执行该层时仅从磁盘加载必要的层，进行所有计算，然后完全释放内存。这样，单个层所需的 GPU 内存仅约为该 Transformer 层的参数大小，即完整模型的 1/80，即约 2GB。

AirLLM 背后的主要思想确实是将原始的 LLM 分割成更小的子模型，每个子模型包含一层或几层，并在推理过程中按需加载它们。这样，在任何给定时间只有必要的子模型保留在内存中，其余的则存储在磁盘上。它还应用分块量化来进一步压缩子模型，从而减少磁盘加载时间和内存使用量。

AirLLM 支持 Hugging Face 开放 LLM 排行榜中的大多数顶级模型，例如 Platypus2、LLaMa2、Mistral、Mixtral、SOLAR、StellarBright 等。

如何使用AirLLM？

使用 AirLLM 非常简单且直观。你只需安装airllm pip 包，然后使用AutoModel 类从Hugging Face hub 或本地路径加载你选择的LLM。然后，你可以使用生成方法执行类似于常规变压器模型的推理。例如，以下代码片段展示了如何使用AirLLM加载和使用Platypus2–70B-instruct模型，该模型可以回答自然语言问题并遵循指令。

pip install airllm

from airllm import AutoModel

MAX_LENGTH = 128

# load the model from the Hugging Face hub
model = AutoModel.from_pretrained("garage-bAInd/Platypus2-70B-instruct")

# or load the model from a local path
# model = AutoModel.from_pretrained("/home/ubuntu/.cache/huggingface/hub/models--garage-bAInd--Platypus2-70B-instruct/snapshots/b585e74bcaae02e52665d9ac6d23f4d0dbc81a0f")

# prepare the input text
input_text = [
    'What is the capital of United States?',
]

# tokenize the input text
input_tokens = model.tokenizer(input_text,
    return_tensors="pt",
    return_attention_mask=False,
    truncation=True,
    max_length=MAX_LENGTH,
    padding=False)

# generate the output text
generation_output = model.generate(
    input_tokens['input_ids'].cuda(),
    max_new_tokens=20,
    use_cache=True,
    return_dict_in_generate=True)

# decode the output text
output = model.tokenizer.decode(generation_output.sequences[0])

# print the output text
print(output)

该代码片段的输出是：

What is the capital of United States?
The capital of the United States is Washington, D.C.

请注意，在第一次推理过程中，AirLLM 会逐层分解并保存原始 LLM，因此请确保有足够的磁盘空间。之后，AirLLM 将按需加载子模型，并以更少的内存更快地执行推理。

使用AirLLM有哪些优势？

通过使用AirLLM，有以下优势：

访问最先进的LLMs：你可以使用最先进的LLMs来完成你的NLP任务，例如问答、文本生成、文本摘要、文本分类、和更多。你可以从适合你的需求和偏好的各种模型中进行选择，例如特定领域、多语言或指令调整的模型。
内存要求低：你无需担心内存不足错误或昂贵的云计算资源。你可以在单个 4GB GPU 上运行推理，甚至可以在 CPU 或 Mac 设备上运行推理。
简单直观的使用：你可以使用 AirLLM 作为常规变压器模型的直接替代品，只需最少的代码更改。