lrm/flux/trainer/accelerators/base_accelerator.py

import abc
import hashlib
import json
import math
import os
import shutil
import time
from dataclasses import field, dataclass
from glob import glob
from typing import List, Optional

import datasets
import torch
import transformers
from accelerate.logging import get_logger
from accelerate.utils import set_seed as accelerate_set_seed, PrecisionType
from accelerate.utils.dataclasses import BaseEnum, LoggerType, DynamoBackend
from omegaconf import DictConfig, OmegaConf, II
from tqdm import tqdm

from trainer.accelerators.utils import get_nvidia_smi_gpu_memory_stats_str, print_config, _flatten_dict

logger = get_logger(__name__)

TRAINING_STAGE_PATH = "training_stage.json"


def debug(port):
    logger.info("Connecting to debugger...")
    import pydevd_pycharm
    pydevd_pycharm.settrace('localhost', port=port, stdoutToServer=True, stderrToServer=True)


@dataclass
class DebugConfig:
    activate: bool = False
    port: int = 5900


class TrainingMode(BaseEnum):
    SKIPPING = "skipping"
    TRAINING = "training"


class MetricMode(BaseEnum):
    MAX = "max"
    MIN = "min"


@dataclass
class BaseAcceleratorConfig:
    _target_: str = "trainer.accelerators.base_accelerator.Accelerator"
    output_dir: str = II("output_dir")
    mixed_precision: PrecisionType = PrecisionType.NO
    gradient_accumulation_steps: int = 1
    log_with: Optional[LoggerType] = LoggerType.WANDB
    debug: DebugConfig = field(default_factory=DebugConfig)
    seed: int = 42
    resume_from_checkpoint: bool = True
    max_steps: int = 4000
    num_epochs: int = 10
    validate_steps: int = 100
    generalization_validate_steps: int = 500
    eval_on_start: bool = True
    project_name: str = "reward"
    run_name: str = "default"
    max_grad_norm: float = 1.0
    save_steps: int = 100
    metric_name: str = "accuracy"
    metric_mode: MetricMode = MetricMode.MAX
    limit_num_checkpoints: int = 1
    save_only_if_best: bool = True
    dynamo_backend: DynamoBackend = DynamoBackend.NO
    keep_best_ckpts: bool = True
    progress_log_interval: int = 50


