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
• 🤖 DimensionGPT
  • 🏗️ Architecture
  • 💾 Data
  • 🦾 Training
  • 🪛 Fine-tuning
  • 🧪 Tests
  • 🎛️ Weights
• 📦 Dependencies
• 🦾 Training
• ⚗️ Testing
• 🙏 Credits

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.

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.

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.

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 :

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.

Here are some examples of the model outputs:

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

• 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

• 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)

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

• Angel Uriot : Creator of the project.