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:

• pretrain.py - Serves as entry point for pre-training
• finetune.py - Serves as entry point for fine-tuning
• run_eval.py - Serves as entry point for inference
• Dockerfile - Container with the basic set of dependencies to run BART

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

• run_summarization.sh - Runs summarization finetuning followed by inference using the finetune.py and run_eval.py files.
• run_summarization_eval.sh - Runs inference on fine tuned checkpoint using the run_eval.py file.
• get_data.sh - Preprocesses CNN-DM dataset as well as downloads and preprocesses XSum dataset.
• get_pretraining_data.sh - Downloads pre-train dataset.
• preprocess_pretrain_data.sh - Preprocesses pre-train dataset.

Other folders included in the root directory are:

• data/ - Necessary folder to download datasets required for fine tuning of BART.
• src/ - Modeling, tokenization and configuration functionality files for implementing the BART model.
• utils/ - Necessary utility files for BART model.

Parameters

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

--config_path: The configuration file corresponding to BART Model
--warmup_steps: Number of WARMUP_STEPS
--max_steps: Number of MAX_STEPS
--data_dir: Location to DATA_DIR
--learning_rate: Learning Rate
--n_val: Number of validation examples to test for early stopping
--train_batch_size: Train batch size
--gradient_accumulation_steps: Number of accumulation steps
--max_source_length: Maximum source length
--max_target_length: Maximum target length
--val_max_target_length: Maximum length of validation tokens
--eval_max_gen_length: Maximum length while generating validation tokens
--weight_decay: weight decay
--dropout: drop out
--lamb: Whether to use LAMB optimizer
--pre_ln: Whether to use Pre-LN architecture
--allreduce_post_accumulation_half_precision: Whether to do fp16/bf16 allreduce post accumulation

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

--config_path: The configuration file corresponding to BART Model
--warmup_steps: Number of WARMUP_STEPS
--max_steps: Number of MAX_STEPS
--data_dir: Location to DATA_DIR
--learning_rate: Learning Rate
--n_val: Number of validation examples to test for early stopping
--train_batch_size: Train batch size
--gradient_accumulation_steps: Number of accumulation steps
--max_source_length: Maximum source length
--max_target_length: Maximum target length
--val_max_target_length: Maximum length of validation tokens
--eval_max_gen_length: Maximum length while generating validation tokens
--weight_decay: weight decay
--dropout: drop out
--pre_ln: Whether to use Pre-LN architecture
--allreduce_post_accumulation_half_precision: Whether to do fp16/bf16 allreduce post accumulation

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 pretrain.py --help
python finetune.py --help
python run_eval.py --help

Getting the data

For pre-training BART, we use the concatenation of Wikipedia, Common Crawl, and OpenWebTextCorpus.

Common Crawl is an archieve of news articles from small and major publishers world wide, which is provided from commoncrawl.org.

OpenWebTextCorpus is an open source effort to reproduce OpenAI's WebText dataset. The distribution was created by Aaron Gokaslan and Vanya Cohen of Brown University.

For fine-tuning BART, we have tested fine tuning the BART model on summarization benchmarks such as CNN-DM and XSum.

CNN-DM is a concatenation of CNN Stories as well as Daily Mail Stories. CNN consists of approximately 90k documents whereas Daily Mail consists of 197k documents.

These documents are preprocessed to have two features:

• Article: text of news article, used as the document to be summarized
• Highlights: joined text of highlights with and around each highlight, which is the target summary

XSum, on the other hand, is also a single-document summarization task dataset but one that favors abstractive modeling. It consists of BBC articles and single sentence summaries. It consists of approximately 230k articles.

Dataset guidelines

The repository contains scripts to preprocess and download data. It can be run as:

bash scripts/get_data.sh <path to output data folder>

The script downloads CNN and DM raw data from here. The raw data is preprocessed using scripts from repository. The stories are first tokenized, written to serialized binary files and split into train, test and validation sets.

The script also downloads the XSum dataset from the HuggingFace storage.

bash scripts/get_pretraining_data.sh <path to data folder>

The script uses the LDDL downloader to download Wikipedia, Common Crawl, and OpenWebTextCorpus dataset. The Common Crawl is downloaded by news-please. And OpenWebTextCorpus is downloaded from here

For downloading less dataset, you can change the date period of Common Crawl archive in the script to take less time. For example:

download_common_crawl \
    --outdir $data_folder/common_crawl \
    --warc-files-start-date 2016-09-01 \
    --warc-files-end-date 2016-10-31 \
    --start-date 2016-09-01 \
    --end-date 2016-10-31

bash scripts/preprocess_pretrain_data.sh <path to Wikipedia> <path to Common Crawl> <path to OpenWebTextCorpus> <path to data folder>

The script uses the LDDL preprocessor and load balancer to preprocess the pre-training dataset into Parquet shards which are then streamed during the pre-training by the LDDL data loader.

The script by default stores the data into the /workspace/bart/data folder.

Training process

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

Pre-training

