Transformer Lab

Python

PyTorch

Lightning

🤗 Datasets

SentencePiece

Repository

Downloads

PyPI

Transformer Lab is a collection of Transformer model and dataset implementations, an integrated tokenizer, Lightning training code, and an inference library. It facilitates the development of language models and rapid experimentation with them.

Features:

• Simple, efficient implementations and modular design
• Seamless parallel and reproducible training on multiple GPUs
• Easily configurable architecture modifications and training procedure
• Multiple generation strategies during inference

Model

A Transformer decoder architecture [Vaswani et al. 2017] with the following performance-improving modifications:

• Normalization is performed before the Transformer sub-layers. An extra normalization layer is added after the last feed-forward sub-layer. These changes also improve training stability.
• The GELU activation function is used.
• Dropout is applied only in the attention and feed-forward sub-layers (to the product of the queries and keys, and to the hidden activations, respectively).

Sinusoidal positional encodings are used because learned positional embeddings degrade performance in the current setup.

Tokenizer

The following criteria were used to select a tokenization algorithm:

• Able to encode any Unicode string without resorting to “unknown” tokens
• Perfectly reversible (no information loss during preprocessing)

The selected algorithm is BPE (byte-pair encoding) from SentencePiece. This implementation operates on characters but is able to encode out-of-vocabulary symbols as bytes. SentencePiece is mostly reversible but loses information by replacing spaces with the meta symbol “▁” (U+2581). To ensure perfect reversibility, during preprocessing Transformer Lab replaces the meta symbol with a rare sequence and unambiguously escapes occurrences of the replacement. The vocabulary size is 32K tokens, chosen as a reasonable trade-off between context compression and model head size.

Dataset

wikimedia/wikipedia, 20231101.en, split into:

• training: 90%
• validation: 5%
• test: 2.5%
• predict: 2.5%

Articles are chunked into sequences without overlap. When (the remainder of) a sequence is shorter than the context length, it is padded. Experiments with concatenating such sequences led to reduced training speed and only marginal improvements in validation performance.

Training

The model is trained on sequences with a maximum length of 512 tokens. To speed up training and reduce memory usage, FlashAttention-2 and 16-bit mixed precision are used. To mitigate stability issues, the bfloat16 data type is used, along with gradient clipping at a norm of 1.0. The AdamW optimizer is used with a learning rate of 3e-4 and a weight decay of 0.1. The reported results come from training on a single A100 GPU for 3 epochs (144K steps) with a batch size of 256.

Results

Version	Checkpoint	Loss (test)	Accuracy (test)
0.1	last	3.18	40.0%
0.2	last	2.79	47.0%

Getting started

Installation

Transformer Lab can be installed as a Python package or by cloning its repository. The former lets you use the models for inference, but to train them you need your own code. The latter comes with pre-built training and inference scripts, as well as configuration.

To install Transformer Lab as a Python package, run:

pip install xlabml

Note that this will also install Lightning and its dependencies. Given enough interest, future package versions may come without Lightning to minimize the number of dependencies.

To install Transformer Lab from its repository, clone it and check out a release tag (because the main branch may be unstable). Next, create and activate a virtual environment, then run:

pip install -r requirements.txt

Obtaining weights

If you need pre-trained weights for inference or fine-tuning, download the .pt file(s) from Transformer Lab’s releases.

Using in your code

import torch
from xlabml.tokenizer import Tokenizer
from xlabml.datamodules import XLabDataModule
from xlabml.models import XLabModel
from xlabml import inference

# adjust these
checkpoint_path = 'path/to/weights.pt'
prompt = 'april'
device = 'cuda'
limit = 10

tokenizer_path = XLabDataModule.load_from_checkpoint(checkpoint_path, map_location=device).tokenizer_path
tokenizer = Tokenizer.load(tokenizer_path)
model = XLabModel.load_from_checkpoint(checkpoint_path, map_location=device).eval().requires_grad_(False)
inputs = torch.tensor([tokenizer[tokenizer.sos_token]] + tokenizer.encode(prompt), device=model.device)
outputs = inference.sample(
    model, inputs, limit,
    block_size=model.hparams['max_len'],
    eos_class=tokenizer[tokenizer.eos_token]
)
output = tokenizer.decode(outputs.tolist())

Note that the resulting model object will be a LightningModule. To decouple the model from Lightning, use the underlying PyTorch module, accessible via the model attribute.

