LLM Foundry is a codebase for training, fine-tuning, evaluating, and deploying large language models for inference, built on Composer and the MosaicML platform, with the goal of being easy to use, efficient, and flexible, supporting rapid experimentation with the latest techniques.
Main directories in the repository: llmfoundry/ contains source code for models, datasets, callbacks, utilities, etc.; scripts/ contains scripts for various LLM workloads: data_prep/ converts raw text into StreamingDataset format; train/ trains or fine-tunes HuggingFace and MPT models, with parameter ranges from 125M to 70B, and train/benchmarking is used to analyze training throughput and MFU; inference/ converts models to HuggingFace or ONNX format and generates responses, and inference/benchmarking is used to analyze inference latency and throughput; eval/ evaluates LLMs on academic or custom in-context learning tasks; mcli/ uses MCLI and the MosaicML platform to launch the above workloads; TUTORIAL.md provides in-depth tutorials, example workflows, and FAQs.
DBRX is the latest open-source LLM trained by the Databricks Mosaic team, using a Mixture-of-Experts (MoE) architecture and trained with optimized versions of Composer, LLM Foundry, and MegaBlocks, with 132B total parameters and 36B active parameters, and two DBRX models have been released. The model weights and code are available under license to both researchers and commercial entities; see LICENSE for licensing details, see the relevant link for the acceptable use policy, and for more information see https://github.com/databricks/dbrx .
Mosaic Pretrained Transformers (MPT) are GPT-style models with Flash Attention (efficiency), ALiBi (context length extrapolation), and stability improvements that mitigate loss spikes. As part of the MosaicML Foundation series, multiple MPT models have been open-sourced. To try locally, follow the instructions in scripts/inference/README.md and use the hf_generate.py or hf_chat.py scripts to prompt HF models. Community contributions include: ReplitLM (replit-code-v1-3b, a 2.7B causal language model focused on code completion, trained on a subset of Stack Dedup v1.2, covering 20 languages including Java, Python, C++, etc.), LLaVa-MPT (visual instruction tuning, giving MPT multimodal capabilities), ggml (an optimized MPT version for efficient inference on consumer hardware), GPT4All (a locally running chat system that now supports MPT), and Q8MPT-Chat (an 8-bit optimized CPU version of MPT provided by Intel). Community tutorial videos include: Using MPT-7B with Langchain, Introduction to MPT-7B StoryWriter, Fine-tuning MPT-7B on a single GPU, and Fine-tuning MPT-7B-Instruct on Google Colab. Contributions via PR are welcome.
Related blogs: Introducing DBRX: A New State-of-the-Art Open LLM; LLM Training and Inference with Intel Gaudi2 AI Accelerators; Training LLMs at Scale with AMD MI250 GPUs; Training LLMs with AMD MI250 GPUs and MosaicML; Announcing MPT-7B-8K: 8K Context Length for Document Understanding; MPT-30B: Raising the bar for open-source foundation models; Introducing MPT-7B; Benchmarking LLMs on H100; Blazingly Fast LLM Evaluation; GPT3 Quality for $500k; Billion parameter GPT training made easy.
Hardware and software requirements: This codebase has been tested on NVIDIA A100 and H100 using PyTorch 2.4. It may also work on other devices such as consumer NVIDIA cards and AMD cards, but has not been actively tested. If you succeed or fail using it on other systems, please let us know in a Github issue so that the support matrix can be updated.
MosaicML Docker images: It is strongly recommended to use the prebuilt images, see https://hub.docker.com/orgs/mosaicml/repositories The mosaicml/pytorch image pins specific PyTorch and CUDA versions, is stable, and is rarely updated. The mosaicml/llm-foundry image builds a new tag on every main branch commit; you can choose a specific commit hash, such as mosaicml/llm-foundry:2.7.0_cu128-9867a7b, or use the latest version mosaicml/llm-foundry:2.7.0_cu128-latest. Note: the mosaicml/llm-foundry image does not come with the llm-foundry package preinstalled, only its dependencies, so you still need to run pip install llm-foundry from PyPi or from source. It is assumed that PyTorch, CMake, and packaging are already installed; otherwise, you can install them with pip install cmake packaging torch.
To get started, clone the repository and set up the environment. Using Docker (recommended): it is strongly recommended to work inside a Docker container (see the recommended images above), and follow the steps to clone the repository and install dependencies. Not using Docker (not recommended): you should create and use a virtual environment.
TransformerEngine and amp_fp8 support: NVIDIA H100 GPUs support FP8; Flash Attention and Transformer are already installed in the Docker image. If you do not use the Docker image, you can install these packages. For more details on enabling TransformerEngine layers and amp_fp8, see the relevant links.
AMD (BETA support): in AMD GPU testing, the environment setup includes... Finally, install the ROCm-enabled flash attention (see the link for instructions). There is not yet a Docker image in which everything works perfectly, and you may need to upgrade/downgrade some packages (for example, you may need to downgrade to numpy==1.23.5) to run without issues.
LLM Foundry support on Intel Gaudi devices is experimental; please use the habana_alpha branch and see that branch's README, which contains installation instructions and known issues. For training and inference performance results on Intel Gaudi2 accelerators, see the blog:https://www.databricks.com/blog/llm-training-and-inference-intel-gaudi2-ai-accelerators Note: be sure to complete the installation steps above before trying the quick start.
End-to-end workflow example: prepare a subset of the C4 dataset, train an MPT-125M model for 10 batches, convert the model to HuggingFace format, evaluate the model on the Winograd challenge, and generate responses to prompts (this is only a quick start to demonstrate the tooling; to get good quality you need to train for more than 10 batches). Note: the composer command used in the above training refers to the distributed launcher of the Composer library. If you have a writable HuggingFace authentication token, you can optionally upload the model to the Hub: export the token, and when calling inference/convert_composer_to_hf.py, uncomment the line containing --hf_repo_for_upload ...
Registry: you can use the registry to customize workflows without forking the library. Some components of LLM Foundry can be registered, such as models, loggers, and callbacks. This means you can register new options for these components and then use them in yaml configuration. To discover registrable components: you can use the llmfoundry registry cli command. Two commands are currently provided: llmfoundry registry get [--group] lists all registries and their components, and you can optionally specify a particular registry, for example: llmfoundry registry get --group loggers or llmfoundry registry get; llmfoundry registry find <group> <name> gets information about a specific registered component, for example: llmfoundry registry find loggers wandb. Use --help with any command for more information. These commands can also help you understand the composition of each registry, because each registry contains a docstring that will be printed. The overall concept is that each registry defines an interface, and components registered to that registry must implement that interface. If some part of the library is currently not extensible but you think it should be, please open an issue!
How to register: there are several ways to register new components:
- Python entry points: if you build your own package with registered components, you can specify them through Python entry points. If you are building a large extension and overriding many components, this is the intended usage. Note that anything registered through entry points will override directComponents registered in the code. For example, the following registers the MyLogger class with the llm_foundry.loggers registry under the key my_logger. If you develop a new component via an entry point, note that Python entry points are global to the Python environment, which means that if multiple packages register a component with the same key, the last one installed will be used. This can be useful for overriding components in LLM Foundry, but can also lead to unexpected behavior if you are not careful. In addition, if you change pyproject.toml, you need to reinstall the package for the changes to take effect. You can quickly install with pip install -e . --no-deps to avoid reinstalling dependencies.
- Call register directly: You can also register a component directly in code.
- Decorator: You can also use a decorator to register a component directly from code. For both direct calls and the decorator approach, if you use the LLM Foundry training/evaluation scripts, you need to provide the code_paths parameter, which is the list of files that need to be executed to register the component. For example, you might have a file named foundry_imports.py that contains the relevant content. Then provide code_paths for the training/evaluation script in the yaml configuration. This is the intended usage if you are building a small extension, overriding only a few components, and do not want to create an entire package.
Learn more: See TUTORIAL.md to continue learning about LLM Foundry; the tutorial focuses on example workflows, points to other resources in the repository, and answers common questions. If you encounter problems with the code, please file a Github issue directly in this repository. If you want to train LLMs on the MosaicML platform, contact demo@mosaicml.com.
