MosaicML composer is an open-source PyTorch training library built around one idea: if you can observe training as a sequence of well-defined moments, you can intervene at those moments to make models train faster without changing the math of your model. The library packages dozens of researched “speedup methods” — ColOut, BlurPool, SelectiveBackprop, EMA, SeqLengthWarmup, and more — as first-class objects that plug into a shared life cycle. Rather than scattering hacks through your training script, composer lets you list methods by name and trust the framework to fire them at exactly the right instant.
Why does the source deserve a tour? Because composer is one of the cleanest public examples of event-driven training architecture in the PyTorch ecosystem. The Trainer owns the loop; the Engine walks a strict Event enum and asks every Algorithm and Callback whether it wants to act; a State object carries the model, optimizers, schedulers, and timers so every participant reads and writes the same truth. Once you see how that event chain is wired, you will recognize the same pattern in larger systems, and you will understand exactly where your own hooks would go. It also pairs naturally with the tooling we have covered before: PyTorch supplies the tensors and autograd underneath, DeepSpeed and Accelerate cover other sharding and launch styles, and Flash-Attention is one kernel-level optimization that composer-style methods complement rather than replace.
In this post we read the actual code from the main branch: the trainer loop in composer/trainer/trainer.py, the dispatch machinery in composer/core/engine.py, real algorithm implementations under composer/algorithms/, the callback and logger packages, and the distributed sync helpers. Two diagrams map the territory — one coarse, one detailed.
The composer system at a glance: a Trainer loop feeds an event engine that dispatches to algorithms, callbacks, and distributed sync helpers.
Reading the overview from left to right:
- The Trainer, L526 class is the front door. Its constructor at init, L1003 assembles the
State, the logger, the engine, and every callback you registered, so by the time training starts the whole object graph exists. - fit, L1933 is the public entry point. It normalizes durations, validates your configuration, and then launches the private loop below it.
- The real rhythm lives in _train_loop, L2383, which emits
FIT_START,EPOCH_START,BATCH_START,ITERATION_END, and their siblings — onerun_eventcall per life-cycle moment. - Those calls land in the Engine, L175, whose run_event, L227 method iterates over registered algorithms and callbacks for the given Event, L10 value.
- The State, L257 object is the shared blackboard: model, optimizers, learning-rate schedulers, dataloaders, and timestamps all live here, which is why a single
Algorithm.applycan reshape training mid-flight. - On the speedup side, the Algorithm, L20 base class defines the
match/applycontract, and concrete methods such as ColOut, L162 and BlurPool, L81 implement it. - Observability comes from SpeedMonitor, L171, which reads timings from state and pushes them through the Logger, L26 facade into any destination you configured.
- Finally, on multiple workers, ddp_sync_context, L74 decides when gradient all-reduce happens, trading communication frequency for throughput.
From Install to Your First Model
Composer installs like any PyTorch companion package. pip install mosaicml pulls the library plus its core dependencies; pip install 'mosaicml[tensorboard]' or pip install 'mosaicml[wandb]' adds logger extras. There is no separate compile step and no plugin registration ritual — speedup methods and callbacks are ordinary Python classes exported from the composer namespace.
A minimal fit looks like this:
from composer import Trainer, algorithms, callbacks
trainer = Trainer(
model=model,
train_dataloader=train_dl,
eval_dataloader=eval_dl,
max_duration="2ep",
algorithms=[
algorithms.ColOut(),
algorithms.BlurPool(),
algorithms.SeqLengthWarmup(duration=0.3),
],
callbacks=[callbacks.SpeedMonitor(window_size=100)],
loggers=[],
)
trainer.fit()
The interesting part is that nothing above mentions when each method runs. ColOut declares internally that it acts on Event.FIT_START (it shrinks both image rows/columns and the batch size proportionally), BlurPool asks to be re-applied after model construction, and SeqLengthWarmup edits sequence length during the early fraction of training. The Trainer collects their preferences at init, and the engine honors them for the rest of the run. That declaration-over-orchestration style is the library’s signature: your script lists what, the framework owns when. If you have used Hugging Face tooling, the mental model transfers directly — see our Transformers tour for that side of the fence.
The Event Chain: How Algorithms Meet Training
The detail diagram zooms into the dispatch machinery, which is the heart of the codebase.
Detailed view: run_event walks the Event enum, filters algorithms through match(), applies them to State, and notifies callbacks that fan results into loggers.
Every call to Engine.run_event performs two passes in a fixed order. First _run_algorithms (line 381 of engine.py) consults each algorithm’s match(event, state) hook — declared at match, L71 — and invokes apply (defined at apply, L103) on everything that matches. Second, _run_callbacks (line 464) calls the method named after the event on each Callback — the base class at Callback, L19 provides an empty default for all of them, so subclasses override only what they need. Algorithms may mutate the model graph; callbacks are strictly observers and reporters.
The loop code in _train_loop shows the chain in context. Around BATCH_START the trainer wraps the batch in a ddp_sync_context chosen by the DDPSyncStrategy enum at DDPSyncStrategy, L44 — single_auto syncs once per batch, while multi-microbatch strategies defer the all-reduce so gradient accumulation stays cheap. Then BEFORE_FORWARD, AFTER_FORWARD, BEFORE_LOSS, AFTER_LOSS, BEFORE_BACKWARD, and AFTER_BACKWARD bracket the actual tensor work, with AFTER_TRAIN_BATCH and BATCH_END closing each step. Eval reuses the identical machinery with an EVAL_* family of events driven from _eval_loop (line 3314 of trainer.py). Checkpointing is just another event pair — BATCH_CHECKPOINT and EPOCH_CHECKPOINT — which is why resuming an interrupted run needs no special-case code.
Time itself is first-class here. The Time class at Time, L83 models durations in epochs, batches, samples, or tokens, and Timestamp, L461 snapshots how much of each unit has elapsed. Because durations can be declared as strings like '2ep' or '5000ba' and converted internally, the trainer decides when to stop without the loop knowing which unit you meant. This is also what makes fractional durations such as the 0.3 in SeqLengthWarmup work: the engine compares fractional progress on a shared clock rather than guessing from raw iteration counts, so a warmup stays correct whether the run lasts two epochs or twenty.
Speedup Methods Worth Knowing
The composer/algorithms/ directory is the library’s research catalog, each subpackage with its own README, metadata, and tests. A few illustrate the range:
- ColOut, L162 drops a small percentage of rows and columns from every image and the same percentage of examples from each batch, shrinking total FLOPs by roughly the square of the image reduction. Accuracy stays close because information loss is spread thinly.
- BlurPool, L81 rewrites strided convolutions to blur before subsampling, preserving shift invariance — a classic method from the “Making Convolutional Networks Shift-Invariant Again” paper, productized here as an apply-once model edit.
- SelectiveBackprop, L171 ranks the examples in a batch by current loss, runs the backward pass only on the hardest fraction, and skips the easy ones — cutting gradient compute when the dataset has many trivially learnable samples.
- EMA, L96 maintains an exponential moving average of weights on
AFTER_BACKWARD, usually improving final eval metrics for a tiny memory cost. - GatedLinearUnits, L169 and SeqLengthWarmup, L152 target transformer training: the former swaps FFN blocks for gated variants, the latter starts training on short sequences and ramps up, front-loading the cheapest tokens first.
- LayerFreezing, L92 freezes low-value layers as training progresses, shrinking backward cost late in the run.
When the trainer shortens a schedule — for example because ColOut effectively reduced total work — it calls scale_pytorch_scheduler, L16 so native PyTorch LR schedulers see a rescaled timeline instead of drifting out of sync. Small detail, but it shows how much of the library’s correctness lives in bookkeeping.
Callbacks, Logging, and Distributed Sync
Callbacks are the observability half of the event chain. SpeedMonitor, L171 accumulates wall-clock and throughput statistics in batch_end and reports them over a rolling window; LRMonitor, L11 logs the learning rate of every optimizer each batch; and CheckpointSaver, L45 persists full training state on a cadence you configure, delegating to the trainer’s save_checkpoint method (line 3661 of trainer.py). All of them subclass the same passive Callback base and stay out of the training path unless fired.
Logging is a two-layer design. The Logger facade fans every call out to LoggerDestination objects; composer ships destinations for TensorBoard, Weights & Biases, Comet, MLflow, Neptune, Slack, console output, and a handy InMemoryLogger, L27 that stores timeseries in a dict — ideal for tests, notebooks, and ConsoleLogger, L27-style quick inspection. Because callbacks receive the same Logger, a single speed_monitor line of logging surfaces in every destination at once. On the distributed side, helpers in composer/utils/dist.py (such as get_world_size, line 212) abstract rank and world-size queries so the trainer code reads the same on one GPU or sixty-four.
The profiler follows the same pattern. The Profiler class at Profiler, L30 is itself a callback: it schedules profiling windows, wraps torch.profiler, and writes Chrome-style traces through handlers living in the same package. Because it hooks the shared event chain, you enable it with one constructor argument and read results without ever touching the loop. Other ready-made observers in composer/callbacks/ include a memory monitor and a system metrics collector, plus an early stopper and a threshold stopper — each a small file that reacts to events and nothing more. If you ever need your own observability, copying one of these compact callbacks is the gentlest possible introduction to extending the framework.
Honest Limits
Composer’s strengths are real, but so are the caveats. First, the speedup numbers are workload-dependent: methods tuned for vision CNNs (BlurPool) do nothing for a transformer, and text-only methods (SeqLengthWarmup) are irrelevant to image classifiers — you must pick a coherent subset, not enable everything. Second, some algorithms change semantics (SelectiveBackprop skips gradients for easy examples), so loss curves will differ from a vanilla run and ablations remain your responsibility. Third, the library has evolved quickly — DeepSpeed-specific integration once lived under Enjoyed this post? Never miss out on future posts by following us composer/trainer/ but the distributed strategies now live in composer/distributed/, so older blog snippets and tutorials drift out of date; always check the branch you are reading. Fourth, composer is heavier than a bare PyTorch loop: the event machinery is worth it at scale, but for a five-minute experiment on one GPU, plain PyTorch with a hand-written loop can be simpler. Finally, the Trainer is opinionated about orchestration — if your project needs a fully custom optimization regime, you may spend more time fighting State assumptions than writing your own loop. None of this diminishes the architecture: as a study of how to make training extensible without chaos, composer remains one of the best codebases in the ecosystem. 
