💬 Language model

This repository contains the code to train and test autoregressive language models like ChatGPT from scratch. I also used it to train the french open-source DimensionGPT models.

📋 Summary

• 📋 Summary
• 🤖 DimensionGPT
  • 🏗️ Architecture
  • 💾 Data
  • 🦾 Training
  • 🪛 Fine-tuning
  • 🧪 Tests
  • 🎛️ Weights
• 📦 Dependencies
• 🦾 Training
• ⚗️ Testing
• 🙏 Credits

🤖 DimensionGPT

Using this repository, I trained DimensionGPT-0.2B, a small 0.2B language model on 50B tokens with my personal RTX 3090 GPU during ≈570 hours.

🏗️ Architecture

The model is based on the transformer architecture (only the decoder part) from the paper Attention is All You Need by Google Brain (2017), with a few improvements:

• I replaced the default normalization layer by the Root Mean Square Layer Normalization (RMSNorm) from the paper Root Mean Square Layer Normalization by Edinburgh University (2019)

• I moved the normalization layers before the transformer blocks (instead of after) like in the paper On Layer Normalization in the Transformer Architecture by Microsoft Research (2020)

• I replaced the ReLU activation by the SwiGLU activation from the paper GLU Variants Improve Transformer by Google (2020)

• I implemented Grouped-Query Attention (GQA) from the paper GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints by Google Research (2023)

• I replaced the absolute positional embedding by the Rotary Position Embedding (RoPE) from the paper RoFormer: Enhanced Transformer with Rotary Position Embedding by Zhuiyi Technology (2023)

• I implemented the Sliding Window Attention (SWA) from the paper Longformer: The Long-Document Transformer by Allen Institute (2020)

Here are the main parameters of the architecture:

The resulting model has 208,929,792 trainable parameters and fits on a single RTX 3090 GPU with a batch size of 16 for training using mixed precision. For inference only, the model will probably fit on any modern GPU.

💾 Data

The dataset used to train this model is exclusively in french and is a mix of multiple sources:

For the tokenization, I created my own tokenizer that starts by cleaning the text to keep only a predefined set of characters, then it uses the Byte Pair Encoding (BPE) algorithm to create the vocabulary. I trained the tokenizer on a 300 million characters subset of the dataset to get my 32,000 tokens vocabulary.

🦾 Training

For the training I used stochastic gradient descent with warmup and cosine decay learning rate schedules, here are the main hyperparameters:

I trained the model on my personal RTX 3090 GPU for 1 epoch on the full dataset (13 times the Chinchilla optimal) using mixed precision and gradient accumulation to increase the speed and reduce the memory usage :

🪛 Fine-tuning

I fine-tuned the model on the french instructions dataset I made for this project to create DimensionGPT-0.2B-Chat, a 0.2B language model trained to follow instructions and answer questions in french.

🧪 Tests

Here are some examples of the model outputs:

🎛️ Weights

The trained weights of the different models are available on Google Drive, you just need to:

• Download the .pt file of the model you want to use and put it in the models folder
• Download the vocab.txt file and put it in the data folder

📦 Dependencies

• Python
• PyTorch
• Flash Attention
• Datasets 🤗
• Tokenizers 🤗
• Unidecode
• Regex
• Tqdm
• PSUtil

Run the following command to install the dependencies:

$ pip install -r requirements.txt

⚠️ You may need to use a specific command for PyTorch if you want to use CUDA

⚠️ You way need to manually install a Flash Attention release for Windows

🦾 Training

• Run the create_data.ipynb file to create the tokenizer and the dataset (it may take an entire day and consume a few hundred gigabytes of disk space)

• Run the training.ipynb file (you can stop the training at any time and resume it later thanks to the checkpoints)

• If you don't have an overpriced 24GB GPU like me, the default settings (those used to train DimensionGPT) may not work for you. You can try to:

  • Reduce the batch size (less stable and worse lowest point)
  • Increase the accumulation steps (fix previous problems but slower)
  • Reduce some architecture parameters (worse lowest point)

⚗️ Testing

• Run the testing.ipynb file to use the models you downloaded or trained

🙏 Credits

• Angel Uriot : Creator of the project.