从云存储中快速、准确地流式传输训练数据

[官网] - [快速入门] - [文档] - [招聘！]

👋 欢迎

我们构建了StreamingDataset，旨在使基于云存储的大型数据集训练尽可能快速、经济且可扩展。

它专为大型模型的多节点分布式训练而设计，最大限度地保证正确性、性能和易用性。现在，您可以不受训练数据位置的限制，在任何地方高效地进行训练。只需在需要时按需流式加载所需数据即可。要了解我们为何构建StreamingDataset，请阅读我们的公告博客。

StreamingDataset兼容任何数据类型，包括图像、文本、视频和多模态数据。

StreamingDataset支持主要的云存储提供商（AWS、OCI、GCS、Azure、Databricks，以及任何与S3兼容的对象存储，如Cloudflare R2、Coreweave、Backblaze b2等），并且被设计为PyTorch IterableDataset 类的即插即用替代品，能够无缝集成到您现有的训练工作流中。

🚀 快速开始

💾 安装

可以通过pip安装Streaming：

pip install mosaicml-streaming

🏁 快速入门

1. 准备您的数据

将原始数据集转换为我们的支持的流式格式之一：

• MDS（Mosaic Data Shard）格式，可编码和解码任何Python对象
• CSV / TSV
• JSONL

import numpy as np
from PIL import Image
from streaming import MDSWriter

# 用于存储压缩后输出文件的本地或远程目录
data_dir = '数据集路径'

# 将输入字段映射到其数据类型的字典
columns = {
    'image': 'jpeg',
    'class': 'int'
}

# 如果有分片压缩，则指定压缩算法
compression = 'zstd'

# 使用MDSWriter将样本保存为分片
with MDSWriter(out=data_dir, columns=columns, compression=compression) as out:
    for i in range(10000):
        sample = {
            'image': Image.fromarray(np.random.randint(0, 256, (32, 32, 3), np.uint8)),
            'class': np.random.randint(10),
        }
        out.write(sample)

2. 将您的数据上传到云存储

将您的流式数据集上传到您选择的云存储中（AWS、OCI 或 GCP）。以下是使用AWS CLI将目录上传到S3存储桶的一个示例。

$ aws s3 cp --recursive 数据集路径 s3://my-bucket/数据集路径

3. 构建StreamingDataset和DataLoader

from torch.utils.data import DataLoader
from streaming import StreamingDataset

# 存储完整数据集的远程路径
remote = 's3://my-bucket/数据集路径'

# 数据集在运行期间缓存的本地工作目录
local = '/tmp/数据集路径'

# 创建流式数据集
dataset = StreamingDataset(local=local, remote=remote, shuffle=True)

# 让我们看看第1337个样本里有什么...
sample = dataset[1337]
img = sample['image']
cls = sample['class']

# 创建PyTorch DataLoader
dataloader = DataLoader(dataset)

📚 下一步？

入门指南、示例、API 参考以及其他有用的信息，都可以在我们的文档中找到。

我们提供了针对以下数据集训练模型的端到端教程：

• CIFAR-10
• FaceSynthetics
• SyntheticNLP

此外，我们还为以下常用数据集提供了入门代码，这些代码可以在 streaming 目录中找到：

数据集	任务	读取	写入
LAION-400M	文本和图像	读取	写入
WebVid	文本和视频	读取	写入
C4	文本	读取	写入
EnWiki	文本	读取	写入
Pile	文本	读取	写入
ADE20K	图像分割	读取	写入
CIFAR10	图像分类	读取	写入
COCO	图像分类	读取	写入
ImageNet	图像分类	读取	写入

要开始使用这些数据集进行训练：

1. 使用 convert 目录中的相应脚本将原始数据转换为 .mds 格式。

例如：

$ python -m streaming.multimodal.convert.webvid --in <CSV文件> --out <MDS输出目录>

2. 导入数据集类以开始训练模型。

from streaming.multimodal import StreamingInsideWebVid
dataset = StreamingInsideWebVid(local=local, remote=remote, shuffle=True)

🔑 关键特性

---

无缝的数据混合

借助 Stream，您可以轻松地试验不同的数据集组合。数据采样既可以通过相对比例控制，也可以通过绝对数量（重复次数或样本数）来控制。在流式处理过程中，不同数据集会实时无缝地被加载、打乱并混合。

# 混合 C4、GitHub 代码库和内部数据集
streams = [
  Stream(remote='s3://datasets/c4', proportion=0.4),
  Stream(remote='s3://datasets/github', proportion=0.1),
  Stream(remote='gcs://datasets/my_internal', proportion=0.5),
]

