wehub/nvlabs--longlive
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:31:40 +08:00
commit 4df75a30cd
234 changed files with 41407 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/diffusion.py
+574
View File
@@ -0,0 +1,574 @@
# Adopted from https://github.com/guandeh17/Self-Forcing
# SPDX-License-Identifier: Apache-2.0
from typing import Tuple
import random
import torch
from model.base import BaseModel
from pipeline import CausalDiffusionInferencePipeline
from utils.i2v_conditioning import _overwrite_i2v_context, _zero_i2v_context_timestep
from utils.wan_5b_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
class CausalDiffusion(BaseModel):
def __init__(self, args, device):
"""
Initialize the Diffusion loss module.
"""
super().__init__(args, device)
self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
if self.num_frame_per_block > 1:
self.generator.model.num_frame_per_block = self.num_frame_per_block
self.independent_first_frame = getattr(args, "independent_first_frame", False)
if self.independent_first_frame and not getattr(args, "i2v", False):
self.generator.model.independent_first_frame = True
if args.gradient_checkpointing:
self.generator.enable_gradient_checkpointing()
# Step 2: Initialize all hyperparameters
self.num_train_timestep = args.num_train_timestep
self.min_step = int(0.02 * self.num_train_timestep)
self.max_step = int(0.98 * self.num_train_timestep)
self.guidance_scale = args.guidance_scale
self.timestep_shift = getattr(args, "timestep_shift", 1.0)
self.teacher_forcing = getattr(args, "teacher_forcing", False)
# Noise augmentation in teacher forcing, we add small noise to clean context latents
self.noise_augmentation_max_timestep = getattr(args, "noise_augmentation_max_timestep", 0)
self.args = args
self.device = device
self.inference_pipeline = None
# Error recycling (SVI-style error buffer)
# When ``enable_position_bucketing`` is true, each rank holds a 2D
# buffer ``(local_block_position × timestep)``. The pos dimension only
# covers the LOCAL slice of the sequence this rank is responsible for
# (no cross-SP-rank pos sharing — those positions are simply not
# reachable by this rank during forward), so memory cost scales as
# ``num_blocks_global / sp_size`` instead of ``num_blocks_global``.
# ``global_block_offset`` is recorded for logging only.
# During the first ``buffer_warmup_iter`` global steps, errors are
# all-gathered across the DP group (ranks with the same SP rank but
# different DP replicas), so each rank's local pos buckets fill up
# ``dp_size`` × faster without any wasted bandwidth.
self.error_buffer = None
self.noise_error_buffer = None
self.er_num_blocks = 0 # local; >0 means 2D position-bucketed
self.er_block_offset = 0 # global block offset for THIS rank
er_cfg = getattr(args, "error_recycling", None)
if er_cfg is not None and getattr(er_cfg, "enabled", False):
from utils.error_buffer import build_error_buffer
cfg_dict = er_cfg if isinstance(er_cfg, dict) else dict(er_cfg)
cfg_dict.setdefault("num_train_timesteps", self.num_train_timestep)
sp_size = int(getattr(args, "sequence_parallel_size", 1) or 1)
if cfg_dict.get("enable_position_bucketing", False):
shape = list(getattr(args, "image_or_video_shape", [1, 0]))
total_frames = int(shape[1]) if len(shape) > 1 else 0
assert total_frames > 0 and self.num_frame_per_block > 0, (
"enable_position_bucketing=true requires "
"image_or_video_shape[1] and num_frame_per_block to be set."
)
num_blocks_global = total_frames // self.num_frame_per_block
assert num_blocks_global % sp_size == 0, (
f"num_blocks_global ({num_blocks_global}) must be divisible "
f"by sequence_parallel_size ({sp_size})."
)
self.er_num_blocks = num_blocks_global // sp_size # local
# Determine this rank's SP index → global block offset.
if sp_size > 1:
import torch.distributed as dist
if dist.is_initialized():
sp_rank = dist.get_rank() % sp_size
else:
sp_rank = 0
else:
sp_rank = 0
self.er_block_offset = sp_rank * self.er_num_blocks
# Shard timestep buckets across SP ranks: each SP rank only
# stores t_bucket % sp_size == sp_rank, cutting per-rank CPU
# memory by ~sp_size. This uses the same SP dimension that
# already splits positions in 2D mode, so both 1D and 2D
# follow one save/load pattern (per sp_rank).
import torch.distributed as dist
if sp_size > 1 and dist.is_initialized():
er_shard_rank = dist.get_rank() % sp_size
er_shard_size = sp_size
else:
er_shard_rank = 0
er_shard_size = 1
self.error_buffer = build_error_buffer(
cfg_dict, num_blocks=self.er_num_blocks,
global_block_offset=self.er_block_offset,
shard_rank=er_shard_rank, shard_size=er_shard_size,
)
self.noise_error_buffer = build_error_buffer(
cfg_dict, num_blocks=self.er_num_blocks,
global_block_offset=self.er_block_offset,
shard_rank=er_shard_rank, shard_size=er_shard_size,
)
self.er_context_inject_prob = float(cfg_dict.get("context_inject_prob", 0.9))
self.er_latent_inject_prob = float(cfg_dict.get("latent_inject_prob", 0.0))
self.er_noise_inject_prob = float(cfg_dict.get("noise_inject_prob", 0.0))
self.er_clean_prob = float(cfg_dict.get("clean_prob", 0.0))
self.er_clean_buffer_update_prob = float(cfg_dict.get("clean_buffer_update_prob", 0.1))
self.er_start_step = int(cfg_dict.get("start_step", 0))
self.er_buffer_warmup_iter = int(cfg_dict.get("buffer_warmup_iter", 50))
self.er_skip_block_0 = bool(cfg_dict.get("skip_block_0", False))
def _initialize_models(self, args, device):
model_name = getattr(args.model_kwargs, "model_name", "Wan2.2-TI2V-5B")
if "5B" not in model_name:
raise ValueError(f"Only Wan2.2-TI2V-5B is supported in this release, got {model_name}")
self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
self.generator.model.requires_grad_(True)
self.text_encoder = WanTextEncoder()
self.text_encoder.requires_grad_(False)
self.vae = WanVAEWrapper()
self.vae.requires_grad_(False)
self.scheduler = self.generator.get_scheduler()
self.scheduler.timesteps = self.scheduler.timesteps.to(device)
def generator_loss(
self,
image_or_video_shape,
conditional_dict: dict,
unconditional_dict: dict,
clean_latent: torch.Tensor,
initial_latent: torch.Tensor = None,
loss_mask: torch.Tensor = None,
loss_mask_global_valid_count: torch.Tensor = None,
global_step: int = None,
) -> Tuple[torch.Tensor, dict]:
"""
Generate image/videos from noise and compute the DMD loss.
The noisy input to the generator is backward simulated.
This removes the need of any datasets during distillation.
See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
Input:
- image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
- conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
- unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
- clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
- loss_mask: optional tensor of shape [B, F] with 1.0 for valid frames and 0.0 for padded frames.
Under Sequence Parallel this is already the local chunk.
- loss_mask_global_valid_count: optional scalar tensor with the total valid count across all SP ranks.
When provided (SP mode), used as the denominator instead of the local loss_mask.sum().
- global_step: current training step, used for error recycling delayed start.
Output:
- loss: a scalar tensor representing the generator loss.
- generator_log_dict: a dictionary containing the intermediate tensors for logging.
"""
batch_size, num_frame = image_or_video_shape[:2]
noise = torch.randn_like(clean_latent)
# Step 2: Randomly sample a timestep and add noise to denoiser inputs
index = self._get_timestep(
0,
self.scheduler.num_train_timesteps,
image_or_video_shape[0],
image_or_video_shape[1],
self.num_frame_per_block,
uniform_timestep=False
)
timestep = self.scheduler.timesteps[index].to(dtype=self.dtype, device=self.device)
context_latent = (
initial_latent
if getattr(self.args, "i2v", False) and initial_latent is not None
else None
)
context_frames = int(context_latent.shape[1]) if context_latent is not None else 0
if context_frames > 0:
if context_frames >= num_frame:
raise ValueError(
f"initial_latent has {context_frames} frames but training clip has {num_frame}."
)
timestep[:, :context_frames] = 0
# Step 2.5 & 3.5: Error recycling — clean_prob acts as a master switch.
# When clean_prob fires, skip ALL error injection and use pristine input;
# otherwise each injection type rolls its own probability and is gated
# only by whether the corresponding buffer has any samples (SVI behavior).
#
# NOTE on rank-sync: random.random() below is INTENTIONALLY independent
# across ranks. None of these decisions guard a collective call (we use
# SVI's pattern of unconditional all_gather + local random replay in
# Step 5), so per-rank divergence here only affects which slice of data
# gets corrupted on which rank — perfectly safe under DP+SP, and matches
# SVI's behavior exactly.
er_latent_injected = False
er_noise_injected = False
er_injected = False
er_use_clean = False
er_ready = (
self.error_buffer is not None
and (global_step is None or global_step >= self.er_start_step)
)
if er_ready and self.er_clean_prob > 0 and random.random() < self.er_clean_prob:
er_use_clean = True
# Noise error injection (SVI's noise_prob): corrupt noise input.
# training_target is then computed with the corrupted noise so the
# model learns to predict a self-correcting velocity (SVI Eq. logic).
noise_for_train = noise
if (
er_ready and not er_use_clean
and self.er_noise_inject_prob > 0
and not self.noise_error_buffer.is_empty()
and random.random() < self.er_noise_inject_prob
):
noise_for_train = self._inject_noise_error_buffer(
noise, index, batch_size, num_frame
)
er_noise_injected = True
# Latent error injection (SVI's latent_prob): corrupt clean_latent
# before noising. training_target keeps pointing to ORIGINAL clean_latent.
clean_latent_for_noise = clean_latent
if (
er_ready and not er_use_clean
and self.er_latent_inject_prob > 0
and not self.error_buffer.is_empty()
and random.random() < self.er_latent_inject_prob
):
clean_latent_for_noise = self._inject_latent_error_buffer(
clean_latent, index, batch_size, num_frame
)
er_latent_injected = True
noisy_latents = self.scheduler.add_noise(
clean_latent_for_noise.flatten(0, 1),
noise_for_train.flatten(0, 1),
timestep.flatten(0, 1)
).unflatten(0, (batch_size, num_frame))
training_target = self.scheduler.training_target(clean_latent, noise_for_train, timestep)
if context_frames > 0:
noisy_latents[:, :context_frames] = context_latent.to(
device=noisy_latents.device,
dtype=noisy_latents.dtype,
)
training_target[:, :context_frames] = 0
# Step 3: Noise augmentation, also add small noise to clean context latents
if self.noise_augmentation_max_timestep > 0:
index_clean_aug = self._get_timestep(
0,
self.noise_augmentation_max_timestep,
image_or_video_shape[0],
image_or_video_shape[1],
self.num_frame_per_block,
uniform_timestep=False
)
timestep_clean_aug = self.scheduler.timesteps[index_clean_aug].to(dtype=self.dtype, device=self.device)
clean_latent_aug = self.scheduler.add_noise(
clean_latent.flatten(0, 1),
noise.flatten(0, 1),
timestep_clean_aug.flatten(0, 1)
).unflatten(0, (batch_size, num_frame))
else:
clean_latent_aug = clean_latent
timestep_clean_aug = None
# Step 3.5: Error recycling — inject sampled errors into clean prefix.
# 2D mode: per-position (random timestep). 1D mode: SVI global sampling.
if (
er_ready and not er_use_clean
and not self.error_buffer.is_empty()
and random.random() < self.er_context_inject_prob
):
clean_latent_aug = self._inject_error_buffer(
clean_latent_aug, index, batch_size, num_frame
)
er_injected = True
if context_frames > 0:
clean_latent_aug = _overwrite_i2v_context(
clean_latent_aug, context_latent, context_frames
)
if timestep_clean_aug is not None:
timestep_clean_aug = _zero_i2v_context_timestep(
timestep_clean_aug, context_frames
)
# Compute loss
flow_pred, x0_pred = self.generator(
noisy_image_or_video=noisy_latents,
conditional_dict=conditional_dict,
timestep=timestep,
clean_x=clean_latent_aug if self.teacher_forcing else None,
aug_t=timestep_clean_aug if self.teacher_forcing else None,
)
loss = torch.nn.functional.mse_loss(
flow_pred.float(), training_target.float(), reduction='none'
).mean(dim=(2, 3, 4))
loss = loss * self.scheduler.training_weight(timestep).unflatten(0, (batch_size, num_frame))
if context_frames > 0:
if loss_mask is None:
loss_mask = torch.ones(
(batch_size, num_frame),
device=loss.device,
dtype=loss.dtype,
)
loss_mask[:, :context_frames] = 0
if loss_mask is not None:
loss = loss * loss_mask
valid_count = loss_mask_global_valid_count if loss_mask_global_valid_count is not None else loss_mask.sum()
loss = loss.sum() / valid_count.clamp(min=1.0)
else:
loss = loss.mean()
log_dict = {
"x0": clean_latent.detach(),
"x0_pred": x0_pred.detach()
}
# Step 5: Store prediction errors into error buffer.
#
# SVI-style two-phase pattern (avoids the rank-divergent collective
# deadlock that an ``if random.random() < p: all_gather()`` would
# introduce):
# PHASE A (collective, UNCONDITIONAL): every rank reaches the
# all_gather call together so NCCL stays in sync.
# PHASE B (local, GATED): each rank independently decides whether
# to actually replay the gathered items into its buffer. The
# gate uses random.random() per rank — divergence here is fine
# because no further collective follows.
if self.error_buffer is not None:
with torch.no_grad():
# latent error: x0_pred - clean_latent ≡ -σ(v_pred - v_gt) — used for context/latent injection
latent_err = x0_pred.detach() - clean_latent.detach()
# noise error: SVI definition is (1-σ)(v_pred - v_gt) so the buffer entry,
# when later added directly to noise, equals ε_pred - ε_gt.
sigma = self.scheduler.sigmas.to(flow_pred.device)[index].reshape(
batch_size, num_frame, 1, 1, 1
).to(flow_pred.dtype)
noise_err = (flow_pred.detach() - training_target.detach()) * (1.0 - sigma)
use_distributed = (
global_step is not None
and global_step <= self.er_buffer_warmup_iter
)
# === PHASE A: collective — runs on EVERY rank, no gating ===
if use_distributed:
lat_items = self._gather_errors_for_buffer(
self.error_buffer, latent_err, index, batch_size, num_frame
)
noise_items = self._gather_errors_for_buffer(
self.noise_error_buffer, noise_err, index, batch_size, num_frame
)
else:
lat_items = self._collect_local_items(
self.error_buffer, latent_err, index, batch_size, num_frame
)
noise_items = self._collect_local_items(
self.noise_error_buffer, noise_err, index, batch_size, num_frame
)
# === PHASE B: local replay — random.random() per-rank is OK ===
# When the input was clean (low-error), only update buffer with
# small probability to avoid flooding it with near-zero samples
# (SVI: clean_buffer_update_prob).
should_update = True
if er_use_clean and random.random() >= self.er_clean_buffer_update_prob:
should_update = False
if should_update:
self._apply_gathered_items(self.error_buffer, lat_items)
self._apply_gathered_items(self.noise_error_buffer, noise_items)
buf_stats = self.error_buffer.stats()
noise_buf_stats = self.noise_error_buffer.stats()
log_dict["er_total_added"] = buf_stats["total_added"]
log_dict["er_filled_buckets"] = buf_stats["filled_buckets"]
log_dict["er_total_entries"] = buf_stats["total_entries"]
log_dict["er_noise_total_entries"] = noise_buf_stats["total_entries"]
log_dict["er_injected"] = er_injected
log_dict["er_latent_injected"] = er_latent_injected
log_dict["er_noise_injected"] = er_noise_injected
return loss, log_dict
def _inject_error_buffer(self, clean_latent_aug, index, batch_size, num_frame):
"""Inject errors into the clean prefix (E_img).
2D (position-bucketed): the i-th LOCAL prefix block draws from
``buckets[(i, *)]`` with a RANDOM timestep — the clean prefix is
the product of full ODE integration so its accumulated error can
come from any noise level, but its magnitude scales with the
block's global position. Note ``skip_block_0`` is interpreted in
the GLOBAL frame: only the very first SP rank may skip its block 0.
1D (timestep-bucketed): falls back to SVI ``sample_global``.
"""
block_size = self.num_frame_per_block
num_blocks = num_frame // block_size
result = clean_latent_aug.clone()
for b in range(batch_size):
for blk in range(num_blocks):
if self.er_skip_block_0 and (self.er_block_offset + blk) == 0:
continue
if self.er_num_blocks > 0:
err = self.error_buffer.sample_pos_any_t(
blk, device=result.device, dtype=result.dtype
)
else:
err = self.error_buffer.sample_global(
device=result.device, dtype=result.dtype
)
if err is not None:
start = blk * block_size
end = start + block_size
result[b, start:end] = result[b, start:end] + err
return result
def _inject_latent_error_buffer(self, clean_latent, index, batch_size, num_frame):
"""Inject errors into clean_latent before noising (E_vid).
Matches BOTH block_position (LOCAL) and timestep when the buffer is
2D, else only timestep (SVI default).
"""
block_size = self.num_frame_per_block
num_blocks = num_frame // block_size
index_per_block = index[:, ::block_size]
result = clean_latent.clone()
for b in range(batch_size):
for blk in range(num_blocks):
t_idx = index_per_block[b, blk].item()
pos = blk if self.er_num_blocks > 0 else None
err = self.error_buffer.sample(
t_idx, device=result.device, dtype=result.dtype,
block_pos=pos,
)
if err is not None:
start = blk * block_size
end = start + block_size
result[b, start:end] = result[b, start:end] + err
return result
def _inject_noise_error_buffer(self, noise, index, batch_size, num_frame):
"""Inject errors into the noise (E_noise).
Same matching strategy as ``_inject_latent_error_buffer`` but reads
from the dedicated noise buffer.
"""
block_size = self.num_frame_per_block
num_blocks = num_frame // block_size
index_per_block = index[:, ::block_size]
result = noise.clone()
for b in range(batch_size):
for blk in range(num_blocks):
t_idx = index_per_block[b, blk].item()
pos = blk if self.er_num_blocks > 0 else None
err = self.noise_error_buffer.sample(
t_idx, device=result.device, dtype=result.dtype,
block_pos=pos,
)
if err is not None:
start = blk * block_size
end = start + block_size
result[b, start:end] = result[b, start:end] + err
return result
def _gather_errors_for_buffer(
self, buffer, error, index, batch_size, num_frame
):
"""All-gather errors/timesteps across the appropriate group and return
a list of ready-to-add ``(err_block, t_idx, pos_or_None)`` items.
★ This is a COLLECTIVE — every rank MUST reach this call together.
The caller is responsible for invoking it unconditionally during the
warmup window (just like SVI's ``all_gather`` outside the random
``if`` blocks). Random decisions about whether to actually consume
the returned items belong to ``_apply_gathered_items`` instead.
Group selection mirrors SVI's intent:
* **2D (num_blocks > 0)** — DP group only. Other SP ranks' samples
map to positions unreachable by this rank, so cross-SP gather
wastes bandwidth.
* **1D (num_blocks == 0)** — WORLD group (SVI default). Buckets
are pos-agnostic so every rank's errors are valid samples.
"""
import torch.distributed as dist
if not dist.is_initialized() or dist.get_world_size() <= 1:
return self._collect_local_items(buffer, error, index, batch_size, num_frame)
if buffer.num_blocks > 0:
from wan_5b.distributed.sp_training import get_data_parallel_group
comm_group = get_data_parallel_group()
if comm_group is None:
return self._collect_local_items(buffer, error, index, batch_size, num_frame)
comm_size = dist.get_world_size(comm_group)
else:
comm_group = None
comm_size = dist.get_world_size()
if comm_size <= 1:
return self._collect_local_items(buffer, error, index, batch_size, num_frame)
err_local = error.detach().contiguous()
idx_local = index.detach().contiguous()
err_list = [torch.empty_like(err_local) for _ in range(comm_size)]
idx_list = [torch.empty_like(idx_local) for _ in range(comm_size)]
if comm_group is None:
dist.all_gather(err_list, err_local)
dist.all_gather(idx_list, idx_local)
else:
dist.all_gather(err_list, err_local, group=comm_group)
dist.all_gather(idx_list, idx_local, group=comm_group)
block_size = self.num_frame_per_block
num_blocks = num_frame // block_size
items = []
for err_r, idx_r in zip(err_list, idx_list):
idx_per_block = idx_r[:, ::block_size]
err_blocks = err_r.reshape(
batch_size, num_blocks, block_size, *err_r.shape[2:]
)
for b in range(batch_size):
for blk in range(num_blocks):
pos = blk if buffer.num_blocks > 0 else None
items.append((err_blocks[b, blk], idx_per_block[b, blk].item(), pos))
return items
def _collect_local_items(self, buffer, error, index, batch_size, num_frame):
"""Same item-list format as ``_gather_errors_for_buffer`` but with no
collective — used outside the warmup window or when distributed is off."""
block_size = self.num_frame_per_block
num_blocks = num_frame // block_size
idx_per_block = index[:, ::block_size]
error_blocks = error.reshape(
batch_size, num_blocks, block_size, *error.shape[2:]
)
items = []
for b in range(batch_size):
for blk in range(num_blocks):
pos = blk if buffer.num_blocks > 0 else None
items.append((error_blocks[b, blk], idx_per_block[b, blk].item(), pos))
return items
def _apply_gathered_items(self, buffer, items):
"""Pure local: drop ``items`` into ``buffer``. No collective, no
cross-rank coordination — each rank may invoke this independently
(or skip it entirely) without risking a deadlock."""
for err_block, t_idx, pos in items:
buffer.add(err_block, t_idx, block_pos=pos)
def _initialize_inference_pipeline(self):
"""
Lazy initialize the inference pipeline during the first backward simulation run.
Here we encapsulate the inference code with a model-dependent outside function.
We pass our FSDP-wrapped modules into the pipeline to save memory.
"""
self.inference_pipeline = CausalDiffusionInferencePipeline(
args=self.args,
device=self.device,
generator=self.generator,
text_encoder=self.text_encoder,
vae=self.vae
)