Ode init running for hy15

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[@]}"

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")

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,
+    )

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

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(