MosaicML

MosaicML developed software and infrastructure for organizations to train and deploy AI models on their own data. Its May 2023 MPT-7B release combined open model weights with training and evaluation code, demonstrating the platform customers could use to build custom models. Databricks acquired the company in July 2023; the Mosaic team subsequently led development of Databricks’ DBRX model.[3][6][8]

Chief Scientist Jonathan Frankle described the business in a May 2023 interview: “Our business is built around customers creating their own models.” MPT served as a demonstration of those tools, whether customers fine-tuned a released checkpoint or used its configuration to train from scratch. Research Scientist Abhinav Venigalla spearheaded the MPT-7B training run, while MosaicML’s data and NLP teams contributed to the broader release.[1][3]

That custom-model platform followed an earlier focus on training efficiency. MosaicML’s launch-era products supplied tools for composing training methods and comparing the resulting tradeoffs.[2][3]

MosaicML emerged from stealth on October 13, 2021 with $37 million in funding, according to contemporaneous reporting by Nicole Hemsoth Prickett. Named investors included Lux Capital, Future Ventures and E14. Naveen Rao and Hanlin Tang had previously worked at Nervana Systems and Intel; Rao was MosaicML’s co-founder and CEO.[2][5]

The early products addressed different parts of the optimization problem. Composer combined training methods into recipes based on benchmarked findings and published research. Explorer let users choose tradeoffs among accuracy, cost and speed, then visualize those tradeoffs across thousands of benchmark training runs. By May 2023, Composer remained at the heart of MosaicML’s stack, with LLM Foundry providing model configurations and code for pretraining, fine-tuning and evaluation.[2][1][3]

Released on March 9, 2023, MosaicBERT adapted BERT for efficient pretraining on English C4 using masked-language modeling: the model predicted hidden tokens in a text sequence. It combined FlashAttention, ALiBi, unpadding, low-precision LayerNorm and gated linear units. Unpadding avoided computation on empty padding tokens; FlashAttention reduced transfers between GPU memory levels. The FlashAttention implementation came from Hazy Research using OpenAI’s Triton.[4]

MosaicML reported that MosaicBERT-Base reached the original BERT’s 79.6 average GLUE score after 1.28 hours of pretraining on eight A100-40GB GPUs—about $20 at an assumed $2 per GPU-hour. Eight A100-80GB GPUs took 1.13 hours, or about $22 at $2.50 per GPU-hour. GLUE scores followed task-specific fine-tuning; the quoted prices covered pretraining to that target. MosaicBERT-Base had 137 million parameters, compared with BERT-Base’s 110 million.[4]

The report was authored by Jacob Portes, Alex Trott, Daniel King, Sam Havens and Erica Ji Yuen. It credited Trott, Portes and Havens with experiments and Trott with the final architecture, with additional improvement contributions from Abhinav Venigalla, Moin Nadeem, Nikhil Sardana, Vlad Ivanchuk, Ishana Shastri, Davis Blalock and Daya Khudia.[4]

MosaicML released MPT-7B on May 5, 2023 as the first entry in its Foundation Series. The base model was a 6.7-billion-parameter decoder-style transformer pretrained on one trillion tokens of text and code. Venigalla said the team trained beyond the compute-optimal Chinchilla point to make a relatively small model useful for inference. MosaicML also released Instruct, Chat and StoryWriter variants, alongside the LLM Foundry codebase.[3][1]

Training used Composer, PyTorch Fully Sharded Data Parallelism, LLM Foundry and StreamingDataset on Oracle Cloud A100 GPUs and object storage. StreamingDataset streamed data without requiring a complete download and supported resumption from any point, preserving sample order across changes in GPU, node or worker counts. MosaicML used EleutherAI’s GPT-NeoX-20B tokenizer and acknowledged AI2’s assistance with data curation and tokenizer selection.[3]

MosaicML reported that the base run took 9.5 days on 440 A100-40GB GPUs, with no human intervention despite four hardware failures. The approximately $200,000 compute cost was calculated at reserved platform pricing of $2 per GPU-hour; runtime included checkpointing, evaluation and restarts. NodeDoctor monitored and resolved hardware failures, while JobMonitor resumed the jobs. Separately, MosaicML attributed resistance to catastrophic loss spikes to architectural and optimization changes, including ALiBi and the Lion optimizer.[3]

MosaicML reported similar MPT-7B and LLaMA-7B performance across its twelve zero-shot academic tasks, evaluated through the same LLM Foundry framework without model-specific prompt tuning. Frankle nevertheless argued that multiple-choice benchmarks did not adequately capture open-ended generation. During training, the team also watched responses to recurring prompts, including one asking for games a three-year-old and seven-year-old could play together. In the May interview, he estimated the research team at about fifteen researchers.[3][1]

StoryWriter-65k+ was fine-tuned for 2,500 steps on 65,000-token fiction excerpts from books3. ALiBi replaced positional embeddings with biases in the attention operation, allowing extrapolation beyond the training length. MosaicML demonstrated this by supplying the full text of The Great Gatsby—about 68,000 tokens—followed by “Epilogue.” The company reported testing up to 84,000 tokens; this was a StoryWriter demonstration rather than the base model’s trained context length.[3][1]

The release listed Base and StoryWriter under Apache-2.0, Instruct under CC-By-SA-3.0, and Chat under CC-By-NC-SA-4.0, permitting non-commercial use only. Frankle recalled briefly changing StoryWriter’s license after criticism of books3 and internal discussion about releasing creative-generation software during the Hollywood writers’ strike. He said the team restored Apache-2.0 and decided to leave the model available. Those license listings did not settle the legal questions around training data.[3][1]

Frankle said MosaicML initially concentrated on efficient training because it believed inference was already well served. By May 2023, customer requests had persuaded it to add inference so users could serve models within the same ecosystem. The service offered managed model endpoints and deployment of custom models; Venigalla described per-token charging for its starter tier.[1][3]

On June 26, 2023, Databricks announced an agreement to acquire MosaicML at an approximately $1.3 billion transaction value, including retention packages. Databricks reported more than 3.3 million MPT-7B downloads and described MPT-30B as recently released. It named AI2, Generally Intelligent, Hippocratic AI, Replit and Scatter Labs as MosaicML customers; the download figure measured downloads rather than unique users.[5]

Databricks completed the acquisition on July 19, 2023. At closing, it announced that MosaicML’s product would remain available for direct purchase and that the teams would work to simplify custom-model training for Databricks Lakehouse customers.[6]

Databricks released DBRX on March 27, 2024. Development was led by the Mosaic team that had built MPT, working with contributors across Databricks. The training stack used optimized versions of MegaBlocks, LLM Foundry, Composer and Streaming. Databricks acknowledged collaborators including Trevor Gale and MegaBlocks, PyTorch/FSDP, NVIDIA/TensorRT-LLM, vLLM, EleutherAI, Lilac AI’s Daniel Smilkov and Nikhil Thorat, and AI2.[8]

DBRX used a mixture-of-experts architecture with 132 billion total parameters and 36 billion active parameters, selecting four of sixteen experts for an input. Databricks reported pretraining on twelve trillion tokens and training on 3,072 NVIDIA H100 GPUs. Its reported three-month main development process included pretraining, post-training, evaluation, red-teaming and refinement. This was post-acquisition Databricks model scale, extending the earlier MPT work and software stack.[8]