Composer: 加速深度学习模型训练的开源框架

Composer简介

Composer是MosaicML公司开发的一个开源深度学习训练库,旨在简化大规模分布式训练工作流程,提高模型训练效率。它基于PyTorch构建,集成了许多最佳实践,使研究人员和工程师能够更轻松地进行大规模模型训练。

Composer Logo

Composer的设计理念是"可扩展性和易用性的优化"。它抽象了底层的复杂性,如并行化技术、分布式数据加载和内存优化等,让用户可以专注于模型训练和实验,而不会因技术细节而减慢速度。

Composer的主要特性

1. 卓越的可扩展性

无论您是在单个GPU上训练,还是在512个GPU上训练,无论是处理50MB还是10TB的数据,Composer都能保持工作流程的简单性。它支持以下关键功能:

FSDP(完全分片数据并行): 对于大型模型,Composer集成了PyTorch的FullyShardedDataParallelism,使自定义模型的并行化变得简单高效。
弹性分片检查点: 支持在不同数量的GPU之间保存和恢复模型状态,提高了灵活性。
数据流式处理: 与MosaicML StreamingDataset集成,支持从云存储中动态下载和处理大型数据集。

2. 高度可定制性

Composer提供了简单的方法来定制训练过程,以满足特定需求:

回调系统: 允许用户在训练循环的任何点插入自定义逻辑。例如,可以用于监控内存使用、记录和可视化图像等。
加速算法: 提供了一系列算法加速技术,可以组合使用以提高训练速度。例如,在ImageNet上训练ResNet-50的速度提高了7倍。

3. 优化的工作流程

Composer旨在自动化低级痛点,让用户专注于深度学习的重要(和有趣)部分:

自动恢复: 训练中断?只需重新运行代码,Composer会自动从最新的检查点恢复。
CUDA内存溢出预防: 将microbatch大小设置为"auto",Composer会自动选择适合GPU的最大批量。
时间抽象: 使用自定义单位(epochs、batches、samples和tokens)在训练循环中指定持续时间。

4. 广泛的集成

Composer与许多流行的工具和平台集成,用于实验跟踪和数据流处理:

云集成: 检查点和日志功能原生支持远程存储和从云存储桶加载(OCI、GCP、AWS S3)。
实验跟踪: 支持Weights and Biases、MLFlow、CometML和neptune.ai等平台。

使用Composer

安装

Composer可以通过pip安装:

pip install mosaicml

为简化环境设置,MosaicML还提供了一系列预构建的Docker镜像。强烈建议使用这些Docker镜像。

快速开始

以下是一个使用Composer训练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):
    """用于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()

这个例子展示了Composer的核心功能:

定义一个简单的PyTorch模型
设置数据加载器
使用Composer的Trainer类配置训练过程
应用一些加速算法
开始训练

Composer在实际项目中的应用

Composer已在多个大型项目中得到应用,展示了其在实际场景中的强大能力:

MPT基础系列: 这是一系列商业可用的开源大语言模型(LLMs),使用Composer进行优化训练。包括MPT-7B、MPT-7B-8k和MPT-30B等模型。
Mosaic扩散模型: 使用Composer,研究人员能够以不到$50,000的成本从头训练一个稳定扩散模型,相比原始成本节省了近70%。
replit-code-v1-3b: 这是一个专注于代码补全的2.7B因果语言模型,由Replit使用Mosaic AI训练平台在10天内训练完成。
BabyLLM: 这是第一个同时支持阿拉伯语和英语的LLM。这个7B参数的模型由MetaDialog在世界最大的阿拉伯语/英语数据集上训练,旨在改善客户支持工作流程。
BioMedLM: 这是MosaicML和斯坦福CRFM合作开发的生物医学领域专用LLM。