Skip to content

ddp_posthoc_ema

Official nitrous PostHocEMA lifecycle for native DDP training.

Classes

fastvideo.train.callbacks.ddp_posthoc_ema.DDPPostHocEMACallback 

DDPPostHocEMACallback(*, sigma_rels: list[float] | tuple[float, ...] = (0.05, 0.1), update_every: int = 1, checkpoint_every: int = 5000, checkpoint_folder: str | None = None, start_iter: int = 0, default_output_sigma: float = 0.05, step_size_correction: bool = True)

Bases: Callback

Run the exact nitrous_ema.PostHocEMA used by MMAudio.

Official MMAudio owns two complete CUDA EMA models on local rank zero and checkpoints them every 5,000 optimizer steps. Other DDP ranks do not hold EMA copies. This callback preserves that lifecycle while the ordinary FastVideo FSDP callback continues to maintain rank-local CPU shards.

Source code in fastvideo/train/callbacks/ddp_posthoc_ema.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def __init__(
    self,
    *,
    sigma_rels: list[float] | tuple[float, ...] = (0.05, 0.1),
    update_every: int = 1,
    checkpoint_every: int = 5000,
    checkpoint_folder: str | None = None,
    start_iter: int = 0,
    default_output_sigma: float = 0.05,
    step_size_correction: bool = True,
) -> None:
    self.sigma_rels = tuple(float(value) for value in sigma_rels)
    self.update_every = max(1, int(update_every))
    self.checkpoint_every = max(1, int(checkpoint_every))
    self.checkpoint_folder = checkpoint_folder
    self.start_iter = int(start_iter)
    self.default_output_sigma = float(default_output_sigma)
    self.step_size_correction = bool(step_size_correction)

    self._ema: Any | None = None
    self._calls = 0
    self._rank = 0
    self._checkpoint_root: Path | None = None

Functions:

fastvideo.train.callbacks.ddp_posthoc_ema.method_device 

method_device(module: Module) -> device

Return the device of a nitrous EMA module.

Source code in fastvideo/train/callbacks/ddp_posthoc_ema.py 
168
169
170
def method_device(module: torch.nn.Module) -> torch.device:
    """Return the device of a nitrous EMA module."""
    return next(module.parameters()).device