The following sections provide greater details of the datasets, running training and inference, and the training results.

Scripts and sample code

Descriptions of the key scripts and folders are provided below.

• data/ - Contains scripts for downloading and preparing individual datasets, and will contain downloaded and processed datasets.
• scripts/ - Contains shell scripts to launch the Docker container, data download, pre-training, fine-tuning and inference.
• results/ - Folder where all training and inference results get stored by default.
• run_squad.sh - Interface for launching question answering fine-tuning with run_tf_squad.py.
• run_pretraining.sh - Interface for launching ELECTRA pre-training with run_pretraining.py.
• finetune_ckpts_on_squad.sh - Interface for extracting and saving discriminator and generator from the pretrained checkpoint and run SQuAD fine-tuning on discriminator.
• build_pretraining_dataset.py - Creates tfrecord files from shared text files in the final step of dataset creation.
• postprocess_pretrained_ckpt.py - Converts pretrained checkpoint to discriminator checkpoint and generator checkpoint which can be fed into run_tf_squad.py.
• modeling.py - Implements the ELECTRA pre-training and fine-tuning model architectures with TensorFlow2.
• optimization.py - Implements the Adam optimizer, LAMB and the learning rate schedule with TensorFlow2.
• configuration.py - Implements parent class for model config.
• tokenization.py - Implements the ELECTRA tokenizer.
• run_pretraining.py - Implements ELECTRA pre-training.
• pretrain_utils.py - Utilities required for pre-training such as dynamic masking etc.,
• run_tf_squad.py - Implements fine-tuning training and evaluation for question answering on the SQuAD dataset.
• inference.py - Implements interactive question answering.
• postprocess_pretrained_ckpt.py - Implements extracting and saving the discriminator and the generator from the pretrained checkpoint.

Parameters

Pre-training parameters

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

In the parameters expected by scripts/run_pretraining.sh, p1 stands for phase 1 whereas p2 stands for phase 2 training. They are as follows:

• <training_batch_size_p1> is per-GPU batch size used for training. Larger batch sizes run more efficiently, but require more GPU memory. Default is 176.

• <learning_rate_p1> is the base learning rate for training. Default is 6e-3.

• <precision> is the type of math in your model, can be either fp32 or amp. Default is amp. The options mean:

  • FP32: 32-bit IEEE single precision float format.
  • AMP: Automatic mixed precision 16 and 32-bit float format.

• <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. Default is 8.

• <warmup_steps_p1> is the percentage of training steps used for warm-up at the start of training. Default is 2000.

• <train_steps_p1> is the total number of training steps. Default is 10000.

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

• <resume_training> if set to true, training should resume from the latest model in /results/checkpoints. Default is false.

• <accumulate_gradient> a flag indicating whether a larger batch should be simulated with gradient accumulation. Default is true.

• <gradient_accumulation_steps_p1> an integer indicating the number of steps to accumulate gradients over. Effective batch size / GPU = training_batch_size x gradient_accumulation_steps. Default is 48.

• <seed> random seed for the run.

• <training_batch_size_p2> is per-GPU batch size used for training in phase 2. Larger batch sizes run more efficiently, but require more memory. Default is 24.

• <learning_rate_p2> is the base learning rate for training phase 2. Default is 4e-3.

• <warmup_steps_p2> is the percentage of training steps used for warm-up at the start of training. Default is 200.

• <training_steps_p2> is the total number of training steps for phase 2, to be continued in addition to phase 1. Default is 930.

• <gradient_accumulation_steps_p2> an integer indicating the number of steps to accumulate gradients over in phase 2. Effective batch size / GPU = training_batch_size_p2 * gradient_accumulation_steps_p2. Default is 144.

• <init_checkpoint> A checkpoint to start the pre-training routine on (Usually a ELECTRA pretrained checkpoint). Default is None.

