To train your model using mixed or TF32 precision with Tensor Cores or using FP32, perform the following steps using the default parameters of the BART model on the CNN-DM/XSum datasets. For the specifics concerning training and inference, see the Advanced section.

1. Clone the repository.

git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/PyTorch/LanguageModeling/BART

2. Build the BART PyTorch container.

bash scripts/docker/build.sh

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

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

bash scripts/docker/launch.sh

The launch.sh script, by default, mounts the current directory to /workspace/bart.

4. Download and preprocess the dataset.

Use the following script to download and preprocess CNN DM data as well as XSum dataset.

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

Use the script to download Wikipedia, Common Crawl, and OpenWebTextCorpus for pre-training dataset

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

The pretraining dataset is 200GB+ and takes 24+ hours to download.

For downloading less dataset, you can change the date period of Common Crawl archive 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

Use the script to preprocess the pre-training dataset into LDDL Parquet shards

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

By default, the path to the data folder is set to /workspace/bart/data for ease of use in all the scripts.

5. Start pre-training

BART is designed to pre-train language representations. The following scripts are to replicate pre-training on Wikipedia, Common Crawl, and OpenWebTextCorpus 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.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>

6. Start summarizing.

Pretrained BART representations can be fine tuned for a state-of-the-art summarization system. From within the container, you can use the following script to run summarization on CNN DM dataset.

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>

This repository contains a number of predefined configurations to run the CNN+DM fine tuning on NVIDIA DGX-1 V100 or NVIDIA DGX A100 nodes in scripts/params/cnn_dm_params.sh. For example, to use the default DGX A100 8 gpu config, run:

bash scripts/run_summarization.sh $(source scripts/params/cnn_dm_params.sh && dgxa100_8gpu_bf16)

Similarly, configurations for XSum dataset are available in scripts/params/xsum_params.sh.

7. Start inference/predictions.

You can run the following script to run inference summarization using a fine-tuned checkpoint:

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>

This repository contains multiple predefined configurations in scripts/params/cnn_dm_params.sh and scripts/params/xsum_params.sh. For example, to run inference on CNN-DM with a checkpoint run:

bash scripts/run_eval_summarization.sh <INIT_CKPT> $(source scripts/params/cnn_dm_params.sh && dgxa100_8gpu_bf16_eval)

Now that you have your model trained and evaluated, you can choose to compare your training results with our Training accuracy results. You can also choose to benchmark yours performance to Training performance benchmark, or Inference performance benchmark. Following the steps in these sections will ensure that you achieve the same accuracy and performance results as stated in the Results section.

8. Run Custom Inference with the fine-tuned checkpoint We can write a simple few lines of code to run custom inference with the fine-tuned checkpoint.

from bart.modeling.modeling_bart import BartForConditionalGeneration
from bart.tokenization.tokenization_bart import BartTokenizer
from bart.configuration.configuration_bart import BartConfig
import json
config = BartConfig(**json.load(open('configs/config.json', "r")))
config.dtype = None
config.pre_ln = True
model_path = 'results/_epoch1_step2000.ckpt' # The fine-tuned checkpoint path
model = BartForConditionalGeneration.from_pretrained(model_path, config=config)
tokenizer = BartTokenizer.from_pretrained('facebook/bart-large-cnn')
ARTICLE_TO_SUMMARIZE = "NVIDIA Geforce Won't Run or Uninstall"
inputs = tokenizer([ARTICLE_TO_SUMMARIZE], max_length=1024, truncation=True, return_tensors='pt')
summary_ids = model.generate(inputs['input_ids'],
              num_beams=4,
              max_length=50,
              num_beam_groups=1,
              output_scores=False,
              return_dict_in_generate=False,
              encoder_no_repeat_ngram_size=0,
              diversity_penalty=0.0,
              early_stopping=True)
print([tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g in summary_ids])