From f934efc58dbd7e5f7494490d5f85c64e7ec7fd2d Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Sat, 27 Jun 2026 08:02:19 +0000 Subject: [PATCH] add cf --- IMPLEMENTATION_NOTES_causal_forcing.md | 113 ++++++++ .../consistency_model/wan/causal_cd_t2v.yaml | 78 ++++++ .../fine_tuning/wan/tfsft_causal_t2v.yaml | 74 ++++++ fastvideo/models/dits/causal_wanvideo.py | 121 ++++++++- .../wan_causal_t2v_causal_cd_min.yaml | 56 ++++ .../fixtures/wan_causal_t2v_tfsft_min.yaml | 45 ++++ .../tests/train/methods/test_wan_causal_cd.py | 148 +++++++++++ .../train/methods/test_wan_causal_tfsft.py | 113 ++++++++ fastvideo/train/methods/__init__.py | 9 + .../methods/consistency_model/__init__.py | 21 +- .../methods/consistency_model/causal_cd.py | 248 ++++++++++++++++++ .../train/methods/fine_tuning/__init__.py | 7 + fastvideo/train/methods/fine_tuning/dfsft.py | 27 +- fastvideo/train/methods/fine_tuning/tfsft.py | 39 +++ fastvideo/train/models/wan/wan.py | 17 +- 15 files changed, 1099 insertions(+), 17 deletions(-) create mode 100644 IMPLEMENTATION_NOTES_causal_forcing.md create mode 100644 examples/train/configs/consistency_model/wan/causal_cd_t2v.yaml create mode 100644 examples/train/configs/fine_tuning/wan/tfsft_causal_t2v.yaml create mode 100644 fastvideo/tests/train/fixtures/wan_causal_t2v_causal_cd_min.yaml create mode 100644 fastvideo/tests/train/fixtures/wan_causal_t2v_tfsft_min.yaml create mode 100644 fastvideo/tests/train/methods/test_wan_causal_cd.py create mode 100644 fastvideo/tests/train/methods/test_wan_causal_tfsft.py create mode 100644 fastvideo/train/methods/consistency_model/causal_cd.py create mode 100644 fastvideo/train/methods/fine_tuning/tfsft.py diff --git a/IMPLEMENTATION_NOTES_causal_forcing.md b/IMPLEMENTATION_NOTES_causal_forcing.md new file mode 100644 index 000000000..15dd83113 --- /dev/null +++ b/IMPLEMENTATION_NOTES_causal_forcing.md @@ -0,0 +1,113 @@ +# Causal-Forcing in FastVideo: Clean-History Teacher Forcing + Causal Consistency Distillation + +Adds two of the critical Wan-parity gaps identified vs `thu-ml/Causal-Forcing`: +**Stage-1 clean-history Teacher Forcing (TF)** and **Stage-2b Causal Consistency +Distillation (CD / Causal-Forcing++)**. Implemented in the modular `fastvideo/train/` +stack, reusing existing abstractions; the legacy `fastvideo/training/` stack is +untouched. + +## What was added + +### 1. DiT: clean-history teacher-forcing path +`fastvideo/models/dits/causal_wanvideo.py` +- `_prepare_teacher_forcing_mask(...)` — ported from Causal-Forcing + `wan/modules/causal_model.py:569`. Builds a FlexAttention mask over the + concatenated `[clean | noisy]` sequence (length `2 * num_frames * frame_seqlen`): + a noisy block attends to its own noisy tokens plus the **clean** context of all + strictly previous blocks (+ diagonal). +- `_forward_train(..., clean_x=None, aug_t=None)` — when `clean_x` is provided: + patch-embeds the clean latents, prepends them to the noisy tokens, tiles RoPE and + the per-frame timestep modulation so clean frame *i* and noisy frame *i* share the + same position / embedding (clean frames are time-embedded at `aug_t`, default 0), + runs the blocks under the TF mask, then drops the clean half before `norm_out` + (loss is denoise-only on the current/noisy tokens). The diffusion-forcing path is + unchanged. A separate `teacher_forcing_block_mask` is cached so the two masks never + collide. + +### 2. Wrapper plumbing +`fastvideo/train/models/wan/wan.py` +- `predict_noise(..., clean_x=None, aug_t=None)` and `_build_distill_input_kwargs(...)` + forward `clean_x`/`aug_t` to the transformer (permuting `clean_x` to `[B,C,F,H,W]`). + The streaming/KV-cache path is untouched. + +### 3. Teacher-Forcing SFT method (Stage 1) +`fastvideo/train/methods/fine_tuning/tfsft.py` — `TeacherForcingSFTMethod`, a thin +subclass of `DiffusionForcingSFTMethod`. Identical objective (inhomogeneous per-chunk +timesteps, flow MSE, bsmntw weighting) except it passes `clean_x=clean_latents`. The +only change to `dfsft.py` is extracting the predict call into an overridable +`_predict_noise(...)` hook (behavior-preserving for DFSFT). + +### 4. Causal Consistency Distillation method (Stage 2b) +`fastvideo/train/methods/consistency_model/causal_cd.py` — +`CausalConsistencyDistillationMethod`, ported from Causal-Forcing +`model/naive_consistency.py`. Per step: +1. Sample a discrete index over a `discrete_cd_N`-step flow-match schedule → `(t, t_next)`. +2. **Teacher** (frozen, teacher-forced) does a single CFG Euler step + `latent_t → latent_t_next = latent_t - dt·[v_uncond + g·(v_cond - v_uncond)]`, + `dt = (t - t_next)/num_train_timesteps`. +3. **Student** (trainable, teacher-forced) predicts `x0` at `t`. +4. **EMA** student (frozen, teacher-forced) predicts `x0` at `t_next`. +5. Loss `= MSE(x0_t, x0_t_next)`; EMA is lerp-updated after each optimizer step. + +Roles: `student` (trainable), `teacher` and `ema` (frozen) — all initialized from the +same checkpoint, matching the reference's three-network setup. + +### 5. Configs, tests, registration +- Smoke tests (mirror `test_wan_causal_dfsft.py`, real Wan-1.3B + tiny synthetic data): + `fastvideo/tests/train/methods/test_wan_causal_tfsft.py`, + `.../test_wan_causal_cd.py`, with fixtures + `fastvideo/tests/train/fixtures/wan_causal_t2v_{tfsft,causal_cd}_min.yaml`. +- Example run configs: `examples/train/configs/fine_tuning/wan/tfsft_causal_t2v.yaml`, + `examples/train/configs/consistency_model/wan/causal_cd_t2v.yaml`. +- Registered in `fastvideo/train/methods/__init__.py`, + `.../fine_tuning/__init__.py`, `.../consistency_model/__init__.py`. + +## Running the smoke tests +```bash +conda activate FastVideo_kaiqin +export PYTHONPATH=$PWD # from the FastVideo_cf clone +pytest fastvideo/tests/train/methods/test_wan_causal_tfsft.py -q +pytest fastvideo/tests/train/methods/test_wan_causal_cd.py -q +``` +Both require CUDA and the cached `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` weights. Each runs +one real train step and asserts finite loss + nonzero student gradients reaching the +first transformer block (CD additionally asserts a frozen teacher and a working EMA +update). **Status: both pass on GB200.** + +For W&B logging on a real run, set `WANDB_API_KEY` in the environment (do not commit +keys) and add a `training.tracker` block (see the example configs). + +## Assumptions & deviations from the reference +- **Discrete CD schedule.** The reference builds its own `FlowMatchScheduler(shift=5, + extra_one_step=True).set_timesteps(discrete_cd_N)`. Here we sub-sample `discrete_cd_N` + points from the student's `FlowMatchEulerDiscreteScheduler.timesteps` and reuse the + reference's Euler convention `dt = (t - t_next)/num_train_timesteps`. The + shift-vs-linear sigma parameterization is therefore approximate in exactly the same + way the reference's `/1000` Euler step is — internally consistent, but not a + bit-exact reproduction of the reference scheduler. Tune `discrete_cd_N`/`flow_shift` + for a real run. +- **EMA as an explicit role, not a deepcopy.** The reference keeps `generator_ema` and + updates it via an external EMA tracker (`ema_model.copy_to`). FSDP `fully_shard` + forbids `deepcopy`, so the EMA is a third role model loaded from the same checkpoint + and lerp-updated in place (`optimizers_schedulers_step`). This is FSDP-safe and the + EMA shards align with the student's. It does cost a third full model copy (as the + reference also effectively does). It is **not** integrated with FastVideo's + `EMACallback` (which is export/inference-only); a production setup may want to unify + these. +- **`aug_t` noise augmentation** of the clean context (reference + `noise_augmentation_max_timestep`) is plumbed through the DiT (`aug_t`) but the TF + method passes `aug_t=None` (clean context at timestep 0), matching the reference's + default. CD also uses clean (un-augmented) context. +- **Chunk vs frame granularity** is the existing `num_frame_per_block` knob (1 = frame, + 3 = chunk); both TF and CD honor it. There is still no shipped Wan *framewise* + checkpoint/recipe (orthogonal to this change). +- **CFG uncond text** for the teacher comes from the student's negative-prompt + conditioning (`unconditional_dict`), as in the reference. + +## Not implemented (out of scope for this change) +- Stage-2a Causal **ODE** init / `blockwise_kv` data curation (the CD branch is the + ODE-data-free alternative). +- Stage 3 asymmetric **DMD** changes, 1/2-step few-step (ASD first-chunk), SiD/GAN. +- Native frame-wise **I2V** architecture, Rolling-Forcing long video, TAEHV/TRT decode. +- `--use_ema` selection at standalone inference. +See `../CAUSAL_FORCING_PARITY_GAP_ANALYSIS.md` for the full gap list and priorities. diff --git a/examples/train/configs/consistency_model/wan/causal_cd_t2v.yaml b/examples/train/configs/consistency_model/wan/causal_cd_t2v.yaml new file mode 100644 index 000000000..bbb6a7ea9 --- /dev/null +++ b/examples/train/configs/consistency_model/wan/causal_cd_t2v.yaml @@ -0,0 +1,78 @@ +# Causal Consistency Distillation (Causal-Forcing++ Stage-2b): Wan 2.1 T2V 1.3B Causal +# +# ODE-data-free initialization for asymmetric DMD. A frozen teacher takes a +# single CFG Euler step; the student matches an EMA copy of itself at the next +# timestep, all under clean-history teacher forcing. +# +# All three roles initialize from the SAME checkpoint (the Stage-1 AR-diffusion +# / teacher-forcing model, e.g. checkpoints/{chunkwise,framewise}/ar_diffusion). +# Point init_from at that checkpoint for a real run. + +models: + student: + _target_: fastvideo.train.models.wan.WanCausalModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: true + teacher: + _target_: fastvideo.train.models.wan.WanCausalModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: false + ema: + _target_: fastvideo.train.models.wan.WanCausalModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: false + +method: + _target_: fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod + discrete_cd_N: 48 + guidance_scale: 3.0 + ema_decay: 0.95 + +training: + distributed: + num_gpus: 8 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 8 + + data: + data_path: data/Wan-Syn_77x448x832_600k + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 18 + num_height: 448 + num_width: 832 + num_frames: 69 + + optimizer: + learning_rate: 2.0e-6 + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 3000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/wan2.1_causal_cd + training_state_checkpointing_steps: 1000 + checkpoints_total_limit: 3 + + tracker: + project_name: distillation_wan_r + run_name: wan2.1_causal_cd_shift5 + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + max_grad_norm: 1.0 + +pipeline: + flow_shift: 5 diff --git a/examples/train/configs/fine_tuning/wan/tfsft_causal_t2v.yaml b/examples/train/configs/fine_tuning/wan/tfsft_causal_t2v.yaml new file mode 100644 index 000000000..4c181365f --- /dev/null +++ b/examples/train/configs/fine_tuning/wan/tfsft_causal_t2v.yaml @@ -0,0 +1,74 @@ +# TFSFT (Teacher-Forcing SFT): Wan 2.1 T2V 1.3B Causal +# Causal-Forcing Stage-1 AR-diffusion training. +# +# - Student: trainable causal Wan model +# - Training: inhomogeneous timesteps per chunk, but the causal transformer +# denoises the current block while attending to *clean* history (clean_x), +# not its own noisy rollout. This is the "Causal Forcing" namesake mechanism +# and the recommended init for the ODE / consistency (Stage-2) stages. + +models: + student: + _target_: fastvideo.train.models.wan.WanCausalModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.tfsft.TeacherForcingSFTMethod + chunk_size: 3 + +training: + distributed: + num_gpus: 8 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 8 + + data: + data_path: data/Wan-Syn_77x448x832_600k + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 18 + num_height: 448 + num_width: 832 + num_frames: 69 + + optimizer: + learning_rate: 2.0e-6 + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 4000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/wan2.1_causal_tfsft + training_state_checkpointing_steps: 1000 + checkpoints_total_limit: 3 + + tracker: + project_name: distillation_wan_r + run_name: wan2.1_causal_tfsft_shift5_gauss_weight + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + max_grad_norm: 1.0 + validation: + pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_pipeline.WanCausalPipeline + dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json + every_steps: 50 + sampling_steps: [40] + guidance_scale: 6.0 + num_frames: 69 + +pipeline: + flow_shift: 5 diff --git a/fastvideo/models/dits/causal_wanvideo.py b/fastvideo/models/dits/causal_wanvideo.py index 32599ec82..a21bc0b6a 100644 --- a/fastvideo/models/dits/causal_wanvideo.py +++ b/fastvideo/models/dits/causal_wanvideo.py @@ -437,6 +437,7 @@ class CausalWanTransformer3DModel(BaseDiT): # Causal-specific self.block_mask = None + self.teacher_forcing_block_mask = None self.num_frame_per_block = config.arch_config.num_frames_per_block assert self.num_frame_per_block <= 3 self.independent_first_frame = False @@ -500,6 +501,68 @@ class CausalWanTransformer3DModel(BaseDiT): return block_mask + @staticmethod + def _prepare_teacher_forcing_mask( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=1 + ) -> BlockMask: + """Teacher-forcing attention mask (ported from Causal-Forcing). + + The token sequence is the concatenation ``[clean | noisy]`` (length + ``2 * num_frames * frame_seqlen``). A clean token attends block-wise + causally to the clean context; a noisy token attends to the noisy + tokens of its own block plus the clean context of all strictly + previous blocks (and itself). This lets the model denoise the current + block while conditioning on *clean* history rather than its own noisy + rollout. + """ + total_length = num_frames * frame_seqlen * 2 + padded_length = math.ceil(total_length / 128) * 128 - total_length + + clean_ends = num_frames * frame_seqlen + context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + + attention_block_size = frame_seqlen * num_frame_per_block + frame_indices = torch.arange( + start=0, end=num_frames * frame_seqlen, + step=attention_block_size, device=device, dtype=torch.long + ) + for start in frame_indices: + context_ends[start:start + attention_block_size] = start + attention_block_size + + noisy_image_start_list = torch.arange( + num_frames * frame_seqlen, total_length, + step=attention_block_size, device=device, dtype=torch.long + ) + noisy_image_end_list = noisy_image_start_list + attention_block_size + for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)): + noise_noise_starts[start:end] = start + noise_noise_ends[start:end] = end + noise_context_ends[start:end] = block_index * attention_block_size + + def attention_mask(b, h, q_idx, kv_idx): + clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx]) + c1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx]) + c2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx]) + noise_mask = (q_idx >= clean_ends) & (c1 | c2) + eye_mask = q_idx == kv_idx + return eye_mask | clean_mask | noise_mask + + block_mask = create_block_mask( + attention_mask, B=None, H=None, + Q_LEN=total_length + padded_length, KV_LEN=total_length + padded_length, + _compile=False, device=device) + + if not dist.is_initialized() or dist.get_rank() == 0: + print(f" cache a teacher-forcing mask with block size of {num_frame_per_block} frames") + print(block_mask) + + return block_mask + def _forward_inference( self, hidden_states: torch.Tensor, @@ -628,9 +691,12 @@ class CausalWanTransformer3DModel(BaseDiT): encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, start_frame: int = 0, + clean_x: torch.Tensor | None = None, + aug_t: torch.Tensor | None = None, **kwargs) -> torch.Tensor: orig_dtype = hidden_states.dtype + teacher_forcing = clean_x is not None if not isinstance(encoder_hidden_states, torch.Tensor): encoder_hidden_states = encoder_hidden_states[0] if isinstance(encoder_hidden_states_image, @@ -663,15 +729,27 @@ class CausalWanTransformer3DModel(BaseDiT): freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None - # Construct blockwise causal attn mask - if self.block_mask is None: - self.block_mask = self._prepare_blockwise_causal_attn_mask( - device=hidden_states.device, - num_frames=num_frames, - frame_seqlen=post_patch_height * post_patch_width, - num_frame_per_block=self.num_frame_per_block, - local_attn_size=self.local_attn_size - ) + # Construct attention mask. Teacher forcing concatenates a clean + # context copy of every frame, so it needs its own mask. + if teacher_forcing: + if self.teacher_forcing_block_mask is None: + self.teacher_forcing_block_mask = self._prepare_teacher_forcing_mask( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * post_patch_width, + num_frame_per_block=self.num_frame_per_block, + ) + block_mask = self.teacher_forcing_block_mask + else: + if self.block_mask is None: + self.block_mask = self._prepare_blockwise_causal_attn_mask( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size + ) + block_mask = self.block_mask hidden_states = self.patch_embedding(hidden_states) grid_sizes = torch.stack( @@ -679,6 +757,7 @@ class CausalWanTransformer3DModel(BaseDiT): hidden_states = hidden_states.flatten(2).transpose(1, 2) encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states.new_zeros(1, self.text_len - encoder_hidden_states.size(1), encoder_hidden_states.size(2))], dim=1) + encoder_hidden_states_text = encoder_hidden_states temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( timestep.flatten(), encoder_hidden_states, encoder_hidden_states_image) @@ -694,18 +773,38 @@ class CausalWanTransformer3DModel(BaseDiT): assert encoder_hidden_states.dtype == orig_dtype + if teacher_forcing: + # Prepend clean-context tokens; time-embed them at aug_t (default 0). + # RoPE and the per-frame modulation are tiled so clean frame i and + # noisy frame i share the same position / timestep embedding. + clean_tokens = self.patch_embedding(clean_x).flatten(2).transpose(1, 2) + hidden_states = torch.cat([clean_tokens, hidden_states], dim=1) + if aug_t is None: + aug_t = torch.zeros_like(timestep) + _, timestep_proj_clean, _, _ = self.condition_embedder( + aug_t.flatten(), encoder_hidden_states_text, None) + timestep_proj_clean = timestep_proj_clean.unflatten( + 1, (6, self.hidden_size)).unflatten(dim=0, sizes=timestep.shape) + timestep_proj = torch.cat([timestep_proj_clean, timestep_proj], dim=1) + freqs_cis = (torch.cat([freqs_cos, freqs_cos], dim=0), + torch.cat([freqs_sin, freqs_sin], dim=0)) + # 4. Transformer blocks if torch.is_grad_enabled() and self.gradient_checkpointing: for block in self.blocks: hidden_states = self._gradient_checkpointing_func( block, hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, - block_mask=self.block_mask) + block_mask=block_mask) else: for block in self.blocks: hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, - block_mask=self.block_mask) + block_mask=block_mask) + + if teacher_forcing: + # Drop the clean half; keep the denoised (noisy) tokens. + hidden_states = hidden_states[:, hidden_states.shape[1] // 2:] # 5. Output norm, projection & unpatchify temb = temb.unflatten(dim=0, sizes=timestep.shape).unsqueeze(2) diff --git a/fastvideo/tests/train/fixtures/wan_causal_t2v_causal_cd_min.yaml b/fastvideo/tests/train/fixtures/wan_causal_t2v_causal_cd_min.yaml new file mode 100644 index 000000000..f8a1459ad --- /dev/null +++ b/fastvideo/tests/train/fixtures/wan_causal_t2v_causal_cd_min.yaml @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# Minimum config to run a single training step of +# CausalConsistencyDistillationMethod (Causal-Forcing++ Stage-2b) on +# WanCausalModel for the per-method smoke test. Uses the real Wan 2.1 +# 1.3B checkpoint for both the trainable student and the frozen teacher +# (AR Euler-step target), with tiny synthetic latents. + +models: + student: + _target_: fastvideo.train.models.wan.WanCausalModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: true + teacher: + _target_: fastvideo.train.models.wan.WanCausalModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: false + ema: + _target_: fastvideo.train.models.wan.WanCausalModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: false + +method: + _target_: fastvideo.train.methods.consistency_model.causal_cd.CausalConsistencyDistillationMethod + discrete_cd_N: 12 + guidance_scale: 3.0 + ema_decay: 0.95 + +training: + dit_precision: bf16 + + distributed: + num_gpus: 1 + sp_size: 1 + tp_size: 1 + + data: + seed: 42 + train_batch_size: 1 + training_cfg_rate: 0.0 + num_latent_t: 6 + num_height: 64 + num_width: 64 + num_frames: 21 + + optimizer: + learning_rate: 2.0e-6 + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 1 + gradient_accumulation_steps: 1 + +pipeline: {} diff --git a/fastvideo/tests/train/fixtures/wan_causal_t2v_tfsft_min.yaml b/fastvideo/tests/train/fixtures/wan_causal_t2v_tfsft_min.yaml new file mode 100644 index 000000000..dd6b876f7 --- /dev/null +++ b/fastvideo/tests/train/fixtures/wan_causal_t2v_tfsft_min.yaml @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: Apache-2.0 +# Minimum config to run a single training step of +# TeacherForcingSFTMethod on WanCausalModel for the per-method smoke +# test. Identical to wan_causal_t2v_dfsft_min.yaml except the method, +# which feeds clean history to the causal transformer (teacher forcing). + +models: + student: + _target_: fastvideo.train.models.wan.WanCausalModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: true + +method: + _target_: fastvideo.train.methods.fine_tuning.tfsft.TeacherForcingSFTMethod + chunk_size: 3 + +training: + dit_precision: bf16 + + distributed: + num_gpus: 1 + sp_size: 1 + tp_size: 1 + + data: + seed: 42 + train_batch_size: 1 + training_cfg_rate: 0.0 + num_latent_t: 6 + num_height: 64 + num_width: 64 + num_frames: 21 + + optimizer: + learning_rate: 2.0e-6 + betas: [0.9, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 1 + gradient_accumulation_steps: 1 + +pipeline: {} diff --git a/fastvideo/tests/train/methods/test_wan_causal_cd.py b/fastvideo/tests/train/methods/test_wan_causal_cd.py new file mode 100644 index 000000000..42b9beb2d --- /dev/null +++ b/fastvideo/tests/train/methods/test_wan_causal_cd.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-method GPU smoke test: ``WanCausalModel`` + ``CausalConsistencyDistillationMethod``. + +Mirrors ``test_wan_causal_dfsft.py``. Causal Consistency Distillation +(Causal-Forcing++ Stage-2b) bootstraps a consistency MSE between the +student's ``x0`` at ``t`` and an EMA copy of the student at ``t_next``, +where ``t_next`` is produced online by a single CFG Euler step of a frozen +teacher (all under clean-history teacher forcing). This test exercises the +full step: finite loss, nonzero student gradients, frozen teacher, and a +post-step EMA update. +""" + +from __future__ import annotations + +import os + +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29519") + +from pathlib import Path + +import pytest +import torch + +from fastvideo.train.methods.consistency_model.causal_cd import ( + CausalConsistencyDistillationMethod, ) +from fastvideo.train.models.wan import WanCausalModel +from fastvideo.train.utils.config import load_run_config + + +_FIXTURE = str( + Path(__file__).resolve().parent.parent / "fixtures" + / "wan_causal_t2v_causal_cd_min.yaml") + + +def _build_synthetic_batch( + device: torch.device, + dtype: torch.dtype, +) -> dict[str, torch.Tensor]: + batch_size = 1 + return { + "text_embedding": + torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), + "text_attention_mask": + torch.ones(batch_size, 16, device=device), + "vae_latent": + torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype), + } + + +@pytest.mark.usefixtures("distributed_setup") +def test_wan_causal_cd_single_train_step( + monkeypatch: pytest.MonkeyPatch) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + cfg = load_run_config(_FIXTURE) + + device = torch.device("cuda:0") + dtype = torch.bfloat16 + + monkeypatch.setattr( + "fastvideo.train.utils.dataloader." + "build_parquet_t2v_train_dataloader", + lambda *args, **kwargs: None, + ) + + student = WanCausalModel( + init_from=cfg.models["student"]["init_from"], + training_config=cfg.training, + trainable=True, + ) + student.transformer = student.transformer.to(device=device, dtype=dtype) + + teacher = WanCausalModel( + init_from=cfg.models["teacher"]["init_from"], + training_config=cfg.training, + trainable=False, + ) + teacher.transformer = teacher.transformer.to(device=device, dtype=dtype) + + ema = WanCausalModel( + init_from=cfg.models["ema"]["init_from"], + training_config=cfg.training, + trainable=False, + ) + ema.transformer = ema.transformer.to(device=device, dtype=dtype) + + method = CausalConsistencyDistillationMethod( + cfg=cfg, + role_models={"student": student, "teacher": teacher, "ema": ema}, + ) + method.on_train_start() + + batch = _build_synthetic_batch(device, dtype) + loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0) + + loss = loss_map["total_loss"] + assert torch.is_tensor(loss), "total_loss must be a torch.Tensor" + assert torch.isfinite(loss).item(), ( + f"total_loss is not finite: {loss.item()}") + + method.backward(loss_map, outputs, grad_accum_rounds=1) + + blocks = student.transformer.blocks + assert blocks is not None and len(blocks) > 0 + layer0 = blocks[0] + trainable = [p for p in layer0.parameters() if p.requires_grad] + assert len(trainable) > 0, "student layer 0 has no trainable parameters" + for i, p in enumerate(trainable): + assert p.grad is not None, f"student layer 0 param[{i}] has None grad" + assert torch.isfinite(p.grad).all().item(), ( + f"student layer 0 param[{i}] grad contains NaN/Inf") + assert any(p.grad.detach().float().norm().item() > 0.0 for p in trainable), ( + "all student layer-0 grads are exactly zero; consistency loss " + "did not reach the first transformer block") + + # Teacher must stay frozen. + assert all(not p.requires_grad for p in teacher.transformer.parameters()), ( + "teacher must be frozen for Causal-CD") + + # The EMA model and student start from the same checkpoint, so the first + # parameter must match before any update. FSDP fully_shard params are + # DTensors; compare the local shards (torch.equal is unsupported on + # DTensor). + def _local(p: torch.Tensor) -> torch.Tensor: + return p.to_local() if hasattr(p, "to_local") else p + + ema_param = next(ema.transformer.parameters()) + student_param = next(student.transformer.parameters()) + assert torch.equal(_local(ema_param), _local(student_param)), ( + "EMA model should start identical to the student (same checkpoint)") + + # The EMA update must move EMA toward the student. Apply a visibly large + # perturbation so the bf16 lerp is well above rounding noise (the real + # optimizer step at lr=2e-6 would be sub-ULP in bf16). + with torch.no_grad(): + student_param.add_(1.0) + before = _local(ema_param).detach().float().clone() + method._update_ema() + after = _local(ema_param).detach().float() + assert not torch.equal(before, after), ( + "EMA weights did not move after _update_ema") + # EMA = decay*ema + (1-decay)*student moves ~ (1-decay) of the gap. + expected = before + (1.0 - method._ema_decay) * ( + _local(student_param).detach().float() - before) + assert torch.allclose(after, expected, atol=1e-2), ( + "EMA update did not follow the expected lerp") diff --git a/fastvideo/tests/train/methods/test_wan_causal_tfsft.py b/fastvideo/tests/train/methods/test_wan_causal_tfsft.py new file mode 100644 index 000000000..36c7d8e9d --- /dev/null +++ b/fastvideo/tests/train/methods/test_wan_causal_tfsft.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Per-method GPU smoke test: ``WanCausalModel`` + ``TeacherForcingSFTMethod``. + +Mirrors ``test_wan_causal_dfsft.py``. Teacher forcing concatenates a clean +context copy of every frame inside the causal transformer (``clean_x``) and +denoises the current block while attending to *clean* history. This test +exercises the ``clean_x`` path end-to-end: forward, finite loss, and nonzero +gradients reaching the first transformer block. +""" + +from __future__ import annotations + +import os + +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29518") + +from pathlib import Path + +import pytest +import torch + +from fastvideo.train.methods.fine_tuning.tfsft import ( + TeacherForcingSFTMethod, ) +from fastvideo.train.models.wan import WanCausalModel +from fastvideo.train.utils.config import load_run_config + + +_FIXTURE = str( + Path(__file__).resolve().parent.parent / "fixtures" + / "wan_causal_t2v_tfsft_min.yaml") + + +def _build_synthetic_batch( + device: torch.device, + dtype: torch.dtype, +) -> dict[str, torch.Tensor]: + batch_size = 1 + return { + "text_embedding": + torch.randn(batch_size, 16, 4096, device=device, dtype=dtype), + "text_attention_mask": + torch.ones(batch_size, 16, device=device), + "vae_latent": + torch.randn(batch_size, 16, 6, 8, 8, device=device, dtype=dtype), + } + + +@pytest.mark.usefixtures("distributed_setup") +def test_wan_causal_tfsft_single_train_step( + monkeypatch: pytest.MonkeyPatch) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + + cfg = load_run_config(_FIXTURE) + + device = torch.device("cuda:0") + dtype = torch.bfloat16 + + monkeypatch.setattr( + "fastvideo.train.utils.dataloader." + "build_parquet_t2v_train_dataloader", + lambda *args, **kwargs: None, + ) + + model = WanCausalModel( + init_from=cfg.models["student"]["init_from"], + training_config=cfg.training, + trainable=True, + ) + model.transformer = model.transformer.to(device=device, dtype=dtype) + + method = TeacherForcingSFTMethod( + cfg=cfg, + role_models={"student": model}, + ) + method.on_train_start() + + batch = _build_synthetic_batch(device, dtype) + loss_map, outputs, _metrics = method.single_train_step(batch, iteration=0) + + loss = loss_map["total_loss"] + assert torch.is_tensor(loss), "total_loss must be a torch.Tensor" + assert torch.isfinite(loss).item(), ( + f"total_loss is not finite: {loss.item()}") + + method.backward(loss_map, outputs, grad_accum_rounds=1) + + blocks = getattr(model.transformer, "blocks", None) + assert blocks is not None and len(blocks) > 0, ( + "CausalWanTransformer is expected to expose ``.blocks``") + layer0 = blocks[0] + + trainable = [p for p in layer0.parameters() if p.requires_grad] + assert len(trainable) > 0, "layer 0 has no trainable parameters" + + for i, p in enumerate(trainable): + assert p.grad is not None, f"layer 0 param[{i}] has None grad" + assert torch.isfinite(p.grad).all().item(), ( + f"layer 0 param[{i}] grad contains NaN/Inf") + + any_nonzero = any( + p.grad.detach().float().norm().item() > 0.0 for p in trainable) + assert any_nonzero, ( + "all layer-0 grads are exactly zero; backward did not " + "reach the first transformer block") + + # Teacher forcing must build its own (concatenated) attention mask and + # must not have constructed the diffusion-forcing mask. + assert model.transformer.teacher_forcing_block_mask is not None, ( + "teacher-forcing mask was not constructed") + assert model.transformer.block_mask is None, ( + "diffusion-forcing mask should not be built on the TF path") diff --git a/fastvideo/train/methods/__init__.py b/fastvideo/train/methods/__init__.py index c119550e6..0be93931e 100644 --- a/fastvideo/train/methods/__init__.py +++ b/fastvideo/train/methods/__init__.py @@ -9,6 +9,8 @@ __all__ = [ "KDMethod", "SelfForcingMethod", "DiffusionForcingSFTMethod", + "TeacherForcingSFTMethod", + "CausalConsistencyDistillationMethod", ] @@ -28,4 +30,11 @@ def __getattr__(name: str) -> object: if name == "DiffusionForcingSFTMethod": from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod return DiffusionForcingSFTMethod + if name == "TeacherForcingSFTMethod": + from fastvideo.train.methods.fine_tuning.tfsft import TeacherForcingSFTMethod + return TeacherForcingSFTMethod + if name == "CausalConsistencyDistillationMethod": + from fastvideo.train.methods.consistency_model.causal_cd import ( + CausalConsistencyDistillationMethod, ) + return CausalConsistencyDistillationMethod raise AttributeError(name) diff --git a/fastvideo/train/methods/consistency_model/__init__.py b/fastvideo/train/methods/consistency_model/__init__.py index 324710b84..227c26284 100644 --- a/fastvideo/train/methods/consistency_model/__init__.py +++ b/fastvideo/train/methods/consistency_model/__init__.py @@ -1,3 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 -__all__: list[str] = [] +from __future__ import annotations + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from fastvideo.train.methods.consistency_model.causal_cd import ( + CausalConsistencyDistillationMethod, ) + +__all__ = [ + "CausalConsistencyDistillationMethod", +] + + +def __getattr__(name: str) -> object: + if name == "CausalConsistencyDistillationMethod": + from fastvideo.train.methods.consistency_model.causal_cd import ( + CausalConsistencyDistillationMethod, ) + + return CausalConsistencyDistillationMethod + raise AttributeError(name) diff --git a/fastvideo/train/methods/consistency_model/causal_cd.py b/fastvideo/train/methods/consistency_model/causal_cd.py new file mode 100644 index 000000000..f509b38ed --- /dev/null +++ b/fastvideo/train/methods/consistency_model/causal_cd.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Causal Consistency Distillation (Causal-Forcing++ Stage-2b). + +ODE-data-free initialization for asymmetric DMD. Ported from +Causal-Forcing ``model/naive_consistency.py``. + +For a clean latent and a discrete flow-match schedule of ``discrete_cd_N`` +steps, sample an index ``i`` and form ``(t, t_next)``. A frozen *teacher* +takes a single CFG Euler step from ``latent_t`` to ``latent_t_next``. The +trainable *student* predicts ``x0`` at ``t`` and an EMA copy of the student +predicts ``x0`` at ``t_next``; the loss is their MSE. All three forwards run +under clean-history teacher forcing (``clean_x``), so the target is produced +online from ground-truth latents only -- no precomputed ODE pairs. + +Roles (all initialized from the same checkpoint, as in the reference): + * ``student`` -- trainable generator. + * ``teacher`` -- frozen; the single-step Euler ODE target. + * ``ema`` -- frozen; EMA of the student, the consistency target. +""" + +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn.functional as F + +from fastvideo.models.utils import pred_noise_to_pred_video +from fastvideo.train.methods.base import LogScalar, TrainingMethod +from fastvideo.train.models.base import ModelBase +from fastvideo.train.utils.optimizer import build_optimizer_and_scheduler + + +class CausalConsistencyDistillationMethod(TrainingMethod): + """Causal consistency distillation (Causal-Forcing++).""" + + def __init__( + self, + *, + cfg: Any, + role_models: dict[str, ModelBase], + ) -> None: + super().__init__(cfg=cfg, role_models=role_models) + + for role in ("student", "teacher", "ema"): + if role not in role_models: + raise ValueError(f"Causal-CD requires role {role!r} " + "(student trainable; teacher + ema frozen, " + "both initialized from the student's " + "checkpoint)") + if not self.student._trainable: + raise ValueError("Causal-CD requires student to be trainable") + self.teacher = role_models["teacher"] + self.ema_model = role_models["ema"] + + self._attn_kind = self._infer_attn_kind() + self._guidance_scale = float(self.method_config.get("guidance_scale", 3.0)) + self._discrete_cd_n = int(self.method_config.get("discrete_cd_N", 48)) + if self._discrete_cd_n < 2: + raise ValueError("method.discrete_cd_N must be >= 2") + self._ema_decay = float(self.method_config.get("ema_decay", 0.95)) + + self.student.init_preprocessors(self.training_config) + self._init_optimizers_and_schedulers() + + # ------------------------------------------------------------------ + + @property + def _optimizer_dict(self) -> dict[str, Any]: + return {"student": self._student_optimizer} + + @property + def _lr_scheduler_dict(self) -> dict[str, Any]: + return {"student": self._student_lr_scheduler} + + def get_optimizers(self, iteration: int) -> list[torch.optim.Optimizer]: + del iteration + return [self._student_optimizer] + + def get_lr_schedulers(self, iteration: int) -> list[Any]: + del iteration + return [self._student_lr_scheduler] + + # ------------------------------------------------------------------ + + def single_train_step( + self, + batch: dict[str, Any], + iteration: int, + ) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, LogScalar]]: + del iteration + training_batch = self.student.prepare_batch( + batch, + generator=self.cuda_generator, + latents_source="data", + ) + clean_latents = training_batch.latents + if not torch.is_tensor(clean_latents) or clean_latents.ndim != 5: + raise ValueError("Causal-CD expects [B, T, C, H, W] latents") + + batch_size, num_latents = int(clean_latents.shape[0]), int(clean_latents.shape[1]) + device = clean_latents.device + + cd_timesteps = self._cd_timesteps(device) + idx = int(torch.randint(0, self._discrete_cd_n - 1, (1, ), generator=self.cuda_generator, device=device).item()) + t = cd_timesteps[idx] + t_next = cd_timesteps[idx + 1] + t_pf = t * torch.ones(batch_size, num_latents, device=device) + t_next_pf = t_next * torch.ones(batch_size, num_latents, device=device) + + noise = torch.randn( + clean_latents.shape, + generator=self.cuda_generator, + device=device, + dtype=clean_latents.dtype, + ) + latent_t = self.student.add_noise(clean_latents, noise, t_pf.flatten()) + + # Teacher: single CFG Euler step latent_t -> latent_t_next (no grad). + with torch.no_grad(): + v_cond = self._predict_flow(self.teacher, + latent_t, + t_pf, + training_batch, + conditional=True, + clean_x=clean_latents) + v_uncond = self._predict_flow(self.teacher, + latent_t, + t_pf, + training_batch, + conditional=False, + clean_x=clean_latents) + v_pred = v_uncond + self._guidance_scale * (v_cond - v_uncond) + dt = ((t - t_next) / float(self.student.num_train_timesteps)) + latent_t_next = latent_t - dt * v_pred + + # Student x0 at t (with grad). + training_batch.timesteps = t_pf + flow_student = self._predict_flow(self.student, + latent_t, + t_pf, + training_batch, + conditional=True, + clean_x=clean_latents) + x0_t = self._to_x0(flow_student, latent_t, t_pf) + + # EMA-student x0 at t_next (no grad) -- the consistency target. + with torch.no_grad(): + flow_ema = self._predict_flow(self.ema_model, + latent_t_next, + t_next_pf, + training_batch, + conditional=True, + clean_x=clean_latents) + x0_t_next = self._to_x0(flow_ema, latent_t_next, t_next_pf) + + loss = F.mse_loss(x0_t.float(), x0_t_next.float()) + + loss_map = {"total_loss": loss, "causal_cd_loss": loss} + outputs: dict[str, Any] = {"_fv_backward": (t_pf, training_batch.attn_metadata)} + metrics: dict[str, LogScalar] = {} + return loss_map, outputs, metrics + + # ------------------------------------------------------------------ + + def backward( + self, + loss_map: dict[str, torch.Tensor], + outputs: dict[str, Any], + *, + grad_accum_rounds: int = 1, + ) -> None: + grad_accum_rounds = max(1, int(grad_accum_rounds)) + ctx = outputs.get("_fv_backward") + if ctx is None: + super().backward(loss_map, outputs, grad_accum_rounds=grad_accum_rounds) + return + self.student.backward(loss_map["total_loss"], ctx, grad_accum_rounds=grad_accum_rounds) + + def optimizers_schedulers_step(self, iteration: int) -> None: + super().optimizers_schedulers_step(iteration) + self._update_ema() + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _predict_flow( + self, + model: ModelBase, + latents: torch.Tensor, + timestep: torch.Tensor, + batch: Any, + *, + conditional: bool, + clean_x: torch.Tensor, + ) -> torch.Tensor: + """Run a (teacher-forced) flow prediction on ``model``.""" + return model.predict_noise(latents, + timestep, + batch, + conditional=conditional, + cfg_uncond=None, + attn_kind=self._attn_kind, + clean_x=clean_x) + + def _to_x0( + self, + flow: torch.Tensor, + latents: torch.Tensor, + timestep: torch.Tensor, + ) -> torch.Tensor: + return pred_noise_to_pred_video( + pred_noise=flow.flatten(0, 1), + noise_input_latent=latents.flatten(0, 1), + timestep=timestep, + scheduler=self.student.noise_scheduler, + ).unflatten(0, flow.shape[:2]) + + def _cd_timesteps(self, device: torch.device) -> torch.Tensor: + full = self.student.noise_scheduler.timesteps.to(device=device, dtype=torch.float32) + idx = torch.linspace(0, full.shape[0] - 1, self._discrete_cd_n, device=device).round().long() + return full[idx] + + @torch.no_grad() + def _update_ema(self) -> None: + decay = self._ema_decay + for ema_p, p in zip(self.ema_model.transformer.parameters(), self.student.transformer.parameters(), + strict=True): + ema_p.mul_(decay).add_(p.detach().to(ema_p.dtype), alpha=1.0 - decay) + + def _init_optimizers_and_schedulers(self) -> None: + tc = self.training_config + student_lr = float(tc.optimizer.learning_rate) + if student_lr <= 0.0: + raise ValueError("training.learning_rate must be > 0 for causal-cd") + student_params = [p for p in self.student.transformer.parameters() if p.requires_grad] + ( + self._student_optimizer, + self._student_lr_scheduler, + ) = build_optimizer_and_scheduler( + params=student_params, + optimizer_config=tc.optimizer, + loop_config=tc.loop, + learning_rate=student_lr, + betas=tc.optimizer.betas, + scheduler_name=str(tc.optimizer.lr_scheduler), + ) diff --git a/fastvideo/train/methods/fine_tuning/__init__.py b/fastvideo/train/methods/fine_tuning/__init__.py index 6f862df4e..96da179f7 100644 --- a/fastvideo/train/methods/fine_tuning/__init__.py +++ b/fastvideo/train/methods/fine_tuning/__init__.py @@ -7,10 +7,12 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod + from fastvideo.train.methods.fine_tuning.tfsft import TeacherForcingSFTMethod __all__ = [ "DiffusionForcingSFTMethod", "FineTuneMethod", + "TeacherForcingSFTMethod", ] @@ -25,4 +27,9 @@ def __getattr__(name: str) -> object: from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod return FineTuneMethod + if name == "TeacherForcingSFTMethod": + from fastvideo.train.methods.fine_tuning.tfsft import ( + TeacherForcingSFTMethod, ) + + return TeacherForcingSFTMethod raise AttributeError(name) diff --git a/fastvideo/train/methods/fine_tuning/dfsft.py b/fastvideo/train/methods/fine_tuning/dfsft.py index 7cfe5f4df..9ebe5f448 100644 --- a/fastvideo/train/methods/fine_tuning/dfsft.py +++ b/fastvideo/train/methods/fine_tuning/dfsft.py @@ -135,12 +135,11 @@ class DiffusionForcingSFTMethod(TrainingMethod): t_inhom.flatten(), ) - pred = self.student.predict_noise( + pred = self._predict_noise( noisy_latents, t_inhom, training_batch, - conditional=True, - attn_kind=self._attn_kind, + clean_latents, ) if bool(self.training_config.model.precondition_outputs): @@ -178,6 +177,28 @@ class DiffusionForcingSFTMethod(TrainingMethod): metrics: dict[str, LogScalar] = {} return loss_map, outputs, metrics + def _predict_noise( + self, + noisy_latents: torch.Tensor, + timestep: torch.Tensor, + training_batch: Any, + clean_latents: torch.Tensor, + ) -> torch.Tensor: + """Predict flow for the noisy latents. + + Diffusion forcing conditions on the *noisy* rollout only, so + ``clean_latents`` is unused here. The teacher-forcing variant + overrides this to feed clean context via ``clean_x``. + """ + del clean_latents + return self.student.predict_noise( + noisy_latents, + timestep, + training_batch, + conditional=True, + attn_kind=self._attn_kind, + ) + # TrainingMethod override: backward def backward( self, diff --git a/fastvideo/train/methods/fine_tuning/tfsft.py b/fastvideo/train/methods/fine_tuning/tfsft.py new file mode 100644 index 000000000..fe9211cd4 --- /dev/null +++ b/fastvideo/train/methods/fine_tuning/tfsft.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Teacher-forcing SFT method (TFSFT; algorithm layer). + +Stage-1 AR-diffusion training from Causal-Forcing. Identical to +:class:`DiffusionForcingSFTMethod` (inhomogeneous per-chunk timesteps, +flow loss, bsmntw weighting) except the causal transformer denoises the +current block while attending to *clean* history (``clean_x``) instead of +its own noisy rollout. This is the "Causal Forcing" namesake mechanism and +the recommended initialization for the downstream ODE / consistency stages. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from fastvideo.train.methods.fine_tuning.dfsft import ( + DiffusionForcingSFTMethod, ) + + +class TeacherForcingSFTMethod(DiffusionForcingSFTMethod): + """AR-diffusion SFT with clean-history teacher forcing.""" + + def _predict_noise( + self, + noisy_latents: torch.Tensor, + timestep: torch.Tensor, + training_batch: Any, + clean_latents: torch.Tensor, + ) -> torch.Tensor: + return self.student.predict_noise( + noisy_latents, + timestep, + training_batch, + conditional=True, + attn_kind=self._attn_kind, + clean_x=clean_latents, + ) diff --git a/fastvideo/train/models/wan/wan.py b/fastvideo/train/models/wan/wan.py index b3a202082..7e11c88d6 100644 --- a/fastvideo/train/models/wan/wan.py +++ b/fastvideo/train/models/wan/wan.py @@ -320,6 +320,8 @@ class WanModel(ModelBase): conditional: bool, cfg_uncond: dict[str, Any] | None = None, attn_kind: Literal["dense", "vsa"] = "dense", + clean_x: torch.Tensor | None = None, + aug_t: torch.Tensor | None = None, ) -> torch.Tensor: device_type = self.device.type dtype = self._get_training_dtype() @@ -347,7 +349,11 @@ class WanModel(ModelBase): current_timestep=batch.timesteps, attn_metadata=attn_metadata, ): - input_kwargs = (self._build_distill_input_kwargs(noisy_latents, timestep, text_dict)) + input_kwargs = (self._build_distill_input_kwargs(noisy_latents, + timestep, + text_dict, + clean_x=clean_x, + aug_t=aug_t)) transformer = self._get_transformer(timestep) pred_noise = transformer(**input_kwargs).permute(0, 2, 1, 3, 4) return pred_noise @@ -530,17 +536,24 @@ class WanModel(ModelBase): noise_input: torch.Tensor, timestep: torch.Tensor, text_dict: dict[str, torch.Tensor] | None, + clean_x: torch.Tensor | None = None, + aug_t: torch.Tensor | None = None, ) -> dict[str, Any]: if text_dict is None: raise ValueError("text_dict cannot be None for " "Wan distillation") - return { + kwargs: dict[str, Any] = { "hidden_states": noise_input.permute(0, 2, 1, 3, 4), "encoder_hidden_states": text_dict["encoder_hidden_states"], "encoder_attention_mask": text_dict["encoder_attention_mask"], "timestep": timestep, "return_dict": False, } + if clean_x is not None: + # Teacher forcing: clean context latents (+ optional aug timestep). + kwargs["clean_x"] = clean_x.permute(0, 2, 1, 3, 4) + kwargs["aug_t"] = aug_t + return kwargs def _get_transformer(self, timestep: torch.Tensor) -> torch.nn.Module: return self.transformer