horovod_train

import glob
import horovod.tensorflow as hvd  # isort:skip
from tqdm import tqdm
#import tensorflow_addons as tfa

from src.model_mod.opt import LAMB, AdamWeightDecay, WarmUp
from src.model_mod.lookahead_opt import Lookahead
from src.prep_data_mod.create_pretraining_data import GeneratorBERT
from src.prep_data_mod.build_dataset import Dataset
from src.confs.base_conf import logger_obj, fdir_obj, input_obj
from src.model_mod.transformer_lm_models.transformer_model_lm_pre import TransfomerMaskedLM
from src.confs.nn_config import NNConfig
from src.confs.training_config import TrainConfig
from src.model_mod.ga import GradientAccumulator
from src.utils_mod.train_utils import read_meta_data_csv, fill_meta_data_csv, write_meta_data_csv



def mlm_loss_fn(logits, label_ids, label_weights, voc_size):
    """ function to compute loss and acc based on logits results
        Args:
            logits
            labels_ids
            label_weights
            voc_size
        Returns:
            lm_loss
            acc
    """
    logits = tf.nn.log_softmax(logits, axis=-1) # get log_
    lm_label_ids_one_hot = tf.keras.backend.one_hot(label_ids, voc_size)
    #lm_label_weights = tf.keras.backend.reshape(label_weights, [-1])
    lm_per_example_loss = -tf.keras.backend.sum(logits * lm_label_ids_one_hot, axis=[-1])
    loss_numerator = tf.keras.backend.sum(label_weights * lm_per_example_loss)
    denominator = tf.keras.backend.sum(label_weights) + 1e-5
    preds = tf.math.argmax(logits, -1)
    preds = tf.cast(preds, dtype=tf.int32)
    acc_per_token = label_weights * tf.cast(tf.math.equal(label_ids, preds), tf.float32)  # [b, num_tokens]
    acc_numerator = tf.reduce_sum(acc_per_token)
    lm_loss = loss_numerator / denominator
    acc = acc_numerator / denominator
    return lm_loss, acc

