To pretrain or fine tune your model for Question Answering using mixed precision with Tensor Cores or using FP32/TF32, perform the following steps using the default parameters of the BERT model.

1. Clone the repository.

git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/TensorFlow2/LanguageModeling/BERT

2. Build the BERT TensorFlow NGC container.

bash scripts/docker/build.sh

3. Download and preprocess the dataset.

This repository provides scripts to download, verify, extract the SQuAD dataset and pretrained weights for fine tuning as well as Wikipedia and BookCorpus dataset for pre-training.

To download, verify, and extract the required datasets, run:

bash scripts/data_download.sh

The script launches a Docker container with the current directory mounted and downloads the datasets to a data/ folder on the host.

• Download datasets for fine tuning and pretraining

bash scripts/data_download.sh all

• Download datasets for fine tuning only

bash scripts/data_download.sh squad

• Download Wikipedia only for pretraining

The pretraining dataset is 170GB+ and takes 15+ hours to download. The BookCorpus server most of the times get overloaded and also contains broken links resulting in HTTP 403 and 503 errors. Hence, it is recommended to skip downloading BookCorpus data by running:

bash scripts/data_download.sh pretrained wiki_only

Note: Wikipedia dataset is 17GB and takes about one hour to download.

• Download Wikipedia and BookCorpus

Users can download BookCorpus from other sources to match our accuracy, or repeatedly try our script until the required number of files are downloaded by running the following:

bash scripts/data_download.sh pretrained wiki_books

Note: Not using BookCorpus can potentially change final accuracy on a few downstream tasks.

4. Download the pretrained models from NGC.

We have uploaded checkpoints for fine tuning with various configurations on the NGC Model Registry. You can download them directly from the NGC model catalog. Download them to the results/models/ to easily access them in your scripts.

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

After you build the container image and download the data, you can start an interactive CLI session as follows:

bash scripts/docker/launch.sh

6. Start pre-training.

BERT is designed to pre-train deep bidirectional representations for language representations. The following scripts are to replicate pre-training on Wikipedia and BookCorpus from the LAMB paper. These scripts are general and can be used for pre-training language representations on any corpus of choice.

From within the container, you can use the following script to run pre-training using LAMB.

bash 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>

For BERT-Large FP16 training with XLA using a DGX-2H, run:

bash scripts/run_pretraining_lamb.sh 60 10 8 7.5e-4 5e-4 fp16 true 8 2000 200 7820 100 64 192 large

This repository also contains a number of predefined configurations to run the LAMB pretraining on NVIDIA DGX-1, NVIDIA DGX-2H or NVIDIA DGX A100 nodes in scripts/configs/pretrain_config.sh. For example, to use the default DGX A100 8 GPU config, run:

bash scripts/run_pretraining_lamb.sh $(source scripts/configs/pretrain_config.sh && dgxa100_8gpu_fp16)

Alternatively, to run pre-training with Adam as in the original BERT paper from within the container, run:

bash scripts/run_pretraining_adam.sh <train_batch_size_per_gpu> <eval_batch_size> <learning_rate_per_gpu> <precision> <use_xla> <num_gpus> <warmup_steps> <train_steps> <save_checkpoint_steps>

7. Start fine tuning.

The above pretrained BERT representations can be fine tuned with just one additional output layer for a state-of-the-art Question Answering system. From within the container, you can use the following script to run fine-training for SQuAD.

bash scripts/run_squad.sh <num_gpus> <batch_size_per_gpu> <learning_rate_per_gpu> <precision> <use_xla> <bert_model> <squad_version> <epochs>

For SQuAD 1.1 FP16 training with XLA using a DGX-1 V100 32G, run:

bash scripts/run_squad.sh 8 12 5e-6 fp16 true large 1.1 2

For SQuAD 2.0 FP32 training without XLA using a DGX-1 V100 32G, run:

bash scripts/run_squad.sh 8 8 5e-6 fp32 false large 2.0 2

The fine-tuned checkpoint will save to /results/tf_bert_finetuning_squad_xxxxxx/ctl_step_xxx.ckpt-x

8. Start validation/evaluation.

The run_squad_inference.sh script runs inference on a checkpoint fine tuned for SQuAD and evaluates the validity of predictions on the basis of exact match and F1 score.

bash scripts/run_squad_inference.sh <init_checkpoint> <batch_size> <precision> <use_xla> <bert_model> <squad_version>

The init_checkpoint is the fine-tuned checkpoint path. For example, we take the checkpoint from previous step /results/tf_bert_finetuning_squad_xxxxxx/ctl_step_xxx.ckpt-x and rename to /results/model.ckpt. SQuAD 1.1 and SQuAD 2.0 should be different checkpoints.

For SQuAD 2.0 FP16 inference with XLA using a DGX-1 V100 32G, run:

bash scripts/run_squad_inference.sh /results/model.ckpt 8 fp16 true large 2.0

For SQuAD 1.1 FP32 inference without XLA using a DGX-1 V100 32G, run:

bash scripts/run_squad_inference.sh /results/model.ckpt 8 fp32 false large 1.1