AWS Trainium & Inferentia 文档
在 AWS Inferentia 上使用 llama-2-13B 创建您自己的聊天机器人
并获得增强的文档体验
开始使用
在 AWS Inferentia 上使用 llama-2-13B 创建您自己的聊天机器人
本教程有一个 notebook 版本，点击此处.
本指南将详细介绍如何在 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 不仅约束输入上下文的长度，还约束键/值缓存的长度，从而约束输出长度。
根据您选择的参数和 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 parametersLlama 2 13B 模型使用 float16 权重（存储在 2 个字节上）并具有 130 亿个参数，这意味着它至少需要 2 * 130 亿或约 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 库生成文本，正如 此帖子 中详细描述的那样。
如果您像建议的那样跳过了第一部分，请不要担心：我们将使用 hub 上已存在的预编译模型。
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 上使用它们的副作用。
</Warning>