Hugging Face's logo

andrewdalpino
/
LightGPT

Text Generation
PyTorch
TensorBoard
ONNX
Safetensors
English
NoPE
Model card Files Metrics Community
LightGPT / README.md
Andrew DalPino
A little nicer
fc4824e
preview code
|
raw
history blame
13.2 kB
metadata
license: apache-2.0
datasets:
  - HuggingFaceFW/fineweb
  - tatsu-lab/alpaca
language:
  - en
metrics:
  - perplexity
pipeline_tag: text-generation
tags:
  - LightGPT

LightGPT

LightGPT is a lightweight generative pre-trained Transformer (GPT) model for the people! Built using pure PyTorch, LightGPT can answer questions, follow instructions, summarize documents, chat, and much more. Best of all, the model weights and code are fully open-source for you to customize, improve upon, and share with the world.

Features

  • Parameter-efficiency: LightGPT aims to be a more parsimonious model by only training parameters that are absolutely necessary. As such, biases and positional embeddings have been completely removed from the architecture. Despite having no positional embeddings (NoPE), LightGPT performs better at context-length generalization than relative embeddings offering good performance even at 2X of the trained context window.

  • Low Memory Utilization: LightGPT lets you progressively employ training-time memory optimizations such as fully-sharded data-parallel (FSDP), activation checkpointing, mixed precision, and low-memory optimizer updates that allow you to train larger models on smaller hardware.

  • Fully Open-source: Unlike closed-source LLMs, LightGPT provides both the model weights and the source code to train, fine-tune, export, and generate text from the model using your own hardware. With the help of the open-source software community, we aim to democratize access to AI and continually improve the models.

Suggested Pre-training Configurations

Below is a table of some suggested pre-training configurations but feel free to experiment with settings on your own. See the model_sizing.ipynb notebook to estimate the memory and compute requirements for your model configuration.

Name Vocab. Size Block Size Embedding Dim. Attn. Heads Layers Parameters Training Tokens
Small 50,257 1024 1024 16 24 353M 7B
Medium 50,257 1024 2048 32 32 1.7B 34B
Large 100,275 2048 4096 64 32 6.8B 132B
X-large 100,275 2048 4096 64 64 13B 262B
XX-large 200,017 4096 8192 128 64 53B 1T
XXX-large 200,017 4096 8192 128 128 105B 2T

Install Project Dependencies

Project dependencies are specified in the requirements.txt file. You can install them with pip using the following command from the project root. We recommend using a virtual environment such as venv to keep package dependencies on your system tidy.

python -m venv ./.venv

source ./.venv/bin/activate

pip install -r requirements.txt

Pre-training

For the pre-training corpus we use the Fineweb dataset which consists of about 15T high-quality tokens gathered from the worldwide web. The dataset has been split into 3 subsets (10BT, 100BT, and 350BT versions) for training smaller models. If you'd like to start training right away, the default settings should work on most single-GPU systems with 12G of VRAM or more.

python pre-train.py

Note that it will take a while to download and pre-process the dataset the first time that the training script is run.

To customize the default "Small" architecture you can adjust the block_size, embedding_dimensions, num_hidden_layers, and num_attention_heads arguments of the pre-training script.

python pre-train.py --block_size=2048 --embedding_dimensions=4096 --num_hidden_layers=64 --num_attention_heads=64

You can also adjust the batch_size, learning_rate, and gradient_accumulation_steps to suite your training setup.

python pre-train.py --batch_size=32 --learning_rate=0.01 --gradient_accumulation_steps=128

For distributed training, use PyTorch's torchrun extension to launch a distributed data parallel (DDP) session. The example below is for executing the training script on a single node with 8 individual GPUs.

torchrun --standalone --nnodes=1 --nproc-per-node=8 pre-train.py --batch_size=16 --gradient_accumulation_steps=128

Note that when training in data-parallel mode it's important that the gradient_accumulation_steps divides evenly into the world size for maximum performance. For example, if we have an 8 GPU cluster, we could perform 32 gradient accumulation steps in exactly 4 passes over the network.

Pre-training Arguments

Argument Default Type Description
--dataset_subset "sample-10BT" str The subset of the Fineweb dataset to train on. Options are sample-10BT, sample-100BT, and sample-350BT. Set to None to train on the full 15T token dataset.
--token_encoding "r50k_base" str The Tiktoken encoding scheme to use when tokenizing the dataset. Options include r50k_base, p50k_base, cl100k_base, and o200k_base.
--dataset_path "./datasets" str The path to the preprocessed dataset files on disk.
--num_dataset_processes 8 int The number of processes (CPUs) to use to process the dataset.
--batch_size 1 int The number of samples to pass through the network at a time.
--gradient_accumulation_steps 128 int The number of batches to pass through the network before updating the weights.
--samples_per_epoch 4096 int The number of training samples to pass through the network every epoch.
--num_epochs 1686 int The number of epochs to train for.
--learning_rate 1e-2 float The learning rate of the Adafactor optimizer.
--rms_decay -0.8 float The decay rate of the RMS coefficient of the Adafactor optimizer.
--low_memory_optimizer False bool Should the optimizer reduce its memory consumption in exchange for a slightly slower runtime?
--max_gradient_norm 1.0 float Clip gradients above this threshold before stepping.
--eval_interval 10 int Evaluate the model after this many epochs on the testing set.
--block_size 1024 int The number of tokens within the context window for every sample.
--embedding_dimensions 1024 int The dimensionality of the token embeddings.
--num_attention_heads 16 int The number of attention heads within every block.
--num_hidden_layers 24 int The number of attention/MLP blocks within the hidden layer of the network.
--feed_forward_ratio 4 (1, 2, 4) The ratio of hidden neurons to embedding dimensions in the MLP layers of the network.
--dropout 0.1 float The proportion of signals to send to zero during training as regularization.
--activation_checkpointing False bool Should we use activation checkpointing? This will reduce drastically memory utilization during training at the cost of needing to recompute the forward pass.
--ddp_sharding_level 2 int The level of sharding to use for DDP training. Options are 2 or 3 for partial and full sharding respectively, or 0 for no sharding.
--checkpoint_interval 20 int Save the model parameters to disk every this many epochs.
--checkpoint_path "./checkpoints/checkpoint.pt" str The path to the base checkpoint file on disk.
--resume False bool Should we resume training from the last checkpoint?
--device "cuda" str The device to run the computation on.
--seed None int The seed for the random number generator.

