YAML Metadata Warning: empty or missing yaml metadata in repo card (https://huggingface.co/docs/hub/model-cards#model-card-metadata)

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

  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)

Loss Curves (Train and Val)


Output

  • Prompt: Once upon a time

Prompt: Once upon a time


Local setup

Requirements

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

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 download_model_weight.py

Running

Training a model

  • Kindly change 'device' to any of your available cuda gpus.

To run:

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 inference.py --prompt "Once upon a time" --max_length 100 --temperature 0.8 --topk 50 
Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model is not currently available via any of the supported Inference Providers.
The model cannot be deployed to the HF Inference API: The model has no library tag.