Amazon SageMaker 与 Hugging Face 的合作伙伴关系

TL;DR

Hugging Face 与 Amazon 建立了战略合作伙伴关系，通过将 Hugging Face Transformers 集成到 Amazon SageMaker 中，简化了最先进 NLP 模型的训练和部署。此次合作引入了专门的 Hugging Face 深度学习容器 (DLCs) 和 SageMaker Python SDK 的一流扩展，将设置和运行实验的时间从几天缩短至几分钟。

Hugging Face 深度学习容器 (DLCs)

Amazon SageMaker 现在提供 Hugging Face DLCs，这些是预装了深度学习框架和库（如 transformers、datasets 和 tokenizers）的 Docker 镜像。这些容器消除了用户从头开始构建和优化环境的需求。

DLCs 的主要特征包括：

  • 框架支持：提供针对 TensorFlow 和 PyTorch 的优化变体。
  • 基础设施灵活性：支持单 GPU、单节点多 GPU 以及多节点集群。
  • 性能：内置针对 PyTorch 和 TensorFlow 的优化，以提高训练速度并优化性价比。
  • 开源：DLCs 采用 Apache 2.0 许可。

SageMaker Python SDK 集成

为了加速从数据科学到生产环境的过渡，Hugging Face 开发了 SageMaker Python SDK 的一流扩展。这种集成允许用户创建 HuggingFace Estimator 来管理端到端的训练过程。

核心能力

  • 自动模型调优：与 SageMaker 的自动模型调优集成，以优化超参数并提高模型准确性。
  • 实验追踪：与 SageMaker Studio IDE 兼容，用于追踪和比较训练产物和实验。
  • 管理式基础设施：SageMaker 会自动管理 Amazon EC2 实例，上传训练脚本（例如 train.py），并处理来自 Amazon S3 的数据 I/O。

高级训练功能

此次合作伙伴关系引入了多种高性能训练能力，以处理大规模 NLP 工作负载。

分布式训练

用户可以通过定义分布参数，直接通过 HuggingFace Estimator 利用 SageMaker 分布式训练库：

  • 数据并行：集成到 Trainer API 中，以简化在大规模数据集上的训练。
  • 模型并行：通过将模型划分到多个 GPU/节点上，支持训练具有数十亿参数的模型。

成本与工作流优化

  • 管理式 Spot 实例：用户可以使用 EC2 spot 实例来降低高达 90% 的训练成本。对于长时间运行的任务，建议结合使用 checkpoint_s3_uri。
  • Git 仓库集成：Estimator 的 entry_point 和 source_dir 可以直接指向一个 GitHub 仓库，从而无需在本地下载脚本。
  • SageMaker Metrics：系统可以使用正则表达式自动解析日志，并将训练指标（例如 eval_loss、eval_accuracy）发送到 Amazon CloudWatch。

部署与推理

虽然最初的重点是训练，但该集成在 2021 年 7 月 8 日进行了扩展，以包含在 Amazon SageMaker 中对 Transformers 模型进行简便的部署和推理。

优势总结

根据 Hugging Face FAQ，使用 SageMaker 运行 Hugging Face 模型具有三个主要优势：

  1. 成本效益：实例按秒计费，并在任务结束时自动终止。
  2. MLOps 自动化：训练元数据、日志以及用于检查点和产物的 S3 I/O 是完全托管的。
  3. 可扩展性：支持启动多个并行训练任务和大规模分布式训练。

Sources