Instruction-tuning

Instruction-tuning Arguments

Argument Default Type Description
--base_model_path "./checkpoints/checkpoint.pt" string The path to the base checkpoint on disk.
--batch_size 1 int The number of samples to pass through the network at a time.
--gradient_accumulation_steps 64 int The number of batches to pass through the network before updating the weights.
--learning_rate 5e-4 float The learning rate of the Adafactor optimizer.
--rms_decay -0.8 float The decay rate of the RMS coefficient of the Adafactor optimizer.
--optimizer_low_memory True bool Should the optimizer reduce its memory consumption in exchange for a slightly slower runtime?
--mask_input False bool Should we mask the input part of the sample i.e. only train on the output?
--rank 8 int The rank of the LoRA decomposition matrices.
--alpha 1.0 float The strength of the LoRA signal.
--dropout 0.05 float The proportion of signals to send to zero during training as regularization.
--num_epochs 4 int The number of epochs to train for.
--activation_checkpointing False bool Should we use activation checkpointing? This will reduce drastically memory utilization during training at the cost of needing to recompute the forward pass.
--eval_interval 1 int Evaluate the model after this many epochs on the testing set.
--checkpoint_interval 1 int Save the model parameters to disk every this many epochs.
--checkpoint_path "./checkpoints/lora_instruction.pt" string The path to the LoRA checkpoint.
--resume False bool Should we resume training from the last checkpoint?
--device "cuda" string The device to run the computation on.
--seed None int The seed for the random number generator.

Text Generation

After training, you can generate text from the model by running the generate.py script from the commandline. This inference script samples tokens from the model one at a time conditioned on a prompt and any previously generated tokens, together referred to as the context window. In the example below we are choosing to only sample from the top_k predicted tokens that have at least top_p cumulative probability mass when ordered descending by predicted probability.

python generate.py --top_k=500 --top_p=0.9

Generation Arguments

Argument Default Type Description
--checkpoint_path "./checkpoints/checkpoint.pt" string The path to the base checkpoint file on disk.
--lora_path None string The path to the LoRA checkpoint.
--max_tokens 500 int The maximum number of tokens that the model should generate per sample.
--temperature 1.0 float The amount of regularization applied to the candidate token probabilities.
--top_k 500 int Only sample from this many candidate tokens with the highest probabilities.
--top_p 0.9 float Of the top_k tokens, drop all but the top_p portion of the cumulative probability distribution.
--device "cuda" string The device to run the computation on.
--seed None int The seed for the random number generator.

We also provide a script that samples entire sequences rather than single tokens independently which we call beam_search.py. Beam Search maintains a list of the top beam_width sequence candidates and outputs the top num_candidates completed sequences with the highest overall priority. It is a form of greedy search that works well for some things like text summarization and translation but often results in less natural responses as natural language follows a more stochastic process.

python beam_search.py --beam_width=16 --num_candidates=3

Beam Search Arguments

Argument Default Type Description
--checkpoint_path "./checkpoints/checkpoint.pt" string The path to the base checkpoint file on disk.
--lora_path None string The path to the LoRA checkpoint.
--max_tokens 200 int The maximum number of tokens that the model should generate per sample.
--num_candidates 3 int The number of candidate sequences to output.
--beam_width 16 int The number of candidate sequences to keep track of during search.
--device "cuda" string The device to run the computation on.
--seed None int The seed for the random number generator.

References:

  • G. Penedo, et al. The FineWeb Datasts: Decanting the Web for the Finest Text Data at Scale, 38th Conference on Neural Information Processing Systems (NeurIPS 2024) Track on Datasets and Benchmarks.
  • A. Radford, et al. Language Models are Unsupervised Multitask Learners, OpenAI, 2019.
  • T. Brown, et al. Language Models are Few-Shot Learners. OpenAI, 2020.
  • A. Kazemnejad, et al. The Impact of Positional Encoding on Length Generalization in Transformers, 37th Conference on Neural Information Processing Systems (NeurIPS 2023).
  • S. Rajbhandari, et al. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, 2020.
  • J. R. Hermans, et al. Accumulated Gradient Normalization, JMLR: Workshop and Conference Proceedings, 2017.
  • T. Chen, et al. Training Deep Nets with Sublinear Memory Cost. MIT, 2019.
  • B. Zhang, et al. Root Mean Square Layer Normalization. 33rd Conference on Neural Information Processing Systems, NeurIPS 2019.
  • J. Kaplan, et al. Scaling Laws for Neural Language Models, OpenAI, 2020.
  • J. Hoffman, et al. Training Compute-Optimal Large Language Models, Deep Mind, 2022.