Transformer-XL for PyTorch

NVIDIA Deep Learning Examples

Resource

NVIDIA Deep Learning Examples

Transformer-XL for PyTorch

Transformer-XL is a transformer-based language model with a segment-level recurrence and a novel relative positional encoding.

The following sections provide greater details of the dataset, running training and inference, and the training results.

Scripts and sample code

In the root directory, the most important files are:

Dockerfile: container with the basic set of dependencies to run Transformer-XL
requirements.txt: set of extra requirements for running Transformer-XL
getdata.sh: script for downloading datasets

In the pytorch directory, the most important files are:

data_utils.py: data loading utilities
eval.py: serves as the entry point to launch the evaluation and inference
lamb.py: implementation of LAMB optimizer
mem_transformer.py: implementation of the Transformer-XL model
train.py: serves as the entry point to launch the training
run.sub: Slurm batch script for launching multi-node training

The pytorch/utils directory contains the following additional modules:

adaptive_softmax.py: implementation of adaptive softmax
data_parallel.py: implementation of BalancedDataParallel class
distributed.py: utility functions for running distributed training
exp_utils.py: utility functions for running training and benchmarking
log_uniform_sampler.py: implementation of log-uniform sampler
proj_adaptive_softmax.py: implementation of projected adaptive softmax
vocabulary.py: implementation of word-level vocabulary and BPE-based vocabulary

The pytorch/inference directory contains modules optimized for running inference with TorchScript:

mem_transformer_jit.py: implementation of TorchScript-compatible Transformer-XL model
proj_adaptive_softmax_jit.py: implementation of TorchScript-compatible projected adaptive softmax

Parameters

Training

The complete list of available parameters for the pytorch/train.py training script contains:

general setup:
  --work_dir WORK_DIR   Directory for the results
  --append_dataset      Automatically append dataset name to work_dir
  --append_time         Automatically append current time to work_dir
  --cuda                Run training on a GPU using CUDA
  --fp16                Run training in fp16/mixed precision
  --restart RESTART     Restart training from the saved checkpoint
  --debug               Run in debug mode (do not create exp dir)
  --log_all_ranks       Enable logging from all distributed ranks
  --dllog_file DLLOG_FILE
                        Name of the DLLogger output file
  --txtlog_file TXTLOG_FILE
                        Name of the txt log file
  --save_all            Save all checkpoints
  --no_env              Do not print info on execution env
  --no_eval             Disable model evaluation
  --log_interval LOG_INTERVAL
                        Report interval
  --target_throughput TARGET_THROUGHPUT
                        Target training throughput (for benchmarking)
  --target_perplexity TARGET_PERPLEXITY
                        Target validation perplexity (for benchmarking)
  --amp_mode {O0,O1,O2,O3}
                        Optimization level for apex amp

dataset setup:
  --data DATA           Location of the data corpus
  --dataset {wt103,lm1b,enwik8,text8}
                        Dataset name
  --vocab {word,bpe}    Type of vocabulary

model setup:
  --n_layer N_LAYER     Number of total layers
  --n_head N_HEAD       Number of heads
  --d_head D_HEAD       Head dimension
  --d_embed D_EMBED     Embedding dimension
  --d_model D_MODEL     Model dimension
  --d_inner D_INNER     Inner dimension in feedforward layer
  --dropout DROPOUT     Global dropout rate
  --dropatt DROPATT     Attention probability dropout rate
  --pre_lnorm           Apply LayerNorm to the input instead of the output
  --attn_type ATTN_TYPE
                        Attention type. 0 for ours, 1 for Shaw et al,2 for
                        Vaswani et al, 3 for Al Rfou et al.
  --not_tied            Do not tie the word embedding and softmax weights
  --clamp_len CLAMP_LEN
                        Use the same pos embeddings after clamp_len
  --adaptive            Use adaptive softmax
  --div_val DIV_VAL     Dividend value for adaptive input and softmax
  --sample_softmax SAMPLE_SOFTMAX
                        Number of samples in sampled softmax
  --init INIT           Parameter initializer to use
  --emb_init EMB_INIT   Parameter initializer to use
  --init_range INIT_RANGE
                        Parameters initialized by U(-init_range, init_range)
  --emb_init_range EMB_INIT_RANGE
                        Parameters initialized by U(-init_range, init_range)
  --init_std INIT_STD   Parameters initialized by N(0, init_std)
  --proj_init_std PROJ_INIT_STD
                        Parameters initialized by N(0, init_std)

