File size: 6,583 Bytes
b66f552 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | <div align="center">
# 🔥 Flame: Flash Linear Attention Made Easy
</div>
> [!IMPORTANT]
> The `flame` project has been migrated to a new project built on torchtitan.
> Please visit the [new repository](https://github.com/fla-org/flame) for details and updates.
>
> The code here is now **archived as legacy**, and no future updates will be synchronized here.
A minimal framework for training FLA models, whether from scratch or through finetuning.
Built on the robust infrastructure of 🤗, `flame` enables you to train large language models with just a few lines of code:
we use `datasets` for data processing, `transformers` for model definitions, and `accelerate`[^1] for seamless distributed training.
In this README, we will guide you through the process of using `flame` to train GLA models.
## Setup
To get started, you'll need to install the required packages.
Both `fla` and `flame` have minimal dependencies.
Clone the `fla` repository and install the necessary packages as follows:
```bash
git clone https://github.com/sustcsonglin/flash-linear-attention.git
pip install .
pip install accelerate
```
> [!CAUTION]
> The 🤗 `tokenizers` have some [memory leak issues](https://github.com/huggingface/tokenizers/issues/1539) when processing very long documents.
> To address this, please ensure you install `tokenizers>=0.20.4`.
## Preprocessing
Before training, you need to download and pre-tokenize your dataset.
We provide a straightforward script for this.
For instance, to tokenize a 10B sample of the `fineweb-edu` dataset, run:
```bash
python preprocess.py \
--dataset HuggingFaceFW/fineweb-edu \
--name sample-10BT \
--split train \
--context_length 2048
```
This will cache the processed dataset at `data/HuggingFaceFW/fineweb-edu/sample-10BT/train`.
GLA utilizes a subset of Slimpajama for pretraining [in the paper](https://proceedings.mlr.press/v235/yang24ab.html).
Given the size of the dataset, the fastest way to download it is using `git lfs` (refer to [this issue](https://huggingface.co/datasets/cerebras/SlimPajama-627B/discussions/2)).
```bash
git lfs install
git clone https://huggingface.co/datasets/cerebras/SlimPajama-627B --depth 1
python preprocess.py \
--dataset SlimPajama-627B \
--split train \
--context_length 2048
```
## Training from scratch
To train your 340M model from scratch, execute the following command:
```bash
bash train.sh \
type=gla \
lr=3e-4 \
scheduler=cosine_with_min_lr \
batch=32 \
update=1 \
warmup=1024 \
steps=20480 \
context=2048 \
gpus=8 \
nodes=1 \
path=exp/gla-340M-10B \
project=fla \
model=configs/gla_340M.json \
data=HuggingFaceFW/fineweb-edu \
name=sample-10BT \
cache=data/HuggingFaceFW/fineweb-edu/sample-10BT/train
```
Key parameters:
| | Description | Default |
| :-------- | :---------------------------- | -------------------- |
| lr | `learning_rate` | `3e-4` |
| scheduler | `lr_scheduler_type` | `cosine_with_min_lr` |
| batch | `batch_size` | `32` |
| update | `gradient_accumulation_steps` | `1` |
| context | `context_length` | `2048` |
| gpus | `num_gpus_per_node` | `8` |
| nodes | `num_nodes` | `1` |
| warmup | `warmup_steps` | `1024` |
| steps | `max_steps` | `20480` |
The learning rate is set to `3e-4` by default, equipped with a cosine scheduler.
Other scheduler types like WSD (`warmup_stable_decay`)[^2] are also supported.
The total number of tokens processed per batch, referred to as `global_batch_size`, is calculated as
`batch_size × gradient_accumulation_steps × context_length × num_gpus_per_node × num_nodes`.
For instance, in the 340M model example, the `global_batch_size` calculates to $32 \times 1 \times 2048 \times 8 \times 1 = 524,288$ (0.5M tokens).
The `warmup_steps` parameter indicates the number of steps for the learning rate warmup phase, while `max_steps` represents the maximum number of training steps.
Each step processes `global_batch_size` tokens.
Consequently, `512` and `20480` correspond to processing 0.5B and 10B tokens, respectively.
:warning: Monitor the value of `global_batch_size`, `warmup_steps`, and `max_steps` carefully when modifying any of the hyperparameters!!
`flame` also supports resuming interrupted training by specifying the checkpoint path.
Simply use the following command:
```bash
bash train.sh \
type=gla \
lr=3e-4 \
steps=20480 \
batch=32 \
update=1 \
warmup=1024 \
context=2048 \
gpus=8 \
nodes=1 \
path=exp/gla-340M-10B \
project=fla \
model=configs/gla_340M.json \
data=HuggingFaceFW/fineweb-edu \
name=sample-10BT \
cache=data/HuggingFaceFW/fineweb-edu/sample-10BT/train \
checkpoint=exp/gla-340M-10B/checkpoint-8192
```
You can also use `wandb` to monitor your training process effectively.
## Continual Pretraining
`flame` supports continual training from a pretrained checkpoint.
Below, we provide an example of how to finetune Mistral-7B to GLA.
You can follow similar steps to reproduce the results in the [GSA paper](https://arxiv.org/abs/2409.07146):
1. Initialize a brand-new GLA-7B model from the config and copy the mathced pretrained weights from Mistral-7B:
```bash
cd ../utils
python convert_from_llama.py \
--model mistralai/Mistral-7B-v0.1 \
--config ../training/configs/gla_7B.json \
--output ../training/converted/gla-7B
cd -
```
2. Directly launch training from the converted checkpoint:
```bash
bash train.sh \
type=gla \
lr=3e-5 \
steps=10240 \
batch=4 \
update=8 \
warmup=512 \
context=2048 \
path=exp/gla-7B-20B \
project=fla \
model=converted/gla-7B \
data=SlimPajama-627B \
cache=data/SlimPajama-627B/train
```
Please be aware that finetuning on a single node may not be the most efficient approach.
If available, consider leveraging multi-node GPUs for optimal performance.
You can find guidance on how to launch a multi-node job in the [accelerate tutorial](https://github.com/huggingface/accelerate/blob/main/examples/slurm/submit_multinode.sh).
[^1]: The `accelerate` library supports various distributed frameworks, like `deepspeed` and `megatron` for large-scale training. We use `deepspeed` in our case.
[^2]: https://arxiv.org/abs/2404.06395
|