宣布基于 C++ 的 S3 IO DataPipes

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

在这篇文章中，我们将介绍 PyTorch 的新 S3 IO DataPipes：S3FileLister 和 S3FileLoader。为了内存效率和快速运行，新的 DataPipes 使用 C++ 扩展来访问 Amazon S3。基准测试表明，S3FileLoader 在从 Amazon S3 下载自然语言处理 (NLP) 数据集时比 FSSpecFileOpener 快 59.8%。您可以使用新的 DataPipes 构建 IterDataPipe 训练管道。我们还证明了新的 DataPipe 可以将 Bert 和 ResNet50 的总训练时间减少 7%。新的 DataPipes 已上传到开源的 TorchData 0.4.0，并与 PyTorch 1.12.0 兼容。

概述

Amazon S3 是一种可扩展的云存储服务，数据量没有限制。从 Amazon S3 加载数据并将其馈送到 NVIDIA A100 等高性能 GPU 可能具有挑战性。它需要一个高效的数据管道，能够满足 GPU 的数据处理速度。为了解决这个问题，我们为 PyTorch 发布了一个新的高性能工具：S3 IO DataPipes。DataPipes 是从 torchdata.datapipes.iter.IterDataPipe 子类化的，因此它们可以与 IterableDataPipe 接口交互。开发人员可以快速构建他们的 DataPipe DAG，以访问、转换和操作具有洗牌、分片和批量功能的数据。

新的 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 支持全局存储桶。但是，存储桶是在区域中创建的。您可以使用 __init__() 将区域传递给 DataPipes。或者，您可以在 shell 中 export AWS_REGION=us-west-2，或者在代码中使用 os.environ['AWS_REGION'] = 'us-east-1' 设置环境变量。

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

安装和配置 AWS 命令行界面 (AWS CLI) 并使用 AWS configure
在本地系统上的 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 的总训练时间。

针对 FSSpec 的独立 DataLoader 性能评估

FSSpecFileOpener 是另一个 PyTorch S3 DataPipe。它使用 botocore 和 aiohttp/asyncio 访问 S3 数据。以下是性能测试设置和结果（引自基于原生 AWSSDK 和 FSSpec (boto3) 的 DataPipes 之间的性能比较）。

测试中的 S3 数据是一个分片文本数据集。每个分片大约有 100,000 行，每行大约 1.6 KB，因此每个分片大约 156 MB。此基准测试中的测量值是 1,000 个批次上的平均值。没有执行洗牌、采样或转换。

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

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

针对 Boto3 训练 ResNet50 模型

对于以下图表，我们在由 4 个 p3.16xlarge 实例组成的集群上训练了一个 ResNet50 模型，总共有 32 个 GPU。训练数据集是 ImageNet，包含 120 万张图像，组织成 1,000 张图像分片。训练批处理大小为 64。训练时间以秒为单位测量。对于八个 epoch，S3FileLoader 比 Boto3 快 7.5%。

针对 Boto3 训练 Bert 模型

对于以下购物车，我们在由 4 个 p3.16xlarge 实例组成的集群上训练了一个 Bert 模型，总共有 32 个 GPU。训练语料库有 1474 个文件。每个文件大约有 150,000 个样本。为了运行更短的 epoch，我们每个文件使用 0.05%（大约 75 个样本）。批处理大小为 2,048。训练时间以秒为单位测量。对于一个 epoch，S3FileLoader 比 Boto3 快 7%。

与原始 PyTorch S3 插件的比较

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

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

最佳实践

训练大型深度学习模型可能需要一个包含数十甚至数百个节点的大型计算集群。集群中的每个节点可能会生成大量数据加载请求，这些请求会命中特定的 S3 分片。为了避免节流，我们建议在 S3 存储桶和 S3 文件夹之间分片训练数据。

为了获得良好的性能，文件大小应足够大以在给定文件上并行化，但又不能太大以至于根据训练作业达到 Amazon S3 上该对象的吞吐量限制。最佳大小可以在 50-200 MB 之间。

结论和下一步

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

接下来，我们计划通过关注以下功能来改进可用性、性能和功能

使用 IAM 角色进行 S3 授权 – 目前，S3 DataPipes 支持显式访问凭据、实例配置文件和 S3 存储桶策略。但是，有些用例更喜欢 IAM 角色。
双缓冲 – 我们计划提供双缓冲以支持多工作程序下载。
本地缓存 – 我们计划使模型训练能够对训练数据集进行多次遍历。在第一个 epoch 之后进行本地缓存可以消除 Amazon S3 的飞行时间延迟，这可以显着加快后续 epoch 的数据检索时间。
可定制配置 – 我们计划公开更多参数，例如内部缓冲区大小、多部分块大小和执行器数量，并允许用户进一步调整数据加载效率。
Amazon S3 上传 – 我们计划扩展 S3 DataPipes 以支持检查点上传。
与 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 的审阅。

概述

用法