class BaseAccelerator(abc.ABC):

    def __init__(self, cfg: BaseAcceleratorConfig):
        self.cfg = cfg
        self.accelerator = None
        self.epoch = 0
        self.step = 0
        self.global_step = 0
        self.step_loss = 0.0
        self.lr = None
        self.metrics = {}
        self.progress_bar = None
        self.mode = TrainingMode.TRAINING
        self.num_update_steps_per_epoch = None
        self.num_steps_per_epoch = None
        self.training_start_time = None

    def post_init(self):
        self.set_seed()
        self.debug()
        logger.info(f"Initialized accelerator: rank={self.accelerator.process_index}", main_process_only=False)
        self.set_logging_level()

    def set_logging_level(self):
        if self.accelerator.is_local_main_process:
            datasets.utils.logging.set_verbosity_warning()
            transformers.utils.logging.set_verbosity_warning()
        else:
            datasets.utils.logging.set_verbosity_error()
            transformers.utils.logging.set_verbosity_error()

    def debug(self):
        if self.accelerator.is_main_process and self.cfg.debug.activate:
            debug(self.cfg.debug.port)

    def set_seed(self):
        logger.info(f"Setting seed {self.cfg.seed}")
        accelerate_set_seed(self.cfg.seed, device_specific=True)

    def prepare(self, *args, device_placement=None):
        return self.accelerator.prepare(*args, device_placement=device_placement)

    def get_latest_checkpoint(self):
        all_ckpts = list(glob(os.path.join(self.cfg.output_dir, "checkpoint-*")))
        if len(all_ckpts) == 0:
            return
        all_ckpts.sort(key=os.path.getctime)
        if "final" in all_ckpts[-1]:
            all_ckpts.pop()
        return all_ckpts[-1] if len(all_ckpts) > 0 else None

    def load_state_if_needed(self):
        if not self.cfg.resume_from_checkpoint:
            return
        ckpt_path = self.get_latest_checkpoint()

        if ckpt_path is None:
            logger.info("No checkpoint found, training from scratch")
            return

        stage = json.load(open(os.path.join(ckpt_path, TRAINING_STAGE_PATH)))
        self.epoch, self.step, self.global_step, self.metrics = stage["epoch"], stage["step"], stage["global_step"], \
            stage["metrics"]
        logger.info(
            f"Resuming from checkpoint: {ckpt_path} | epoch={self.epoch} step={self.step} gstep={self.global_step}")
        self.accelerator.load_state(ckpt_path)
        logger.info("Checkpoint loaded")

    @property
    def is_main_process(self):
        return self.accelerator.is_main_process

    @property
    def num_processes(self):
        return self.accelerator.num_processes

    def pre_training_log(self, cfg: DictConfig):
        total_batch_size = cfg.dataset.batch_size * self.num_processes * self.cfg.gradient_accumulation_steps
        logger.info("***** Running training *****")
        logger.info(f"  Instantaneous batch size per device = {cfg.dataset.batch_size}")
        logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
        logger.info(f"  Gradient Accumulation steps = {self.cfg.gradient_accumulation_steps}")
        logger.info(f"  Total warmup steps = {cfg.lr_scheduler.lr_warmup_steps}")
        logger.info(f"  Total training steps = {self.cfg.max_steps * self.cfg.gradient_accumulation_steps}")
        logger.info(f"  Total epochs = {self.cfg.num_epochs}")
        logger.info(f"  Steps per epoch = {self.num_steps_per_epoch}")
        logger.info(f"  Update steps per epoch = {self.num_update_steps_per_epoch}")
        logger.info(f"  Total optimization steps = {self.cfg.max_steps}")
        logger.info(f"  Mixed precision = {self.cfg.mixed_precision}")
        logger.info(f"  World size = {self.accelerator.num_processes}")

    def init_training(self, cfg: DictConfig):
        if self.is_main_process:
            yaml = OmegaConf.to_yaml(cfg, resolve=True, sort_keys=True)
            log_cfg = _flatten_dict(OmegaConf.create(yaml))
            logger.info("Initializing trackers")
            self.accelerator.init_trackers(
                self.cfg.project_name, 
                log_cfg,
                init_kwargs={"wandb": {
                    "name": self.cfg.run_name, 
                    "entity": None,
                }}
            )
            logger.info("Training config:")
            print_config(cfg)
        logger.info(get_nvidia_smi_gpu_memory_stats_str())
        self.pre_training_log(cfg)
        self.training_start_time = time.time()
        self.progress_bar = tqdm(range(self.cfg.max_steps * self.cfg.gradient_accumulation_steps), disable=not self.accelerator.is_main_process)
        self.progress_bar.set_description("Steps")

    @staticmethod
    def _format_seconds(total_seconds: float) -> str:
        total_seconds = max(0, int(total_seconds))
        hours, rem = divmod(total_seconds, 3600)
        minutes, seconds = divmod(rem, 60)
        if hours > 0:
            return f"{hours:02d}:{minutes:02d}:{seconds:02d}"
        return f"{minutes:02d}:{seconds:02d}"

    def maybe_log_progress_eta(self):
        if not self.is_main_process:
            return

        interval = int(getattr(self.cfg, "progress_log_interval", 0) or 0)
        if interval <= 0:
            return

        if self.global_step <= 0:
            return

        should_log = self.global_step == 1 or self.global_step % interval == 0 or self.global_step >= self.cfg.max_steps
        if not should_log:
            return

        if not self.training_start_time:
            return

        elapsed = time.time() - self.training_start_time
        if elapsed <= 0:
            return

        steps_per_second = self.global_step / elapsed
        if steps_per_second <= 0:
            return

        remaining_steps = max(self.cfg.max_steps - self.global_step, 0)
        eta_seconds = remaining_steps / steps_per_second
        progress_pct = 100.0 * self.global_step / max(self.cfg.max_steps, 1)

        logger.info(
            "Training progress: step=%s/%s (%.2f%%), epoch=%s, lr=%s, speed=%.3f step/s, elapsed=%s, eta=%s",
            self.global_step,
            self.cfg.max_steps,
            progress_pct,
            self.epoch,
            self.lr,
            steps_per_second,
            self._format_seconds(elapsed),
            self._format_seconds(eta_seconds),
        )

    def should_skip(self, epoch, step):
        should = epoch < self.epoch or (epoch == self.epoch and step < self.step)
        if should:
            self.mode = TrainingMode.SKIPPING
            self.progress_bar.set_postfix(**{"status": TrainingMode.SKIPPING})
        else:
            self.mode = TrainingMode.TRAINING
        return should

    def update_progbar_step(self):
        self.progress_bar.update(1)

    def log(self, data):
        if self.is_main_process:
            self.accelerator.log(data, step=self.global_step)

    def recalc_train_length_after_prepare(self, num_batches):
        num_update_steps_per_epoch = math.ceil(num_batches / self.cfg.gradient_accumulation_steps)
        if self.cfg.max_steps is None:
            self.cfg.max_steps = self.cfg.num_epochs * num_update_steps_per_epoch
        self.num_update_steps_per_epoch = num_update_steps_per_epoch
        self.num_steps_per_epoch = num_batches
        self.cfg.num_epochs = math.ceil(self.cfg.max_steps / num_update_steps_per_epoch)
        
        logger.info(f"num_update_steps_per_epoch = {num_update_steps_per_epoch}")
        logger.info(f"num_batches = {num_batches}")
        logger.info(f"num_epochs = {self.cfg.num_epochs}")
        
        

    def accumulate(self, model):
        return self.accelerator.accumulate(model)

    def gather(self, data):
        return self.accelerator.gather(data)

    @property
    def sync_gradients(self):
        return self.accelerator.sync_gradients

    def update_step_loss(self, loss):
        self.step_loss = loss

    def update_global_step(self, loss):
        self.global_step += 1
        self.log({
            "lr": self.lr,
            "step": self.step,
            "epoch": self.epoch,
            "global_step": self.global_step,
            "loss": loss,
        })

    def get_allocated_cuda_memory(self):
        return round(torch.cuda.max_memory_allocated(self.accelerator.device) / 1024 / 1024 / 1024, 2)

    def update_step(self, loss, lr):
        self.step += 1
        self.lr = lr
        logs = {
            "stl": loss,
            "gstl": loss,
            "mem": self.get_allocated_cuda_memory(),
            "st": self.step,
            "ep": self.epoch,
            "gst": self.global_step,
            "lr": self.lr,
        }
        self.progress_bar.set_postfix(**logs)
        self.maybe_log_progress_eta()
        self.update_progbar_step()

    def wait_for_everyone(self):
        self.accelerator.wait_for_everyone()

    def update_epoch(self):
        if self.mode == TrainingMode.SKIPPING:
            return
        logger.info(f"Epoch {self.epoch} finished")
        self.epoch += 1
        self.step = 0

    def update_metrics(self, metrics):
        self.metrics.update(metrics)
        logger.info(f"Metrics: {self.metrics}")
        self.log(metrics)

    def end_training(self):
        self.accelerator.wait_for_everyone()
        self.accelerator.end_training()

    def unwrap_and_save(self, model):
        if not self.is_main_process:
            return
        model = self.accelerator.unwrap_model(model)
        save_dir = os.path.join(self.cfg.output_dir, f"checkpoint-final")
        logger.info(f"Saving final checkpoint to {save_dir}")
        model.save(save_dir)
        self.save_training_stage(save_dir)
        logger.info(f"Saved checkpoint to {save_dir}")

    def should_end(self):
        return self.global_step >= self.cfg.max_steps

    def backward(self, loss):
        self.accelerator.backward(loss)

    def clip_grad_norm_(self, params):
        self.accelerator.clip_grad_norm_(params, self.cfg.max_grad_norm)

    def should_eval(self):
        if not self.mode == TrainingMode.TRAINING:
            return False
        if self.step == 0 and self.global_step == 0 and self.cfg.eval_on_start:
            return True
        if self.global_step > 0 and self.sync_gradients and self.global_step % self.cfg.validate_steps == 0:
            return True
        return False
    
    def should_generalization_eval(self):
        if not self.mode == TrainingMode.TRAINING:
            return False
        if self.step == 0 and self.global_step == 0 and self.cfg.eval_on_start:
            return True
        if self.global_step > 0 and self.sync_gradients and self.global_step % self.cfg.generalization_validate_steps == 0:
            return True
        return False

    def should_save(self):
        return self.sync_gradients and self.global_step > 0 and self.cfg.save_steps > 0 and self.global_step % self.cfg.save_steps == 0

    @property
    def training_stage(self):
        return {
            "epoch": self.epoch,
            "step": self.step,
            "global_step": self.global_step,
            "step_loss": self.step_loss,
            "lr": self.lr,
            "metrics": self.metrics,
        }

    def save_training_stage(self, save_dir):
        json.dump(self.training_stage, open(os.path.join(save_dir, TRAINING_STAGE_PATH), "w"), indent=4)

    def save_checkpoint(self):
        if self.cfg.save_only_if_best:
            all_ckpts = self.get_all_ckpts()
            for ckpt in all_ckpts:
                training_stage = json.load(open(os.path.join(ckpt, TRAINING_STAGE_PATH)))
                metric_val = training_stage["metrics"][self.cfg.metric_name]
                cur_metric_val = self.training_stage["metrics"][self.cfg.metric_name]
                if (self.cfg.metric_mode == MetricMode.MIN and metric_val < cur_metric_val) or \
                        (self.cfg.metric_mode == MetricMode.MAX and metric_val > cur_metric_val):
                    logger.info(
                        f"Metric {self.cfg.metric_name}={cur_metric_val} is not better than {metric_val} of {ckpt}, skipping checkpoint")
                    return
        self.cleanup_checkpoints()
        self.accelerator.wait_for_everyone()
        save_dir = os.path.join(self.cfg.output_dir, f"checkpoint-gstep{self.global_step}")
        logger.info(f"Saving checkpoint to {save_dir}")
        self.accelerator.save_state(save_dir)
        if self.accelerator.is_main_process:
            self.save_training_stage(save_dir)
        # self.save_training_stage(save_dir)
        logger.info(f"Saved checkpoint to {save_dir}")

    @property
    def gradient_state(self):
        return self.accelerator.gradient_state

    def get_all_ckpts(self):
        return list(glob(os.path.join(self.cfg.output_dir, f"checkpoint-*")))

    def load_best_checkpoint(self):
        all_ckpts = self.get_all_ckpts()
        if not self.cfg.keep_best_ckpts:
            all_ckpts.sort(key=os.path.getctime, reverse=True)
            logger.info(f"Returning the most recent checkpoint: {all_ckpts[0]}")
            return all_ckpts[0]
        logger.info(f"Found {len(all_ckpts)} checkpoints in {self.cfg.output_dir}")
        logger.info(all_ckpts)
        if len(all_ckpts) == 0:
            logger.info(f"No checkpoint found in {self.cfg.output_dir} to load. Keeping current model.")
            return
        best_ckpt, best_metric_val = None, math.inf if self.cfg.metric_mode == MetricMode.MIN else -math.inf
        for ckpt in all_ckpts:
            training_stage = json.load(open(os.path.join(ckpt, TRAINING_STAGE_PATH)))
            metric_val = training_stage["metrics"][self.cfg.metric_name]
            if (self.cfg.metric_mode == MetricMode.MIN and metric_val < best_metric_val) or \
                    (self.cfg.metric_mode == MetricMode.MAX and metric_val > best_metric_val):
                best_ckpt, best_metric_val = ckpt, metric_val
        logger.info(f"Loading best checkpoint from {best_ckpt} with metric {self.cfg.metric_name}={best_metric_val}")
        self.accelerator.load_state(best_ckpt)

    @property
    def device(self):
        return self.accelerator.device

    def cleanup_checkpoints(self):
        if self.cfg.limit_num_checkpoints <= 0 or not self.accelerator.is_main_process:
            logger.info(f"Not cleaning up checkpoints as limit_num_checkpoints={self.cfg.limit_num_checkpoints}")
            return

        all_ckpts = self.get_all_ckpts()
        if len(all_ckpts) <= self.cfg.limit_num_checkpoints:
            logger.info(f"Not cleaning up checkpoints as only {len(all_ckpts)} checkpoints found")
            return

        logger.info(f"Found {len(all_ckpts)} checkpoints in {self.cfg.output_dir}")
        ckpts_to_delete = self.get_ckpts_to_delete()
        ckpts_to_delete.sort(key=os.path.getctime)

        ckpts_to_delete = ckpts_to_delete[:-1]
        for ckpt in ckpts_to_delete:
            logger.info(f"Deleting checkpoint {ckpt}")
            shutil.rmtree(ckpt)

    def get_ckpts_to_delete(self):
        all_ckpts = self.get_all_ckpts()
        if self.cfg.keep_best_ckpts:
            metric_vals = []
            for ckpt in all_ckpts:
                training_stage = json.load(open(os.path.join(ckpt, TRAINING_STAGE_PATH)))
                metric_val = training_stage["metrics"][self.cfg.metric_name]
                metric_vals.append(metric_val)
            metric_ckpt = list(zip(metric_vals, all_ckpts))
            metric_ckpt.sort(key=lambda x: x[0], reverse=self.cfg.metric_mode == MetricMode.MAX)
            ckpts_to_delete = [ckpt for _, ckpt in metric_ckpt[self.cfg.limit_num_checkpoints:]]
        else:
            all_ckpts.sort(key=os.path.getctime, reverse=True)
            ckpts_to_delete = all_ckpts[self.cfg.limit_num_checkpoints:]
        return ckpts_to_delete