From fcb5b465c5bec94e7e65333612994e14ea2fd0d0 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 5 Jul 2026 15:00:27 -0700 Subject: [PATCH] [bugfix]: checkpoint causal-CD EMA consistency target on save/resume The base TrainingMethod.checkpoint_state() only persists trainable roles, so CausalConsistencyDistillationMethod's EMA target (frozen but mutated by _update_ema every step) was never saved; a resume reloaded it from init_from, snapping the consistency target back to the base checkpoint. Persist roles.ema.transformer via the same full-state DCP wrapper DiffusionNFT already uses for its frozen 'old' role, moving _FullModelState into fastvideo/train/utils/checkpoint.py so both methods share it. --- .../methods/consistency_model/causal_cd.py | 10 ++++++ fastvideo/train/methods/rl/diffusion_nft.py | 32 +----------------- fastvideo/train/utils/checkpoint.py | 33 +++++++++++++++++++ 3 files changed, 44 insertions(+), 31 deletions(-) diff --git a/fastvideo/train/methods/consistency_model/causal_cd.py b/fastvideo/train/methods/consistency_model/causal_cd.py index 563ad0604..70b7ec949 100644 --- a/fastvideo/train/methods/consistency_model/causal_cd.py +++ b/fastvideo/train/methods/consistency_model/causal_cd.py @@ -11,6 +11,7 @@ import torch.nn.functional as F from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import (SelfForcingFlowMatchScheduler) from fastvideo.train.methods.base import LogScalar, TrainingMethod from fastvideo.train.models.base import ModelBase +from fastvideo.train.utils.checkpoint import _FullModelState from fastvideo.train.utils.optimizer import build_optimizer_and_scheduler @@ -75,6 +76,15 @@ class CausalConsistencyDistillationMethod(TrainingMethod): del iteration return [self._student_lr_scheduler] + def checkpoint_state(self) -> dict[str, Any]: + # The EMA role is frozen (so the base class skips it) but mutated by + # _update_ema every step; without persisting it a resume reloads the + # EMA from init_from and the consistency target snaps back to the + # base checkpoint. Mirrors DiffusionNFT's frozen "old" role. + states = super().checkpoint_state() + states["roles.ema.transformer"] = _FullModelState(self.ema_model.transformer) + return states + # ------------------------------------------------------------------ def single_train_step( diff --git a/fastvideo/train/methods/rl/diffusion_nft.py b/fastvideo/train/methods/rl/diffusion_nft.py index fb9488c0b..d2ca98e3b 100644 --- a/fastvideo/train/methods/rl/diffusion_nft.py +++ b/fastvideo/train/methods/rl/diffusion_nft.py @@ -10,12 +10,6 @@ from typing import Any import torch import torch.distributed as dist -from torch.distributed.checkpoint.state_dict import ( - StateDictOptions, - get_model_state_dict, - set_model_state_dict, -) -from torch.distributed.checkpoint.stateful import Stateful from tqdm.auto import tqdm from fastvideo.dataset.parquet_dataset_map_style import ( @@ -37,6 +31,7 @@ from fastvideo.train.methods.rl.common import ( validation_caption, validation_shard_indices, ) +from fastvideo.train.utils.checkpoint import _FullModelState from fastvideo.train.utils.config import ( get_optional_float, get_optional_int, @@ -78,31 +73,6 @@ class _DiffusionNFTEMAState: self._method._ema_update_count = int(update_count) -class _FullModelState(Stateful): - """DCP wrapper that saves frozen model parameters too. - - The shared ``ModelWrapper`` intentionally filters to ``requires_grad`` - parameters. DiffusionNFT's old policy is frozen but must be restored on - resume, so it needs full model state. - """ - - def __init__(self, model: torch.nn.Module) -> None: - self.model = model - - def state_dict(self) -> dict[str, Any]: - return get_model_state_dict(self.model) # type: ignore[no-any-return] - - def load_state_dict( - self, - state_dict: dict[str, Any], - ) -> None: - set_model_state_dict( - self.model, - model_state_dict=state_dict, - options=StateDictOptions(strict=False), - ) - - class DiffusionNFTMethod(TrainingMethod): """DiffusionNFT-style RL for diffusion models. diff --git a/fastvideo/train/utils/checkpoint.py b/fastvideo/train/utils/checkpoint.py index 94453a196..7a3d6c105 100644 --- a/fastvideo/train/utils/checkpoint.py +++ b/fastvideo/train/utils/checkpoint.py @@ -15,6 +15,13 @@ import numpy as np import torch import torch.distributed as dist import torch.distributed.checkpoint as dcp +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + set_model_state_dict, +) +from torch.distributed.checkpoint.stateful import Stateful + from fastvideo.logger import init_logger logger = init_logger(__name__) @@ -131,6 +138,32 @@ class _RoleModuleContainer(torch.nn.Module): self.add_module(name, module) +class _FullModelState(Stateful): + """DCP wrapper that saves frozen model parameters too. + + The shared ``ModelWrapper`` intentionally filters to ``requires_grad`` + parameters. Frozen-but-mutated roles (e.g. DiffusionNFT's old policy, + causal-CD's EMA target) must still be restored on resume, so they need + full model state. + """ + + def __init__(self, model: torch.nn.Module) -> None: + self.model = model + + def state_dict(self) -> dict[str, Any]: + return get_model_state_dict(self.model) # type: ignore[no-any-return] + + def load_state_dict( + self, + state_dict: dict[str, Any], + ) -> None: + set_model_state_dict( + self.model, + model_state_dict=state_dict, + options=StateDictOptions(strict=False), + ) + + class _CallbackStateWrapper: """Wraps a CallbackDict for DCP save/load."""