class PreTrainAlbertLM():
    """
    model_name,
    albert_config,
    train_config,
    load_model=False,
    load_weights=False,
    checkpoint_path=None
    """
    def __init__(self, model_name, albert_config, train_config, load_optimizer=False):
        self.config = train_config  # training config
        self.model_name = model_name    # model_name
        self.load_optimizer = load_optimizer    # load optimizer or not
        self.voc_size = albert_config.voc_size  # vocabulary size
        # path to model directory to hold output results
        self.model_dir = fdir_obj.path2maindir + '/models/{}/'.format(model_name)
        # path 2 store checkpoints of model
        self.path2model_ckpt = self.model_dir + '/ckpt/'
        self.logger = tf.get_logger()   # logger_object
        assert self.config.optimizer in ["adamw", "lamb"], \
            logger_obj.logger.error('The optimizer type is not supported: %s' % self.config.optimizer)

        # value for gradient clipping
        self.max_grad_norm = self.config.max_grad_norm
        # max steps to train
        self.total_train_steps = self.config.total_train_steps
        # number of warmup steps for linear loss
        self.num_warmup_steps = self.config.num_warmup_steps
        # number of accumulate gradients step
        self.gradient_accumulation_steps = self.config.gradient_accumulation_steps
        # compute global batches size
        self.global_train_batch_size = self.config.train_batch_size * hvd.size() * self.gradient_accumulation_steps
        self.global_test_batch_size = self.config.test_batch_size * hvd.size() * self.gradient_accumulation_steps
        self.voc_size = albert_config.voc_size  # size of vocab with tokens
        self.max_seq_length = input_obj.max_num_phrase  # maximum length
        self.path2metadata = self.model_dir + '/info/'   # directory to output files with information of train
        # fils to hold result info about training proc
        self.meta_test_sum = self.path2metadata + 'test_summary.csv'
        self.meta_test_av = self.path2metadata + 'test_average.csv'
        self.meta_train_sum = self.path2metadata + 'train_summary.csv'
        self.meta_train_av = self.path2metadata + 'train_average.csv'
        self.meta_csv = self.path2metadata + "meta.csv"
        # number steps for one epoch
        self.train_steps_per_epoch = math.ceil(self.total_train_steps / self.config.num_epochs)
        #self.steps_for_save = math.ceil(self.train_steps_per_epoch / self.freq_save)
        self.path2tb = self.path2metadata + '/tb/'  # path to tensorboadd directory
        # arr drrs to hold results
        arr_dirs = [self.model_dir, self.path2model_ckpt, self.path2metadata, self.path2tb]
        # create dir trees if absent
        for dir in arr_dirs:
            if not Path(dir).is_dir():
                Path(dir).mkdir(parents=True, exist_ok=True)
        # initialize tensorboard
        self.summary_writer = tf.summary.create_file_writer(self.path2tb)
        # initaalize csv files to hold result
        fill_meta_data_csv(meta_train_sum=self.meta_train_sum, meta_test_sum=self.meta_test_sum,
                           meta_test_av=self.meta_test_av, meta_train_av=self.meta_train_av)
        self.loaded_optimizer_weights = None    # object to hold optimizer weights
        self.model = TransfomerMaskedLM(config=albert_config)  # our model
        # initialize learning rate schedule
        learning_rate_fn = tf.keras.optimizers.schedules.PolynomialDecay(
                initial_learning_rate=self.config.initial_lr * hvd.size(),
                decay_steps=int(self.total_train_steps),
                end_learning_rate=0.0)
        learning_rate_fn = WarmUp(initial_learning_rate=self.config.initial_lr * hvd.size(),
                                  decay_schedule_fn=learning_rate_fn,
                                  warmup_steps=self.num_warmup_steps)
        # create optimizer
        if self.config.optimizer == "lamb":
            optimizer_fn = LAMB
        else:
            optimizer_fn = AdamWeightDecay
        self.optimizer = optimizer_fn(learning_rate=learning_rate_fn, weight_decay_rate=self.config.weight_decay,
                                      beta_1=self.config.opt_beta1, beta_2=self.config.opt_beta2,
                                      epsilon=self.config.opt_epsilon,
                                      exclude_from_weight_decay=["layer_norm", "bias"],
                                      experimental_aggregate_gradients=False)
        #if self.config.use_lookahead_opt:
        self.optimizer = Lookahead(self.optimizer, sync_period=6, slow_step_size=0.5,
                                   experimental_aggregate_gradients=False)

        self.optimizer = tf.train.experimental.enable_mixed_precision_graph_rewrite(self.optimizer, loss_scale="dynamic")
        self.optimizer._HAS_AGGREGATE_GRAD = False
        # create gradient accumulator object
        self.gradient_accumulator = GradientAccumulator()
        self.ckpt = tf.train.Checkpoint(model=self.model, optimizer=self.optimizer)
        self.ckpt_manager = tf.train.CheckpointManager(self.ckpt, directory=self.path2model_ckpt,
                                                             max_to_keep=100)
        self.first_batch = True
        # print out info about training parameters
        logger_obj.logger.info('Num epochs:{}\nTrain batch size:{}\n'
                               'Global train batch size:{}\n'
                               'Train steps per epoch:{}\n'
                               'Num train step:{}\n Num warmup steps:{}\n'
                               'Optimizer:{}\n'.format(self.config.num_epochs,
                                                       self.config.train_batch_size,
                                                       self.global_train_batch_size,
                                                       self.train_steps_per_epoch,
                                                       self.total_train_steps,
                                                       self.num_warmup_steps,
                                                       self.config.optimizer))

    def allreduce(self, mlm_loss, mlm_acc):
        """
        function to compile gradients and using horovod functional to output res
        Args:
            mlm_loss: loss per gpu
            mlm_acc: accuracy per gpu
        Returns:
            mlm_loss: tf.float: val loss
            mlm_acc: tf.float: val accuracy
            grad_norm: tf.float: norm values of gradient
            weight_norm: tf.float: norm values of weights
        """
        scaled_grads = self.gradient_accumulator.gradients
        grads = self.optimizer.get_unscaled_gradients(scaled_grads)
        # This, which is equivalent to sparse_as_dense=True, gives a mild 2% speedup
        grads = [
            tf.convert_to_tensor(grad)
            if grad is not None and isinstance(grad, tf.IndexedSlices)
            else grad
            for grad in grads
        ]
        # !todo: Does placing this clip before or after allreduce affect accuracy?
        # Placing before has a regularization effect, no single example can contribute as much.
        # Placing before also gives a 20% speedup when training BERT-large, probably because the
        # gradient operations can be fused by XLA.
        if self.config.use_clip_grad:
            (grads, grad_norm) = tf.clip_by_global_norm(grads, clip_norm=self.max_grad_norm)
        else:
            grad_norm = grads
        weight_norm = tf.math.sqrt(tf.math.reduce_sum([tf.norm(var, ord=2) ** 2 for var in self.model.trainable_variables]))
        grads = [
            hvd.allreduce(grad, compression=hvd.Compression.fp16) if grad is not None else None
            for grad in grads
        ]
        self.optimizer.apply_gradients(
            [
                (tf.cast(grad, var.dtype), var)
                for (grad, var) in zip(grads, self.model.trainable_variables)
                if grad is not None
            ], experimental_aggregate_gradients=False
        )
        # Clear the gradient accumulator
        self.gradient_accumulator.reset()
        mlm_loss = hvd.allreduce(mlm_loss)
        mlm_acc = hvd.allreduce(mlm_acc)
        return mlm_loss, mlm_acc, grad_norm, weight_norm

    def train_grad(self, input_ids, masked_lm_pos, label_ids, label_weights):
        """
        Args:
            input_ids: masked tokens ids
            masked_lm_pos: positions of masked tokens
            label_ids: ids of masked tokens
            label_weights: masked weights (1.0 if not masked, 0 if masked)
        Returns:
            mlm_loss: loss
            mlm_acc: acc
        """
        with tf.GradientTape() as tape:
            mlm_logits = self.model(input_ids, masked_lm_pos, training=True)
            mlm_loss, mlm_acc = mlm_loss_fn(logits=mlm_logits,
                                            label_ids=label_ids, label_weights=label_weights, voc_size=self.voc_size)
            mlm_loss = tf.cast(mlm_loss, dtype=tf.float32)
            scaled_loss = self.optimizer.get_scaled_loss(mlm_loss)
        # !todo: On iteration 0, loss=11 and loss_scale()=32768, so scaled_loss=inf.
        # But scaled_grads is not inf, how? tape.gradient() must not be using direct backprop calc
        #tape = hvd.DistributedGradientTape(tape)
        scaled_grads = tape.gradient(scaled_loss, self.model.trainable_variables)
        self.gradient_accumulator(scaled_grads)
        #self.optimizer.apply_gradients(zip(scaled_grads, self.model.trainable_variables))
        if self.first_batch:
            hvd.broadcast_variables(self.model.variables, root_rank=0)
            hvd.broadcast_variables(self.optimizer.variables(), root_rank=0)
            self.first_batch = False
        return mlm_loss, mlm_acc

    def train_step(self, inputs):
        """
        functional of train step using input da of batches of batches
        through horovod system
        Args:
            inputs: [grad_step, batch, seq_length] : input values of batch
        Returns:
            mlm_loss
            mlm_acc
            grad_norm
            weight_norm
        """
        # initialize values to hold results throug batches of batches
        total_mlm_loss, total_mlm_acc = (tf.constant(0, dtype=tf.float32), tf.constant(0, dtype=tf.float32),)
        # loop of batches of batches
        print(inputs)
        for step in range(self.gradient_accumulation_steps):
            input_ids, masked_lm_pos, label_ids, label_weights = inputs
            input_ids = input_ids[step]
            masked_lm_pos = masked_lm_pos[step]
            label_ids = label_ids[step]
            label_weights = label_weights[step]
            mlm_loss, mlm_acc = self.train_grad(input_ids, masked_lm_pos, label_ids, label_weights)
            total_mlm_loss += tf.cast(mlm_loss, total_mlm_loss.dtype)
            total_mlm_acc += tf.cast(mlm_acc, total_mlm_acc.dtype)
        total_mlm_loss /= self.gradient_accumulation_steps
        total_mlm_acc /= self.gradient_accumulation_steps
        return total_mlm_loss, total_mlm_acc #self.allreduce(mlm_loss=total_mlm_loss, mlm_acc=total_mlm_acc, )

    def test_step(self, inputs):
        """validation step
        Args:
            inputs:
        Returns:
            loss, acc
        """
        input_word_ids, masked_positions, mlm_ids, label_weights = inputs
        mlm_logits = self.model(input_word_ids=input_word_ids, masked_positions=masked_positions,
                                training=tf.convert_to_tensor(False))
        loss, acc = mlm_loss_fn(logits=mlm_logits,
                                label_ids=mlm_ids, label_weights=label_weights, voc_size=self.voc_size)
        return loss, acc

    def test_op(self, num_epoch):
        """validation procedure
        Args:
            num_epoch: number of epoch (current)
        Returns:
            val_mlm_loss
            val_mlm_acc
        """
        count = 0
        val_mlm_loss, val_mlm_acc= (0, 0)
        start_eval_time = time.time()
        for (test_batch, (inputs)) in enumerate(self.test_dataset):
            loss, acc = self.test_step(inputs)
            val_mlm_loss += loss
            val_mlm_acc += acc
            loss_ = loss.numpy()
            acc_ = acc.numpy()
            count += 1
            if test_batch % self.config.show_per_step == 0:
                logger_obj.logger.info(f'Test --> Epoch {num_epoch+1} | Batch {test_batch} | Loss {loss_:.4f} | Acc {acc_:4f}')
            df_test_sum = pd.DataFrame({'epoch_num': [num_epoch], 'loss': [loss], 'acc' : [acc]})
            df_test_sum.to_csv(self.meta_test_sum, sep='\t', index=False, columns=None, mode='a')
        #
        val_mlm_loss /= count
        val_mlm_acc /= count
        logger_obj.logger.info(f'Test --> Epoch {num_epoch+1} | Av Loss {val_mlm_loss:.4f} | Av Acc {val_mlm_acc:.4f}')
        df_test_av = pd.DataFrame({'epoch_num': [num_epoch], 'av_loss': [val_mlm_loss], 'av_acc' : val_mlm_acc})
        df_test_av.to_csv(self.meta_test_av, sep='\t', index=False, columns=None, mode='a')
        end_eval_epoch_time = time.time()
        logger_obj.logger.info(f'Time taken for test 1 epoch: {end_eval_epoch_time - start_eval_time} secs')
        return val_mlm_loss, val_mlm_acc

    def train_op(self):
        """
        train procedure
        """
        # initialize generator function to generate input data
        gen_train = GeneratorBERT(input_dir=fdir_obj.dir_train_data, random_seed=self.config.random_seed)
        # and format input data from arrat to Tensorflow format
        ds_train_obj = Dataset(generator=gen_train, batch_size=self.config.train_batch_size,
                               buffer_size=self.config.train_batch_size, is_train=True, len2=input_obj.max_num_bpe_mask,
                               )
        self.train_dataset = ds_train_obj.next_element
        if self.gradient_accumulation_steps > 1:
            self.train_dataset = self.train_dataset.batch(self.gradient_accumulation_steps)
            # One iteration with 10 dupes, 8 nodes seems to be 60-70k steps.
            self.train_dataset = self.train_dataset.prefetch(buffer_size=8)
        if input_obj.mode == "train_eval":
            #
            gen_test = GeneratorBERT(input_dir=fdir_obj.dir_dev_data, random_seed=self.config.random_seed)
            ds_test_obj = Dataset(generator=gen_test, batch_size=self.config.test_batch_size,
                                  buffer_size=self.config.test_batch_size, is_train=False,
                                  len2=input_obj.max_num_bpe_mask)
            self.test_dataset = ds_test_obj.next_element
            if self.gradient_accumulation_steps > 1:
                self.test_dataset = self.test_dataset.prefetch(buffer_size=8)

        start_train_time = time.time()
        # read info from checkpoint file
        if self.ckpt_manager.latest_checkpoint:
            # self.model.load_weights(self.path2model_ckpt)
            self.ckpt.restore(self.ckpt_manager.latest_checkpoint).expect_partial()
            logger_obj.logger.info("Restored from {}".format(self.ckpt_manager.latest_checkpoint))
        else:
            logger_obj.logger.info("Initializing from scratch.")
        # read meta info
        sum_acc_train = 0
        start_epoch, sum_loss_train, count_train = read_meta_data_csv(meta_csv=self.meta_csv)
        i = (self.train_steps_per_epoch * start_epoch) + 1
        self.logger.info("Start process of training...")
        # train operation
        for epoch in range(start_epoch, self.config.num_epochs):
            #pbar = tqdm(total=self.train_steps_per_epoch)
            self.logger.info(f"Start epoch {epoch}")
            for (train_batch, (inputs)) in enumerate(self.train_dataset):
                #pbar.set_description(f"Current batch #{train_batch+1}")
                self.logger.info(f"Current batch #{train_batch+1}")
                learning_rate = self.optimizer.learning_rate(step=tf.constant(i, dtype=tf.float32))
                loss_scale = self.optimizer.loss_scale()
                #mlm_loss, mlm_acc, grad_norm, weight_norm = self.train_step(inputs)
                mlm_loss, mlm_acc = self.train_step(inputs)
                loss = mlm_loss.numpy()
                acc = mlm_acc.numpy()
                # define inputs of model to save later
                if train_batch % self.config.show_per_step == 0:
                    description = f"Epoch:{epoch+1} | Loss:{mlm_loss:.4f} | Accuracy:{mlm_acc:.4f}"
                    #pbar.set_description(description)
                    self.logger.info(description)
                #pbar.update(1)
                sum_loss_train += loss
                sum_acc_train += acc
                df_train_sum = pd.DataFrame({'epoch_num' : [epoch], 'loss': [loss], 'acc' : [acc]})
                df_train_sum.to_csv(self.meta_train_sum, mode='a', sep='\t', index=False, columns=None)

                train_metrics = {
                    #"weight_norm": weight_norm,
                    #"grad_norm": grad_norm,
                    "loss_scale": loss_scale,
                    "learning_rate": learning_rate,
                    "train/mlm_loss": mlm_loss,
                    "train/mlm_acc": mlm_acc,
                }
                all_metrics = {**train_metrics}
                is_final_step = train_batch >= self.train_steps_per_epoch
                i += 1
                count_train += 1
                if is_final_step:
                    break
            av_loss_train = sum_loss_train / count_train
            av_acc_train = sum_acc_train / count_train
            logger_obj.logger.info(f'Train --> Epoch {epoch+1} | Av Loss {av_loss_train:.4f} | Av Acc {av_acc_train:.4f}')
            write_meta_data_csv(epoch=epoch, sum_loss_train=sum_loss_train, count_train=count_train,
                                meta_csv=self.meta_csv)
            end_train_epoch_time = time.time()
            logger_obj.logger.info(f'Time taken for train 1 epoch: {end_train_epoch_time - start_train_time} secs\n')
            ckpt_save_path = self.ckpt_manager.save()
            logger_obj.logger.info(f'Saving checkpoint for epoch {epoch} to {ckpt_save_path}\n')
            if input_obj.mode == "train_eval":
                val_mlm_loss, val_mlm_acc = self.test_op(num_epoch=epoch)
                val_metrics = {"val/mlm_loss": val_mlm_loss,
                               "val/mlm_acc": val_mlm_acc,
                               }
                all_metrics = {**all_metrics, **val_metrics}
            with self.summary_writer.as_default():
                for name, val in all_metrics.items():
                    tf.summary.scalar(name, val, step=epoch)
            sum_loss_train = 0.0
            sum_acc_train = 0.0
            count_train = 0
        self.logger.info('End of process of training!')

