使用 Hugging Face Transformers 和 Amazon SageMaker 部署 GPT-J 6B 进行推理

距今大约 6 个月前，EleutherAI 发布了 GPT-J 6B，这是 OpenAI GPT-3 的一个开源替代品。GPT-J 6B 是 EleutherAI 的 GPT-NEO 系列的 60 亿参数后继者，该系列是基于 GPT 架构用于文本生成的 transformer 语言模型家族。

EleutherAI 的主要目标是训练一个与 GPT-3 规模相当的模型，并以开放许可证向公众提供。

在过去的 6 个月里，`GPT-J` 引起了研究人员、数据科学家甚至软件开发人员的极大兴趣，但将其部署到生产环境中用于实际用例和产品仍然非常具有挑战性。

有一些托管解决方案可将 `GPT-J` 用于生产工作负载，例如 Hugging Face 推理 API，或者使用 EleutherAI 的 6b playground 进行实验，但关于如何轻松地将其部署到您自己环境中的示例较少。

在这篇博文中，您将学习如何使用 Amazon SageMaker 和 Hugging Face 推理工具包，通过几行代码轻松部署 `GPT-J`，从而使用常规大小的 GPU 实例（NVIDIA T4，约 500 美元/月）实现可扩展、可靠和安全的实时推理。

但在我们深入探讨之前，我想解释一下为什么将 `GPT-J` 部署到生产环境中具有挑战性。

---

背景

这个 60 亿参数模型的权重占用约 24GB 内存。要以 float32 格式加载它，至少需要 2 倍模型大小的 CPU 内存：1 倍用于初始权重，另 1 倍用于加载检查点。因此，对于 `GPT-J`，仅加载模型就需要至少 48GB 的 CPU 内存。

为了让模型更易于使用，EleutherAI 也提供了 float16 权重，并且 `transformers` 有了新的选项来减少加载大型语言模型时的内存占用。综合所有这些，加载模型大约需要 12.1GB 的 CPU 内存。

from transformers import GPTJForCausalLM
import torch

model = GPTJForCausalLM.from_pretrained(
    "EleutherAI/gpt-j-6B",
        revision="float16",
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True
)

这个例子的一个问题是，模型加载到内存并准备好使用需要很长时间。在我的实验中，在 `P3.2xlarge` AWS EC2 实例上使用上面的代码片段加载模型需要 `3 分 32 秒`（模型未存储在磁盘上）。通过将模型预先存储在磁盘上，可以缩短这个时间，加载时间减少到 `1 分 23 秒`，这对于需要考虑可扩展性和可靠性的生产工作负载来说仍然很长。

例如，Amazon SageMaker 对请求响应有 60 秒的限制，这意味着模型加载和预测运行必须在 60 秒内完成，在我看来，这对于保持模型/端点的可扩展性和工作负载的可靠性非常有意义。如果您有更长的预测任务，可以使用批量转换。

在 Transformers 中，使用 `from_pretrained` 方法加载的模型遵循 PyTorch 的推荐做法，对于 BERT 来说大约需要 `1.97 秒` [参考]。PyTorch 提供了另一种保存和加载模型的方法，即使用 `torch.save(model, PATH)` 和 `torch.load(PATH)`。

“用这种方式保存模型将使用 Python 的 pickle 模块保存整个模块。这种方法的缺点是，序列化的数据与保存模型时使用的特定类和确切的目录结构绑定在一起。”

这意味着，当我们使用 `transformers==4.13.2` 保存模型时，在尝试使用 `transformers==4.15.0` 加载时可能会不兼容。然而，以这种方式加载模型可将加载时间减少约 **12 倍**，对于 BERT 来说，时间降至 `0.166 秒`。

将此方法应用于 `GPT-J` 意味着我们可以将加载时间从 `1 分 23 秒` 减少到 `7.7 秒`，速度提高了约 10.5 倍。

教程

