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:

• run_pretraining.py - Serves as entry point for pre-training
• run_squad.py - Serves as entry point for SQuAD training
• run_classifier.py - Serves as entry point for GLUE training
• Dockerfile - Container with the basic set of dependencies to run BERT

The scripts/ folder encapsulates all the one-click scripts required for running various functionalities supported such as:

• run_squad.sh - Runs SQuAD training and inference using run_squad.py file
• run_glue.sh - Runs GLUE training and inference using the run_classifier.py file
• run_pretraining_adam.sh - Runs pre-training with Adam optimizer using the run_pretraining.py file
• run_pretraining_lamb.sh - Runs pre-training with LAMB optimizer using the run_pretraining.py file in two phases. Phase 1 does 90% of training with sequence length = 128. In phase 2, the remaining 10% of the training is done with sequence length = 512.
• data_download.sh - Downloads datasets using files in the data/ folder
• finetune_train_benchmark.sh - Captures performance metrics of training for multiple configurations
• finetune_inference_benchmark.sh - Captures performance metrics of inference for multiple configurations

Other folders included in the root directory are:

• data/ - Necessary folders and scripts to download datasets required for fine tuning and pre-training BERT.
• utils/ - Necessary files for preprocessing data before feeding into BERT and hooks for obtaining performance metrics from BERT.

Parameters

Aside from the options to set hyperparameters, the relevant options to control the behaviour of the run_pretraining.py script are:

  --bert_config_file: The config json file corresponding to the pre-trained BERT model. This specifies the model architecture.
  --init_checkpoint: Initial checkpoint (usually from a pre-trained BERT model).
  --[no]do_eval: Whether to run evaluation on the dev set.(default: 'false')
  --[no]do_train: Whether to run training.(evaluation: 'false')
  --eval_batch_size: Total batch size for eval.(default: '8')(an integer)
  --[no]horovod: Whether to use Horovod for multi-gpu runs(default: 'false')
  --[no]amp: Whether to enable AMP ops. When false, uses TF32 on A100 and FP32 on V100 GPUS.(default: 'True')
  --[no]use_xla: Whether to enable XLA JIT compilation.(default: 'True')
  --input_files_dir: Input TF example files (can be a dir or comma separated).
  --output_dir: The output directory where the model checkpoints will be    written.
  --optimizer_type: Optimizer used for training - LAMB or ADAM
  --num_accumulation_steps: Number of accumulation steps before gradient update. Global batch size = num_accumulation_steps * train_batch_size
  --allreduce_post_accumulation: Whether to all reduce after accumulation of N steps or after each step

Aside from the options to set hyperparameters, some relevant options to control the behaviour of the run_squad.py script are:

  --bert_config_file: The config json file corresponding to the pre-trained BERT model. This specifies the model architecture.
  --output_dir: The output directory where the model checkpoints will be written.
  --[no]do_predict: Whether to run evaluation on the dev set. (default: 'false')
  --[no]do_train: Whether to run training. (default: 'false')
  --learning_rate: The initial learning rate for Adam.(default: '5e-06')(a number)
  --max_answer_length: The maximum length of an answer that can be generated. This is needed because the start and end predictions are not conditioned on one another.(default: '30')(an integer)
  --max_query_length: The maximum number of tokens for the question. Questions longer than this will be truncated to this length.(default: '64')(an integer)
  --max_seq_length: The maximum total input sequence length after WordPiece tokenization. Sequences longer than this will be truncated, and sequences shorter than this will be padded.(default: '384')(an integer)
  --predict_batch_size: Total batch size for predictions.(default: '8')(an integer)
  --train_batch_size: Total batch size for training.(default: '8')(an integer)
  --[no]amp: Whether to enable AMP ops. When false, uses TF32 on A100 and FP32 on V100 GPUS.(default: 'True')
  --[no]use_xla: Whether to enable XLA JIT compilation.(default: 'True')
  --[no]version_2_with_negative: If true, the SQuAD examples contain some that do not have an answer.(default: 'false')

