MosaicML

Overview

MosaicML provides a platform for training large deep learning models, with a specific emphasis on large language models (LLMs). The platform aims to reduce the time and cost associated with training complex models by offering optimized algorithms and infrastructure. It was acquired by Databricks in 2023, integrating its capabilities into the broader Databricks Machine Learning Platform product offering. The core of MosaicML's approach lies in its software stack, which includes Composer, an open-source library for training efficiency, and LLM Foundry, a toolkit for pre-training and fine-tuning LLMs.

Developers use MosaicML to manage the entire lifecycle of deep learning model training, from data preparation to model deployment. The platform supports various neural network architectures and is designed for GPU-accelerated workloads. It provides programmatic access via a Python SDK, allowing users to define training runs, manage datasets, and monitor performance. The primary use cases for MosaicML revolve around organizations that require significant computational resources for deep learning development, such as those involved in generative AI, natural language processing, and computer vision research and application.

The technology focuses on practical optimizations that translate into faster training times and lower infrastructure costs. For example, techniques like mixed-precision training and gradient accumulation are integrated into the platform to improve computational efficiency. This focus on efficiency is particularly relevant for the pre-training of LLMs, which often require extensive datasets and prolonged training durations, as detailed in research on foundational models such as LLaMA. MosaicML aims to provide a managed environment where these optimizations are applied automatically or with minimal configuration, allowing developers to concentrate on model architecture and data quality rather than low-level infrastructure management. Its integration with Databricks further extends its capabilities by offering unified data and AI governance, MLOps tooling, and enterprise-grade security within a single platform.

Key features

• LLM Pre-training and Fine-tuning: Tools and optimized workflows specifically for training and adapting large language models, including support for various architectures and datasets.
• Deep Learning Training Efficiency: Incorporates techniques like gradient accumulation, mixed-precision training, and model parallelism to accelerate training and reduce GPU hours.
• Composer Library: An open-source PyTorch library that integrates performance-enhancing methods directly into training loops, improving model convergence and throughput as documented.
• LLM Foundry: A toolkit built on Composer for pre-training, fine-tuning, and evaluating LLMs, including pre-built recipes for common models and datasets with detailed guides.
• Scalable Infrastructure: Provides managed GPU clusters and elasticity to scale training jobs from single-node experiments to distributed training across hundreds of GPUs.
• Experiment Tracking and MLOps: Integrates with Databricks MLOps capabilities for tracking experiments, managing model versions, and deploying models into production.
• Cost Optimization: Designed to minimize cloud infrastructure costs by enhancing training efficiency and optimizing resource utilization.
• Python SDK: Offers programmatic control over the platform, allowing developers to define and execute training jobs, manage datasets, and interact with models.

Pricing

MosaicML, as part of Databricks, operates under a custom enterprise pricing model. Specific costs are determined based on an organization's usage, compute consumption (measured in Databricks Units or DBUs), and feature requirements. Direct self-service pricing is not publicly available for the MosaicML components; interested parties are directed to contact Databricks sales for a customized quote. There is no free tier available for the platform's services.

For detailed pricing inquiries, refer to the Databricks Platform Pricing page or contact their sales team directly.

Common integrations

• Databricks Lakehouse Platform: Deep integration with Databricks for unified data, analytics, and AI, leveraging Delta Lake for data management as part of the offering.
• MLflow: Utilized for experiment tracking, model lifecycle management, and MLOps workflows within the Databricks environment via Databricks.
• Major Cloud Providers: Deploys on AWS, Azure, and Google Cloud, abstracting away underlying infrastructure management.
• PyTorch: Composer is built on PyTorch, enabling direct use of the PyTorch ecosystem for model development as an open-source framework.
• Hugging Face: Can integrate with Hugging Face models and datasets, particularly for LLM development and fine-tuning for pre-trained models.

Alternatives

• AWS SageMaker: A comprehensive cloud-based machine learning platform offering tools for building, training, and deploying ML models at scale.
• Google Cloud AI Platform: Provides managed services for developing and deploying ML models, including tools for data preparation, model training, and prediction.
• Azure Machine Learning: A cloud service for accelerating ML development, offering MLOps features, automated ML, and support for various frameworks.
• Hugging Face Transformers/Accelerate: Open-source libraries that provide models, tools, and training utilities for LLMs and other transformer-based architectures.
• Custom Kubernetes/GPU Clusters: Building and managing bespoke deep learning infrastructure on cloud providers or on-premises, offering maximum control but requiring significant operational overhead.

Getting started

To begin using MosaicML within the Databricks environment, you would typically interact with it through a Databricks Notebook or a Python script executed on a Databricks cluster. The following example demonstrates a basic training script using the Composer library, a core component of MosaicML, to train a simple neural network. This example assumes you have a Databricks workspace configured with the necessary compute resources and the Composer library installed.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

from composer import Trainer
from composer.models import ComposerClassifier
from composer.metrics import Accuracy
from composer.loggers import InMemoryLogger

# 1. Define a simple neural network
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 5)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(5, 2) # Two output classes

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

# 2. Wrap the model for Composer
# ComposerClassifier automatically handles common tasks like loss calculation and metric reporting
model = ComposerClassifier(module=SimpleNet(), num_classes=2)

# 3. Create dummy data
X = torch.randn(100, 10)
y = torch.randint(0, 2, (100,))
dataset = TensorDataset(X, y)
dataloader = DataLoader(dataset, batch_size=16)

# 4. Define optimizer and scheduler (Composer can optimize these too)
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 5. Initialize the Composer Trainer
# Use an InMemoryLogger for simplicity in a quick example
logger = InMemoryLogger()

trainer = Trainer(
    model=model,
    train_dataloader=dataloader,
    eval_dataloader=dataloader, # Using train_dataloader for eval in this example
    optimizers=optimizer,
    max_duration="1ep", # Train for 1 epoch
    loggers=logger,
    device='cpu', # Or 'gpu' if available and properly configured
    # Add more callbacks or algorithms for efficiency, e.g., gradient_clipping, mixed_precision
)

# 6. Train the model
trainer.fit()

print("Training complete.")
# Access logged metrics:
# print(logger.state.logged_metrics)

This snippet demonstrates setting up a basic training loop using Composer. For more advanced scenarios, such as distributed training, mixed-precision, or specific LLM workflows, MosaicML's documentation provides further guidance on configuring trainers with various algorithms and callbacks to optimize performance and cost through their SDK. You would typically run this code within a Databricks Notebook attached to a compute cluster with appropriate GPU resources for real-world deep learning tasks.