AWS Trainium & Inferentia 文档

在 AWS Inferentia 上使用 Llama-2-13B 创建你自己的聊天机器人

Hugging Face's logo
加入 Hugging Face 社区

并获得增强文档体验

开始

在 AWS Inferentia 上使用 Llama-2-13B 创建你自己的聊天机器人

该教程有一个笔记本版本 这里.

本指南将详细介绍如何在 AWS Inferentia 上导出、部署和运行 **Llama-2 13B** 聊天模型。

您将学习如何

  • 将 Llama-2 模型导出为 Neuron 格式,
  • 将导出的模型推送到 Hugging Face Hub,
  • 部署模型并在聊天应用程序中使用它。

注意:本教程是在 inf2.48xlarge AWS EC2 实例上创建的。

1. 将 Llama 2 模型导出到 Neuron

在本指南中,我们将使用非门控的 NousResearch/Llama-2-13b-chat-hf 模型,其功能等效于原始的 meta-llama/Llama-2-13b-chat-hf

此模型是 **Llama 2** 模型系列的一部分,并且经过微调以识别 *用户* 和 *助手* 之间的聊天交互(稍后将详细介绍)。

optimum-neuron 文档 中所述,在 Neuron 设备上运行模型之前,需要将其编译并导出为序列化格式。

导出模型时,我们将指定两组参数

  • 使用 *compiler_args*,我们指定希望模型部署在多少个核心上(每个 Neuron 设备有两个核心),以及使用哪种精度(此处为 *float16*),
  • 使用 *input_shapes*,我们设置模型的静态输入和输出维度。所有模型编译器都需要静态形状,Neuron 也不例外。请注意,*sequence_length* 不仅限制输入上下文的长度,还限制 Key/Value 缓存的长度,因此也限制输出长度。

根据您选择的参数和 Inferentia 主机,这可能需要几分钟到一个多小时。

为了方便起见,我们在 Hugging Face Hub 上托管了该模型的预编译版本,因此您可以跳过导出并立即在第 2 节中开始使用该模型。

from optimum.neuron import NeuronModelForCausalLM

compiler_args = {"num_cores": 24, "auto_cast_type": 'fp16'}
input_shapes = {"batch_size": 1, "sequence_length": 2048}
model = NeuronModelForCausalLM.from_pretrained(
        "NousResearch/Llama-2-13b-chat-hf",
        export=True,
        **compiler_args,
        **input_shapes)

这可能需要一段时间。

幸运的是,您只需要执行一次此操作,因为您可以保存模型并在以后重新加载它。

model.save_pretrained("llama-2-13b-chat-neuron")

更好的是,您可以将其推送到 Hugging Face Hub

为此,您需要登录到 HuggingFace 帐户

在终端中,只需键入以下命令并在请求时粘贴您的 Hugging Face 令牌

huggingface-cli login

默认情况下,模型将上传到您的帐户(组织等于您的用户名)。

如果您想将模型上传到特定的 Hugging Face 组织,请随时编辑下面的代码。

from huggingface_hub import whoami

org = whoami()['name']

repo_id = f"{org}/llama-2-13b-chat-neuron"

model.push_to_hub("llama-2-13b-chat-neuron", repository_id=repo_id)

关于导出参数的更多说明。

加载模型所需的最小内存可以通过以下方式计算

   memory = bytes per parameter * number of parameters

**Llama 2 13B** 模型使用float16权重(存储在 2 个字节中)并具有 130 亿个参数,这意味着至少需要 2 * 13B 或约 26GB 的内存来存储其权重。

每个 NeuronCore 具有 16GB 的内存,这意味着 26GB 的模型无法放入单个 NeuronCore 中。

实际上,由于缓存注意力层投影(KV 缓存),所需的总空间远大于参数数量。这种缓存机制使内存分配随序列长度和批次大小线性增长。

在这里,我们将batch_size设置为 1,这意味着我们只能并行处理一个输入提示。我们将sequence_length设置为 2048,这对应于模型最大容量(4096)的一半。

评估 KV 缓存大小的公式更为复杂,因为它还取决于与模型架构相关的参数,例如嵌入的宽度和解码器块的数量。

底线是,为了使非常大的语言模型适应,使用张量并行来跨多个 NeuronCore 分割权重、数据和计算,同时记住每个核心上的内存不能超过 16GB。

请注意,将核心数量增加到超出最低要求几乎总是会导致模型运行更快。增加张量并行度可以提高内存带宽,从而提高模型性能。

