[WIP][Feat] distillation
This commit is contained in:
Jinzhe Pan
@@ -34,7 +34,7 @@ class FastVideoArgs:
|
||||
# Distributed executor backend
|
||||
distributed_executor_backend: str = "mp"
|
||||
|
||||
inference_mode: bool = True # if False == training mode
|
||||
mode: str = "inference" # Options: "inference", "training", "distill"
|
||||
|
||||
# HuggingFace specific parameters
|
||||
trust_remote_code: bool = False
|
||||
@@ -111,7 +111,15 @@ class FastVideoArgs:
|
||||
|
||||
@property
|
||||
def training_mode(self) -> bool:
|
||||
return not self.inference_mode
|
||||
return self.mode == "training"
|
||||
|
||||
@property
|
||||
def distill_mode(self) -> bool:
|
||||
return self.mode == "distill"
|
||||
|
||||
@property
|
||||
def inference_mode(self) -> bool:
|
||||
return self.mode == "inference"
|
||||
|
||||
def __post_init__(self):
|
||||
self.check_fastvideo_args()
|
||||
@@ -146,10 +154,11 @@ class FastVideoArgs:
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--inference-mode",
|
||||
action=StoreBoolean,
|
||||
default=FastVideoArgs.inference_mode,
|
||||
help="Whether to use inference mode",
|
||||
"--mode",
|
||||
type=str,
|
||||
default=FastVideoArgs.mode,
|
||||
choices=["inference", "training", "distill"],
|
||||
help="The mode to use",
|
||||
)
|
||||
|
||||
# HuggingFace specific parameters
|
||||
|
||||
@@ -94,6 +94,12 @@ class ComposedPipelineBase(ABC):
|
||||
self.initialize_validation_pipeline(self.training_args)
|
||||
self.initialize_training_pipeline(self.training_args)
|
||||
|
||||
if fastvideo_args.distill_mode:
|
||||
self.initialize_distillation_pipeline(fastvideo_args)
|
||||
|
||||
if fastvideo_args.log_validation:
|
||||
self.initialize_validation_pipeline(fastvideo_args)
|
||||
|
||||
self.initialize_pipeline(fastvideo_args)
|
||||
|
||||
if not fastvideo_args.training_mode:
|
||||
@@ -109,6 +115,10 @@ class ComposedPipelineBase(ABC):
|
||||
"if log_validation is True, the pipeline must implement this method"
|
||||
)
|
||||
|
||||
def initialize_distillation_pipeline(self, fastvideo_args: FastVideoArgs):
|
||||
raise NotImplementedError(
|
||||
"if distill_mode is True, the pipeline must implement this method")
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(cls,
|
||||
model_path: str,
|
||||
@@ -148,9 +158,11 @@ class ComposedPipelineBase(ABC):
|
||||
config_args = shallow_asdict(config)
|
||||
config_args.update(kwargs)
|
||||
|
||||
if args is None or args.inference_mode:
|
||||
fastvideo_args = FastVideoArgs(model_path=model_path, **config_args)
|
||||
|
||||
if args.mode == "inference":
|
||||
fastvideo_args = FastVideoArgs(model_path=model_path,
|
||||
device_str=device or "cuda" if
|
||||
torch.cuda.is_available() else "cpu",
|
||||
**config_args)
|
||||
fastvideo_args.model_path = model_path
|
||||
for key, value in config_args.items():
|
||||
setattr(fastvideo_args, key, value)
|
||||
@@ -164,7 +176,7 @@ class ComposedPipelineBase(ABC):
|
||||
|
||||
fastvideo_args.use_cpu_offload = False
|
||||
# make sure we are in training mode
|
||||
fastvideo_args.inference_mode = False
|
||||
fastvideo_args.mode = args.mode
|
||||
# we hijack the precision to be the master weight type so that the
|
||||
# model is loaded with the correct precision. Subsequently we will
|
||||
# use FSDP2's MixedPrecisionPolicy to set the precision for the
|
||||
|
||||
@@ -0,0 +1,766 @@
|
||||
import gc
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import deque
|
||||
from copy import deepcopy
|
||||
|
||||
import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision
|
||||
import wandb
|
||||
from diffusers.optimization import get_scheduler
|
||||
from einops import rearrange
|
||||
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
|
||||
from torch.distributed.fsdp import ShardingStrategy
|
||||
from torchdata.stateful_dataloader import StatefulDataLoader
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from fastvideo.distill.solver import EulerSolver, extract_into_tensor
|
||||
from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_dit_input
|
||||
from fastvideo.utils.checkpoint import save_checkpoint_v1
|
||||
from fastvideo.v1.configs.sample import SamplingParam
|
||||
from fastvideo.v1.dataset.parquet_datasets import ParquetVideoTextDataset
|
||||
from fastvideo.v1.distributed import cleanup_dist_env_and_memory, get_sp_group
|
||||
from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
|
||||
from fastvideo.v1.forward_context import set_forward_context
|
||||
from fastvideo.v1.logger import init_logger
|
||||
from fastvideo.v1.pipelines import ComposedPipelineBase
|
||||
from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.v1.pipelines.training_utils import (
|
||||
_clip_grad_norm_while_handling_failing_dtensor_cases)
|
||||
from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
# from: https://github.com/genmoai/models/blob/075b6e36db58f1242921deff83a1066887b9c9e1/src/mochi_preview/infer.py#L77
|
||||
def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None):
|
||||
if linear_steps is None:
|
||||
linear_steps = num_steps // 2
|
||||
linear_sigma_schedule = [
|
||||
i * threshold_noise / linear_steps for i in range(linear_steps)
|
||||
]
|
||||
threshold_noise_step_diff = linear_steps - threshold_noise * num_steps
|
||||
quadratic_steps = num_steps - linear_steps
|
||||
quadratic_coef = threshold_noise_step_diff / (linear_steps *
|
||||
quadratic_steps**2)
|
||||
linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (
|
||||
quadratic_steps**2)
|
||||
const = quadratic_coef * (linear_steps**2)
|
||||
quadratic_sigma_schedule = [
|
||||
quadratic_coef * (i**2) + linear_coef * i + const
|
||||
for i in range(linear_steps, num_steps)
|
||||
]
|
||||
sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule
|
||||
sigma_schedule = [1.0 - x for x in sigma_schedule]
|
||||
return sigma_schedule
|
||||
|
||||
|
||||
def reshard_fsdp(model):
|
||||
"""Reshard FSDP model for EMA updates."""
|
||||
for m in FSDP.fsdp_modules(model):
|
||||
if m._has_params and m.sharding_strategy is not ShardingStrategy.NO_SHARD:
|
||||
torch.distributed.fsdp._runtime_utils._reshard(m, m._handle, True)
|
||||
|
||||
|
||||
def get_norm(model_pred, norms, gradient_accumulation_steps):
|
||||
"""Calculate and aggregate model prediction norms."""
|
||||
fro_norm = (
|
||||
torch.linalg.matrix_norm(model_pred, ord="fro") / # codespell:ignore
|
||||
gradient_accumulation_steps)
|
||||
largest_singular_value = (torch.linalg.matrix_norm(model_pred, ord=2) /
|
||||
gradient_accumulation_steps)
|
||||
absolute_mean = torch.mean(
|
||||
torch.abs(model_pred)) / gradient_accumulation_steps
|
||||
absolute_max = torch.max(
|
||||
torch.abs(model_pred)) / gradient_accumulation_steps
|
||||
|
||||
sp_group = get_sp_group()
|
||||
sp_group.all_reduce(fro_norm, op=torch.distributed.ReduceOp.AVG)
|
||||
sp_group.all_reduce(largest_singular_value,
|
||||
op=torch.distributed.ReduceOp.AVG)
|
||||
sp_group.all_reduce(absolute_mean, op=torch.distributed.ReduceOp.AVG)
|
||||
|
||||
norms["fro"] += torch.mean(fro_norm).item() # codespell:ignore
|
||||
norms["largest singular value"] += torch.mean(largest_singular_value).item()
|
||||
norms["absolute mean"] += absolute_mean.item()
|
||||
norms["absolute max"] += absolute_max.item()
|
||||
|
||||
|
||||
class DistillationPipeline(ComposedPipelineBase, ABC):
|
||||
"""
|
||||
A pipeline for distillation training. All distillation pipelines should inherit from this class.
|
||||
"""
|
||||
_required_config_modules = ["scheduler", "transformer"]
|
||||
|
||||
def initialize_distillation_pipeline(self, fastvideo_args: TrainingArgs):
|
||||
logger.info("Initializing distillation pipeline...")
|
||||
self.device = fastvideo_args.device
|
||||
self.sp_group = get_sp_group()
|
||||
self.world_size = self.sp_group.world_size
|
||||
self.rank = self.sp_group.rank
|
||||
self.local_rank = self.sp_group.local_rank
|
||||
|
||||
# Initialize student model
|
||||
self.transformer = self.get_module("transformer")
|
||||
assert self.transformer is not None
|
||||
|
||||
self.transformer.requires_grad_(True)
|
||||
self.transformer.train()
|
||||
|
||||
# Initialize teacher model
|
||||
self.teacher_transformer = deepcopy(self.transformer)
|
||||
self.teacher_transformer.requires_grad_(False)
|
||||
|
||||
# Initialize EMA model if needed
|
||||
if fastvideo_args.use_ema:
|
||||
self.ema_transformer = deepcopy(self.transformer)
|
||||
self.ema_transformer.requires_grad_(False)
|
||||
else:
|
||||
self.ema_transformer = None
|
||||
|
||||
args = fastvideo_args
|
||||
noise_scheduler = self.get_module("scheduler")
|
||||
assert noise_scheduler is not None
|
||||
|
||||
# Initialize solver for distillation
|
||||
if args.scheduler_type == "pcm_linear_quadratic":
|
||||
linear_steps = int(noise_scheduler.config.num_train_timesteps *
|
||||
args.linear_range)
|
||||
sigmas = linear_quadratic_schedule(
|
||||
noise_scheduler.config.num_train_timesteps,
|
||||
args.linear_quadratic_threshold,
|
||||
linear_steps,
|
||||
)
|
||||
sigmas = torch.tensor(sigmas).to(dtype=torch.float32)
|
||||
else:
|
||||
sigmas = noise_scheduler.sigmas
|
||||
|
||||
self.solver = EulerSolver(
|
||||
sigmas.numpy()[::-1],
|
||||
noise_scheduler.config.num_train_timesteps,
|
||||
euler_timesteps=args.num_euler_timesteps,
|
||||
)
|
||||
self.solver.to(self.device)
|
||||
|
||||
# Setup optimizer
|
||||
params_to_optimize = self.transformer.parameters()
|
||||
params_to_optimize = list(
|
||||
filter(lambda p: p.requires_grad, params_to_optimize))
|
||||
|
||||
optimizer = torch.optim.AdamW(
|
||||
params_to_optimize,
|
||||
lr=args.learning_rate,
|
||||
betas=(0.9, 0.999),
|
||||
weight_decay=args.weight_decay,
|
||||
eps=1e-8,
|
||||
)
|
||||
|
||||
init_steps = 0
|
||||
logger.info("optimizer: %s", optimizer)
|
||||
|
||||
# Setup lr scheduler
|
||||
lr_scheduler = get_scheduler(
|
||||
args.lr_scheduler,
|
||||
optimizer=optimizer,
|
||||
num_warmup_steps=args.lr_warmup_steps * self.world_size,
|
||||
num_training_steps=args.max_train_steps * self.world_size,
|
||||
num_cycles=args.lr_num_cycles,
|
||||
power=args.lr_power,
|
||||
last_epoch=init_steps - 1,
|
||||
)
|
||||
|
||||
# Setup dataset
|
||||
train_dataset = ParquetVideoTextDataset(
|
||||
args.data_path,
|
||||
batch_size=args.train_batch_size,
|
||||
rank=self.rank,
|
||||
world_size=self.world_size,
|
||||
cfg_rate=args.cfg,
|
||||
num_latent_t=args.num_latent_t)
|
||||
|
||||
train_dataloader = StatefulDataLoader(
|
||||
train_dataset,
|
||||
batch_size=args.train_batch_size,
|
||||
num_workers=args.dataloader_num_workers,
|
||||
prefetch_factor=2,
|
||||
shuffle=False,
|
||||
pin_memory=True,
|
||||
drop_last=True)
|
||||
|
||||
self.lr_scheduler = lr_scheduler
|
||||
self.train_dataset = train_dataset
|
||||
self.train_dataloader = train_dataloader
|
||||
self.init_steps = init_steps
|
||||
self.optimizer = optimizer
|
||||
self.noise_scheduler = noise_scheduler
|
||||
|
||||
# Get unconditional embeddings
|
||||
self.uncond_prompt_embed = train_dataset.uncond_prompt_embed
|
||||
self.uncond_prompt_mask = train_dataset.uncond_prompt_mask
|
||||
|
||||
if self.rank <= 0:
|
||||
project = args.tracker_project_name or "fastvideo"
|
||||
wandb.init(project=project, config=args)
|
||||
|
||||
@abstractmethod
|
||||
def initialize_validation_pipeline(self, fastvideo_args: FastVideoArgs):
|
||||
raise NotImplementedError(
|
||||
"Distillation pipelines must implement this method")
|
||||
|
||||
@abstractmethod
|
||||
def distill_one_step(self, transformer, model_type, teacher_transformer,
|
||||
ema_transformer, optimizer, lr_scheduler, loader_iter,
|
||||
noise_scheduler, solver, noise_random_generator,
|
||||
gradient_accumulation_steps, sp_size, max_grad_norm,
|
||||
uncond_prompt_embed, uncond_prompt_mask,
|
||||
num_euler_timesteps, multiphase, not_apply_cfg_solver,
|
||||
distill_cfg, ema_decay, pred_decay_weight,
|
||||
pred_decay_type, hunyuan_teacher_disable_cfg):
|
||||
"""
|
||||
Distill one step of the model.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Distillation pipeline must implement this method")
|
||||
|
||||
def log_validation(self, transformer, fastvideo_args, global_step):
|
||||
"""Log validation results during training."""
|
||||
fastvideo_args.inference_mode = True
|
||||
fastvideo_args.use_cpu_offload = False
|
||||
if not fastvideo_args.log_validation:
|
||||
return
|
||||
if self.validation_pipeline is None:
|
||||
raise ValueError("Validation pipeline is not set")
|
||||
|
||||
# Create sampling parameters if not provided
|
||||
sampling_param = SamplingParam.from_pretrained(
|
||||
fastvideo_args.model_path)
|
||||
|
||||
# Prepare validation prompts
|
||||
validation_dataset = ParquetVideoTextDataset(
|
||||
fastvideo_args.validation_prompt_dir,
|
||||
batch_size=1,
|
||||
rank=0,
|
||||
world_size=1,
|
||||
cfg_rate=0,
|
||||
num_latent_t=fastvideo_args.num_latent_t)
|
||||
|
||||
validation_dataloader = StatefulDataLoader(validation_dataset,
|
||||
batch_size=1,
|
||||
num_workers=1,
|
||||
prefetch_factor=2,
|
||||
shuffle=False,
|
||||
pin_memory=True,
|
||||
drop_last=False)
|
||||
|
||||
transformer.requires_grad_(False)
|
||||
for p in transformer.parameters():
|
||||
p.requires_grad = False
|
||||
transformer.eval()
|
||||
|
||||
# Add the transformer to the validation pipeline
|
||||
self.validation_pipeline.add_module("transformer", transformer)
|
||||
self.validation_pipeline.latent_preparation_stage.transformer = transformer
|
||||
self.validation_pipeline.denoising_stage.transformer = transformer
|
||||
|
||||
# Process validation prompts
|
||||
videos = []
|
||||
captions = []
|
||||
for _, embeddings, masks, infos in validation_dataloader:
|
||||
logger.info(f"infos: {infos}")
|
||||
caption = infos['caption']
|
||||
captions.append(caption)
|
||||
prompt_embeds = embeddings.to(fastvideo_args.device)
|
||||
prompt_attention_mask = masks.to(fastvideo_args.device)
|
||||
|
||||
# Calculate sizes
|
||||
latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
|
||||
sampling_param.height // 8,
|
||||
sampling_param.width // 8]
|
||||
n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
|
||||
|
||||
# Prepare batch for validation
|
||||
batch = ForwardBatch(
|
||||
data_type="video",
|
||||
latents=None,
|
||||
prompt_embeds=[prompt_embeds],
|
||||
prompt_attention_mask=[prompt_attention_mask],
|
||||
height=fastvideo_args.num_height,
|
||||
width=fastvideo_args.num_width,
|
||||
num_frames=fastvideo_args.num_frames,
|
||||
num_inference_steps=10,
|
||||
guidance_scale=1,
|
||||
n_tokens=n_tokens,
|
||||
do_classifier_free_guidance=False,
|
||||
eta=0.0,
|
||||
extra={},
|
||||
)
|
||||
|
||||
# Run validation inference
|
||||
with torch.inference_mode():
|
||||
output_batch = self.validation_pipeline.forward(
|
||||
batch, fastvideo_args)
|
||||
samples = output_batch.output
|
||||
|
||||
# Process outputs
|
||||
video = rearrange(samples, "b c t h w -> t b c h w")
|
||||
frames = []
|
||||
for x in video:
|
||||
x = torchvision.utils.make_grid(x, nrow=6)
|
||||
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
|
||||
frames.append((x * 255).numpy().astype(np.uint8))
|
||||
videos.append(frames)
|
||||
|
||||
# Log validation results
|
||||
if self.rank == 0:
|
||||
video_filenames = []
|
||||
video_captions = []
|
||||
for i, video in enumerate(videos):
|
||||
caption = captions[i]
|
||||
filename = os.path.join(
|
||||
fastvideo_args.output_dir,
|
||||
f"validation_step_{global_step}_video_{i}.mp4")
|
||||
imageio.mimsave(filename, video, fps=sampling_param.fps)
|
||||
video_filenames.append(filename)
|
||||
video_captions.append(caption)
|
||||
|
||||
logs = {
|
||||
"validation_videos": [
|
||||
wandb.Video(filename,
|
||||
caption=caption) for filename, caption in zip(
|
||||
video_filenames, video_captions)
|
||||
]
|
||||
}
|
||||
wandb.log(logs, step=global_step)
|
||||
|
||||
# Re-enable gradients for training
|
||||
transformer.requires_grad_(True)
|
||||
transformer.train()
|
||||
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
class WanDistillationPipeline(DistillationPipeline):
|
||||
"""
|
||||
A distillation pipeline for Wan.
|
||||
"""
|
||||
_required_config_modules = ["scheduler", "transformer"]
|
||||
|
||||
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
|
||||
pass
|
||||
|
||||
def create_training_stages(self, fastvideo_args: FastVideoArgs):
|
||||
pass
|
||||
|
||||
def initialize_validation_pipeline(self, fastvideo_args: FastVideoArgs):
|
||||
logger.info("Initializing validation pipeline...")
|
||||
args_copy = deepcopy(fastvideo_args)
|
||||
|
||||
args_copy.mode = "inference"
|
||||
args_copy.vae_config.load_encoder = False
|
||||
validation_pipeline = WanValidationPipeline.from_pretrained(
|
||||
fastvideo_args.model_path, args=args_copy)
|
||||
|
||||
self.validation_pipeline = validation_pipeline
|
||||
|
||||
def distill_one_step(
|
||||
self,
|
||||
transformer,
|
||||
model_type,
|
||||
teacher_transformer,
|
||||
ema_transformer,
|
||||
optimizer,
|
||||
lr_scheduler,
|
||||
loader_iter,
|
||||
noise_scheduler,
|
||||
solver,
|
||||
noise_random_generator,
|
||||
gradient_accumulation_steps,
|
||||
sp_size,
|
||||
max_grad_norm,
|
||||
uncond_prompt_embed,
|
||||
uncond_prompt_mask,
|
||||
num_euler_timesteps,
|
||||
multiphase,
|
||||
not_apply_cfg_solver,
|
||||
distill_cfg,
|
||||
ema_decay,
|
||||
pred_decay_weight,
|
||||
pred_decay_type,
|
||||
hunyuan_teacher_disable_cfg,
|
||||
):
|
||||
"""Perform one step of distillation training."""
|
||||
total_loss = 0.0
|
||||
optimizer.zero_grad()
|
||||
model_pred_norm = {
|
||||
"fro": 0.0, # codespell:ignore
|
||||
"largest singular value": 0.0,
|
||||
"absolute mean": 0.0,
|
||||
"absolute max": 0.0,
|
||||
}
|
||||
|
||||
for _ in range(gradient_accumulation_steps):
|
||||
(
|
||||
latents,
|
||||
encoder_hidden_states,
|
||||
encoder_attention_mask,
|
||||
infos,
|
||||
) = next(loader_iter)
|
||||
|
||||
latents = latents.to(self.device, dtype=torch.bfloat16)
|
||||
encoder_hidden_states = encoder_hidden_states.to(
|
||||
self.device, dtype=torch.bfloat16)
|
||||
|
||||
model_input = normalize_dit_input(model_type, latents)
|
||||
noise = torch.randn_like(model_input)
|
||||
bsz = model_input.shape[0]
|
||||
index = torch.randint(0,
|
||||
num_euler_timesteps, (bsz, ),
|
||||
device=model_input.device).long()
|
||||
if sp_size > 1:
|
||||
self.sp_group.broadcast(index, src=0)
|
||||
|
||||
# Add noise according to flow matching
|
||||
sigmas = extract_into_tensor(solver.sigmas, index,
|
||||
model_input.shape)
|
||||
sigmas_prev = extract_into_tensor(solver.sigmas_prev, index,
|
||||
model_input.shape)
|
||||
|
||||
timesteps = (sigmas *
|
||||
noise_scheduler.config.num_train_timesteps).view(-1)
|
||||
timesteps_prev = (
|
||||
sigmas_prev *
|
||||
noise_scheduler.config.num_train_timesteps).view(-1)
|
||||
noisy_model_input = sigmas * noise + (1.0 - sigmas) * model_input
|
||||
|
||||
# Get student model prediction
|
||||
with torch.autocast("cuda", dtype=torch.bfloat16):
|
||||
input_kwargs = {
|
||||
"hidden_states": noisy_model_input,
|
||||
"encoder_hidden_states": encoder_hidden_states,
|
||||
"timestep": timesteps,
|
||||
"encoder_attention_mask": encoder_attention_mask,
|
||||
"return_dict": False,
|
||||
}
|
||||
if hunyuan_teacher_disable_cfg:
|
||||
input_kwargs["guidance"] = torch.tensor(
|
||||
[1000.0],
|
||||
device=noisy_model_input.device,
|
||||
dtype=torch.bfloat16)
|
||||
|
||||
with set_forward_context(current_timestep=timesteps,
|
||||
attn_metadata=None):
|
||||
model_pred = transformer(**input_kwargs)[0]
|
||||
|
||||
# Apply multi-phase prediction
|
||||
model_pred, end_index = solver.euler_style_multiphase_pred(
|
||||
noisy_model_input, model_pred, index, multiphase)
|
||||
|
||||
# Get teacher model guidance
|
||||
with torch.no_grad():
|
||||
w = distill_cfg
|
||||
with torch.autocast("cuda", dtype=torch.bfloat16):
|
||||
with set_forward_context(current_timestep=timesteps,
|
||||
attn_metadata=None):
|
||||
cond_teacher_output = teacher_transformer(
|
||||
noisy_model_input,
|
||||
encoder_hidden_states,
|
||||
timesteps,
|
||||
encoder_attention_mask,
|
||||
return_dict=False,
|
||||
)[0].float()
|
||||
|
||||
if not_apply_cfg_solver:
|
||||
uncond_teacher_output = cond_teacher_output
|
||||
else:
|
||||
# Get teacher model prediction on unconditional embedding
|
||||
with torch.autocast("cuda", dtype=torch.bfloat16):
|
||||
with set_forward_context(current_timestep=timesteps,
|
||||
attn_metadata=None):
|
||||
uncond_teacher_output = teacher_transformer(
|
||||
noisy_model_input,
|
||||
uncond_prompt_embed.unsqueeze(0).expand(
|
||||
bsz, -1, -1),
|
||||
timesteps,
|
||||
uncond_prompt_mask.unsqueeze(0).expand(bsz, -1),
|
||||
return_dict=False,
|
||||
)[0].float()
|
||||
|
||||
teacher_output = uncond_teacher_output + w * (
|
||||
cond_teacher_output - uncond_teacher_output)
|
||||
x_prev = solver.euler_step(noisy_model_input, teacher_output,
|
||||
index)
|
||||
|
||||
# Get target prediction
|
||||
with torch.no_grad():
|
||||
with torch.autocast("cuda", dtype=torch.bfloat16):
|
||||
if ema_transformer is not None:
|
||||
with set_forward_context(
|
||||
current_timestep=timesteps_prev,
|
||||
attn_metadata=None):
|
||||
target_pred = ema_transformer(
|
||||
x_prev.float(),
|
||||
encoder_hidden_states,
|
||||
timesteps_prev,
|
||||
encoder_attention_mask,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
else:
|
||||
with set_forward_context(
|
||||
current_timestep=timesteps_prev,
|
||||
attn_metadata=None):
|
||||
target_pred = transformer(
|
||||
x_prev.float(),
|
||||
encoder_hidden_states,
|
||||
timesteps_prev,
|
||||
encoder_attention_mask,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
target, end_index = solver.euler_style_multiphase_pred(
|
||||
x_prev, target_pred, index, multiphase, True)
|
||||
|
||||
# Calculate loss
|
||||
huber_c = 0.001
|
||||
loss = (torch.mean(
|
||||
torch.sqrt((model_pred.float() - target.float())**2 +
|
||||
huber_c**2) - huber_c) / gradient_accumulation_steps)
|
||||
|
||||
if pred_decay_weight > 0:
|
||||
if pred_decay_type == "l1":
|
||||
pred_decay_loss = (
|
||||
torch.mean(torch.sqrt(model_pred.float()**2)) *
|
||||
pred_decay_weight / gradient_accumulation_steps)
|
||||
loss += pred_decay_loss
|
||||
elif pred_decay_type == "l2":
|
||||
pred_decay_loss = (torch.mean(model_pred.float()**2) *
|
||||
pred_decay_weight /
|
||||
gradient_accumulation_steps)
|
||||
loss += pred_decay_loss
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"pred_decay_type is not implemented")
|
||||
|
||||
# Calculate model prediction norms
|
||||
get_norm(model_pred.detach().float(), model_pred_norm,
|
||||
gradient_accumulation_steps)
|
||||
loss.backward()
|
||||
|
||||
avg_loss = loss.detach().clone()
|
||||
self.sp_group.all_reduce(avg_loss,
|
||||
op=torch.distributed.ReduceOp.AVG)
|
||||
total_loss += avg_loss.item()
|
||||
|
||||
# Update EMA
|
||||
if ema_transformer is not None:
|
||||
reshard_fsdp(ema_transformer)
|
||||
for p_averaged, p_model in zip(ema_transformer.parameters(),
|
||||
transformer.parameters()):
|
||||
with torch.no_grad():
|
||||
p_averaged.copy_(
|
||||
torch.lerp(p_averaged.detach(), p_model.detach(),
|
||||
1 - ema_decay))
|
||||
|
||||
# Gradient clipping and optimization step
|
||||
model_parts = [transformer]
|
||||
grad_norm = _clip_grad_norm_while_handling_failing_dtensor_cases(
|
||||
[p for m in model_parts for p in m.parameters()],
|
||||
max_grad_norm,
|
||||
foreach=None,
|
||||
)
|
||||
|
||||
optimizer.step()
|
||||
lr_scheduler.step()
|
||||
|
||||
return total_loss, grad_norm.item(), model_pred_norm
|
||||
|
||||
def forward(
|
||||
self,
|
||||
batch: ForwardBatch,
|
||||
fastvideo_args: TrainingArgs,
|
||||
):
|
||||
args = fastvideo_args
|
||||
train_dataloader = self.train_dataloader
|
||||
init_steps = self.init_steps
|
||||
lr_scheduler = self.lr_scheduler
|
||||
optimizer = self.optimizer
|
||||
noise_scheduler = self.noise_scheduler
|
||||
solver = self.solver
|
||||
noise_random_generator = None
|
||||
uncond_prompt_embed = self.uncond_prompt_embed
|
||||
uncond_prompt_mask = self.uncond_prompt_mask
|
||||
|
||||
# Train!
|
||||
total_batch_size = (self.world_size * args.gradient_accumulation_steps /
|
||||
args.sp_size * args.train_sp_batch_size)
|
||||
logger.info("***** Running distillation training *****")
|
||||
logger.info(f" Resume training from step {init_steps}")
|
||||
logger.info(
|
||||
f" Instantaneous batch size per device = {args.train_batch_size}")
|
||||
logger.info(
|
||||
f" Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}"
|
||||
)
|
||||
logger.info(
|
||||
f" Gradient Accumulation steps = {args.gradient_accumulation_steps}"
|
||||
)
|
||||
logger.info(f" Total optimization steps = {args.max_train_steps}")
|
||||
logger.info(
|
||||
f" Total training parameters per FSDP shard = {sum(p.numel() for p in self.transformer.parameters() if p.requires_grad) / 1e9} B"
|
||||
)
|
||||
logger.info(
|
||||
f" Master weight dtype: {self.transformer.parameters().__next__().dtype}"
|
||||
)
|
||||
|
||||
# Potentially load in the weights and states from a previous save
|
||||
if args.resume_from_checkpoint:
|
||||
raise NotImplementedError(
|
||||
"resume_from_checkpoint is not supported now.")
|
||||
|
||||
progress_bar = tqdm(
|
||||
range(0, args.max_train_steps),
|
||||
initial=init_steps,
|
||||
desc="Steps",
|
||||
disable=self.local_rank > 0,
|
||||
)
|
||||
|
||||
loader_iter = iter(train_dataloader)
|
||||
step_times = deque(maxlen=100)
|
||||
|
||||
# Skip steps if resuming
|
||||
for i in range(init_steps):
|
||||
next(loader_iter)
|
||||
|
||||
def get_num_phases(multi_phased_distill_schedule, step):
|
||||
# step-phase,step-phase
|
||||
multi_phases = multi_phased_distill_schedule.split(",")
|
||||
phase = multi_phases[-1].split("-")[-1]
|
||||
for step_phases in multi_phases:
|
||||
phase_step, phase = step_phases.split("-")
|
||||
if step <= int(phase_step):
|
||||
return int(phase)
|
||||
return int(phase)
|
||||
|
||||
for step in range(init_steps + 1, args.max_train_steps + 1):
|
||||
start_time = time.perf_counter()
|
||||
|
||||
assert args.multi_phased_distill_schedule is not None
|
||||
num_phases = get_num_phases(args.multi_phased_distill_schedule,
|
||||
step)
|
||||
|
||||
loss, grad_norm, pred_norm = self.distill_one_step(
|
||||
self.transformer,
|
||||
"wan", # model_type
|
||||
self.teacher_transformer,
|
||||
self.ema_transformer,
|
||||
optimizer,
|
||||
lr_scheduler,
|
||||
loader_iter,
|
||||
noise_scheduler,
|
||||
solver,
|
||||
noise_random_generator,
|
||||
args.gradient_accumulation_steps,
|
||||
args.sp_size,
|
||||
args.max_grad_norm,
|
||||
uncond_prompt_embed,
|
||||
uncond_prompt_mask,
|
||||
args.num_euler_timesteps,
|
||||
num_phases,
|
||||
args.not_apply_cfg_solver,
|
||||
args.distill_cfg,
|
||||
args.ema_decay,
|
||||
args.pred_decay_weight,
|
||||
args.pred_decay_type,
|
||||
args.hunyuan_teacher_disable_cfg,
|
||||
)
|
||||
|
||||
step_time = time.perf_counter() - start_time
|
||||
step_times.append(step_time)
|
||||
avg_step_time = sum(step_times) / len(step_times)
|
||||
|
||||
progress_bar.set_postfix({
|
||||
"loss": f"{loss:.4f}",
|
||||
"step_time": f"{step_time:.2f}s",
|
||||
"grad_norm": grad_norm,
|
||||
"phases": num_phases,
|
||||
})
|
||||
progress_bar.update(1)
|
||||
|
||||
if self.rank <= 0:
|
||||
wandb.log(
|
||||
{
|
||||
"train_loss":
|
||||
loss,
|
||||
"learning_rate":
|
||||
lr_scheduler.get_last_lr()[0],
|
||||
"step_time":
|
||||
step_time,
|
||||
"avg_step_time":
|
||||
avg_step_time,
|
||||
"grad_norm":
|
||||
grad_norm,
|
||||
"pred_fro_norm":
|
||||
pred_norm["fro"], # codespell:ignore
|
||||
"pred_largest_singular_value":
|
||||
pred_norm["largest singular value"],
|
||||
"pred_absolute_mean":
|
||||
pred_norm["absolute mean"],
|
||||
"pred_absolute_max":
|
||||
pred_norm["absolute max"],
|
||||
"phases":
|
||||
num_phases,
|
||||
},
|
||||
step=step,
|
||||
)
|
||||
|
||||
if step % args.checkpointing_steps == 0:
|
||||
if args.use_lora:
|
||||
raise NotImplementedError("LoRA is not supported now")
|
||||
else:
|
||||
if args.use_ema:
|
||||
save_checkpoint_v1(self.ema_transformer, self.rank,
|
||||
args.output_dir, step)
|
||||
else:
|
||||
save_checkpoint_v1(self.transformer, self.rank,
|
||||
args.output_dir, step)
|
||||
self.sp_group.barrier()
|
||||
|
||||
if args.log_validation and step % args.validation_steps == 0:
|
||||
self.log_validation(self.transformer, args, step)
|
||||
|
||||
# Final checkpoint
|
||||
if args.use_lora:
|
||||
raise NotImplementedError("LoRA is not supported now")
|
||||
else:
|
||||
save_checkpoint_v1(self.transformer, self.rank, args.output_dir,
|
||||
args.max_train_steps)
|
||||
|
||||
if get_sp_group():
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
def main(args):
|
||||
logger.info("Starting distillation pipeline...")
|
||||
|
||||
pipeline = WanDistillationPipeline.from_pretrained(
|
||||
args.pretrained_model_name_or_path, args=args)
|
||||
|
||||
args = pipeline.fastvideo_args
|
||||
pipeline.forward(None, args)
|
||||
logger.info("Distillation pipeline done")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
argv = sys.argv
|
||||
from fastvideo.v1.fastvideo_args import TrainingArgs
|
||||
from fastvideo.v1.utils import FlexibleArgumentParser
|
||||
parser = FlexibleArgumentParser()
|
||||
parser = TrainingArgs.add_cli_args(parser)
|
||||
parser = FastVideoArgs.add_cli_args(parser)
|
||||
args = parser.parse_args()
|
||||
args.use_cpu_offload = False
|
||||
print(args)
|
||||
main(args)
|
||||
@@ -0,0 +1,820 @@
|
||||
import gc
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from abc import ABC, abstractmethod
|
||||
from collections import deque
|
||||
from copy import deepcopy
|
||||
|
||||
import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision
|
||||
# import torch.distributed as dist
|
||||
import wandb
|
||||
from diffusers.optimization import get_scheduler
|
||||
from einops import rearrange
|
||||
from torchdata.stateful_dataloader import StatefulDataLoader
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_dit_input
|
||||
from fastvideo.utils.checkpoint import save_checkpoint_v1
|
||||
from fastvideo.v1.configs.sample import SamplingParam
|
||||
from fastvideo.v1.dataset.parquet_datasets import ParquetVideoTextDataset
|
||||
from fastvideo.v1.distributed import cleanup_dist_env_and_memory, get_sp_group
|
||||
from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
|
||||
from fastvideo.v1.forward_context import set_forward_context
|
||||
from fastvideo.v1.logger import init_logger
|
||||
from fastvideo.v1.pipelines import ComposedPipelineBase
|
||||
from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
|
||||
from fastvideo.v1.pipelines.training_utils import (
|
||||
_clip_grad_norm_while_handling_failing_dtensor_cases,
|
||||
compute_density_for_timestep_sampling, get_sigmas)
|
||||
from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Manual gradient checking flag - set to True to enable gradient verification
|
||||
ENABLE_GRADIENT_CHECK = False
|
||||
GRADIENT_CHECK_DTYPE = torch.bfloat16
|
||||
|
||||
|
||||
class TrainingPipeline(ComposedPipelineBase, ABC):
|
||||
"""
|
||||
A pipeline for training a model. All training pipelines should inherit from this class.
|
||||
All reusable components and code should be implemented in this class.
|
||||
"""
|
||||
_required_config_modules = ["scheduler", "transformer"]
|
||||
|
||||
def initialize_training_pipeline(self, fastvideo_args: TrainingArgs):
|
||||
logger.info("Initializing training pipeline...")
|
||||
self.device = fastvideo_args.device
|
||||
self.sp_group = get_sp_group()
|
||||
self.world_size = self.sp_group.world_size
|
||||
self.rank = self.sp_group.rank
|
||||
self.local_rank = self.sp_group.local_rank
|
||||
self.transformer = self.get_module("transformer")
|
||||
assert self.transformer is not None
|
||||
|
||||
self.transformer.requires_grad_(True)
|
||||
self.transformer.train()
|
||||
|
||||
args = fastvideo_args
|
||||
|
||||
noise_scheduler = self.modules["scheduler"]
|
||||
params_to_optimize = self.transformer.parameters()
|
||||
params_to_optimize = list(
|
||||
filter(lambda p: p.requires_grad, params_to_optimize))
|
||||
|
||||
optimizer = torch.optim.AdamW(
|
||||
params_to_optimize,
|
||||
lr=args.learning_rate,
|
||||
betas=(0.9, 0.999),
|
||||
weight_decay=args.weight_decay,
|
||||
eps=1e-8,
|
||||
)
|
||||
|
||||
init_steps = 0
|
||||
logger.info("optimizer: %s", optimizer)
|
||||
|
||||
# todo add lr scheduler
|
||||
lr_scheduler = get_scheduler(
|
||||
args.lr_scheduler,
|
||||
optimizer=optimizer,
|
||||
num_warmup_steps=args.lr_warmup_steps * self.world_size,
|
||||
num_training_steps=args.max_train_steps * self.world_size,
|
||||
num_cycles=args.lr_num_cycles,
|
||||
power=args.lr_power,
|
||||
last_epoch=init_steps - 1,
|
||||
)
|
||||
|
||||
train_dataset = ParquetVideoTextDataset(
|
||||
args.data_path,
|
||||
batch_size=args.train_batch_size,
|
||||
rank=self.rank,
|
||||
world_size=self.world_size,
|
||||
cfg_rate=args.cfg,
|
||||
num_latent_t=args.num_latent_t)
|
||||
|
||||
train_dataloader = StatefulDataLoader(
|
||||
train_dataset,
|
||||
batch_size=args.train_batch_size,
|
||||
num_workers=args.
|
||||
dataloader_num_workers, # Reduce number of workers to avoid memory issues
|
||||
prefetch_factor=2,
|
||||
shuffle=False,
|
||||
pin_memory=True,
|
||||
drop_last=True)
|
||||
|
||||
self.lr_scheduler = lr_scheduler
|
||||
self.train_dataset = train_dataset
|
||||
self.train_dataloader = train_dataloader
|
||||
self.init_steps = init_steps
|
||||
self.optimizer = optimizer
|
||||
self.noise_scheduler = noise_scheduler
|
||||
# self.noise_random_generator = noise_random_generator
|
||||
|
||||
# num_update_steps_per_epoch = math.ceil(
|
||||
# len(train_dataloader) / args.gradient_accumulation_steps *
|
||||
# args.sp_size / args.train_sp_batch_size)
|
||||
# args.num_train_epochs = math.ceil(args.max_train_steps /
|
||||
# num_update_steps_per_epoch)
|
||||
|
||||
if self.rank <= 0:
|
||||
project = args.tracker_project_name or "fastvideo"
|
||||
wandb.init(project=project, config=args)
|
||||
|
||||
@abstractmethod
|
||||
def initialize_validation_pipeline(self, fastvideo_args: FastVideoArgs):
|
||||
raise NotImplementedError(
|
||||
"Training pipelines must implement this method")
|
||||
|
||||
@abstractmethod
|
||||
def train_one_step(self, transformer, model_type, optimizer, lr_scheduler,
|
||||
loader, noise_scheduler, noise_random_generator,
|
||||
gradient_accumulation_steps, sp_size,
|
||||
precondition_outputs, max_grad_norm, weighting_scheme,
|
||||
logit_mean, logit_std, mode_scale):
|
||||
"""
|
||||
Train one step of the model.
|
||||
"""
|
||||
raise NotImplementedError(
|
||||
"Training pipeline must implement this method")
|
||||
|
||||
def log_validation(self, transformer, fastvideo_args, global_step):
|
||||
fastvideo_args.inference_mode = True
|
||||
fastvideo_args.use_cpu_offload = False
|
||||
if not fastvideo_args.log_validation:
|
||||
return
|
||||
if self.validation_pipeline is None:
|
||||
raise ValueError("Validation pipeline is not set")
|
||||
|
||||
# Create sampling parameters if not provided
|
||||
sampling_param = SamplingParam.from_pretrained(
|
||||
fastvideo_args.model_path)
|
||||
|
||||
# Prepare validation prompts
|
||||
print('fastvideo_args.validation_prompt_dir',
|
||||
fastvideo_args.validation_prompt_dir)
|
||||
validation_dataset = ParquetVideoTextDataset(
|
||||
fastvideo_args.validation_prompt_dir,
|
||||
batch_size=1,
|
||||
rank=0,
|
||||
world_size=1,
|
||||
cfg_rate=0,
|
||||
num_latent_t=args.num_latent_t)
|
||||
|
||||
validation_dataloader = StatefulDataLoader(
|
||||
validation_dataset,
|
||||
batch_size=1,
|
||||
num_workers=1, # Reduce number of workers to avoid memory issues
|
||||
prefetch_factor=2,
|
||||
shuffle=False,
|
||||
pin_memory=True,
|
||||
drop_last=False)
|
||||
|
||||
transformer.requires_grad_(False)
|
||||
for p in transformer.parameters():
|
||||
p.requires_grad = False
|
||||
transformer.eval()
|
||||
|
||||
# Add the transformer to the validation pipeline
|
||||
self.validation_pipeline.add_module("transformer", transformer)
|
||||
self.validation_pipeline.latent_preparation_stage.transformer = transformer
|
||||
self.validation_pipeline.denoising_stage.transformer = transformer
|
||||
|
||||
# Process each validation prompt
|
||||
videos = []
|
||||
captions = []
|
||||
for _, embeddings, masks, infos in validation_dataloader:
|
||||
logger.info(f"infos: {infos}")
|
||||
caption = infos['caption']
|
||||
captions.append(caption)
|
||||
prompt_embeds = embeddings.to(fastvideo_args.device)
|
||||
prompt_attention_mask = masks.to(fastvideo_args.device)
|
||||
|
||||
# Calculate sizes
|
||||
latents_size = [(sampling_param.num_frames - 1) // 4 + 1,
|
||||
sampling_param.height // 8,
|
||||
sampling_param.width // 8]
|
||||
n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
|
||||
|
||||
# Prepare batch for validation
|
||||
# print('shape of embeddings', prompt_embeds.shape)
|
||||
batch = ForwardBatch(
|
||||
# **shallow_asdict(sampling_param),
|
||||
data_type="video",
|
||||
latents=None,
|
||||
# seed=sampling_param.seed,
|
||||
# data_type="video",
|
||||
prompt_embeds=[prompt_embeds],
|
||||
prompt_attention_mask=[prompt_attention_mask],
|
||||
# make sure we use the same height, width, and num_frames as the training pipeline
|
||||
height=args.num_height,
|
||||
width=args.num_width,
|
||||
num_frames=args.num_frames,
|
||||
# num_inference_steps=fastvideo_args.validation_sampling_steps,
|
||||
num_inference_steps=10,
|
||||
# guidance_scale=fastvideo_args.validation_guidance_scale,
|
||||
guidance_scale=1,
|
||||
n_tokens=n_tokens,
|
||||
do_classifier_free_guidance=False,
|
||||
eta=0.0,
|
||||
extra={},
|
||||
)
|
||||
|
||||
# Run validation inference
|
||||
with torch.inference_mode():
|
||||
output_batch = self.validation_pipeline.forward(
|
||||
batch, fastvideo_args)
|
||||
samples = output_batch.output
|
||||
|
||||
# Process outputs
|
||||
video = rearrange(samples, "b c t h w -> t b c h w")
|
||||
frames = []
|
||||
for x in video:
|
||||
x = torchvision.utils.make_grid(x, nrow=6)
|
||||
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
|
||||
frames.append((x * 255).numpy().astype(np.uint8))
|
||||
videos.append(frames)
|
||||
|
||||
# Log validation results
|
||||
rank = int(os.environ.get("RANK", 0))
|
||||
|
||||
if rank == 0:
|
||||
video_filenames = []
|
||||
video_captions = []
|
||||
for i, video in enumerate(videos):
|
||||
caption = captions[i]
|
||||
filename = os.path.join(
|
||||
fastvideo_args.output_dir,
|
||||
f"validation_step_{global_step}_video_{i}.mp4")
|
||||
imageio.mimsave(filename, video, fps=sampling_param.fps)
|
||||
video_filenames.append(filename)
|
||||
video_captions.append(
|
||||
caption) # Store the caption for each video
|
||||
|
||||
logs = {
|
||||
"validation_videos": [
|
||||
wandb.Video(filename,
|
||||
caption=caption) for filename, caption in zip(
|
||||
video_filenames, video_captions)
|
||||
]
|
||||
}
|
||||
wandb.log(logs, step=global_step)
|
||||
|
||||
# Re-enable gradients for training
|
||||
transformer.requires_grad_(True)
|
||||
transformer.train()
|
||||
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def gradient_check_parameters(self,
|
||||
transformer,
|
||||
latents,
|
||||
encoder_hidden_states,
|
||||
encoder_attention_mask,
|
||||
timesteps,
|
||||
target,
|
||||
eps=5e-2,
|
||||
max_params_to_check=2000):
|
||||
"""
|
||||
Verify gradients using finite differences for FSDP models with GRADIENT_CHECK_DTYPE.
|
||||
Uses standard tolerances for GRADIENT_CHECK_DTYPE precision.
|
||||
"""
|
||||
# Move all inputs to CPU and clear GPU memory
|
||||
inputs_cpu = {
|
||||
'latents': latents.cpu(),
|
||||
'encoder_hidden_states': encoder_hidden_states.cpu(),
|
||||
'encoder_attention_mask': encoder_attention_mask.cpu(),
|
||||
'timesteps': timesteps.cpu(),
|
||||
'target': target.cpu()
|
||||
}
|
||||
del latents, encoder_hidden_states, encoder_attention_mask, timesteps, target
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def compute_loss():
|
||||
# Move inputs to GPU, compute loss, cleanup
|
||||
inputs_gpu = {
|
||||
k:
|
||||
v.to(self.fastvideo_args.device,
|
||||
dtype=GRADIENT_CHECK_DTYPE
|
||||
if k != 'encoder_attention_mask' else None)
|
||||
for k, v in inputs_cpu.items()
|
||||
}
|
||||
|
||||
# Use GRADIENT_CHECK_DTYPE for more accurate gradient checking
|
||||
# with torch.autocast(enabled=False, device_type="cuda"):
|
||||
with torch.autocast("cuda", dtype=GRADIENT_CHECK_DTYPE):
|
||||
with set_forward_context(
|
||||
current_timestep=inputs_gpu['timesteps'],
|
||||
attn_metadata=None):
|
||||
model_pred = transformer(
|
||||
hidden_states=inputs_gpu['latents'],
|
||||
encoder_hidden_states=inputs_gpu[
|
||||
'encoder_hidden_states'],
|
||||
timestep=inputs_gpu['timesteps'],
|
||||
encoder_attention_mask=inputs_gpu[
|
||||
'encoder_attention_mask'],
|
||||
return_dict=False)[0]
|
||||
|
||||
if self.fastvideo_args.precondition_outputs:
|
||||
sigmas = get_sigmas(self.noise_scheduler,
|
||||
inputs_gpu['latents'].device,
|
||||
inputs_gpu['timesteps'],
|
||||
n_dim=inputs_gpu['latents'].ndim,
|
||||
dtype=inputs_gpu['latents'].dtype)
|
||||
model_pred = inputs_gpu['latents'] - model_pred * sigmas
|
||||
target_adjusted = inputs_gpu['target']
|
||||
else:
|
||||
target_adjusted = inputs_gpu['target']
|
||||
|
||||
loss = torch.mean((model_pred - target_adjusted)**2)
|
||||
|
||||
# Cleanup and return
|
||||
loss_cpu = loss.cpu()
|
||||
del inputs_gpu, model_pred, target_adjusted
|
||||
if 'sigmas' in locals(): del sigmas
|
||||
torch.cuda.empty_cache()
|
||||
return loss_cpu.to(self.fastvideo_args.device)
|
||||
|
||||
try:
|
||||
# Get analytical gradients
|
||||
transformer.zero_grad()
|
||||
analytical_loss = compute_loss()
|
||||
analytical_loss.backward()
|
||||
|
||||
# Check gradients for selected parameters
|
||||
absolute_errors = []
|
||||
param_count = 0
|
||||
|
||||
for name, param in transformer.named_parameters():
|
||||
if not (param.requires_grad and param.grad is not None
|
||||
and param_count < max_params_to_check
|
||||
and param.grad.abs().max() > 5e-4):
|
||||
continue
|
||||
|
||||
# Get local parameter and gradient tensors
|
||||
local_param = param._local_tensor if hasattr(
|
||||
param, '_local_tensor') else param
|
||||
local_grad = param.grad._local_tensor if hasattr(
|
||||
param.grad, '_local_tensor') else param.grad
|
||||
|
||||
# Find first significant gradient element
|
||||
flat_param = local_param.data.view(-1)
|
||||
flat_grad = local_grad.view(-1)
|
||||
check_idx = next((i for i in range(min(10, flat_param.numel()))
|
||||
if abs(flat_grad[i]) > 1e-4), 0)
|
||||
|
||||
# Store original values
|
||||
orig_value = flat_param[check_idx].item()
|
||||
analytical_grad = flat_grad[check_idx].item()
|
||||
|
||||
# Compute numerical gradient
|
||||
for delta in [eps, -eps]:
|
||||
with torch.no_grad():
|
||||
flat_param[check_idx] = orig_value + delta
|
||||
loss = compute_loss()
|
||||
if delta > 0: loss_plus = loss.item()
|
||||
else: loss_minus = loss.item()
|
||||
|
||||
# Restore parameter and compute error
|
||||
with torch.no_grad():
|
||||
flat_param[check_idx] = orig_value
|
||||
|
||||
numerical_grad = (loss_plus - loss_minus) / (2 * eps)
|
||||
abs_error = abs(analytical_grad - numerical_grad)
|
||||
rel_error = abs_error / max(abs(analytical_grad),
|
||||
abs(numerical_grad), 1e-3)
|
||||
absolute_errors.append(abs_error)
|
||||
|
||||
logger.info(
|
||||
f"{name}[{check_idx}]: analytical={analytical_grad:.6f}, "
|
||||
f"numerical={numerical_grad:.6f}, abs_error={abs_error:.2e}, rel_error={rel_error:.2%}"
|
||||
)
|
||||
|
||||
# param_count += 1
|
||||
|
||||
# Compute and log statistics
|
||||
if absolute_errors:
|
||||
min_err, max_err, mean_err = min(absolute_errors), max(
|
||||
absolute_errors
|
||||
), sum(absolute_errors) / len(absolute_errors)
|
||||
logger.info(
|
||||
f"Gradient check stats: min={min_err:.2e}, max={max_err:.2e}, mean={mean_err:.2e}"
|
||||
)
|
||||
|
||||
if self.rank <= 0:
|
||||
wandb.log({
|
||||
"grad_check/min_abs_error":
|
||||
min_err,
|
||||
"grad_check/max_abs_error":
|
||||
max_err,
|
||||
"grad_check/mean_abs_error":
|
||||
mean_err,
|
||||
"grad_check/analytical_loss":
|
||||
analytical_loss.item(),
|
||||
})
|
||||
return max_err
|
||||
|
||||
return float('inf')
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Gradient check failed: {e}")
|
||||
traceback.print_exc()
|
||||
return float('inf')
|
||||
|
||||
def setup_gradient_check(self, args, loader_iter, noise_scheduler,
|
||||
noise_random_generator):
|
||||
"""
|
||||
Setup and perform gradient check on a fresh batch.
|
||||
Args:
|
||||
args: Training arguments
|
||||
loader_iter: Data loader iterator
|
||||
noise_scheduler: Noise scheduler for diffusion
|
||||
noise_random_generator: Random number generator for noise
|
||||
Returns:
|
||||
float or None: Maximum gradient error or None if check is disabled/fails
|
||||
"""
|
||||
if not ENABLE_GRADIENT_CHECK:
|
||||
return None
|
||||
|
||||
try:
|
||||
# Get a fresh batch and process it exactly like train_one_step
|
||||
check_latents, check_encoder_hidden_states, check_encoder_attention_mask, check_infos = next(
|
||||
loader_iter)
|
||||
|
||||
# Process exactly like in train_one_step but use GRADIENT_CHECK_DTYPE
|
||||
check_latents = check_latents.to(self.fastvideo_args.device,
|
||||
dtype=GRADIENT_CHECK_DTYPE)
|
||||
check_encoder_hidden_states = check_encoder_hidden_states.to(
|
||||
self.fastvideo_args.device, dtype=GRADIENT_CHECK_DTYPE)
|
||||
check_latents = normalize_dit_input("wan", check_latents)
|
||||
batch_size = check_latents.shape[0]
|
||||
check_noise = torch.randn_like(check_latents)
|
||||
|
||||
check_u = compute_density_for_timestep_sampling(
|
||||
weighting_scheme=args.weighting_scheme,
|
||||
batch_size=batch_size,
|
||||
generator=noise_random_generator,
|
||||
logit_mean=args.logit_mean,
|
||||
logit_std=args.logit_std,
|
||||
mode_scale=args.mode_scale,
|
||||
)
|
||||
check_indices = (check_u *
|
||||
noise_scheduler.config.num_train_timesteps).long()
|
||||
check_timesteps = noise_scheduler.timesteps[check_indices].to(
|
||||
device=check_latents.device)
|
||||
|
||||
check_sigmas = get_sigmas(
|
||||
noise_scheduler,
|
||||
check_latents.device,
|
||||
check_timesteps,
|
||||
n_dim=check_latents.ndim,
|
||||
dtype=check_latents.dtype,
|
||||
)
|
||||
check_noisy_model_input = (
|
||||
1.0 - check_sigmas) * check_latents + check_sigmas * check_noise
|
||||
|
||||
# Compute target exactly like train_one_step
|
||||
if args.precondition_outputs:
|
||||
check_target = check_latents
|
||||
else:
|
||||
check_target = check_noise - check_latents
|
||||
|
||||
# Perform gradient check with the exact same inputs as training
|
||||
max_grad_error = self.gradient_check_parameters(
|
||||
transformer=self.transformer,
|
||||
latents=
|
||||
check_noisy_model_input, # Use noisy input like in training
|
||||
encoder_hidden_states=check_encoder_hidden_states,
|
||||
encoder_attention_mask=check_encoder_attention_mask,
|
||||
timesteps=check_timesteps,
|
||||
target=check_target,
|
||||
max_params_to_check=100 # Check more parameters
|
||||
)
|
||||
|
||||
if max_grad_error > 5e-2:
|
||||
logger.error(
|
||||
f"❌ Large gradient error detected: {max_grad_error:.2e}")
|
||||
else:
|
||||
logger.info(
|
||||
f"✅ Gradient check passed: max error {max_grad_error:.2e}")
|
||||
|
||||
return max_grad_error
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Gradient check setup failed: {e}")
|
||||
traceback.print_exc()
|
||||
return None
|
||||
|
||||
|
||||
class WanTrainingPipeline(TrainingPipeline):
|
||||
"""
|
||||
A training pipeline for Wan.
|
||||
"""
|
||||
_required_config_modules = ["scheduler", "transformer"]
|
||||
|
||||
def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
|
||||
pass
|
||||
|
||||
def create_training_stages(self, fastvideo_args: FastVideoArgs):
|
||||
pass
|
||||
|
||||
def initialize_validation_pipeline(self, fastvideo_args: FastVideoArgs):
|
||||
logger.info("Initializing validation pipeline...")
|
||||
args_copy = deepcopy(fastvideo_args)
|
||||
|
||||
args_copy.mode = "inference"
|
||||
args_copy.vae_config.load_encoder = False
|
||||
validation_pipeline = WanValidationPipeline.from_pretrained(
|
||||
args.model_path, args=args_copy)
|
||||
|
||||
self.validation_pipeline = validation_pipeline
|
||||
|
||||
def train_one_step(
|
||||
self,
|
||||
transformer,
|
||||
model_type,
|
||||
optimizer,
|
||||
lr_scheduler,
|
||||
loader_iter,
|
||||
noise_scheduler,
|
||||
noise_random_generator,
|
||||
gradient_accumulation_steps,
|
||||
sp_size,
|
||||
precondition_outputs,
|
||||
max_grad_norm,
|
||||
weighting_scheme,
|
||||
logit_mean,
|
||||
logit_std,
|
||||
mode_scale,
|
||||
):
|
||||
self.modules["transformer"].requires_grad_(True)
|
||||
self.modules["transformer"].train()
|
||||
|
||||
total_loss = 0.0
|
||||
optimizer.zero_grad()
|
||||
for _ in range(gradient_accumulation_steps):
|
||||
(
|
||||
latents,
|
||||
encoder_hidden_states,
|
||||
encoder_attention_mask,
|
||||
infos,
|
||||
) = next(loader_iter)
|
||||
latents = latents.to(self.fastvideo_args.device,
|
||||
dtype=torch.bfloat16)
|
||||
encoder_hidden_states = encoder_hidden_states.to(
|
||||
self.fastvideo_args.device, dtype=torch.bfloat16)
|
||||
latents = normalize_dit_input(model_type, latents)
|
||||
batch_size = latents.shape[0]
|
||||
noise = torch.randn_like(latents)
|
||||
u = compute_density_for_timestep_sampling(
|
||||
weighting_scheme=weighting_scheme,
|
||||
batch_size=batch_size,
|
||||
generator=noise_random_generator,
|
||||
logit_mean=logit_mean,
|
||||
logit_std=logit_std,
|
||||
mode_scale=mode_scale,
|
||||
)
|
||||
indices = (u * noise_scheduler.config.num_train_timesteps).long()
|
||||
timesteps = noise_scheduler.timesteps[indices].to(
|
||||
device=latents.device)
|
||||
if sp_size > 1:
|
||||
# Make sure that the timesteps are the same across all sp processes.
|
||||
sp_group = get_sp_group()
|
||||
sp_group.broadcast(timesteps, src=0)
|
||||
sigmas = get_sigmas(
|
||||
noise_scheduler,
|
||||
latents.device,
|
||||
timesteps,
|
||||
n_dim=latents.ndim,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
noisy_model_input = (1.0 - sigmas) * latents + sigmas * noise
|
||||
print('device before forward ',
|
||||
next(transformer.named_parameters())[1].device)
|
||||
with torch.autocast("cuda", dtype=torch.bfloat16):
|
||||
input_kwargs = {
|
||||
"hidden_states": noisy_model_input,
|
||||
"encoder_hidden_states": encoder_hidden_states,
|
||||
"timestep": timesteps,
|
||||
"encoder_attention_mask": encoder_attention_mask, # B, L
|
||||
"return_dict": False,
|
||||
}
|
||||
if 'hunyuan' in model_type:
|
||||
input_kwargs["guidance"] = torch.tensor(
|
||||
[1000.0],
|
||||
device=noisy_model_input.device,
|
||||
dtype=torch.bfloat16)
|
||||
with set_forward_context(current_timestep=timesteps,
|
||||
attn_metadata=None):
|
||||
model_pred = transformer(**input_kwargs)[0]
|
||||
|
||||
if precondition_outputs:
|
||||
model_pred = noisy_model_input - model_pred * sigmas
|
||||
if precondition_outputs:
|
||||
target = latents
|
||||
else:
|
||||
target = noise - latents
|
||||
|
||||
loss = (torch.mean((model_pred.float() - target.float())**2) /
|
||||
gradient_accumulation_steps)
|
||||
print('device before backwardin context',
|
||||
next(transformer.named_parameters())[1].device)
|
||||
|
||||
print('device before backward out context',
|
||||
next(transformer.named_parameters())[1].device)
|
||||
loss.backward()
|
||||
print('device after backward out context',
|
||||
next(transformer.named_parameters())[1].device)
|
||||
|
||||
avg_loss = loss.detach().clone()
|
||||
sp_group = get_sp_group()
|
||||
sp_group.all_reduce(avg_loss, op=torch.distributed.ReduceOp.AVG)
|
||||
total_loss += avg_loss.item()
|
||||
|
||||
model_parts = [self.transformer]
|
||||
grad_norm = _clip_grad_norm_while_handling_failing_dtensor_cases(
|
||||
[p for m in model_parts for p in m.parameters()],
|
||||
max_grad_norm,
|
||||
foreach=None,
|
||||
)
|
||||
|
||||
optimizer.step()
|
||||
print('device after optimizer step',
|
||||
next(transformer.named_parameters())[1].device)
|
||||
lr_scheduler.step()
|
||||
print('device after scheduler step',
|
||||
next(transformer.named_parameters())[1].device)
|
||||
return total_loss, grad_norm.item()
|
||||
|
||||
def forward(
|
||||
self,
|
||||
batch: ForwardBatch,
|
||||
fastvideo_args: FastVideoArgs,
|
||||
):
|
||||
args = fastvideo_args
|
||||
self.fastvideo_args = args
|
||||
train_dataloader = self.train_dataloader
|
||||
init_steps = self.init_steps
|
||||
lr_scheduler = self.lr_scheduler
|
||||
optimizer = self.optimizer
|
||||
noise_scheduler = self.noise_scheduler
|
||||
noise_random_generator = None
|
||||
|
||||
from diffusers import FlowMatchEulerDiscreteScheduler
|
||||
noise_scheduler = FlowMatchEulerDiscreteScheduler()
|
||||
|
||||
# Train!
|
||||
total_batch_size = (self.world_size * args.gradient_accumulation_steps /
|
||||
args.sp_size * args.train_sp_batch_size)
|
||||
logger.info("***** Running training *****")
|
||||
# logger.info(f" Num examples = {len(train_dataset)}")
|
||||
# logger.info(f" Dataloader size = {len(train_dataloader)}")
|
||||
# logger.info(f" Num Epochs = {args.num_train_epochs}")
|
||||
logger.info(f" Resume training from step {init_steps}")
|
||||
logger.info(
|
||||
f" Instantaneous batch size per device = {args.train_batch_size}")
|
||||
logger.info(
|
||||
f" Total train batch size (w. data & sequence parallel, accumulation) = {total_batch_size}"
|
||||
)
|
||||
logger.info(
|
||||
f" Gradient Accumulation steps = {args.gradient_accumulation_steps}"
|
||||
)
|
||||
logger.info(f" Total optimization steps = {args.max_train_steps}")
|
||||
logger.info(
|
||||
f" Total training parameters per FSDP shard = {sum(p.numel() for p in self.transformer.parameters() if p.requires_grad) / 1e9} B"
|
||||
)
|
||||
# print dtype
|
||||
logger.info(
|
||||
f" Master weight dtype: {self.transformer.parameters().__next__().dtype}"
|
||||
)
|
||||
|
||||
# Potentially load in the weights and states from a previous save
|
||||
if args.resume_from_checkpoint:
|
||||
assert NotImplementedError(
|
||||
"resume_from_checkpoint is not supported now.")
|
||||
# TODO
|
||||
|
||||
progress_bar = tqdm(
|
||||
range(0, args.max_train_steps),
|
||||
initial=init_steps,
|
||||
desc="Steps",
|
||||
# Only show the progress bar once on each machine.
|
||||
disable=self.local_rank > 0,
|
||||
)
|
||||
|
||||
loader_iter = iter(train_dataloader)
|
||||
|
||||
step_times = deque(maxlen=100)
|
||||
|
||||
# todo future
|
||||
for i in range(init_steps):
|
||||
next(loader_iter)
|
||||
# get gpu memory usage
|
||||
gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
|
||||
logger.info(
|
||||
f"GPU memory usage before train_one_step: {gpu_memory_usage} MB")
|
||||
|
||||
for step in range(init_steps + 1, args.max_train_steps + 1):
|
||||
start_time = time.perf_counter()
|
||||
|
||||
loss, grad_norm = self.train_one_step(
|
||||
self.transformer,
|
||||
# args.model_type,
|
||||
"wan",
|
||||
optimizer,
|
||||
lr_scheduler,
|
||||
loader_iter,
|
||||
noise_scheduler,
|
||||
noise_random_generator,
|
||||
args.gradient_accumulation_steps,
|
||||
args.sp_size,
|
||||
args.precondition_outputs,
|
||||
args.max_grad_norm,
|
||||
args.weighting_scheme,
|
||||
args.logit_mean,
|
||||
args.logit_std,
|
||||
args.mode_scale,
|
||||
)
|
||||
gpu_memory_usage = torch.cuda.memory_allocated() / 1024**2
|
||||
logger.info(
|
||||
f"GPU memory usage after train_one_step: {gpu_memory_usage} MB")
|
||||
|
||||
step_time = time.perf_counter() - start_time
|
||||
step_times.append(step_time)
|
||||
avg_step_time = sum(step_times) / len(step_times)
|
||||
|
||||
# Manual gradient checking - only at first step
|
||||
if step == 1 and ENABLE_GRADIENT_CHECK:
|
||||
logger.info(f"Performing gradient check at step {step}")
|
||||
self.setup_gradient_check(args, loader_iter, noise_scheduler,
|
||||
noise_random_generator)
|
||||
|
||||
progress_bar.set_postfix({
|
||||
"loss": f"{loss:.4f}",
|
||||
"step_time": f"{step_time:.2f}s",
|
||||
"grad_norm": grad_norm,
|
||||
})
|
||||
progress_bar.update(1)
|
||||
if self.rank <= 0:
|
||||
wandb.log(
|
||||
{
|
||||
"train_loss": loss,
|
||||
"learning_rate": lr_scheduler.get_last_lr()[0],
|
||||
"step_time": step_time,
|
||||
"avg_step_time": avg_step_time,
|
||||
"grad_norm": grad_norm,
|
||||
},
|
||||
step=step,
|
||||
)
|
||||
if step % args.checkpointing_steps == 0:
|
||||
if args.use_lora:
|
||||
raise NotImplementedError("LoRA is not supported now")
|
||||
# Save LoRA weights
|
||||
# save_lora_checkpoint(transformer, optimizer, rank,
|
||||
# args.output_dir, step, pipe)
|
||||
else:
|
||||
# Your existing checkpoint saving code
|
||||
save_checkpoint_v1(self.transformer, self.rank,
|
||||
args.output_dir, step)
|
||||
self.transformer.train()
|
||||
self.sp_group.barrier()
|
||||
if args.log_validation and step % args.validation_steps == 0:
|
||||
self.log_validation(self.transformer, args, step)
|
||||
|
||||
if args.use_lora:
|
||||
raise NotImplementedError("LoRA is not supported now")
|
||||
# save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps, pipe)
|
||||
else:
|
||||
save_checkpoint_v1(self.transformer, self.rank, args.output_dir,
|
||||
args.max_train_steps)
|
||||
|
||||
if get_sp_group():
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
def main(args):
|
||||
logger.info("Starting training pipeline...")
|
||||
|
||||
pipeline = WanTrainingPipeline.from_pretrained(
|
||||
args.pretrained_model_name_or_path, args=args)
|
||||
args = pipeline.fastvideo_args
|
||||
pipeline.forward(None, args)
|
||||
logger.info("Training pipeline done")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
argv = sys.argv
|
||||
from fastvideo.v1.fastvideo_args import TrainingArgs
|
||||
from fastvideo.v1.utils import FlexibleArgumentParser
|
||||
parser = FlexibleArgumentParser()
|
||||
parser = TrainingArgs.add_cli_args(parser)
|
||||
parser = FastVideoArgs.add_cli_args(parser)
|
||||
args = parser.parse_args()
|
||||
args.use_cpu_offload = False
|
||||
print(args)
|
||||
main(args)
|
||||
@@ -0,0 +1,47 @@
|
||||
export WANDB_BASE_URL="https://api.wandb.ai"
|
||||
export WANDB_MODE=online
|
||||
|
||||
DATA_DIR=data/HD-Mixkit-Finetune-Wan/combined_parquet_dataset
|
||||
VALIDATION_DIR=data/HD-Mixkit-Finetune-Wan/validation_parquet_dataset
|
||||
num_gpus=1
|
||||
# IP=[MASTER NODE IP]
|
||||
|
||||
# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t
|
||||
# --gradient_checkpointing\
|
||||
# --pretrained_model_name_or_path hunyuanvideo-community/HunyuanVideo \
|
||||
# --pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
|
||||
torchrun --nnodes 1 --nproc_per_node $num_gpus\
|
||||
fastvideo/v1/pipelines/distillation_pipeline.py\
|
||||
--mode distill\
|
||||
--pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
|
||||
--cache_dir "/home/test/.cache"\
|
||||
--data_path "$DATA_DIR"\
|
||||
--validation_prompt_dir "$VALIDATION_DIR"\
|
||||
--train_batch_size=1 \
|
||||
--num_latent_t 1 \
|
||||
--sp_size $num_gpus \
|
||||
--train_sp_batch_size 1\
|
||||
--dataloader_num_workers $num_gpus\
|
||||
--gradient_accumulation_steps=1\
|
||||
--max_train_steps=320\
|
||||
--learning_rate=1e-6\
|
||||
--mixed_precision="bf16"\
|
||||
--checkpointing_steps=64\
|
||||
--validation_steps 20\
|
||||
--validation_sampling_steps "2,4,8" \
|
||||
--checkpoints_total_limit 3\
|
||||
--allow_tf32\
|
||||
--ema_start_step 0\
|
||||
--cfg 0.0\
|
||||
--log_validation\
|
||||
--output_dir="$DATA_DIR/outputs/hy_phase1_shift17_bs_16_HD"\
|
||||
--tracker_project_name Hunyuan_Distill \
|
||||
--num_height 720 \
|
||||
--num_width 1280 \
|
||||
--num_frames 125 \
|
||||
--shift 17 \
|
||||
--validation_guidance_scale "1.0" \
|
||||
--num_euler_timesteps 50 \
|
||||
--multi_phased_distill_schedule "4000-1" \
|
||||
--not_apply_cfg_solver \
|
||||
--master_weight_type "bf16"
|
||||
@@ -12,7 +12,7 @@ NUM_GPUS=4
|
||||
torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
|
||||
fastvideo/v1/training/wan_training_pipeline.py\
|
||||
--model_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
|
||||
--inference_mode False\
|
||||
--mode training\
|
||||
--pretrained_model_name_or_path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
|
||||
--cache_dir "/home/ray/.cache"\
|
||||
--data_path "$DATA_DIR"\
|
||||
|
||||
Reference in New Issue
Block a user