modalities 0.3.1

Modalities is a PyTorch-native framework for distributed training of Large Language Models (LLMs) and Foundation Models (FMs) at scale. Given the complexity of distributed training and rapid advancements in the field, we aim to provide a flexible and easy-to-use framework that enables researchers and practitioners to train and evaluate LLMs and FMs efficiently. Modalities is built on top of PyTorch and leverages the latest advancements in distributed training, such as Fully Sharded Data Parallel (FSDP), mixed precision training, Flash Attention and many more, to achieve state-of-the-art performance and throughput.

We successfully scaled Modalities up to 2048 GPUs on two HPC centers, namely Leonardo Booster and MareNostrum 5, featuring Nvidia A100 and H100 GPUs, respectively. The results of our scaling experiments can be found here.

Besides its scalabilty, Modalities allows to seamlessly integrate new components and features, such as custom attention mechanisms, loss functions, optimizers or models. We provide a series of tutorials to help you get started with training and evaluating models using Modalities. We achieve this level of extensibility by having clear interfaces for each component type (e.g., model, optimizer, etc.), that a component must implement to be registered within Modalities at runtime.

Getting Started

For training and evaluation of a model, feel free to checkout this getting started tutorial, in which we train a small, 60M-parameter GPT model on a tiny subset of the Redpajama V2 dataset.

Installation

There are two ways to install Modalities. If you want to use the latest nightly version, or if you want to modify the code base itself, we recommend installing Modalities directly from source.

If you want to use Modalities as a library and register your custom components with Modalities, you can install it directly via pip which provides you with the latest stable version.

Option 1: Installation from source

Create a conda environment and activate it via

conda create -n modalities python=3.10
conda activate modalities

Either clone the repository via

git clone git@github.com:Modalities/modalities.git

or download the repository as a zip file and extract it.

wget https://github.com/Modalities/modalities/archive/refs/heads/main.zip
unzip main.zip

Currently, the flash attention dependency cannot be installed without torch being installed beforehand. Until the flash attention developers fix this, we have to run

pip install torch~=2.4.1

beforehand.

Afterwards, Modalities can be installed via

cd modalities
pip install -e . 

Option 2: Installation via pip

To install Modalities via pip, run

pip install torch~=2.4.1
pip install modalities

Note, that also here, torch has to be installed before installing Modalities due to flash attention's dependency management.

Usage

Modalities provides several entry points to interact with the framework. The following section lists the available entry points and their respective functionalities.

Model Training

For model pretraining, we have to pass a configuration file that specifies the model architecture, optimizer, dataset, dataloader, and other training components. Additionally, we specify the number of nodes, the number of processes per node, and the rendezvous endpoint.

Example:

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun --rdzv-endpoint localhost:29515 \
                                        --nnodes 1 \
                                        --nproc_per_node 4 \
                                        $(which modalities) run --config_file_path configs/pretraining_config.yaml

Explanation:

• CUDA_VISIBLE_DEVICES=0,1,2,3: This environment variable specifies which GPUs will be used for the job. In the example above, the four GPUs with IDs 0, 1, 2, 3 are selected for training.

• torchrun: This is a utility from PyTorch used to launch distributed training. It automatically manages multiple processes for distributed training.

• --rdzv-endpoint localhost:29515: Specifies the rendezvous endpoint. Here, localhost is the machine's address, and 29515 is the port. The rendezvous endpoint coordinates the processes involved in distributed training.

• --nnodes 1: Specifies the number of nodes to be used in the distributed setup. In the example above, a single-node setup is used.

• --nproc_per_node 4: This argument tells torchrun how many processes to launch on each node. In the example above, 4 processes are launched per node, corresponding to the 4 GPUs (IDs 0, 1, 2, 3) specified by CUDA_VISIBLE_DEVICES.

• $(which modalities) run: This part dynamically finds the path to the Modalities executable and runs it. The run command triggers the main process to start the training.

• --config_file_path configs/pretraining_config.yaml: The --config_file_path argument provides the path to the configuration file for the training job. In the example above, it is given by configs/pretraining_config.yaml. A configuraton file contains an exhaustive parameterization for all the training components (e.g., dataset, model, optimizer, etc.), making training fully reproducible. An example configuration file can be found here, and a complete list of components available in Modalities is provided here.

If you are a VSCode user, you may want to add this to your launch.json:

        {
            "name": "Torchrun Main",
            "type": "python",
            "request": "launch",
            "module": "torch.distributed.run",
            "env": {
                "CUDA_VISIBLE_DEVICES": "0,1,2,3"
            },
            "args": [
                "--nnodes",
                "1",
                "--nproc_per_node",
                "4",
                "--rdzv-endpoint=0.0.0.0:29515",
                "src/modalities/__main__.py",
                "run",
                "--config_file_path",
                "config_files/pretraining_config.yaml",
            ],
            "console": "integratedTerminal",
            "justMyCode": true,
            "envFile": "${workspaceFolder}/.env"
        }

It will allow you to run the training endpoint directly from VSCode and debug it.

Raw Training Dataset Indexation

The goal of the indexation process is to determine the starting byte position and length of each document in the raw data file. Subsequently, the index file is used to efficiently access the raw data during tokenization.

Example:

modalities data create_raw_index --index_path data/preprocessed/fineweb_edu_num_docs_483606.idx \
                                               data/raw/fineweb_edu_num_docs_483606.jsonl

Explanation:

The modalities data create_raw_index command triggers the process of creating the index from the raw data. The --index_path argument specifies the location where the generated index file will be saved. In this example, the index will be stored at data/preprocessed/fineweb_edu_num_docs_483606.idx. The last part, i.e., data/raw/fineweb_edu_num_docs_483606.jsonl is the input file in JSONL (JSON Lines) format containing the raw data. The command will process this file to create the index.

Raw Training Dataset Tokenization

Tokenization is the process of converting raw text data into a sequence of tokens that can be used as input to the model. The tokenization requires a configuration file, fully describing the tokenization process, making it fully reproducible. An example tokenization config can be found here.

Example:

modalities data pack_encoded_data configs/tokenization_config.yaml

Inference

For inference on a model checkpoint, we have to pass a configuration file that specifies the full inference setup. An example inference config can be found here.

Example:

modalities generate_text --config_file_path example_text_generation_config.yaml 

Tutorials

Even though Modalities significantly simplifies LLM training, there is still some technical complexity left. We provide a series of tutorials to help you get started with training and evaluating models using Modalities.

• Modalities in 15mins
  Train a dense model with Modalities in 15 minutes

• Getting Started
  Brief overview on how to get started with Modalities by training a small GPT model on a tiny subset of the Redpajama V2 dataset.

• Warmstart
  Continue the training from a checkpoint, e.g., after the training was interrupted or had crashed.

• Library Usage
  How to use Modalities as a library and register custom components with Modalities.

Supported Features

In the following, we list the most important features of Modalities.

Throughput Features

Downstream Performance Features

Scaling Experiments

In the following, you can find the results of our scaling experiments performed on two HPC centers, namely Leonardo Booster and MareNostrum 5.

In a first step, we explored a limited set of different configurations (batch size, gradient accumulation steps, etc.) to get our baseline results. In a second step, we will focus on optimizing these configurations to maximize throughput.

Leonardo Booster - NVIDIA A100 64GB

Further scaling results can be found at Leonardo Booster Scaling Experiments

MareNostrum 5 - NVIDIA H100 64GB

Further scaling results can be found at MareNostrum5 Scaling Experiments

Contributing

Modalities welcomes your contributions! Please check out our contributing guidelines regarding the details on formatting, testing, etc.


Thanks so much to all of our amazing contributors!