为了优化性能,建议使用实例上所有可用的核心。

在本指南中,我们使用了inf2.48xlarge的所有 24 个核心,但如果您使用的是inf2.24xlarge实例,则应将其更改为 12。

2. 在 AWS Inferentia2 上使用 Llama 2 生成文本

导出模型后,您可以使用 transformers 库生成文本,如这篇文章中详细描述

如果您按照建议跳过了第一部分,请不要担心:我们将使用中心上已有的预编译模型。

from optimum.neuron import NeuronModelForCausalLM

try:
    model
except NameError:
    # Edit this to use another base model
    model = NeuronModelForCausalLM.from_pretrained('aws-neuron/Llama-2-13b-chat-hf-neuron-latency')

我们需要一个Llama 2分词器将提示字符串转换为文本标记。

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("NousResearch/Llama-2-13b-chat-hf")

支持以下生成策略

  • 贪婪搜索,
  • 具有 top-k 和 top-p(带温度)的多项式采样。

大多数 logits 预处理/过滤器(如重复惩罚)都受支持。

inputs = tokenizer("What is deep-learning ?", return_tensors="pt")
outputs = model.generate(**inputs,
                         max_new_tokens=128,
                         do_sample=True,
                         temperature=0.9,
                         top_k=50,
                         top_p=0.9)
tokenizer.batch_decode(outputs, skip_special_tokens=True)

3. 在 AWS Inferentia2 上使用 Llama 创建聊天应用程序

我们特别选择了Llama 2聊天变体来说明导出模型在编码上下文长度增长时的出色表现。

模型期望提示按照特定模板格式化,该模板对应于用户角色和助手角色之间的交互。

每个聊天模型都有自己的编码此类内容的约定,在本指南中我们不会详细介绍,因为我们将直接使用与我们的模型相对应的Hugging Face 聊天模板

下面的实用程序函数将用户和模型之间的一系列交换转换为格式良好的聊天提示。

def format_chat_prompt(message, history, max_tokens):
    """ Convert a history of messages to a chat prompt


    Args:
        message(str): the new user message.
        history (List[str]): the list of user messages and assistant responses.
        max_tokens (int): the maximum number of input tokens accepted by the model.

    Returns:
        a `str` prompt.
    """
    chat = []
    # Convert all messages in history to chat interactions
    for interaction in history:
        chat.append({"role": "user", "content" : interaction[0]})
        chat.append({"role": "assistant", "content" : interaction[1]})
    # Add the new message
    chat.append({"role": "user", "content" : message})
    # Generate the prompt, verifying that we don't go beyond the maximum number of tokens
    for i in range(0, len(chat), 2):
        # Generate candidate prompt with the last n-i entries
        prompt = tokenizer.apply_chat_template(chat[i:], tokenize=False)
        # Tokenize to check if we're over the limit
        tokens = tokenizer(prompt)
        if len(tokens.input_ids) <= max_tokens:
            # We're good, stop here
            return prompt
    # We shall never reach this line
    raise SystemError

现在,我们具备了构建简单聊天应用程序的能力。

我们只需将用户和助手之间的交互存储在一个列表中,然后使用该列表生成输入提示。

history = []
max_tokens = 1024

def chat(message, history, max_tokens):
    prompt = format_chat_prompt(message, history, max_tokens)
    # Uncomment the line below to see what the formatted prompt looks like
    #print(prompt)
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs,
                             max_length=2048,
                             do_sample=True,
                             temperature=0.9,
                             top_k=50,
                             repetition_penalty=1.2)
    # Do not include the input tokens
    outputs = outputs[0, inputs.input_ids.size(-1):]
    response = tokenizer.decode(outputs, skip_special_tokens=True)
    history.append([message, response])
    return response

要测试聊天应用程序,您可以例如使用以下提示序列

print(chat("My favorite color is blue. My favorite fruit is strawberry.", history, max_tokens))
print(chat("Name a fruit that is on my favorite colour.", history, max_tokens))
print(chat("What is the colour of my favorite fruit ?", history, max_tokens))
<警告>

虽然功能强大,但大型语言模型有时会产生幻觉。我们将幻觉称为模型生成的与主题无关或虚构的内容,但将其呈现为准确的内容。这是 LLM 的一个缺陷,而不是在 Trainium/Inferentia 上使用它们产生的副作用。

</警告>