Aside from the options to set hyperparameters, some relevant options to control the behaviour of the run_classifier.py script are:

  --bert_config_file: The config json file corresponding to the pre-trained BERT model. This specifies the model architecture.
  --data_dir: The input data dir. Should contain the .tsv files (or other data files) for the task.
  --[no]do_eval: Whether to run eval on the dev set.
    (default: 'false')
  --[no]do_predict: Whether to run the model in inference mode on the test set.(default: 'false')
  --[no]do_train: Whether to run training.(default: 'false')
  --[no]horovod: Whether to use Horovod for multi-gpu runs(default: 'false')
  --init_checkpoint: Initial checkpoint (usually from a pre-trained BERT model).
  --max_seq_length: The maximum total input sequence length after WordPiece tokenization. Sequences longer than this will be truncated, and sequences shorter than this will be padded.(default: '128')(an integer)
  --num_train_epochs: Total number of training epochs to perform.(default: '3.0')(a number)
  --output_dir: The output directory where the model checkpoints will be written.
  --task_name: The name of the task to train.
  --train_batch_size: Total batch size for training.(default: '32')(an integer)
  --[no]amp: Whether to enable AMP ops. When false, uses TF32 on A100 and FP32 on V100 GPUS.(default: 'True')
  --[no]use_xla: Whether to enable XLA JIT compilation.(default: 'True')
  --vocab_file: The vocabulary file that the BERT model was trained on.
  --warmup_proportion: Proportion of training to perform linear learning rate warmup for. E.g., 0.1 = 10% of training.(default: '0.1')(a number)

Note: When initializing from a checkpoint using --init_checkpoint and a corpus of your choice, keep in mind that bert_config_file and vocab_file should remain unchanged.

Command-line options

To see the full list of available options and their descriptions, use the -h or --help command-line option with the Python file, for example:

python run_pretraining.py --help
python run_squad.py --help
python run_classifier.py --help

Getting the data

For pre-training BERT, we use the concatenation of Wikipedia (2500M words) as well as BookCorpus (800M words). For Wikipedia, we extract only the text passages from here and ignore headers list and tables. It is structured as a document level corpus rather than a shuffled sentence level corpus because it is critical to extract long contiguous sentences.

The next step is to run create_pretraining_data.py with the document level corpus as input, which generates input data and labels for the masked language modeling and next sentence prediction tasks. Pre-training can also be performed on any corpus of your choice. The collection of data generation scripts are intended to be modular to allow modifications for additional preprocessing steps or to use additional data. They can hence easily be modified for an arbitrary corpus.

The preparation of an individual pre-training dataset is described in the create_datasets_from_start.sh script found in the data/ folder. The component steps to prepare the datasets are as follows:

1. Data download and extract - the dataset is downloaded and extracted.
2. Clean and format - document tags, etc. are removed from the dataset. The end result of this step is a {dataset_name_one_article_per_line}.txt file that contains the entire corpus. Each line in the text file contains an entire document from the corpus. One file per dataset is created in the formatted_one_article_per_line folder.
3. Sharding - the sentence segmented corpus file is split into a number of smaller text documents. The sharding is configured so that a document will not be split between two shards. Sentence segmentation is performed at this time using NLTK.
4. TFRecord file creation - each text file shard is processed by the create_pretraining_data.py script to produce a corresponding TFRecord file. The script generates input data and labels for masked language modeling and sentence prediction tasks for the input text shard.

For fine tuning BERT for the task of Question Answering, we use SQuAD and GLUE. SQuAD v1.1 has 100,000+ question-answer pairs on 500+ articles. SQuAD v2.0 combines v1.1 with an additional 50,000 new unanswerable questions and must not only answer questions but also determine when that is not possible. GLUE consists of single-sentence tasks, similarity and paraphrase tasks and inference tasks. We support one of each: CoLA, MNLI and MRPC.

Dataset guidelines

The procedure to prepare a text corpus for pre-training is described in the previous section. This section provides additional insight into how exactly raw text is processed so that it is ready for pre-training.

First, raw text is tokenized using WordPiece tokenization. A [CLS] token is inserted at the start of every sequence, and the two sentences in the sequence are separated by a [SEP] token.

Note: BERT pre-training looks at pairs of sentences at a time. A sentence embedding token [A] is added to the first sentence and token [B] to the next.

BERT pre-training optimizes for two unsupervised classification tasks. The first is Masked Language Modelling (Masked LM). One training instance of Masked LM is a single modified sentence. Each token in the sentence has a 15% chance of being replaced by a [MASK] token. The chosen token is replaced with [MASK] 80% of the time, 10% with another random token and the remaining 10% with the same token. The task is then to predict the original token.

The second task is next sentence prediction. One training instance of BERT pre-training is two sentences (a sentence pair). A sentence pair may be constructed by simply taking two adjacent sentences from a single document, or by pairing up two random sentences with equal probability. The goal of this task is to predict whether or not the second sentence followed the first in the original document.

The create_pretraining_data.py script takes in raw text and creates training instances for both pre-training tasks.

Multi-dataset

We are able to combine multiple datasets into a single dataset for pre-training on a diverse text corpus. Once TFRecords have been created for each component dataset, you can create a combined dataset by adding the directory to SOURCES in run_pretraining_*.sh. This will feed all matching files to the input pipeline in run_pretraining.py. However, in the training process, only one TFRecord file is consumed at a time, therefore, the training instances of any given training batch will all belong to the same source dataset.

Training process

The training process consists of two steps: pre-training and fine tuning.

Pre-training

