# Copyright (c) 2021 Graphcore Ltd. All rights reserved. # Copyright 2020 The HuggingFace Team All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # This file has been modified by Graphcore Ltd. import argparse import logging import math from datetime import datetime from datasets import load_metric import tensorflow as tf from tensorflow.python import ipu from tensorflow.python.ipu.gradient_accumulation import GradientAccumulationReductionMethod from transformers import BertConfig, TFBertForQuestionAnswering import wandb from data_utils.batch_config import BatchConfig, Task from data_utils.squad_v1.load_squad_data import format_raw_data_for_metric, get_prediction_dataset, get_squad_data from data_utils.squad_v1.postprocess_squad_predictions import postprocess_qa_predictions 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.losses import QuestionAnsweringLossFunction from model.convert_bert_model import convert_tf_bert_model, post_process_bert_input_layer from model.pipeline_stage_names import PIPELINE_ALLOCATE_PREVIOUS, PIPELINE_NAMES from utilities.argparser import add_shared_arguments, add_squad_arguments, combine_config_file_with_args from utilities.assign_pipeline_stages import PipelineStagesAssigner 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_stateful_metric_in_enqueuer from utilities.options import SQuADOptions def fine_tune_squad(config): """Main function to run fine tuning BERT for SQuAD. :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( "SQuAD 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:" f" `bert_model_name`={config.bert_model_name}, and" f" `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"-{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_replica = max(config.ipu_config.pipeline_device_mapping) + 1 # Load training and validation data # ================================= train_dataset, eval_dataset, num_train_samples, num_eval_samples, raw_datasets = get_squad_data( config.global_batch.micro_batch_size, config.dataset_dir, config.generated_dataset, config.max_seq_length ) 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_batch_config = BatchConfig( micro_batch_size=config.global_batch.micro_batch_size, num_replicas=config.global_batch.replicas, total_num_train_samples=total_num_train_samples, 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, ) # Create model # ============ 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_replica, num_replicas=config.global_batch.replicas, enable_stochastic_rounding=config.enable_stochastic_rounding, enable_recomputation=config.enable_recomputation, compile_only=config.compile_only, ) set_random_seeds(config.seed) with strategy.scope(): # Instantiate the pretrained model given in the config. if config.bert_model_name is not None: model = TFBertForQuestionAnswering.from_pretrained(config.bert_model_name) else: bert_config = BertConfig(**config.bert_config.dict(), hidden_act=ipu.nn_ops.gelu) model = TFBertForQuestionAnswering(config=bert_config) # Convert subclass model to functional, expand main layers to enable pipelining, and replace some layers to # optimise performance. model = convert_tf_bert_model( model, train_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, rename_outputs={"tf.compat.v1.squeeze": "start_positions", "tf.compat.v1.squeeze_1": "end_positions"}, 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, ) # 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=config.do_training ) # Configure pipeline stages # ========================= if num_pipeline_stages > 1: pipeline_assigner = PipelineStagesAssigner(PIPELINE_ALLOCATE_PREVIOUS, PIPELINE_NAMES) assignments = model.get_pipeline_stage_assignment() assignments = pipeline_assigner.assign_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_replica, 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, ) if config.do_training: # Compile the model for training # ============================== # Wrap loss in an out-feed queue. loss = wrap_loss_in_enqueuer( QuestionAnsweringLossFunction, ["end_positions_loss", "start_positions_loss"] )() # Wrap accuracy in an out-feed queue. accuracy = wrap_stateful_metric_in_enqueuer( tf.keras.metrics.SparseCategoricalAccuracy, ["end_positions_accuracy", "start_positions_accuracy"] )() # Define optimiser with polynomial decay learning rate. 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_batch_config.num_train_steps, queue=lr_outfeed_queue, ) # 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 the model. model.compile( optimizer=optimizer, loss={"end_positions": loss, "start_positions": loss}, metrics={"end_positions": accuracy, "start_positions": accuracy}, steps_per_execution=train_batch_config.steps_per_execution, ) # Train the model # =============== # Set up callbacks callbacks = CallbackFactory.get_callbacks( universal_run_name=universal_run_name, batch_config=train_batch_config, model=model, checkpoint_path=config.save_ckpt_path, ckpt_every_n_steps_per_execution=config.ckpt_every_n_steps_per_execution, outfeed_queues=[lr_outfeed_queue, loss.outfeed_queue, accuracy.outfeed_queue], 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_batch_config}") # Train the model # Set verbose to 0 so the default progress bar, which is unreliable # with `steps_per_execution > 1`, is hidden in favour of using a # logging callback included in callbacks dir model.fit( train_dataset, steps_per_epoch=( train_batch_config.total_num_micro_batches if train_batch_config.epochs < 1 else train_batch_config.num_micro_batches_per_epoch ), epochs=math.ceil(train_batch_config.epochs), callbacks=callbacks, verbose=0, ) if config.do_validation: # Evaluate the model on the validation set # ======================================== # Prepare the dataset to be evaluated in the IPU. eval_batch_config = BatchConfig( micro_batch_size=config.global_batch.micro_batch_size, total_num_train_samples=num_eval_samples.numpy(), gradient_accumulation_count=config.global_batch.grad_acc_steps_per_replica, dataset_size=num_eval_samples.numpy(), task=Task.OTHER, ) eval_pred_dataset = get_prediction_dataset(eval_dataset, eval_batch_config.num_micro_batches_per_epoch) with strategy.scope(): # Re-compile the model for prediction if needed. if ( train_batch_config.steps_per_execution != eval_batch_config.steps_per_execution or not config.do_training ): model.compile(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...") # Set verbose to 0 so the default progress bar is hidden in # favour of using a logging callback included in callbacks dir predictions = model.predict( eval_pred_dataset, batch_size=eval_batch_config.micro_batch_size, verbose=0, ) # The predictions for the end position is first in the model outputs tuple (note the output of model.summary()). end_predictions, start_predictions = predictions # Match the predictions to answers in the original context. # This will also write out the predictions to a json file in the directory given by `output_dir`. final_predictions = postprocess_qa_predictions( list(raw_datasets["validation"].as_numpy_iterator()), list(eval_dataset.take(eval_batch_config.num_micro_batches_per_epoch).unbatch().as_numpy_iterator()), (start_predictions, end_predictions), output_dir=config.output_dir, ) # Format the predictions and the actual labels as expected by the metric. formatted_predictions = [{"id": k, "prediction_text": v} for k, v in final_predictions.items()] formatted_labels = format_raw_data_for_metric(raw_datasets["validation"]) metric = load_metric("squad") metrics = metric.compute(predictions=formatted_predictions, references=formatted_labels) logging.info("Evaluation metrics:") for key, value in metrics.items(): logging.info(f"{key}: {value:.3f}") if __name__ == "__main__": # Combine arguments and config file parser = argparse.ArgumentParser( description="TF2 BERT SQuAD Fine Tuning", formatter_class=argparse.ArgumentDefaultsHelpFormatter ) parser = add_shared_arguments(parser) parser = add_squad_arguments(parser) args = parser.parse_args() config = combine_config_file_with_args(args, SQuADOptions) # 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 SQuAD fine tuning fine_tune_squad(config)