diff --git a/examples/training/consistency_finetune/causal_ode_init/hy15_finetune_ode_init.slurm b/examples/training/consistency_finetune/causal_ode_init/hy15_finetune_ode_init.slurm new file mode 100644 index 000000000..b5fa4d395 --- /dev/null +++ b/examples/training/consistency_finetune/causal_ode_init/hy15_finetune_ode_init.slurm @@ -0,0 +1,123 @@ +#!/bin/bash +#SBATCH --job-name=hy15_ode_1333 +#SBATCH --partition=main +#SBATCH --nodes=8 +#SBATCH --ntasks=8 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=hy15_ode_1333_output/hy15_ode_1333_%j.out +#SBATCH --error=hy15_ode_1333_output/hy15_ode_1333_%j.err +#SBATCH --exclusive + +# Basic Info +export NCCL_P2P_DISABLE=1 +export TORCH_NCCL_ENABLE_MONITORING=0 +export NCCL_DEBUG_SUBSYS=INIT,NET +# different cache dir for different processes +export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} +export MASTER_PORT=29500 +export NODE_RANK=$SLURM_PROCID +nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) +export MASTER_ADDR=${nodes[0]} +export TOKENIZERS_PARALLELISM=false +# export WANDB_API_KEY="8d9f4b39abd68eb4e29f6fc010b7ee71a2207cde" +export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45" +# export WANDB_API_KEY='your_wandb_api_key_here' +export WANDB_BASE_URL="https://api.wandb.ai" +export WANDB_MODE=online +export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN + +source ~/conda/miniconda/bin/activate +conda activate wei-fv-distill +export HOME="/mnt/weka/home/hao.zhang/wei" + +# Configs +NUM_GPUS=8 +NUM_TOTAL_GPUS=64 + +MODEL_PATH="hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v" +DATA_DIR="/mnt/weka/home/hao.zhang/wl/sharefs/wl/release/fv-ode-preprocessing-16k-hy15-121" +# IP=[MASTER NODE IP] + +# Training arguments +training_args=( + --tracker_project_name "hy15_ode_init" + --output_dir "/mnt/weka/home/hao.zhang/wei/hy15_ode_init_1333" + --wandb_run_name "hy15_ode_init_1333" + # --resume_from_checkpoint "ode_init_diffusers/" + # --warp_denoising_step + --log_visualization + --visualization_steps 100 + --max_train_steps 6001 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 2 + --num_latent_t 31 + --num_height 480 + --num_width 848 + --num_frames 121 + --dmd_denoising_steps "1000,750,500,250" + --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus $NUM_TOTAL_GPUS # 64 + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 # 64 + --hsdp_shard_dim $NUM_TOTAL_GPUS +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 4 +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 1e-5 + --mixed_precision "bf16" + --training_state_checkpointing_steps 400 + --weight_only_checkpointing_steps 200 + --weight_decay 0.01 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 6 + --training_cfg_rate 0.1 + --multi_phased_distill_schedule "4000-1" + --not_apply_cfg_solver + --dit_precision "fp32" + --num_euler_timesteps 50 + --ema_start_step 0 + # --enable_gradient_checkpointing_type "full" +) + +# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t +srun torchrun \ +--nnodes $SLURM_JOB_NUM_NODES \ +--nproc_per_node $NUM_GPUS \ +--node_rank $SLURM_PROCID \ +--rdzv_backend=c10d \ +--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + fastvideo/training/ode_causal_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" \ No newline at end of file diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index 33c06efcc..18bc5e532 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -859,6 +859,7 @@ class TrainingArgs(FastVideoArgs): validation_sampling_steps: str = "" validation_guidance_scale: str = "" validation_steps: float = 0.0 + visualization_steps: float = 0.0 log_validation: bool = False trackers: list[str] = dataclasses.field(default_factory=list) tracker_project_name: str = "" @@ -1091,6 +1092,9 @@ class TrainingArgs(FastVideoArgs): parser.add_argument("--validation-steps", type=float, help="Number of validation steps") + parser.add_argument("--visualization-steps", + type=float, + help="Number of visualization steps") parser.add_argument("--log-validation", action=StoreBoolean, help="Whether to log validation results") diff --git a/fastvideo/optim/muon.py b/fastvideo/optim/muon.py new file mode 100644 index 000000000..c5dd150ca --- /dev/null +++ b/fastvideo/optim/muon.py @@ -0,0 +1,229 @@ +import math +import torch +try: + from torch.distributed.tensor import DTensor + from torch.distributed.tensor.placement_types import ( + Partial, + Placement, + Replicate, + Shard, + ) +except ImportError: + # handle old pytorch versions + Dtensor = None + +# This code is modified from the GitHub repository of KellerJordan: +# https://github.com/KellerJordan/Muon/blob/master/muon.py +def zeropower_via_newtonschulz5(G, steps = 5): + """ + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a + quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose + of minimizing steps, it turns out to be empirically effective to keep increasing the slope at + zero even beyond the point where the iteration no longer converges all the way to one everywhere + on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T + where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model + performance at all relative to UV^T, where USV^T = G is the SVD. + """ + if isinstance(G, DTensor): + device_mesh = G.device_mesh + G = G.full_tensor() + else: + device_mesh = None + + assert len(G.shape) >= 2 + a, b, c = (3.4445, -4.7750, 2.0315) + X = G + if G.size(-2) > G.size(-1): + X = X.mT + # Ensure spectral norm is at most 1 + X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7) + # Perform the NS iterations + for _ in range(steps): + A = X @ X.T + B = b * A + c * A @ A # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng + X = a * X + B @ X + + if G.size(-2) > G.size(-1): + X = X.mT + + if device_mesh != None: + return DTensor.from_local(X, device_mesh) + else: + return X + +class Muon(torch.optim.Optimizer): + """ + Muon - MomentUm Orthogonalized by Newton-schulz + + Arguments: + muon_params: The parameters to be optimized by Muon. + lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default) + momentum: The momentum used by the internal SGD. (0.95 is a good default) + nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended) + ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough) + adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are + {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well. + adamw_lr: The learning rate for the internal AdamW. + adamw_betas: The betas for the internal AdamW. + adamw_eps: The epsilon for the internal AdamW. + adamw_wd: The weight decay for the internal AdamW. + """ + + def __init__( + self, + lr=1e-3, + wd=0.1, + muon_params=None, + momentum=0.95, + nesterov=True, + ns_steps=5, + adamw_params=None, + adamw_betas=(0.95, 0.95), + adamw_eps=1e-8, + ): + + defaults = dict( + lr=lr, + wd=wd, + momentum=momentum, + nesterov=nesterov, + ns_steps=ns_steps, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + ) + + params = list(muon_params) + adamw_params = list(adamw_params) if adamw_params is not None else [] + params.extend(adamw_params) + super().__init__(params, defaults) + # Sort parameters into those for which we will use Muon, and those for which we will not + for p in muon_params: + # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer + assert p.ndim >= 2, p.ndim + self.state[p]["use_muon"] = True + for p in adamw_params: + # Do not use Muon for parameters in adamw_params + self.state[p]["use_muon"] = False + + def adjust_lr_for_muon(self, lr, param_shape): + A, B = param_shape[:2] + # We adjust the learning rate and weight decay based on the size of the parameter matrix + # as describted in the paper + adjusted_ratio = 0.2 * math.sqrt(max(A, B)) + adjusted_lr = lr * adjusted_ratio + return adjusted_lr + + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + + ############################ + # Muon # + ############################ + + params = [p for p in group["params"] if self.state[p]["use_muon"]] + + lr = group["lr"] + wd = group["wd"] + momentum = group["momentum"] + + # generate weight updates in distributed fashion + for p in params: + # sanity check + g = p.grad + if g is None: + continue + if g.ndim > 2: + g = g.view(g.size(0), -1) + assert g is not None + + # calc update + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(g) + buf = state["momentum_buffer"] + buf.mul_(momentum).add_(g) + if group["nesterov"]: + g = g.add(buf, alpha=momentum) + else: + g = buf + g = g.bfloat16() + u = zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) + + # scale update + adjusted_lr = self.adjust_lr_for_muon(lr, p.shape) + + # apply weight decay + p.data.mul_(1 - lr * wd) + + # apply update + p.data.add_(u.view(p.shape), alpha=-adjusted_lr) + + ############################ + # AdamW backup # + ############################ + + params = [p for p in group["params"] if not self.state[p]["use_muon"]] + lr = group['lr'] + beta1, beta2 = group["adamw_betas"] + eps = group["adamw_eps"] + weight_decay = group["wd"] + + for p in params: + g = p.grad + if g is None: + continue + state = self.state[p] + if "step" not in state: + state["step"] = 0 + state["moment1"] = torch.zeros_like(g) + state["moment2"] = torch.zeros_like(g) + state["step"] += 1 + step = state["step"] + buf1 = state["moment1"] + buf2 = state["moment2"] + buf1.lerp_(g, 1 - beta1) + buf2.lerp_(g.square(), 1 - beta2) + + g = buf1 / (eps + buf2.sqrt()) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + scale = bias_correction1 / bias_correction2**0.5 + p.data.mul_(1 - lr * weight_decay) + p.data.add_(g, alpha=-lr / scale) + + return loss + +# help function to create the Muon optimizer +def get_muon_optimizer(model, lr=1e-3, weight_decay=0.1, momentum=0.95, adamw_betas=(0.95, 0.95), adamw_eps=1e-8): + muon_params = [ + p + for name, p in model.named_parameters() + if p.requires_grad and p.ndim >= 2 + ] + adamw_params = [ + p + for name, p in model.named_parameters() + if p.requires_grad and not (p.ndim >= 2) + ] + + return Muon( + lr=lr, + wd=weight_decay, + muon_params=muon_params, + momentum=momentum, + adamw_params=adamw_params, + adamw_betas=adamw_betas, + adamw_eps=adamw_eps, + ) \ No newline at end of file diff --git a/fastvideo/training/ode_causal_pipeline.py b/fastvideo/training/ode_causal_pipeline.py index 8f4f6f245..df2e3b120 100644 --- a/fastvideo/training/ode_causal_pipeline.py +++ b/fastvideo/training/ode_causal_pipeline.py @@ -22,6 +22,8 @@ from fastvideo.pipelines.pipeline_batch_info import TrainingBatch from fastvideo.training.training_pipeline import TrainingPipeline from fastvideo.training.training_utils import ( clip_grad_norm_while_handling_failing_dtensor_cases) +from fastvideo.pipelines.stages.text_encoding import TextEncodingStage +from fastvideo.pipelines.stages.decoding import DecodingStage logger = init_logger(__name__) @@ -36,7 +38,7 @@ class ODEInitTrainingPipeline(TrainingPipeline): - minimizing MSE to the stored next latent at timestep t_next """ - _required_config_modules = ["scheduler", "transformer", "vae"] + _required_config_modules = ["scheduler", "transformer", "vae", "text_encoder", "tokenizer"] # def initialize_pipeline(self, fastvideo_args: FastVideoArgs): # # Match the preprocess/generation scheduler for consistent stepping @@ -57,6 +59,18 @@ class ODEInitTrainingPipeline(TrainingPipeline): self.vae = self.get_module("vae") self.vae.requires_grad_(False) + self.text_encoder = self.get_module("text_encoder") + self.text_encoder.requires_grad_(False) + + self.add_stage(stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder"), self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer"), self.get_module("tokenizer")], + )) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + self.timestep_shift = self.training_args.pipeline_config.flow_shift assert self.timestep_shift == 5.0, "flow_shift must be 5.0" # self.noise_scheduler = SelfForcingFlowMatchScheduler( @@ -65,7 +79,7 @@ class ODEInitTrainingPipeline(TrainingPipeline): logger.info("dmd_denoising_steps: %s", self.training_args.pipeline_config.dmd_denoising_steps) - self.dmd_denoising_steps = torch.tensor([0, 12, 24, 36, 49], + self.dmd_denoising_steps = torch.tensor([0, 12, 24, 36, 50], dtype=torch.long, device=get_local_torch_device()) if training_args.warp_denoising_step: # Warp the denoising step according to the scheduler time shift @@ -134,6 +148,16 @@ class ODEInitTrainingPipeline(TrainingPipeline): ] infos = batch['info_list'] + if encoder_hidden_states.dim() < 3: + prompt_embeds_list, prompt_masks_list = self.prompt_encoding_stage.encode_text( + infos[0]["caption"], + self.training_args, + encoder_index=[0], + return_attention_mask=True, + ) + encoder_hidden_states = prompt_embeds_list[0].to(device, dtype=torch.bfloat16) + encoder_attention_mask = prompt_masks_list[0].to(device, dtype=torch.bfloat16) + # Trajectory tensors may include a leading singleton batch dim per row trajectory_latents = batch['trajectory_latents'] if trajectory_latents.dim() == 7: @@ -159,7 +183,7 @@ class ODEInitTrainingPipeline(TrainingPipeline): f"Unexpected trajectory_timesteps dim: {trajectory_timesteps.dim()}" ) # [B, S, C, T, H, W] -> [B, S, T, C, H, W] to match self-forcing - trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5)[:, :, :19] + trajectory_latents = trajectory_latents.permute(0, 1, 3, 2, 4, 5) # Move to device training_batch.encoder_hidden_states = [encoder_hidden_states, encoder_hidden_states_2] @@ -281,6 +305,7 @@ class ODEInitTrainingPipeline(TrainingPipeline): latent_vis_dict["x0"] = target_latent.permute(0, 2, 1, 3, 4).detach().clone().cpu() + logger.info("timestep: %s", timestep) input_kwargs = { "hidden_states": noisy_input.permute(0, 2, 1, 3, 4), "encoder_hidden_states": encoder_hidden_states, @@ -378,14 +403,14 @@ class ODEInitTrainingPipeline(TrainingPipeline): assert latent_key in latents_vis_dict and latents_vis_dict[ latent_key] is not None latent = latents_vis_dict[latent_key] - pixel_latent = self.validation_pipeline.decoding_stage.decode( + pixel_latent = self.decoding_stage.decode( latent, training_args) video = pixel_latent.cpu().float() video = video.permute(0, 2, 1, 3, 4) video = (video * 255).numpy().astype(np.uint8) video_artifact = self.tracker.video( - video, fps=16, format="mp4") # change to 16 for Wan2.1 + video, fps=24, format="mp4") # change to 16 for Wan2.1 if video_artifact is not None: tracker_loss_dict[latent_key] = video_artifact # Clean up references diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index bc23a3b07..fdf482faa 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -47,6 +47,7 @@ from fastvideo.training.training_utils import ( shard_latents_across_sp) from fastvideo.utils import (is_vmoba_available, is_vsa_available, set_random_seed, shallow_asdict) +from fastvideo.optim.muon import get_muon_optimizer vsa_available = is_vsa_available() vmoba_available = is_vmoba_available() @@ -130,12 +131,18 @@ class TrainingPipeline(LoRAPipeline, ABC): betas_str = training_args.betas betas = tuple(float(x.strip()) for x in betas_str.split(",")) - self.optimizer = torch.optim.AdamW( - params_to_optimize, + # self.optimizer = torch.optim.AdamW( + # params_to_optimize, + # lr=training_args.learning_rate, + # betas=betas, + # weight_decay=training_args.weight_decay, + # eps=1e-8, + # ) + + self.optimizer = get_muon_optimizer( + self.transformer, lr=training_args.learning_rate, - betas=betas, weight_decay=training_args.weight_decay, - eps=1e-8, ) self.init_steps = 0 @@ -663,7 +670,7 @@ class TrainingPipeline(LoRAPipeline, ABC): "grad_norm": grad_norm, "vsa_sparsity": current_vsa_sparsity, } - metrics["batch_size"] = int(training_batch.raw_latent_shape[0]) + # metrics["batch_size"] = int(training_batch.raw_latent_shape[0]) patch_t, patch_h, patch_w = self.training_args.pipeline_config.dit_config.patch_size seq_len = (training_batch.raw_latent_shape[2] // patch_t) * ( @@ -675,14 +682,14 @@ class TrainingPipeline(LoRAPipeline, ABC): else: context_len = 0 - metrics["dit_seq_len"] = int(seq_len) - metrics["context_len"] = context_len + # metrics["dit_seq_len"] = int(seq_len) + # metrics["context_len"] = context_len - arch_config = self.training_args.pipeline_config.dit_config.arch_config + # arch_config = self.training_args.pipeline_config.dit_config.arch_config - metrics["hidden_dim"] = arch_config.hidden_size - metrics["num_layers"] = arch_config.num_layers - metrics["ffn_dim"] = arch_config.ffn_dim + # metrics["hidden_dim"] = arch_config.hidden_size + # metrics["num_layers"] = arch_config.num_layers + # metrics["ffn_dim"] = arch_config.ffn_dim self.tracker.log(metrics, step) if step % self.training_args.training_state_checkpointing_steps == 0: @@ -695,12 +702,13 @@ class TrainingPipeline(LoRAPipeline, ABC): self.noise_random_generator) self.transformer.train() self.sp_group.barrier() + + if self.training_args.log_visualization and step % self.training_args.visualization_steps == 0: + self.visualize_intermediate_latents(training_batch, self.training_args, step) + if self.training_args.log_validation and step % self.training_args.validation_steps == 0: with self.profiler_controller.region( "profiler_region_training_validation"): - if self.training_args.log_visualization: - self.visualize_intermediate_latents( - training_batch, self.training_args, step) self._log_validation(self.transformer, self.training_args, step) gpu_memory_usage = current_platform.get_torch_device(