def main_train():
    # Users should always run this script under TF 2.x
    assert tf.version.VERSION.startswith('2.')
    tf.random.set_seed(12345)
    tf.autograph.set_verbosity(0)
    #
    #
    model_name = 'albert_pretrained'
    path2cfgs = fdir_obj.path2maindir + '/models/{}/configs/'.format(model_name)
    if not Path(path2cfgs).is_dir():
        Path(path2cfgs).mkdir(parents=True, exist_ok=True)
    if len(os.listdir(path2cfgs)) == 0:
        for filename in glob.glob(os.path.join(fdir_obj.path2cfgs, '*.*')):
            shutil.copy(filename, path2cfgs)
    nn_gin_cfg = path2cfgs + 'nn.gin'
    train_gin_cfg = path2cfgs + 'train.gin'
    #
    gin.parse_config_file(nn_gin_cfg)
    gin.parse_config_file(train_gin_cfg)

    albert_config = NNConfig()
    train_config = TrainConfig()
    #
    # XLA, AutoGraph
    #tf.config.optimizer.set_jit(True)
    #tf.config.experimental_run_functions_eagerly(False)
    # Horovod init
    hvd.init()
    gpus = tf.config.list_physical_devices("GPU")
    print(gpus)
    # set mem growth for nvidia cuda
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)
    if gpus:
        tf.config.set_visible_devices(gpus[hvd.local_rank()], "GPU")
    trainer_obj = PreTrainAlbertLM(model_name=model_name, albert_config=albert_config, train_config=train_config,
                                   load_optimizer=True)
    trainer_obj.train_op()
    #trainer_obj.eval_op()


if __name__ == "__main__":
    #set_memory_growth()
    #app.run(main_train)
    main_train()