利用 S3 Connector for Pytorch 在训练代码中实现 S3 的流式读取

摘要

在大模型的训练过程时，将源数据存储在 Amazon S3，并从 S3 下载至本地节点进行训练是一种常见的做法。然而，这种方法常常伴随着一系列挑战，例如数据集的下载时间过长造成 GPU 等待，每次训练任务都需要重复下载，为了实现数据的随机采样不得不在每个训练节点上下载完整数据集等问题。为了解决这些问题，训练数据流式读取变得至关重要。

在“利用 Mountpoint for S3 在 Kubernetes 上加速 LLM 的训练”一文中，我们讨论了在利用 MountPoint for S3 进行大模型训练的方法。基于 Mountpoint for Amazon S3 CSI driver，训练节点可以利用 CSI 挂载的方式，对 S3 实现类文件接口访问。而在“构建基于 Amazon SageMaker FastFileMode 的 LLM 流式训练”一文中我们探讨了利用 SageMaker FastFileMode 来实现流式训练。FastFileMode 是一个高效的解决方案，它可以显著减少数据处理时间，提升训练效率。

现在，在最近的亚马逊云科技 re:Invent 大会上发布的新功能——Amazon S3 Connector for PyTorch 为我们提供了一种利用在 PyTorch 代码中直接流式读取 S3 文件的能力。PyTorch 作为一个广受欢迎的大模型训练框架，Amazon S3 Connector for PyTorch 是亚马逊云科技开源通过在训练代码中引入 Amazon S3 Connector for PyTorch 包，实现简易的 S3 数据流式读取。本文将深入探讨 S3 Connector for PyTorch 的使用方法，旨在帮助用户快速理解并运用这一工具，以优化他们的模型训练流程。

S3 Connecotr fo PyTorch 介绍

PyTorch 是一个开源机器学习框架，亚马逊云科技客户广泛将其用于构建和训练机器学习模型。适用于 PyTorch 的 Amazon S3 Connector for PyTorch 会自动优化 S3 读取和列出请求，以改善训练工作负载的数据加载和检查点性能。Amazon S3 Connector for PyTorch 通过底层调用 Amazon Common Runtime（CRT）能够实现 40% 的训练加载效果。

适用于 PyTorch 的 Amazon S3 连接器提供了 PyTorch 数据集基元的新实现，您可以使用它来从 Amazon S3 加载训练数据。它既支持用于随机数据访问模式的地图样式数据集，也支持用于顺序数据访问模式的可迭代样式数据集。适用于 PyTorch 的 Amazon S3 连接器还包括一个检查点接口，用于将检查点直接保存和加载到 Amazon S3，而无需先保存到本地存储，也无需编写自定义代码上传到 Amazon S3。

利用 S3 Connecotr fo PyTorch 进行存储测试

下面是 S3 Connector for PyTorch的示例代码，我们将用这个示例进行 S3 访问测试。

import s3Torchconnector
def load_image(object):
    img = Image.open(object)
    return (object.key, Torchvision.transforms.functional.pil_to_tensor(img))

s3dataset = s3Torchconnector.S3MapDataset.from_prefix(<DATA_S3_URI>, 
                                                        region='us-east-1', 
                                                        transform=load_image
                                                       )

dataloader = Torch.utils.data.DataLoader(s3dataset, batch_size=<BATCH_SIZE>, num_workers=<N>)

for step, batch in enumerate(dataloader): 
    # batch processing
    # model training

从这里可以看到 S3 Connector for PyTorch 对于 S3 与 S3 Express One Zone 都能很好适配，代码不需要任何改动。代码的关键部分有：

1. 直接使用 S3MapDataset，基于给定的 S3 URI 构建一个与 Torch Dataset 使用范式一致的 dataset；
2. 使用该 dataset，构建 Torch 的标准 dataloader ，并通过 loop dataloader 进行标准的训练流程，如对当前 batch 的处理、模型训练等。

同时可以看到，在进行训练的实例上不无需要使用 Amazon boto3 SDK 进行数据下载，不需要首先落盘至本地存储的过程，实现按需加载。

我们准备一份数据集进行实际测试，10000 张，每张 280KB 图片，分别将其存储在 S3 Standard 和 S3 Express One Zone，并通过 S3 Connector for PyTorch 分别对其进行读取测试。本次测试采用 ml.g5.2xlarge，num_worker 设置为 4，Batch Size 设置为 128。

测试结果数据如下：

从结果数据可以看到，S3 Express One Zone 性能大约是 S3 Standard 的 3 倍以上，S3 Connector for PyTorch 能够通过 Dataloader multi-workers 很容易地利用高性能存储 S3 Express One Zone 实现高并发、高带宽访问。总结来说，通过在 Python 代码中设置 S3 Connector for PyTorch 就能很容易实现 S3 对象的流式访问，同时可以通过 PyTorch 中 Torch.utils.data.DataLoader 类，实现并发访问控制。

接下来，我们会利用 S3 Connector for PyTorch 进行实际模型训练。

利用 S3 Connecoter fo PyTorch 进行模型训练

