run_seq_classification.py

# Copyright (c) 2022 Graphcore Ltd. All rights reserved.

import argparse
from datetime import datetime
import logging
import math

import tensorflow as tf
from tensorflow.python import ipu
from tensorflow.python.ipu.gradient_accumulation import GradientAccumulationReductionMethod
from transformers import BertConfig, TFBertForSequenceClassification
import wandb

from data_utils.batch_config import BatchConfig, Task
from data_utils.glue.load_glue_data import get_glue_data
from keras_extensions.callbacks.callback_factory import CallbackFactory
from keras_extensions.learning_rate.scheduler_builder import get_lr_scheduler
from keras_extensions.optimization import get_optimizer
from model.convert_bert_model import convert_tf_bert_model, post_process_bert_input_layer
from model.accuracy import classification_accuracy_fn
from model.losses import ClassificationLossFunction, ClassificationLossFunctionRegression
from model.pipeline_stage_names import PIPELINE_ALLOCATE_PREVIOUS, PIPELINE_NAMES
from utilities.argparser import add_shared_arguments, add_glue_arguments, combine_config_file_with_args
from utilities.assign_pipeline_stages import GluePipelineStagesAssigner
from utilities.checkpoint_utility import load_checkpoint_into_model
from utilities.ipu_utils import create_ipu_strategy, get_poplar_options_per_pipeline_stage, set_random_seeds
from utilities.metric_enqueuer import wrap_loss_in_enqueuer, wrap_stateless_metric_in_enqueuer
from utilities.options import GLUEOptions


