To train your model using mixed precision with Tensor Cores or using FP32, perform the following steps using the default parameters of the Transformer-XL base model on the WikiText-103 dataset.

For the specifics concerning training and inference, see the Advanced section.

1. Clone the repository.

git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/TensorFlow/LanguageModeling/Transformer-XL

2. Download and preprocess the dataset.

bash getdata.sh

3. Build the Transformer-XL TensorFlow NGC container.

bash tf/scripts/docker/build.sh

4. Start an interactive session in the NGC container to run training/inference.

bash tf/scripts/docker/interactive.sh

5. Create tfrecords before your first training/evaluation for a given batch size per GPU. Use same --batch_chunk and --training_batch_size flags as in the training.

For training on DGX-A100 without gradient accumulation:

bash run_wt103_base.sh train_data

For training on DGX-1 with gradient accumulation in 2 steps:

bash run_wt103_base.sh train_data --batch_chunk 2

For single GPU training with gradient accumulation in 16 steps:

bash run_wt103_base.sh train_data --batch_chunk 16

For evaluation:

bash run_wt103_base.sh test_data

6. Start training.

To start TF32 training on 8 GPUs on DGX-A100, run:

bash run_wt103_base.sh train 8

To start mixed precision training on 8 GPUs on DGX-1, run:

bash run_wt103_base.sh train 8 --amp --batch_chunk 2

To start FP32 training on single GPU, run:

bash run_wt103_base.sh train 1 --batch_chunk 16

To start mixed precision training on 16 GPUs on DGX-2, run:

bash run_wt103_base.sh train 16 --amp

To start FP32 training on 16 GPUs on DGX-2, run:

bash run_wt103_base.sh train 16

For more information on the available options, and for an explanation of what happens at the end of training, refer to the Training process section.

7. Start evaluation.

To start mixed precision inference on the test set, run:

bash run_wt103_base.sh eval [--amp]

The --amp flag is optional, however, if it's set, then the script launches mixed precision inference with Tensor Cores. If the flag is not present, then the script launches FP32 inference. By default, the script is loading the checkpoint from LM-TFM/model.ckpt, which contains the model corresponding to the last checkpoint from the previous training run. The path to the checkpoint can be customized by setting the --model_dir flag.

For more information on the available options, refer to the Inference process section.