dataset = StreamingDataset(
  streams=streams,
  samples_per_epoch=1e8,
)

真正的确定性

我们解决方案的一个独特优势：无论 GPU 数量、节点数或 CPU 工作进程如何，样本的顺序始终保持一致。这使得以下操作更加容易：

• 复现和调试训练过程及损失突增问题
• 加载在 64 个 GPU 上训练的检查点，并在 8 个 GPU 上进行可复现的调试

请参见下图——无论使用 1、8、16、32 或 64 个 GPU 训练模型，都会得到完全相同的损失曲线（受限于浮点运算精度）！

无需等待，即可在训练中途恢复

当硬件故障或损失突增导致数据加载器卡住时，等待作业恢复可能既昂贵又令人沮丧。得益于我们确定性的样本排序机制，StreamingDataset 允许您在长时间训练过程中随时快速恢复训练，而无需等待数小时。

与现有方案相比，这种低延迟的恢复方式可以节省数千美元的流量费用和空闲 GPU 计算时间。

高吞吐量

我们的 MDS 格式最大限度地减少了冗余工作，从而实现了极低的样本延迟和更高的吞吐量，特别适合那些受数据加载瓶颈限制的工作负载。

工具	吞吐量
StreamingDataset	~19000 张/秒
ImageFolder	~18000 张/秒
WebDataset	~16000 张/秒

以上结果来自 ImageNet + ResNet-50 的训练，数据在第一个 epoch 缓存后，经过 5 次重复实验得出。

平等的收敛性

由于我们采用了独特的洗牌算法，使用 StreamingDataset 进行模型训练的收敛效果与直接使用本地磁盘无异。

以下是 ImageNet + ResNet-50 训练的结果，共进行了 5 次重复实验。

工具	Top-1 准确率
StreamingDataset	76.51% ± 0.09
ImageFolder	76.57% ± 0.10
WebDataset	76.23% ± 0.17

StreamingDataset 会在分配给某个节点的所有样本中进行洗牌，而其他方案则仅在一个较小的范围内（单个进程内）对样本进行洗牌。更广泛的洗牌范围能够更好地分散相邻样本。此外，我们的洗牌算法还能最大限度地减少样本丢失。我们发现，这两种洗牌特性都有助于提升模型的收敛性能。

随机访问

在需要时访问所需的数据。

即使某个样本尚未下载，您仍然可以通过 dataset[i] 来获取第 i 个样本。系统会立即启动下载，并在完成后返回结果——这与按顺序编号且可任意访问的 PyTorch map-style 数据集类似。

dataset = StreamingDataset(...)
sample = dataset[19543]

无数据集大小可除性要求

StreamingDataset 可以愉快地遍历任意数量的样本。您无需永久删除部分样本，以使数据集大小能被预设的设备数量整除。相反，在每个 epoch 中，都会重复不同的样本子集（不丢弃任何样本），从而确保每台设备处理的样本数量相同。

dataset = StreamingDataset(...)
dl = DataLoader(dataset, num_workers=...)

磁盘使用限制

为了将磁盘使用量控制在指定的上限内，系统会动态删除最近最少使用的分片。此功能可通过设置 StreamingDataset 的 cache_limit 参数来启用。更多详细信息请参阅 洗牌 指南。

dataset = StreamingDataset(
    cache_limit='100gb',
    ...
)

🏆 项目展示

以下是一些使用了 StreamingDataset 的项目和实验。您也有想分享的吗？请发送邮件至 mcomm@databricks.com 或加入我们的 Community Slack。

• BioMedLM：由 MosaicML 和斯坦福 CRFM 共同推出的生物医学领域专用大型语言模型
• Mosaic 扩散模型：从零开始训练 Stable Diffusion 的成本低于 16 万美元
• Mosaic LLM：以不到 50 万美元的成本实现 GPT-3 水平的性能
• Mosaic ResNet：借助 Mosaic ResNet 和 Composer，实现超快速的计算机视觉训练
• Mosaic DeepLabv3：使用 MosaicML Recipes，图像分割训练速度提升 5 倍
• …更多内容即将发布！敬请期待！

💫 贡献者

我们欢迎任何形式的贡献、拉取请求或问题反馈。

如需开始贡献，请参阅我们的 贡献指南。

附注：我们正在招聘！

如果您喜欢这个项目，请为我们点亮星标 ⭐，并查看我们的其他项目：