通过这种保存和加载模型的方法，我们为 `GPT-J` 实现了与生产场景兼容的模型加载性能。但我们需要记住，我们需要对齐

在使用 `torch.save(model,PATH)` 保存模型和使用 `torch.load(PATH)` 加载模型时，对齐 PyTorch 和 Transformers 的版本，以避免不兼容性。

使用 `torch.save` 保存 `GPT-J`

为了创建我们与 `torch.load()` 兼容的模型文件，我们使用 Transformers 和 `from_pretrained` 方法加载 `GPT-J`，然后使用 `torch.save()` 保存它。

from transformers import AutoTokenizer,GPTJForCausalLM
import torch

# load fp 16 model
model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", revision="float16", torch_dtype=torch.float16)
# save model with torch.save
torch.save(model, "gptj.pt")

现在我们能够使用 `torch.load()` 加载我们的 `GPT-J` 模型来运行预测。

from transformers import pipeline
import torch

# load model
model = torch.load("gptj.pt")
# load tokenizer
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")

# create pipeline
gen = pipeline("text-generation",model=model,tokenizer=tokenizer,device=0)

# run prediction
gen("My Name is philipp")
#[{'generated_text': 'My Name is philipp k. and I live just outside of Detroit....

---

为 Amazon SageMaker 实时端点创建 `model.tar.gz`

既然我们可以快速加载模型并进行推理，让我们把它部署到 Amazon SageMaker 上。

有两种方式可以部署 transformers 到 Amazon SageMaker。您可以直接“从 Hugging Face Hub 部署模型”，或者“使用存储在 S3 上的 `model_data` 部署模型”。由于我们不使用默认的 Transformers 方法，我们需要选择第二种方案，并使用存储在 S3 上的模型来部署我们的端点。

为此，我们需要创建一个 `model.tar.gz` 文件，其中包含我们的模型权重和推理所需的其他文件，例如 `tokenizer.json`。

我们提供了已上传且可公开访问的 `model.tar.gz` 文件，这些文件可以与 `HuggingFaceModel` 一起使用，以将 `GPT-J` 部署到 Amazon SageMaker。

请参阅“将 `GPT-J` 部署为 Amazon SageMaker 端点”部分了解如何使用它们。

如果您仍然想或需要创建自己的 `model.tar.gz`，例如出于合规性准则，您可以使用辅助脚本 convert_gpt.py 来实现此目的，该脚本会创建 `model.tar.gz` 并将其上传到 S3。

# clone directory
git clone https://github.com/philschmid/amazon-sagemaker-gpt-j-sample.git

# change directory to amazon-sagemaker-gpt-j-sample
cd amazon-sagemaker-gpt-j-sample

# create and upload model.tar.gz
pip3 install -r requirements.txt
python3 convert_gptj.py --bucket_name {model_storage}

`convert_gpt.py` 应该会打印出一个类似于此的 S3 URI。`s3://hf-sagemaker-inference/gpt-j/model.tar.gz`。

将 `GPT-J` 部署为 Amazon SageMaker 端点

为了部署我们的 Amazon SageMaker 端点，我们将使用 Amazon SageMaker Python SDK 和 `HuggingFaceModel` 类。

下面的代码片段使用了 `get_execution_role`，它仅在 Amazon SageMaker Notebook 实例或 Studio 中可用。如果您想在这些环境之外部署模型，请查看文档。

`model_uri` 定义了我们的 `GPT-J` 模型文件的位置。我们将使用由我们提供的公开可用的文件。

from sagemaker.huggingface import HuggingFaceModel
import sagemaker

# IAM role with permissions to create endpoint
role = sagemaker.get_execution_role()

# public S3 URI to gpt-j artifact
model_uri="s3://huggingface-sagemaker-models/transformers/4.12.3/pytorch/1.9.1/gpt-j/model.tar.gz"

# create Hugging Face Model Class
huggingface_model = HuggingFaceModel(
    model_data=model_uri,
    transformers_version='4.12.3',
    pytorch_version='1.9.1',
    py_version='py38',
    role=role, 
)