optimizer setup:
  --optim {adam,sgd,adagrad,lamb,jitlamb}
                        Optimizer to use
  --lr LR               Initial learning rate
  --mom MOM             Momentum for sgd
  --scheduler {cosine,inv_sqrt,dev_perf,constant}
                        LR scheduler to use
  --max_step_scheduler MAX_STEP_SCHEDULER
                        Max number of training steps for LR scheduler
  --warmup_step WARMUP_STEP
                        Number of iterations for LR warmup
  --decay_rate DECAY_RATE
                        Decay factor when ReduceLROnPlateau is used
  --lr_min LR_MIN       Minimum learning rate during annealing
  --clip CLIP           Gradient clipping
  --weight_decay WEIGHT_DECAY
                        Weight decay for adam|lamb
  --clip_nonemb         Only clip the gradient of non-embedding params
  --patience PATIENCE   Patience
  --eta_min ETA_MIN     Min learning rate for cosine scheduler

training setup:
  --max_step MAX_STEP   Max number of training steps
  --batch_size BATCH_SIZE
                        Global batch size
  --local_batch_size LOCAL_BATCH_SIZE
                        Local (per-device) batch size, this setting overrides
                        global --batch_size and sets batch_size to
                        local_batch_size * world_size
  --batch_chunk BATCH_CHUNK
                        Split batch into chunks and train with gradient
                        accumulation
  --roll                Enable random shifts within each data stream
  --tgt_len TGT_LEN     Number of tokens to predict
  --ext_len EXT_LEN     Length of the extended context
  --mem_len MEM_LEN     Length of the retained previous heads
  --seed SEED           Random seed
  --multi_gpu {ddp,dp}  Use multiple GPU
  --gpu0_bsz GPU0_BSZ   Batch size on gpu 0 (for "dp" backend)
  --same_length         Use the same attn length for all tokens
  --varlen              Use variable length

validation setup:
  --eval_tgt_len EVAL_TGT_LEN
                        Number of tokens to predict for evaluation
  --eval_batch_size EVAL_BATCH_SIZE
                        Eval batch size
  --eval_max_steps EVAL_MAX_STEPS
                        Max eval steps
  --eval_interval EVAL_INTERVAL
                        Evaluation interval

Inference

The complete list of available parameters for the eval.py inference script contains:

  --work_dir WORK_DIR   experiment directory
  --debug               run in debug mode (do not create exp dir)
  --data DATA           location of the data corpus
  --manual MANUAL [MANUAL ...]
                        run model on raw input data
  --dataset {wt103,lm1b,enwik8,text8}
                        dataset name
  --split {all,valid,test}
                        which split to evaluate
  --type {pytorch,torchscript}
                        type of runtime to use
  --batch_size BATCH_SIZE
                        batch size
  --tgt_len TGT_LEN     number of tokens to predict
  --ext_len EXT_LEN     length of the extended context
  --mem_len MEM_LEN     length of the retained previous heads
  --seed SEED           Random seed
  --clamp_len CLAMP_LEN
                        max positional embedding index
  --cuda                Run evaluation on a GPU using CUDA
  --model MODEL         path to the checkpoint
  --manual_config MANUAL_CONFIG
                        Manually specify config for the model
  --manual_vocab {word,bpe}
                        Manually specify type of vocabulary
  --fp16                Run training in fp16/mixed precision
  --log_all_ranks       Enable logging for all distributed ranks
  --dllog_file DLLOG_FILE
                        Name of the DLLogger output file
  --same_length         set same length attention with masking
  --no_env              Do not print info on execution env
  --log_interval LOG_INTERVAL
                        Report interval
  --target_perplexity TARGET_PERPLEXITY
                        target perplexity
  --target_throughput TARGET_THROUGHPUT
                        target throughput
  --save_data           save latency and throughput data to a file
  --repeat REPEAT       loop over the dataset REPEAT times
  --max_size MAX_SIZE   run inference on up to MAX_SIZE batches
  --percentiles PERCENTILES [PERCENTILES ...]
                        percentiles for latency confidence intervals
  --save_torchscript SAVE_TORCHSCRIPT
                        save torchscript model to a file
  --load_torchscript LOAD_TORCHSCRIPT
                        load torchscript model from a file