• Composer —— 一个现代化的 PyTorch 库，让大规模、高效的神经网络训练变得简单易行
• MosaicML 示例 —— 快速且高精度训练机器学习模型的参考示例，包含 GPT / 大型语言模型、Stable Diffusion、BERT、ResNet-50 和 DeepLabV3 的入门代码
• MosaicML Cloud —— 我们的训练平台专为降低大型语言模型、扩散模型及其他大型模型的训练成本而设计，提供多云编排、轻松的多节点扩展以及底层优化，以加速训练过程

✍️ 引用

@misc{mosaicml2022streaming,
    author = {The Mosaic ML Team},
    title = {streaming},
    year = {2022},
    howpublished = {\url{<https://github.com/mosaicml/streaming/>}},
}

MosaicML Streaming 快速上手指南

MosaicML Streaming 是一个专为大规模分布式训练设计的开源工具，支持从云存储（如 AWS S3、阿里云 OSS 等兼容接口）高速流式加载训练数据。它兼容 PyTorch IterableDataset，支持图像、文本、视频及多模态数据，能够显著降低数据加载延迟并节省存储成本。

环境准备

• 系统要求：Linux 或 macOS（推荐 Linux 环境用于生产训练）。
• Python 版本：Python 3.8 及以上。
• 核心依赖：
  • PyTorch
  • Pillow (处理图像数据)
  • NumPy
• 云存储访问：若使用云端数据，需配置相应的云厂商 CLI 工具（如 awscli）或确保运行环境具有对应的访问凭证（Access Key/Secret Key）。

提示：国内开发者若访问 PyPI 较慢，建议使用清华或阿里镜像源进行安装。

安装步骤

使用 pip 直接安装官方发布包：

pip install mosaicml-streaming

推荐使用国内镜像源加速安装：

pip install mosaicml-streaming -i https://pypi.tuna.tsinghua.edu.cn/simple

如需安装额外依赖（如特定编解码器），可参考官方文档追加 extras，基础使用上述命令即可。

基本使用

使用流程分为三步：数据转换、上传云端（可选）、构建 DataLoader。

1. 准备数据 (转换为 MDS 格式)

首先将原始数据转换为 Streaming 支持的 MDS (Mosaic Data Shard) 格式。以下示例演示如何生成包含图像和标签的随机数据：

import numpy as np
from PIL import Image
from streaming import MDSWriter

# 本地或远程目录，用于存储压缩后的输出文件
data_dir = 'path-to-dataset'

# 定义字段及其数据类型
columns = {
    'image': 'jpeg',
    'class': 'int'
}

# 压缩算法 (可选，推荐 zstd)
compression = 'zstd'

# 使用 MDSWriter 写入分片数据
with MDSWriter(out=data_dir, columns=columns, compression=compression) as out:
    for i in range(10000):
        sample = {
            'image': Image.fromarray(np.random.randint(0, 256, (32, 32, 3), np.uint8)),
            'class': np.random.randint(10),
        }
        out.write(sample)

2. 上传数据到云存储 (可选)

如果需要在集群间共享或直接从云端训练，将生成的数据集目录上传至对象存储（如 AWS S3、阿里云 OSS、腾讯云 COS 等）。

以 AWS S3 为例：

aws s3 cp --recursive path-to-dataset s3://my-bucket/path-to-dataset

注意：国内云厂商（阿里云/腾讯云）通常提供 S3 兼容接口，只需配置对应的 endpoint 即可使用类似的上传逻辑。

3. 构建 StreamingDataset 和 DataLoader

在训练代码中，直接使用 StreamingDataset 替换原有的 Dataset 类。它会自动处理数据的下载、缓存和流式读取。

from torch.utils.data import DataLoader
from streaming import StreamingDataset

# 云端数据存储的完整路径 (远程)
remote = 's3://my-bucket/path-to-dataset'

# 本地缓存目录 (数据会被流式下载到此地)
local = '/tmp/path-to-dataset'

# 创建流式数据集
# shuffle=True 开启打乱，支持多卡确定性打乱
dataset = StreamingDataset(local=local, remote=remote, shuffle=True)

# 验证读取单个样本
sample = dataset[1337]
img = sample['image']
cls = sample['class']

# 创建标准的 PyTorch DataLoader
dataloader = DataLoader(dataset, batch_size=32, num_workers=4)

现在，dataloader 即可直接用于您的 PyTorch 训练循环中。该方案支持断点续训、多节点确定性采样以及无缝混合多个数据集。