Pre-training is performed using the run_pretraining.py script along with parameters defined in the scripts/run_pretraining_lamb.sh.

The run_pretraining_lamb.sh script runs a job on a single node that trains the BERT-large model from scratch using the Wikipedia and BookCorpus datasets as training data. By default, the training script:

• Runs on 8 GPUs.

• Has FP16 precision enabled.

• Is XLA enabled.

• Creates a log file containing all the output.

• Saves a checkpoint every 100 iterations (keeps only the latest checkpoint) and at the end of training. All checkpoints, evaluation results and training logs are saved to the /results directory (in the container which can be mounted to a local directory).

• Evaluates the model at the end of each phase.

• Phase 1

  • Runs 7038 steps with 2000 warmup steps
  • Sets Maximum sequence length as 128
  • Sets Global Batch size as 64K

• Phase 2

  • Runs 1564 steps with 200 warm-up steps
  • Sets Maximum sequence length as 512
  • Sets Global Batch size as 32K
  • Starts from Phase1's final checkpoint

These parameters train Wikipedia and BookCorpus with reasonable accuracy on a DGX-1 with 32GB V100 cards.

For example:

scripts/run_pretraining_lamb.sh <train_batch_size_phase1> <train_batch_size_phase2> <eval_batch_size> <learning_rate_phase1> <learning_rate_phase2> <precision> <use_xla> <num_gpus> <warmup_steps_phase1> <warmup_steps_phase2> <train_steps> <save_checkpoint_steps> <num_accumulation_phase1> <num_accumulation_steps_phase2> <bert_model>

Where:

• <training_batch_size_phase*> is per-GPU batch size used for training in the respective phase. Batch size varies with precision, larger batch sizes run more efficiently, but require more memory.

• <eval_batch_size> is per-GPU batch size used for evaluation after training.

• <learning_rate_phase1> is the default rate of 1e-4 is good for global batch size 256.

• <learning_rate_phase2> is the default rate of 1e-4 is good for global batch size 256.

• <precision> is the type of math in your model, can be either fp32 or fp16. Specifically:

  • fp32 is 32-bit IEEE single precision floats. Is enabled by default on V100.
  • fp16 is Automatic rewrite of TensorFlow compute graph to take advantage of 16-bit arithmetic whenever it is safe.
  • tf32 uses same 10 bit mantissa as fp16 and 8 bit exponent as fp32. Is enabled by default on A100.

• <num_gpus> is the number of GPUs to use for training. Must be equal to or smaller than the number of GPUs attached to your node.

• <warmup_steps_phase*> is the number of warm-up steps at the start of training in the respective phase.

• <training_steps> is the total number of training steps in both phases combined.

• <save_checkpoint_steps> controls how often checkpoints are saved. Default is 100 steps.

• <num_accumulation_phase*> is used to mimic higher batch sizes in the respective phase by accumulating gradients N times before weight update.

• <bert_model> is used to indicate whether to pretrain BERT Large or BERT Base model

The following sample code trains BERT-large from scratch on a single DGX-2 using FP16 arithmetic. This will take around 4.5 days.

bert_tf/scripts/run_pretraining_lamb.sh 32 8 8 3.75e-4 2.5e-4 fp16 true 16 2000 200 7820 100 128 256 large

Fine tuning

Fine tuning is performed using the run_squad.py script along with parameters defined in scripts/run_squad.sh.

The run_squad.sh script trains a model and performs evaluation on the SQuAD dataset. By default, the training script:

• Trains for SQuAD v1.1 dataset.
• Trains on BERT Large Model.
• Uses 8 GPUs and batch size of 10 on each GPU.
• Has FP16 precision enabled.
• Is XLA enabled.
• Runs for 2 epochs.
• Saves a checkpoint every 1000 iterations (keeps only the latest checkpoint) and at the end of training. All checkpoints, evaluation results and training logs are saved to the /results directory (in the container which can be mounted to a local directory).
• Evaluation is done at the end of training. To skip evaluation, modify --do_predict to False.

This script outputs checkpoints to the /results directory, by default, inside the container. Mount point of /results can be changed in the scripts/docker/launch.sh file. The training log contains information about:

• Loss for the final step
• Training and evaluation performance
• F1 and exact match score on the Dev Set of SQuAD after evaluation.

The summary after training is printed in the following format:

I0312 23:10:45.137036 140287431493376 run_squad.py:1332] 0 Total Training Time = 3007.00 Training Time W/O start up overhead = 2855.92 Sentences processed = 175176
I0312 23:10:45.137243 140287431493376 run_squad.py:1333] 0 Training Performance = 61.3378 sentences/sec
I0312 23:14:00.550846 140287431493376 run_squad.py:1396] 0 Total Inference Time = 145.46 Inference Time W/O start up overhead = 131.86 Sentences processed = 10840
I0312 23:14:00.550973 140287431493376 run_squad.py:1397] 0 Inference Performance = 82.2095 sentences/sec
{"exact_match": 83.69914853358561, "f1": 90.8477003317459}

