Quick Start Guide

Perform the following steps to run the training using the default parameters of the GNMT v2 model on the WMT16 English-German dataset.

1. Build and launch the GNMT Docker container

bash scripts/docker/build.sh
bash scripts/docker/interactive.sh

2. Download the training dataset

Download and preprocess the WMT16 English-German dataset. Data will be downloaded to the data directory (on the host). The data directory is mounted to the /workspace/gnmt/data location in the Docker container.

bash scripts/wmt16_en_de.sh

3. Run training

By default, the training script will use all available GPUs. The training script saves only one checkpoint with the lowest value of the loss function on the validation dataset. All results and logs are saved to the results directory (on the host) or to the /workspace/gnmt/results directory (in the container). By default, the train.py script will launch mixed precision training with Tensor Cores. You can change this behaviour by setting the --math fp32 flag for the train.py training script.

Launching training on 1, 4 or 8 GPUs:

python3 -m launch train.py --seed 2 --train-global-batch-size 1024

Launching training on 16 GPUs:

python3 -m launch train.py --seed 2 --train-global-batch-size 2048

By default the training script will launch training with batch size 128 per GPU. If specified --train-global-batch-size is larger than 128 times the number of GPUs available for the training then the training script will accumulate gradients over consecutive iterations and then perform the weight update. For example 1 GPU training with --train-global-batch-size 1024 will accumulate gradients over 8 iterations before doing the weight update with accumulated gradients.

The training script automatically runs the validation and testing after each training epoch. The results from the validation and testing are printed to the standard output (stdout) and saved to log files.

The summary after each training epoch is printed in the following format:

Summary: Epoch: 3	Training Loss: 3.1735	Validation Loss: 3.0511	Test BLEU: 21.89
Performance: Epoch: 3  Training: 300155 Tok/s  Validation: 156066 Tok/s

The training loss is averaged over an entire training epoch, the validation loss is averaged over the validation dataset and the BLEU score is computed by the SacreBLEU package on the test dataset. Performance is reported in total tokens per second. The result is averaged over an entire training epoch and summed over all GPUs participating in the training.

Details

Getting the data

The GNMT v2 model was trained on the WMT16 English-German dataset. Concatenation of the newstest2015 and newstest2016 test sets are used as a validation dataset and the newstest2014 is used as a testing dataset.

This repository contains the scripts/wmt16_en_de.sh download script which will automatically download and preprocess the training, validation and test datasets. By default, data will be downloaded to the data directory.

Our download script is very similar to the wmt16_en_de.sh script from the tensorflow/nmt repository. Our download script contains an extra preprocessing step, which discards all pairs of sentences which can't be decoded by latin-1 encoder. The scripts/wmt16_en_de.sh script uses the subword-nmt package to segment text into subword units (BPE). By default, the script builds the shared vocabulary of 32,000 tokens.

In order to test with other datasets, scripts need to be customized accordingly.

Training process

The default training configuration can be launched by running the train.py training script. By default, the training script saves only one checkpoint with the lowest value of the loss function on the validation dataset, an evaluation is performed after each training epoch. Results are stored in the results/gnmt_wmt16 directory.

The training script launches data-parallel training with batch size 128 per GPU on all available GPUs. We have tested reliance on up to 16 GPUs on a single node. After each training epoch, the script runs an evaluation on the validation dataset and outputs a BLEU score on the test dataset (newstest2014). BLEU is computed by the SacreBLEU package. Logs from the training and evaluation are saved to the results directory.

Even though the training script uses all available GPUs, you can change this behavior by setting the CUDA_VISIBLE_DEVICES variable in your environment or by setting the NV_GPU variable at the Docker container launch (see section "GPU isolation").

By default, the train.py script will launch mixed precision training with Tensor Cores. You can change this behaviour by setting the --math fp32 flag for the train.py script.

To view all available options for training, run python3 train.py --help.

Inference process

Inference can be run by launching the translate.py inference script, although, it requires a pre-trained model checkpoint and tokenized input.

The inference script, translate.py, supports batched inference. By default, it launches beam search with beam size of 5, coverage penalty term and length normalization term. Greedy decoding can be enabled by setting the beam size to 1.

To view all available options for inference, run python3 translate.py --help.