Connector for PyTorch

Cloud Storage 的 Connector for PyTorch 是 Google 支持的开源产品，可实现 Cloud Storage 与 PyTorch 的直接集成。

概览

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

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

  • Connector for PyTorch 包含优化，可在主要由小于 1MB 的文件组成的数据集中将训练速度提高至默认 PyTorch 的 3 倍。
  • Connector for 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 进行异步检查点保存。

如需了解详情，请参阅 Connector for PyTorch 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）将这些权限授予 Connector for PyTorch 将用于身份验证的账号。

安装

如需安装 Connector for PyTorch，请使用以下命令：

pip install gcs-torch-dataflux

配置

必须提供身份验证，才能通过以下某种方法使用 Connector for PyTorch 的应用默认凭据：

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

示例

您可以在 Connector for PyTorch GitHub 代码库的演示目录中找到一整套有关如何使用 Connector for PyTorch 的示例。例如：

性能

Connector for PyTorch 针对机器学习工作负载进行了特定优化，性能相较直接向 Cloud Storage 发出 API 调用得到显著提升：

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

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

  • Lightning 文本训练
  • 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 等工具的高度分布式训练模型。这些模型与 Connector for PyTorch 兼容，但通常会触发 Cloud Storage 的速率限制。
  • 如果 429 错误伴随以下消息，例如“此工作负载占用了 Cloud Storage 中过多的出站带宽”或“此工作负载触发了 Cloud Storage 出站带宽上限”，则表明工作负载的数据吞吐率超过了 Google Cloud 项目的最大容量。要解决这些问题，请按以下步骤操作：
  • QPS 限制可能会触发 429 错误，并在正文消息中指明 TooManyRequests，但更常见的情况是执行时间比预期更长。在处理大量小文件时，QPS 瓶颈更为常见。存储桶 QPS 上限会随着时间的推移而自然调整，因此设定一个预热期通常可以提高速度。如需详细了解目标存储桶的性能，请在 Google Cloud 控制台中查看存储桶时，查看“可观测性”标签页。
  • 如果您的工作负载失败，并且错误消息中存在包含关键字 dataflux-composed-objects 的 TooManyRequests 错误，那么停用组合对象是最佳的初步问题排查步骤。这样做可以减少大规模使用组合操作时所带来的 QPS 负载。

内存消耗

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

获取支持

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

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