Pre-training BART is done using scripts/run_pretraining.sh script that, in turn, uses the pretrain.py file to perform training.

For example, it can be invoked by calling:

bash scripts/run_pretraining.sh <train_batch_size_phase1> <train_batch_size_phase2> <learning_rate_phase1> <learning_rate_phase2> <precision> <use_preln> <num_gpus> <warmup_steps_phase1> <warmup_steps_phase2> <train_steps_phase1> <train_steps_phase2> <save_checkpoint_steps> <num_accumulation_phase1> <num_accumulation_steps_phase2> <config_path>

Where:

• train_batch_size_phase* - per-GPU batch size used for training in the respective phase
• learning_rate_phase* - Learning rate in the respective phase
• precision - fp16/bf16/fp32/tf32 precision for training
• use_preln - Whether to use Pre-LN architecture
• num_gpus - number of GPUs to run training with
• warmup_steps_phase* - Number of warmup steps for learning rate scheduler in the respective phase
• train_steps_phase* - Number of training steps in the respective phase
• save_checkpoint_steps - Number of steps for saving checkpoint
• num_accumulation_phase* - Number of accumulation steps for an effective larger training batch size in the respective phase
• config_path - path to configuration file of BART Model

By default, the training script stores results to results/bart_pyt_pretraining and runs with:

bash scripts/run_pretraining.sh 200 32 5e-3 4e-3 bf16 true 8 2166 200 95040 7560 100 40 120 configs/config.json

Fine-tuning

Training BART for summarization is done using scripts/run_summarization.sh script that, in turn, uses the finetune.py file to perform training.

For example, it can be invoked by calling:

bash scripts/run_summarization.sh <DATA_DIR> <CKPT_PATH> <CONFIG_PATH> <NUM_GPU> <LR> <BS> <ACCUM> <PREC> <TRAIN_STEPS> <WARMUP_STEPS> <MAX_SOURCE_LEN> <MAX_TARGET_LEN> <EVAL_BEAMS> <EVAL_BS> <PRED_BS> <PRELN>

Where:

• DATA_DIR - path to data directory with train/test/val files.
• CONFIG_PATH - path to configuration file of BART Model
• NUM_GPU - number of GPUs to run training with
• LR - Learning rate for training process
• BS - Training batch size
• ACCUM - Number of accumulation steps for an effective larger training batch size
• PREC - fp16/fp32/tf32 precision for training and inference
• TRAIN_STEPS - Maximum number of training steps
• WARMUP_STEPS - Number of warmup steps for learning rate scheduler
• MAX_SOURCE_LEN - Maximum source length of articles
• MAX_TARGET_LEN - Maximum target length of summaries
• EVAL_BEAMS - Number of beams to run during inference
• EVAL_BS - Batch size for inference during validation
• PRED_BS - Batch size for inference on test data
• PRELN - Whether to use Pre-LN architecture

By default, the training script stores results to results/bart_pyt_${DATESTAMP} and runs with:

bash scripts/run_summarization.sh data/cnn_dm/ data/nvidia_pretrained/bart_large/ configs/config.json 8 1.25e-4 40 1 bf16 2000 50 1024 142 4 128 true

These parameters train CNN-DM with reasonable rouge scores on a LUNA with 80GB A100 cards. Other tested configurations are available under scripts/params/cnn_dm_params.sh for CNN-DM and scripts/params/xsum_params.sh for XSum datasets.

Inference process

Evaluating BART for summarization is done using scripts/run_eval_summarization.sh script that, in turn, uses the run_eval.py file to perform inference.

For example, it can be invoked by calling:

bash scripts/run_eval_summarization.sh <INIT_CKPT> <PRED_BS> <NUM_GPU> <PRECISION> <EVAL_BEAMS> <MAX_SOURCE_LEN> <MAX_TARGET_LEN> <DATA_DIR> <CONFIG_PATH> <PRELN>

Where:

• INIT_CKPT - Model name or path to initialize BART Model weights with.
• PRED_BS - Batch size for inference
• NUM_GPU - number of GPUs to run training with
• PRECISION - FP16/FP32/TF32 precision for training and inference
• EVAL_BEAMS - Number of beams to run during inference
• MAX_SOURCE_LEN - Maximum source length of articles
• MAX_TARGET_LEN - Maximum target length of summaries
• DATA_DIR - path to data directory with train/test/val files.
• CONFIG_PATH - path to configuration file of BART Model
• PRELN - Whether to use Pre-LN architecture

By default, the training script stores results to results/bart_pyt_inference_${DATESTAMP} and runs with:

bash scripts/run_eval_summarization.sh data/nvidia-pretrained/model.ckpt128 8 fp16 4 1024 142 data/cnn_dm/ configs/config.json

These parameters run inference on CNN-DM on a DGX A100 with 80GB A100 cards. For XSum, try EVAL_BEAMS=6, MAX_SOURCE_LEN=1024 and MAX_TARGET_LEN=60. For other GPUS/precisions, reduce PRED_BS as indicated in scripts/params.