为你的模型训练赋能

大规模训练深度学习（Deep Learning）框架

[网站] - [快速入门] - [文档] - [我们正在招聘！]

👋 欢迎

Composer 是由 MosaicML 开发的开源深度学习（Deep Learning）训练库。基于 PyTorch 构建，Composer 库使得在大规模集群上实现分布式训练（Distributed Training）工作流变得更加容易。

我们构建 Composer 是为了针对可扩展性（Scalability）和易用性（Usability）进行优化，集成了高效多节点训练的最佳实践。通过抽象掉并行化技术、分布式数据加载和内存优化等底层复杂性，您可以专注于训练现代机器学习（ML）模型和运行实验，而无需担心速度变慢。

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

• 大型语言模型（Large Language Models, LLMs）
• 扩散模型（Diffusion models）
• 嵌入模型（Embedding models，例如 BERT）
• 基于 Transformer 的模型
• 卷积神经网络（Convolutional Neural Networks, CNNs）

MosaicML 研究团队广泛使用 Composer 来训练最先进（State-of-the-art）的模型，如 MPT，我们将此库开源以使机器学习社区也能做到同样的事情。该框架被科技行业和学术界的组织使用，并不断更新新功能、错误修复和稳定性改进，以适应生产负载。

🔑 主要特性

我们从零开始设计 Composer，以适应现代深度学习工作负载。AlexNet 和 ResNet 的时代已经过去，那时最先进的模型可以在几块桌面 GPU（图形处理器）上训练。如今，开发最新最好的深度学习模型通常需要集群级硬件——但在 Composer 的帮助下，您几乎不会注意到其中的差异。

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

可扩展性

无论您是在 1 个 GPU（图形处理器）还是 512 个 GPU 上训练，处理 50MB 还是 10TB 的数据 - Composer 旨在保持您的工作流简单。

• FSDP：对于大到无法放入 GPU 的大型模型，Composer 已将 PyTorch 完全分片数据并行（FullyShardedDataParallelism） 集成到我们的 Trainer 中，并使得高效并行化自定义模型变得简单。我们发现 FSDP 在性能上与更复杂的并行策略具有竞争力。或者，Composer 也支持标准的 PyTorch 分布式数据并行（Distributed Data Parallelism, DDP）执行。
• 弹性分片检查点（Elastic sharded checkpointing）：在八个 GPU 上保存，在十六个 GPU 上恢复。Composer 支持弹性分片检查点，因此您无需担心分片保存的状态是否与新的硬件设置兼容。
• 数据流式传输（Data streaming）： 处理大型数据集？通过在模型训练期间集成 MosaicML StreamingDataset，按需从云对象存储（Cloud blob storage）下载数据集。

可定制性

其他高级深度学习训练器（deep learning trainers）往往以牺牲灵活性为代价来换取简便性。当您想要添加自己的功能时，它们的抽象层反而会成为阻碍。相比之下，Composer 提供了简单的方法，让您可以根据需求定制我们的 Trainer（训练器）。

图 1： Composer 的训练循环在训练过程的每个阶段都有一系列事件发生。回调（Callbacks）是用户编写的函数，用于在特定事件运行时执行。例如，我们的 Learning Rate Monitor Callback 会在每个 BATCH_END 事件记录学习率。

• 回调（Callbacks）：Composer 的回调系统允许您在训练循环的任何位置插入自定义逻辑。我们编写了回调来监控内存使用情况、记录和可视化图像，以及估算模型的剩余训练时间等。此功能在想要实现和实验自定义训练技术的研究人员中很受欢迎。
• 加速算法（Speedup algorithms）：我们借鉴最新研究创建了一系列算法加速方法。将这些加速方法堆叠到 MosaicML 配方（recipes）中以提升训练速度。我们的团队已经开源了针对不同类型模型的最佳加速组合。
  • 8 倍加速：Stable Diffusion
    • $200k 原始 SD2 成本 —> $50k (博客)
  • 7 倍加速：ImageNet 上的 ResNet-50
    • 3 小时 33 分 —> 25 分钟（8xA100 上）(博客)
  • 8.8 倍加速：BERT-Base 预训练
    • 10 小时 —> 1.13 小时（8xA100 上）(博客)
  • 5.4 倍加速：ADE20K 上的 DeepLab v3
    • 3 小时 30 分 —> 39 分钟（8xA100 上）(博客)

更优的工作流

Composer 旨在自动化消除低级痛点和问题，以便您可以专注于深度学习的重要（且有趣）部分并更快地迭代。

• 自动恢复（Auto-resumption）：训练运行失败？别担心——只需重新运行代码，Composer 就会自动从您最新保存的检查点（checkpoint）恢复。
• CUDA 内存不足预防（CUDA OOM Prevention）：告别内存不足（OOM）错误。将您的微批次大小（microbatch size）设置为"auto"，Composer 将自动选择适合您图形处理器（GPU）的最大值。
• 时间抽象（Time Abstractions）： 是否曾搞混过更新步骤、轮次（epochs）、样本（samples）和令牌（tokens）之间的转换？使用我们的 Time 类，在训练循环中指定具有自定义单位（epochs、batches、samples 和 tokens）的训练持续时间。

集成

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

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

🚀 快速开始

📍 前置条件

Composer 专为熟悉 Python 并对深度学习基础知识和 PyTorch 有基本了解的用户设计。

软件要求： 最新版本的 PyTorch。

硬件要求： 具有 CUDA 兼容图形处理器（GPU）的系统（AMD + RoCM 即将推出！）。Composer 可以在中央处理器（CPU）上运行，但为了获得全部益处，我们建议在硬件加速器上使用它。