# deploy model to SageMaker Inference
predictor = huggingface_model.deploy(
    initial_instance_count=1, # number of instances
    instance_type='ml.g4dn.xlarge' #'ml.p3.2xlarge' # ec2 instance type
)

如果您想使用自己的 `model.tar.gz`，只需将 `model_uri` 替换为您的 S3 URI 即可。

部署过程大约需要 3-5 分钟。

运行预测

我们可以使用由 `.deploy` 方法创建的 `predictor` 实例来运行预测。为了向我们的端点发送请求，我们使用 `predictor.predict` 和我们的 `inputs`。

predictor.predict({
    "inputs": "Can you please let us know more details about your "
})

如果您想使用额外的 `kwargs` (如 `min_length`) 来自定义您的预测，请查看下面的“使用最佳实践”部分。

使用最佳实践

在使用生成模型时，大多数情况下您希望配置或自定义您的预测以满足您的需求，例如使用束搜索、配置生成序列的最大或最小长度，或调整温度以减少重复。Transformers 库提供了不同的策略和 `kwargs` 来实现这一点，Hugging Face 推理工具包通过请求负载中的 `parameters` 属性提供了相同的功能。下面您可以找到如何生成不带参数的文本、使用束搜索以及使用自定义配置的示例。如果您想了解不同的解码策略，请查看这篇博客文章。

默认请求

这是一个使用 `greedy` 搜索的默认请求示例。

第一次请求后的推理时间：`3 秒`

predictor.predict({
    "inputs": "Can you please let us know more details about your "
})

束搜索请求

这是一个使用 `beam` 搜索（5 个束）的请求示例。

首次请求后的推理时间：`3.3 秒`

predictor.predict({
    "inputs": "Can you please let us know more details about your ",
  "parameters" : {
    "num_beams": 5,
  }
})

参数化请求

这是一个使用自定义参数的请求示例，例如使用 `min_length` 生成至少 512 个 token。

第一次请求后的推理时间：`38 秒`

predictor.predict({
    "inputs": "Can you please let us know more details about your ",
  "parameters" : {
    "max_length": 512,
    "temperature": 0.9,
  }
})

少样本示例（高级）

这是一个如何使用 `eos_token_id` 在特定标记处停止生成的示例，例如 `\n`、`.` 或 `###` 用于少样本预测。下面是一个为关键词生成推文的少样本示例。

首次请求后的推理时间：`15-45 秒`

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")

end_sequence="###"
temperature=4
max_generated_token_length=25
prompt= """key: markets
tweet: Take feedback from nature and markets, not from people.
###
key: children
tweet: Maybe we die so we can come back as children.
###
key: startups
tweet: Startups shouldn’t worry about how to put out fires, they should worry about how to start them.
###
key: hugging face
tweet:"""

predictor.predict({
    'inputs': prompt,
  "parameters" : {
    "max_length": int(len(prompt) + max_generated_token_length),
    "temperature": float(temperature),
    "eos_token_id": int(tokenizer.convert_tokens_to_ids(end_sequence)),
    "return_full_text":False
  }
})

---

要删除您的端点，您可以运行。

predictor.delete_endpoint()

结论

我们成功地使用 Amazon SageMaker 部署了 `GPT-J`，这是一个由 EleutherAI 创建的 60 亿参数语言模型。我们将模型加载时间从 3.5 分钟减少到 8 秒，以便能够运行可扩展、可靠的推理。

请记住，使用 `torch.save()` 和 `torch.load()` 可能会产生不兼容问题。如果您想了解更多关于扩展 Amazon SageMaker 端点的信息，请查看我的另一篇博客文章：“MLOps：使用 Hub 和 SageMaker Pipelines 实现端到端的 Hugging Face Transformers”。

---

感谢阅读！如果您有任何问题，请随时通过 Github 或在 论坛上与我联系。您也可以在 Twitter 或 LinkedIn 上与我联系。