import copy import gc import json import logging import os import pickle import socket import subprocess import time from collections import defaultdict from pathlib import Path from typing import Dict, List, Union import accelerate import psutil import torch import torch.optim as torch_optim import transformers.optimization as transformers_optim import wandb from packaging import version from m4.training.config import ( DataParams, GlobalBatchSizeRampUpRunningParams, Hparams, OptimizerParams, Parameters, ResumeParams, ) from m4.training.dataset import DatasetNames from m4.training.debug_utils import validate_optim_states_are_reset from m4.training.utils import ( # deepspeed_gathered_parameters_context_manager, IMAGE_TOKEN, JSONEncoderForDataclasses, LoggingTypes, SigtermListener, get_deepspeed_engine, get_stats, get_stats_format, is_deepspeed_used, is_deepspeed_zero3_used, is_deepspeed_zero_init_enabled, lora_unload, mem_usage_formatted, pynmvl_handle, pynvml_get_total_energy_in_joules, ) from m4.utils.activation_tracker import ActivationTracker from m4.utils.debug import printflock as print from m4.utils.progress import BarColumn, MofNCompleteColumn, Progress, TaskProgressColumn, TimeElapsedColumn from m4.utils.training.timer import DeviceAgnosticTimer, Timer, format_secs_to_sec_fractions, format_secs_to_time logger = logging.getLogger(__name__) fqdn = socket.getfqdn() if "compute.internal" in fqdn: # hfc: 1.1TB RAM _MEMORY_EXPLOSION_THRESHOLD = 93.0 elif "idris.fr" in fqdn or "idrsrv" in fqdn: # jz: 0.5TB RAM _MEMORY_EXPLOSION_THRESHOLD = 90.0 else: _MEMORY_EXPLOSION_THRESHOLD = 90.0 METRICS_TO_DEFAULT_VALUE_FN = { "lr": lambda: None, "num_opt_steps": lambda: 0, "num_epochs": lambda: 0, "per_token_loss": lambda: defaultdict(lambda: None), "z_loss": lambda: defaultdict(lambda: None), "watt/s": lambda: defaultdict(lambda: None), "tflops": lambda: defaultdict(lambda: None), "tflop_counter": lambda: defaultdict(lambda: None), "fwd_bwd_time": lambda: defaultdict(lambda: None), "tflops_acc": lambda: defaultdict(lambda: None), "num_per_device_batches": lambda: defaultdict(lambda: None), "num_images": lambda: defaultdict(lambda: None), "num_image_tokens": lambda: defaultdict(lambda: None), "num_tokens": lambda: defaultdict(lambda: None), "image_to_text_ratio": lambda: defaultdict(lambda: None), "pixel_values_sum": lambda: defaultdict(lambda: None), "num_padding": lambda: defaultdict(lambda: None), "num_per_device_batches_in_curr_epoch": lambda: defaultdict(lambda: None), "num_batches": lambda: defaultdict(int), "num_batches_in_curr_epoch": lambda: defaultdict(int), "per_token_loss_acc": lambda: defaultdict(lambda: None), "z_loss_acc": lambda: defaultdict(lambda: None), "num_batches_since_training_logged": lambda: defaultdict(int), "num_per_device_batches_since_training_logged": lambda: defaultdict(lambda: None), "tflop_counter_since_training_logged": lambda: defaultdict(lambda: None), "total_energy_delta_since_training_logged": lambda: defaultdict(lambda: None), "fwd_bwd_time_since_training_logged": lambda: defaultdict(lambda: None), } METRICS_TO_RESET_AFTER_LOGGING = [ "per_token_loss_acc", "z_loss_acc", "num_batches_since_training_logged", "num_per_device_batches_since_training_logged", "tflop_counter_since_training_logged", "total_energy_delta_since_training_logged", "fwd_bwd_time_since_training_logged", ] class Trainer(object): """ Trainer object to monitor training and validation :- config -: json config file """ def __init__(self, accelerator, vl_model, tokenizer, train_loader, val_loader, config): # Initialize params self.config: Parameters = config self.optim_param: OptimizerParams = config.optim_param self.hparams: Hparams = config.hparams self.resume_param: ResumeParams = config.resume_param self.data_param: DataParams = config.data_param # Initialize last step directory self.last_opt_step_dir = "" # Initialize the model self.vl_model = vl_model # Gradient checkpointing if self.hparams.gradient_checkpointing: self.vl_model.gradient_checkpointing_enable() # Debug if accelerator.is_main_process and self.hparams.train_logging_activations: self.activation_tracker = ActivationTracker(self.vl_model) else: self.activation_tracker = None # Initialize tokenizer self.tokenizer = tokenizer self.image_token_id = self.tokenizer.convert_tokens_to_ids(IMAGE_TOKEN) # Initialize accelerator self.accelerator = accelerator # Initialize loaders self.train_loader = train_loader self.val_loader = val_loader # Checks self._compatibility_checks() # Initialize everything related to distributed training self._configure_optimizer_and_scheduler() # Prepare and/or register model, optimizer, dataloaders and scheduler self._prepare_register() # now that we have num_processes, figure out batch_size-related variables self.setup_batch_size_related_configs() # Compute useful variables self.optim_param.opt_batch_size = self.hparams.global_batch_size if self.hparams.max_num_opt_steps is None and self.hparams.max_num_epochs is None: if hasattr(self.train_loader, "__len__") and self.hparams.global_batch_size_ramp_up.start is not None: raise ValueError("Currently global batch size ramp up doesn't work with MappedDataset") try: self.hparams.max_num_opt_steps = int(len(self.train_loader) // self.hparams.grad_acc_size) except TypeError: raise ValueError("max_num_opt_steps or max_num_epochs must be defined if you use IterableDataset") # self._set_model_tflops_per_batch_per_gpu() # Init trackers self._init_trackers() # Handle jz timing and memory self.jz_training_time_over = [False] self.memory_explosion = False # Stopping on demand self.kill_switch_activated = False # Sigterm signal listener self.sigterm_signal_received = False self.sigterm_listener = SigtermListener() sizes = defaultdict(int) trainable_params = [] numel_fn = lambda p: p.ds_numel if is_deepspeed_zero_init_enabled() else p.numel() # noqa for name, param in self.accelerator.unwrap_model(self.vl_model).named_parameters(): numel = numel_fn(param) sizes["total"] += numel sizes["total_lora"] += numel if "lora_" in name else 0 if "vision_model" in name: sizes["vision_model"] += numel sizes["vision_model_lora"] += numel if "lora_" in name else 0 if "perceiver_resampler" in name: sizes["perceiver_resampler"] += numel if "modality_projection" in name: sizes["modality_projection"] += numel if param.requires_grad: sizes["trainable"] += numel sizes["trainable_lora"] += numel if "lora_" in name else 0 trainable_params += [name] if self.accelerator.is_main_process: logger.info(f""" ------------------------------------- Model: - Total size: {sizes["total"]} ---- Lora size: {sizes["total_lora"]} - Vision encoder size: {sizes["vision_model"]} ---- Lora size: {sizes["vision_model_lora"]} - Perceiver resampler size: {sizes["perceiver_resampler"]} - Modality projection: {sizes["modality_projection"]} - Number of trainable parameters: {sizes["trainable"]} ---- Lora trainable parameters: {sizes["trainable_lora"]} Maximum number of optimizer steps: {self.hparams.max_num_opt_steps} Maximum number of epochs: {self.hparams.max_num_epochs if self.hparams.max_num_epochs else "N/A"} Number of gradient accumulation steps: {self.hparams.grad_acc_size} Number of processes: {self.accelerator.num_processes} Batch sizes: - Per device batch size: {self.hparams.batch_size_per_gpu} - Optimizer batch size: {self.optim_param.opt_batch_size} ------------------------------------- """) logger.info("Trainable/non-trainable parameters:") for name, param in vl_model.named_parameters(): logger.info(f" Name: {name} | Trainable: {param.requires_grad}") if len(self.hparams.train_logging_grad_param_deepspeed) > 0: from deepspeed.utils import safe_get_full_fp32_param, safe_get_full_grad, safe_get_full_optimizer_state self.safe_get_full_fp32_param = safe_get_full_fp32_param self.safe_get_full_grad = safe_get_full_grad self.safe_get_full_optimizer_state = safe_get_full_optimizer_state self.float_placeholder_tensor = torch.tensor(-1.0, device=self.accelerator.device, dtype=torch.float) self.long_placeholder_tensor = torch.tensor(-1, device=self.accelerator.device, dtype=torch.long) def setup_batch_size_related_configs(self): """ batch_size-related configs are processed here. All this work is done here because it requires knowing the value of num_processes """ hparams = self.hparams if hparams.global_batch_size_ramp_up.start is not None: # case 1. global batch size ramp up # 1a. ramp up constraints if ( hparams.global_batch_size_ramp_up.finish is None or hparams.global_batch_size_ramp_up.increment is None or hparams.global_batch_size_ramp_up.samples is None ): raise ValueError( "When using batch size ramp up hparam config entries global_batch_size_ramp_up.start," " global_batch_size_ramp_up.finish, global_batch_size_ramp_up.increment and" " global_batch_size_ramp_up.samples have to be defined." ) # range checks ramp_up_range = hparams.global_batch_size_ramp_up.finish - hparams.global_batch_size_ramp_up.start if ramp_up_range < hparams.global_batch_size_ramp_up.increment: raise ValueError( f"{hparams.global_batch_size_ramp_up.start=} has to be smaller than" f" {hparams.global_batch_size_ramp_up.finish=}." ) if not (ramp_up_range / hparams.global_batch_size_ramp_up.increment).is_integer(): raise ValueError( f"({hparams.global_batch_size_ramp_up.finish=} -" f" {hparams.global_batch_size_ramp_up.start=}) /" f" {hparams.global_batch_size_ramp_up.increment=} has to be a whole number" ) if not ( hparams.global_batch_size_ramp_up.increment / (hparams.batch_size_per_gpu * self.accelerator.num_processes) ).is_integer(): raise ValueError( f"{hparams.global_batch_size_ramp_up.increment=} has to be a multiple of" f" {hparams.batch_size_per_gpu * self.accelerator.num_processes=}" ) if self.accelerator.is_main_process: logger.info( "Will ramp up global batch size from" f" {hparams.global_batch_size_ramp_up.start} to" f" {hparams.global_batch_size_ramp_up.finish}, in increments of" f" {hparams.global_batch_size_ramp_up.increment} every" f" {hparams.global_batch_size_ramp_up.samples} samples." ) # 1b. hparam.grad_acc_size constraints and derivations if hparams.grad_acc_size > 1: raise ValueError("When using batch size ramp up hparam.grad_acc_size must be None or 1.") hparams.grad_acc_size = hparams.global_batch_size_ramp_up.start / ( hparams.batch_size_per_gpu * self.accelerator.num_processes ) if not hparams.grad_acc_size.is_integer(): raise ValueError( f"{hparams.global_batch_size_ramp_up.start=} has to be a multiple of" f" {hparams.batch_size_per_gpu * self.accelerator.num_processes}" ) hparams.grad_acc_size = int(hparams.grad_acc_size) logger.info(f"Derived {hparams.grad_acc_size=}") # 1c. hparams.global_batch_size constraints and derivation if hparams.global_batch_size is not None: raise ValueError("When using batch size ramp up hparam.global_batch_size must be None.") # in the first run we start at global_batch_size == global_batch_size_ramp_up.start hparams.global_batch_size = hparams.global_batch_size_ramp_up.start else: # case 2. fixed global batch size or fixed grad_acc_size # 2a. constraints if hparams.grad_acc_size > 1: # when global_batch_size is used grad_acc_size will be derived automatically from global_batch_size and n_gpus if hparams.global_batch_size is not None and hparams.global_batch_size > 1: raise ValueError("set either hparams.grad_acc_size>1 or hparams.global_batch_size>1, but not both") if hparams.global_batch_size is not None and hparams.global_batch_size > 1: # 2b. have global_batch_size need to derive grad_acc_size hparams.grad_acc_size = hparams.global_batch_size / ( hparams.batch_size_per_gpu * self.accelerator.num_processes ) if not hparams.grad_acc_size.is_integer(): raise ValueError( f"The derived {hparams.grad_acc_size=} is not an integer," f" {hparams.global_batch_size=} / ({hparams.batch_size_per_gpu=} *" f" {self.accelerator.num_processes=})" ) hparams.grad_acc_size = int(hparams.grad_acc_size) logger.info(f"Derived {hparams.grad_acc_size=}") else: # 2c. have grad_acc_size need to derive global_batch_size hparams.global_batch_size = ( hparams.batch_size_per_gpu * hparams.grad_acc_size * self.accelerator.num_processes ) logger.info(f"Derived {hparams.global_batch_size=}") def update_gas_and_gbs(self, grad_acc_size_current, global_batch_size_current): """ Update m4, deepspeed and accelerate with the derived global_batch_size and grad_acc_size """ self.hparams.grad_acc_size = grad_acc_size_current self.hparams.global_batch_size = global_batch_size_current self.accelerator.gradient_accumulation_steps = grad_acc_size_current if is_deepspeed_used(): get_deepspeed_engine(self.accelerator).set_train_batch_size(global_batch_size_current) def _compatibility_checks(self): # BF16 requires cuda>=11 and nccl>=2.10.3 if self.accelerator.mixed_precision == "bf16": if version.parse(torch.version.cuda) < version.parse("11.0"): raise ValueError(f"mixed precision dtype BF16 requires cuda>=11, but got {torch.version.cuda}") if torch.cuda.nccl.version() < (2, 10, 3): raise ValueError( f"mixed precision dtype BF16 requires NCCL>=2.10.3, but got {'.'.join(torch.cuda.nccl.version())}" ) def _init_trackers(self): if self.hparams.wandb_enable and self.accelerator.is_main_process: # Initialize wandb_run_id if self.hparams.resume_run and self.hparams.wandb_run_id != "": pass else: if self.hparams.resume_run and not self.hparams.wandb_run_id != "": logger.warning( "** The run you are resuming was not logging into wandb. Therefore a new wandb_run_id has" " been generated **" ) self.hparams.wandb_run_id = wandb.util.generate_id() logger.info(f"** `wandb_run_id`: {self.hparams.wandb_run_id} **") # Initialize all trackers run_name = self.hparams.save_dir.name wdb_config = {} for k, v in vars(self.config).items(): if not hasattr(v, "__dict__"): wdb_config[k] = v continue for key, value in vars(v).items(): wdb_config[f"{k}-{key}"] = str(value) wandb_logger = wandb.init( config=wdb_config, id=self.hparams.wandb_run_id, resume="allow", project=self.hparams.wandb_project, entity=self.hparams.wandb_entity, name=run_name, allow_val_change=True, tags=self.hparams.wandb_tags, ) self.dummy_module = torch.nn.LayerNorm(1) wandb.watch( self.dummy_module, log="all", log_freq=self.hparams.wandb_log_freq * self.hparams.grad_acc_size, idx=0, # To avoid creating a new panel each we un-watch and then re-watch ) tb_run_name = "tb_run_" + self.hparams.wandb_run_id tensorboard_tracker = accelerate.tracking.TensorBoardTracker( run_name=tb_run_name, logging_dir=self.hparams.save_dir ) logger.info(f"** TensorBoardTracker logging into { self.hparams.save_dir / tb_run_name } **") self.accelerator.trackers = [tensorboard_tracker, wandb_logger] # Alert WB in replacement of slurm notifications and emails wandb.alert( title="Training has either started or resumed", text=( f"Run name = {run_name}, Jobid = {self.hparams.job_id}, Resuming = {self.hparams.resume_run}," f" Experiment folder = {str(self.hparams.save_dir)}" ), ) def _configure_optimizer_and_scheduler(self): """defines model optimizer and lr scheduler""" vl_optim = getattr(torch_optim, self.optim_param.vl_optim) if issubclass(vl_optim, torch_optim.AdamW): no_decay = self.optim_param.vl_optim_params.pop("no_decay", []) weight_decay = self.optim_param.vl_optim_params.pop("weight_decay", 0.0) optim_grouped_params = [ # Group: tunable parameters with weight decay { "params": [ p for n, p in self.vl_model.named_parameters() if not any(nd in n for nd in no_decay) and p.requires_grad ], "weight_decay": weight_decay, }, # Group: tunable parameters without weight decay at all { "params": [ p for n, p in self.vl_model.named_parameters() if any(nd in n for nd in no_decay) and p.requires_grad ], "weight_decay": 0.0, }, ] vl_optim = vl_optim( optim_grouped_params, **self.optim_param.vl_optim_params, ) else: vl_optim = vl_optim( self.vl_model.parameters(), **self.optim_param.vl_optim_params, ) try: vl_scheduler_class = getattr(torch_optim.lr_scheduler, self.optim_param.vl_lr_scheduler) except AttributeError: vl_scheduler_class = getattr(transformers_optim, self.optim_param.vl_lr_scheduler) else: raise ValueError( f"Could not find {self.optim_param.vl_lr_scheduler} type of LR Scheduler in neither `torch.optim` nor" " `transformers.optimization`" ) vl_scheduler = vl_scheduler_class( optimizer=vl_optim, **self.optim_param.vl_lr_scheduler_params, ) self.vl_optim = vl_optim self.vl_scheduler = vl_scheduler def _prepare_register(self): """ Prepare model, optimizer and dataloader if necessary. Register the scheduler for checkpointing. """ if isinstance(self.train_loader.dataset, torch.utils.data.IterableDataset): # `dummy_dataloader`: trick as suggested here: https://discuss.huggingface.co/t/when-using-deepspeed-why-do-i-need-to-pass-dataloaders-to-the-accelerator-prepare/22432/2?u=victorsanh # In our IterableDataset, *WE* handle dispatch (instead of `Accelerate`) for each process ourselves as we need # better shuffling support. # =>> See `DATA_PROCESSING.md` for more information! dummy_dataloader = torch.utils.data.DataLoader( [0 for _ in range(20)], batch_size=self.hparams.batch_size_per_gpu ) # important: do note add lr scheduler in `prepare`. We are doing something non canonical # with our custom data loaders that leads to not being able to use the standard # arguments in the accelerator (typically split_batches which would have ensure the LR # was increased every x steps, where x>1 and would be the correct value for grad acc). # context: https://github.com/huggingface/m4/pull/386 self.vl_model, self.vl_optim, dummy_dataloader = self.accelerator.prepare( self.vl_model, self.vl_optim, dummy_dataloader ) else: ( self.vl_model, self.vl_optim, self.train_loader, self.val_loader, ) = self.accelerator.prepare(self.vl_model, self.vl_optim, self.train_loader, self.val_loader) self.accelerator.register_for_checkpointing(self.vl_scheduler) def _set_up_training(self): """ Prepare variables for trainings. """ if self.hparams.resume_run: # 1. resume train_logs = self.resume_param.train_logs curr_opt_step = self.resume_param.resume_opt_step curr_epoch = self.resume_param.resume_epoch gbs_running = self.resume_param.gbs_running # This check is necessary because the info is saved as json in the checkpoint # and when it is loaded back it is converted to a normal dictionary which can # fail downstream in case one of the dataset keys were missing in the saved info train_logs = self._check_default_dict_in_train_logs(train_logs) self.train_loader.load_state(self.resume_param.opt_step_dir / "resumable_states") if self.hparams.load_optimizer_states: self.accelerator.load_state(self.resume_param.accelerator_state_dir) else: # don't load the optimizer states and start with a fresh optimizer self.accelerator.load_state(self.resume_param.accelerator_state_dir, load_optimizer_states=False) validate_optim_states_are_reset(self) self.accelerator.wait_for_everyone() opt_step_is_saved = True eval_is_done = True else: # 2. non-resume (first run) train_logs = self._reset_train_logs(None) curr_opt_step = 0 curr_epoch = 0 opt_step_is_saved = False eval_is_done = False gbs_running = GlobalBatchSizeRampUpRunningParams( global_seen_samples=0, global_batch_size_current=self.hparams.global_batch_size, next_goal_samples=self.hparams.global_batch_size_ramp_up.samples, grad_acc_size_current=self.hparams.grad_acc_size, ) self.train_loader.reset_state() # rng = np.random.RandomState(np.random.MT19937(np.random.SeedSequence(self.hparams.seed))) # self.main_rng_seed = rng.get_state() # self.main_rng_seed = rng.RandomState.get_state() self.update_gas_and_gbs(gbs_running.grad_acc_size_current, gbs_running.global_batch_size_current) max_num_epochs = self.hparams.max_num_epochs try: num_batches = int(len(self.train_loader) // self.hparams.grad_acc_size) max_num_updates = min(self.hparams.max_num_opt_steps, num_batches) if max_num_epochs is not None: logger.info( "** Setting `max_num_updates` to `max_num_epochs * num_batches` since `max_num_epochs` " "was specified and `max_num_epochs * num_batches` is smaller than `max_num_updates`. **" ) max_num_updates = min(max_num_updates, max_num_epochs * num_batches) except TypeError: # For iterable datasets len(dataset) is not defined max_num_updates = self.hparams.max_num_opt_steps if self.hparams.max_num_opt_steps_this_run is not None: self.max_num_updates_this_run = min( max_num_updates, curr_opt_step + self.hparams.max_num_opt_steps_this_run ) else: self.max_num_updates_this_run = max_num_updates progress_columns = ( "[progress.description]{task.description}", BarColumn(), TaskProgressColumn(), "Time Elapsed:", TimeElapsedColumn(), "Steps Completed", MofNCompleteColumn(), ) return ( progress_columns, train_logs, max_num_epochs, max_num_updates, curr_opt_step, curr_epoch, opt_step_is_saved, eval_is_done, gbs_running, ) def _do_batch(self, batch, curr_opt_step, dataset_name=None, dataset_idx=None, validation=False): # Use effective max_num_images per batch. ie: if the max_num_images of this batch is 3, the pixel_values and image mask are truncated accordingly. # Same for max_height and max_width effective_max_num_images = max(batch["num_images"]) if effective_max_num_images > 0: images_heights = batch["pixel_attention_mask"][:, :, :, 0].sum(dim=-1) images_widths = batch["pixel_attention_mask"][:, :, 0].sum(dim=-1) effective_max_height = images_heights.max().int() effective_max_width = images_widths.max().int() batch["pixel_values"] = batch["pixel_values"][ :, :effective_max_num_images, :, :effective_max_height, :effective_max_width ] batch["pixel_attention_mask"] = batch["pixel_attention_mask"][ :, :effective_max_num_images, :effective_max_height, :effective_max_width ] else: # This case is a security check: if there are no images, then it should not appear in `batch` in the first place batch.pop("pixel_values", None) batch.pop("pixel_attention_mask", None) effective_max_num_tokens = max(batch["attention_mask"].sum(dim=-1)) batch["input_ids"] = batch["input_ids"][:, :effective_max_num_tokens] if "labels" in batch: batch["labels"] = batch["labels"][:, :effective_max_num_tokens] batch["attention_mask"] = batch["attention_mask"][:, :effective_max_num_tokens] batch = accelerate.utils.operations.send_to_device(batch, self.accelerator.device) num_images = batch["num_images"].sum() num_image_tokens = (batch["input_ids"] == self.image_token_id).sum() num_text_tokens = batch["num_text_tokens"].sum() total_tokens = batch["attention_mask"].numel() num_padding = total_tokens - batch["attention_mask"].sum() if "pixel_values" in batch: pixel_values_sum = batch["pixel_values"].sum() else: pixel_values_sum = torch.tensor(0.0, device=self.accelerator.device) image_to_text_ratio = torch.div(num_images, num_text_tokens) vl_output = self.vl_model( input_ids=batch["input_ids"], attention_mask=batch["attention_mask"], pixel_values=batch["pixel_values"] if "pixel_values" in batch else None, pixel_attention_mask=batch["pixel_attention_mask"] if "pixel_attention_mask" in batch else None, labels=batch["labels"] if "labels" in batch else batch["input_ids"], ) per_token_loss = vl_output.loss if validation: return ( per_token_loss, num_images, num_text_tokens, image_to_text_ratio, num_padding, pixel_values_sum, ) else: if self.hparams.loss_weights_per_dataset is not None: per_token_loss *= self.hparams.loss_weights_per_dataset[dataset_idx] if self.optim_param.z_loss > 0.0: logits = vl_output.logits attention_mask = batch["attention_mask"] * (1 - (batch["input_ids"] == self.image_token_id).long()) log_z = torch.logsumexp(logits, dim=-1) * attention_mask z_loss = log_z**2 z_loss = z_loss.sum() / attention_mask.sum() combined_loss = per_token_loss + self.optim_param.z_loss * z_loss else: z_loss = torch.tensor(0.0, device=self.accelerator.device) combined_loss = per_token_loss sync_gradients = self.accelerator.sync_gradients deepspeed = hasattr(self.accelerator, "deepspeed_engine_wrapped") # accelerate's deepspeed `backward` calls `engine.step`, which is a problem if we want # to investigate things before step, so override with just a backward call and then call # `engine.step` along with `optim.step` a bit lower if deepspeed: self.accelerator.deepspeed_engine_wrapped.engine.backward(combined_loss) else: self.accelerator.backward(combined_loss) if sync_gradients: self.accelerator.clip_grad_norm_(self.vl_model.parameters(), self.hparams.grad_clip) if ( len(self.hparams.train_logging_grad_param_deepspeed) != 0 and (curr_opt_step + 1) % self.hparams.train_logging_grad_param_deepspeed_opt_steps == 0 ): self._log_deepspeed_training_stats(curr_opt_step=curr_opt_step) if deepspeed: self.accelerator.deepspeed_engine_wrapped.engine.step() self.vl_optim.step() # 1. sync_gradients is used for this dirty trick: https://github.com/huggingface/m4/pull/386 # 2. since we don't accelerate prepare the lr scheduler we need to manually skip it if # optimizer skipped (otherwise accelerate would do that for us) if sync_gradients and not self.accelerator.optimizer_step_was_skipped: self.vl_scheduler.step() self.vl_optim.zero_grad(set_to_none=True) tflops_per_batch_per_gpu = self.vl_model.get_model_tflops_per_batch_per_gpu( hparams=self.hparams, data_param=getattr(self.data_param, dataset_name), tokenizer=self.tokenizer, max_num_images=effective_max_num_images, max_num_tokens=effective_max_num_tokens, ).to(self.accelerator.device) # Reset batch return ( per_token_loss, z_loss, num_images, num_image_tokens, num_text_tokens, image_to_text_ratio, num_padding, pixel_values_sum, tflops_per_batch_per_gpu, ) def _log_deepspeed_training_stats(self, curr_opt_step): if self.hparams.job_id is not None: log_stats_file = self.hparams.save_dir / "logs" / f"{self.hparams.job_id}_logs_params_grads_stats.jsonl" else: log_stats_file = self.hparams.save_dir / "logs" / "logs_params_grads_stats.jsonl" beta1 = self.optim_param.vl_optim_params["betas"][0] beta2 = self.optim_param.vl_optim_params["betas"][1] eps = self.optim_param.vl_optim_params.get("eps", 1e-8) step = self.vl_optim.optimizer.state[list(self.vl_optim.optimizer.state.keys())[0]]["step"] bias_correction_1 = 1 / (1 - beta1**step) bias_correction_2 = 1 / (1 - beta2**step) for n, lp in self.vl_model.named_parameters(): self.accelerator.wait_for_everyone() if not lp.requires_grad: continue hp = self.safe_get_full_fp32_param(lp) exp_avg = self.safe_get_full_optimizer_state(lp, "exp_avg") exp_avg_sq = self.safe_get_full_optimizer_state(lp, "exp_avg_sq") hp_grad = self.safe_get_full_grad(lp) if not self.accelerator.is_main_process: continue if exp_avg_sq is not None and exp_avg is not None: effective_update = exp_avg * bias_correction_1 / (torch.sqrt(exp_avg_sq * bias_correction_2) + eps) else: effective_update = None grad_param_logs = { "name": n, "step": step.item(), **get_stats(hp, "hp"), **get_stats(exp_avg, "exp_avg"), **get_stats(exp_avg_sq, "exp_avg_sq"), **get_stats(hp_grad, "hp_grad"), **get_stats(effective_update, "effective_update"), } grad_param_format = { "step": "", "name": "", **get_stats_format("hp"), **get_stats_format("exp_avg"), **get_stats_format("exp_avg_sq"), **get_stats_format("hp_grad"), **get_stats_format("effective_update"), } if LoggingTypes.JSONL in self.hparams.train_logging_grad_param_deepspeed: with open(log_stats_file, "a") as f: f.write(json.dumps(grad_param_logs) + "\n") if LoggingTypes.WANDB in self.hparams.train_logging_grad_param_deepspeed and self.hparams.wandb_enable: self.accelerator.log({**grad_param_logs, **self._get_additional_step_logs()}, step=curr_opt_step + 1) if LoggingTypes.PRINT in self.hparams.train_logging_grad_param_deepspeed: log = "Grads and params stats: " log += self.format_print_logs(grad_param_logs, grad_param_format) print(log) def gather_metrics( self, local_metric_list: List[Dict[str, torch.Tensor]], placeholder_tensor: torch.Tensor, reduce_op_list, ds_name_suffix: str, ) -> List[Dict[str, torch.Tensor]]: """ Collating all metrics to gather into ONE call to `torch.distributed.all_gather` instead of doing one per metric x dataset_name. """ if self.accelerator.num_processes == 1: for local_metric in local_metric_list: for ds_name, tensor in local_metric.items(): if tensor is not None: local_metric[ds_name] = tensor.item() return local_metric_list dataset_names = sorted([f"{e.value}{ds_name_suffix}" for e in DatasetNames] + ["all"]) for local_metric in local_metric_list: for ds_name in dataset_names: if local_metric[ds_name] is None: local_metric[ds_name] = placeholder_tensor.clone() collated_local_metrics = torch.stack( [torch.stack([local_metric[ds_name] for ds_name in dataset_names]) for local_metric in local_metric_list] ) # Size len(local_metric_list) x len(dataset_names) broadcasted_placeholder_tensor = torch.stack( [torch.stack([placeholder_tensor for _ in dataset_names]) for _ in local_metric_list] ) # Size len(local_metric_list) x len(dataset_names) output_objects = [broadcasted_placeholder_tensor.clone() for _ in range(self.accelerator.num_processes)] torch.distributed.all_gather(output_objects, collated_local_metrics) gathered_metrics = torch.stack( output_objects ) # Size num_processes x len(local_metric_list) x len(dataset_names) gathered_metric_list = [] for metric_idx, (_, reduce_op) in enumerate(zip(local_metric_list, reduce_op_list)): result = {} for ds_idx, ds_name in enumerate(dataset_names): metrics = gathered_metrics[:, metric_idx, ds_idx] metrics = metrics[metrics != placeholder_tensor] if metrics.numel() == 0: result[ds_name] = None else: result[ds_name] = reduce_op(metrics).item() gathered_metric_list.append(result) return gathered_metric_list def _update_logs( self, curr_opt_step, curr_epoch, global_batch_size_current, train_logs, per_token_loss, z_loss, num_tokens, num_images, num_image_tokens, image_to_text_ratio, num_padding, fwd_bwd_time, pixel_values, tflops_per_batch_per_gpu, total_energy_delta_per_gpu, dataset_name, ds_name_suffix="", ): def init_dict(): return {f"{e.value}{ds_name_suffix}": None for e in DatasetNames} local_per_token_loss = init_dict() local_z_loss = init_dict() local_num_tokens = init_dict() local_num_images = init_dict() local_num_image_tokens = init_dict() local_image_to_text_ratio = init_dict() local_fwd_bwd_time = init_dict() local_pixel_values = init_dict() local_tflops_per_batch_per_gpu = init_dict() local_total_energy_delta_per_gpu = init_dict() local_num_padding = init_dict() local_num_batches = init_dict() for key_name in [f"{dataset_name}{ds_name_suffix}", "all"]: local_per_token_loss[key_name] = per_token_loss local_z_loss[key_name] = z_loss local_num_tokens[key_name] = num_tokens local_num_images[key_name] = num_images local_num_image_tokens[key_name] = num_image_tokens local_image_to_text_ratio[key_name] = image_to_text_ratio local_fwd_bwd_time[key_name] = fwd_bwd_time local_pixel_values[key_name] = pixel_values local_tflops_per_batch_per_gpu[key_name] = tflops_per_batch_per_gpu local_total_energy_delta_per_gpu[key_name] = total_energy_delta_per_gpu local_num_padding[key_name] = num_padding local_num_batches[key_name] = torch.tensor(1.0, device=self.accelerator.device, dtype=torch.long) [ gathered_per_token_loss, gathered_z_loss, gathered_image_to_text_ratio, gathered_fwd_bwd_time, gathered_pixel_values, gathered_tflops_per_batch_per_gpu, gathered_total_energy_delta_per_gpu, ] = self.gather_metrics( local_metric_list=[ local_per_token_loss, local_z_loss, local_image_to_text_ratio, local_fwd_bwd_time, local_pixel_values, local_tflops_per_batch_per_gpu, local_total_energy_delta_per_gpu, ], reduce_op_list=[torch.sum, torch.sum, torch.mean, torch.sum, torch.sum, torch.sum, torch.sum], placeholder_tensor=self.float_placeholder_tensor, ds_name_suffix=ds_name_suffix, ) [ gathered_num_padding, gathered_num_tokens, gathered_num_batches, gathered_num_images, gathered_num_image_tokens, ] = self.gather_metrics( local_metric_list=[ local_num_padding, local_num_tokens, local_num_batches, local_num_images, local_num_image_tokens, ], reduce_op_list=[torch.sum, torch.sum, torch.sum, torch.sum, torch.sum], placeholder_tensor=self.long_placeholder_tensor, ds_name_suffix=ds_name_suffix, ) for ds_name in local_per_token_loss.keys(): for metric_name, new_value in [ ("per_token_loss_acc", gathered_per_token_loss[ds_name]), ("z_loss_acc", gathered_z_loss[ds_name]), ("num_images", gathered_num_images[ds_name]), ("num_image_tokens", gathered_num_image_tokens[ds_name]), ("num_tokens", gathered_num_tokens[ds_name]), ("num_padding", gathered_num_padding[ds_name]), ("pixel_values_sum", gathered_pixel_values[ds_name]), ("tflop_counter_since_training_logged", gathered_tflops_per_batch_per_gpu[ds_name]), ("fwd_bwd_time_since_training_logged", gathered_fwd_bwd_time[ds_name]), ("total_energy_delta_since_training_logged", gathered_total_energy_delta_per_gpu[ds_name]), ("fwd_bwd_time", gathered_fwd_bwd_time[ds_name]), ("tflop_counter", gathered_tflops_per_batch_per_gpu[ds_name]), ("num_per_device_batches_since_training_logged", gathered_num_batches[ds_name]), ("num_per_device_batches", gathered_num_batches[ds_name]), ("num_per_device_batches_in_curr_epoch", gathered_num_batches[ds_name]), ]: if new_value is None: continue if train_logs[metric_name][ds_name] is None: train_logs[metric_name][ds_name] = new_value else: train_logs[metric_name][ds_name] += new_value if gathered_image_to_text_ratio[ds_name] is not None: train_logs["image_to_text_ratio"][ds_name] = gathered_image_to_text_ratio[ds_name] if gathered_fwd_bwd_time[ds_name] is not None: train_logs["tflops_acc"][ds_name] = ( train_logs["tflop_counter"][ds_name] / train_logs["fwd_bwd_time"][ds_name] ) train_logs["num_batches_since_training_logged"]["all"] += 1 train_logs["num_batches"]["all"] += 1 train_logs["num_batches_in_curr_epoch"]["all"] += 1 train_logs["lr"] = self.vl_scheduler.get_last_lr()[0] train_logs["num_opt_steps"] = curr_opt_step train_logs["num_epochs"] = curr_epoch train_logs["global_batch_size_current"] = global_batch_size_current return train_logs def _update_datasets_states(self, dataset_idx, dataset_state): # TODO: This step will go away in future PRs. The dataloader already knows the state when it # sends it to the trainer. There is no need to send it to trainer and send it back. Let's # simplify this as well in the future self.train_loader.update_state(dataset_idx, dataset_state) def _get_additional_step_logs(self): if self.config.hparams.job_id is not None: return {"job_id": self.config.hparams.job_id, "commit": self.config.hparams.repo_commit_id} else: return {"commit": self.config.hparams.repo_commit_id} def format_print_logs(self, dict_logs, keys_known_formats, skip_keys=[]): """ compact formatting of the logs with pre-specified formatter for each log entry, plus a catch-all if new log entries are added but forgotten to be added in keys_known_formats the keys order is the one that controls how the logs are printed (py37+). even if there is no formatter there is still an empty value entry here as it tells use the order of keys. """ log = "" for key in keys_known_formats.keys(): if key in dict_logs: if isinstance(dict_logs[key], dict): for sub_key in dict_logs[key].keys(): prefix = f"{key}" if sub_key != "all": if LoggingTypes.PRINT not in self.hparams.train_logging_per_dataset_info: continue prefix += f"/{sub_key}" log += f"{prefix}: {dict_logs[key][sub_key]:{keys_known_formats[key]}} | " else: log += f"{key}: {dict_logs[key]:{keys_known_formats[key]}} | " # in case some new log entries were added that don't yet have the formatter string we dump them as is for key in set(dict_logs.keys() - set(skip_keys) - set(keys_known_formats.keys())): if key in dict_logs: log += f"{key}: {dict_logs[key]} | " return log def format_jsonl_logs(self, dict_logs): """ Similar to format_print_logs but for jsonl logs """ log = {} for key in dict_logs: # We don't want to log the accumulated values if "_acc" in key: continue elif isinstance(dict_logs[key], dict): for sub_key in dict_logs[key].keys(): prefix = f"{key}" if sub_key != "all": if LoggingTypes.JSONL not in self.hparams.train_logging_per_dataset_info: continue prefix += f"/{sub_key}" log[prefix] = dict_logs[key][sub_key] else: log[key] = dict_logs[key] return log def format_val_logs(self, val_logs, logger_type=LoggingTypes.PRINT): keys_known_formats = { "val_per_token_loss": ".4f", "val_num_images": "", "val_num_tokens": "", "val_num_padding": "", "val_image_to_text_ratio": ".4f", } if logger_type == LoggingTypes.PRINT: return self.format_print_logs(val_logs, keys_known_formats) elif logger_type == LoggingTypes.JSONL: return self.format_jsonl_logs(val_logs) else: raise ValueError(f"Unknown logger type: {logger_type}") def format_train_logs(self, train_logs, logger_type=LoggingTypes.PRINT) -> Union[str, Dict]: keys_known_formats = { "per_token_loss": ".4f", "lr": ".3E", "global_batch_size": "", "num_tokens": "", "num_images": "", "num_image_tokens": "", "num_padding": "", "fwd_bwd_time": ".1f", "image_to_text_ratio": ".4f", "num_batches": "", "num_batches_in_curr_epoch": "", "num_batches_since_training_logged": "", "num_per_device_batches": "", "num_per_device_batches_in_curr_epoch": "", "num_per_device_batches_since_training_logged": "", "tflops": ".1f", "watt/s": ".1f", "fwd_bwd_time_since_training_logged": ".1f", "num_epochs": "", "num_opt_steps": "", "z_loss": ".4f", "pixel_values_sum": ".5E", "tflop_counter": ".3E", "tflops_acc": ".1f", } # intermediary state accumulate keys skip_keys = [ "per_token_loss_acc", "z_loss_acc", "tflop_counter_since_training_logged", "num_batches_since_training_logged", "num_per_device_batches_since_training_logged", "total_energy_delta_since_training_logged", ] if logger_type == LoggingTypes.PRINT: return self.format_print_logs(train_logs, keys_known_formats, skip_keys) elif logger_type == LoggingTypes.JSONL: return self.format_jsonl_logs(train_logs) else: raise ValueError(f"Unknown logger type: {logger_type}") def _log_training(self, curr_opt_step, train_task, train_logs): for key in train_logs["per_token_loss_acc"].keys(): if train_logs["num_per_device_batches_since_training_logged"][key] is not None: train_logs["per_token_loss"][key] = ( train_logs["per_token_loss_acc"][key] / train_logs["num_per_device_batches_since_training_logged"][key] ) train_logs["z_loss"][key] = ( train_logs["z_loss_acc"][key] / train_logs["num_per_device_batches_since_training_logged"][key] ) else: train_logs["per_token_loss"][key] = None train_logs["z_loss"][key] = None if train_logs["fwd_bwd_time_since_training_logged"][key] is not None: train_logs["tflops"][key] = ( train_logs["tflop_counter_since_training_logged"][key] / train_logs["fwd_bwd_time_since_training_logged"][key] ) train_logs["watt/s"][key] = ( train_logs["total_energy_delta_since_training_logged"][key] / train_logs["fwd_bwd_time_since_training_logged"][key] ) else: train_logs["tflops"][key] = None train_logs["watt/s"][key] = None if self.accelerator.is_main_process: print_log = "" progress = f"{str(MofNCompleteColumn().render(train_task)):>12} {TaskProgressColumn().render(train_task)}" print_log += f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] iteration: {progress} | " elapsed_time = TimeElapsedColumn().render(train_task) print_log += f"elapsed time: {elapsed_time} | " print_log += self.format_train_logs(train_logs, logger_type=LoggingTypes.PRINT) # TODO: Allow mem usage to be logged according to LogginTypes passed in hparams if self.hparams.train_log_mem_usage: print_log += mem_usage_formatted(LoggingTypes.PRINT) print(print_log) jsonl_logs = { "iteration": progress.strip(), "elapsed_time": str(elapsed_time), "set": "train", } jsonl_logs.update(self.format_train_logs(train_logs, logger_type=LoggingTypes.JSONL)) if self.hparams.train_log_mem_usage: jsonl_logs.update(mem_usage_formatted(LoggingTypes.JSONL)) if self.hparams.job_id is not None: log_jsonl_file = self.hparams.save_dir / "logs" / f"{self.hparams.job_id}_logs.jsonl" else: log_jsonl_file = self.hparams.save_dir / "logs" / "logs.jsonl" log_jsonl_file.parent.mkdir(parents=True, exist_ok=True) with open(log_jsonl_file, "a") as f: f.write(json.dumps(jsonl_logs) + "\n") if self.hparams.wandb_enable: filtered_train_logs = train_logs if LoggingTypes.WANDB not in self.hparams.train_logging_per_dataset_info: filtered_train_logs = {} for key in train_logs.keys(): if isinstance(train_logs[key], dict): filtered_train_logs[key] = train_logs[key]["all"] else: filtered_train_logs[key] = train_logs[key] # remove nested None values as wandb doesn't support them filtered_train_logs = {k: v for k, v in filtered_train_logs.items() if v is not None} for k, v in filtered_train_logs.items(): if isinstance(v, dict): filtered_train_logs[k] = {k2: v2 for k2, v2 in v.items() if v2 is not None} self.accelerator.log({**filtered_train_logs, **self._get_additional_step_logs()}, step=curr_opt_step) train_logs = self._reset_train_logs(train_logs) return train_logs def _log_activations(self, curr_opt_step): if not self.activation_tracker.jsonl_stats: return if LoggingTypes.JSONL in self.hparams.train_logging_activations: if self.hparams.job_id is not None: log_activations_filename = ( self.hparams.save_dir / "logs" / f"{self.hparams.job_id}_logs_activations.jsonl" ) else: log_activations_filename = self.hparams.save_dir / "logs" / "logs_activations.jsonl" self.activation_tracker.dump_stats(log_activations_filename, curr_opt_step=curr_opt_step) # if LoggingTypes.WANDB in self.hparams.train_logging_activations and self.hparams.wandb_enable: # for stats in self.activation_tracker.jsonl_stats: # self.accelerator.log({**stats, **self._get_additional_step_logs()}, step=curr_opt_step) if LoggingTypes.PRINT in self.hparams.train_logging_activations: for stats in self.activation_tracker.jsonl_stats: stats["step"] = curr_opt_step activation_format = { k: "" if ("nonzero" in k or "step" in k or "name" in k or "type" in k or "batches" in k) else "e" for k in stats.keys() } log = "Activation stats: " log += self.format_print_logs(stats, activation_format) print(log) self.activation_tracker.reset_jsonl_stats() def _check_kill_switch(self): if self.hparams.kill_switch_path is not None and self.hparams.kill_switch_path.exists(): self.kill_switch_activated = True def _check_jz_time_and_memory(self, curr_opt_step): # From https://github.com/wandb/wandb/blob/9c777265f8cea1eaeb0407dd37ab889aeea81114/wandb/sdk/internal/stats.py#L263 if self.accelerator.is_local_main_process: self.memory_value = torch.tensor(psutil.virtual_memory().percent).to(self.accelerator.device) else: self.memory_value = torch.tensor(0.0).to(self.accelerator.device) self.accelerator.wait_for_everyone() memory_value_max = self.accelerator.gather(self.memory_value) memory_value_max = memory_value_max.max().item() self.memory_explosion = memory_value_max >= _MEMORY_EXPLOSION_THRESHOLD if self.hparams.jz_job_time_sec is not None: if self.accelerator.is_main_process: overall_time = time.time() - self.hparams.jz_start_time self.jz_training_time_over[0] = overall_time >= self.hparams.jz_job_time_sec self.accelerator.wait_for_everyone() accelerate.utils.broadcast_object_list(self.jz_training_time_over) if self.accelerator.is_main_process and self.hparams.wandb_enable: system_metrics_logs = self._get_system_metrics_logs(memory_value_max) self.accelerator.log({**system_metrics_logs, **self._get_additional_step_logs()}, step=curr_opt_step) def _check_sigterm_signal(self): if self.sigterm_listener.kill_now: self.sigterm_signal_received = True def _save( self, train_logs, curr_opt_step, curr_epoch, gbs_running, ): self.accelerator.wait_for_everyone() # create directory and file names self.last_opt_step_dir = self.hparams.save_dir / f"opt_step-{curr_opt_step}" # Make directory for this step self.last_opt_step_dir.mkdir(parents=True, exist_ok=True) self.train_loader.save_state(self.last_opt_step_dir / "resumable_states") # XXX: why is there a hardcoded accelerator_state path? should be coming from config, no? self.accelerator.save_state(self.last_opt_step_dir / "accelerator_state") if self.accelerator.is_main_process: # Save model and accelerator state unwrapped_model = self.accelerator.unwrap_model(self.vl_model) # fix the model class name to be of VLOOOM type the first time it's saved unwrapped_model.config.architectures = [unwrapped_model.__class__.__name__] # deepspeed doesn't need the overhead of gathering the model from all gpus if not is_deepspeed_zero3_used(): if self.hparams.use_lora: unwrapped_model.save_pretrained(self.last_opt_step_dir / "unwrapped_adapter") # Manual unloading with a simple PeftMixin to avoid having to deal with PeftModel state dict base_model = lora_unload(copy.deepcopy(unwrapped_model)) # Save pretrained with _hf_peft_config_loaded=True will save the adapters only. So we set it manually to False base_model._hf_peft_config_loaded = False base_model.save_pretrained(self.last_opt_step_dir / "unwrapped_model") del base_model else: unwrapped_model.save_pretrained(self.last_opt_step_dir / "unwrapped_model") else: # For deepspeed, `save_checkpoint` done by accelerate takes care of saving the model in a # special format per gpu which on resume will be used to load the model - so we don't need to # save pytorch state_dict separately, which can be costly or impossible if there not enough CPU RAM. # We only need to save the config unwrapped_model.config.save_pretrained( self.last_opt_step_dir / "unwrapped_model", ) if self.hparams.use_lora: unwrapped_model.peft_config["default"].save_pretrained( self.last_opt_step_dir / "unwrapped_adapter", ) # Save tokenizer directly into the same dir self.tokenizer.save_pretrained( self.last_opt_step_dir / "tokenizer", ) # infos to resume run at this step data = { "train_logs": train_logs, "wandb_run_id": self.hparams.wandb_run_id, "seed": self.hparams.seed, "resume_opt_step": curr_opt_step, "resume_epoch": curr_epoch, "gbs_running": gbs_running, } with open(self.last_opt_step_dir / "resume_run_infos.json", "w") as fp: # json.dump(data, fp, indent=2) json.dump(data, fp, indent=2, cls=JSONEncoderForDataclasses) self.accelerator.wait_for_everyone() if self.accelerator.is_main_process: # mark this checkpoint as finished - needed for async slurm jobs like s3 uploader latest_path_finished = self.last_opt_step_dir / "finished-saving" latest_path_finished.touch() # mark which is latest saved checkpoint for correct resume latest_path = self.hparams.save_dir / "latest_opt_step_dir" with open(latest_path, "w") as fd: fd.write(str(self.last_opt_step_dir)) logger.info(f"** Saving finished at `{self.last_opt_step_dir}` **") if self.accelerator.is_main_process and self.hparams.upload_to_s3: # We keep around the last checkpoint (which was saved just above) locally, and delete the previous to last one. locally_present_saved_steps_inds = [ int(os.path.split(dir)[-1].split("opt_step-")[-1]) for dir in self.hparams.save_dir.iterdir() if (dir.is_dir() and "opt_step" in str(dir)) ] if len(locally_present_saved_steps_inds) >= 2: previous_to_last_saved_step = sorted(locally_present_saved_steps_inds)[-2] previous_to_last_folder = f"opt_step-{previous_to_last_saved_step}" else: previous_to_last_folder = "" # Subprocess command is inspired from https://stackoverflow.com/questions/5772873/python-spawn-off-a-child-subprocess-detach-and-exit/64145368#64145368 # `stdout` and `stderr` are supressed to avoid polluting the logs with the output of the command. cmd = ( str(Path(__file__).resolve().parents[2]) + "/experiments/pretraining/vloom/common/sync_and_upload.sh", os.path.split(self.hparams.save_dir)[0], os.path.split(self.hparams.save_dir)[1], previous_to_last_folder, "opt_step-" + str(curr_opt_step), ) subprocess.Popen(cmd, start_new_session=True, stdout=subprocess.DEVNULL, stderr=subprocess.STDOUT) def _save_batch(self, batch, curr_idx): dir_path = self.hparams.save_dir / "batches" dir_path.mkdir(parents=True, exist_ok=True) with open(dir_path / f"batch_idx_{curr_idx}_proc_{self.accelerator.process_index}.pkl", "wb") as file: pickle.dump(batch, file) def _check_if_training_is_over(self, curr_opt_step, max_num_updates): self._check_kill_switch() self._check_jz_time_and_memory(curr_opt_step=curr_opt_step) self._check_sigterm_signal() finished_training = True if self.kill_switch_activated: logger.info(f"** Kill switch activated (Don't forget to remove {self.hparams.kill_switch_path}) **") elif self.sigterm_signal_received: logger.info("** SIGTERM signal received. Please restart training **") elif self.jz_training_time_over[0]: logger.info("** Training time is over **") elif self.memory_explosion: logger.info("** CPU memory is close to explosion. Please restart training **") elif curr_opt_step >= max_num_updates: logger.info("** Maximum number of steps has been reached for this training **") elif curr_opt_step >= self.max_num_updates_this_run: logger.info("** Maximum number of steps has been reached for this run **") else: finished_training = False return finished_training def _check_if_training_is_over_and_maybe_save_model( self, curr_opt_step, curr_epoch, gbs_running, max_num_updates, train_logs, opt_step_is_saved, ): # check if should finish the training finished_training = self._check_if_training_is_over(curr_opt_step, max_num_updates) # save the model # 1. either because it's a scheduled saving is_opt_to_save_model = curr_opt_step != 0 and curr_opt_step % self.hparams.train_saving_opt_steps == 0 # 2. it was requested via the save switch if self.hparams.save_switch_path is not None and self.hparams.save_switch_path.exists(): is_opt_to_save_model = True logger.info("** Save switch activated - forcing checkpoint save **") self.hparams.save_switch_path.unlink() if not opt_step_is_saved and (finished_training or is_opt_to_save_model): self._save( train_logs, curr_opt_step, curr_epoch, gbs_running, ) opt_step_is_saved = True return finished_training, opt_step_is_saved def _reset_train_logs(self, train_logs): if train_logs is None: train_logs = {} for metric_name, default_value_fn in METRICS_TO_DEFAULT_VALUE_FN.items(): if metric_name in METRICS_TO_RESET_AFTER_LOGGING: continue train_logs[metric_name] = default_value_fn() # reset counters that get zeroed after every log event train_logs.update( {metric_name: METRICS_TO_DEFAULT_VALUE_FN[metric_name]() for metric_name in METRICS_TO_RESET_AFTER_LOGGING} ) return train_logs def _check_default_dict_in_train_logs(self, train_logs): for metric_name, default_value_fn in METRICS_TO_DEFAULT_VALUE_FN.items(): value = default_value_fn() # if the metric was initialized as a defaultdict, we need to convert it again to a defaultdict if not isinstance(value, defaultdict): continue value.update(train_logs.get(metric_name, {})) train_logs[metric_name] = value return train_logs def _end_of_epoch_reset_train_logs(self, train_logs): if train_logs is None: raise ValueError("`train_logs` should not be `None` at the end of an epoch.") train_logs.update( { metric_name: METRICS_TO_DEFAULT_VALUE_FN[metric_name]() for metric_name in ["num_batches_in_curr_epoch", "num_per_device_batches_in_curr_epoch"] } ) return train_logs def _do_validation(self, progress, curr_opt_step): try: val_len = len(self.val_loader) except TypeError: val_len = None self.vl_model.eval() val_loop = progress.add_task( f"[cyan]Validation step-{curr_opt_step}", disable=not self.accelerator.is_main_process, total=val_len, visible=False, ) curr_val_task = progress.tasks[-1] starter_dict = {e.value: 0 for e in DatasetNames} starter_dict["all"] = 0 val_per_token_loss_acc = copy.deepcopy(starter_dict) val_steps = copy.deepcopy(starter_dict) val_num_images = copy.deepcopy(starter_dict) val_num_tokens = copy.deepcopy(starter_dict) val_num_padding = copy.deepcopy(starter_dict) val_image_to_text_ratio = {e.value: [] for e in DatasetNames} val_image_to_text_ratio["all"] = [] for _, dataset_name, _, batch in self.val_loader: with torch.no_grad(): ( curr_val_per_token_loss, curr_val_num_images, curr_val_num_tokens, curr_val_image_to_text_ratio, curr_val_num_padding, _, ) = self._do_batch(batch, curr_opt_step, validation=True) val_per_token_loss_acc["all"] += curr_val_per_token_loss val_num_images["all"] += curr_val_num_images val_num_tokens["all"] += curr_val_num_tokens val_num_padding["all"] += curr_val_num_padding val_image_to_text_ratio["all"].append(curr_val_image_to_text_ratio) val_steps["all"] += 1 val_per_token_loss_acc[dataset_name] += curr_val_per_token_loss val_num_images[dataset_name] += curr_val_num_images val_num_tokens[dataset_name] += curr_val_num_tokens val_num_padding[dataset_name] += curr_val_num_padding val_image_to_text_ratio[dataset_name].append(curr_val_image_to_text_ratio) val_steps[dataset_name] += 1 progress.update(val_loop, advance=1) if ( curr_val_task.completed % self.hparams.val_inline_logging_opt_steps == 0 and self.accelerator.is_main_process ): logger.info( "Validation" f" step-{curr_opt_step} state:{TaskProgressColumn().render(curr_val_task)} Time" f" Elapsed: {TimeElapsedColumn().render(curr_val_task)} Steps" f" Completed:{MofNCompleteColumn().render(curr_val_task)}" ) self.vl_model.train() return ( val_steps, val_per_token_loss_acc, val_num_images, val_num_tokens, val_image_to_text_ratio, val_num_padding, ) def _log_validation( self, val_steps, curr_opt_step, val_per_token_loss_acc, val_num_images, val_num_tokens, val_image_to_text_ratio, val_num_padding, ): def convert_to_tensor(x): if not torch.is_tensor(x): return torch.tensor(x, device=self.accelerator.device) return x for key, value in val_image_to_text_ratio.items(): if len(value) != 0: val_image_to_text_ratio[key] = sum(value) / len(value) else: val_image_to_text_ratio[key] = 0.0 gathered_val_per_token_loss = {} gathered_val_num_images = {} gathered_val_num_tokens = {} gathered_val_num_padding = {} gathered_val_image_to_text_ratio = {} gathered_val_steps = {} for key in val_image_to_text_ratio.keys(): ( gathered_val_per_token_loss[key], gathered_val_num_images[key], gathered_val_num_tokens[key], gathered_val_num_padding[key], gathered_val_image_to_text_ratio[key], gathered_val_steps[key], ) = self.accelerator.gather( ( convert_to_tensor(val_per_token_loss_acc[key]), convert_to_tensor(val_num_images[key]), convert_to_tensor(val_num_tokens[key]), convert_to_tensor(val_num_padding[key]), convert_to_tensor(val_image_to_text_ratio[key]), convert_to_tensor(val_steps[key]), ) ) # No overall steps so we should skip this if gathered_val_steps[key].sum() == 0: gathered_val_per_token_loss.pop(key) gathered_val_num_images.pop(key) gathered_val_num_tokens.pop(key) gathered_val_num_padding.pop(key) gathered_val_image_to_text_ratio.pop(key) continue gathered_val_per_token_loss[key] = ( gathered_val_per_token_loss[key].sum().item() / gathered_val_steps[key].sum().item() ) gathered_val_num_images[key] = gathered_val_num_images[key].sum().item() gathered_val_num_tokens[key] = gathered_val_num_tokens[key].sum().item() gathered_val_num_padding[key] = gathered_val_num_padding[key].sum().item() gathered_val_image_to_text_ratio[key] = ( gathered_val_image_to_text_ratio[key][gathered_val_steps[key] != 0.0].mean().item() ) val_logs = { "val_per_token_loss": gathered_val_per_token_loss, "val_num_images": gathered_val_num_images, "val_num_tokens": gathered_val_num_tokens, "val_num_padding": gathered_val_num_padding, "val_image_to_text_ratio": gathered_val_image_to_text_ratio, } if self.accelerator.is_main_process: print(f"Validation logs: {self.format_val_logs(val_logs, LoggingTypes.PRINT)}") jsonl_logs = {"current step": curr_opt_step, "set": "validation"} jsonl_logs.update(self.format_val_logs(val_logs, LoggingTypes.JSONL)) if self.hparams.job_id is not None: log_jsonl_file = self.hparams.save_dir / "logs" / f"{self.hparams.job_id}_logs.jsonl" else: log_jsonl_file = self.hparams.save_dir / "logs" / "logs.jsonl" with open(log_jsonl_file, "a") as f: f.write(json.dumps(jsonl_logs) + "\n") if self.hparams.wandb_enable: self.accelerator.log({**val_logs, **self._get_additional_step_logs()}, step=curr_opt_step) def train(self, maybe_torch_profile_scheduler=None): # timing_break_down = self.accelerator.is_main_process and self.hparams.timing_break_down if self.accelerator.is_main_process: logger.info(f"** Global main process pid={os.getpid()} **") elif self.accelerator.is_local_main_process: logger.info(f"** Local main process pid={os.getpid()} **") # -------------------- # Set-up everything needed for training # -------------------- ( progress_columns, train_logs, max_num_epochs, max_num_updates, curr_opt_step, curr_epoch, opt_step_is_saved, eval_is_done, gbs_running, ) = self._set_up_training() # -------------------- # Training loop # -------------------- self.vl_model.train() pynvml_handle = pynmvl_handle(self.accelerator) with Progress(*progress_columns, refresh_per_second=5, disable=True) as progress: progress_bar = progress.add_task( "[red]Training", disable=not self.accelerator.is_main_process, total=max_num_updates, visible=False ) train_task = progress.tasks[-1] progress.update(progress_bar, advance=curr_opt_step) finished_training = False training_logged = True timer = DeviceAgnosticTimer() timer.start() if self.hparams.timing_break_down: self.accelerator.wait_for_everyone() if self.accelerator.is_main_process: timer2 = Timer() time_deltas = {} while not finished_training: self.train_loader.set_epoch(curr_epoch) # Handle resume based on `realtime_processing` if self.hparams.resume_run: if not self.data_param.realtime_processing: # When preparing, `accelerate` has a bug that overwrites the top-level `sampler`... but this is aesthetic! # The "actual" sampler that the DataLoader uses is the one that's tucked inside the `batch_sampler` # attribute when "single-process" (world_size = 1), or the `batch_sampler.batch_sampler` when # "multi-process" (world_size > 1) -- both of which point to our specialized ResumableSampler! # # This is pretty annoying and nuanced; should PR into `accelerate` to fix this... logger.warning("NOT realtime processing has not been extensively tested yet") if self.accelerator.num_processes == 1: # TODO; Fix this self.train_loader.batch_sampler.sampler.set_state(self.train_loader.get_resume_state(0)) else: # TODO :: This is actually broken and not respected by `accelerate` - fails! # self.train_loader.batch_sampler.batch_sampler.sampler.set_state( # self.resumable_state.get_resume_state() # ) raise NotImplementedError("Map Dataset Resume w/ DDP not yet implemented!") else: self.train_loader.load_resume_states() if self.hparams.timing_break_down: self.accelerator.wait_for_everyone() if self.accelerator.is_main_process: timer2.start() for curr_idx, (dataset_idx, dataset_name, dataset_state, batch) in enumerate(self.train_loader): # -------------------- # Check if the training is over and if so may be save the model before training batch # -------------------- if self.hparams.timing_break_down: self.accelerator.wait_for_everyone() if self.accelerator.is_main_process: time_deltas["dl"] = timer2.delta() finished_training, opt_step_is_saved = self._check_if_training_is_over_and_maybe_save_model( curr_opt_step, curr_epoch, gbs_running, max_num_updates, train_logs, opt_step_is_saved, ) if finished_training: break # -------------------- # Activate/deactivate hooks for logging activations or not # -------------------- # We are logging everything at `curr_opt_step`, but `curr_opt_step` is incremented a few lines later, so activating # the activation tracking hooks based on `curr_opt_step + 1`. See `_log_activations` for more details. if self.activation_tracker: if (curr_opt_step + 1) % self.hparams.train_logging_activations_opt_steps == 0 and ( curr_idx + 1 ) % self.hparams.grad_acc_size == 0: self.activation_tracker.activate_hooks() else: self.activation_tracker.deactivate_hooks() if ( self.hparams.save_batch_max_idx is not None and self.hparams.save_batch_min_idx is not None and curr_idx <= self.hparams.save_batch_max_idx and curr_idx >= self.hparams.save_batch_min_idx ): self._save_batch(batch, curr_idx) # right before fwd-bwd-step if self.hparams.timing_break_down: self.accelerator.wait_for_everyone() if self.accelerator.is_main_process: time_deltas["between_dl_fwd_bwd"] = timer2.delta() total_energy_start = pynvml_get_total_energy_in_joules(pynvml_handle) with self.accelerator.accumulate(self.vl_model): ( per_token_loss, z_loss, num_images, num_image_tokens, num_tokens, image_to_text_ratio, num_padding, pixel_values_sum, tflops_per_batch_per_gpu, ) = self._do_batch( batch, curr_opt_step=curr_opt_step, dataset_name=dataset_name, dataset_idx=dataset_idx, validation=False, ) # right after fwd-bwd-step if self.hparams.timing_break_down: self.accelerator.wait_for_everyone() if self.accelerator.is_main_process: time_deltas["fwd-bwd-step"] = timer2.delta() fwd_bwd_time = timer.stop() fwd_bwd_time = torch.tensor(fwd_bwd_time, device=self.accelerator.device) total_energy_delta_per_gpu = torch.tensor( pynvml_get_total_energy_in_joules(pynvml_handle) - total_energy_start, device=self.accelerator.device, ) if (curr_idx + 1) % self.hparams.grad_acc_size == 0: curr_opt_step += 1 opt_step_is_saved, eval_is_done, training_logged = False, False, False progress.update(progress_bar, advance=1) # -------------------- # Update logs # -------------------- train_logs = self._update_logs( curr_opt_step, curr_epoch, gbs_running.global_batch_size_current, train_logs, per_token_loss, z_loss, num_tokens, num_images, num_image_tokens, image_to_text_ratio, num_padding, fwd_bwd_time, pixel_values_sum, tflops_per_batch_per_gpu, total_energy_delta_per_gpu, dataset_name, self.hparams.train_logging_per_dataset_suffix, ) timer = DeviceAgnosticTimer() timer.start() # -------------------- # Update datasets states # -------------------- self._update_datasets_states(dataset_idx, dataset_state) # -------------------- # Log training infos # -------------------- if curr_opt_step % self.hparams.train_logging_opt_steps == 0 and not training_logged: train_logs = self._log_training(curr_opt_step, train_task, train_logs) training_logged = True # -------------------- # Log activations # -------------------- if self.activation_tracker: batch_idx = train_logs["num_batches"]["all"] self.activation_tracker.fill_in_batch_idx(batch_idx=batch_idx) if curr_opt_step % self.hparams.train_logging_activations_opt_steps == 0: self._log_activations(curr_opt_step=curr_opt_step) # --------------------------- # Global batch size ramp up # --------------------------- # # This logic needs to happen after the batch has been processed and results # logged, but before the model is saved for resume, so that the updated ramup up # variables will have the correct values on resume gbs_running.global_seen_samples += self.accelerator.num_processes * self.hparams.batch_size_per_gpu if ( self.hparams.global_batch_size_ramp_up.start is not None and self.hparams.global_batch_size_ramp_up.finish > gbs_running.global_batch_size_current and gbs_running.global_seen_samples >= gbs_running.next_goal_samples ): gbs_running.next_goal_samples += self.hparams.global_batch_size_ramp_up.samples gbs_running.global_batch_size_current += self.hparams.global_batch_size_ramp_up.increment gbs_running.grad_acc_size_current = int( gbs_running.global_batch_size_current / (self.hparams.batch_size_per_gpu * self.accelerator.num_processes) ) self.update_gas_and_gbs( gbs_running.grad_acc_size_current, gbs_running.global_batch_size_current ) # -------------------- # Check if the training is over and if so may be save the model before validation # -------------------- finished_training, opt_step_is_saved = self._check_if_training_is_over_and_maybe_save_model( curr_opt_step, curr_epoch, gbs_running, max_num_updates, train_logs, opt_step_is_saved, ) if finished_training: break if self.hparams.timing_break_down: self.accelerator.wait_for_everyone() if self.accelerator.is_main_process: time_deltas["post_fwd"] = timer2.delta() time_deltas["iteration"] = timer2.elapsed() # -------------------- # Validation loop # -------------------- if ( self.config.hparams.do_validation and not eval_is_done and curr_opt_step != 0 and curr_opt_step % self.hparams.val_logging_opt_steps == 0 ): if self.hparams.timing_break_down: self.accelerator.wait_for_everyone() if self.accelerator.is_main_process: timer3 = Timer() timer3.start() gc.collect() if self.accelerator.is_main_process and self.hparams.wandb_enable: wandb.unwatch(self.dummy_module) if self.activation_tracker: self.activation_tracker.is_eval() logger.info("** Starting validation **") ( val_steps, val_per_token_loss_acc, val_num_images, val_num_tokens, val_image_to_text_ratio, val_num_padding, ) = self._do_validation(progress, curr_opt_step) # -------------------- # Log validation infos # -------------------- self._log_validation( val_steps, curr_opt_step, val_per_token_loss_acc, val_num_images, val_num_tokens, val_image_to_text_ratio, val_num_padding, ) eval_is_done = True logger.info("** Finished validation **") if self.accelerator.is_main_process and self.hparams.wandb_enable: wandb.watch( self.dummy_module, log="all", log_freq=self.hparams.wandb_log_freq * self.hparams.grad_acc_size, idx=0, ) if self.activation_tracker: self.activation_tracker.is_train() gc.collect() self.accelerator.wait_for_everyone() if self.hparams.timing_break_down: self.accelerator.wait_for_everyone() if self.accelerator.is_main_process: print(f"[TIME] Validation: {format_secs_to_time(timer3.stop())}") # restart timer from zero to avoid accounting for validation timer = DeviceAgnosticTimer() timer.start() if self.hparams.timing_break_down: self.accelerator.wait_for_everyone() if self.accelerator.is_main_process: # Finalize print( f'[TIME] Iteration {format_secs_to_sec_fractions(time_deltas["iteration"])}:' f' {format_secs_to_sec_fractions(time_deltas["dl"]):>6} dl |' f' {format_secs_to_sec_fractions(time_deltas["between_dl_fwd_bwd"]):>6} between' f' dl/fwd-bwd | {format_secs_to_sec_fractions(time_deltas["fwd-bwd-step"]):>6} fwd/bwd' f' | {format_secs_to_sec_fractions(time_deltas["post_fwd"]):>6} post' ) # restart for __iter__ timer2.stop() timer2.start() if maybe_torch_profile_scheduler is not None and self.accelerator.is_main_process: maybe_torch_profile_scheduler.step() if not finished_training: curr_epoch += 1 train_logs = self._end_of_epoch_reset_train_logs(train_logs) self.train_loader.reset_state() if curr_epoch == max_num_epochs: self._save( train_logs, curr_opt_step, curr_epoch, gbs_running, ) finished_training = True logger.info("** Maximum number of epochs has been reached **") break if self.hparams.wandb_enable: self.accelerator.end_training() return train_logs def _get_system_metrics_logs(self, memory_value_max): return {"memory_max_over_nodes": memory_value_max}