为模型训练注入超级动力

面向规模化训练的深度学习框架

[官网] - [快速开始] - [文档] - [我们正在招聘！]

👋 欢迎

Composer 是 MosaicML 开发的开源深度学习训练库。它构建于 PyTorch 之上，旨在简化大规模集群上的分布式训练工作流。

我们构建 Composer 时，特别针对可扩展性和易用性进行了优化，集成了高效多节点训练的最佳实践。通过抽象掉并行技术、分布式数据加载和内存优化等底层复杂性，您可以专注于训练现代 ML 模型和运行实验，而不会降低速度。

如果您正在训练任何规模的神经网络，我们推荐使用 Composer 来加速您的实验工作流，包括：

• 大语言模型 (LLMs)
• 扩散模型
• 嵌入模型 (例如 BERT)
• 基于 Transformer 的模型
• 卷积神经网络 (CNNs)

MosaicML 研究团队广泛使用 Composer 来训练像 MPT 这样的前沿模型，我们将此库开源，以使整个 ML 社区也能做到这一点。该框架被科技行业和学术界的组织所使用，并持续更新新功能、错误修复和生产工作负载的稳定性改进。

🔑 核心特性

我们为现代深度学习工作负载从头设计了 Composer。AlexNet 和 ResNet 时代已经过去，那时最先进的模型可以在几台桌面 GPU 上训练。如今，开发最新、最强大的深度学习模型通常需要集群规模的硬件——但在 Composer 的帮助下，您几乎不会感觉到差异。

Composer 的核心是我们的 Trainer 抽象：一个高度优化的 PyTorch 训练循环，旨在让您和您的模型都能更快地迭代。我们的训练器提供了简单的方法来配置并行化方案、数据加载器、指标、记录器等。

可扩展性

无论您是在 1 个 GPU 还是 512 个 GPU 上训练，处理 50MB 还是 10TB 的数据，Composer 都能让您的工作流保持简单。

• FSDP：对于过大而无法放入 GPU 的大型模型，Composer 已将 PyTorch 的 FullyShardedDataParallelism 集成到我们的训练器中，并简化了高效并行化自定义模型的过程。我们发现 FSDP 在性能上可与更复杂的并行策略相媲美。此外，Composer 也支持标准的 PyTorch 分布式数据并行 (DDP) 执行。
• 弹性分片检查点：在八个 GPU 上保存，在十六个 GPU 上恢复。Composer 支持弹性分片检查点，因此您永远不必担心分片保存的状态是否与新的硬件设置兼容。
• 数据流式传输：处理大型数据集？通过与 MosaicML StreamingDataset 集成，在模型训练期间动态从云存储下载数据集。

可定制性

其他高级深度学习训练器以牺牲灵活性为代价提供简单性。当您想添加自己的功能时，它们的抽象会成为障碍。而 Composer 则提供了简单的方法，让您可以根据需要定制我们的训练器。

图 1： Composer 的训练循环在训练过程的每个阶段都有一系列事件发生。回调是用户在特定事件时编写的函数。例如，我们的 学习率监控回调 会在每个 BATCH_END 事件记录学习率。

• 回调：Composer 的回调系统允许您在训练循环的任何位置插入自定义逻辑。我们已经编写了用于监控内存使用、记录和可视化图像以及估计模型剩余训练时间等功能的回调。这个功能深受希望实现和实验自定义训练技术的研究人员欢迎。
• 加速算法：我们从最新研究中汲取灵感，创建了一系列算法加速方法。将这些加速方法组合成 MosaicML 配方，可以提升您的训练速度。我们的团队已经开源了针对不同类型模型的最佳加速组合。
  • 8 倍加速：Stable Diffusion
    • 原始 SD2 成本 $200k —> $50k (博客)
  • 7 倍加速：ImageNet 上的 ResNet-50
    • 8xA100 上从 3h33m —> 25m (博客)
  • 8.8 倍加速：BERT-Base 预训练
    • 8xA100 上从 10h —> 1.13h (博客)
  • 5.4 倍加速：ADE20K 上的 DeepLab v3
    • 8xA100 上从 3h30m —> 39m (博客)

更好的工作流

Composer 旨在自动化处理底层的痛点，让您可以专注于深度学习重要（且有趣）的部分，并更快地迭代。

• 自动恢复：训练运行失败了？别担心——只需重新运行您的代码，Composer 将自动从最新的保存检查点恢复。
• CUDA OOM 预防：告别内存不足错误。将您的微批次大小设置为“auto”，Composer 将自动选择适合您 GPU 的最大批次大小。
• 时间抽象： 是否曾搞混更新步数、周期数、样本数和词元数之间的转换？使用我们的 Time 类，在训练循环中用自定义单位（周期、批次、样本、词元）指定训练时长。

