Spaces:
Paused
Paused
File size: 8,598 Bytes
0feb305 4c291d2 0feb305 5bb6ad4 65b3f00 5bb6ad4 65b3f00 5bb6ad4 4c291d2 |
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 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 |
---
title: StoryLlama
emoji: 📖
colorFrom: blue
colorTo: purple
sdk: gradio
sdk_version: 5.21.0
app_file: app.py
pinned: false
---
# Introducing StoryLlama - A Smaller Language Model for Bedtime Stories!
- So, I trained a Llama a 88M architecture I coded from ground up to build a small instruct model, going through the below-mentioned stages from scratch.
- Trained on TiyStories dataset form HuggingFace consisting of 4B tokens for a total of 5000 steps
### Pretraining
#### Dataset
- I used the [TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories) dataset from HuggingFace.
1) Train dataset - 2 M records approx
2) Val dataset - 26K records approx
---
#### ModelArgs (Hyperparameters)
Below is a table summarizing the configuration parameters for the model:
| Parameter | Description | Default Value | Type |
|--------------------------------|-----------------------------------------------------------------------------|-----------------------------------|-----------|
| `epochs` | Number of training epochs | `4` | `int` |
| `block_size` | Size of each block (context length) | `512` | `int` |
| `batch_size` | Batch size for training | `64` | `int` |
| `inference` | Inference mode (not specified) | `None` | `None` |
| `embeddings_dims` | Dimensionality of embeddings | `512` | `int` |
| `attn_dropout` | Dropout rate for attention layers | `0.1` | `float` |
| `no_of_heads` | Number of attention heads | `8` | `int` |
| `dropout` | Dropout rate for the model | `0.1` | `float` |
| `val_epochs` | Number of validation epochs | `2` | `int` |
| `max_lr` | Maximum learning rate | `6e-4` | `float` |
| `no_of_decoder_layers` | Number of decoder layers | `8` | `int` |
| `weight_decay_optim` | Weight decay for the optimizer | `0.1` | `float` |
| `beta_1` | Beta 1 for Adam optimizer | `0.9` | `float` |
| `beta_2` | Beta 2 for Adam optimizer | `0.95` | `float` |
| `clip` | Gradient clipping value | `1.0` | `float` |
| `device` | Device to run the model (`cuda` or `cpu`) | `'cuda'` | `str` |
| `no_kv_heads` | Number of key-value heads | `2` | `int` |
| `vocab_size` | Size of the vocabulary | `50304` | `int` |
| `eps` | Epsilon value for numerical stability | `1e-5` | `float` |
| `dtype` | Data type for tensors (`bfloat16` if supported, else `float16`) | `'bfloat16'` or `'float16'` | `str` |
| `save_checkpoint_dir` | Directory to save model checkpoints | `"checkpoints"` | `str` |
| `prompt` | Default prompt for inference | `"Once upon a time"` | `str` |
| `save_checkpoint_iter` | Save checkpoint every N iterations | `50` | `int` |
| `total_iters` | Total number of training iterations | `10000` | `int` |
| `eval_iters` | Evaluate model every N iterations | `50` | `int` |
| `eval_check` | Check evaluation metrics every N iterations | `100` | `int` |
| `warmup_iters` | Number of warmup iterations for learning rate scheduling | `700` | `int` |
| `min_lr` | Minimum learning rate (10% of `max_lr`) | `0.1 * max_lr` | `float` |
| `lr_decay_iters` | Number of iterations for learning rate decay | `10000` | `int` |
| `total_batch_size` | Total batch size across all devices | `524288` | `int` |
| `micro_batch_size` | Micro batch size per device | `batch_size` | `int` |
| `gradient_accumulation_steps` | Gradient accumulation steps | 524288 | `int` |
---
#### Hardware Setup
- Used DPP using Pytorch torchrun consisting of 2x GeForce RTX A100 AXM (80gb VRAM each) rented on runpod.io
- The model is a 0.768GB in size but needs around 4 GB of VRAM when loaded in fp32 precision
---
#### Frameworks:
**Pytorch**
---
#### Epochs/Steps
- Iterations (train) = 5k
- Val iterations = every 50 steps
---
#### Losses
- Train loss - 1.43
- Val loss - 1.45
---
#### Screenshots of the loss curves
- Loss Curves (Train and Val)
---
#### Output
- Prompt: Once upon a time
---
### Local setup
### Requirements
```python
git [clone the repo](https://github.com/YuvrajSingh-mist/StoryLlama.git)
cd StoryLlama
bash ./install.sh
```
- A wandb.ai account for plotting graphs for your loss curves
- On your terminal run
```python
wandb login
```
- Enter the api key and follow the instructions and once you are succesfully logged in follow the given steps
- Download the model
```python
python download_model_weight.py
```
---
### Running
#### Training a model
- Kindly change 'device' to any of your available cuda gpus.
To run:
```python
bash ./install.sh
```
```python
torchrun --standalone --nproc_per_node=gpu trainer.py \
--epochs 10 \
--block_size 256 \
--batch_size 128 \
--embeddings_dims 768 \
--attn_dropout 0.2 \
--no_of_heads 12 \
--dropout 0.2 \
--val_epochs 3 \
--max_lr 5e-4 \
--no_of_decoder_layers 6 \
--weight_decay_optim 0.01 \
--beta_1 0.85 \
--beta_2 0.99 \
--clip 0.5 \
--device "cuda" \
--no_kv_heads 4 \
--vocab_size 50257 \
--eps 1e-6 \
--dtype "float16" \
--save_checkpoint_dir "model_checkpoints" \
--prompt "Once upon a time" \
--save_checkpoint_iter 100 \
--total_iters 5000 \
--eval_iters 200 \
--eval_check 500 \
--warmup_iters 1000 \
--min_lr 1e-5 \
--lr_decay_iters 2000 \
--total_batch_size 262144 \
--micro_batch_size 128 \
--gradient_accumulation_steps 4
```
--standalone - if all the gpu are on one server
--npro_per_node - number of gpus available and use the keyword gpu to use all
#### Inference on a model
```python
python inference.py --prompt "Once upon a time" --max_length 100 --temperature 0.8 --topk 50
``` |