mosaicml 0.4.0

composing methods for ML training efficiency

Reproducible State-Of-The-Art Training Efficiency:
Giving You Lower Cost and Better Accuracy

[Website] - [Getting Started] - [Examples] - [Docs] - [Methods] - [Explorer] - [We're Hiring!]

---

Composer provides well-engineered implementations of efficient training methods to give the tools that help you train a better model for cheaper.

Using Composer, you can:

• Train an ImageNet model to 76.1% accuracy for $37 (with vanilla PyTorch: $127)
• Train a GPT-2 125M to a perplexity of 23.9 for $148 (with vanilla PyTorch: $255)
• Use start-of-the-art implementations of methods to speed up your own training loop.

At MosaicML, we are focused on making training ML models accessible. To do this, we continually productionize state-of-the-art academic research on efficient model training, and also study the combinations of these methods in order to ensure that model training is ✨ as efficient as possible ✨.

Everyone has their own priorities: best accuracy, cheapest cost, and somewhere in between. Composer provides novel recipes that push the boundary of both cost and accuracy. Composer allows you to choose the best model for your real-world constraints.

Composer features:

• 20+ efficient training methods for training a better language and vision models! Don't waste hours trying to reproduce research papers when Composer has done the work for you.
• Easy-to-use (optional) Trainer interface written to be as performant as possible, and integrated best practices.
• Easy-to-use Functional forms that allow you to integrate efficient training methods into your training loop!
• Strong, reproducible baselines to get you started as 💨 fast 💨 as possible

Quickstart

Installation

Composer is available with Pip

pip install mosaicml

Alternatively install Composer with Conda

conda install mosaicml

---

Usage

Composer provides both a Functional API (similar to torch.nn.functional) and a Trainer (that abstracts away the training loop) to provide flexibility to users.

Example: Functional API

For users who choose to use their own training loop, we provide state-less functional implementations of our algorithms for a end-user to integrate.

The following example highlights using BlurPool, which applies an anti-aliasing filter before every downsampling operation.

from composer import functional as cf
import torchvision

model = torchvision.models.resnet50()

# Apply model surgery before training by replacing eligible layers
# with a BlurPool-enabled layer (Zhang, 2019)
model = cf.apply_blurpool(model)

# Start your training loop here
for epoch in range(NUM_EPOCHS):
    for input, labels in dataloader:
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

See the official Composer Functional API Colab notebook for more.

Example: Trainer

For maximal speedups, we recommend using our Trainer, which manages handling user state, performant algorithm implementations, and provides useful engineering abstractions to permit rapid experimentation.

import composer

# Normalization constants
mean = (0.507, 0.487, 0.441)
std = (0.267, 0.256, 0.276)

batch_size = 1024

# setup data
data_directory = "data"
cifar10_transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean, std)]
train_dataset = datasets.CIFAR10(data_directory, train=True, download=True, transform=cifar10_transforms)
test_dataset = datasets.CIFAR10(data_directory, train=False, download=True, transform=cifar10_transforms)

# setup model & optimization
model = composer.models.CIFAR10_ResNet56()

optimizer = composer.optim.DecoupledSGDW(
    model.parameters(), # Model parameters to update
    lr=0.05, # Peak learning rate
    momentum=0.9,
    weight_decay=2.0e-3 # If this looks large, it's because its not scaled by the LR as in non-decoupled weight decay
)

warmup = composer.optim.WarmUpLR(
    optimizer, # Optimizer
    warmup_iters=25, # Number of iterations to warmup over. 50k samples * 1 batch/2048 samples
    warmup_method="linear", # Linear warmup
    warmup_factor=1e-4, # Initial LR = LR * warmup_factor
    interval="step", # Update LR with stepwise granularity for superior results
)

# setup algorithm in one line
blurpool = composer.algorithms.BlurPool()  # credit:  (Zhang, 2019)

# for brevity, we hardcode some argument values
trainer = composer.trainer.Trainer(model=model,
                                   train_dataloader=train_dataloader,
                                   eval_dataloader=test_dataloader,
                                   max_duration="3ep", # Train for 3 epochs because we're assuming Colab environment and hardware
                                   optimizers=optimizer,
                                   schedulers=[warmup],
                                   algorithms=[blurpool] # Adding BlurPool via model surgery, can just add more algorithms here!
                                   device="gpu" # Train on the GPU,
                                   seed=42) # the meaning to life, the universe, and everything

# start training!
trainer.fit()

Using the Composer Trainer allows you to add multiple efficient training methods in a single line of code! Trying out new methods or combinations of methods is as easy as adding another line! As Composer gets better and we implement more methods and quality of life improvements, the savings are directly passed to you.

For concrete examples of methods in Composer, here's some (see here for all) efficiency methods currently in Composer:

Name	Functional	Attribution	tl;dr
Alibi	cf.apply_alibi	(Press et al, 2021)	Replace attention with AliBi
AugMix	cf.augmix_image	(Hendrycks et al, 2020)	Image-perserving data augmentations
BlurPool	cf.apply_blurpool	(Zhang, 2019)	applies blur before pooling
ChannelsLast	cf.apply_channels_last	PyTorch	Uses channels last memory format (NHWC)
ColOut	cf.colout_batch	Many	Removes columns and rows from the image for augmentation and efficiency.
CutMix	cf.cutmix_batch	(Yun et al, 2019)	Combines pairs of examples in non-overlapping regions and mixes labels
CutOut	cf.cutout_batch	(DeVries et al, 2017)	Randomly erases rectangular blocks from the image.
Factorize	cf.apply_factorization	MosaicML	Factorize GEMMs into smaller GEMMs
GhostBatchNorm	cf.apply_ghost_batchnorm	(Dimitriou et al, 2020)	Use smaller samples to compute batchnorm
LabelSmoothing	cf.smooth_labels	(Szegedy et al, 2015)	Smooths the labels with a uniform prior
LayerFreezing	cf.freeze_layers	Many (Raghu et al, 2017)	Progressively freezes layers during training.
MixUp	cf.mixup_batch	(Zhang et al, 2017)	Blends pairs of examples and labels
ProgressiveResizing	cf.resize_batch	Fast AI	Increases the input image size during training
RandAugment	cf.randaugment_image	(Cubuk et al, 2020)	Applies a series of random augmentations
SAM	N/A	(Foret et al, 2021)	SAM optimizer measures sharpness of optimization space
ScaleSchedule	N/A	Many	Scale the learning rate schedule by a factor
SelectiveBackprop	cf.selective_backprop	(Jiang et al, 2019)	Drops examples with small loss contributions.
SeqLengthWarmup	cf.set_batch_sequence_length	(Li et al, 2021)	Progressively increase sequence length.
SqueezeExcite	cf.apply_squeeze_excite	Hu et al, 2017	Replaces eligible layers with Squeeze-Excite layers
StochasticDepth	cf.apply_stochastic_depth	(Huang et al, 2016)	Replaces a specified layer with a stochastic verion that randomly drops the layer or samples during training
SWA	N/A	(Izmailov et al, 2018)	Computes running average of model weights.

Speedups are measured based on time to train to iso-accuracy.

What benchmarks does Composer support?

MosaicML uses a benchmark as a term to denote an reproducible standard within the machine learning community. A benchmark is a specific model trained for a task, where a task is defined as a specific dataset with a specific loss function.

Composer is currently focused on supporting computer vision and natural language processing use cases. We currently support the following combinations of models, datasets, and loss functions.

Model	Dataset	Loss	Task	Evaluation Metrics
Computer Vision
ResNet Family	CIFAR-10	Cross Entropy	Image Classification	Classification Accuracy
ResNet Family	ImageNet	Cross Entropy	Image Classification	Classification Accuracy
EfficientNet Family	ImageNet	Cross Entropy	Image Classification	Classification Accuracy
UNet	BraTS	Dice Loss	Image Segmentation	Dice Coefficient
DeepLab v3	ADE20K	Cross Entropy	Image Segmentation	mIoU
Natural Language Processing
BERT Family	{OpenWebText, C4}	Cross Entropy	Masked Language Modeling	GLUE
GPT Family	{Wikipedia & BooksCorpus, C4}	Cross Entropy	Language Modeling	Perplexity

Why use Composer?

The compute required to train a state-of-the-art machine learning model is doubling every 6 months, subsequently making machine learning less accessible for the broader community. Composer shifts the focus to efficiency, and contains reproducible versions of cutting-edge algorithms that help reduce the compute and cost required to train state-of-the-art models. While every paper will claim state-of-the-art efficiency results, Composer will be your source well-engineered implementations of efficient training methods that actually work in practice.

Furthermore, combining these efficiency methods together isn't a piece of cake. When exploring how to seamlessly combine different efficient training methods, we that existing Trainers failed to provide a flexible design that would be needed to interleave and inject many different methods into the training loop. To fulfill this need, we designed a Trainer built for efficiency from first-principles. In Composer, we carefully designed new abstractions to allow us to have flexibility in all necesssary parts of the training loop.

Composer is designed with two-way callbacks (Howard et al, 2020) as a first-class citizen. This enables easy injection of efficiency methods throughout the entire training loop. Not only does the two-way callback system trigger at every part of the training loop, each callback is designed to pass the entire training state designed so that any part of training can be modified.

Composer leverages the two-way callback system to seamlessly integrate our methods into the training pipeline, for examples:

• Composer modifies data inputs for batches (data augmentations, sequence length warmup, skipping examples, etc)
• Composer modifies the neural network (layer freezing, pruning, model surgery, etc)
• Composer modifies the loss function (label smoothing, MixUp, CutMix, etc)
• Composer modifies the optimizer (Sharpness Aware Minimization)

With Composer, we make it easy to add your own methods or callbacks to easily instrument any part of the training loop.

Why shouldn’t I use Composer?

Composer is a framework built with a focus on training the most efficient model for your downstream task.

• If training efficiency is not a concern, then Composer may not be as well suited for your task. Inference efficiency is not in the current roadmap for Composer support.
• Composer currently supports most computer vision (CV) and natural language processing (NLP) use cases. We will support most industry applications, but may not support highly custom or novel architectures.

Composer is an active and ongoing project. Since Composer is still in alpha, our API may not be stable. We recommend pegging your work to a Composer version.

We welcome any comments, feedback, or contributions to Composer! Please do not hesitate to file an issue or pull request 🤩.

Learn More

Here's some resources actively maintained by the Composer community to help you get started:

Resource	Details
Founder's Blog	A blog post by our founders highlighting why MosaicML exists
Getting started with our Trainer	An interactive Colab Notebook aimed at teaching users about our Trainer
Getting started with our Functional API	An interactive Colab Notebook aimed at teaching users about our Functional API
PyTorch Lightning Migration Guide	An interactive Colab Notebook aimed at helping users migrate from PTL to Composer
We're Hiring!	Join us! 🤩

If you have any questions, please feel free to reach out to us on Twiter, email, or our Community Slack!

Contributors

Composer is part of the broader Machine Learning community, and we welcome any contributions, pull requests, or issues!

Citation

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