Magix is a mininalist toolkit for training LLM with flexible data and model parallel.
- Training Billion-scale LLM on GPUs and TPUs.
- Familiar Huggingface model interfaces and eco-system (dataset, hub, etc.).
- Pre-defined model parallel (sharding) rules for popular models like Llama, Mistral, Gemma, etc.
- Acceleration with flash attention and operation fusion.
- Fast checkpoint save/restore with arbirary device and para 7D75 llism design.
If you have ever used Huggingface Flax transformers, using magix is as simple as adding several magic functions into the common worflow.
- We start by importing necessary dependencies,
import magix
from magix.models.llama_model import FlaxLlamaForCausalLM
- We will explicitly reason about all the GPU(TPU) devices available to us. We will place the GPUs in a grid (aka mesh) using the
magix.create_device_meshfunction.
# Assume we have 4 GPUs in total; we can arrange them arbitrarily.
# Say, we arrange them into 2x2 mesh and name the first axis `data` and the second axis `model`.
# These axes will be responsible for data and model parallelisms respectively.
mesh = magix.create_device_mesh((2,2), names=('data', 'model'))
- For the next step we will load our model onto the mesh, each device will hold a part (shard) of the full model. Instead of the familiar
from_pretrained, we will use the functionmagix.load_model_hubfunction which will callfrom_pretrainedinternally but also place the model correctly.
model, params = magix.load_model_hub(
FlaxLlamaForCausalLM,
'meta-llama/Llama-2-13b',
FlaxLlamaForCausalLM.partition_rules, # use the pre-defined partitioning
mesh
)
Here params is partitioned and placed on to the mesh. As a side note, JAX will reason about model definition and parameter seperately, analogous to y = f(x|θ).
- For training, you will also need to do something simlar and build the optimizer states onto the mesh,
opt_state = magix.initialize_opt_state(optimizer, params, sharding_config, mesh)
- You may have seen tutorial using
jax.pmap. For our case with both data and model parallelism, we will use the more powerfuljax.jit,
train_step = jax.jit(
train_step, # or generate_step
donate_argnums=... # set based on the actual function input
out_shardings=(magix.item_sharding(params), magix.item_sharding(opt_state),... # set based on the actual function output
)
With all these, you are ready to start your training/inference loop.
Take a look at the complete scripts in train.py, train_lora.py and generate.py.
Assume we have 4 GPUs. Let's train mistral-7b on UltraChat with data and tensor parallism, dp=2 and tp=2 (mesh_shape=2 2):
python train_lora.py \
--checkpoint_dir /absolute/path/to/checkpoint \
--model_type mistral \
--model_name mistralai/Mistral-7B-v0.1 \
--tokenizer_name mistralai/Mistral-7B-v0.1 \
--train_file HuggingFaceH4/ultrachat_200k \
--split train_sft \
--train_data_field messages \
--use_chat_template \
--batch_size 32 \
--num_epochs 1 \
--learning_rate 5e-5 \
--seed 12345 \
--mesh_shape 2 2 \
--weight_decay 0.001 \
--max_length 1024
After training, let's solve some math problems. Do generation with full tensor parallel tp=4 (mesh_shape=1 -1):
python generate.py \
--prompts gsm8k \
--hf_data_config main \
--hf_data_split test \
--use_chat_template \
--data_field question \
--output_file generation.jsonl \
--mesh_shape 1 -1 \
--model_type mistral \
--model_name_or_path mistralai/Mistral-7B-v0.1 \
--tokenizer_name_or_path mistralai/Mistral-7B-v0.1 \
--model_config_name mistralai/Mistral-7B-v0.1 \
--batch_size 32 \
--pad_to_multiple_of 64 \
--max_length 512 \
--lora /absolute/path/to/checkpoint/EVALUATION_STEP/lora
We recommend using the jax-toolbox jax container image from nvidia. We have example Dockerfile and Singulrity Definition File.
Install appropriate jax build, torch-cpu and then the rest of the dependencies.
pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
# get torch-cpu for model conversion
pip install torch --index-url https://download.pytorch.org/whl/cpu
git clone https://github.com/luyug/magix.git
cd magix
pip install -e .