From 1a383f3f664bc226f202f5b484aade5656e25321 Mon Sep 17 00:00:00 2001 From: Zhang Peiyuan Date: Mon, 9 Mar 2026 18:58:56 -0700 Subject: [PATCH] [refactor] train v1 clean up Co-authored-by: alexzms <3036648523@qq.com> Co-authored-by: Peiyuan Zhang --- docs/training/train_infra.md | 2 +- .../configs/dfsft_wan_causal_t2v_1.3B.yaml | 1 - examples/train/configs/example.yaml | 3 --- ...2.1_t2v_1.3B_vsa_phase3.4_0.9sparsity.yaml | 12 +++------ examples/train/run.sh | 2 +- .../Wan2.1-VSA/Wan-Syn-Data/validation_4.json | 20 --------------- .../python/fastvideo_kernel/ops.py | 17 +++---------- .../dataset/parquet_dataset_map_style.py | 1 - fastvideo/train/callbacks/grad_clip.py | 4 +-- fastvideo/train/callbacks/validation.py | 2 +- fastvideo/train/entrypoint/train.py | 7 ++++++ fastvideo/train/methods/base.py | 25 ++++++------------- .../methods/distribution_matching/dmd2.py | 3 --- fastvideo/train/methods/fine_tuning/dfsft.py | 5 +--- .../train/methods/fine_tuning/finetune.py | 5 +--- fastvideo/train/models/base.py | 1 - fastvideo/train/models/wan/wan.py | 6 ++--- fastvideo/train/trainer.py | 18 +------------ fastvideo/train/utils/checkpoint.py | 18 ++++--------- fastvideo/train/utils/config.py | 7 +----- fastvideo/train/utils/moduleloader.py | 2 +- fastvideo/train/utils/training_config.py | 9 +------ 22 files changed, 39 insertions(+), 131 deletions(-) diff --git a/docs/training/train_infra.md b/docs/training/train_infra.md index 597bb6b6c..bab3a57db 100644 --- a/docs/training/train_infra.md +++ b/docs/training/train_infra.md @@ -591,5 +591,5 @@ fastvideo/train/ rationale, model/method abstractions, and open questions. - [Training Overview](overview.md) — data requirements and preprocessing. - [Data Preprocessing](data_preprocess.md) — how to prepare datasets. -- [Config Reference](../../examples/train/example.yaml) — fully-commented +- [Config Reference](../../examples/train/configs/example.yaml) — fully-commented YAML config with all fields and defaults. diff --git a/examples/train/configs/dfsft_wan_causal_t2v_1.3B.yaml b/examples/train/configs/dfsft_wan_causal_t2v_1.3B.yaml index a2e3903be..59b15172c 100644 --- a/examples/train/configs/dfsft_wan_causal_t2v_1.3B.yaml +++ b/examples/train/configs/dfsft_wan_causal_t2v_1.3B.yaml @@ -11,7 +11,6 @@ models: method: _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod - attn_kind: dense chunk_size: 3 training: diff --git a/examples/train/configs/example.yaml b/examples/train/configs/example.yaml index 5c9f226d1..c01546bcb 100644 --- a/examples/train/configs/example.yaml +++ b/examples/train/configs/example.yaml @@ -67,7 +67,6 @@ method: # --- FineTuneMethod keys (if using finetune instead) --- # _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod - # attn_kind: vsa # "dense" or "vsa" # ------------------------------------------------------------------------------ # training: [TYPED] -> TrainingConfig @@ -130,8 +129,6 @@ training: # --- training.vsa [TYPED] -> VSAConfig --- vsa: sparsity: 0.0 # default: 0.0 (0.0 = disabled) - decay_rate: 0.0 # default: 0.0 - decay_interval_steps: 0 # default: 0 # --- training.model [TYPED] -> ModelTrainingConfig --- model: diff --git a/examples/train/configs/finetune_wan2.1_t2v_1.3B_vsa_phase3.4_0.9sparsity.yaml b/examples/train/configs/finetune_wan2.1_t2v_1.3B_vsa_phase3.4_0.9sparsity.yaml index 0e418d680..210fd1592 100644 --- a/examples/train/configs/finetune_wan2.1_t2v_1.3B_vsa_phase3.4_0.9sparsity.yaml +++ b/examples/train/configs/finetune_wan2.1_t2v_1.3B_vsa_phase3.4_0.9sparsity.yaml @@ -11,21 +11,20 @@ models: method: _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod - attn_kind: vsa training: distributed: - num_gpus: 4 + num_gpus: 8 sp_size: 1 tp_size: 1 - hsdp_replicate_dim: 4 + hsdp_replicate_dim: 8 hsdp_shard_dim: 1 data: data_path: data/Wan-Syn_77x448x832_600k dataloader_num_workers: 4 train_batch_size: 1 - training_cfg_rate: 0.1 + training_cfg_rate: 0.0 seed: 1000 num_latent_t: 20 num_height: 448 @@ -58,16 +57,11 @@ training: vsa: sparsity: 0.9 - decay_rate: 0.03 - decay_interval_steps: 1 callbacks: grad_clip: _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback max_grad_norm: 1.0 - ema: - _target_: fastvideo.train.callbacks.ema.EMACallback - decay: 0.9999 validation: _target_: fastvideo.train.callbacks.validation.ValidationCallback pipeline_target: fastvideo.pipelines.basic.wan.wan_pipeline.WanPipeline diff --git a/examples/train/run.sh b/examples/train/run.sh index 09a2a3f63..263a6713a 100755 --- a/examples/train/run.sh +++ b/examples/train/run.sh @@ -49,7 +49,7 @@ echo "Extra args: $*" echo "Log file: ${LOG_FILE}" echo "==============================" -torchrun \ +python -m torch.distributed.run \ --nnodes "${NNODES}" \ --node_rank "${NODE_RANK}" \ --nproc_per_node "${NUM_GPUS}" \ diff --git a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json b/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json index 23ecf8a67..b834d2fdc 100644 --- a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json +++ b/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json @@ -2,35 +2,15 @@ "data": [ { "caption": "In the video, a woman is elegantly showcasing her earrings, bringing attention to their intricate design with a gentle touch of her fingers. She is bathed in ambient purple and pink lighting, which casts a soft glow on her delicate features and enhances the vivid tones of her lipstick and eye makeup. Her hair is styled to frame her face smoothly, emphasizing the contours of her jawline and cheekbones. The background features a blurred neon light, adding an artistic and modern touch to the overall aesthetic.", - "video_path": "Fashion/mixkit-face-of-an-elegant-and-captivating-woman-41914_clip_1.mp4", - "num_inference_steps": 3, - "height": 448, - "width": 832, - "num_frames": 61 }, { "caption": "In the video, a lone rider guides a majestic horse across an expansive, open field as the sun sets in the background. The rider, dressed in a classic blue shirt and wide-brimmed hat, sits confidently in the saddle, silhouetted against the warm glow of the evening sky. The horse moves gracefully, its mane and tail flowing with each step, creating a sense of harmony between horse and rider. Surrounding the pair, towering trees form a natural border, their leaves gently rustling in the breeze. The shadows lengthen on the ground, accentuating the serene and timeless feel of the scene. The distant hills and wooden fences frame the horizon, adding depth to the tranquil landscape. A few horses graze peacefully in the background, blending into the pastoral setting. The overall ambiance evokes a sense of calmness and quietude, capturing a perfect moment in the golden light of dusk.", - "video_path": "Man/mixkit-a-rancher-riding-a-horse-at-sunset-1143_clip_1.mp4", - "num_inference_steps": 3, - "height": 448, - "width": 832, - "num_frames": 61 }, { "caption": "In a dimly lit, eerie setting, a mysterious pink bottle labeled \"Authentic 100% organic POISON\" sits prominently in the foreground, casting a menacing aura. The bottle is accentuated by green fog, which swirls lightly around it, enhancing its sinister allure. Behind it, a shadowy golden bottle adorned with a spider emblem subtly emerges, adding an extra layer of mystery to the scene. Dim candles provide faint, flickering light, which complements the dark atmosphere, making the setting ideal for an illusion of hidden dangers.", - "video_path": "smoke/mixkit-poison-in-halloween-ritual-33879_clip_1.mp4", - "num_inference_steps": 3, - "height": 448, - "width": 832, - "num_frames": 61 }, { "caption": "The video opens with a tranquil scene in the heart of a dense forest, emphasizing two large, textured tree trunks in the foreground framing the view. Sunlight filters through the canopy above, casting intricate patterns of light and shadow on the trees and the ground. Between the tree trunks, a clear view of a calm, muddy river unfolds, its surface shimmering under the gentle sunlight. The riverbank is decorated with a variety of small bushes and vibrant foliage, subtly transitioning into the deep greens of tall, leafy plants. In the background, the dense forest looms, filled with dark, towering trees, their branches intertwining to form an intricate canopy. The scene is bathed in the soft glow of the sun, creating a serene and picturesque setting. Occasional sunbeams pierce through the foliage, adding a magical aura to the landscape. The vibrant reds and oranges of the smaller plants add contrast, bringing warmth to the earthy tones of the scenery. Overall, this harmonious blend of natural elements creates a peaceful and idyllic forest setting.", - "video_path": "forest/mixkit-view-of-a-river-between-two-old-trees-560_clip_1.mp4", - "num_inference_steps": 3, - "height": 448, - "width": 832, - "num_frames": 61 } ] } diff --git a/fastvideo-kernel/python/fastvideo_kernel/ops.py b/fastvideo-kernel/python/fastvideo_kernel/ops.py index a0924d4ce..72b49bb3c 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/ops.py +++ b/fastvideo-kernel/python/fastvideo_kernel/ops.py @@ -1,20 +1,15 @@ import math import torch from .block_sparse_attn import block_sparse_attn -from .triton_kernels.block_sparse_attn_triton import triton_block_sparse_attn_forward from .triton_kernels.st_attn_triton import sliding_tile_attention_triton -from .triton_kernels.index import map_to_index # Try to load the C++ extension try: from fastvideo_kernel._C import fastvideo_kernel_ops sta_fwd = getattr(fastvideo_kernel_ops, "sta_fwd", None) - block_sparse_fwd = getattr(fastvideo_kernel_ops, "block_sparse_fwd", None) - block_sparse_bwd = getattr(fastvideo_kernel_ops, "block_sparse_bwd", None) except ImportError: sta_fwd = None - block_sparse_fwd = None - block_sparse_bwd = None + def sliding_tile_attention( q: torch.Tensor, @@ -135,14 +130,8 @@ def video_sparse_attn( mask = torch.zeros_like(scores, dtype=torch.bool).scatter_(-1, topk_idx, True) - idx, num = map_to_index(mask) - - if block_sparse_fwd is not None: - # Use autograd-enabled wrapper so backward works (and still uses SM90 kernel when available) - out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0] - else: - # Triton-only forward (kept for environments without the wrapper deps) - out_s, _ = triton_block_sparse_attn_forward(q, k, v, idx, num, variable_block_sizes) + # out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0] + out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0] if compress_attn_weight is not None: return out_c * compress_attn_weight + out_s diff --git a/fastvideo/dataset/parquet_dataset_map_style.py b/fastvideo/dataset/parquet_dataset_map_style.py index d499ccfbf..c14eee93e 100644 --- a/fastvideo/dataset/parquet_dataset_map_style.py +++ b/fastvideo/dataset/parquet_dataset_map_style.py @@ -392,7 +392,6 @@ def build_parquet_map_style_dataloader( collate_fn=passthrough, num_workers=num_data_workers, pin_memory=True, - pin_memory_device=current_platform.device_name, persistent_workers=num_data_workers > 0, ) return dataset, loader diff --git a/fastvideo/train/callbacks/grad_clip.py b/fastvideo/train/callbacks/grad_clip.py index 4e8aa248a..e15761062 100644 --- a/fastvideo/train/callbacks/grad_clip.py +++ b/fastvideo/train/callbacks/grad_clip.py @@ -31,8 +31,8 @@ class GradNormClipCallback(Callback): def __init__( self, *, - max_grad_norm: float = 0.0, - log_grad_norms: bool = False, + max_grad_norm: float = 1.0, + log_grad_norms: bool = True, ) -> None: self._max_grad_norm = float(max_grad_norm) self._log_grad_norms = bool(log_grad_norms) diff --git a/fastvideo/train/callbacks/validation.py b/fastvideo/train/callbacks/validation.py index 29225a579..446243314 100644 --- a/fastvideo/train/callbacks/validation.py +++ b/fastvideo/train/callbacks/validation.py @@ -361,7 +361,7 @@ class ValidationCallback(Callback): generator=self.validation_random_generator, n_tokens=n_tokens, eta=0.0, - VSA_sparsity=tc.vsa.sparsity, + VSA_sparsity=tc.vsa_sparsity, timesteps=sampling_timesteps_tensor, ) batch._inference_args = inference_args # type: ignore[attr-defined] diff --git a/fastvideo/train/entrypoint/train.py b/fastvideo/train/entrypoint/train.py index 789a7fc59..bd20f0bd2 100644 --- a/fastvideo/train/entrypoint/train.py +++ b/fastvideo/train/entrypoint/train.py @@ -54,6 +54,13 @@ def run_training_from_config( cfg = load_run_config(config_path, overrides=overrides) tc = cfg.training + # Auto-set attention backend for VSA when sparsity is configured. + if tc.vsa_sparsity > 0.0: + os.environ.setdefault( + "FASTVIDEO_ATTENTION_BACKEND", + "VIDEO_SPARSE_ATTN", + ) + maybe_init_distributed_environment_and_model_parallel( tc.distributed.tp_size, tc.distributed.sp_size, diff --git a/fastvideo/train/methods/base.py b/fastvideo/train/methods/base.py index 0f2898028..8561f167f 100644 --- a/fastvideo/train/methods/base.py +++ b/fastvideo/train/methods/base.py @@ -4,10 +4,11 @@ from __future__ import annotations from abc import ABC, abstractmethod from collections.abc import Sequence -from typing import Any, Literal, TypeAlias, cast +from typing import Any, Literal, TypeAlias import torch +from fastvideo import envs from fastvideo.logger import init_logger from fastvideo.train.models.base import ModelBase from fastvideo.train.utils.checkpoint import _RoleModuleContainer @@ -77,8 +78,6 @@ class TrainingMethod(torch.nn.Module, ABC): self, batch: dict[str, Any], iteration: int, - *, - current_vsa_sparsity: float = 0.0, ) -> tuple[ dict[str, torch.Tensor], dict[str, Any], @@ -236,18 +235,10 @@ class TrainingMethod(torch.nn.Module, ABC): self.student.on_train_start() - def get_rng_generators(self, ) -> dict[str, torch.Generator]: - generators: dict[str, torch.Generator] = {} - if self.cuda_generator is not None: - generators["cuda"] = self.cuda_generator - return generators - @staticmethod - def _parse_attn_kind(raw: Any, ) -> Literal["dense", "vsa"]: - if raw in (None, ""): - return "dense" - kind = str(raw).strip().lower() - if kind not in {"dense", "vsa"}: - raise ValueError("method_config.attn_kind must be one of " - f"{{'dense', 'vsa'}}, got {raw!r}.") - return cast(Literal["dense", "vsa"], kind) + def _infer_attn_kind() -> Literal["dense", "vsa"]: + """Derive attn_kind from ``FASTVIDEO_ATTENTION_BACKEND``.""" + backend = envs.FASTVIDEO_ATTENTION_BACKEND + if backend == "VIDEO_SPARSE_ATTN": + return "vsa" + return "dense" diff --git a/fastvideo/train/methods/distribution_matching/dmd2.py b/fastvideo/train/methods/distribution_matching/dmd2.py index 9162bae63..6bfbb44ad 100644 --- a/fastvideo/train/methods/distribution_matching/dmd2.py +++ b/fastvideo/train/methods/distribution_matching/dmd2.py @@ -81,8 +81,6 @@ class DMD2Method(TrainingMethod): self, batch: dict[str, Any], iteration: int, - *, - current_vsa_sparsity: float = 0.0, ) -> tuple[ dict[str, torch.Tensor], dict[str, Any], @@ -95,7 +93,6 @@ class DMD2Method(TrainingMethod): training_batch = self.student.prepare_batch( batch, generator=self.cuda_generator, - current_vsa_sparsity=current_vsa_sparsity, latents_source=latents_source, ) diff --git a/fastvideo/train/methods/fine_tuning/dfsft.py b/fastvideo/train/methods/fine_tuning/dfsft.py index a0f778886..e4eda24d4 100644 --- a/fastvideo/train/methods/fine_tuning/dfsft.py +++ b/fastvideo/train/methods/fine_tuning/dfsft.py @@ -31,7 +31,7 @@ class DiffusionForcingSFTMethod(TrainingMethod): raise ValueError("DFSFT requires role 'student'") if not self.student._trainable: raise ValueError("DFSFT requires student to be trainable") - self._attn_kind: Literal["dense", "vsa"] = (self._parse_attn_kind(self.method_config.get("attn_kind", None))) + self._attn_kind: Literal["dense", "vsa"] = (self._infer_attn_kind()) self._chunk_size = self._parse_chunk_size(self.method_config.get("chunk_size", None)) self._timestep_index_range = (self._parse_timestep_index_range()) @@ -54,8 +54,6 @@ class DiffusionForcingSFTMethod(TrainingMethod): self, batch: dict[str, Any], iteration: int, - *, - current_vsa_sparsity: float = 0.0, ) -> tuple[ dict[str, torch.Tensor], dict[str, Any], @@ -65,7 +63,6 @@ class DiffusionForcingSFTMethod(TrainingMethod): training_batch = self.student.prepare_batch( batch, generator=self.cuda_generator, - current_vsa_sparsity=current_vsa_sparsity, latents_source="data", ) diff --git a/fastvideo/train/methods/fine_tuning/finetune.py b/fastvideo/train/methods/fine_tuning/finetune.py index 624a5ecab..ba6d52e88 100644 --- a/fastvideo/train/methods/fine_tuning/finetune.py +++ b/fastvideo/train/methods/fine_tuning/finetune.py @@ -30,7 +30,7 @@ class FineTuneMethod(TrainingMethod): if not self.student._trainable: raise ValueError("FineTuneMethod requires student to be " "trainable") - self._attn_kind: Literal["dense", "vsa"] = (self._parse_attn_kind(self.method_config.get("attn_kind", None))) + self._attn_kind: Literal["dense", "vsa"] = (self._infer_attn_kind()) # Initialize preprocessors on student. self.student.init_preprocessors(self.training_config) @@ -50,8 +50,6 @@ class FineTuneMethod(TrainingMethod): self, batch: dict[str, Any], iteration: int, - *, - current_vsa_sparsity: float = 0.0, ) -> tuple[ dict[str, torch.Tensor], dict[str, Any], @@ -61,7 +59,6 @@ class FineTuneMethod(TrainingMethod): training_batch = self.student.prepare_batch( batch, generator=self.cuda_generator, - current_vsa_sparsity=current_vsa_sparsity, latents_source="data", ) diff --git a/fastvideo/train/models/base.py b/fastvideo/train/models/base.py index 6b97a9d0b..f56e8c9b1 100644 --- a/fastvideo/train/models/base.py +++ b/fastvideo/train/models/base.py @@ -69,7 +69,6 @@ class ModelBase(ABC): raw_batch: dict[str, Any], *, generator: torch.Generator, - current_vsa_sparsity: float = 0.0, latents_source: Literal["data", "zeros"] = "data", ) -> TrainingBatch: """Convert a dataloader batch into forward primitives.""" diff --git a/fastvideo/train/models/wan/wan.py b/fastvideo/train/models/wan/wan.py index 61e157ca8..57240f047 100644 --- a/fastvideo/train/models/wan/wan.py +++ b/fastvideo/train/models/wan/wan.py @@ -192,7 +192,6 @@ class WanModel(ModelBase): raw_batch: dict[str, Any], *, generator: torch.Generator, - current_vsa_sparsity: float = 0.0, latents_source: Literal["data", "zeros"] = "data", ) -> TrainingBatch: self.ensure_negative_conditioning() @@ -202,7 +201,7 @@ class WanModel(ModelBase): dtype = self._get_training_dtype() device = self.device - training_batch = TrainingBatch(current_vsa_sparsity=current_vsa_sparsity) + training_batch = TrainingBatch() encoder_hidden_states = raw_batch["text_embedding"] encoder_attention_mask = raw_batch["text_attention_mask"] infos = raw_batch.get("info_list") @@ -460,7 +459,6 @@ class WanModel(ModelBase): patch_size = ( tc.pipeline_config.dit_config.patch_size # type: ignore[union-attr] ) - current_vsa_sparsity = (training_batch.current_vsa_sparsity) assert latents_shape is not None assert training_batch.timesteps is not None @@ -474,7 +472,7 @@ class WanModel(ModelBase): raw_latent_shape=latents_shape[2:5], current_timestep=(training_batch.timesteps), patch_size=patch_size, - VSA_sparsity=current_vsa_sparsity, + VSA_sparsity=tc.vsa_sparsity, device=self.device, ) elif (envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN"): diff --git a/fastvideo/train/trainer.py b/fastvideo/train/trainer.py index 4a033b2c9..d814a6cfe 100644 --- a/fastvideo/train/trainer.py +++ b/fastvideo/train/trainer.py @@ -36,7 +36,6 @@ def _coerce_log_scalar(value: Any, *, where: str) -> float: class TrainLoopState: step: int accum_iter: int - current_vsa_sparsity: float class Trainer: @@ -73,19 +72,6 @@ class Trainer: batch = next(data_iter) yield batch - def _get_current_vsa_sparsity(self, step: int) -> float: - tc = self.training_config - vsa_sparsity = tc.vsa.sparsity - vsa_decay_rate = tc.vsa.decay_rate - vsa_decay_interval_steps = (tc.vsa.decay_interval_steps) - if vsa_decay_interval_steps > 1: - current_decay_times = min( - step // vsa_decay_interval_steps, - int(vsa_sparsity // vsa_decay_rate), - ) - return current_decay_times * vsa_decay_rate - return vsa_sparsity - def run( self, method: TrainingMethod, @@ -136,7 +122,6 @@ class Trainer: ) for step in progress: t0 = time.perf_counter() - current_vsa_sparsity = (self._get_current_vsa_sparsity(step)) loss_sums: dict[str, float] = {} metric_sums: dict[str, float] = {} @@ -145,7 +130,6 @@ class Trainer: loss_map, outputs, step_metrics = method.single_train_step( batch, step, - current_vsa_sparsity=(current_vsa_sparsity), ) method.backward( @@ -179,7 +163,7 @@ class Trainer: metrics = {k: v / grad_accum for k, v in loss_sums.items()} metrics.update({k: v / grad_accum for k, v in metric_sums.items()}) metrics["step_time_sec"] = (time.perf_counter() - t0) - metrics["vsa_sparsity"] = float(current_vsa_sparsity) + metrics["vsa_sparsity"] = float(tc.vsa_sparsity) if self.global_rank == 0 and metrics: self.tracker.log(metrics, step) diff --git a/fastvideo/train/utils/checkpoint.py b/fastvideo/train/utils/checkpoint.py index eca95d396..94453a196 100644 --- a/fastvideo/train/utils/checkpoint.py +++ b/fastvideo/train/utils/checkpoint.py @@ -275,12 +275,8 @@ class CheckpointManager: "python_rng": random.getstate(), "numpy_rng": np.random.get_state(), } - if torch.cuda.is_available(): - rng["cuda_rng"] = torch.cuda.get_rng_state() - generators = (self.method.get_rng_generators() if hasattr(self.method, "get_rng_generators") else {}) - for name, gen in (generators or {}).items(): - if gen is not None: - rng[f"gen_{name}"] = gen.get_state() + rng["cuda_rng"] = torch.cuda.get_rng_state() + rng["gen_cuda"] = self.method.cuda_generator.get_state() torch.save( rng, checkpoint_dir / f"rng_state_rank{rank}.pt", @@ -326,13 +322,9 @@ class CheckpointManager: random.setstate(rng["python_rng"]) if "numpy_rng" in rng: np.random.set_state(rng["numpy_rng"]) - if torch.cuda.is_available() and "cuda_rng" in rng: - torch.cuda.set_rng_state(rng["cuda_rng"]) - generators = (self.method.get_rng_generators() if hasattr(self.method, "get_rng_generators") else {}) - for name, gen in (generators or {}).items(): - key = f"gen_{name}" - if key in rng and gen is not None: - gen.set_state(rng[key]) + + torch.cuda.set_rng_state(rng["cuda_rng"]) + self.method.cuda_generator.set_state(rng["gen_cuda"]) logger.info( "Restored RNG snapshot from %s", rng_path, diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py index 40da284fc..86e9ebaf3 100644 --- a/fastvideo/train/utils/config.py +++ b/fastvideo/train/utils/config.py @@ -20,7 +20,6 @@ from fastvideo.train.utils.training_config import ( OptimizerConfig, TrackerConfig, TrainingLoopConfig, - VSAConfig, ) logger = init_logger(__name__) @@ -354,11 +353,7 @@ def _build_training_config( project_name=str(tr.get("project_name", "fastvideo") or "fastvideo"), run_name=str(tr.get("run_name", "") or ""), ), - vsa=VSAConfig( - sparsity=float(vs.get("sparsity", 0.0) or 0.0), - decay_rate=float(vs.get("decay_rate", 0.0) or 0.0), - decay_interval_steps=int(vs.get("decay_interval_steps", 0) or 0), - ), + vsa_sparsity=float(vs.get("sparsity", 0.0) or 0.0), model=ModelTrainingConfig( weighting_scheme=str(m.get("weighting_scheme", "uniform") or "uniform"), logit_mean=float(m.get("logit_mean", 0.0) or 0.0), diff --git a/fastvideo/train/utils/moduleloader.py b/fastvideo/train/utils/moduleloader.py index d775db9bd..c1c92d875 100644 --- a/fastvideo/train/utils/moduleloader.py +++ b/fastvideo/train/utils/moduleloader.py @@ -68,7 +68,7 @@ def make_inference_args( args.inference_mode = True args.mode = ExecutionMode.INFERENCE args.dit_cpu_offload = True - args.VSA_sparsity = tc.vsa.sparsity + args.VSA_sparsity = tc.vsa_sparsity return args diff --git a/fastvideo/train/utils/training_config.py b/fastvideo/train/utils/training_config.py index 0167751db..8dd0f440b 100644 --- a/fastvideo/train/utils/training_config.py +++ b/fastvideo/train/utils/training_config.py @@ -66,13 +66,6 @@ class TrackerConfig: run_name: str = "" -@dataclass(slots=True) -class VSAConfig: - sparsity: float = 0.0 - decay_rate: float = 0.0 - decay_interval_steps: int = 0 - - @dataclass(slots=True) class ModelTrainingConfig: weighting_scheme: str = "uniform" @@ -92,7 +85,7 @@ class TrainingConfig: loop: TrainingLoopConfig = field(default_factory=TrainingLoopConfig) checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig) tracker: TrackerConfig = field(default_factory=TrackerConfig) - vsa: VSAConfig = field(default_factory=VSAConfig) + vsa_sparsity: float = 0.0 model: ModelTrainingConfig = field(default_factory=ModelTrainingConfig) pipeline_config: PipelineConfig | None = None model_path: str = ""