DOCSEARCHbyAlgolia
Submit Docs

Composer Documentation

repository·main·Indexed 26 days ago

https://github.com/mosaicml/composer

An open-source deep learning training library built on PyTorch designed to simplify and accelerate large-scale distributed training for LLMs, diffusion models, and large-scale neural networks. It provides high-level abstractions for parallelism, data streaming, and training loop management, featuring a Trainer class and a functional API for algorithms such as ALiBi, AugMix, BlurPool, and ChannelsLast.

Tokens
64.7K
Snippets
156
Records
266
Agent score
90%

What's inside Composer

  1. Overview of Composer

    main
    Composer is an open-source deep learning training library built on top of PyTorch. It is designed to simplify and optimize distributed training workflows on large-scale clusters by abstracting low-level complexities such as parallelism techniques, distributed data loading, and memory optimization. It is suitable for training various model architectures, including Large Language Models (LLMs), Diffusion models, Embedding models (e.g., BERT), Transformers, and Convolutional Neural Networks (CNNs).

  2. Overview of Composer Training Methods

    main

    Composer provides a variety of training methods and augmentations categorized by their application domain: Computer Vision (CV) and Natural Language Processing (NLP). These methods can be used to improve scalability, customizability, and training efficiency.

    Computer Vision (CV) Methods

    • AugMix: Image-preserving data augmentations.
    • BlurPool: Applies blur before pooling or downsampling.
    • ChannelsLast: Uses channels last memory format (NHWC).
    • ColOut: Removes columns and rows from the image for augmentation and efficiency.
    • CutMix: Combines pairs of examples in non-overlapping regions and mixes labels.
    • CutOut: Randomly erases rectangular blocks from the image.
    • GhostBatchNorm: Uses smaller number of samples to compute batchnorm.
    • GyroDropout: (Note: Source text indicates this clips gradients, similar to GradientClipping).
    • LabelSmoothing: Smooths labels with a uniform prior.
    • MixUp: Blends pairs of examples and labels.
    • ProgressiveResizing: Increases the input image size during training.
    • RandAugment: Applies a series of random augmentations.
    • SqueezeExcite: Replaces eligible layers with Squeeze-Excite layers.
    • StochasticDepth: Replaces a specified layer with a stochastic version that randomly drops the layer or samples during training.
    • Weight Standardization: Ensures convolution weights always have zero mean and unit variance.

    Natural Language Processing (NLP) Methods

    • Alibi: Replaces attention with AliBi.
    • GatedLinearUnits: Swaps the building block from a Linear layer to a Gated Linear layer.
    • SeqLengthWarmup: Progressively increases sequence length.

    Cross-Domain Methods (CV & NLP)

    • EMA: Maintains an exponential moving average of model weights for evaluation.
    • Factorize: Factorizes GEMMs into smaller GEMMs.
    • GradientClipping: Clips all gradients in the model based on a specified clipping_type.
    • LayerFreezing: Progressively freezes layers during training.
    • LowPrecisionGroupNorm: Forces GroupNorm to run in lower precision.
    • LowPrecisionLayerNorm: Forces LayerNorm to run in lower precision.
    • SAM: SAM optimizer measures sharpness of optimization space.
    • SelectiveBackprop: Drops examples with small loss contributions.
    • SWA: Computes running average of model weights.

  3. Overview of Composer features

    main

    Composer is an open-source deep learning training library by MosaicML that optimizes PyTorch for scalability and usability. Key capabilities include:

    • Scalability: Supports PyTorch FullyShardedDataParallelism (FSDP) for large models, standard Distributed Data Parallelism (DDP), and elastic sharded checkpointing (allowing you to save on one GPU count and resume on another).
    • Data Streaming: Integrates with MosaicML StreamingDataset to stream data from cloud blob storage during training.
    • Customizability: Features a callback system to insert custom logic at specific training loop events (e.g., BATCH_END) and a collection of algorithmic speedup recipes.
    • Workflow Automation: Includes auto-resumption from checkpoints, CUDA OOM prevention via automatic microbatch size selection, and Time abstractions for specifying training duration in epochs, batches, samples, or tokens.
    • Integrations: Supports cloud storage (OCI, GCP, AWS S3) for checkpointing/logging and experiment tracking via Weights and Biases, MLFlow, CometML, and neptune.ai.

  4. Composer Ecosystem Integrations

    main

    For optimal performance, Composer is designed to work with the following MosaicML ecosystem components:

    • Mosaic AI training (MCLI): A CLI and Python SDK for orchestrating and scaling GPU training.
    • MosaicML LLM Foundry: An open-source repository for training, finetuning, and evaluating LLMs using Composer.
    • MosaicML StreamingDataset: An open-source library for high-speed streaming from cloud storage.
    • MosaicML Diffusion: Open-source code for training Stable Diffusion models.

  5. Understand Composer training events

    main
    Events represent specific points in the training loop where a composer.core.Algorithm or a composer.core.Callback can execute. These events allow you to hook into the lifecycle of the training process to perform custom actions, logging, or modifications to the training state.

  6. Inject custom logic with Events and State

    main
    The Composer trainer uses an event-driven system to allow injecting custom logic at specific points in the training loop (e.g., before_forward, after_forward). Algorithms and callbacks register to these events. The State object stores the trainer's current state (model, optimizers, dataloader, current batch, etc.), which can be modified by algorithms during events.

  7. Save only model weights for inference

    main

    To save a checkpoint containing only model weights (plus metadata and integrations) instead of the full training state, set save_weights_only=True in the Trainer.

    from composer.trainer import Trainer