我们选取 Lavis BLIP2 [https://blog.salesforceairesearch.com/blip-2/] 模型进行本次模型训练的演示，其他模型训练的 S3 Connector for PyTorch 使用类似。存储采用 S3 Express One Zone，训练实例类型为 SageMaker Training ml.p3dn.24xlarge，SageMaker ml.p4d.24xlarge。

关键代码如下：https://github.com/haozhx23/Lavis-Blip2-on-SageMaker。

LAVIS/lavis/datasets/datasets/caption_datasets.py，重写 CaptionDataset 类。

class CaptionDataset(BaseDataset, __DisplMixin):
    def __init__(self, vis_processor, text_processor, vis_root, ann_paths):
        super().__init__(vis_processor, text_processor, vis_root, ann_paths)

        self.fn_idx_map = _genS3FileMaps(<BUCKET>, <PREFIX>)
        self.s3dataset = s3Torchconnector.S3MapDataset.from_prefix(<DATA_S3_URI>, region=region)
        

    def __getitem__(self, index):
        ann = self.annotation[index]
        fidx = self.fn_idx_map[ann["image"]]
        imgobj = self.s3dataset[fidx]
        
        image = Image.open(imgobj).convert("RGB")
        
        image = self.vis_processor(image)
        caption = self.text_processor(ann["caption"])

        return {
            "image": image,
            "text_input": caption,
        }

由于训练数据中，Image 需要跟其对应的 Caption 配对，这里定义_genS3FileMaps 方法来为每张图片生成一个全局 ID，同时获得图片路径及 ID 的映射关系。

def _genS3FileMaps(bucket, prefix):
    s3 = boto3.client('s3')
    paginator = s3.get_paginator('list_objects_v2')
    pages = paginator.paginate(Bucket=bucket, Prefix=prefix)
    filename_list = list(set([o['Key'].replace(prefix,'') for p in pages for o in p['Contents']]))
    filename_list.sort()
    fn_idx_map = dict(zip(filename_list, range(len(filename_list))))
    return fn_idx_map

同时，重载 CaptionDataset 中的__getitem__ 方法，在数据的随机读取过程中，先通过 index 获取到一条标注（包含 Image Name 及对应的 Caption），后通过图片名称获取其全局 index 并使用 S3Dataset 直接进行加载。

def __getitem__(self, index):   
   ann = self.annotation[index]
    fidx = self.fn_idx_map[ann["image"]]
    imgobj = self.s3dataset[fidx]

运行 1000 Iterations。测试结果如下：

通过上面演示，一般在训练中，我们需要对数据集进行随机采样来实现模型更好的泛化，所以我们需要通过 S3 List 操作获取数据集文件列表，再重载 Dataset 类实现文件名到 Index 标注的对应关系，最后实现随机采样，完全不需要进行数据集的下载即可开启训练任务。

S3 Connector for S3 的特性总结

根据上面总结 S3 Connecoter for S3 特性，利用 S3 Connector for PyTorch 在训练过程中访问 S3/S3 Express One Zone 的优势在于：

1. 流式数据加载，无需提前下载 S3 数据到本地存储或 EBS 上，节省了 GPU 等待时间和昂贵的成本
2. 源代码开源，简单易用，直接对齐 PyTorch Dataset 使用范式
3. S3 存储的访问配置在训练脚本中，无需要关心训练的底层环境，训练脚本的可移植性更强
4. 不依赖具体训练平台，这种方式既可以在 SageMaker 的 training job 中使用，也可以 EKS 或 EC2 平台上训练使用
5. 模型开发人员不需要关心存储底层细节和掌握底层存储知识，不依赖底层存储的配置和操作（如 Linux mount/umount 或 S3 API）

在模型训练中，选择合适的存储方案

在模型训练中，存储往往成为瓶颈，由此产生的 GPU 浪费，我们认为是值得注意和优化的。在本文中以及之前博客中，我们探讨了不同的平台，不同场景下模型训练的存储优化方法。

• 利用 Mountpoint for S3 在 Kubernetes 上加速 LLM 的训练
• 构建基于 Amazon SageMaker FastFileMode 的 LLM 流式训练
• Speed up training on Amazon SageMaker using Amazon FSx for Lustre and Amazon EFS file systems

一般有如下建议：

• 从性能角度来说，FSx for Lustre 能够提供最低延迟（亚毫秒级）和高达数百 GB 带宽，是针对分布式、大规模训练，训练数据集在 TB 级甚至 PB 级场景最佳选择。
• S3 提供了最优的存储训练成本，往往用在单机或小集群训练、数据集较小（如 GB 或 TB 级）模型训练、多账号共享需求的训练场景、LLM 模型训练等，传统上采用先下载再训练的方式，但这增加了 GPU 等待时间和训练脚本的复杂度，所以我们推荐流式读取的方式。
  • 如果您在 SageMaker 上进行模型训练，并熟悉 SageMaker 训练方式，建议采用 FastFileMode 方法
  • 如果您基于 EC2（如 P4de/P5）或 EKS 进行训练，并习惯操作系统环境或 EKS CSI 配置，建议采用 Mountpoint for S3 方法
  • 如果您希望保持训练代码高可移植性，同时保持训练存储基础设施配置最小化（如 S3 API 操作），建议采用 S3 Connector for S3 方法

总结

本文探讨了利用 S3 Connector for PyTorch 直接从 Amazon S3 中流式读取训练数据的方法。测试结果显示，相比先下载再训练的方式，该方法可以简化训练代码，提高 GPU 利用率，降低存储成本。S3 Connector for PyTorch 为基于 PyTorch 的模型训练提供了高效的数据加载接口，可简化代码，优化性能。文中还归纳了模型训练中存储的几种选择，提出根据训练规模、数据集大小、代码可移植性等因素，在SageMaker FastFileMode、Mountpoint for S3 以及 S3 Connector 之间进行选择。

参考文献

• 利用 Mountpoint for S3 在 Kubernetes 上加速 LLM 的训练
• 构建基于 Amazon SageMaker FastFileMode 的 LLM 流式训练
• Speed up training on Amazon SageMaker using Amazon FSx for Lustre and Amazon EFS file systems
• Amazon S3 Connector for PyTorch
• BLIP-2: Scalable Pre-training of Multimodal Foundation Models for the World’s First Open-source Multimodal Chatbot
• Lavis-Blip2-on-SageMaker

本篇作者