集成

与您熟悉和喜爱的实验跟踪和数据流工具集成。

• 云集成：我们的检查点和日志记录功能对远程存储和从云存储桶（OCI、GCP、AWS S3）加载提供一流支持。
• 实验跟踪：Weights and Biases、MLFlow、CometML 和 neptune.ai——选择权在您，轻松将数据记录到您喜欢的平台。

🚀 快速开始

📍 先决条件

Composer 面向熟悉 Python 并对深度学习基础和 PyTorch 有基本了解的用户。

软件要求： 需要较新版本的 PyTorch。

硬件要求： 需要支持 CUDA 的 GPU 系统（AMD + RoCM 即将支持！）。Composer 可以在 CPU 上运行，但要获得全部优势，我们建议在硬件加速器上使用。

💾 安装

可以通过 pip 安装 Composer：

pip install mosaicml

为了简化 Composer 的环境设置，我们还提供了一套预构建的 Docker 镜像。我们强烈推荐您使用我们的 Docker 镜像。

🏁 快速入门

以下是在 MNIST 数据集上使用我们的训练器的代码片段。

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

from composer import Trainer
from composer.models import ComposerClassifier
from composer.algorithms import LabelSmoothing, CutMix, ChannelsLast

class Model(nn.Module):
    """Toy convolutional neural network architecture in pytorch for MNIST."""

    def __init__(self, num_classes: int = 10):
        super().__init__()

        self.num_classes = num_classes

        self.conv1 = nn.Conv2d(1, 16, (3, 3), padding=0)
        self.conv2 = nn.Conv2d(16, 32, (3, 3), padding=0)
        self.bn = nn.BatchNorm2d(32)
        self.fc1 = nn.Linear(32 * 16, 32)
        self.fc2 = nn.Linear(32, num_classes)

    def forward(self, x):
        out = self.conv1(x)
        out = F.relu(out)
        out = self.conv2(out)
        out = self.bn(out)
        out = F.relu(out)
        out = F.adaptive_avg_pool2d(out, (4, 4))
        out = torch.flatten(out, 1, -1)
        out = self.fc1(out)
        out = F.relu(out)
        return self.fc2(out)

transform = transforms.Compose([transforms.ToTensor()])
dataset = datasets.MNIST("data", train=True, download=True, transform=transform)
train_dataloader = DataLoader(dataset, batch_size=128)

trainer = Trainer(
    model=ComposerClassifier(module=Model(), num_classes=10),
    train_dataloader=train_dataloader,
    max_duration="2ep",
    algorithms=[
        LabelSmoothing(smoothing=0.1),
        CutMix(alpha=1.0),
        ChannelsLast(),
    ],
)
trainer.fit()

接下来，查看我们的 快速入门 Colab，了解 Composer 主要功能的逐步讲解。在本教程中，我们将介绍 Composer 训练器的基础知识：

• 数据加载器
• 训练器
• 优化器和调度器
• 日志记录
• 训练基线模型
• 加速训练

📚 了解更多

完成快速入门后，您可以浏览以下教程或我们的文档，以进一步熟悉 Composer。

如果您有任何问题，请随时通过我们的 社区 Slack 联系我们！

以下是由 Composer 社区积极维护的一些资源，帮助您入门：

🛠️ 为获得最佳效果，请在 Databricks & MosaicML 生态系统中使用

Composer 可以单独使用，但为了获得最顺畅的体验，我们建议将其与 MosaicML 生态系统的其他组件结合使用：

• Mosaic AI 训练 (MCLI) - 我们专有的命令行界面 (CLI) 和 Python SDK，用于编排、扩展和监控执行训练和部署的 GPU 节点和容器镜像。我们的客户使用它来训练他们自己的生成式 AI 模型。
  • 要开始使用，请点击此处联系我们 并查看我们的 训练 产品页面
• MosaicML LLM Foundry - 这个开源仓库包含使用 Composer 训练、微调、评估和为推理准备 LLM 的代码。旨在易于使用、高效且灵活，此代码库旨在支持使用最新技术进行快速实验。
• MosaicML StreamingDataset - 用于从云存储快速、准确流式传输的开源库。
• MosaicML Diffusion - 在您自己的数据上训练您自己的 Stable Diffusion 模型的开源代码。通过我们的博客了解更多：(结果 , 加速详情)

🏆 项目展示

以下是一些使用 Composer 的项目和实验。