nodes/aigc-apps-VideoX-Fun
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update self forcing code

Browse Source
This commit is contained in:
bubbliiiing
2026-10-23 13:20:38 +08:00
parent e7afc604e1
commit 03ef7ec62b
5 changed files with 219 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/wan2.1_self_forcing/train_distill.py
+108 -49
View File
@@ -3150,41 +3150,24 @@ def main():
critic_timestep
)
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
if use_gan:
# Generated and real latents share critic_noise / critic_timestep. The critic forward is
# split into three passes so the peak activation stays at 2x batch instead of 4x:
# A) fake + real for the denoising and GAN loss, B) real + R1 perturbation, C) fake + R2.
real_latents = clean_latents[:, :, :fake_score_denoised_pred.size(2)].to(fake_score_denoised_pred.dtype)
real_noisy = add_noise(real_latents, critic_noise, critic_timestep)
if args.r1_weight > 0:
r1_perturbation = args.r1_sigma * torch.randn(
real_noisy.shape, dtype=real_noisy.dtype, device=real_noisy.device, generator=torch_rng
)
if args.r2_weight > 0:
r2_perturbation = args.r2_sigma * torch.randn(
fake_score_denoised_input.shape, dtype=fake_score_denoised_input.dtype,
device=fake_score_denoised_input.device, generator=torch_rng
)
fake_score_denoised_output, gan_logits = fake_score_transformer3d(
x=torch.cat([fake_score_denoised_input, real_noisy]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
# Generated and real latents share critic_noise / critic_timestep. Prepare the
# real input and the R1/R2 perturbations once; they are consumed by the critic
# forward below and by the R1/R2 penalties after the main backward(s).
real_noisy = None
r1_perturbation = None
r2_perturbation = None
if use_gan:
real_latents = clean_latents[:, :, :fake_score_denoised_pred.size(2)].to(fake_score_denoised_pred.dtype)
real_noisy = add_noise(real_latents, critic_noise, critic_timestep)
del real_latents
if args.r1_weight > 0:
r1_perturbation = args.r1_sigma * torch.randn(
real_noisy.shape, dtype=real_noisy.dtype, device=real_noisy.device, generator=torch_rng
)
fake_score_denoised_output = fake_score_denoised_output.chunk(2)[0]
gan_logits = gan_logits.float().chunk(2)
else:
fake_score_denoised_output = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
if args.r2_weight > 0:
r2_perturbation = args.r2_sigma * torch.randn(
fake_score_denoised_input.shape, dtype=fake_score_denoised_input.dtype,
device=fake_score_denoised_input.device, generator=torch_rng
)
def custom_mse_loss(noise_pred, target, weighting=None, threshold=50):
@@ -3199,23 +3182,99 @@ def main():
final_loss = masked_loss.mean()
return final_loss
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
if use_gan and not args.relativistic_discriminator:
# Memory-optimized critic: the non-relativistic gan_d_loss =
# softplus(-real) + softplus(fake) is separable, so run fake and real as two
# sequential 1x-batch forwards and backward the fake pass first. Backwarding
# fake frees its gradient-checkpointing block inputs before the real forward,
# dropping the peak activation from 2x batch to 1x batch. Both backward(s)
# accumulate into fake_score.grad, giving the identical update as the combined
# 2x-batch pass. R1/R2 below stay as their own 2x-batch re-forwards.
# -- Fake pass (1x batch): denoising output + fake discriminator logit --
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
fake_score_denoised_output, fake_gan_logits = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
gan_classify=True,
)
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
critic_loss = denoising_loss
if use_gan:
fake_logit, real_logit = gan_logits[0], gan_logits[1]
if args.relativistic_discriminator:
gan_d_loss = F.softplus(-(real_logit - fake_logit)).mean()
else:
gan_d_loss = F.softplus(-real_logit).mean() + F.softplus(fake_logit).mean()
gan_d_loss = gan_d_loss * args.gan_d_weight
train_gan_d_loss += accelerator.gather(gan_d_loss.detach().repeat(args.train_batch_size)).mean().item()
critic_loss = critic_loss + gan_d_loss
fake_logit = fake_gan_logits.float()
fake_d_loss = F.softplus(fake_logit).mean() * args.gan_d_weight
critic_loss = denoising_loss + fake_d_loss
accelerator_fake_score_transformer3d.backward(critic_loss)
# Fake graph freed: the transformer blocks' stored inputs are released here.
fake_d_detached = fake_d_loss.detach()
del critic_loss, fake_gan_logits, fake_logit, fake_d_loss, fake_score_denoised_output
accelerator_fake_score_transformer3d.backward(critic_loss)
del critic_loss
# -- Real pass (1x batch): discriminator logit only --
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
_, real_gan_logits = fake_score_transformer3d(
x=real_noisy,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
gan_classify=True,
)
real_logit = real_gan_logits.float()
real_d_loss = F.softplus(-real_logit).mean() * args.gan_d_weight
accelerator_fake_score_transformer3d.backward(real_d_loss)
# Real graph freed. Combine the two detached halves for logging; gan_d_loss is
# kept alive because the progress-bar logging below reads it.
gan_d_loss = fake_d_detached + real_d_loss.detach()
train_gan_d_loss += accelerator.gather(gan_d_loss.repeat(args.train_batch_size)).mean().item()
del fake_d_detached, real_gan_logits, real_logit, real_d_loss
else:
# Combined path: 2x-batch forward for relativistic GAN (both logits coupled by
# softplus(-(real - fake))), or 1x-batch forward when GAN is disabled.
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
if use_gan:
fake_score_denoised_output, gan_logits = fake_score_transformer3d(
x=torch.cat([fake_score_denoised_input, real_noisy]),
context=repeat_batch(prompt_embeds, 2),
t=critic_timestep.repeat(2),
seq_len=seq_len,
y=repeat_batch(inpaint_latents, 2) if args.train_mode != "normal" else None,
clip_fea=repeat_batch(clip_context, 2) if args.train_mode != "normal" else None,
gan_classify=True,
)
fake_score_denoised_output = fake_score_denoised_output.chunk(2)[0]
gan_logits = gan_logits.float().chunk(2)
else:
fake_score_denoised_output = fake_score_transformer3d(
x=fake_score_denoised_input,
context=prompt_embeds,
t=critic_timestep,
seq_len=seq_len,
y=inpaint_latents if args.train_mode != "normal" else None,
clip_fea=clip_context if args.train_mode != "normal" else None,
)
denoising_loss = custom_mse_loss(fake_score_denoised_output, critic_noise - fake_score_denoised_pred)
avg_denoising_loss = accelerator.gather(denoising_loss.repeat(args.train_batch_size)).mean()
train_denoising_loss += avg_denoising_loss.item() / args.gradient_accumulation_steps
critic_loss = denoising_loss
if use_gan:
fake_logit, real_logit = gan_logits[0], gan_logits[1]
if args.relativistic_discriminator:
gan_d_loss = F.softplus(-(real_logit - fake_logit)).mean()
else:
gan_d_loss = F.softplus(-real_logit).mean() + F.softplus(fake_logit).mean()
gan_d_loss = gan_d_loss * args.gan_d_weight
train_gan_d_loss += accelerator.gather(gan_d_loss.detach().repeat(args.train_batch_size)).mean().item()
critic_loss = critic_loss + gan_d_loss
accelerator_fake_score_transformer3d.backward(critic_loss)
del critic_loss
if use_gan and args.r1_weight > 0:
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=accelerator.device):
videox_fun/pipeline/pipeline_wan_self_forcing.py
+8 -7
View File
@@ -868,13 +868,14 @@ class WanSelfForcingPipeline(DiffusionPipeline):
end_idx = start_idx + current_num_frames
noisy_input = noise[:, :, start_idx:end_idx]
# Denoising loop for current block
# Reset scheduler state for each block (required for causal generation)
# For Euler scheduler, resetting _step_index is sufficient.
# For multi-step schedulers (UniPC, DPM++), also clear accumulated model outputs.
self.scheduler._step_index = None
if hasattr(self.scheduler, 'model_outputs'):
self.scheduler.model_outputs = []
if not stochastic_sampling and isinstance(self.scheduler, FlowUniPCMultistepScheduler):
self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
timesteps = self.scheduler.timesteps
elif not stochastic_sampling and isinstance(self.scheduler, FlowDPMSolverMultistepScheduler):
self.scheduler.set_timesteps(device=device, sigmas=sampling_sigmas)
timesteps = self.scheduler.timesteps
else:
self.scheduler._step_index = None
denoise_timesteps = timesteps[:-1] if stochastic_sampling else timesteps
with self.progress_bar(total=len(denoise_timesteps)) as progress_bar:
videox_fun/utils/__init__.py
+1
View File
@@ -17,6 +17,7 @@ from .fp8_optimization import (autocast_model_forward,
convert_weight_dtype_wrapper,
replace_parameters_by_name)
from .fsdp_ema import FSDPEMA, LORAFSDPEMA
from .fsdp_utils import checkpoint_fsdp_forward
from .gpu_memory_mode import (GPU_MEMORY_MODES, QUANT_MODES,
apply_gpu_memory_mode, register_quant_mode,
split_gpu_memory_mode)
videox_fun/utils/fsdp_ema.py
+1 -6
View File
@@ -56,18 +56,13 @@ class FSDPEMA:
with torch.no_grad():
module.load_state_dict(source.state_dict(), strict=True)
# Wrap with the same kwargs as the live model goes through in
# `accelerator.prepare_model`, minus the mixed-precision policy: the
# EMA never runs forward/backward, and a `param_dtype` cast would
# quantize the shadow copy away from fp32 (see the note in the class
# docstring), so no casting is wanted here.
fsdp_plugin.set_auto_wrap_policy(module)
self.module = FSDP(
module,
sharding_strategy=fsdp_plugin.sharding_strategy or fsdp_plugin.reshard_after_forward,
cpu_offload=fsdp_plugin.cpu_offload,
auto_wrap_policy=fsdp_plugin.auto_wrap_policy,
mixed_precision=None,
mixed_precision=fsdp_plugin.mixed_precision_policy,
sync_module_states=fsdp_plugin.sync_module_states,
backward_prefetch=fsdp_plugin.backward_prefetch,
forward_prefetch=fsdp_plugin.forward_prefetch,
videox_fun/utils/fsdp_utils.py
+101
View File
@@ -0,0 +1,101 @@
"""Blockwise autograd replay for memory-bounded FSDP training."""
from contextlib import nullcontext
import torch
__all__ = ["checkpoint_fsdp_forward"]
def checkpoint_fsdp_forward(model, accumulated, *, context_fn=None, **kwargs):
"""Replay independent model forwards and accumulate parameter gradients.
Share a fresh ``accumulated`` dictionary across all blocks of one rollout.
Call ``loss.backward()`` once, then clip gradients and update the optimizer.
Existing model gradients are cleared during replay; do not mix this helper
with other gradient accumulation or direct parameter paths in the same loss.
Distributed use requires FSDP1 FULL_SHARD with CPUOffload,
``use_orig_params=True``, and gradient synchronization enabled. FSDP2,
``no_sync``, higher-order gradients, and repeated backward are unsupported.
CPU modules can use the same replay semantics without distributed wrapping.
Inputs are keyword arguments containing tensors or immutable metadata.
Tensors must be passed directly, not inside containers. Outputs must be a
tensor or a tuple of tensors. All CUDA inputs must use one device.
``context_fn`` optionally snapshots mutable forward state and returns
``(forward_context, recompute_context)``, as in non-reentrant checkpointing.
The model must not mutate parameters or buffers during replay.
"""
keys = tuple(kwargs)
values = tuple(kwargs.values())
tensors = [value for value in values if isinstance(value, torch.Tensor)]
if not tensors:
raise ValueError("FSDP replay requires at least one tensor input.")
device = next((value.device for value in tensors if value.is_cuda), tensors[0].device)
if device.type not in ("cpu", "cuda"):
raise ValueError(f"FSDP replay supports CPU and CUDA inputs, got {device}.")
if any(value.is_cuda and value.device != device for value in tensors):
raise ValueError("FSDP replay requires all CUDA inputs on one device.")
cuda_devices = [device.index] if device.type == "cuda" else []
class Replay(torch.autograd.Function):
@staticmethod
def forward(ctx, anchor, *args):
ctx.tensor_indices = [index for index, value in enumerate(args) if isinstance(value, torch.Tensor)]
ctx.other_args = {index: value for index, value in enumerate(args) if not isinstance(value, torch.Tensor)}
tensors = [args[index] for index in ctx.tensor_indices]
ctx.input_requires_grad = [value.requires_grad for value in tensors]
ctx.save_for_backward(*tensors)
forward_context, ctx.replay_context = context_fn() if context_fn is not None else (nullcontext(), nullcontext())
ctx.cpu_rng = torch.get_rng_state()
ctx.cuda_rng = torch.cuda.get_rng_state(device) if cuda_devices else None
ctx.autocast = torch.is_autocast_enabled(device.type)
ctx.autocast_dtype = torch.get_autocast_dtype(device.type)
ctx.set_materialize_grads(False)
with forward_context:
outputs = model(**dict(zip(keys, args)))
if not isinstance(outputs, torch.Tensor) and not (
isinstance(outputs, tuple) and all(isinstance(output, torch.Tensor) for output in outputs)
):
raise TypeError("FSDP replay requires a tensor or a tuple of tensors as output.")
return outputs
@staticmethod
def backward(ctx, *output_grads):
parameters = list(model.parameters())
# Clear gradient slots; the separate accumulator preserves earlier block gradients.
previous_grads = [accumulated.get(parameter) for parameter in parameters]
model.zero_grad(set_to_none=True)
tensors = [value.detach().requires_grad_(requires_grad)
for value, requires_grad in zip(ctx.saved_tensors, ctx.input_requires_grad)]
args = dict(ctx.other_args)
args.update(zip(ctx.tensor_indices, tensors))
with torch.random.fork_rng(devices=cuda_devices), torch.enable_grad(), \
torch.autocast(device.type, enabled=ctx.autocast, dtype=ctx.autocast_dtype):
torch.set_rng_state(ctx.cpu_rng)
if cuda_devices:
torch.cuda.set_rng_state(ctx.cuda_rng, device)
with ctx.replay_context:
outputs = model(**{key: args[index] for index, key in enumerate(keys)})
outputs = (outputs,) if isinstance(outputs, torch.Tensor) else outputs
active = [(output, grad) for output, grad in zip(outputs, output_grads)
if grad is not None and output.requires_grad]
if active:
torch.autograd.backward([output for output, _ in active], [grad for _, grad in active])
# FSDP completes reduce-scatter and offload before accumulating local shards.
for parameter, previous in zip(parameters, previous_grads):
if previous is not None:
if parameter.grad is None:
parameter.grad = previous
else:
parameter.grad.add_(previous)
if parameter.grad is not None:
accumulated[parameter] = parameter.grad.detach().clone()
result = [None] * len(keys)
for index, tensor in zip(ctx.tensor_indices, tensors):
result[index] = tensor.grad
return (None, *result)
with torch.autograd.graph.save_on_cpu():
return Replay.apply(torch.empty(0, device=device, requires_grad=True), *values)