Press "Enter" to skip to content

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

将近6个月前的这一天,EleutherAI发布了GPT-J 6B,这是OpenAI GPT-3的一个开源替代品。GPT-J 6B是EleutherAI的GPT-NEO系列的继任者,该系列是基于GPT架构的基于Transformer的语言模型,用于文本生成。

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

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

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

在本博客文章中,您将学习如何使用Amazon SageMaker和Hugging Face推理工具包轻松部署GPT-J,只需几行代码即可进行可扩展、可靠和安全的实时推理,使用常规大小的GPU实例与NVIDIA T4(约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分钟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倍。

图1. BERT和GPTJ的模型加载时间

教程

通过这种保存和加载模型的方法,我们实现了与生产场景兼容的模型加载性能。但是我们需要记住需要对齐以下内容:

在使用torch.save(model,PATH)保存模型和使用torch.load(PATH)加载模型时,需要保持PyTorch和Transformers版本的一致,以避免不兼容。

使用torch.save保存GPT-J

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

from transformers import AutoTokenizer,GPTJForCausalLM
import torch

# 加载fp16模型
model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", revision="float16", torch_dtype=torch.float16)
# 保存模型
torch.save(model, "gptj.pt")

现在我们可以使用torch.load()加载我们的模型进行预测。

from transformers import pipeline
import torch

# 加载模型
model = torch.load("gptj.pt")
# 加载分词器
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")

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

# 进行预测
gen("我的名字是菲利普")

为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部署到Amazon SageMaker。

有关如何使用它们,请参见“将部署为Amazon SageMaker端点”。

如果您仍然希望或需要创建自己的model.tar.gz,例如出于合规指南的原因,您可以使用辅助脚本convert_gpt.py来完成此目的,该脚本将创建model.tar.gz并将其上传到S3。

# 克隆目录
git clone https://github.com/philschmid/amazon-sagemaker-gpt-j-sample.git

# 切换到amazon-sagemaker-gpt-j-sample目录
cd amazon-sagemaker-gpt-j-sample

# 创建并上传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

将部署为Amazon SageMaker端点

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

下面的代码片段使用get_execution_role,该函数仅在Amazon SageMaker笔记本实例或Studio中可用。如果您想在外部部署模型,请查看文档。

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

from sagemaker.huggingface import HuggingFaceModel
import sagemaker

# 具有创建端点权限的 IAM 角色
role = sagemaker.get_execution_role()

# GPT-J 模型文件的公开 S3 URI
model_uri="s3://huggingface-sagemaker-models/transformers/4.12.3/pytorch/1.9.1/gpt-j/model.tar.gz"

# 创建 Hugging Face 模型类
huggingface_model = HuggingFaceModel(
    model_data=model_uri,
    transformers_version='4.12.3',
    pytorch_version='1.9.1',
    py_version='py38',
    role=role, 
)

# 部署模型到 SageMaker 推理
predictor = huggingface_model.deploy(
    initial_instance_count=1, # 实例数量
    instance_type='ml.g4dn.xlarge' #'ml.p3.2xlarge' # EC2 实例类型
)

如果您想使用自己的model.tar.gz文件,只需将model_uri替换为您的 S3 Uri。

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

运行预测

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

predictor.predict({
    "inputs": "请告诉我们更多关于您的细节"
})

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

使用最佳实践

在使用生成模型时,大多数情况下您希望配置或自定义预测以适应您的需求,例如使用波束搜索、配置生成序列的最大或最小长度,或者调整温度以减少重复。Transformers 库提供了不同的策略和kwargs来实现这一点,Hugging Face 推理工具包使用请求有效载荷的parameters属性提供了相同的功能。下面是一些使用不同参数进行文本生成的示例。如果您想了解不同的解码策略,请查看这篇博文。

默认请求

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

首次请求后的推理时间:3秒

predictor.predict({
    "inputs": "请告诉我们更多关于您的细节"
})

波束搜索请求

这是一个使用波束搜索和5个波束的请求示例。

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

predictor.predict({
    "inputs": "请告诉我们更多关于您的细节",
  "parameters" : {
    "num_beams": 5,
  }
})

参数化请求

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

首次请求后的推理时间:38秒

predictor.predict({
    "inputs": "请告诉我们更多关于您的细节",
  "parameters" : {
    "max_length": 512,
    "temperature": 0.9,
  }
})

Few-Shot 示例(高级)

这是一个示例,演示如何使用eos_token_id来在某个特定标记(例如\n.###)处停止生成,用于进行少样本预测。以下是为关键词生成推文的少样本示例。

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

从transformers库中导入AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")

end_sequence="###"
temperature=4
max_generated_token_length=25
prompt= """key: markets
tweet: 从自然和市场中获取反馈,而不是从人们那里获取。
###
key: children
tweet: 也许我们死了,才能重生为孩子。
###
key: startups
tweet: 创业公司不应该担心如何扑灭火焰,而应该担心如何引燃它们。
###
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部署了由EleutherAI创建的6亿参数语言模型GPT-J。我们将模型加载时间从3.5分钟缩短到8秒,以便能够运行可扩展、可靠的推理。

请记住,使用torch.save()和torch.load()可能会导致不兼容问题。如果您想了解更多关于扩展Amazon SageMaker端点的内容,请查看我其他的博客文章:“MLOps: 使用Hub和SageMaker Pipelines进行端到端的Hugging Face Transformers”。


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

Leave a Reply

Your email address will not be published. Required fields are marked *