发布基于 C++ 的 S3 IO DataPipes

作者：John He, Khaled ElGalaind, Roshani Nagmote, Daiming Yang

训练大型深度学习模型需要大量数据集。Amazon Simple Storage Service (Amazon S3) 是一种可扩展的云对象存储服务，用于存储大型训练数据集。机器学习 (ML) 从业者需要一种高效的数据管道，能够从 Amazon S3 下载数据、转换数据，并将数据馈送到 GPU 以实现高吞吐量和低延迟的模型训练。

在本文中，我们介绍了 PyTorch 新的 S3 IO DataPipes，即 S3FileLister 和 S3FileLoader。为了提高内存效率和运行速度，新的 DataPipes 使用 C++ 扩展来访问 Amazon S3。基准测试表明，对于从 Amazon S3 下载自然语言处理 (NLP) 数据集，S3FileLoader 比 FSSpecFileOpener 快 59.8%。您可以使用新的 DataPipes 构建 IterDataPipe 训练管道。我们还演示了新的 DataPipe 可以将 Bert 和 ResNet50 的总体训练时间减少 7%。新的 DataPipes 已与 PyTorch 1.12.0 一起被合入到开源项目 TorchData 0.4.0 中。

概述

Amazon S3 是一种可扩展的云存储服务，对数据量没有限制。从 Amazon S3 加载数据并将其馈送到高性能 GPU（例如 NVIDIA A100）可能具有挑战性。这需要一个高效的数据管道，能够满足 GPU 的数据处理速度。为了解决这个问题，我们发布了 PyTorch 的一款新的高性能工具：S3 IO DataPipes。DataPipes 是从 torchdata.datapipes.iter.IterDataPipe 继承的子类，因此它们可以与 IterableDataPipe 接口交互。开发者可以快速构建其 DataPipe DAG 以使用 shuffle、sharding 和 batch 功能访问、转换和处理数据。

新的 DataPipes 被设计为与文件格式无关，并且 Amazon S3 数据以二进制大对象 (BLOB) 的形式下载。它可以作为可组合的构建块，用于组装 DataPipe 图，将表格数据、NLP 数据和计算机视觉 (CV) 数据加载到您的训练管道中。

在底层，新的 S3 IO DataPipes 使用 AWS C++ SDK 的 C++ S3 处理器。通常，与 Python 相比，C++ 实现更节省内存，并在多线程中具有更好的 CPU 核心利用率（无全局解释器锁）。建议将新的 C++ S3 IO DataPipes 用于训练大型深度学习模型时的高吞吐量、低延迟数据加载。

新的 S3 IO DataPipes 提供两个一级 API

S3FileLister – 可迭代对象，用于列出给定 S3 前缀内的 S3 文件 URL。此 API 的函数名称是 list_files_by_s3。
S3FileLoader – 可迭代对象，用于从给定 S3 前缀加载 S3 文件。此 API 的函数名称是 load_files_by_s3。

用法

在本节中，我们提供了使用新的 S3 IO DataPipes 的说明。我们还提供了 load_files_by_s3() 的代码片段。

从源代码构建

新的 S3 IO DataPipes 使用 C++ 扩展。默认情况下，它已构建到 torchdata 包中。但是，如果您的环境中无法使用新的 DataPipes（例如 Conda 上的 Windows），则需要从源代码构建。欲了解更多信息，请参阅 Iterable Datapipes。

配置

Amazon S3 支持全球 bucket。但是，bucket 是在特定区域内创建的。您可以通过使用 __init__() 将区域传递给 DataPipes。或者，您可以通过在 shell 中 export AWS_REGION=us-west-2 或在代码中设置环境变量 os.environ['AWS_REGION'] = 'us-east-1'。

要读取非公开访问的 bucket 中的对象，您必须通过以下方法之一提供 AWS 凭证：

使用 AWS configure 安装和配置 AWS 命令行界面 (AWS CLI)
在本地系统上的 AWS 凭证配置文件中设置凭证，该文件位于 Linux、macOS 或 Unix 上的 ~/.aws/credentials
设置 AWS_ACCESS_KEY_ID 和 AWS_SECRET_ACCESS_KEY 环境变量
如果您在 Amazon Elastic Compute Cloud (Amazon EC2) 实例上使用此库，请指定一个 AWS Identity and Access Management (IAM) 角色，然后授予 EC2 实例访问该角色的权限

示例代码

以下代码片段提供了 load_files_by_s3() 的典型用法：

from torch.utils.data import DataLoader 
from torchdata.datapipes.iter import IterableWrapper  

s3_shard_urls = IterableWrapper(["s3://bucket/prefix/",]) .list_files_by_s3()
s3_shards = s3_shard_urls.load_files_by_s3() 
# text data 
training_data = s3_shards.readlines(return_path=False) 
data_loader = DataLoader(
      training_data,
      batch_size=batch_size,
      num_workers=num_workers, 
) # training loop 
for epoch in range(epochs):     
      # training step     
      for bach_data in data_loader:         
         # forward pass, backward pass, model update  