💾 安装

Composer 可以通过 pip 安装：

pip install mosaicml

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

🏁 快速入门

这是一个代码片段，展示了我们在 MNIST 数据集上的 Trainer（训练器）。

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 Trainer 的基础知识：

• 数据加载器（Dataloader）
• 训练器（Trainer）
• 优化器（Optimizer）和调度器（Scheduler）
• 日志记录
• 训练基线模型
• 加速训练

📚 了解更多

完成快速入门 (Quick Start) 后，您可以浏览下面的教程或我们的 文档 以进一步熟悉 Composer (训练框架)。

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

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

资源	详情
使用 Composer 和 🤗 训练 BERT	一个 Colab 笔记本，展示如何使用 Composer 和 🤗 训练 BERT 模型！
大语言模型 (LLM) 预训练和微调教程	来自 MosaicML LLM Foundry 的教程，使用 MosaicML Composer、StreamingDataset 和 MCLI 进行大语言模型的训练和评估。
从 PyTorch Lightning 迁移	本教程说明了从 PyTorch Lightning 工作流迁移到 Composer 的路径。
微调和预训练 HuggingFace 模型	想在 Composer 中使用 Hugging Face 模型？没问题。在这里，我们将逐步介绍如何使用 Composer 微调预训练的 Hugging Face BERT 模型。
构建加速方法	一个 Colab 笔记本，展示如何在 Composer 之上构建新的训练修改方法

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

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

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

🏆 项目展示

以下是一些使用 Composer 的项目和实验。有什么要补充的吗？请在我们的 社区 Slack 中分享！

• MPT Foundation Series: 可用于商业目的的开源大语言模型 (LLM)，针对快速训练和推理进行了优化，并使用 Composer 训练。
  • MPT-7B 博客
  • MPT-7B-8k 博客
  • MPT-30B 博客
• Mosaic Diffusion 模型: 看看我们如何以低于 5 万美元的成本从头训练 Stable Diffusion 模型
• replit-code-v1-3b: 一个 2.7B 因果语言模型，专注于 代码补全，由 Replit 在 Mosaic AI training 上训练，耗时 10 天。
• BabyLLM: 第一个支持阿拉伯语和英语的大语言模型。这个 7B 模型由 MetaDialog 在世界上最大的阿拉伯语/英语数据集上训练，以改善客户支持工作流 (博客)
• BioMedLM: 由 MosaicML 和 Stanford CRFM 构建的生物医药领域专用大语言模型

💫 贡献者

Composer 是更广泛的机器学习 (Machine Learning) 社区的一部分，我们欢迎任何贡献、Pull requests 或问题！

开始贡献，请参阅我们的 贡献 页面。

附注：我们正在招聘!

❓常见问题解答 (FAQ)

• 训练大型模型时，您推荐的最佳技术栈是什么？
  • 我们建议用户结合 MosaicML 生态系统的组件以获得最流畅的体验：
    • Composer
    • StreamingDataset
    • MCLI (Databricks Mosaic AI Training)
• 如何获得使用 Composer 的社区支持？
  • 您可以加入我们的 社区 Slack!
• Composer 与 NeMo Megatron 和 PyTorch Lightning 等其他训练器相比如何？
  • 我们构建 Composer 是为了优化简洁性和效率。社区用户分享说，与其他库相比，他们喜欢 Composer 的功能和易用性。
• 如何使用 Composer 训练图神经网络 (GNNs)、生成对抗网络 (GANs) 或强化学习 (RL) 模型？
  • 如果您想训练这些类型的模型，我们建议您使用其他库 - 我们在设计 Composer 时所做的许多假设对于 GNNs、RL 和 GANs 来说并不是最优的
• 如何加快 HuggingFace 下载速度？
  • 您可以使用 hf transfer (pip install hf-transfer) 并设置环境变量 HF_HUB_ENABLE_HF_TRANSFER=1

✍️ 引用

@misc{mosaicml2022composer,
    author = {The Mosaic ML Team},
    title = {composer},
    year = {2021},
    howpublished = {\url{https://github.com/mosaicml/composer/}},
}

Composer 快速上手指南

简介

Composer 是由 MosaicML 开发的开源深度学习训练库，基于 PyTorch 构建。它专为大规模集群上的分布式训练设计，通过抽象底层复杂性（如并行技术、分布式数据加载和内存优化），帮助用户更高效地训练现代 ML 模型（如 LLM、Diffusion、Transformer 等）。

环境准备

软件要求：

• Python
• 最近版本的 PyTorch

硬件要求：

• 推荐使用配备 CUDA 兼容 GPU 的系统（以获得最佳性能）。
• 支持 CPU 运行，但无法发挥全部优势。
• （AMD + RoCM 支持即将推出）

安装步骤

您可以通过 pip 直接安装：

pip install mosaicml

推荐方式： 为了简化环境配置，官方提供了一套 预构建的 Docker 镜像，强烈建议使用 Docker 进行部署。

基本使用

以下是一个基于 MNIST 数据集的最小化训练示例，展示了如何使用 Trainer 以及集成算法（如标签平滑、CutMix 等）。

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()

核心功能说明：

• Trainer: 高度优化的 PyTorch 训练循环抽象。
• Algorithms: 轻松集成加速算法（如示例中的 LabelSmoothing）。
• Time Abstractions: 支持使用自定义单位（如 2ep 表示 2 个 epoch）指定训练时长。

更多详细信息请参考 官方文档。