Connector for PyTorch

Cloud Storage 的 PyTorch 连接器是一款由 Google 支持的开源产品，可实现 Cloud Storage 与 PyTorch 的直接集成。

概览

Connector for PyTorch 在训练中的数据加载以及检查点和模型加载方面都具有优势：

对于训练中的加载数据，Connector for PyTorch 具有以下优势：

  • PyTorch 连接器包含优化，可在主要由小于 1MB 的文件组成的数据集中将训练速度提高多达 3 倍，相较于默认的 PyTorch。
  • PyTorch 连接器实现了 PyTorch 的数据集基元，可用于帮助高效地从 Cloud Storage 存储分区加载训练数据。
  • 支持针对随机数据访问模式的映射式数据集，以及针对流式数据访问模式的可迭代式数据集。
  • 能够将下载的原始字节数据转换为您选择的格式，从而让 PyTorch DataLoader 能够灵活地处理 NumPy 数组或 PyTorch 张量。

对于检查点和模型加载，Connector for PyTorch 具有以下优势：

  • 一个检查点接口，可让您轻松直接将模型检查点保存到 Cloud Storage 存储桶，并从存储桶加载模型检查点。
  • Connector for PyTorch 通过使用 PyTorch Lightning 的 CheckpointIO 的 DatafluxLightningCheckpoint 实现来支持 PyTorch Lightning 检查点。
  • Connector for PyTorch 提供了 StorageWriter 和 StorageReader 实现，可与 PyTorch 分布式检查点结合使用。Connector for PyTorch 演示库包含在 PyTorch Lightning FSDP 工作负载中使用此库的示例代码。
  • 连接器检查点功能支持使用 Lightning 和基本 PyTorch 进行异步检查点保存。

如需了解详情，请参阅 PyTorch Connector 的 GitHub 首页。

框架

Connector for PyTorch 支持以下框架版本：

  • Python 3.8 或更高版本
  • PyTorch Lightning 2.0 或更高版本
  • PyTorch 2.3.1 或更高版本

使用入门

如需使用 Connector for PyTorch，您必须具备以下条件：

  • 包含要处理的数据的 Cloud Storage 存储桶。
  • 用于处理存储在存储桶中的数据的以下权限：
    • storage.objects.create
    • storage.objects.list
    • storage.objects.get
    • storage.objects.delete，如果您打算使用组合下载

您必须使用 IAM 角色（例如 Storage Object User）向 PyTorch 连接器将用于身份验证的账号授予这些权限。

安装

如需安装 PyTorch 连接器，请使用以下命令：

pip install gcs-torch-dataflux

配置

必须提供身份验证，才能通过以下任一方法使用适用于 PyTorch 的连接器应用默认凭据：

  • 在 Compute Engine 虚拟机上运行 PyTorch 连接器时，应用默认凭据默认会自动使用虚拟机的附加服务账号。如需了解详情，请参阅选择工作负载身份验证方法。
  • 您也可以手动配置应用默认凭据。您可以使用 Google Cloud CLI 直接登录：
gcloud auth application-default login

示例

如需查看有关使用 Connector for PyTorch 的完整示例，请参阅 Connector for PyTorch GitHub 代码库的 demo 目录。例如：

性能

Connector for PyTorch 针对 ML 工作负载进行了专门优化，其性能比直接向 Cloud Storage 发出 API 调用要高得多：

  • 为了优化列表性能，适用于 PyTorch 的连接器利用了一种快速列表算法，该算法旨在平衡并行对象列表进程之间的列表工作负载。
  • 为了优化小文件的下载性能，PyTorch 连接器使用组合操作将一组较小的对象串联成单个较大的对象。这些新的复合对象与源对象存储在同一存储桶中，并且名称中带有前缀 dataflux-composed-objects/。
  • 与标准检查点上传相比，针对检查点写入的分块上传可将性能提升高达 10 倍。

您可以在 GitHub 上找到以下各项的性能数据：

  • 基于文字的 Lightning 培训
  • 闪电图片训练
  • 单节点检查点
  • 多节点检查点

注意事项

应根据每个工作负载考虑以下事项。

快速列出操作

Connector for PyTorch 的快速列表算法会导致 Connector for PyTorch 使用比常规顺序列表更多的列表操作。列表操作的计费方式为 A 类操作。

复合对象用法

为避免在处理临时复合对象时产生过多存储费用和提前删除费用，您应确保您的存储桶使用以下设置：

Connector for PyTorch 创建的复合对象通常会在训练循环结束时自动移除，但在极少数情况下，可能不会移除。如需确保从存储桶中移除对象，您可以运行以下命令：

gcloud storage rm gs://<my-bucket>/dataflux-composed-objects/ --recursive

您可以在构建的数据集的配置部分中添加 disable_compose=True 或 max_composite_object_size=0，以停用复合对象的使用。不过，关闭此行为可能会导致训练循环花费更长时间，尤其是在处理小文件时。

与直接下载文件相比，使用复合对象会导致 Cloud Storage 在较小规模下就达到 QPS 和吞吐量限制。在高多节点规模（即使不使用复合对象也会达到项目 QPS 或吞吐量限制）下运行时，您应停用复合对象。

429 错误和性能下降

在使用 Connector for PyTorch 时，您可能会收到 429 错误，或者执行时间比预期要长。造成这种情况的常见原因有：

  • 许多机器学习项目都选择利用 PyTorch Lightning 和 Ray 等工具采用高度分布式训练模型。这些模型与适用于 PyTorch 的连接器兼容，但通常会触发 Cloud Storage 的速率限制。
  • 如果 429 错误伴随“此工作负载从 Cloud Storage 提取了太多的出站带宽”或“此工作负载触发了 Cloud Storage 出站带宽上限”等消息，则表示工作负载的数据吞吐量速率超出了 Google Cloud 项目的最大容量。如需解决这些问题，请按以下步骤操作：
  • QPS 限制可能会触发 429 错误，并在正文消息中指明 TooManyRequests，但更常见的表现是执行时间比预期长。在处理大量小型文件时，QPS 瓶颈更为常见。令牌桶 QPS 限制会随着时间的推移而自然扩缩，因此允许预热期通常可以提高性能。如需详细了解目标存储桶的性能，请在 Google Cloud 控制台中查看存储桶时查看“可观测性”标签页。
  • 如果您的工作负载因 TooManyRequests 错误而失败，并且错误消息中包含关键字 dataflux-composed-objects，则最好先停用组合对象。这样做可以减少大规模使用时由 Compose 操作带来的 QPS 负载。

内存消耗

为了优化上传和下载性能，所有检查点写入和加载（包括用于推理的最终模型）都完全在内存中暂存。每台机器都必须有足够的可用 RAM 才能在内存中分阶段存储其检查点，以便充分利用这些性能改进。

获取支持

您可以使用 Google Cloud 的官方支持渠道获取支持、提交常见问题，以及请求新功能。您还可以通过在 GitHub 中提交问题来获取支持。

PyTorch、PyTorch 徽标及任何相关标志都是 The Linux Foundation 的商标。