Model Fine-Tuning

NVIDIA provides an easy-to-use tool to help you fine-tune the trained checkpoints on XQuAD for mT5 models.

XQuAD Fine-Tuning

To fine-tune the checkpoint on an XQuAD task, set the fine_tuning configuration in conf/config.yaml to mt5/xsquad, which specifies the fine-tuning configuration file as conf/fine_tuning/mt5/xsquad.yaml. Include the fine_tuning value in stages to run the fine-tuning pipeline.

Common

To specify the configuration for what tasks to run for fine-tuning, use the run.task_name configuration. And use all of The run configurations to define the job specific config:

run:
  name: ${.task_name}_${.model_train_name}
  time_limit: "04:00:00"
  dependency: "singleton"
  convert_name: convert_nemo
  model_train_name: mt5_220m
  task_name: "xquad"
  results_dir: ${base_results_dir}/${.model_train_name}/${.task_name}

Use the model configuration to specify the model checkpoint to load and its definition:

model:
  restore_from_path: ${base_results_dir}/${fine_tuning.run.model_train_name}/${fine_tuning.run.convert_name}/megatron_mt5.nemo # Path to a trained mt5 .nemo file
  tensor_model_parallel_size: 1
  pipeline_model_parallel_size: 1

Slurm

Set the configuration for a Slurm cluster in the conf/cluster/bcm.yaml file:

partition: null
account: null
exclusive: True
gpus_per_task: null
gpus_per_node: 8
mem: 0
overcommit: False
job_name_prefix: "nemo-megatron-"

Example

To run only the evaluation pipeline and not the data preparation, training, conversion, or inference pipelines, set the the stages section of conf/config.yaml to:

stages:
  - fine_tuning

Then enter:

python3 main.py

Base Command Platform

To run the fine-tuning script on Base Command Platform, set the cluster_type configuration in conf/config.yaml to bcp. This configuration can be overridden from the command line using hydra. This script must be launched in a multi-node job.

To run the fine-tuning pipeline to fine-tune a 390M mT5 model converted checkpoint stored in /mount/results/mt5_390m/convert_nemo, enter:

python3 /opt/NeMo-Framework-Launcher/launcher_scripts/main.py  fine_tuning=mt5/xquad stages=fine_tuning \
cluster_type=bcp \
launcher_scripts_path=/opt/NeMo-Framework-Launcher/launcher_scripts data_dir=/mount/data base_results_dir=/mount/results \
fine_tuning.run.model_train_name=mt5_390m \
fine_tuning.model.restore_from_path=/mount/results/mt5_390m/convert_nemo/results/megatron_mt5_xquad.nemo \
>> /results/finetune_mt5_log.txt 2>&1

The command above assumes that you mounted the data workspace in /mount/data, and the results workspace in /mount/results. stdout and stderr are redirected to the file /results/data_gpt3_log.txt, which you can download from NGC. Any other required configuration may be added to modify the command’s behavior.

Fine-tuning on Custom Tasks

NVIDIA supports fine-tuning on custom downstream tasks in T5 and mT5. To run a benchmark on your own dataset, split the original dataset into two files, i.e. a TXTfile corresponding to the source (context) data, and a TXT file corresponding to the target data. Each pair of corresponding lines of these two files forms a fine-tuning sample.

Custom fine-tuning configuration files are in conf/fine_tuning/t5/custom_task.yaml for T5 models and conf/fine_tuning/mt5/custom_task.yaml for mT5 models. The essential configurations are listed below. The dataset paths and preferred benchmark metrics must be specified.

data:
  train_ds:
    src_file_name: ??? # Path to the txt file corresponding to the source data.
    tgt_file_name: ??? # Path to the txt file corresponding to the target data.

  validation_ds:
    src_file_name: ??? # Path to the txt file corresponding to the source data.
    tgt_file_name: ??? # Path to the txt file corresponding to the target data.
    metric:
      name: "exact_string_match" # Name of the evaluation metric to use.
      average: null # Average the metric over the dataset. Options: ['macro', 'micro']. Works only for 'F1', 'accuracy' etc. Refer to torchmetrics for metrics where this is supported.
      num_classes: null

To submit a custom task job, follow the instructions in section XQuAD Fine-Tuning.