The complete list of the available parameters for the run_pretraining.py script are:

  --model_name MODEL_NAME
                              - Model name, used to define the name of the results folder.
 
  --pretrain_tfrecords PRETRAIN_TFRECORDS
                              - Specifies tfrecord files used for pretraining.
 
  --max_seq_length 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.
 
  --mask_prob MASK_PROB       - Percentage of input tokens to mask out / replace.
 
  --disc_weight DISC_WEIGHT
                              - Ratio of discriminator loss over generator loss.
 
  --generator_hidden_size GENERATOR_HIDDEN_SIZE
                              - Fraction of discriminator hidden size for generator.
 
  --train_batch_size TRAIN_BATCH_SIZE
                              - Batch size per GPU for training.
 
  --learning_rate LEARNING_RATE
                              - The initial learning rate for the optimizer.
 
  --num_train_steps NUM_TRAIN_STEPS
                              - Total number of training steps to perform.
 
  --num_warmup_steps NUM_WARMUP_STEPS
                              - Number of steps of training to perform linear learning
                                rate warmup for. For example, 0.1 = 10% of training.
 
  --seed SEED                 - Sets the seed to use for random number generation.
 
  --gradient_accumulation_steps GRADIENT_ACCUMULATION_STEPS
                              - Number of update steps to accumulate before
                                performing a backward/update pass.
 
  --fp16_compression          - Whether to use 16-bit all reduce
 
  --amp                       - If set, will perform computations using
                                automatic mixed precision.
 
  --log_freq LOG_FREQ         - If set, the script will output the training
                                loss every LOG_FREQ steps.
 
  --save_checkpoints_steps SAVE_CHECKPOINTS_STEPS
                              - Checkpoints saving frequency.
 
  --keep_checkpoint_max KEEP_CHECKPOINT_MAX
                              - Maximum number of checkpoints to keep.
 
  --restore_checkpoint RESTORE_CHECKPOINT
                              - Whether to restore from a checkpoint; if specified,
                                set to `path-to-checkpoint` or `latest`
 
  --phase2                    - Specified if training on phase 2 only. 
                                If not specified, default pre-training is on phase 1.
 
  --optimizer OPTIMIZER       - Specifies optimizer, `adam` or `lamb`.
 
  --skip_adaptive             - Whether to apply adaptive learning rate on LayerNorm and biases.
 
  --gradient_accumulation_steps GRADIENT_ACCUMULATION_STEPS
                              - Number of steps to accumulate gradients across before
                                performing an update.
 
  --lr_decay_power LR_DECAY_POWER
                              - Learning rate polynomial decay power.
 
  --opt_beta_1 OPT_BETA_1     - beta2 of optimizer.
 
  --opt_beta_2 OPT_BETA_2     - beta2 of optimizer.
 
  --end_lr END_LR             - Ending learning rate.
 

Fine-tuning parameters

Default arguments are listed below in the order scripts/run_squad.sh expects:

• ELECTRA MODEL - The default is "google/electra-base-discriminator".
• Number of training Epochs - The default is 2.
• Batch size - The default is 16.
• Learning rate - The default is 4e-4.
• Precision (either amp, tf32 or fp32) - The default is amp.
• Number of GPUs - The default is 8.
• Seed - The default is 1.
• SQuAD version - The default is 1.1
• SQuAD directory - The default is /workspace/electra/data/download/squad/v$SQUAD_VERSION.
• Output directory for result - The default is results/.
• Initialize checkpoint - The default is "None"
• Mode (train, eval, train_eval, prediction) - The default is train_eval.

The script saves the checkpoint at the end of each epoch to the checkpoints/ folder.

The main script run_tf_squad.py specific parameters are:

 --electra_model ELECTRA_MODEL     - Specifies the type of ELECTRA model to use;
                                     should be the discriminator of a pretrained checkpoint(output of postprocess_pretrained_ckpt.py)
                                     or one of the following:
              google/electra-small-generator
                google/electra-base-generator
                google/electra-large-generator
                google/electra-small-discriminator
                google/electra-base-discriminator
                google/electra-large-discriminator
 
 --amp                        - If set, will perform computations using
                                automatic mixed precision.
 
 --data_dir DATA_DIR          - Path to the SQuAD json for training and evaluation.
 
 --max_seq_length 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.
 
 --doc_stride DOC_STRIDE      - When splitting up a long document into chunks
                                this parameters sets how much stride to take
                                between chunks of tokens.
 
 --max_query_length MAX_QUERY_LENGTH
                              - The maximum number of tokens for the question.
                                Questions longer than <max_query_length>
                                will be truncated to the value specified.
 
 --n_best_size N_BEST_SIZE       - The total number of n-best predictions to
                                generate in the nbest_predictions.json
                                output file.
 
 --max_answer_length 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.
    
 --joint_head <True|False>    - If true, beam search will be used to jointly predict
                                the start and end positions. Default is True.
 
 --beam_size BEAM_SIZE        - The beam size used to do joint predictions. The default value is 5.
 
 --verbose_logging            - If true, all the warnings related to data
                                processing will be printed. A number of warnings
                                are expected for a normal SQuAD evaluation.
 
 --do_lower_case              - Whether to lower case the input text. Set to
                                true for uncased models and false for cased models.
 
 --version_2_with_negative       - If true, the SQuAD examples contain questions
                                that do not have an answer.
 
 --null_score_diff_threshold NULL_SCORE_DIFF_THRES HOLD
                              - A null answer will be predicted if null_score
                                is greater than NULL_SCORE_DIFF_THRESHOLD.