Training the model

From the command line, run:

./xlab.py fit -c conf/xlab.yaml

Specifying -c conf/xlab.yaml configures the training script to use a larger dataset and model (the default ones are intended for quick experiments). This first downloads the dataset and tokenizer and then encodes the dataset, which takes about an hour in total. If you can’t wait, you can try Transformer Lab with a smaller dataset and model – simply omit -c conf/xlab.yaml from every command.

The actual training takes about 6 hours per epoch on an A100 GPU. For hardware with less memory, you may need to modify the configuration and decrease the context size, the batch size, or both.

By default, the following are saved during training:

• the tokenizer, in the tokenizers directory
• the configuration and hyperparameters, in YAML format under logs/version_*
• learning curves, in TensorBoard format under logs/version_*
• model checkpoints, in logs/version_*/checkpoints

You can view the learning curves by running:

tensorboard --logdir .

Keep this running, point your browser to http://localhost:6006/, and click the Scalars tab.

Validation

./xlab.py validate -c conf/xlab.yaml --ckpt_path PATH

Here, PATH points to a checkpoint downloaded from Transformer Lab’s releases or saved during training. To evaluate on the test set, replace validate with test in the above command.

Inference

For command line inference, run:

./infer.py CHECKPOINT_PATH "PROMPT"

The default generation strategy is multinomial sampling, but greedy search or beam search can also be selected. Options exist for configuring the algorithms (e.g., temperature, top-k, top-p, beam width). To see all inference options, run ./infer.py --help.

Exporting model weights and hyperparameters

Checkpoints created during training contain not only the model weights but also the optimizer state and other information needed to resume training. This makes the checkpoints three times larger than the actual model weights.

To export a “clean” checkpoint, containing only the weights and hyperparameters, run:

./manage.py export-checkpoint CHECKPOINT_PATH

Other actions

Validating training data

./xlab.py validate_data -c conf/xlab.yaml

Training the tokenizer

Delete the old tokenizer and then run:

./xlab.py train_tokenizer -c conf/xlab.yaml

Computing dataset statistics

./xlab.py compute_stats -c conf/xlab.yaml

Configuration

Configuration and hyperparameters are read from multiple sources in the following order, with later sources overriding earlier ones:

• defaults in the code
• conf/defaults.yaml, implicitly
• YAML files, given with the -c PATH option
• command line options

See conf/xlab.yaml for the configuration used to train the current release model. Examples of additional options are in conf/extra. To see the full list, run ./xlab.py --help.

Hyperparameters are saved in checkpoints and automatically restored when loading the model.

Generations

Prompt: April is

April is a municipality in the district of Neuchâtel, within the canton of Lucerne, the canton of Thurgau, Switzerland.

History April racing dinners

April generally placed in cultivation with the ‘Michael Mayer Cup Dunroipel Limestone Witige Witige Klice Vereu Pino Annarum’ on January 2, 1999.

Geography Amongst the top strata of April 2, 1949

Prompt: The world

The world record in Statistics and Statistics is a record-setting record for any years recorded by the Republic of South Africa. Released September 2014, racing dataset has had its generally placed in the United States, most of them generally after reports showed that Witwatersrand took the position, and thus concluded a decline in record earnings. It was first held on September 21, 2014, as a measure of the number of admissions of the nation’s

Prompt: Cats and dogs

Cats and dogs were described as the “major antagonistic activists”. The pro-heroutee activist she was an early Patty Hippolyte Laird and later Frank Hewlett. She had married the archaeologist Robert Johnson Parr, at the time Dunrood Monument. Their daughter crowned women winners, Peck Pino and Pedra Acres. It was first reported that of a white man. She first planned to marry a friend, Frank Lloyd Taylor in 1866. She later

All of these were produced from version 0.2 and can be reproduced with the included inference script, using random seed 42 and a limit of 100.

Possible future work

• Add cosine or inverse-square-root learning rate scheduling after the warmup
• Train on progressively longer sequences
• Use rotary positional embeddings or ALiBi
• Use RMSNorm
• Use the SwiGLU activation
• Train a larger model on a larger, more diverse dataset
• Fine-tune for a downstream task
• Add quantization
• Compile the model for training and inference
• Cache key-value pairs during inference
• Try multi-query or grouped-query attention

---

For source code, please visit Transformer Lab’s repository.