使用 Hugging Face Transformers 在 Amazon SageMaker 上部署 GPT-J 6B

Hugging Face 已详细介绍了一种使用 Amazon SageMaker 和 Hugging Face Inference Toolkit 部署 EleutherAI 的 GPT-J 6B（一个开源的 60 亿参数语言模型）进行实时推理的方法。此方法解决了模型加载时间过长的关键挑战，否则可能超过 Amazon SageMaker 施加的 60 秒请求限制。

优化生产环境中的模型加载时间

由于内存占用和加载速度，将 GPT-J 6B 部署到生产环境具有挑战性。60 亿参数模型的权重大约占用 24GB 内存。以 float32 加载模型至少需要 48GB 的 CPU RAM（一个用于初始权重，一个用于检查点）。

为了提高可访问性，EleutherAI 提供了 float16 权重。结合 transformers 中用于减少内存占用的选项，CPU RAM 需求降至大约 12.1GB。然而，标准加载方法仍然较慢：

标准加载： 在 P3.2xlarge AWS EC2 实例上加载模型大约需要 3 分 32 秒。
磁盘存储加载： 将模型存储在磁盘上可将此时间降至 1 分 23 秒。

由于 Amazon SageMaker 对请求响应有 60 秒的限制，这些加载时间使得模型不适合用于可扩展、可靠的生产工作负载。

使用 torch.save 加速加载时间

使用 torch.save(model, PATH) 和 torch.load(PATH) 而不是推荐的 from_pretrained 方法，可以将 GPT-J 的加载时间从 1 分 23 秒缩短至 7.7 秒——速度提升约 10.5 倍。

关键要求： 为了避免不兼容，用户必须在保存模型时使用的 PyTorch 和 Transformers 版本与加载模型时使用的版本保持一致。

在 Amazon SageMaker 上的部署工作流程

要将 GPT-J 6B 部署用于实时推理，可使用以下工作流程：

模型序列化： 使用 from_pretrained 加载 GPT-J，然后使用 torch.save() 将其保存为 .pt 文件。
制品创建： 创建一个包含模型权重和必要文件（如 tokenizer.json）的 model.tar.gz 归档。然后将此制品上传到 S3 存储桶。
端点部署： 使用 Amazon SageMaker Python SDK 中的 HuggingFaceModel 类部署模型。

对于此部署，Hugging Face 推荐使用 ml.g4dn.xlarge 实例类型（NVIDIA T4 GPU），每月费用约为 500 美元。

推理和使用最佳实践

端点部署后，可通过 predictor.predict 方法运行预测。Hugging Face Inference 工具包允许通过请求负载中的 parameters 属性自定义生成过程。

生成策略

贪婪搜索： 默认请求方法。第一次请求后，推理时间大约为 3 秒。
束搜索： 通过在参数中添加 `