如何用 Hugging Face Transformers 和 Amazon SageMaker 部署 GPT-J 6B 推理

差不多正好 6 个月前，EleutherAI 发布了 GPT-J 6B，这是 OpenAI 的 GPT-3 的开源替代方案。GPT-J 6B 是 EleutherAI GPT-NEO 系列（一个基于 GPT 架构、用于文本生成的 transformer 语言模型系列）的 60 亿参数继任者。

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

在过去的 6 个月里，GPT-J 引起了研究人员、数据科学家甚至软件开发者的广泛关注，但要将 GPT-J 部署到真实用例和产品的生产环境中仍然非常具有挑战性。

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

在这篇博文中，你将学习如何使用 Amazon SageMaker 和 Hugging Face Inference Toolkit 轻松部署 GPT-J，只需几行代码，即可使用常规大小的 NVIDIA T4 GPU 实例（约 500 美元/月）实现可扩展、可靠且安全的实时推理。

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

背景

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

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

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 minutes and 32 seconds（模型未存储在磁盘上）。如果将模型预先存储在磁盘上，这个时间可以缩短到 1 minute and 23 seconds，但对于需要考虑扩展性和可靠性的生产工作负载来说，这仍然非常长。

例如，Amazon SageMaker 对请求响应有 60 秒的限制，这意味着模型需要在 60 秒内加载完成并运行预测，在我看来，这对于保持模型/端点可扩展且可靠地服务于你的工作负载非常合理。如果你有更长的预测，可以使用 batch-transform。

在 Transformers 中，使用 from_pretrained 方法加载的模型遵循 PyTorch 的 推荐实践，对于 BERT 大约需要 1.97 seconds [REF]。PyTorch 提供了 另一种保存和加载模型的替代方式，使用 torch.save(model, PATH) 和 torch.load(PATH)。

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

这意味着，当我们使用 transformers==4.13.2 保存模型时，如果尝试用 transformers==4.15.0 加载，可能会不兼容。然而，以这种方式加载模型可将加载时间缩短 约 12 倍，对于 BERT 降至 0.166s。

将此应用于 GPT-J 意味着我们可以将加载时间从 1 minute and 23 seconds 缩短至 7.7 seconds，速度提升约 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 Inference 工具包通过请求负载的 parameters 属性提供相同的功能。下面你可以找到如何在不带参数、使用束搜索以及使用自定义配置的情况下生成文本的示例。如果你想了解不同的解码策略，请查看这篇 博客文章。

默认请求

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

首次请求后的推理时间：3s

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

束搜索请求

这是一个使用 beam 搜索且束数为 5 的请求示例。

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

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

参数化请求

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

首次请求后的推理时间：38s

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

少样本示例（高级）

这是一个示例，展示你可以如何 eos_token_id 来在某个 token 上停止生成，例如用于少样本预测的 \n、. 或 ###。下面是一个为关键词生成推文的少样本示例。

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

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: End-to-End Hugging Face Transformers with the Hub & SageMaker Pipelines”。

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