trainer = Trainer(
    ...,
    save_folder="checkpoints",
    save_weights_only=True,
    save_overwrite=True,
)

trainer.fit()

  8. Configure FSDP Checkpointing (Save and Load)

    main

    Control how FSDP handles sharded checkpoints using the state_dict_type key in fsdp_config:

    1. state_dict_type='full' (Default): Saves one large checkpoint file by gathering the model state to global rank 0. If load_monolith_rank0_only=True is set, rank 0 loads the file and scatters it to other ranks to reduce system memory usage.
    2. state_dict_type='sharded': Each rank saves an unflattened shard. Recommended for PyTorch 2.0.0 or higher.

    Organizing Sharded Checkpoints: To prevent save_folder pollution, use 'sharded_ckpt_prefix_dir' to group shards into subdirectories. The default is 'ep{epoch}-ba{batch}'.

    Loading Sharded Checkpoints:

    • Set load_path to the directory containing the shards, not a specific file.
    • Ensure state_dict_type is set to 'sharded' in your config.
    • Composer with PyTorch 2.0.0+ supports elastic checkpointing (resuming with a different number of ranks).
        fsdp_config = {
        'sharding_strategy': 'FULL_SHARD',
        'state_dict_type': 'sharded',
        'sharded_ckpt_prefix_dir': 'ba{batch}-shards'
    }

    trainer = Trainer(
        model=composer_model,
        max_duration='4ba',
        parallelism_config={'fsdp': fsdp_config},
        save_folder='checkpoints',
        save_interval='2ba',
    )

    # To load:
    trainer = Trainer(
        model=composer_model,
        parallelism_config={'fsdp': {'sharding_strategy': 'FULL_SHARD', 'state_dict_type': 'sharded'}},
        load_path='./checkpoints/ba2-shards' # Path to the directory
    )

  9. Configure automatic checkpointing in Trainer

    main

    To enable automatic checkpointing, provide a save_folder argument when initializing the Trainer. You can customize filenames using save_filename and the symlink for the latest checkpoint using save_latest_filename. These arguments accept Python format strings (e.g., ep{epoch}).

    Available format variables can be found in the CheckpointSaver documentation.

    from composer import Trainer

trainer = Trainer(
    model=model,
    train_dataloader=train_dataloader,
    max_duration="2ep",
    save_folder="./path/to/checkpoints",
    save_filename="ep{epoch}",
    save_latest_filename="latest",
    save_overwrite=True,
)

trainer.fit()

  11. Identify DataLoader bottlenecks using traces

    main

    To determine if your training is bottlenecked by data loading, compare the duration of the dataloader/train event against the event/batch event (which encompasses the forward and backward passes):

    1. Bottlenecked: If dataloader/train takes significantly more time than event/batch, the training is dataloader bottlenecked. Increasing the number of workers in your DataLoader is recommended.
    2. Optimized: In an optimized setup, dataloader/train should be much smaller in duration compared to event/forward and event/backward.

  12. Use DistributedSampler for non-iterable datasets

    main

    If you are using a torch.utils.data.Dataset (that is NOT an IterableDataset) with a torch.utils.data.DataLoader, you must provide a DistributedSampler to ensure different devices receive different batches. Composer will raise an error if it is missing.

    Use composer.utils.dist.get_sampler to create a sampler with the correct parameters.

    from composer.utils import dist

sampler = dist.get_sampler(dataset, shuffle=True)
dataloader = DataLoader(dataset, batch_size=32, sampler=sampler)