def fine_tune_glue(config):
    """Main function to run fine tuning BERT for GLUE.
    :param config: A pydantic model object that contains the configuration to
        options for this application. See utilities/options.py for accepted
        options.
    """
    if config.bert_model_name is None and config.pretrained_ckpt_path is None:
        logging.warning(
            "GLUE requires a pretrained model, either"
            " `bert_model_name` (via config file) or"
            " `pretrained_ckpt_path` (via config file or"
            " `--pretrained-ckpt-path` command line argument)"
            " but none provided."
        )
    if config.bert_model_name is not None and config.pretrained_ckpt_path is not None:
        logging.warning(
            "Only one checkpoint is accepted, but two provided:"
            " `bert_model_name`={config.bert_model_name}, and"
            " `pretrained_ckpt_path`={config.pretrained_ckpt_path}."
        )

    universal_run_name = config.config.stem if config.name is None else config.name
    universal_run_name += f"-{config.glue_task}"
    universal_run_name += f"-{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    logging.info(f"Universal name for run: {universal_run_name}")

    if config.enable_wandb:
        wandb.init(entity="sw-apps", project="TF2-BERT", name=universal_run_name, config=config, tags=config.wandb_tags)

    num_pipeline_stages = len(config.ipu_config.pipeline_device_mapping)
    num_ipus_per_replicas = max(config.ipu_config.pipeline_device_mapping) + 1

    # Load training and validation data
    # =================================
    dataset, eval_dataset, test_dataset, num_train_samples, num_eval_samples, num_test_samples, _ = get_glue_data(
        config.glue_task,
        config.global_batch.micro_batch_size,
        config.dataset_dir,
        config.max_seq_length,
        generated_dataset=config.generated_dataset,
    )

    total_num_train_samples = (
        config.num_epochs * num_train_samples.numpy()
        if config.total_num_train_samples is None
        else config.total_num_train_samples
    )
    train_config = BatchConfig(
        micro_batch_size=config.global_batch.micro_batch_size,
        total_num_train_samples=total_num_train_samples,
        num_replicas=config.global_batch.replicas,
        gradient_accumulation_count=config.global_batch.grad_acc_steps_per_replica,
        num_pipeline_stages=num_pipeline_stages,
        dataset_size=num_train_samples.numpy(),
        global_batches_per_log=config.global_batches_per_log,
        task=Task.OTHER,
    )

    policy = tf.keras.mixed_precision.Policy("float16")
    tf.keras.mixed_precision.set_global_policy(policy)

    strategy = create_ipu_strategy(
        num_ipus_per_replica=num_ipus_per_replicas,
        num_replicas=config.global_batch.replicas,
        enable_recomputation=config.enable_recomputation,
        enable_stochastic_rounding=config.enable_stochastic_rounding,
        compile_only=config.compile_only,
    )
    set_random_seeds(config.seed)
    with strategy.scope():
        if config.bert_model_name is not None:
            # Instantiate the pretrained model given in the config.
            model = TFBertForSequenceClassification.from_pretrained(config.bert_model_name)
        else:
            bert_config = BertConfig(**config.bert_config.dict(), hidden_act=ipu.nn_ops.gelu)
            model = TFBertForSequenceClassification(config=bert_config)
        # Convert to functional model exposing the encoder stages.
        model = convert_tf_bert_model(
            model,
            dataset,
            post_process_bert_input_layer,
            replace_layers=config.replace_layers,
            use_outlining=config.use_outlining,
            enable_recomputation=config.enable_recomputation,
            embedding_serialization_factor=config.embedding_serialization_factor,
            use_prediction_bias=config.use_prediction_bias,
            use_projection_bias=config.use_projection_bias,
            use_cls_layer=config.use_cls_layer,
            use_qkv_bias=config.use_qkv_bias,
            use_qkv_split=config.use_qkv_split,
            rename_outputs={"classifier": "labels"},
        )

        # Load from pretrained checkpoint if requested.
        if config.pretrained_ckpt_path:
            logging.info("Attempting to load pretrained checkpoint from path " f"{config.pretrained_ckpt_path}")
            _ = load_checkpoint_into_model(
                model=model, pretrained_ckpt_path=config.pretrained_ckpt_path, expect_partial=not config.do_training
            )

        # Configure pipeline stages
        # =========================
        if num_pipeline_stages > 1:
            pipeline_assigner = GluePipelineStagesAssigner(PIPELINE_ALLOCATE_PREVIOUS, PIPELINE_NAMES)
            assignments = model.get_pipeline_stage_assignment()
            assignments = pipeline_assigner.assign_glue_pipeline_stages(assignments, config.ipu_config.pipeline_stages)
            model.set_pipeline_stage_assignment(assignments)
            model.print_pipeline_stage_assignment_summary(print_fn=logging.info)
            poplar_options_per_pipeline_stage = get_poplar_options_per_pipeline_stage(
                num_ipus_per_replicas,
                config.ipu_config.pipeline_device_mapping,
                config.ipu_config.matmul_available_memory_proportion_per_pipeline_stage,
                config.matmul_partials_type,
            )
            model.set_pipelining_options(
                gradient_accumulation_steps_per_replica=config.global_batch.grad_acc_steps_per_replica,
                gradient_accumulation_reduction_method=GradientAccumulationReductionMethod.RUNNING_MEAN,
                pipeline_schedule=ipu.ops.pipelining_ops.PipelineSchedule.Grouped,
                device_mapping=config.ipu_config.pipeline_device_mapping,
                offload_weight_update_variables=config.optimizer_state_offchip,
                forward_propagation_stages_poplar_options=poplar_options_per_pipeline_stage,
                backward_propagation_stages_poplar_options=poplar_options_per_pipeline_stage,
                recomputation_mode=ipu.pipelining_ops.RecomputationMode.RecomputeAndBackpropagateInterleaved,
            )

        # Define optimiser with polynomial decay learning rate.
        queues_to_outfeed = []
        lr_outfeed_queue = ipu.ipu_outfeed_queue.IPUOutfeedQueue(outfeed_mode=ipu.ipu_outfeed_queue.IPUOutfeedMode.LAST)
        lr_scheduler = get_lr_scheduler(
            scheduler_name=config.optimizer_opts.learning_rate.schedule_name,
            max_learning_rate=config.optimizer_opts.learning_rate.max_learning_rate,
            warmup_frac=config.optimizer_opts.learning_rate.warmup_frac,
            num_train_steps=train_config.num_train_steps,
            queue=lr_outfeed_queue,
        )
        queues_to_outfeed.append(lr_outfeed_queue)

        loss_function = (
            ClassificationLossFunctionRegression if config.glue_task == "stsb" else ClassificationLossFunction
        )
        glue_loss = wrap_loss_in_enqueuer(loss_function, ["glue_loss"])()
        queues_to_outfeed.append(glue_loss.outfeed_queue)

        metrics = []
        if config.glue_task != "stsb":
            # Note that the stsb task is a regression and the accuracy
            # is not a valid metric
            accuracy = wrap_stateless_metric_in_enqueuer("glue_accuracy", classification_accuracy_fn, ["glue_accuracy"])
            metrics.append(accuracy)
            queues_to_outfeed.append(accuracy.outfeed_queue)

        if config.do_training:
            # Prepare optimizer.
            optimizer = get_optimizer(
                optimizer_name=config.optimizer_opts.name,
                num_replicas=config.global_batch.replicas,
                learning_rate_schedule=lr_scheduler,
                use_outlining=config.use_outlining,
                weight_decay_rate=config.optimizer_opts.weight_decay_rate,
            )
            # Compile and train the functional model
            model.compile(
                optimizer=optimizer,
                loss={"labels": glue_loss},
                metrics=metrics,
                steps_per_execution=train_config.steps_per_execution,
            )
            # Set up callbacks
            callbacks = CallbackFactory.get_callbacks(
                universal_run_name=universal_run_name,
                batch_config=train_config,
                model=model,
                checkpoint_path=config.save_ckpt_path,
                ckpt_every_n_steps_per_execution=config.ckpt_every_n_steps_per_execution,
                outfeed_queues=queues_to_outfeed,
                enable_wandb=config.enable_wandb,
            )
            # Print configs to be logged in wandb's terminal.
            logging.info(f"Application config:\n{config}")
            logging.info(f"Training batch config:\n{train_config}")

            model.fit(
                dataset,
                steps_per_epoch=(
                    train_config.total_num_micro_batches
                    if train_config.epochs < 1
                    else train_config.num_micro_batches_per_epoch
                ),
                batch_size=config.global_batch.micro_batch_size,
                epochs=math.ceil(train_config.epochs),
                callbacks=callbacks,
                verbose=0,
            )
        if config.do_validation:
            eval_batch_config = BatchConfig(
                micro_batch_size=config.global_batch.micro_batch_size,
                total_num_train_samples=num_eval_samples.numpy(),
                num_replicas=1,
                gradient_accumulation_count=10,
                dataset_size=num_eval_samples.numpy(),
                task=Task.OTHER,
            )
            # Re-compile the model for prediction if needed.
            eval_dataset = eval_dataset.take(eval_batch_config.num_micro_batches_per_epoch)

            if train_config.steps_per_execution != eval_batch_config.steps_per_execution or not config.do_training:
                model.compile(metrics=[accuracy], steps_per_execution=eval_batch_config.steps_per_execution)
            # Get predictions for the validation data.
            logging.info(f"Running inference:\nGenerating predictions on the validation data...")
            validation_results = model.evaluate(eval_dataset, verbose=0)
            logging.info(f"Validation performance for {config.glue_task}:\naccuracy = {validation_results[1]:.3f}")

        if config.do_test:
            test_batch_config = BatchConfig(
                micro_batch_size=config.global_batch.micro_batch_size,
                total_num_train_samples=num_test_samples.numpy(),
                num_replicas=1,
                gradient_accumulation_count=10,
                dataset_size=num_test_samples.numpy(),
                task=Task.OTHER,
            )
            # Re-compile the model for prediction if needed.
            test_dataset = test_dataset.take(test_batch_config.num_micro_batches_per_epoch)
            if train_config.steps_per_execution != test_batch_config.steps_per_execution or not config.do_training:
                model.compile(steps_per_execution=test_batch_config.steps_per_execution)
            # Get predictions for the test data.
            logging.info(f"Running inference:\nGenerating predictions on the test data...")
            model.predict(test_dataset, verbose=0)


if __name__ == "__main__":
    # Combine arguments and config file
    parser = argparse.ArgumentParser(
        description="TF2 BERT GLUE Fine Tuning", formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser = add_shared_arguments(parser)
    parser = add_glue_arguments(parser)
    args = parser.parse_args()
    config = combine_config_file_with_args(args, GLUEOptions)

    # Setup logging
    logging.basicConfig(
        format="%(asctime)s %(levelname)-8s %(message)s", level=config.logging, datefmt="%Y-%m-%d %H:%M:%S"
    )
    # Prevent doubling of TF logs.
    tf.get_logger().propagate = False

    # Run GLUE fine tuning
    fine_tune_glue(config)