# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations from typing import Callable, TYPE_CHECKING, Optional import torch import wandb from lipforcing.callbacks.callback import Callback from lipforcing.utils.basic_utils import get_batch_size_total from lipforcing.utils.distributed import synchronize, is_rank0 import lipforcing.utils.logging_utils as logger if TYPE_CHECKING: from lipforcing.methods import FastGenModel class EMACallback(Callback): def __init__( self, type: str = "constant", # params for type=constant beta: float = 0.9999, # params for type=power gamma: float = 16.97, # params for type=halflife ema_halflife_kimg: float = 500, ema_rampup_ratio: Optional[float] = 0.05, start_iter: int = 0, ema_name: str = "ema", batch_size: int = 1, # overwritten by self.config if it exists fsdp: bool = False, # overwritten by self.config if it exists ): self.type = type self.beta = beta self.gamma = gamma self.ema_halflife_kimg = ema_halflife_kimg self.ema_rampup_ratio = ema_rampup_ratio self.start_iter = start_iter self.ema_name = ema_name self.batch_size = batch_size self._is_fsdp = fsdp self._enabled = True def on_app_begin(self) -> None: if hasattr(self, "config"): # override using config self._is_fsdp = self.config.trainer.fsdp self.batch_size = get_batch_size_total(self.config) def on_model_init_end( self, model: FastGenModel | torch.nn.parallel.DistributedDataParallel, iteration: int = 0 ) -> None: # Unwrap DDP if needed to access the original model's attributes if hasattr(model, "module"): model = model.module # check ema initialization ema = getattr(model, self.ema_name, None) if ema is None: self._enabled = False logger.info(f"EMA {self.ema_name} is not enabled, skipping callback.") return assert ema.training is False, f"EMA {self.ema_name} should be in eval mode" for name, p_net in ema.named_parameters(): assert not p_net.requires_grad, f"EMA parameter {name} should not require gradients" def _total_iteration(self, model: FastGenModel, iteration: int) -> int: if hasattr(model, "resume_iter"): assert isinstance(model.resume_iter, int) iteration = iteration + model.resume_iter return iteration def _power_function_beta(self, iteration): beta = (1 - 1 / iteration) ** (self.gamma + 1) return beta def _get_cur_nimg(self, iteration): cur_nimg = iteration * self.batch_size return self.batch_size, cur_nimg def _halflife_beta(self, iteration): ema_halflife_nimg = self.ema_halflife_kimg * 1000 batch_size, cur_nimg = self._get_cur_nimg(iteration) if self.ema_rampup_ratio is not None: ema_halflife_nimg = min(ema_halflife_nimg, cur_nimg * self.ema_rampup_ratio) ema_beta = 0.5 ** (batch_size / max(ema_halflife_nimg, 1e-8)) return ema_beta def on_training_step_end( self, model: FastGenModel, data_batch: dict[str, torch.Tensor], output_batch: dict[str, torch.Tensor | Callable], loss_dict: dict[str, torch.Tensor], iteration: int = 0, ) -> None: del data_batch, output_batch, loss_dict # Check if EMA is enabled if not self._enabled: return # Get total iteration and skip if before start_iter total_iteration = self._total_iteration(model, iteration) if total_iteration < self.start_iter: return elif total_iteration == self.start_iter: logger.info(f"Starting to update {self.ema_name} at iteration {total_iteration}.") if self.type == "constant": beta = self.beta elif self.type == "power": beta = self._power_function_beta(total_iteration) elif self.type == "halflife": beta = self._halflife_beta(total_iteration) else: raise ValueError(f"Invalid {self.ema_name} type: {self.type}") with torch.no_grad(): ema = getattr(model, self.ema_name) ema_state_dict = ema.state_dict() for name, p_net in model.net.named_parameters(): if self._is_fsdp and hasattr(p_net, "full_tensor"): # Gather the full tensor from all ranks if using FSDP with DTensor # When CPU offloading is enabled, we need to move to CUDA first because # full_tensor() performs an all_gather which requires a CUDA backend if p_net.device.type == "cpu": # Move local shard to CUDA, gather, then the result stays on CUDA # which is fine since we'll copy to EMA (which handles device placement) full_tensor = p_net.to("cuda").full_tensor() else: full_tensor = p_net.full_tensor() else: full_tensor = p_net # Strip checkpoint wrapper prefix if present (EMA doesn't have checkpointing) ema_name = name.replace("_checkpoint_wrapped_module.", "") # Cast to EMA dtype and device for lerp_ compatibility if ema_name in ema_state_dict: ema_param = ema_state_dict[ema_name] if total_iteration == self.start_iter: # re-initialize EMA parameter ema_param.copy_(full_tensor.to(device=ema_param.device, dtype=ema_param.dtype)) else: # interpolate EMA parameter ema_param.lerp_(full_tensor.to(device=ema_param.device, dtype=ema_param.dtype), 1.0 - beta) elif iteration == 1: # only warn on first iteration if parameter is not found logger.warning(f"EMA parameter {ema_name} not found in EMA state dict, skipping update.") # FSDP2 doesn't shard buffers, so we can just copy them for name, p_net in model.net.named_buffers(): if name in ema_state_dict: ema_param = ema_state_dict[name] ema_param.copy_(p_net.to(device=ema_param.device, dtype=ema_param.dtype)) elif iteration == 1: # only warn on first iteration if buffer is not found logger.warning(f"EMA buffer {name} not found in EMA state dict, skipping update.") if hasattr(self, "config"): # only wandb log when config exists if iteration % self.config.trainer.logging_iter == 0 and is_rank0(): if wandb.run: wandb.log({f"ema/{self.ema_name}_beta": beta}, step=iteration) synchronize()