train.py

import os

os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import argparse
import gc
import logging
import sys
import time
from distutils import util
from typing import Any, Callable, Dict, Tuple

import deepspeed
import numpy as np
import pandas as pd
import torch
from torch.cuda.amp import GradScaler, autocast
from torch.utils.data import DataLoader
from tqdm import tqdm
from transformers.deepspeed import HfDeepSpeedConfig

from llm_studio.src.loggers import MainLogger
from llm_studio.src.utils.config_utils import (
    load_config_py,
    load_config_yaml,
    save_config_yaml,
)
from llm_studio.src.utils.data_utils import (
    get_data,
    get_inference_batch_size,
    get_train_dataloader,
    get_train_dataset,
    get_val_dataloader,
    get_val_dataset,
)
from llm_studio.src.utils.exceptions import LLMTrainingException
from llm_studio.src.utils.export_utils import save_prediction_outputs
from llm_studio.src.utils.gpu_utils import sync_across_processes
from llm_studio.src.utils.logging_utils import (
    TqdmToLogger,
    initialize_logging,
    log_plot,
    write_flag,
)
from llm_studio.src.utils.modeling_utils import (
    activate_neftune,
    check_disk_space,
    get_ds_config,
    get_number_of_validation_epochs,
    get_optimizer,
    get_scheduler,
    load_checkpoint,
    run_inference,
    save_checkpoint,
    save_predictions,
    wrap_model_distributed,
)
from llm_studio.src.utils.utils import kill_ddp_processes, set_environment, set_seed

logger = logging.getLogger(__name__)


def run_eval(
    cfg,
    model: torch.nn.Module,
    val_dataloader: DataLoader,
    val_df: pd.DataFrame,
    mode: str = "validation",
) -> Tuple:
    """Runs the evaluation loop.

    Args:
        cfg: config object
        model: trained model
        val_dataloader: validation Dataloader
        val_df: validation DataFrame
        mode: validation

    Returns:
        Validation loss
    """
    with torch.no_grad():
        is_training = model.training
        model.eval()
        val_data: Dict[str, Any] = run_inference(
            cfg, model, val_dataloader, mode
        )  # type: ignore
        model.train(is_training)

    # Sync validation predictions across GPUs
    if cfg.environment._distributed and cfg.environment._distributed_inference:
        for key, value in val_data.items():
            val_data[key] = sync_across_processes(
                value, cfg.environment._world_size, group=cfg.environment._cpu_comm
            )

    if cfg.environment._local_rank != 0:
        # data has been synced, so we can return early on other ranks
        if cfg.environment._distributed:
            torch.distributed.barrier()
        return 0, 0

    # Drop any extra observations
    for k, v in val_data.items():
        val_data[k] = v[: len(val_dataloader.dataset)]  # type: ignore

    val_data = val_dataloader.dataset.postprocess_output(  # type: ignore
        cfg=cfg, df=val_df, output=val_data
    )
    val_loss = np.mean(val_data.get("loss", torch.tensor(0)).float().cpu().numpy())
    # postprocess_output only runs on rank 0 to save time/memory
    val_metric = np.mean(val_data["metrics"])
    logger.info(f"{mode.capitalize()} {cfg.prediction.metric}: {val_metric:.5f}")

    for key in val_data:
        if key.startswith("additional_log_") or key == "loss":
            value = np.mean(val_data[key].float().cpu().numpy())
            key = key.replace("additional_log_", "")
            logger.info(f"Mean {mode} {key}: {value:.5f}")
            cfg.logging._logger.log(
                mode,
                key,
                value,
                step=cfg.environment._curr_step,
            )
    cfg.logging._logger.log(
        mode, cfg.prediction.metric, val_metric, step=cfg.environment._curr_step
    )

    # Log plots
    if val_df is not None:
        plot = cfg.logging.plots_class.plot_validation_predictions(
            val_outputs=val_data, cfg=cfg, val_df=val_df, mode="validation"
        )
        log_plot(cfg, plot, "validation_predictions")

    save_predictions(cfg, val_data, val_dataloader, val_df, mode)

    if cfg.environment._distributed:
        torch.distributed.barrier()

    return val_loss, val_metric


