ELECTRA is method of pre-training language representations which outperforms existing techniques on a wide array of NLP tasks.
To train your model using mixed precision or TF32 precision with Tensor Cores or using FP32, perform the following steps using the default parameters of the ELECTRA model. The default parameters for pre-training have been set to run on both 8x A100 40G and 8 x V100 32G GPUs. For the specifics concerning training and inference, see the Advanced section.
- Clone the repository.
git clone https://github.com/NVIDIA/DeepLearningExamples.git
cd DeepLearningExamples/TensorFlow2/LanguageModeling/ELECTRA
- Build ELECTRA on top of the NGC container.
bash scripts/docker/build.sh
- Start an interactive session in the NGC container to run data download, training and inference.
bash scripts/docker/launch.sh
Resultant logs of pre-training and fine-tuning routines are stored in the results/ folder. Checkpoints are stored in the results/<model-name>/ folder.
Required data is downloaded into the data/ directory by default.
- Download and preprocess the dataset.
This repository provides scripts to download, verify, and extract the following datasets:
- SQuAD (fine-tuning for question answering)
- Wikipedia (pre-training)
- BookCorpus (pre-training)
To download, verify, extract the datasets, and create the shards in tfrecord format, run:
/workspace/electra/data/create_datasets_from_start.sh
Note: For fine-tuning only, Wikipedia and Bookscorpus dataset download and preprocessing can be skipped by commenting it out.
- Download Wikipedia only for pretraining
The pre-training dataset is 170GB+ and takes 15+ hours to download. The BookCorpus server most of the time gets overloaded and also contains broken links resulting in HTTP 403 and 503 errors. Hence, it is recommended to skip downloading BookCorpus data by running:
/workspace/electra/data/create_datasets_from_start.sh wiki_only
- Download Wikipedia and BookCorpus
Users are welcome to download the 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:
/workspace/electra/data/create_datasets_from_start.sh wiki_books
Note: Not using the BookCorpus can potentially change the final accuracy on a few downstream tasks.
- Start pretraining.
To run on a single node 8 x V100 32G, from within the container, you can use the following script to run pre-training.
bash scripts/run_pretraining.sh
The default hyperparameters are set to run on both 8 x A100 40G and 8 x V100 32G.
For the other platforms, the configs present in scripts/configs/pretrain_config.sh can be used as shown below:
bash scripts/run_pretraining.sh $(source scripts/configs/pretrain_config.sh && dgxa100_8gpu_amp)
To run pre-training on multiple nodes, see the Multi-node section.
- Postprocess pretrained checkpoint and fine-tune on SQuAD dataset
The above pretrained ELECTRA model representations can be fine-tuned with just one additional output layer for a state-of-the-art question answering system. Running the following script extracts and saves the discriminator and generator from the pretrained checkpoint and fine-tunes the discriminator on SQuAD:
checkpoints=results/base/checkpoints bash scripts/finetune_ckpts_on_squad.sh
It internally runs postprocess_pretrained_ckpt.py which extracts and saves the discriminator and the generator from the pretrained checkpoint.
The default hyperparameters are set to run on 8 x V100 16G.
To run fine-tuning with the SQuAD dataset on Google's pretrained checkpoints, do the following.
bash scripts/run_squad.sh
For other platforms, configs present in scripts/configs/squad_config.sh can be used as shown below:
bash scripts/run_squad.sh $(source scripts/configs/squad_config.sh && dgxa100_8gpu_amp) train_eval
- Start validation/evaluation.
Validation can be performed by running:
bash scripts/run_squad.sh $(source scripts/configs/squad_config.sh && dgxa100_8gpu_amp) eval
Running training first is required to generate needed checkpoints.
- Start inference/predictions.
Inference can be performed by running:
bash scripts/run_squad.sh $(source scripts/configs/squad_config.sh && dgxa100_8gpu_amp) prediction
Inference predictions are saved to <OUTPUT_DIRECTORY>/predictions.json.