# Copyright (c) 2021 Graphcore Ltd. All rights reserved. import argparse from datetime import datetime import logging import popdist.tensorflow import tensorflow as tf from tensorflow.python import ipu from tensorflow.python.ipu import distributed from tensorflow.python.ipu.gradient_accumulation import GradientAccumulationReductionMethod from transformers import BertConfig import wandb from data_utils.batch_config import BatchConfig, get_pretraining_task, Task from data_utils.wikipedia.load_wikipedia_data import get_pretraining_dataset 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.accuracy import pretraining_accuracy_fn from model.convert_bert_model import convert_tf_bert_model, post_process_bert_input_layer from model.ipu_pretraining_model import IpuTFBertForPreTraining from model.losses import MLMLossFunction, NSPLossFunction from model.pipeline_stage_names import PIPELINE_ALLOCATE_PREVIOUS, PIPELINE_NAMES from utilities.argparser import add_shared_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_stateless_metric_in_enqueuer from utilities.options import PretrainingOptions def pretrain(config): """Main function to run pretraining for BERT. :param config: A pydantic model object that contains the configuration to options for this application. See utilities/options.py for accepted options. """ distributed_training = popdist.isPopdistEnvSet() # Check if `poprun` has initiated distributed training. # Set a name for this run, and be sure it is shared by all hosts when using distributed training. if distributed_training: popdist.init() time_now = float( distributed.broadcast(tf.convert_to_tensor(value=datetime.now().timestamp(), dtype=tf.float32), 0) ) else: time_now = datetime.now().timestamp() universal_run_name = config.config.stem if config.name is None else config.name universal_run_name += "-" + datetime.fromtimestamp(time_now).strftime("%Y%m%d_%H%M%S") logging.info(f"Universal name for run: {universal_run_name}") if config.enable_wandb and popdist.getInstanceIndex() == 0: wandb.init( entity=config.wandb_entity_name, project=config.wandb_project_name, name=universal_run_name, config=config, tags=config.wandb_tags, ) # Get config parameters expected by HF model. bert_config = BertConfig(**config.bert_config.dict(), hidden_act=ipu.nn_ops.gelu) # Calculate the number of pipeline stages and the number of required IPUs per replica. num_pipeline_stages = len(config.ipu_config.pipeline_device_mapping) num_ipus_per_replica = max(config.ipu_config.pipeline_device_mapping) + 1 # Load training data # ================== dataset, num_samples = get_pretraining_dataset( micro_batch_size=config.global_batch.micro_batch_size, dataset_dir=config.dataset_dir, max_seq_length=config.max_seq_length, max_predictions_per_seq=config.max_predictions_per_seq, vocab_size=bert_config.vocab_size, seed=config.seed, data_type=tf.float16, distributed_worker_count=popdist.getNumInstances(), distributed_worker_index=popdist.getInstanceIndex(), generated_dataset=config.generated_dataset, ) task = get_pretraining_task(config.max_seq_length) batch_config = BatchConfig( micro_batch_size=config.global_batch.micro_batch_size, 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_samples, global_batches_per_log=config.global_batches_per_log, total_num_train_samples=config.total_num_train_samples, task=task, ) # Create training strategy # ============ strategy = create_ipu_strategy( num_ipus_per_replica=num_ipus_per_replica, num_replicas=config.global_batch.replicas, distributed_training=distributed_training, fp_exceptions=config.fp_exceptions, enable_recomputation=config.enable_recomputation, min_remote_tensor_size=config.min_remote_tensor_size, compile_only=config.compile_only, ) set_random_seeds(config.seed) # Create model # ============ policy = tf.keras.mixed_precision.Policy("float16") tf.keras.mixed_precision.set_global_policy(policy) with strategy.scope(): model = IpuTFBertForPreTraining(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, 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, ) # Load from pretrained checkpoint if requested. ckpt_found = False if config.pretrained_ckpt_path: logging.info("Attempting to load pretrained checkpoint from path " f"{config.pretrained_ckpt_path}") ckpt_found = load_checkpoint_into_model(model, config.pretrained_ckpt_path) if task == Task.PRETRAIN_PHASE_TWO and not ckpt_found: logging.warning( "WARNING: Phase 2 pre-training should be done from a completed Phase 1 checkpoint. " "Please specify the path to the Phase 1 checkpoint with 'pretrained_ckpt_path' in the " "config or as a command line argument." ) # 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=batch_config.gradient_accumulation_count, 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, replicated_optimizer_state_sharding=config.ipu_config.replicated_tensor_sharding, 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, ) # Compile the model for training # ============================== # Wrap losses in an out-feed queue. nsp_loss = wrap_loss_in_enqueuer(NSPLossFunction, ["nsp_loss"])() mlm_loss = wrap_loss_in_enqueuer(MLMLossFunction, ["mlm_loss"])() # Prepare learning rate and wrap it in an out-feed queue. 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=batch_config.num_train_steps, queue=lr_outfeed_queue, ) # Construct outfeed queues outfeed_queues = [lr_outfeed_queue, nsp_loss.outfeed_queue, mlm_loss.outfeed_queue] # Prepare accuracy and wrap it in an out-feed queue. if config.show_accuracy: pretraining_accuracy = wrap_stateless_metric_in_enqueuer( "accuracy", pretraining_accuracy_fn, ["mlm_acc_average", "nsp_acc_average"] ) outfeed_queues.append(pretraining_accuracy.outfeed_queue) # Prepare optimizer. outline_optimizer_apply_gradients = ( False if config.ipu_config.replicated_tensor_sharding else config.use_outlining ) optimizer = get_optimizer( optimizer_name=config.optimizer_opts.name, num_replicas=popdist.getNumTotalReplicas() if distributed_training else config.global_batch.replicas, learning_rate_schedule=lr_scheduler, use_outlining=outline_optimizer_apply_gradients, loss_scaling=config.optimizer_opts.loss_scaling, weight_decay_rate=config.optimizer_opts.weight_decay_rate, ) # Compile the model. model.compile( optimizer=optimizer, loss={"nsp___cls": nsp_loss, "mlm___cls": mlm_loss}, steps_per_execution=batch_config.steps_per_execution, metrics=[pretraining_accuracy if config.show_accuracy else None], ) # Train the model # =============== # Set up callbacks callbacks = CallbackFactory.get_callbacks( universal_run_name=universal_run_name, batch_config=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=outfeed_queues, distributed_training=distributed_training, 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{batch_config}") # Train the model # In order to achieve a specific number of steps, we set the number of # epochs to 1 and the steps per epoch to the number of steps we require. logging.info( "Forcing `model.fit` to run a particular number of steps by" " running a single 'epoch' with the number of steps we" " require. This allows running a fraction of actual" " epochs." ) # 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. history = model.fit( dataset, steps_per_epoch=batch_config.total_num_micro_batches_per_instance, epochs=1, callbacks=callbacks, verbose=0, ) return history if __name__ == "__main__": # Combine arguments and config file parser = argparse.ArgumentParser( description="TF2 BERT Pretraining", formatter_class=argparse.ArgumentDefaultsHelpFormatter ) parser = add_shared_arguments(parser) args = parser.parse_args() cfg = combine_config_file_with_args(args, PretrainingOptions) # Setup logging logging.basicConfig( format="%(asctime)s %(levelname)-8s %(message)s", level=cfg.logging, datefmt="%Y-%m-%d %H:%M:%S" ) # Prevent doubling of TF logs. tf.get_logger().propagate = False # Run pretraining pretrain(cfg)