def run_train(
    cfg: Any,
    model: torch.nn.Module,
    optimizer,
    scheduler,
    epoch_steps,
    train_dataloader,
    val_dataloader,
    val_df: pd.DataFrame,
):
    """Runs the training loop.

    Args:
        cfg: config object
        model: model
        train_dataloader: custom training Dataloader
        train_df: train DataFrame
        val_dataloader: custom validation Dataloader
        val_df: validation DataFrame

    Returns:
        Validation prediction output
        Validation loss
        Validation metric
        Last train batch
    """
    if (
        hasattr(cfg.augmentation, "neftune_noise_alpha")
        and cfg.augmentation.neftune_noise_alpha > 0
    ):
        activate_neftune(model, cfg.augmentation.neftune_noise_alpha)

    scaler: GradScaler | None = None
    if cfg.environment.mixed_precision:
        scaler = GradScaler()

    optimizer.zero_grad(set_to_none=True)

    # Prepare NLP Augmentation
    nlp_augment = None
    if hasattr(cfg.augmentation, "nlp_augmentations_class"):
        nlp_augment = cfg.augmentation.nlp_augmentations_class(cfg=cfg)

    start_epoch = 0

    _, metric_mode, _ = cfg.prediction.metric_class.get(cfg.prediction.metric)
    objective_op: Callable[[float, float], bool]
    if metric_mode == "max":
        best_val_metric = -np.inf
        objective_op = np.greater
    else:
        best_val_metric = np.inf
        objective_op = np.less

    if cfg.training.evaluate_before_training:
        val_loss, val_metric = run_eval(
            cfg=cfg, model=model, val_dataloader=val_dataloader, val_df=val_df
        )

    for epoch in range(start_epoch, cfg.training.epochs):
        set_seed(
            cfg.environment._seed
            + epoch * cfg.environment._world_size * cfg.environment.number_of_workers
            + cfg.environment._local_rank * cfg.environment.number_of_workers
        )
        if cfg.environment._local_rank == 0:
            logger.info(f"Training Epoch: {epoch + 1} / {cfg.training.epochs}")

        if (
            cfg.environment._distributed
            and not cfg.environment.use_deepspeed
            and hasattr(train_dataloader.sampler, "set_epoch")
        ):
            train_dataloader.sampler.set_epoch(epoch)  # type: ignore

        tqdm_out = TqdmToLogger(logger, level=logging.INFO)
        progress_bar = tqdm(
            total=epoch_steps,
            disable=cfg.environment._local_rank != 0,
            file=tqdm_out,
            ascii=True,
            desc="train loss",
            mininterval=0,
        )
        tr_it = iter(train_dataloader)

        losses = []
        model.train()

        log_update_steps = max(epoch_steps // 20, 1)
        evaluation_step = max(int(epoch_steps * cfg.training.evaluation_epochs), 1)
        logger.info(f"Evaluation step: {evaluation_step}")

        for itr, data in enumerate(tr_it):
            cfg.environment._curr_step += (
                cfg.training.batch_size * cfg.environment._world_size
            )

            # Batch to device
            batch = cfg.dataset.dataset_class.batch_to_device(
                data, cfg.environment._device
            )

            # NLP augmentation
            if nlp_augment is not None:
                batch = nlp_augment(batch)

            # Plot first batch
            if epoch == 0 and itr == 0 and cfg.environment._local_rank == 0:
                plot = cfg.logging.plots_class.plot_batch(batch=batch, cfg=cfg)
                log_plot(cfg, plot, "train_data")

            # Forward pass
            with autocast(enabled=cfg.environment.mixed_precision):
                output_dict = model.forward(batch)

            loss = output_dict["loss"]
            if ~np.isfinite(loss.item()) and (epoch > start_epoch or itr > 20):
                raise LLMTrainingException(
                    "NaN caught in loss during training. "
                    "Please, reduce learning rate, change dtype, "
                    "or disable mixed precision. Alternatively, "
                    "gradient clipping may help to stabilize training."
                )
            losses.append(loss.item())

            # loss is a mean loss per batch/sample
            # as grad_accumulations sums up the gradients, this loss must be scaled
            # by the number of grad_accumulations, to have similar behavior for
            # BS * grad_accumulations = const.
            if cfg.training.grad_accumulation != 1:
                loss = loss / cfg.training.grad_accumulation

            # Backward pass
            if (
                cfg.environment.mixed_precision
                and len(cfg.environment.gpus)
                and not cfg.environment.use_deepspeed
            ):
                scaler.scale(loss).backward()  # type: ignore
                if itr % cfg.training.grad_accumulation == 0:
                    if cfg.training.gradient_clip > 0:
                        scaler.unscale_(optimizer)  # type: ignore
                        torch.nn.utils.clip_grad_norm_(
                            model.parameters(), cfg.training.gradient_clip
                        )
                    scaler.step(optimizer)  # type: ignore
                    scaler.update()
                    optimizer.zero_grad(set_to_none=True)
            else:
                if cfg.environment.use_deepspeed:
                    model.backward(loss)  # type: ignore[operator]
                else:
                    loss.backward()
                if itr % cfg.training.grad_accumulation == 0:
                    if cfg.training.gradient_clip > 0:
                        torch.nn.utils.clip_grad_norm_(
                            model.parameters(), cfg.training.gradient_clip
                        )
                    optimizer.step()
                    optimizer.zero_grad(set_to_none=True)

            if cfg.environment._distributed:
                torch.cuda.synchronize(device=cfg.environment._local_rank)

            if scheduler is not None:
                scheduler.step()

            if cfg.environment._local_rank == 0:
                cfg.logging._logger.log(
                    "train", "loss", losses[-1], step=cfg.environment._curr_step
                )
                cfg.logging._logger.log(
                    "meta",
                    "lr",
                    optimizer.param_groups[0]["lr"],
                    step=cfg.environment._curr_step,
                )
                if cfg.training.differential_learning_rate_layers:
                    cfg.logging._logger.log(
                        "meta",
                        "lr_diff",
                        optimizer.param_groups[2]["lr"],
                        step=cfg.environment._curr_step,
                    )

                cfg.logging._logger.log(
                    "internal",
                    "current_step",
                    cfg.environment._curr_step,
                    step=cfg.environment._curr_step,
                )
                for key in output_dict:
                    if key.startswith("additional_log_"):
                        cfg.logging._logger.log(
                            "train",
                            key.replace("additional_log_", ""),
                            output_dict[key].item(),
                            step=cfg.environment._curr_step,
                        )

                # Show logs each 5% of the epoch (only if doing per epoch evaluation)
                if (itr + 1) % log_update_steps == 0 or itr == epoch_steps - 1:
                    progress_bar.set_description(
                        f"train loss: {np.mean(losses[-10:]):.2f}", refresh=False
                    )
                    if (itr + 1) % log_update_steps == 0:
                        progress_bar.update(log_update_steps)
                    else:
                        progress_bar.update(epoch_steps % log_update_steps)

                del output_dict

            # Validation loop
            if (itr + 1) % evaluation_step == 0:
                if cfg.training.evaluation_epochs == 1:
                    progress_bar.close()

                # TODO: Move back after fixing slow generation of deepspeed.
                if not cfg.training.save_best_checkpoint:
                    checkpoint_path = cfg.output_directory
                    if cfg.environment._local_rank == 0:
                        logger.info(
                            f"Saving last model checkpoint to {checkpoint_path}"
                        )
                    save_checkpoint(model=model, path=checkpoint_path, cfg=cfg)

                val_loss, val_metric = run_eval(
                    cfg=cfg, model=model, val_dataloader=val_dataloader, val_df=val_df
                )

                if cfg.training.save_best_checkpoint:
                    if objective_op(val_metric, best_val_metric):
                        checkpoint_path = cfg.output_directory
                        if cfg.environment._local_rank == 0:
                            logger.info(
                                f"Saving best model checkpoint: "
                                f"val_{cfg.prediction.metric} {best_val_metric:.5} -> "
                                f"{val_metric:.5} to {checkpoint_path}"
                            )
                        save_checkpoint(model=model, path=checkpoint_path, cfg=cfg)
                        best_val_metric = val_metric

                model.train()

        progress_bar.close()
        del progress_bar

        if cfg.environment._distributed:
            torch.cuda.synchronize(device=cfg.environment._local_rank)
            torch.distributed.barrier()

        if cfg.environment._local_rank == 0:
            cfg.logging._logger.log(
                "internal", "epoch", epoch + 1, step=cfg.environment._curr_step
            )

    if cfg.environment._distributed:
        torch.distributed.barrier()

    return val_loss, val_metric


def run(cfg: Any) -> None:
    """Runs the routine.

    Args:
        cfg: config object with all the hyperparameters
    """

    if cfg.problem_type == "text_rlhf_language_modeling":
        raise DeprecationWarning(
            "text_rlhf_language_modeling is deprecated. "
            "Please use DPO Modeling instead."
        )

    os.makedirs(cfg.output_directory, exist_ok=True)

    # Force evaluation if user trains 0 epochs
    cfg.training.evaluate_before_training = (
        cfg.training.evaluate_before_training or cfg.training.epochs == 0
    )

    # Set the random seed for reproducibility
    # either random seed when user set it -1 or deterministic user chosen seed
    if cfg.environment.seed < 0:
        cfg.environment._seed = np.random.randint(1_000_000)
    else:
        cfg.environment._seed = cfg.environment.seed

    if (
        cfg.architecture.backbone_dtype in ["int8", "int4"]
        and cfg.environment.use_deepspeed
    ):
        raise ValueError(
            f"Deepspeed do not support backbone type {cfg.architecture.backbone_dtype}."
            + " Please set backbone type to float16 or bfloat16 for using deepspeed."
        )

    # Prepare environment
    if "WORLD_SIZE" in os.environ:
        cfg.environment._distributed = int(os.environ["WORLD_SIZE"]) > 1
    else:
        cfg.environment._distributed = False

    if cfg.environment._distributed:
        cfg.environment._local_rank = int(os.environ["LOCAL_RANK"])
        cfg.environment._device = "cuda:%d" % cfg.environment._local_rank
        if cfg.environment.use_deepspeed:
            deepspeed.init_distributed()
        else:
            torch.distributed.init_process_group(backend="nccl", init_method="env://")
        cfg.environment._cpu_comm = torch.distributed.new_group(backend="gloo")

        cfg.environment._world_size = torch.distributed.get_world_size()
        cfg.environment._rank = torch.distributed.get_rank()
        torch.cuda.set_device(cfg.environment._rank)
        logger.info(
            f"Training in distributed mode with multiple processes, "
            f"1 GPU per process. Process {cfg.environment._rank}, "
            f"total: {cfg.environment._world_size} "
            f"local rank: {cfg.environment._local_rank}."
        )

        # Sync the random seed
        cfg.environment._seed = int(
            sync_across_processes(
                np.array([cfg.environment._seed]),
                cfg.environment._world_size,
                group=cfg.environment._cpu_comm,
            )[0]
        )
    else:
        cfg.environment._local_rank = 0
        cfg.environment._device = (
            "cuda:0"
            if (torch.cuda.is_available() and len(cfg.environment.gpus) > 0)
            else "cpu"
        )
        if cfg.environment._device == "cpu":
            logger.warning("Training on CPU. This will be slow.")

    set_seed(cfg.environment._seed)
    if cfg.environment._local_rank == 0:
        logger.info(f"Problem Type: {cfg.problem_type}")
        logger.info(f"Global random seed: {cfg.environment._seed}")

    cfg = set_environment(cfg)

    # we need to get train dataframe and number of labels if not set or in training mode
    if cfg.environment._local_rank == 0:
        logger.info("Preparing the data...")
    train_df, val_df = get_data(cfg)

    if (
        len(val_df) > int(os.getenv("GPT_EVAL_MAX", 100))
        and "GPT" in cfg.prediction.metric
    ):
        logger.warning(
            f"More than {os.getenv('GPT_EVAL_MAX', 100)} validation records. "
            "Safeguarding against OpenAI API costs. Setting metric to BLEU. "
            "Change GPT_EVAL_MAX to run GPT validation."
        )
        cfg.prediction.metric = "BLEU"

    # prepare data
    if cfg.environment._local_rank == 0:
        logger.info("Preparing train and validation data")
    train_dataset = get_train_dataset(train_df=train_df, cfg=cfg)
    val_dataset = get_val_dataset(val_df=val_df, cfg=cfg)
    train_dataloader = get_train_dataloader(train_ds=train_dataset, cfg=cfg)
    val_dataloader = get_val_dataloader(val_ds=val_dataset, cfg=cfg)

    if cfg.environment._local_rank == 0:
        total_training_steps = (
            cfg.training.epochs
            * len(train_dataloader)
            * cfg.training.batch_size
            * cfg.environment._world_size
        )

        num_eval_epochs = get_number_of_validation_epochs(
            training_epochs=cfg.training.epochs,
            evaluation_epochs=cfg.training.evaluation_epochs,
        )
        val_batch_size = get_inference_batch_size(cfg)

        # if zero shot, validate once before training
        total_validation_steps = (
            len(val_dataloader)
            * (num_eval_epochs + int(cfg.training.evaluate_before_training))
            * val_batch_size
            * cfg.environment._world_size
        )

    # Prepare model and optimizer
    if cfg.environment.use_deepspeed:
        ds_config = get_ds_config(cfg)
        # keep this object alive.
        dschf = HfDeepSpeedConfig(ds_config)  # noqa: F841
    with torch.device(cfg.environment._device):
        model = cfg.architecture.model_class(cfg)
        check_disk_space(model, cfg.output_directory)

        # load model weights
        if cfg.architecture.pretrained_weights != "":
            # Do not load strictly if continue training from the previous experiment
            load_checkpoint(cfg, model, strict=cfg.training.epochs == -1)
    model.to(cfg.environment._device)

    epoch_steps = len(train_dataloader)
    optimizer = get_optimizer(model=model, cfg=cfg)
    scheduler = get_scheduler(cfg=cfg, optimizer=optimizer, epoch_steps=epoch_steps)

    if getattr(cfg.architecture, "force_embedding_gradients"):
        for module in model.modules():
            if isinstance(module, torch.nn.Embedding):
                for param in module.parameters():
                    param.requires_grad = True
                    param.data = param.data.float()

    if cfg.environment._distributed:
        (
            model,
            optimizer,
            train_dataloader,
            val_dataloader,
            scheduler,
        ) = wrap_model_distributed(
            model=model,
            optimizer=optimizer,
            lr_scheduler=scheduler,
            train_dataloader=train_dataloader,
            val_dataloader=val_dataloader,
            cfg=cfg,
        )

    if cfg.environment.compile_model:
        # deepspeed do not support torch.compile
        if cfg.environment.use_deepspeed:
            logger.warning(
                "Deepspeed is active, but it doesn't support torch.compile."
                "Skipping compilation for this experiment."
            )
        else:
            if cfg.environment._distributed:
                model.module.backbone = torch.compile(model.module.backbone)
            else:
                model.backbone = torch.compile(model.backbone)

    # Force settings when saving best checkpoint
    if cfg.training.save_best_checkpoint:
        cfg.training.train_validation_data = False

    # reset steps
    cfg.environment._curr_step = 0
    cfg.environment._curr_val_step = 0

    gc.collect()

    global_start_time = time.time()
    if cfg.environment._local_rank == 0:
        # re-save cfg
        save_config_yaml(f"{cfg.output_directory}/cfg.yaml", cfg)

        cfg.logging._logger = MainLogger(cfg)

        cfg.logging._logger.log(
            "internal", "total_training_steps", total_training_steps, step=0
        )

        cfg.logging._logger.log(
            "internal", "total_validation_steps", total_validation_steps, step=0
        )

        cfg.logging._logger.log(
            "internal",
            "global_start_time",
            global_start_time,
            step=cfg.environment._curr_step,
        )
        # re-save config
        save_config_yaml(f"{cfg.output_directory}/cfg.yaml", cfg)

    val_loss, val_metric = run_train(
        cfg=cfg,
        model=model,
        optimizer=optimizer,
        scheduler=scheduler,
        epoch_steps=epoch_steps,
        train_dataloader=train_dataloader,
        val_dataloader=val_dataloader,
        val_df=val_df,
    )

    # reset external logging
    if cfg.environment._local_rank == 0:
        cfg.logging._logger.reset_external()

    experiment_path = f"{cfg.output_directory}"

    if cfg.training.epochs == 0:
        checkpoint_path = cfg.output_directory
        if cfg.environment._local_rank == 0:
            logger.info(f"Saving last model checkpoint to {checkpoint_path}")
        save_checkpoint(model=model, path=checkpoint_path, cfg=cfg)

    if cfg.environment._local_rank == 0:
        save_config_yaml(f"{cfg.output_directory}/cfg.yaml", cfg)
        save_prediction_outputs(cfg.experiment_name, experiment_path)

        flag_path = os.path.join(cfg.output_directory, "flags.json")
        write_flag(flag_path, "status", "finished")
        time_took = time.time() - global_start_time
        if time_took > 86400:
            # if more than one day, show days
            # need to subtract 1 day from time_took since strftime shows day of year
            # which starts counting at 1
            time_took_formatted = time.strftime(
                "%-jd %H:%M:%S", time.gmtime(float(time_took - 86400))
            )
        else:
            time_took_formatted = time.strftime(
                "%H:%M:%S", time.gmtime(float(time_took))
            )
        write_flag(flag_path, "info", f"Runtime: {time_took_formatted}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="")
    parser.add_argument(
        "-C", "--config", help="config filename", default=argparse.SUPPRESS
    )
    parser.add_argument("-Y", "--yaml", help="yaml filename", default=argparse.SUPPRESS)
    parser_args, unknown = parser.parse_known_args(sys.argv)

    if "config" in parser_args:
        cfg = load_config_py(parser_args.config)
    elif "yaml" in parser_args:
        cfg = load_config_yaml(parser_args.yaml)
    else:
        raise ValueError("Please, provide a configuration file")

    extra_args = []
    for arg_orig in unknown:
        if arg_orig.startswith(("-", "--")):
            arg = arg_orig.replace("-", "").split(".")
            try:
                arg_type = getattr(cfg, arg[0]).get_annotations()[arg[1]]
            except (AttributeError, KeyError):
                continue
            if arg_type == bool:
                parser.add_argument(arg_orig, type=util.strtobool)
            else:
                parser.add_argument(arg_orig, type=arg_type)
            extra_args.append(arg)

    args = parser.parse_args()

    for arg in extra_args:
        value = getattr(args, ".".join(arg))
        setattr(getattr(cfg, arg[0]), arg[1], value)

    out_dir = cfg.output_directory
    os.makedirs(out_dir, exist_ok=True)

    initialize_logging(cfg)

    try:
        run(cfg=cfg)
    except Exception:
        logging.error("Exception occurred during the run:", exc_info=True)
        if ("WORLD_SIZE" in os.environ) and (int(os.environ["WORLD_SIZE"]) > 1):
            kill_ddp_processes()