基准测试

在本节中，我们演示了新的 DataPipe 如何减少 Bert 和 ResNet50 的总体训练时间。

独立 DataLoader 性能评估与 FSSpec 对比

FSSpecFileOpener 是另一个 PyTorch S3 DataPipe。它使用 botocore 和 aiohttp/asyncio 访问 S3 数据。以下是性能测试设置和结果（引自 Performance Comparison between native AWSSDK and FSSpec (boto3) based DataPipes）。

测试中的 S3 数据是分片文本数据集。每个分片约有 10 万行，每行约 1.6 KB，因此每个分片约 156 MB。此基准测试中的测量结果是 1000 个批次的平均值。未执行 shuffling、sampling 或 transforms。

下图报告了 num_workers=0 时不同批次大小的吞吐量比较，数据加载器在主进程中运行。S3FileLoader 具有更高的每秒查询次数 (QPS)。在批次大小为 512 时，它比 fsspec 高 90%。

Batch Sizes 1

下图报告了 num_workers=4 时的结果，数据加载器在主进程中运行。在批次大小为 512 时，S3FileLoader 比 fsspec 高 59.8%。

Batch Sizes 2

ResNet50 模型训练与 Boto3 对比

在下图中，我们在包含 4 个 p3.16xlarge 实例（共 32 个 GPU）的集群上训练了一个 ResNet50 模型。训练数据集为 ImageNet，包含 120 万张图片，组织成 1000 张图片的分片。训练批次大小为 64。训练时间以秒为单位测量。对于八个周期，S3FileLoader 比 Boto3 快 7.5%。

Boto3

Bert 模型训练与 Boto3 对比

在下图（原文本为 cart，根据上下文应为 chart）中，我们在包含 4 个 p3.16xlarge 实例（共 32 个 GPU）的集群上训练了一个 Bert 模型。训练语料库有 1474 个文件。每个文件约有 15 万个样本。为了缩短周期（epoch），我们每个文件使用 0.05%（约 75 个样本）。批次大小为 2048。训练时间以秒为单位测量。对于一个周期，S3FileLoader 比 Boto3 快 7%。

Boto3 2

与原始 PyTorch S3 插件对比

新的 PyTorch S3 DataPipes 的性能明显优于原始 PyTorch S3 插件。我们针对 S3FileLoader 调整了内部缓冲区大小。加载时间以秒为单位测量。

对于 10 个分片的 charades 文件（每个文件约 1.5 GiB），在我们的实验中，S3FileLoader 快了 3.5 倍。

最佳实践

训练大型深度学习模型可能需要一个包含数十甚至数百个节点的大规模计算集群。集群中的每个节点可能会产生大量针对特定 S3 分片的数据加载请求。为避免限流（throttle），我们建议将训练数据分片到不同的 S3 bucket 和 S3 文件夹中。

Best Practices

为了获得良好性能，文件的尺寸应足够大，以便可以在给定文件中进行并行处理，但又不能过大，以免根据训练任务击中 Amazon S3 上该对象的吞吐量限制。最佳尺寸可在 50-200 MB 之间。

结论和后续步骤

在本文中，我们向您介绍了新的 PyTorch IO DataPipes。新的 DataPipes 使用 aws-sdk-cpp，与基于 Boto3 的数据加载器相比，显示出更好的性能。

后续步骤中，我们计划通过关注以下功能来提高易用性、性能和功能性：

使用 IAM 角色进行 S3 授权 – 当前，S3 DataPipes 支持显式访问凭证、实例配置文件和 S3 bucket 策略。但是，在某些使用场景中，更倾向于使用 IAM 角色。
双缓冲区 – 我们计划提供双缓冲区以支持多 worker 下载。
本地缓存 – 我们计划使模型训练能够多次遍历训练数据集。在第一个周期后进行本地缓存可以消除从 Amazon S3 获取数据的时间延迟，从而大幅加快后续周期的检索时间。
可定制配置 – 我们计划公开更多参数，例如内部缓冲区大小、分块大小和执行器数量，并允许用户进一步调整数据加载效率。
Amazon S3 上传 – 我们计划扩展 S3 DataPipes 以支持 checkpointing 的上传功能。
与 fsspec 合并 – fsspec 在 torch.save() 等其他系统中使用。我们可以将新的 S3 DataPipes 与 fsspec 集成，以便它们有更多使用场景。

致谢

我们感谢 Amazon 的 Vijay Rajakumar 和 Kiuk Chung 为 S3 Common RunTime 和 PyTorch DataLoader 提供了指导。我们还要感谢来自 Meta AI/ML 的 Erjia Guan、Kevin Tse、Vitaly Fedyunin、Mark Saroufim、Hamid Shojanazeri、Matthias Reso 和 Geeta Chauhan，以及来自 AWS 的 Joe Evans 审阅了博文和 GitHub PR。

概述

用法