Command-line options

To see the full list of available options and their descriptions, use the -h or --help command-line option. For example, for training:

python3 train.py --help

usage: train.py [-h] [--work_dir WORK_DIR] [--append_dataset] [--append_time]
                [--cuda] [--fp16] [--restart RESTART] [--debug]
                [--log_all_ranks] [--dllog_file DLLOG_FILE]
                [--txtlog_file TXTLOG_FILE] [--save_all] [--no_env]
                [--no_eval] [--log_interval LOG_INTERVAL]
                [--target_throughput TARGET_THROUGHPUT]
                [--target_perplexity TARGET_PERPLEXITY]
                [--amp_mode {O0,O1,O2,O3}] [--data DATA]
                [--dataset {wt103,lm1b,enwik8,text8}] [--vocab {word,bpe}]
                [--n_layer N_LAYER] [--n_head N_HEAD] [--d_head D_HEAD]
                [--d_embed D_EMBED] [--d_model D_MODEL] [--d_inner D_INNER]
                [--dropout DROPOUT] [--dropatt DROPATT] [--pre_lnorm]
                [--attn_type ATTN_TYPE] [--not_tied] [--clamp_len CLAMP_LEN]
                [--adaptive] [--div_val DIV_VAL]
                [--sample_softmax SAMPLE_SOFTMAX] [--init INIT]
                [--emb_init EMB_INIT] [--init_range INIT_RANGE]
                [--emb_init_range EMB_INIT_RANGE] [--init_std INIT_STD]
                [--proj_init_std PROJ_INIT_STD]
                [--optim {adam,sgd,adagrad,lamb,jitlamb}] [--lr LR]
                [--mom MOM] [--scheduler {cosine,inv_sqrt,dev_perf,constant}]
                [--max_step_scheduler MAX_STEP_SCHEDULER]
                [--warmup_step WARMUP_STEP] [--decay_rate DECAY_RATE]
                [--lr_min LR_MIN] [--clip CLIP] [--weight_decay WEIGHT_DECAY]
                [--clip_nonemb] [--patience PATIENCE] [--eta_min ETA_MIN]
                [--max_step MAX_STEP] [--batch_size BATCH_SIZE]
                [--local_batch_size LOCAL_BATCH_SIZE]
                [--batch_chunk BATCH_CHUNK] [--roll] [--tgt_len TGT_LEN]
                [--ext_len EXT_LEN] [--mem_len MEM_LEN] [--seed SEED]
                [--multi_gpu {ddp,dp}] [--gpu0_bsz GPU0_BSZ] [--same_length]
                [--varlen] [--eval_tgt_len EVAL_TGT_LEN]
                [--eval_batch_size EVAL_BATCH_SIZE]
                [--eval_max_steps EVAL_MAX_STEPS]
                [--eval_interval EVAL_INTERVAL] [--local_rank LOCAL_RANK]

For example, for inference:

python3 eval.py --help

usage: eval.py [-h] [--work_dir WORK_DIR] [--debug] [--data DATA]
               [--manual MANUAL [MANUAL ...]]
               [--dataset {wt103,lm1b,enwik8,text8}]
               [--split {all,valid,test}] [--type {pytorch,torchscript}]
               [--batch_size BATCH_SIZE] [--tgt_len TGT_LEN]
               [--ext_len EXT_LEN] [--mem_len MEM_LEN] [--seed SEED]
               [--clamp_len CLAMP_LEN] [--cuda] [--model MODEL]
               [--manual_config MANUAL_CONFIG] [--manual_vocab {word,bpe}]
               [--fp16] [--log_all_ranks] [--dllog_file DLLOG_FILE]
               [--same_length] [--no_env] [--log_interval LOG_INTERVAL]
               [--target_perplexity TARGET_PERPLEXITY]
               [--target_throughput TARGET_THROUGHPUT] [--save_data]
               [--repeat REPEAT] [--max_size MAX_SIZE]
               [--percentiles PERCENTILES [PERCENTILES ...]]
               [--save_torchscript SAVE_TORCHSCRIPT]
               [--load_torchscript LOAD_TORCHSCRIPT] [--local_rank LOCAL_RANK]

