From 03ef7ec62b69fbb0044cc94cec46884675947a8f Mon Sep 17 00:00:00 2001 From: bubbliiiing <3323290568@qq.com> Date: Fri, 23 Oct 2026 13:20:38 +0800 Subject: [PATCH] Update self forcing code --- scripts/wan2.1_self_forcing/train_distill.py | 157 ++++++++++++------ .../pipeline/pipeline_wan_self_forcing.py | 15 +- videox_fun/utils/__init__.py | 1 + videox_fun/utils/fsdp_ema.py | 7 +- videox_fun/utils/fsdp_utils.py | 101 +++++++++++ 5 files changed, 219 insertions(+), 62 deletions(-) create mode 100644 videox_fun/utils/fsdp_utils.py diff --git a/scripts/wan2.1_self_forcing/train_distill.py b/scripts/wan2.1_self_forcing/train_distill.py index 141cfa9..e9dcae1 100644 --- a/scripts/wan2.1_self_forcing/train_distill.py +++ b/scripts/wan2.1_self_forcing/train_distill.py @@ -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): diff --git a/videox_fun/pipeline/pipeline_wan_self_forcing.py b/videox_fun/pipeline/pipeline_wan_self_forcing.py index 2a7c9b2..4336b12 100644 --- a/videox_fun/pipeline/pipeline_wan_self_forcing.py +++ b/videox_fun/pipeline/pipeline_wan_self_forcing.py @@ -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: diff --git a/videox_fun/utils/__init__.py b/videox_fun/utils/__init__.py index 2f83453..1e5c107 100755 --- a/videox_fun/utils/__init__.py +++ b/videox_fun/utils/__init__.py @@ -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) diff --git a/videox_fun/utils/fsdp_ema.py b/videox_fun/utils/fsdp_ema.py index 18749ff..8bf2ece 100644 --- a/videox_fun/utils/fsdp_ema.py +++ b/videox_fun/utils/fsdp_ema.py @@ -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, diff --git a/videox_fun/utils/fsdp_utils.py b/videox_fun/utils/fsdp_utils.py new file mode 100644 index 0000000..b3a15be --- /dev/null +++ b/videox_fun/utils/fsdp_utils.py @@ -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)