在 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，我们指定模型要部署在多少个核心上（每个神经元设备有两个核心），以及精度（此处为 float16），
使用 input_shapes，我们设置模型的静态输入和输出维度。所有模型编译器都需要静态形状，Neuron 也不例外。请注意，sequence_length 不仅限制了输入上下文的长度，还限制了键/值缓存的长度，因此也限制了输出长度。

根据您选择的参数和 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 上使用它们的副作用。

</警告>

AWS Trainium 和 Inferentia

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

1. 将 Llama 2 模型导出到 Neuron

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

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

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