Getting the data

The Transformer-XL model was trained on the WikiText-103 dataset. The WikiText-103 dataset is a collection of over 100 million tokens extracted from the set of verified Good and Featured articles on Wikipedia.

This repository contains the getdata.sh download script which automatically downloads and extracts the training, validation and test datasets. By default, data is downloaded to the data directory.

In order to test with other datasets, the script needs to be customized accordingly.

Dataset guidelines

The WikiText-103 dataset was already pre-tokenized with word-level tokens. The dataset features a large vocabulary of 267,735 tokens and retains the original case, punctuation and numbers.

The getdata.sh script downloads the data, extracts the archive and renames the training, validation, and test set to train.txt, valid.txt, test.txt respectively.

Multi-dataset

Using other datasets requires changes in the following files:

pytorch/train.py:
- the name of the new dataset should be added to the dataset argument in the parse_args() function
- desired values of cutoffs for adaptive softmax should be added in the main() function, after the section which builds train/valid/test data iterators
pytorch/data_utils.py:
- the support for the new dataset needs to be added to the Corpus class: names of files containing training, validation and test data, options for the tokenizer, and dataset iterator

The current codebase supports training with word-level vocabulary (automatically generated based on the provided dataset) and with BPE vocabulary (using pre-built vocabulary from pretrained GPT2 model imported from github.com/huggingface/transformers.

Additionally, using other datasets may require changes in some hyperparameters (for example, batch size, learning rate, number of training steps, and the configuration of learning rate scheduler).

Training process

The default training configuration can be launched by running the run_wt103_base.sh or the run_wt103_large.sh script with the first argument set to train. By default, the training results are saved to the LM-TFM directory; this can be customized by setting the --work_dir parameter.

The training script launches a single-node data-parallel training with a fixed global batch size of 256, optionally with gradient accumulation to allow training on configurations with less than 8 GPUs. Logs from the training are automatically saved to the LM-TFM/train_log.log file.

Command-line

You can launch training of the Transformer-XL base/large model on the WikiText-103 dataset with the word-based vocabulary and adaptive softmax using <#GPUs> GPUs. For example:

bash run_wt103_base.sh train <#GPUs> [--fp16] [--batch_chunk CHUNK]

and

bash run_wt103_large.sh train <#GPUs> [--fp16] [--batch_chunk CHUNK]

The --fp16 flag is optional, however, if it's specified, then the script launches mixed precision training with Tensor Cores; if the flag is not present, then the script launches FP32 training on NVIDIA Volta GPUs and TF32 training on NVIDIA Ampere GPUs.

The --batch_chunk CHUNK parameter controls gradient accumulation. With gradient accumulation, the batch size is split into CHUNK chunks of equal size, the training script executes the forward and backward pass using each chunk and then executes the optimizer using accumulated gradients.

Examples

You can launch mixed precision training of the Transformer-XL base model on the WikiText-103 dataset using 16 GPUs. For example:

bash run_wt103_base.sh train 16 --fp16 --batch_chunk 1

The batch size per GPU is equal to the default global batch size of 256 divided by the product of the number of GPUs times the number of chunks, in this case batch size per GPU is equal to 256 / (16 * 1) = 16.

You can launch FP32 training using 8 GPUs; the batch size per GPU is equal to 16 (--batch_chunk was set to 2 because a local batch size of 32 runs out of memory on a NVIDIA DGX-1 with Tesla V100 16GB in FP32 training). For example:

bash run_wt103_base.sh train 8 --batch_chunk 2

A progress summary of the training progress is printed after every 10 training iterations; this can be customized by setting the --log_interval parameter. The summary is printed in the following format:


| epoch  18 step    36000 | batches    283 / 2101 | lr 1.220e-03 | ms/batch 185.1 | tok/s  265585 | loss  3.12 | ppl     22.71

which contains information about a current training epoch, current training step, number of batches processed within the current epoch, current learning rate, execution time in milliseconds per batch, throughput in tokens per second, current training loss and training perplexity.

The script saves two checkpoints: checkpoint_best.pt which contains the model corresponding to the lowest value of the validation loss and checkpoint_last.pt which contains the model corresponding to the last execution of the validation step. By default, the validation is executed every 5000 training steps, this can be customized by setting the --eval_interval parameter. The summary of results on the validation dataset is printed in the following format:


| Eval   7 at step    35000 | time:  1.37s | valid loss  3.14 | valid ppl    23.132

which contains information about the current epoch, current training step, time needed to execute the validation, current validation loss, and validation perplexity.

At the end of the training, the training script automatically runs evaluation on the test dataset. This automatic evaluation is executed with values of mem_len and tgt_len hyperparameters inherited from the training setup. Evaluation (inference) benefits from longer attention sequences, therefore to reproduce perplexity values reported in the Transformer-XL paper, it's necessary to run the final evaluation with a dedicated inference script. Refer to the Inference process section for more details.

Multi-node

Multi-node runs can be launched on a pyxis/enroot Slurm cluster (see Requirements). To launch a multi-node run, issue the run.sub script with the following command for an 8-node DGX-2H training, for example:

sbatch run.sub all

This repository contains a number of predefined configurations to run the multi-node training on DGX-2H nodes. By default, run.sub launches 8-node training.

To launch multi-node training on <NODES> DGX-2H nodes, run:

CONFIG=<NODES>dgx2_16gpu_{fp16,fp32} sbatch -N <NODES> run.sub all

supported values for <NODES> parameter are: 1, 2, 4, 8
configs with fp16 suffix launch mixed precision training, configs with fp32 suffix launch FP32 training

Examples:

To launch 4-node mixed-precision training, run:

CONFIG=4dgx2_16gpu_fp16 sbatch -N 4 run.sub all

To launch 2-node FP32 training, run:

CONFIG=2dgx2_16gpu_fp32 sbatch -N 2 run.sub all

Note that the run.sub script is a starting point that has to be adapted depending on the environment. In particular, variables such as WORK_DIR handle the location of the workspace in the file system. The variable CONT should point to the location of the Transformer-XL Docker container. It's assumed that the Docker container built with the scripts/docker/build.sh script was pushed to a Docker registry accessible from all compute nodes.

Refer to the contents of the file to see the full list of variables to adjust for your system.

Inference process

Inference can be run by launching the run_wt103_base.sh or the run_wt103_large.sh script with the first argument set to eval. Running inference requires a pre-trained model checkpoint.

The script supports single-node multi-GPU inference, each batch is split equally among all GPUs running the inference and the loss is averaged over the global batch. Logs from the inference are automatically saved to the LM-TFM/eval_log.log file.

Command-line

You can launch inference of the Transformer-XL base/large model on the WikiText-103 dataset with the word-based vocabulary and adaptive softmax using <#GPUs> GPUs. For example:

bash run_wt103_base.sh eval <#GPUs> --model <PATH TO THE CHECKPOINT> [--fp16] [--type {pytorch, torchscript}]

and

bash run_wt103_large.sh eval <#GPUs> --model <PATH TO THE CHECKPOINT> [--fp16] [--type {pytorch, torchscript}]

The --fp16 flag is optional, however, if it's specified, then the script launches inference with Tensor Cores; if the flag is not present, then the script launches FP32 inference on NVIDIA Volta and NVIDIA Turing GPUs and TF32 inference on NVIDIA Ampere GPUs.

The --type flag selects between pure Python PyTorch execution and TorchScript execution.

Supported values for <#GPUs> are: 1, 2, 4, 8 for NVIDIA DGX-1 and NVIDIA DGX A100 and 1, 2, 4, 8, 16 for NVIDIA DGX-2H.

Examples

To launch TorchScript mixed precision inference on 8 GPUs using a checkpoint loaded from LM-TFM/checkpoint_best.pt, run:

bash run_wt103_base.sh eval 8 --model LM-TFM/checkpoint_best.pt --fp16 --type torchscript

To launch pure Python TF32/FP32 inference on a single GPU using a checkpoint loaded from LM-TFM/checkpoint_best.pt, run:

bash run_wt103_base.sh eval 1 --model LM-TFM/checkpoint_best.pt --type pytorch

After the execution, the script prints a summary in the following format:

Evaluating with math fp16 type torchscript bsz 16 tgt_len 64 ext_len 0 mem_len 640 clamp_len 400
Time : 5.29s, 22.05ms/segment
====================================================================================================

| test loss  3.15 | test ppl    23.304
====================================================================================================

which contains information about runtime parameters, execution time, loss and perplexity on the test dataset.