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