Command-line options

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

python run_pretraining.py --help

python run_tf_squad.py --help

Detailed descriptions of command-line options can be found in the Parameters section.

Getting the data

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

The preparation of the pre-training dataset is described in the dataPrep.py script found in the data/ folder. The component steps in the automated scripts 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.

3. Sentence segmentation - the corpus text file is processed into separate sentences.

4. Sharding - the sentence segmented corpus file is split into a number of uniformly distributed smaller text documents.

5. tfrecord file creation - each text file shard is processed by the build_pretraining_dataset.py script to produce a corresponding tfrecord file. The script generates input data for the input text shard.

The tools used for preparing the BookCorpus and Wikipedia datasets can be applied to prepare an arbitrary corpus. The create_datasets_from_start.sh script in the data/ directory applies sentence segmentation, sharding, and tfrecord file creation given an arbitrary text file containing a document-separated text corpus.

For fine-tuning a pre-trained ELECTRA model for specific tasks, by default this repository prepares the following dataset:

• SQuAD: for question answering

Depending on the speed of your internet connection, this process takes about a day to complete. The BookCorpus server could sometimes get overloaded and also contain broken links resulting in HTTP 403 and 503 errors. You can either skip the missing files or retry downloading at a later time.

Multi-dataset

This repository provides functionality to combine multiple datasets into a single dataset for pre-training on a diverse text corpus at the shard level. Currently Wikipedia and BookCorpus get merged in data/create_datasets_from_start.sh. Snippets to download and format more text corpuses can be added to data/dataPrep.py. The sharding scheme combines multiple corpuses together and splits them into the required number of training(90%) and testing(10%) shards. Once the data is sharded, the build_pretraining_dataset.py converts raw text shards to tokenized segments and saves the dataset to the data directory in TFRecord format. This dataset can now be used to pre-train ELECTRA.

Training process

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

Pre-training

Pre-training is performed using run_pretraining.py along with parameters defined in scripts/run_pretraining.sh and scripts/configs/pretrain_configs.sh.

The run_pretraining.sh script runs a job on a single node that trains the ELECTRA-base model from scratch using Wikipedia and BookCorpus datasets as training data using the LAMB optimizer.

Phase 1: (Maximum sequence length of 128)

• Runs on 8 GPUs with training batch size of 176 per GPU
• Uses a learning rate of 6e-3
• Has FP16 precision enabled
• Runs for 10000 steps, where the first 2000 are warm-up steps
• Saves a checkpoint every 500 iterations (keeps only the latest 5 checkpoints) and at the end of training. All checkpoints, and training logs are saved to the results/<model_name> directory.
• Creates a log file containing all the output

Phase 2: (Maximum sequence length of 512)

• Runs on 8 GPUs with training batch size of 24 per GPU
• Uses a learning rate of 4e-3
• Has FP16 precision enabled
• Runs for 930 steps, where the first 200 are warm-up steps
• Saves a checkpoint every 500 iterations (keeps only the latest 5 checkpoints) and at the end of training. All checkpoints, and training logs are saved to the results/<model_name> directory.
• Creates a log file containing all the output

Specific configs available at scripts/configs/pretrain_config.sh can be run as follows:

bash scripts/run_pretraining.sh $(source scripts/configs/pretrain_config.sh && dgxa100_8gpu_amp)
bash scripts/run_pretraining.sh $(source scripts/configs/pretrain_config.sh && dgx2_16gpu_amp)
bash scripts/run_pretraining.sh $(source scripts/configs/pretrain_config.sh && dgx1_8gpu_amp)

The above commands will train ELECTRA based on Wikipedia and BookCorpus to state-of-the-art accuracy on any DGX platform using FP16 arithmetic. Around 96% of the training sequences are of length 128 (phase 1 of training) and less than 4% of the training sequences are of length 512 (phase 2 of training).

In order to run pre-training routine on an initial checkpoint, perform the following in scripts/run_pretraining.sh:

• set restore_checkpoint=<path_to_checkpoint>
• Note: The parameter value assigned to --model_size during training should remain unchanged. Also, to resume pre-training on your corpus of choice, the training dataset should be created using the same vocabulary file used in data/create_datasets_from_start.sh.

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 48-node NVIDIA DGX A100 example for both phase 1 and phase 2:

BATCHSIZE=176 LR=6e-3 GRAD_ACCUM_STEPS=1 PHASE=1 STEPS=10000 WARMUP=2000 b1=0.878 b2=0.974 decay=0.5 skip_adaptive=yes end_lr=0.0 sbatch N48 --ntasks-per-node=8 run.sub
BATCHSIZE=24 LR=4e-3 GRAD_ACCUM_STEPS=3 PHASE=2 STEPS=930 WARMUP=200 b1=0.878 b2=0.974 decay=0.5 skip_adaptive=yes end_lr=0.0 sbatch N48 --ntasks-per-node=8 run.sub

Checkpoint after phase 1 will be saved in <results_dir>/models/<model_name>. The checkpoint will be automatically picked up to resume training on phase 2. Note that phase 2 should be run after phase 1.

The batch variables BATCHSIZE, LR, GRAD_ACCUM_STEPS, PHASE, STEPS, WARMUP, b1, b2, decay, skip_adaptive and end_lr refer to the Python arguments train_batch_size, learning_rate, gradient_accumulation_steps, phase2, num_train_steps, num_warmup_steps, opt_beta_1, opt_beta_2, lr_decay_power, skip_adaptive and end_lr in run_pretraining.py respectively.

Note that the run.sub script is a starting point that has to be adapted depending on the environment. In particular, variables such as docker_image and 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.

Fine-tuning

Fine-tuning is provided for a variety of tasks. The following tasks are included with this repository through the following scripts:

• Question Answering (scripts/run_squad.sh)

By default, each Python script implements fine-tuning a pre-trained ELECTRA model for a specified number of training epochs as well as evaluation of the fine-tuned model. Each shell script invokes the associated Python script with the following default parameters:

• Uses 8 GPUs
• Has FP16 precision enabled
• Has XLA enabled
• Saves a checkpoint at the end of training to the checkpoints/ folder

Specific configs available at scripts/configs/squad_configs.sh can be run as follows:

bash scripts/run_squad.sh $(source scripts/configs/squad_config.sh && dgxa100_8gpu_amp) train_eval
bash scripts/run_squad.sh $(source scripts/configs/squad_config.sh && dgx2_16gpu_amp) train_eval
bash scripts/run_squad.sh $(source scripts/configs/squad_config.sh && dgx1_8gpu_amp) train_eval

Fine-tuning Python scripts implement support for mixed precision and multi-GPU training through Horovod. For a full list of parameters and associated explanations, see the Parameters section.

All fine-tuning shell scripts have the same positional arguments, outlined below:

bash scripts/run_squad.sh <pretrained electra model> <epochs> <batch size> <learning rate> <amp|fp32> <num_gpus> <seed> <SQuAD version> <path to SQuAD dataset> <results directory> <checkpoint_to_load> <mode (either `train`, `eval` or `train_eval`)>

By default, the mode positional argument is set to train_eval. See the Fine-tuning parameters for explanations of each positional argument.

Note: The first positional argument (the path to the checkpoint to load) is required.

Each fine-tuning script assumes that the corresponding dataset files exist in the data/ directory or separate path can be a command-line input to run_squad.sh.

Inference process

Fine-tuning inference

Evaluation fine-tuning is enabled by the same scripts as training:

• Question Answering (scripts/run_squad.sh)

The mode positional argument of the shell script is used to run in evaluation mode. The fine-tuned ELECTRA model will be run on the evaluation dataset, and the evaluation loss and accuracy will be displayed.

Each inference shell script expects dataset files to exist in the same locations as the corresponding training scripts. The inference scripts can be run with default settings. By setting the mode variable in the script to either eval or prediction flag, you can choose between running predictions and evaluating them on a given dataset or just the former.

bash scripts/run_squad.sh <pretrained electra model> <epochs> <batch size> <learning rate> <amp|fp32> <num_gpus> <seed> <SQuAD version> <path to SQuAD dataset> <results directory> <path to fine-tuned model checkpoint> <eval or prediction>

To run inference interactively on question-context pairs, use the script run_inference.py as follows:

python run_inference.py --electra_model <electra_model_type> --init_checkpoint <fine_tuned_checkpoint> --question="What food does Harry like?" --context="My name is Harry and I grew up in Canada. I love apples."