# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import datetime import functools import inspect import os import atexit import shutil import signal import sys import threading import time import traceback import torch import torch.distributed as dist import torch.utils.data from imaginaire.utils.profiling import maybe_enable_memory_snapshot, maybe_enable_profiling # NOTE(bvh): Evaluator import deferred to avoid pulling in vidar/lpips/skimage # when metrics are disabled (metrics=None in experiment config). from tqdm import tqdm try: from megatron.core import parallel_state USE_MEGATRON = True except ImportError: USE_MEGATRON = False print("Megatron-core is not installed.") from imaginaire.lazy_config import LazyConfig, instantiate from imaginaire.model import ImaginaireModel from imaginaire.utils import callback, distributed, log, misc # NOTE(bvh): use cosmos_predict2 checkpointer (subclass), not imaginaire base class from cosmos_predict2.checkpointer import Checkpointer # from imaginaire.utils.checkpointer import Checkpointer from custom.utils.config import Config from custom.wandb import WandbLogger from custom.utils.utils import get_local_root from custom.utils.aws import aws_s3_sync DEFAULT_INSTANCE_TYPE = 'DGX' def gather_loss(loss): """Gather losses from all ranks for logging""" world_size = dist.get_world_size() gathered_loss = [torch.zeros_like(loss) for _ in range(world_size)] dist.all_gather(gathered_loss, loss) gathered_loss_stacked = torch.stack(gathered_loss, 0) return gathered_loss, gathered_loss_stacked.mean() class ImaginaireTrainer: """The base trainer class of Imaginaire. All trainers in Imaginaire should inherit ImaginaireTrainer. It contains the basic functionality for model training (particularly suited for large-scale training), including data parallel (DDP/FSDP), model weight average (EMA), mixed-precision training (fp16/bf16). Attributes: checkpointer (Checkpointer): checkpointer object to save/load model weights and optimizer states. training_timer (misc.Timer): Timer object to time code blocks and functions. """ def __init__(self, config): """Constructor of the trainer. Args: config (Config): The config object for the Imaginaire codebase. """ super().__init__() if not config.job.processed: # Change run name if environment variable is provided (useful for sagemaker) if os.getenv("JOB_NAME", None): config.job.__dict__['name'] = os.getenv("JOB_NAME") # Keep track of the original name, e.g. for visuals. # NOTE(bvh): Config has been frozen by train.py, so we need to use __dict__ to modify it. config.job.__dict__['undated_name'] = config.job.name # Add datetime to the job name (two ways) if os.getenv("PREPEND_DATETIME", config.job.prepend_datetime): # Spot-safe: reuse the submit-time datetime (RUN_DATETIME env) so a restarted # job keeps the same run dir and can resume; fall back to now() off-SageMaker. dt = os.getenv("RUN_DATETIME") or f'{datetime.datetime.now():%m-%d-%H-%M}' config.job.__dict__['name'] = f'{dt}_{config.job.name}' # NOTE(bvh): shorter wandb DISPLAY name (date only, no time). The full dated `name` # above stays the output-folder name + (sanitized) wandb run id for uniqueness. config.job.__dict__['display_name'] = f'{dt[:5]}_{config.job.undated_name}' if os.getenv("APPEND_DATETIME", config.job.append_datetime): now = datetime.datetime.now() config.job.__dict__['name'] += f'--{now.year}y{now.month}m{now.day}d-{now.hour}h{now.minute}m{now.second}s' config.job.__dict__['processed'] = True # Pass job info to model config, e.g. for visuals. config.model.config.__dict__['job'] = config.job self.config = config # Set up the distributed computing environment. with misc.timer("init_distributed"): distributed.init() # Set up parallel states. if hasattr(config.model, "context_parallel_size"): if config.model_parallel.context_parallel_size > 1: raise ValueError( "Both config.model.context_parallel_size and config.model_parallel.context_parallel_size are set. " "config.model.context_parallel_size is deprecated. Please only set config.model_parallel.context_parallel_size." ) else: log.critical( "Using deprecated config.model.context_parallel_size. Please use config.model_parallel.context_parallel_size instead." ) config.model_parallel.context_parallel_size = config.model.context_parallel_size if USE_MEGATRON: if ( "create_gloo_process_groups" in inspect.signature(parallel_state.initialize_model_parallel).parameters ): parallel_state.initialize_model_parallel( pipeline_model_parallel_size=config.model_parallel.pipeline_model_parallel_size, tensor_model_parallel_size=config.model_parallel.tensor_model_parallel_size, context_parallel_size=config.model_parallel.context_parallel_size, create_gloo_process_groups=False, ) else: parallel_state.initialize_model_parallel( pipeline_model_parallel_size=config.model_parallel.pipeline_model_parallel_size, tensor_model_parallel_size=config.model_parallel.tensor_model_parallel_size, context_parallel_size=config.model_parallel.context_parallel_size, ) # `config.model_parallel.sequence_parallel` is a bool that indicates whether to use sequence parallelism. # It is not part of the original `parallel_state` API, so we need to set it manually. parallel_state.sequence_parallel = config.model_parallel.sequence_parallel if parallel_state.sequence_parallel: os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1" # Create the local job directory, save the config file, and pipe to a local log. local_root = get_local_root(config.job) if distributed.is_rank0(): os.makedirs(f"{local_root}/{config.job.path_local}", exist_ok=True) # Save the config as .pkl for reproducibility. LazyConfig.save_pkl(config, f"{local_root}/{config.job.path_local}/config.pkl") # Save the config as .yaml for reading or parsing experiment hyperparameters. LazyConfig.save_yaml(config, f"{local_root}/{config.job.path_local}/config.yaml") dist.barrier() log.init_loguru_file(f"{local_root}/{config.job.path_local}/stdout.log") # Tee stderr to a per-rank log file (Python via TeeWriter, C-level via fd dup2). stderr_dir = f"{local_root}/{config.job.path_local}/stderr" os.makedirs(stderr_dir, exist_ok=True) stderr_path = f"{stderr_dir}/rank{dist.get_rank()}.log" self._stderr_file = open(stderr_path, "a", buffering=1) self._original_stderr_fd = os.dup(2) # save original fd before redirect os.dup2(self._stderr_file.fileno(), 2) # C-level stderr -> file self._original_stderr = os.fdopen(self._original_stderr_fd, "w", closefd=False) sys.stderr = misc.TeeWriter(self._original_stderr, self._stderr_file) # TODO(bvh): how to get date & time to also print in stderr? # On crash, dump tail of stderr log to stdout (which still goes to console). # NOTE(bvh): C-level stderr (fd 2) goes to file only after dup2 above, so # crash tracebacks from torch.distributed are invisible in console. This atexit # handler prints them on exit. Only fires for rank 0 (the one that actually crashes); # SIGTERM'd ranks don't run atexit. def _dump_stderr_on_exit(): try: self._stderr_file.flush() with open(stderr_path, 'r') as f: lines = f.readlines() tail = lines[-50:] if len(lines) > 50 else lines print(f'\n===== STDERR LOG TAIL ({stderr_path}) =====', flush=True) for line in tail: print(line, end='', flush=True) print('===== END STDERR LOG TAIL =====\n', flush=True) except Exception: pass atexit.register(_dump_stderr_on_exit) if distributed.is_local_rank0(): # Print important environment variables and the effective config. log.info("Config:\n" + config.pretty_print(use_color=True)) misc.print_environ_variables(["TORCH_HOME"]) # Set the random seed. If multi-GPU, different ranks are set with different seeds. misc.set_random_seed(seed=config.trainer.seed, by_rank=True) # Initialize cuDNN. torch.backends.cudnn.deterministic = config.trainer.cudnn.deterministic torch.backends.cudnn.benchmark = config.trainer.cudnn.benchmark # Floating-point precision settings. torch.backends.cudnn.allow_tf32 = torch.backends.cuda.matmul.allow_tf32 = True # Initialize the callback functions. self.callbacks = callback.CallBackGroup(config=config, trainer=self) # Initialize the model checkpointer. if config.checkpoint.type is None: self.checkpointer = Checkpointer(config.checkpoint, config.job, callbacks=self.callbacks) else: self.checkpointer: Checkpointer = instantiate( config.checkpoint.type, config.checkpoint, config.job, callbacks=self.callbacks ) # Initialize the timer for speed benchmarking. self.training_timer = misc.TrainingTimer() # Send a TimeoutError if a training step takes over timeout_period seconds. signal.signal(signal.SIGALRM, functools.partial(misc.timeout_handler, config.trainer.timeout_period)) # type: ignore # Initialize wandb logger if config.model.config.wandb.enabled and distributed.is_rank0(): # wandb DISPLAY name = short date_name (falls back to full name when no datetime). config.model.config.wandb.name = getattr(config.job, 'display_name', None) or config.job.name # wandb run id = full (dated) job name, sanitized -> unique + traceable, independent of # the display label. Passed at WandbLogger construction below since the wandb OmegaConf # node is a closed struct (cannot add a new 'id' key to it by assignment). wandb_run_id = ''.join(c if (c.isalnum() or c in '_-') else '-' for c in config.job.name) config.model.config.wandb.dir = f"{local_root}/wandb/{self.config.job.path_local}" params = Config(**{key: Config(**config.params_for_wandb[key]) for key in config.params_for_wandb}) params.model.num_gpus = dist.get_world_size() params.model.instance_type = os.getenv('INSTANCE_TYPE', DEFAULT_INSTANCE_TYPE) params.model.instance_count = os.getenv('INSTANCE_COUNT', '1') params.model.queue = os.getenv('QUEUE', 'local') # Log git info for reproducibility try: import subprocess git_hash = subprocess.check_output( ['git', 'rev-parse', 'HEAD'], stderr=subprocess.DEVNULL).decode().strip() git_branch = subprocess.check_output( ['git', 'rev-parse', '--abbrev-ref', 'HEAD'], stderr=subprocess.DEVNULL).decode().strip() git_dirty = subprocess.check_output( ['git', 'status', '--porcelain'], stderr=subprocess.DEVNULL).decode().strip() params.git = Config( commit=git_hash, branch=git_branch, dirty=bool(git_dirty), ) except Exception: log.warning(f'Failed to get git info') pass os.makedirs(config.model.config.wandb.dir, exist_ok=True) self.wandb = WandbLogger(Config(**config.model.config.wandb, id=wandb_run_id)) # Start wandb logger self.wandb.log_config(params) # Log configuration parameters else: self.wandb = None self.time = None self.num_iters = 0 self.sec_iters = 0 self.local_path = f"{local_root}/{self.config.job.path_local}" s3_root = self.config.job.s3_root self.s3_path = f"{s3_root}/{self.config.job.path_local}" if s3_root else '' self._sync_thread = None self.outstanding_ckpt_save = False self.outstanding_validate = False self.phase = None # Set evaluator (lazy import to avoid vidar/lpips/skimage when metrics disabled) if config.metrics is not None: from custom.eval.evaluator import Evaluator self.evaluator = Evaluator(config.metrics) else: self.evaluator = None self.unroll_val = config.metrics is not None and config.metrics.get('unroll', False) def reset_logs(self, model, dataloader_train, iteration, num_samples, optimizer): self.phase = None if self.wandb is None: return wandb_logs = dict( global_step=iteration, consumed_samples=num_samples, ) for key, val in optimizer.param_groups[0].items(): if key in ['lr','weight_decay']: wandb_logs[key] = val self.time = time.time() return wandb_logs def store_logs_train(self, wandb_logs, output): self.phase = 'train' losses = dict() for key, val in output.items(): if key.startswith('loss') or key.endswith('loss'): val, val_mean = gather_loss(val) for i, v in enumerate(val): losses[f'losses/{i}/{key}'] = v.item() losses[f'losses/{key}'] = val_mean.item() avg_loss = losses['losses/loss'] if wandb_logs is None: return wandb_logs, avg_loss wandb_logs.update(losses) self.sec_iters = (self.sec_iters * self.num_iters + (time.time() - self.time)) / (self.num_iters + 1) self.num_iters += 1 wandb_logs['iters_sec'] = (1.0 / self.sec_iters) wandb_logs['samples_sec'] = (1.0 / self.sec_iters) * self.samples_step return wandb_logs, avg_loss def _log_dataloader_summary(self, phase, dataloader, model): dp_world = model.data_parallel_size def _log_one(name, dl, rank_info=''): ds = dl.dataset bs = dl.batch_size n_samples = len(ds) n_iters = (n_samples + bs * dp_world - 1) // (bs * dp_world) log.success( f'[{phase}] {name}: {n_samples:,} samples, ' f'batch_size={bs}, dp_world={dp_world}, ' f'=> {n_iters:,} batched iterations' f'{rank_info}' ) if isinstance(dataloader, dict): for key, dl in dataloader.items(): _log_one(key, dl, rank_info=' (all ranks)') else: ds = dataloader.dataset # RankDataset: log from ALL ranks so we see per-rank dataset assignment from custom.dataset.rank_dataset import RankDataset if isinstance(ds, RankDataset): rank_info = f' (rank {dist.get_rank()} => {ds.dataset.dataset_dir})' log.success( f'[{phase}] rank {dist.get_rank()}/{dp_world}: ' f'{ds.dataset.dataset_dir}, ' f'batch_size={dataloader.batch_size}', rank0_only=False, ) else: ds_name = str(getattr(ds, 'dataset_dir', type(ds).__name__)) _log_one(ds_name, dataloader, rank_info=f' (rank {dist.get_rank()})') def store_logs_val(self, wandb_logs, output): self.phase = 'val' if wandb_logs is None: return for key, val in output.items(): wandb_logs[key] = val return wandb_logs def _sync_to_s3(self, blocking=False, cleanup=False): """Sync entire run directory to S3. S3 becomes a mirror of local. Non-blocking by default (runs on a background thread). :param blocking: If True, wait for sync to complete before returning. :param cleanup: If True, delete local checkpoint .pt files after sync (for keep_local=False). Only use on final sync to avoid deleting un-synced checkpoints. """ if not self.s3_path or not distributed.is_local_rank0(): return # Wait for any in-progress checkpoint save and previous sync self.checkpointer.finalize() if self._sync_thread is not None: self._sync_thread.join() self._sync_thread = None def _do_sync(): with misc.timer('S3 sync'): aws_s3_sync( self.local_path, self.s3_path, extras='--acl bucket-owner-full-control', robust=True, ) if cleanup and not self.config.checkpoint.keep_local: self._cleanup_local_checkpoints() if blocking: _do_sync() else: self._sync_thread = threading.Thread(target=_do_sync, daemon=False) self._sync_thread.start() def _cleanup_local_checkpoints(self): """Delete local checkpoint .pt files after S3 sync (keep_local=False).""" # NOTE: checkpoints now live directly under self.local_path if not os.path.isdir(self.local_path): return for folder in ['model', 'optim', 'scheduler', 'trainer', 'extra']: folder_path = os.path.join(self.local_path, folder) if not os.path.isdir(folder_path): continue for f in os.listdir(folder_path): if f.endswith('.pt'): try: os.remove(os.path.join(folder_path, f)) except OSError: log.warning(f'Failed to delete checkpoint file: {os.path.join(folder_path, f)}') pass def upload_logs(self, wandb_logs): # Sync entire run directory to S3 after validation (includes visuals, metrics, logs) # NOTE(bvh): _sync_to_s3 already guards on is_local_rank0 internally if self.phase == 'val': self._sync_to_s3() distributed.barrier() # keep ranks in lockstep here (rank0's S3 sync is non-blocking) if wandb_logs is None: return if self.phase is None: return if self.wandb is not None: self.wandb.experiment.log(wandb_logs) def train( self, model: ImaginaireModel, dataloader_train: torch.utils.data.DataLoader, dataloader_val: torch.utils.data.DataLoader, ) -> None: """The training function. Args: model (ImaginaireModel): The PyTorch model. dataloader_train (torch.utils.data.DataLoader): The training data loader. dataloader_val (torch.utils.data.DataLoader): The validation data loader. """ # Leaving this for backward compability for now, but we can think about moving this to model.on_train_start for all models. model = model.to("cuda", memory_format=self.config.trainer.memory_format) # type: ignore model.on_train_start(self.config.trainer.memory_format) # Initialize the optimizer, scheduler, and grad_scaler. self.callbacks.on_optimizer_init_start() optimizer, scheduler = model.init_optimizer_scheduler(self.config.optimizer, self.config.scheduler) grad_scaler = torch.amp.GradScaler("cuda", **self.config.trainer.grad_scaler_args) self.callbacks.on_optimizer_init_end() # Spot resume: local disk is ephemeral, so restore the prior run dir from S3 (each # node) before loading. No-op for a fresh run (empty s3_path). Pairs with _sync_to_s3. restore_ok = True if self.s3_path and distributed.is_local_rank0(): restore_ok = aws_s3_sync(self.s3_path, self.local_path, robust=True) # Abort if any rank's restore failed: silently continuing would train from scratch and # then _sync_to_s3 would overwrite the intended resume state in S3. MIN over all ranks. if self.s3_path and dist.is_initialized() and distributed.get_world_size() > 1: flag = torch.tensor([1 if restore_ok else 0], device="cuda") dist.all_reduce(flag, op=dist.ReduceOp.MIN) restore_ok = bool(flag.item()) if self.s3_path and not restore_ok: raise RuntimeError( f"Spot-resume: S3 restore failed ({self.s3_path} -> {self.local_path}); aborting to " f"avoid training from scratch and overwriting the checkpoint in S3.") distributed.barrier() # The checkpointer only resumes when checkpoint.resume is set; wire it from the restored # run dir's latest_checkpoint.txt so a spot restart resumes (re-read each time = newest). # A run-dir checkpoint also outranks checkpoint.pretrained: it means THIS run already # trained past the pretrained init, so a spot restart must not re-apply it. if self.checkpointer.resume is None: latest = self.checkpointer._read_latest_checkpoint_file() if latest is not None: self.checkpointer.resume = os.path.join(self.local_path, "model", latest) log.critical(f"Spot-resume: auto-resuming from {self.checkpointer.resume}") # Load the model checkpoint and get the starting iteration number. iteration, num_samples = self.checkpointer.load(model, optimizer, scheduler, grad_scaler) grad_accum_iter = 0 log.critical(f"Distributed parallelism mode: {self.config.trainer.distributed_parallelism}") if self.config.trainer.distributed_parallelism == "ddp": # Create a DDP model wrapper. model_ddp = distributed.parallel_model_wrapper(self.config.trainer.ddp, model) elif self.config.trainer.distributed_parallelism == "fsdp": model_ddp = model else: raise ValueError(f"Unknown distributed parallelism mode: {self.config.trainer.distributed_parallelism}") log.info("Starting training...") model.unroll_val = self.unroll_val self.callbacks.on_train_start(model, iteration=iteration) # Log dataset and dataloader summary. self._log_dataloader_summary('train', dataloader_train, model) self._log_dataloader_summary('val', dataloader_val, model) # Initial validation. self.samples_step = dataloader_train.batch_size * model.data_parallel_size # skip_first_validation: False = run before training, True = skip, # 'delay' = run after 10 steps, 'delayX' = run after X steps skip_first_val = getattr(self.config.trainer, 'skip_first_validation', False) if isinstance(skip_first_val, bool): delay_first_val = False first_val_iter = 10 else: delay_first_val = skip_first_val.startswith('delay') first_val_iter = 10 if len(skip_first_val) <= 5 else int(skip_first_val[5:]) # NOTE(bvh): validation_iter <= 0 disables validation entirely. # This is the canonical way to disable validation from CLI overrides, # since SageMaker passes all values as strings (bool 'False' is truthy). val_iter = int(self.config.trainer.validation_iter) _run_validation = self.config.trainer.run_validation and val_iter > 0 if not _run_validation: log.info(f'Validation disabled (run_validation={self.config.trainer.run_validation}, ' f'validation_iter={val_iter})') if _run_validation and not skip_first_val: # log.info(f'[RANK{distributed.get_rank()}] pre-training validation begin', rank0_only=False) # DEBUG wandb_logs = self.reset_logs(model, dataloader_train, iteration, num_samples, optimizer) logs = self.validate(model, dataloader_val, iteration=iteration) # log.info(f'[RANK{distributed.get_rank()}] pre-training validation end', rank0_only=False) # DEBUG wandb_logs = self.store_logs_val(wandb_logs, logs) # Upload wandb logs self.upload_logs(wandb_logs) # Eval-only mode: if max_iter <= 1, skip training loop and checkpoint saving entirely. if self.config.trainer.max_iter <= 1: log.info(f'max_iter={self.config.trainer.max_iter} <= 1: eval-only mode, skipping training loop.') self.callbacks.on_train_end(model, iteration=iteration) distributed.barrier() self.callbacks.on_app_end() return _end_training = False # TODO (vitor): Move this to Any4D/VidarModel train_directives = self.config.model.config.train_directives with maybe_enable_profiling(self.config, global_step=iteration) as torch_profiler, maybe_enable_memory_snapshot( self.config, global_step=iteration ) as memory_profiler: while True: dataloader_train_iter = iter(dataloader_train) progress = range(iteration, self.config.trainer.max_iter) if distributed.is_local_rank0(): progress = tqdm(progress, initial=iteration + 1, total=self.config.trainer.max_iter, ncols=192) for _ in progress: iteration += 1 num_samples += self.samples_step # Any4D: Avoid accidentally skipping checkpoint save because of unlucky # iteration numbers and gradient accumulation. # There is an option to save immediately for error checking as well. if iteration % self.config.checkpoint.save_iter == 0 or \ iteration == self.config.checkpoint.early_sanity_check: log.debug(f'Flagging outstanding_ckpt_save for iteration: {iteration}') if iteration == self.config.checkpoint.early_sanity_check: log.info(f'Will save checkpoint soon as early sanity check') self.outstanding_ckpt_save = True # Any4D: Same for validation. if _run_validation and \ (iteration % val_iter == 0 or iteration == self.config.trainer.max_iter - 1 or (delay_first_val and iteration == first_val_iter)): log.debug(f'Flagging outstanding_validate for iteration: {iteration}') self.outstanding_validate = True self.callbacks.on_before_dataloading(iteration) dataload_start_time = time.time() try: with self.training_timer("dataloader_train"): data_batch = next(dataloader_train_iter) except StopIteration: log.warning(f'StopIteration: Reached end of training dataloader at iteration: {iteration}') break finally: dataload_time = time.time() - dataload_start_time self.callbacks.on_after_dataloading(iteration) # Track AnyData per-sample fetch latency to measure FastFile dataloader impact. sample_fetch_time = None if isinstance(data_batch, dict) and 'sample_fetch_time' in data_batch: fetch_times = data_batch['sample_fetch_time'] if isinstance(fetch_times, torch.Tensor): sample_fetch_time = fetch_times.float().mean().item() elif isinstance(fetch_times, (list, tuple)) and fetch_times: sample_fetch_time = sum(float(t) for t in fetch_times) / len(fetch_times) else: sample_fetch_time = float(fetch_times) # Reset wandb logs wandb_logs = self.reset_logs( model, dataloader_train, iteration, num_samples, optimizer) # If max_iter is reached, exit the training loop. if iteration >= self.config.trainer.max_iter: _end_training = True break # Move all tensors in the data batch to GPU device. data_batch = misc.to(data_batch, device="cuda") # The actual training step. self.callbacks.on_training_step_start(model, data_batch, iteration=iteration) self.callbacks.on_training_step_batch_start(model, data_batch, iteration=iteration) if not model.training: model_ddp.train() assert model_ddp.training, "model_ddp is not in training mode." assert model.training, "model is not in training mode." start_time = time.time() # log.debug(f'[RANK{distributed.get_rank()}] train step {iteration} begin, ' # DEBUG # f'batch keys: {list(data_batch.keys())[:5]}', rank0_only=False) # DEBUG output_batch, loss, grad_accum_iter = self.training_step( model_ddp, optimizer, scheduler, grad_scaler, data_batch, iteration=iteration, grad_accum_iter=grad_accum_iter, directives=train_directives, ) # log.debug(f'[RANK{distributed.get_rank()}] train step {iteration} end, ' # DEBUG # f'loss={loss.item():.4f}', rank0_only=False) # DEBUG end_time = time.time() step_time = end_time - start_time # Save custom wandb logs if self.wandb is not None and 'wandb_logs' in output_batch: wandb_logs.update(output_batch['wandb_logs']) wandb_logs, avg_loss = self.store_logs_train(wandb_logs, output_batch) self.callbacks.on_training_step_batch_end( model, data_batch, output_batch, loss, iteration=iteration ) # If the gradients are still being accumulated, continue to load the next training batch. if grad_accum_iter >= self.config.trainer.grad_accum_iter: log.error(f'Unexpected grad_accum_iter: {grad_accum_iter} ' f'>= max_grad_accum_iter: {self.config.trainer.grad_accum_iter}') if grad_accum_iter != 0: continue # Save checkpoint. if self.outstanding_ckpt_save: # Skip saves in data-only debug mode -- optimizer is a stub and the # ckpt would only contain a dummy param. data_only_save = getattr(getattr(self.config.model, 'config', None), 'data_only_disable_model', False) if not data_only_save: log.debug(f'Saving checkpoint for iteration: {iteration}...') self.checkpointer.save( model, optimizer, scheduler, grad_scaler, iteration=iteration, num_samples=num_samples, ) self._sync_to_s3() distributed.barrier() # keep ranks in lockstep here (rank0's S3 sync is non-blocking) self.outstanding_ckpt_save = False self.callbacks.on_training_step_end(model, data_batch, output_batch, loss, iteration=iteration) # Validation. if self.outstanding_validate: logs = self.validate(model, dataloader_val, iteration=iteration) wandb_logs = self.store_logs_val(wandb_logs, logs) self.outstanding_validate = False # This iteration is successful; reset the timeout signal. signal.alarm(self.config.trainer.timeout_period) if torch_profiler: torch_profiler.step() if memory_profiler: memory_profiler.step() # Upload wandb logs self.upload_logs(wandb_logs) # DEBUG: Track VRAM every 5 training iters # if iteration % 5 == 0 and distributed.is_local_rank0(): # DEBUG # from custom.utils.debug import vram_snapshot # DEBUG # vram_snapshot(f'train iter={iteration} ({data_batch.get("dset_name", "?") if isinstance(data_batch, dict) else "?"})') # DEBUG # DEBUG: Per-rank tensor census every 100 train iters with growth diff. # if iteration > 0 and iteration % 100 == 0: # DEBUG # from custom.utils.debug import gpu_tensor_census # DEBUG # dset = data_batch.get("dset_name", "?") if isinstance(data_batch, dict) else "?" # DEBUG # gpu_tensor_census( # DEBUG # f'train iter={iteration} ({dset})', # DEBUG # top_n=10, rank0_only=False, diff=True, growth_top_n=8) # DEBUG # Update progress bar if distributed.is_local_rank0(): fetch_time = 0.0 if sample_fetch_time is None else sample_fetch_time progress.set_description( f'##### Loss: {avg_loss:.5f} | Samples: {num_samples} | ' f'Data: {dataload_time:.0f}s | Fetch: {fetch_time:.2f}s | Step: {step_time:.2f}s #####' ) if _end_training: break log.success("Done with training.") signal.alarm(0) # Cancel timeout alarm before cleanup. data_only_final = getattr(getattr(self.config.model, 'config', None), 'data_only_disable_model', False) if iteration % self.config.checkpoint.save_iter != 0 and not data_only_final: log.info(f'Saving final checkpoint for iteration: {iteration}...') try: self.checkpointer.save( model, optimizer, scheduler, grad_scaler, iteration=iteration, num_samples=num_samples, ) except Exception as e: log.error(f'Failed to save final checkpoint: {e}') import traceback log.error(traceback.format_exc()) # Final sync: blocking, ensures everything (logs, metrics, checkpoints) is on S3. # cleanup=True: safe to delete local .pt files now since this is the last sync. self._sync_to_s3(blocking=True, cleanup=True) self.callbacks.on_train_end(model, iteration=iteration) self.checkpointer.finalize() # joins save_thread distributed.barrier() self.callbacks.on_app_end() def training_step( self, model_ddp: torch.nn.Module | distributed.DistributedDataParallel, optimizer: torch.optim.Optimizer, scheduler: torch.optim.lr_scheduler.LRScheduler, grad_scaler: torch.amp.GradScaler, data: dict[str, torch.Tensor], iteration: int = 0, grad_accum_iter: int = 0, directives: dict = None, ) -> tuple[dict[str, torch.Tensor], torch.Tensor, int]: """The training step. Args: model_ddp (torch.nn.Module | distributed.DistributedDataParallel): The model with a DDP wrapper or, the bare module, depending on whether distributed training is enabled or not. optimizer (torch.optim.Optimizer): The model optimizer. scheduler (torch.optim.lr_scheduler.LRScheduler): The optimization scheduler. grad_scaler (torch.amp.GradScaler): The gradient scaler (for mixed precision training). data (dict[str, torch.Tensor]): Data batch (dictionary of tensors). iteration (int): Current iteration number. grad_accum_iter (int): Number of gradient accumulation iterations. Returns: output (dict[str, torch.Tensor]): The model output from the training data batch (dictionary of tensors). loss (torch.Tensor): The total loss of the training data batch. """ # Any4D: data-only debug mode -- skip backward / optimizer / scheduler updates. data_only = getattr(getattr(self.config.model, 'config', None), 'data_only_disable_model', False) # Only let DDP sync gradient at the last iteration of the gradient accumulation window with distributed.ddp_sync_grad(model_ddp, grad_accum_iter == self.config.trainer.grad_accum_iter - 1): self.callbacks.on_before_forward(iteration=iteration) # log.debug(f'[RANK{distributed.get_rank()}] forward begin', rank0_only=False) # DEBUG with self.training_timer("forward"): output_batch, loss = model_ddp.training_step( data, iteration, self.local_path, directives=directives) # log.debug(f'[RANK{distributed.get_rank()}] forward end, loss={loss.item():.4f}', rank0_only=False) # DEBUG self.callbacks.on_after_forward(iteration=iteration) self.callbacks.on_before_backward(model_ddp, loss, iteration=iteration) # NOTE(bvh): For debugging, enable through TORCH_ANOMALY flag: if os.environ.get('TORCH_ANOMALY', False): torch.autograd.set_detect_anomaly(True) # log.debug(f'[RANK{distributed.get_rank()}] backward begin', rank0_only=False) # DEBUG with self.training_timer("backward"): if not data_only: loss_scaled = grad_scaler.scale(loss / self.config.trainer.grad_accum_iter) loss_scaled.backward() #@IgnoreException if self.config.trainer.distributed_parallelism == "ddp": model_ddp.module.on_after_backward() else: model_ddp.on_after_backward() # log.debug(f'[RANK{distributed.get_rank()}] backward end', rank0_only=False) # DEBUG self.callbacks.on_after_backward(model_ddp, iteration=iteration) grad_accum_iter += 1 if grad_accum_iter >= self.config.trainer.grad_accum_iter: with self.training_timer("optimizer_step"): if not data_only: self.callbacks.on_before_optimizer_step( model_ddp, optimizer, scheduler, grad_scaler, iteration=iteration ) grad_scaler.step(optimizer) grad_scaler.update() scheduler.step() self.callbacks.on_before_zero_grad(model_ddp, optimizer, scheduler, iteration=iteration) if self.config.trainer.distributed_parallelism == "ddp": model_ddp.module.on_before_zero_grad(optimizer, scheduler, iteration=iteration) else: model_ddp.on_before_zero_grad(optimizer, scheduler, iteration=iteration) optimizer.zero_grad(set_to_none=True) grad_accum_iter = 0 return output_batch, loss, grad_accum_iter @torch.no_grad() def validate( self, model: ImaginaireModel, dataloader_val: torch.utils.data.DataLoader, iteration: int = 0, ) -> dict[str, float]: """Validate on the full validation dataset (one or multiple). Args: model (ImaginaireModel): The PyTorch model. dataloader_val (torch.utils.data.DataLoader): The validation data loader. iteration (int): Current iteration number. """ log.info(f'Starting outer validation loop at iteration: {iteration}, ' f'with max_val_iter: {self.config.trainer.max_val_iter}', rank0_only=False) model.eval() wandb_logs = dict() if hasattr(model, 'on_validation_start'): model.on_validation_start(iteration=iteration) # TODO (vitor): Move this to Any4D/VidarModel val_directives = self.config.model.config.val_directives # has_multiple_val = isinstance(list(val_directives.values())[0], dict) # NOTE(bvh): ^ this line is bad for several reasons, please fix asap @vitor has_multiple_val = False # NOTE(yuzeng): dropped for now to defer to main's disabled state. # per-dataloader dual validation (single-view and view-synth) is needed for vid2vid multi-view. # TODO(yuzeng): add a better approach later (this heuristic revives the multi-val path bvh flagged. # Kept here for reference: # has_multiple_val = ( # isinstance(dataloader_val, dict) # and hasattr(val_directives, 'keys') # and set(map(str, val_directives.keys())) == set(map(str, dataloader_val.keys())) # ) # True only when val_directives is keyed by the SAME keys as dataloader_val # # (robust to OmegaConf DictConfig, for which isinstance(_, dict) is False). # Evaluate on the full validation set. with model.pipe.ema_scope(context="Validation", is_cpu=False): num_errors = 0 if isinstance(dataloader_val, dict): # If there are multiple datasets, they are stored as a dict for dataloader_key_i, dataloader_val_i in dataloader_val.items(): val_directives_i = val_directives[dataloader_key_i] if has_multiple_val else val_directives try: wandb_logs = self.validate_single( wandb_logs, model, dataloader_key=dataloader_key_i, dataloader_val=dataloader_val_i, iteration=iteration, directives=val_directives_i) except Exception as e: log.error(f'Error validating dataloader_key: {dataloader_key_i}: {e}') log.error(f"Stack trace:\n{traceback.format_exc()}") num_errors += 1 if num_errors >= len(dataloader_val) // 4 + 1: log.error(f'Too many errors, raising exception...') raise e else: # Single validation dataset wandb_logs = self.validate_single( wandb_logs, model, dataloader_key=None, dataloader_val=dataloader_val, iteration=iteration, directives=val_directives) return wandb_logs def validate_single(self, wandb_logs, model, dataloader_key, dataloader_val, iteration, directives): self.callbacks.on_validation_start(model, dataloader_val, iteration=iteration) # Prepare evaluation metrics if self.evaluator is not None: self.evaluator.prepare(dataloader_val) # Synchronize iteration count across ranks to prevent FSDP deadlocks. local_len = len(dataloader_val) max_val_iter = self.config.trainer.max_val_iter if max_val_iter is not None: local_len = min(local_len, max_val_iter) global_max = torch.tensor([local_len], device='cuda') dist.all_reduce(global_max, op=dist.ReduceOp.MAX) global_max = int(global_max.item()) if distributed.is_rank0(): progress = tqdm(range(global_max), ncols=128, total=global_max, unit_scale=dist.get_world_size(), desc=dataloader_key, ) else: progress = range(global_max) log.info(f'Starting inner validation loop for dataloader_key: {dataloader_key}, ' f'with validation set size: {len(dataloader_val)}, ' f'local_len: {local_len}, global_max: {global_max}', rank0_only=False) # from custom.utils.debug import vram_snapshot, vram_delta, gpu_tensor_census # DEBUG # vram_before_val, _ = vram_snapshot(f'val_start ({dataloader_key})', rank0_only=False) # DEBUG # Reset the dataset's iterator state before every val round so we # deterministically re-consume the same leading shard subsequence. # No-op on map-style datasets; rebuilds the pipeline iterator on # webbed datasets. See .reset_iter() for details. dataloader_val.dataset.reset_iter() log.info(f'Reset iterator for val dataset ' f'({getattr(dataloader_val.dataset, "dataset_dir", "?")})', rank0_only=False) dataloader_iter = iter(dataloader_val) last_batch = None # Pre-flight: try to get the first batch on each rank. Webdatasets may # not yield data to all ranks (insufficient shards). Ranks without data # receive a broadcast batch from rank 0 so they can participate in FSDP # forward passes (allgather requires matching tensor shapes from same dataset). first_batch = None try: first_batch = next(dataloader_iter) except (StopIteration, RuntimeError) as e: # NOTE(bvh): AnyDataset is robust but StopIteration is intentionally passed through immediately log.warning(f'Pre-flight batch fetch failed for {dataloader_key}: {e}', rank0_only=False) has_data = torch.tensor([1 if first_batch is not None else 0], device='cuda') n_with_data = has_data.clone() dist.all_reduce(n_with_data, op=dist.ReduceOp.SUM) n_with_data = int(n_with_data.item()) if n_with_data == 0: log.warning( f'Skipping validation for {dataloader_key}: 0/{dist.get_world_size()} ' f'ranks have data (webdataset fully exhausted).', rank0_only=False) return wandb_logs # Broadcast a real batch from a rank that has data to ranks that don't. # This ensures all ranks have a same-dataset batch for FSDP forward pass. if n_with_data < dist.get_world_size(): # Find the lowest rank that has data and broadcast from it. src_rank = torch.tensor([dist.get_rank() if first_batch is not None else dist.get_world_size()], device='cuda') dist.all_reduce(src_rank, op=dist.ReduceOp.MIN) src_rank = int(src_rank.item()) batch_list = [first_batch] dist.broadcast_object_list(batch_list, src=src_rank) if first_batch is None: first_batch = batch_list[0] log.warning( f'Received broadcast batch for {dataloader_key} from rank {src_rank} ' f'({n_with_data}/{dist.get_world_size()} ranks had data). ' f'This rank will run padding forward passes only.', rank0_only=False) no_own_data = (has_data.item() == 0) else: no_own_data = False try: for val_iter in progress: # Get next batch, or reuse last one as padding to keep FSDP in sync. is_padding = no_own_data # ranks with no own data are always padding if not no_own_data and val_iter < local_len: if val_iter == 0: data_batch = first_batch # already fetched in pre-flight else: try: data_batch = next(dataloader_iter) except (StopIteration, RuntimeError): log.warning(f'StopIteration at val_iter={val_iter} ' f'(local_len={local_len})', rank0_only=False) data_batch = last_batch is_padding = True last_batch = data_batch else: if val_iter == 0: data_batch = first_batch # broadcast batch for ranks with no data else: data_batch = last_batch is_padding = True last_batch = data_batch data_batch = misc.to(data_batch, device="cuda") data_batch['dl_key'] = dataloader_key data_batch['true_val_size'] = len(dataloader_val.dataset) # before DistributedSampler padding if not is_padding: self.callbacks.on_validation_step_start(model, data_batch, iteration=iteration) log.info(f'Running validation step for dataloader_key: {dataloader_key} ' f'iteration: {iteration}, ' f'val_iter: {val_iter}' f'{" (padding)" if is_padding else ""}', rank0_only=False) # vram_pre_step, _ = vram_snapshot(f'val_iter={val_iter} pre_step ({dataloader_key})', rank0_only=False) # DEBUG (output_batch, loss) = model.validation_step(data_batch, iteration, dataloader_key=dataloader_key, local_path=self.local_path, directives=directives, val_iter=val_iter, is_padding=is_padding) # vram_delta(f'val_iter={val_iter} post_step ({dataloader_key})', vram_pre_step, rank0_only=False) # DEBUG if not is_padding: # Save custom wandb logs if self.wandb is not None and 'wandb_logs' in output_batch: wandb_logs.update(output_batch['wandb_logs']) # Handle metrics if self.evaluator is not None: self.evaluator.step(data_batch, output_batch) self.callbacks.on_validation_step_end(model, data_batch, output_batch, loss, iteration=iteration) # DEBUG: Explicitly free large GPU tensors from output_batch # del output_batch, loss, data_batch # DEBUG # torch.cuda.empty_cache() # DEBUG finally: # vram_delta(f'val_end ({dataloader_key})', vram_before_val, rank0_only=False) # DEBUG # gpu_tensor_census(f'val_end ({dataloader_key})', rank0_only=False) # DEBUG # Always run collective ops to prevent other ranks from deadlocking. distributed.barrier() # Pre-validation synchronization self.callbacks.on_validation_end(model, iteration=iteration) if hasattr(model, 'on_validation_end'): # We have to call this on all ranks to ensure metrics are gathered, # otherwise NCCL desynchronization will occur. wandb_logs_end = model.on_validation_end(iteration=iteration, local_path=self.local_path) wandb_logs.update(wandb_logs_end) # Handle metrics if self.evaluator is not None: self.evaluator.finish(wandb_logs, dataloader_key) distributed.barrier() # Post-validation synchronization # Save evaluator's numeric metrics/ entries to a separate JSON if distributed.is_rank0() and self.local_path is not None: import json eval_metrics = {} for k, v in wandb_logs.items(): if k.startswith('metrics/') and isinstance(v, (int, float, torch.Tensor)): eval_metrics[k] = v.item() if isinstance(v, torch.Tensor) else v if eval_metrics: iter_tag = f'_iter{iteration:06d}' if iteration >= 0 else '' eval_path = os.path.join(self.local_path, f'val_metrics_summary_vitor{iter_tag}.json') try: os.makedirs(self.local_path, exist_ok=True) with open(eval_path, 'w') as f: json.dump(eval_metrics, f, indent=2) log.info(f'Saved evaluator metrics to {eval_path}') except Exception as e: log.warning(f'Could not save evaluator metrics JSON: {e}') # NOTE(bvh): release per-rank caching allocator blocks back to driver # after each val run, regardless of manual_gc config. Caching allocator # is per-process so every rank must call this on its own GPU. torch.cuda.empty_cache() return wandb_logs