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

Scripts and sample code

• Dockerfile: a container with the basic set of dependencies to run Transformer-XL

In the tf directory, the most important files are:

• data_utils.py: data loading utilities
• exp_utils.py: utility functions for running training and benchmarking
• lamb.py: implementation of LAMB optimizer
• main.py: serves as the entry point to launch the training and inference
• model.py: implementation of the Transformer-XL model
• vocabulary.py: implementation of word-level vocabulary

Parameters

The complete list of available parameters for the tf/main.py script contains:

 --batch_chunk: Number of accumulation steps.
 (default: '1')
 (an integer)
 --clamp_len: Clamp length
 (default: '-1')
 (an integer)
 --clip: Gradient clipping value.
 (default: '0.25')
 (a number)
 --corpus_info_path: Path to corpus-info.json file.
 (default: '')
 --d_embed: Dimension of the embeddings.
 (default: '512')
 (an integer)
 --d_head: Dimension of each attention head.
 (default: '64')
 (an integer)
 --d_inner: Dimension of inner hidden size in positionwise feed-forward.
 (default: '2048')
 (an integer)
 --d_model: Dimension of the model.
 (default: '512')
 (an integer)
 --data_dir: Path to tf-records directory.
 (default: '')
 --div_val: Divide the embedding size by this val for each bin
 (default: '1')
 (an integer)
 --[no]do_eval: Whether to run eval on the dev set.
 (default: 'false')
 --[no]do_train: Whether to run training.
 (default: 'true')
 --dropatt: Attention dropout rate.
 (default: '0.0')
 (a number)
 --dropout: Dropout rate.
 (default: '0.1')
 (a number)
 --eval_batch_size: Size of valid batch.
 (default: '16')
 (an integer)
 --eval_ckpt_path: Checkpoint path for do_test evaluation.If set, model_dir will be ignored.If unset, will use the latest ckpt in model_dir.
 --eval_split: Which data split to evaluate.
 (default: 'valid')
 --amp: Whether to enable AMP ops.
 (default: 'false')
 --init: : Initialization method.
 (default: 'normal')
 --init_range: Initialization std when init is uniform.
 (default: '0.1')
 (a number)
 --init_std: Initialization std when init is normal.
 (default: '0.02')
 (a number)
 --learning_rate: Maximum learning rate.
 (default: '0.01')
 (a number)
 --log_interval: Number of iterations per repeat loop.
 (default: '100')
 (an integer)
 --max_eval_batch: Set -1 to turn off. Only used in test mode.
 (default: '-1')
 (an integer)
 --mem_len: Number of steps to cache
 (default: '192')
 (an integer)
 --min_lr_ratio: Minimum ratio learning rate.
 (default: '0.1')
 (a number)
 --model_dir: Estimator model_dir.
 (default: 'LM-TFM')
 --n_head: Number of attention heads.
 (default: '8')
 (an integer)
 --n_layer: Number of layers.
 (default: '16')
 (an integer)
 --num_core_per_host: Number of cores per host
 (default: '8')
 (an integer)
 --percentiles: percentiles for latency confidence intervals
 (default: '90,95,99')
 (a comma separated list)
 --proj_init_std: Initialization std for embedding projection.
 (default: '0.01')
 (a number)
 --[no]proj_same_dim: Project the bin with the same dimension.
 (default: 'true')
 --[no]proj_share_all_but_first: True to share all but first projs, False not to share.
 (default: 'false')
 --record_info_dir: Path to local directory containing filenames.txt.
 (default: '')
 --[no]same_length: Same length attention
 (default: 'false')
 --save_steps: number of steps for model checkpointing.
 (default: '5000')
 (an integer)
 --tgt_len: Number of steps to predict
 (default: '192')
 (an integer)
 --[no]tie_weight: Tie embedding and softmax weight.
 (default: 'true')
 --train_batch_size: Size of train batch.
 (default: '256')
 (an integer)
 --train_steps: Total number of training steps.
 (default: '40000')
 (an integer)
 --[no]untie_r: untie r_w_bias and r_r_bias
 (default: 'false')
 --warmup_steps: Number of steps for linear lr warmup.
 (default: '1000')
 (an integer)

Command-line options

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

python3 main.py --help

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 tf/data_utils.py file:

• the name of the new dataset should be added to the dataset flag
• 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, dataset iterator and desired values of cutoffs for adaptive softmax

The current codebase supports training with word-level vocabulary (automatically generated based on the provided dataset)

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 script with the first argument set to train. By default, the training results are saved to tf/LM-TFM directory, and map to your container's /workspace/transformer-x/tf/LM-TFM directory; this can be customized by setting the --model_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 16 GPUs.

Command-line

You can launch training of the Transformer-XL base 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> [--amp] [--batch_chunk CHUNK]

The --amp flag is optional, however, if it's set, then the script launches mixed precision training with Tensor Cores; if the flag is not present, then the script launches FP32 training.

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.

The --train batch_size GBS parameter can be used to change batch size. Equation below shows the relation between parameters:

global batch size (--train_batch_size) = <#GPU> * gradient accumulation steps (--batch_chunk) * local batch size

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 --amp --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 DGX-1 with Tesla V100 16G in FP32 training). For example:

bash run_wt103_base.sh train 8 --batch_chunk 2

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

step 1300 | lr 0.009998686 | loss 5.09 | pplx 162.70, bpc 7.3461, tok/s 138037

which contains information about a current training step, current learning rate, current training loss, training perplexity, bits per character and throughput in tokens per second.

The script saves one checkpoint: model.ckpt which contains the last saved model. By default, model saving is executed every 5000 training steps, this can be customized by setting the --save_steps parameter.

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.

Inference process

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

The script supports only single-GPU inference.

Command-line

You can launch inference of the Transformer-XL base model on the WikiText-103 dataset with the word-based vocabulary and adaptive softmax.

For example:

bash run_wt103_base.sh eval --model_dir [--amp]

The --amp 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.

Examples

To launch mixed precision inference on a single GPU using a checkpoint loaded from LM-TFM/model.ckpt*, run:

bash run_wt103_base.sh eval --model_dir LM-TFM --amp

To launch FP32 inference on a single GPU using a checkpoint loaded from LM-TFM/model.ckpt*, run:

bash run_wt103_base.sh eval --model_dir LM-TFM

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

I0109 13:02:31.304439 139903273469760 main.py:440] Evaluating with: math fp16
INFO:tensorflow:| loss 3.15 | pplx 23.32, bpc 4.5432, tok/s 9946, ms/batch 102.84

which contains information about loss, perplexity and execution performance on the test dataset.