# mt_dao: Federated LLM Pre-Training

mt_dao is a framework for federated pre-training of Large Language Models (LLMs). This README provides instructions on how to set up and run the mt_dao framework.

<p align="center">
  <a href="https://flower.ai/">
    <img src="https://flower.ai/_next/image/?url=%2F_next%2Fstatic%2Fmedia%2Fflwr-head.4d68867a.png&w=384&q=75" width="140px" alt="Flower Website" />
  </a>
</p>
<p align="center">
    <a href="https://flower.ai/research">Website</a> |
    <a href="https://flower.ai/blog">Blog</a> |
    <a href="https://flower.ai/docs/">Docs</a> |
    <a href="https://flower.ai/events/flower-ai-summit-2025">Summit</a> |
    <a href="https://flower.ai/join-slack">Slack</a>
    <br /><br />
</p>

[![DOI](https://zenodo.org/badge/963528824.svg)](https://doi.org/10.5281/zenodo.15187915)
[![GitHub license](https://img.shields.io/github/license/anon/mt_dao)](https://github.com/anon/mt_dao/blob/main/LICENSE)
[![PRs Welcome](https://img.shields.io/badge/PRs-welcome-brightgreen.svg)](https://github.com/adap/anon/mt_dao/main/CONTRIBUTING.md)
![Build](https://github.com/adap/flower/actions/workflows/framework.yml/badge.svg)
[![Slack](https://img.shields.io/badge/Chat-Slack-red)](https://flower.ai/join-slack)

## Table of Contents

- [mt\_dao: Federated LLM Pre-Training](#mt_dao-federated-llm-pre-training)
  - [Table of Contents](#table-of-contents)
  - [Citations](#citations)
  - [System Requirements](#system-requirements)
  - [Installation](#installation)
    - [System Setup](#system-setup)
    - [Environment Installation](#environment-installation)
  - [Running mt\_dao](#running-mt_dao)
    - [Dataset Preparation](#dataset-preparation)
    - [Federated Training](#federated-training)
    - [Centralized Training](#centralized-training)
    - [Evaluation](#evaluation)

## Citations

```bib
@inproceedings{
anon2025mt_dao,
title={mt_dao: Federated {LLM} Pre-Training},
author={anon anon, anon anon, Zeyu Cao, Royson Lee, Bill Marino, Yan Gao, Wanru Zhao, Dongqi Cai, Zexi Li, Xinchi Qiu and Nicholas Donald Lane },
booktitle={Eighth Conference on Machine Learning and Systems},
year={2025},
url={https://openreview.net/forum?id=AQgYcfg5EI}
}
```


```bib
@inproceedings{
anon2025dept,
title={{DEPT}: Decoupled Embeddings for Pre-training Language Models},
author={anon anon and anon anon and Meghdad Kurmanji and William F. Shen and Xinchi Qiu and Dongqi Cai and Yan Gao and Nicholas Donald Lane},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
url={https://openreview.net/forum?id=vf5aUZT0Fz}
}
```


```bib
@inproceedings{
anon2024worldwide,
title={Worldwide Federated Training of Language Models},
author={anon anon and anon anon and Bill Marino and Preslav Aleksandrov and William F. Shen and Nicholas Donald Lane},
booktitle={International Workshop on Federated Foundation Models in Conjunction with NeurIPS 2024},
year={2024},
url={https://openreview.net/forum?id=YMSLZUmQVV}
}
```

```bib
@inproceedings{
anon2024the,
title={The Future of Large Language Model Pre-training is Federated},
author={anon anon and anon anon and Zeyu Cao and Bill Marino and Yan Gao and Tomas Paulik and Wanru Zhao and William F. Shen and Preslav Aleksandrov and Xinchi Qiu and Nicholas Donald Lane},
booktitle={International Workshop on Federated Foundation Models in Conjunction with NeurIPS 2024},
year={2024},
url={https://openreview.net/forum?id=hfeH5AP9NY}
}
```

## System Requirements

- Ubuntu 22.04 (or compatible Linux distribution)
- NVIDIA GPU with CUDA 12.4 support
- Python 3.11.9
- At least 16GB of GPU memory for 125M parameter models (more for larger models)

## Installation

### System Setup

The `system_setup.sh` script sets up the system environment for mt_dao, including CUDA, CuDNN, and Python:

```bash
cd scripts
. system_setup.sh
```

> Note that you may want to inspect the script prior executing it to remove parts of the installation that are not necessary for your system. For example, if you already have CUDA installed, you may want to remove the CUDA installation part from the script.
>
> Note that the `system_setup.sh` script will define some environmental variables that are going to be used in the next installation scripts, thus it is wise to execute it with the "." in case you run the following scripts in the same bash session. Alternatively, you can execute it and then source the .bashrc file.
>
> Note that the logic for storing AWS-S3 and Wandb credentials is also provided even though it is not strictly necessary for the framework to run. If one prefers to log to Wandb or use AWS-S3, one needs to set up such credentials and modify the configuration parameters in the example scripts accordingly.

This script:

- Installs essential system packages
- Installs CUDA 12.4 and sets up environment variables
- Sets GPU persistence mode
- Installs CuDNN
- Sets up pyenv and installs Python 3.11.9
- Installs necessary monitoring utilities
- [Optional] Sets up AWS S3 credentials for dataset streaming (you'll need to add your credentials)
- [Optional] Sets up Weights & Biases (wandb) for experiment tracking (you'll need to add your credentials)

### Environment Installation

After system setup, install the mt_dao environment:

```bash
. install_env.sh --project_path /path/to/mt_dao
```

Parameters:

- `--project_path` or `-p`: Path to the mt_dao project directory (default: `$HOME/projects/mt_dao`)

> Note that the script will install manually the flash-attention library, because uv has difficulty in installing it by matching the dependencies. This will incur in compiling the library from source, which might take a while depending on the system.
>
> Note that some NVIDIA GPUs don't support the `flash-attention` library. In this case, you can remove the `flash-attention` dependency from the `pyproject.toml` file and run `uv sync` again. You will need to modify accordingly the configuration parameters through hydra replacing `llm_config.model.attn_config.attn_impl=flash` with `llm_config.model.attn_config.attn_impl=torch`.
>
> Note that, similarly to the above, some NVIDIA GPUs don't support executing mixed precision context with bf16. In this case, we recommend switching to standard fp16 or fp32 by replacing `llm_config.precision=amp_bf16` with `llm_config.precision=amp_fp16` or `llm_config.precision=fp32` respectively.

This script:

- Sets up a uv environment with all required dependencies
- Installs flash-attention for GPU acceleration
- Sets up necessary environment variables

## Running mt_dao

### Dataset Preparation

To prepare datasets for training and evaluation:

```bash
# Convert C4 dataset for training
bash convert_c4_dataset.sh --project_path /path/to/mt_dao

# Prepare evaluation datasets
bash prepare_eval_dataset.sh --project_path /path/to/mt_dao
```

You will need to follow a similar procedure to get the smollm mixture, taking the original data-sources and splitting them in IID fashion.

> Note that the `convert_c4_dataset.sh` script will download a small shard of the C4 dataset from Hugging Face and convert it to the mt_dao-compatible format and partition it in 8 clients. This might take a while depending on the internet speed and CPU parallelization capabilities of the machine.
>
> Note that the `prepare_eval_dataset.sh` script will download the Gauntlet dataset from a public fork of [this](https://github.com/mosaicml/llm-foundry) repository and move it to the `dataset_cache` folder. This is executed through the command `git clone git@github.com:<account>/llm-foundry.git`, but one may want to change it to `git clone https://github.com/<account>/llm-foundry.git` instead.

### Federated Training

To run federated training of LLMs with mt_dao, use the `mt_dao_llm_125M.sh` script:

```bash
bash mt_dao_llm_125M.sh --project_path /path/to/mt_dao
```

This script sets up and runs federated training for a 125M parameter model. The script:

- Activates the uv environment
- Sets up the necessary environment variables
- Configures the federated training parameters
- Launches the server and client nodes for training

We recommend running the following example with a more advanced federated training configuration:

```bash
bash fed_125m_example.sh --project_path /path/to/mt_dao
```

### Centralized Training

For centralized (non-federated) training, use:

```bash
bash centralised_training.sh --project_path /path/to/mt_dao [MODEL_SIZE]
```

Parameters:

- `--project_path` or `-p`: Path to the mt_dao project directory
- `MODEL_SIZE`: Size of the model to train (e.g., "125m", "350m")

We recommend running the following example for a small 125M parameter model with centralized training:

```bash
bash cen_125m_example.sh --project_path /path/to/mt_dao
```

### Evaluation

To evaluate trained models:

```bash
bash eval_gauntlet_only.sh --project_path /path/to/mt_dao [MODEL_PATH]
```

Parameters:

- `--project_path` or `-p`: Path to the mt_dao project directory
- `MODEL_PATH`: Path to the trained model checkpoint

We recommend running the following example for evaluating the 125M parameter model on the Gauntlet:

```bash
bash eval_125m_example.sh --project_path /path/to/mt_dao
```


Assuming you have successfuly setup an s3 and downloaded the smollm mixture partitioned in 16 shards per dataset, as indicated by ```bash 
mt_dao/conf/dataset/streams/smollm_corpus_4_clients_proportional.yaml```
then you can use the scripts in the scripts folder, specificaly base_centralised_training and ```scripts/base_launchers/mt_dao_base.sh``` together with the hyperparameters specified in the paper to recreate our experiments, an example of how this may be achieved is shown in scripts/iclr_mt_dao.