Multi-GPU training is enabled with the Horovod TensorFlow module. The following example runs fine tuning on 8 GPUs:

BERT_DIR=data/download/google_pretrained_weights/uncased_L-24_H-1024_A-16
SQUAD_DIR=data/download/squad/v1.1
mpi_command="mpirun -np 8 -H localhost:8 \
    --allow-run-as-root -bind-to none -map-by slot \
    -x NCCL_DEBUG=INFO \
    -x LD_LIBRARY_PATH \
    -x PATH -mca pml ob1 -mca btl ^openib" \
     python run_squad.py --horovod --vocab_file=$BERT_DIR/vocab.txt \
     --bert_config_file=$BERT_DIR/bert_config.json \
     --output_dir=/results --do_train --train_file=$SQUAD_DIR/train-v1.1.json

Multi-node

Multi-node runs can be launched on a pyxis/enroot Slurm cluster (see Requirements) with the run.sub script with the following command for a 4-node DGX1 example for both phase 1 and phase 2:

BATCHSIZE=16 LEARNING_RATE='1.875e-4' NUM_ACCUMULATION_STEPS=128 PHASE=1 sbatch -N4 --ntasks-per-node=8 run.sub
BATCHSIZE=2 LEARNING_RATE='1.25e-4' NUM_ACCUMULATION_STEPS=512 PHASE=1 sbatch -N4 --ntasks-per-node=8 run.sub

Checkpoint after phase 1 will be saved in checkpointdir specified in run.sub. The checkpoint will be automatically picked up to resume training on phase 2. Note that phase 2 should be run after phase 1.

Variables to re-run the Training performance results are available in the scripts/configs/configurations.yml file.

The batch variables BATCHSIZE, LEARNING_RATE, NUM_ACCUMULATION_STEPS refer to the Python arguments train_batch_size, learning_rate, num_accumulation_steps respectively. The variable PHASE refers to phase specific arguments available in run.sub.

Note that the run.sub script is a starting point that has to be adapted depending on the environment. In particular, variables such as datadir handle the location of the files for each phase.

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

Inference process

Inference on a fine tuned Question Answering system is performed using the run_squad.py script along with parameters defined in scripts/run_squad_inference.sh. Inference is supported on a single GPU.

The run_squad_inference.sh script trains a model and performs evaluation on the SQuAD dataset. By default, the inferencing script:

• Uses SQuAD v1.1 dataset
• Has FP16 precision enabled
• Is XLA enabled
• Evaluates the latest checkpoint present in /results with a batch size of 8

This script outputs predictions file to /results/predictions.json and computes F1 score and exact match score using SQuAD's evaluate file. Mount point of /results can be changed in the scripts/docker/launch.sh file.

The output log contains information about: Inference performance Inference Accuracy (F1 and exact match scores) on the Dev Set of SQuAD after evaluation.

The summary after inference is printed in the following format:

I0312 23:14:00.550846 140287431493376 run_squad.py:1396] 0 Total Inference Time = 145.46 Inference Time W/O start up overhead = 131.86 Sentences processed = 10840
I0312 23:14:00.550973 140287431493376 run_squad.py:1397] 0 Inference Performance = 82.2095 sentences/sec
{"exact_match": 83.69914853358561, "f1": 90.8477003317459}

Inference Process With TensorRT

NVIDIA TensorRT is a platform for high-performance deep learning inference. It includes a deep learning inference optimizer and runtime that delivers low latency and high-throughput for deep learning inference applications. More information on how to perform inference using TensorRT can be found in the subfolder ./https://raw.githubusercontent.com/NVIDIA/DeepLearningExamples/master/TensorFlow/LanguageModeling/BERT//trt/README.md

Deploying the BERT model using Triton Inference Server

The NVIDIA Triton Inference Server provides a datacenter and cloud inferencing solution optimized for NVIDIA GPUs. The server provides an inference service via an HTTP or gRPC endpoint, allowing remote clients to request inferencing for any number of GPU or CPU models being managed by the server. More information on how to perform inference using Triton Inference Server can be found in the subfolder ./triton/README.md.

BioBERT

Many works, including BioBERT, SciBERT, NCBI-BERT, ClinicalBERT (MIT), ClinicalBERT (NYU, Princeton), and others at BioNLP'19 workshop, show that pre-training of BERT on large biomedical text corpus such as PubMed results in better performance in biomedical text-mining tasks.

More information on how to download a biomedical corpus and pre-train as well as finetune for biomedical tasks can be found in the subfolder ./biobert/README.md.