diff --git a/examples/training/finetune/MatrixGame2.0/finetune_df.slurm b/examples/training/finetune/MatrixGame2.0/finetune_df.slurm new file mode 100644 index 000000000..d76840f7b --- /dev/null +++ b/examples/training/finetune/MatrixGame2.0/finetune_df.slurm @@ -0,0 +1,125 @@ +#!/bin/bash +#SBATCH --job-name=mg-df +#SBATCH --partition=main +#SBATCH --nodes=4 +#SBATCH --ntasks=4 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=1440G +#SBATCH --output=finetune_df_output/df_%j.out +#SBATCH --error=finetune_df_output/df_%j.err +#SBATCH --exclusive + +set -e -x + +# Environment Setup +source ~/conda/miniconda/bin/activate +conda activate /mnt/weka/home/hao.zhang/conda/miniconda/envs/alexfv +export PYTHONPATH="/mnt/weka/home/hao.zhang/kaiqin/FastVideo_clean:$PYTHONPATH" + +# Basic Info +export WANDB_API_KEY="7ff8b6e8356924f7a6dd51a0342dd1a422ea9352" +export WANDB_MODE="online" +export NCCL_P2P_DISABLE=1 +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 + +echo "MASTER_ADDR: $MASTER_ADDR" +echo "NODE_RANK: $NODE_RANK" + +RUN_NAME=$(date +"%m%d_%H%M") +echo "RUN_NAME: $RUN_NAME" + +# Configs +MODEL_PATH="../mg_models/Matrix-Game-2.0-Foundation-Diffusers" +DATA_DIR="../solaris/datasets/train/preprocessed_0227" +VALIDATION_DATASET_FILE="examples/training/finetune/MatrixGame2.0/validation_mc.json" + +# Training arguments +training_args=( + --tracker_project_name "mg_finetune_solaris_df" + --output_dir "checkpoints/matrixgame_df_${RUN_NAME}" + --wandb_run_name "${RUN_NAME}_df_bs32" + --max_train_steps 12000 + --train_batch_size 1 + --train_sp_batch_size 1 + --gradient_accumulation_steps 1 + --num_latent_t 18 + --num_height 352 + --num_width 640 + --num_frames 69 + --num_frame_per_block 3 + --enable_gradient_checkpointing_type "full" +) + +# Parallel arguments +parallel_args=( + --num_gpus 32 + --sp_size 1 + --tp_size 1 + --hsdp_replicate_dim 1 + --hsdp_shard_dim 32 +) + +# Model arguments +model_args=( + --model_path $MODEL_PATH + --pretrained_model_name_or_path $MODEL_PATH + --override-transformer-cls-name "CausalMatrixGameWanModel" + --flow_shift 5 +) + +# Dataset arguments +dataset_args=( + --data_path "$DATA_DIR" + --dataloader_num_workers 1 +) + +# Validation arguments +validation_args=( + --log_validation + --validation_dataset_file "$VALIDATION_DATASET_FILE" + --validation_steps 500 + --validation_sampling_steps "50" + --validation_guidance_scale "6.0" +) + +# Optimizer arguments +optimizer_args=( + --learning_rate 1e-4 + --betas "0.9,0.95" + --mixed_precision "bf16" + --weight_only_checkpointing_steps 2000 + --training_state_checkpointing_steps 2000 + --weight_decay 0 + --max_grad_norm 1.0 +) + +# Miscellaneous arguments +miscellaneous_args=( + --inference_mode False + --checkpoints_total_limit 3 + --training_cfg_rate 0.1 + --dit_precision "fp32" +) + +mkdir -p finetune_df_output + +srun torchrun \ +--nnodes $SLURM_JOB_NUM_NODES \ +--nproc_per_node 8 \ +--node_rank $SLURM_PROCID \ +--rdzv_backend=c10d \ +--rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ + fastvideo/training/matrixgame_ar_diffusion_pipeline.py \ + "${parallel_args[@]}" \ + "${model_args[@]}" \ + "${dataset_args[@]}" \ + "${training_args[@]}" \ + "${optimizer_args[@]}" \ + "${validation_args[@]}" \ + "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/MatrixGame2.0/finetune_i2v.slurm b/examples/training/finetune/MatrixGame2.0/finetune_i2v.slurm index dfc097b22..256a6cc94 100644 --- a/examples/training/finetune/MatrixGame2.0/finetune_i2v.slurm +++ b/examples/training/finetune/MatrixGame2.0/finetune_i2v.slurm @@ -1,5 +1,5 @@ #!/bin/bash -#SBATCH --job-name=mg-i2v +#SBATCH --job-name=mg-solaris #SBATCH --partition=main #SBATCH --nodes=4 #SBATCH --ntasks=4 @@ -43,16 +43,16 @@ VALIDATION_DATASET_FILE="examples/training/finetune/MatrixGame2.0/validation_mc. training_args=( --tracker_project_name "mg_finetune_solaris" --output_dir "checkpoints/matrixgame_finetune_${RUN_NAME}" - --wandb_run_name "${RUN_NAME}_bs32" + --wandb_run_name "${RUN_NAME}_bs64_more_steps" --max_train_steps 150000 --train_batch_size 1 --train_sp_batch_size 1 - --gradient_accumulation_steps 2 + --gradient_accumulation_steps 1 --num_latent_t 20 --num_height 352 --num_width 640 --num_frames 77 - --enable_gradient_checkpointing_type "full" + # --enable_gradient_checkpointing_type "full" --reinit-action-module True --override-keyboard-dim 23 ) @@ -89,11 +89,12 @@ validation_args=( # Optimizer arguments optimizer_args=( - --learning_rate 2e-5 + --learning_rate 1e-4 + --betas "0.9,0.95" --mixed_precision "bf16" --weight_only_checkpointing_steps 2000 --training_state_checkpointing_steps 2000 - --weight_decay 1e-4 + --weight_decay 0 --max_grad_norm 1.0 ) diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index a22a582e0..d37352ab0 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -93,6 +93,9 @@ _SCHEDULERS = { "FlowMatchEulerDiscreteScheduler": ("schedulers", "scheduling_flow_match_euler_discrete", "FlowMatchEulerDiscreteScheduler"), + "DiffusionForcingScheduler": + ("schedulers", "scheduling_diffusion_forcing", + "DiffusionForcingScheduler"), "UniPCMultistepScheduler": ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"), "FlowUniPCMultistepScheduler": diff --git a/fastvideo/models/schedulers/scheduling_diffusion_forcing.py b/fastvideo/models/schedulers/scheduling_diffusion_forcing.py new file mode 100644 index 000000000..c8fad5b97 --- /dev/null +++ b/fastvideo/models/schedulers/scheduling_diffusion_forcing.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput +import torch + +from fastvideo.models.schedulers.base import BaseScheduler + + +class DiffusionForcingSchedulerOutput(BaseOutput): + prev_sample: torch.FloatTensor + + +class DiffusionForcingScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): + + config_name = "scheduler_config.json" + order = 1 + + @register_to_config + def __init__( + self, + num_inference_steps: int = 100, + num_train_timesteps: int = 1000, + shift: float = 5.0, + sigma_max: float = 1.0, + sigma_min: float = 0.0, + extra_one_step: bool = True, + training: bool = False, + ): + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.extra_one_step = extra_one_step + self.set_timesteps(num_inference_steps, training=training) + + def sigma_from_timestep(self, timestep: torch.Tensor) -> torch.Tensor: + timestep = timestep.to(torch.float32) + sigma = timestep / float(self.num_train_timesteps) + return sigma.clamp_(0.0, 1.0) + + def set_timesteps( + self, + num_inference_steps: int = 100, + denoising_strength: float = 1.0, + training: bool = False, + return_dict: bool = False, + **kwargs, + ): + if training: + self.timesteps = torch.arange( + self.num_train_timesteps, -1, -1, dtype=torch.float32 + ) + self.sigmas = self.sigma_from_timestep(self.timesteps) + self.linear_timesteps_weights = torch.ones_like(self.timesteps) + return + + sigma_start = self.sigma_min + ( + self.sigma_max - self.sigma_min + ) * denoising_strength + if self.extra_one_step: + sigmas = torch.linspace( + sigma_start, self.sigma_min, num_inference_steps + 1 + )[:-1] + else: + sigmas = torch.linspace( + sigma_start, self.sigma_min, num_inference_steps + ) + if self.shift != 1.0: + sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) + self.sigmas = sigmas.to(torch.float32) + self.timesteps = ( + self.sigmas * float(self.num_train_timesteps) + ).to(torch.float32) + + def step( + self, + model_output: torch.FloatTensor, + timestep: torch.FloatTensor, + sample: torch.FloatTensor, + to_final: bool = False, + return_dict: bool = False, + **kwargs, + ): + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + elif timestep.ndim == 0: + timestep = timestep.unsqueeze(0) + + timestep = timestep.to(model_output.device, dtype=torch.float32) + self.timesteps = self.timesteps.to(model_output.device) + + sigma = self.sigma_from_timestep(timestep).reshape(-1, 1, 1, 1) + if to_final: + sigma_next = torch.zeros_like(sigma) + else: + timestep_id = torch.argmin( + (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), + dim=1, + ) + sigma_next = torch.zeros_like(sigma) + valid = timestep_id + 1 < len(self.timesteps) + if valid.any(): + next_timestep = self.timesteps[timestep_id[valid] + 1] + sigma_next[valid] = self.sigma_from_timestep( + next_timestep + ).reshape(-1, 1, 1, 1) + + prev_sample = sample + model_output * (sigma_next - sigma) + if isinstance(prev_sample, torch.Tensor | float) and not return_dict: + return (prev_sample,) + return DiffusionForcingSchedulerOutput(prev_sample=prev_sample) + + @staticmethod + def calculate_alpha_beta_high(sigma, sigma_bound): + alpha = (1 - sigma) / (1 - sigma_bound) + beta = torch.sqrt(sigma**2 - (alpha * sigma_bound) ** 2) + return alpha, beta + + def add_noise(self, original_samples, noise, timestep): + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + sigma = self.sigma_from_timestep(timestep).to(noise.device).reshape( + -1, 1, 1, 1 + ) + sample = (1 - sigma) * original_samples + sigma * noise + return sample.type_as(noise) + + def add_noise_high( + self, original_samples, noise, timestep, boundary_timestep + ): + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + if boundary_timestep.ndim == 2: + boundary_timestep = boundary_timestep.flatten(0, 1) + sigma = self.sigma_from_timestep(timestep).to(noise.device).reshape( + -1, 1, 1, 1 + ) + sigma_boundary = self.sigma_from_timestep(boundary_timestep).to( + noise.device + ).reshape(-1, 1, 1, 1) + alpha, beta = self.calculate_alpha_beta_high(sigma, sigma_boundary) + sample = alpha * original_samples + beta * noise + return sample.type_as(noise) + + def training_target(self, sample, noise, timestep): + return noise - sample + + def training_weight(self, timestep): + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + return torch.ones_like(timestep, dtype=torch.float32) + + def scale_model_input( + self, sample: torch.Tensor, timestep: int | None = None + ) -> torch.Tensor: + return sample + + def set_shift(self, shift: float) -> None: + self.shift = shift + + +EntryClass = DiffusionForcingScheduler diff --git a/fastvideo/models/utils.py b/fastvideo/models/utils.py index 493baeb81..cfade50e0 100644 --- a/fastvideo/models/utils.py +++ b/fastvideo/models/utils.py @@ -173,11 +173,15 @@ def pred_noise_to_pred_video(pred_noise: torch.Tensor, device = pred_noise.device pred_noise = pred_noise.double().to(device) noise_input_latent = noise_input_latent.double().to(device) - sigmas = scheduler.sigmas.double().to(device) - timesteps = scheduler.timesteps.double().to(device) - timestep_id = torch.argmin( - (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) - sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) + if hasattr(scheduler, "sigma_from_timestep"): + sigma_t = scheduler.sigma_from_timestep(timestep).double().to( + device).reshape(-1, 1, 1, 1) + else: + sigmas = scheduler.sigmas.double().to(device) + timesteps = scheduler.timesteps.double().to(device) + timestep_id = torch.argmin( + (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) pred_video = noise_input_latent - sigma_t * pred_noise return pred_video.to(dtype) @@ -217,14 +221,20 @@ def pred_noise_to_x_bound(pred_noise: torch.Tensor, device = pred_noise.device pred_noise = pred_noise.double().to(device) noise_input_latent = noise_input_latent.double().to(device) - sigmas = scheduler.sigmas.double().to(device) - timesteps = scheduler.timesteps.double().to(device) - timestep_id = torch.argmin( - (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) - sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) + if hasattr(scheduler, "sigma_from_timestep"): + sigma_t = scheduler.sigma_from_timestep(timestep).double().to( + device).reshape(-1, 1, 1, 1) + sigma_t_boundary = scheduler.sigma_from_timestep( + boundary_timestep).double().to(device).reshape(-1, 1, 1, 1) + else: + sigmas = scheduler.sigmas.double().to(device) + timesteps = scheduler.timesteps.double().to(device) + timestep_id = torch.argmin( + (timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) + sigma_t = sigmas[timestep_id].reshape(-1, 1, 1, 1) - boundary_timestep_id = torch.argmin( - (timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1) - sigma_t_boundary = sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1) + boundary_timestep_id = torch.argmin( + (timesteps.unsqueeze(0) - boundary_timestep.unsqueeze(1)).abs(), dim=1) + sigma_t_boundary = sigmas[boundary_timestep_id].reshape(-1, 1, 1, 1) pred_video = noise_input_latent - (sigma_t - sigma_t_boundary) * pred_noise return pred_video.to(dtype) \ No newline at end of file diff --git a/fastvideo/training/matrixgame_ar_diffusion_pipeline.py b/fastvideo/training/matrixgame_ar_diffusion_pipeline.py new file mode 100644 index 000000000..1a7bb4a56 --- /dev/null +++ b/fastvideo/training/matrixgame_ar_diffusion_pipeline.py @@ -0,0 +1,543 @@ +# SPDX-License-Identifier: Apache-2.0 + +import sys +from copy import deepcopy +from typing import Any, cast + +import torch +import torch.nn.functional as F + +from fastvideo.configs.sample import SamplingParam +from fastvideo.dataset.dataloader.schema import pyarrow_schema_matrixgame +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs, TrainingArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_diffusion_forcing import ( + DiffusionForcingScheduler, +) +from fastvideo.pipelines.basic.matrixgame.matrixgame_causal_dmd_pipeline import ( + MatrixGameCausalDMDPipeline, +) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch, TrainingBatch +from fastvideo.training.training_pipeline import TrainingPipeline +from fastvideo.training.training_utils import ( + clip_grad_norm_while_handling_failing_dtensor_cases, +) +from fastvideo.utils import shallow_asdict + +logger = init_logger(__name__) + + +class MatrixGameARDiffusionPipeline(TrainingPipeline): + + _required_config_modules = ["scheduler", "transformer", "vae"] + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + self.modules["scheduler"] = DiffusionForcingScheduler( + shift=fastvideo_args.pipeline_config.flow_shift, + sigma_min=0.0, + extra_one_step=True, + ) + self.modules["scheduler"].set_timesteps( + num_inference_steps=1000, training=True + ) + + def set_schemas(self): + self.train_dataset_schema = pyarrow_schema_matrixgame + + def _get_temporal_compression_ratio(self) -> int: + assert self.training_args is not None + return int( + self.training_args.pipeline_config.vae_config.arch_config + .temporal_compression_ratio + ) + + def _resolve_num_frame_per_block( + self, training_args: TrainingArgs + ) -> int: + transformer_num_frame_per_block = getattr( + self.transformer, "num_frame_per_block", None + ) + requested_num_frame_per_block = getattr( + training_args, "num_frame_per_block", None + ) + + if ( + transformer_num_frame_per_block is not None + and requested_num_frame_per_block is not None + and transformer_num_frame_per_block + != requested_num_frame_per_block + ): + raise ValueError( + "num_frame_per_block mismatch between loaded transformer and " + "training args: " + f"transformer={transformer_num_frame_per_block}, " + f"training_args={requested_num_frame_per_block}" + ) + + if transformer_num_frame_per_block is not None: + return int(transformer_num_frame_per_block) + if requested_num_frame_per_block is not None: + return int(requested_num_frame_per_block) + return 3 + + def initialize_training_pipeline(self, training_args: TrainingArgs): + super().initialize_training_pipeline(training_args) + + self.vae = self.get_module("vae") + self.vae.requires_grad_(False) + + self.num_frame_per_block = self._resolve_num_frame_per_block( + training_args + ) + self.timestep_shift = training_args.pipeline_config.flow_shift + self.ar_noise_scheduler = DiffusionForcingScheduler( + shift=self.timestep_shift, sigma_min=0.0, extra_one_step=True + ) + self.ar_noise_scheduler.set_timesteps( + num_inference_steps=1000, training=True + ) + + logger.info( + "MatrixGame AR diffusion pipeline initialized with " + "num_frame_per_block=%d, diffusion_forcing_shift=%.1f", + self.num_frame_per_block, + self.timestep_shift, + ) + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + logger.info("Initializing MatrixGame AR validation pipeline...") + args_copy = deepcopy(training_args) + args_copy.inference_mode = True + + validation_scheduler = DiffusionForcingScheduler( + shift=args_copy.pipeline_config.flow_shift, + sigma_min=0.0, + extra_one_step=True, + ) + validation_scheduler.set_timesteps( + num_inference_steps=1000, training=False + ) + + num_val_steps = int( + training_args.validation_sampling_steps.split(",")[0] + ) + step_size = 1000 // num_val_steps + args_copy.pipeline_config.dmd_denoising_steps = list( + range(1000, 0, -step_size) + ) + args_copy.pipeline_config.warp_denoising_step = True + training_args.pipeline_config.dmd_denoising_steps = ( + args_copy.pipeline_config.dmd_denoising_steps + ) + training_args.pipeline_config.warp_denoising_step = True + + logger.info( + "Validation: %d-step MatrixGame causal denoising, " + "dmd_denoising_steps has %d entries", + num_val_steps, + len(args_copy.pipeline_config.dmd_denoising_steps), + ) + + loaded_modules = { + "transformer": self.get_module("transformer"), + "vae": self.get_module("vae"), + "scheduler": validation_scheduler, + } + image_encoder = self.get_module("image_encoder") + image_processor = self.get_module("image_processor") + if image_encoder is not None: + loaded_modules["image_encoder"] = image_encoder + if image_processor is not None: + loaded_modules["image_processor"] = image_processor + + self.validation_pipeline = MatrixGameCausalDMDPipeline.from_pretrained( + training_args.model_path, + args=args_copy, + inference_mode=True, + loaded_modules=loaded_modules, + tp_size=training_args.tp_size, + sp_size=training_args.sp_size, + num_gpus=training_args.num_gpus, + pin_cpu_memory=training_args.pin_cpu_memory, + dit_cpu_offload=True, + ) + + def _get_timestep( + self, + min_timestep: int, + max_timestep: int, + batch_size: int, + num_frame: int, + num_frame_per_block: int, + uniform_timestep: bool = False, + ) -> torch.Tensor: + """ + Sample latents timesteps. + """ + device = get_local_torch_device() + if uniform_timestep: + timestep = torch.randint( + min_timestep, + max_timestep + 1, + [batch_size, 1], + device=device, + dtype=torch.long, + ).repeat(1, num_frame) + return timestep + + timestep = torch.randint( + min_timestep, + max_timestep + 1, + [batch_size, num_frame], + device=device, + dtype=torch.long, + ) + return timestep + + def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch: + batch = next(self.train_loader_iter, None) # type: ignore + if batch is None: + self.current_epoch += 1 + logger.info("Starting epoch %s", self.current_epoch) + if hasattr(self, "train_dataset") and hasattr( + self.train_dataset, "sampler" + ): + self.train_dataset.sampler.set_epoch(self.current_epoch) + self.train_loader_iter = iter(self.train_dataloader) + batch = next(self.train_loader_iter) + + latents = batch["vae_latent"] + latents = latents[:, :, : self.training_args.num_latent_t] + clip_features = batch["clip_feature"] + image_latents = batch["first_frame_latent"] + image_latents = image_latents[:, :, : self.training_args.num_latent_t] + pil_image = batch["pil_image"] + infos = batch["info_list"] + + training_batch.latents = latents.to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + training_batch.encoder_hidden_states = None + training_batch.encoder_attention_mask = None + training_batch.preprocessed_image = pil_image.to( + get_local_torch_device() + ) + training_batch.image_embeds = clip_features.to(get_local_torch_device()) + training_batch.image_latents = image_latents.to( + get_local_torch_device() + ) + training_batch.infos = infos + + if "mouse_cond" in batch and batch["mouse_cond"].numel() > 0: + training_batch.mouse_cond = batch["mouse_cond"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + else: + training_batch.mouse_cond = None + + if "keyboard_cond" in batch and batch["keyboard_cond"].numel() > 0: + training_batch.keyboard_cond = batch["keyboard_cond"].to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + else: + training_batch.keyboard_cond = None + + temporal_compression_ratio = self._get_temporal_compression_ratio() + expected_num_frames = ( + self.training_args.num_latent_t - 1 + ) * temporal_compression_ratio + 1 + if training_batch.keyboard_cond is not None: + assert training_batch.keyboard_cond.shape[1] >= expected_num_frames, ( + f"keyboard_cond has {training_batch.keyboard_cond.shape[1]} " + f"frames but need at least {expected_num_frames}" + ) + training_batch.keyboard_cond = training_batch.keyboard_cond[ + :, :expected_num_frames + ] + if training_batch.mouse_cond is not None: + assert training_batch.mouse_cond.shape[1] >= expected_num_frames, ( + f"mouse_cond has {training_batch.mouse_cond.shape[1]} frames " + f"but need at least {expected_num_frames}" + ) + training_batch.mouse_cond = training_batch.mouse_cond[ + :, :expected_num_frames + ] + + return training_batch + + def _prepare_dit_inputs( + self, training_batch: TrainingBatch + ) -> TrainingBatch: + """Prepare diffusion-forcing inputs and MatrixGame I2V concat.""" + assert self.training_args is not None + latents = training_batch.latents + assert latents is not None + batch_size = latents.shape[0] + num_latent_t = latents.shape[2] + + latents_btchw = latents.permute(0, 2, 1, 3, 4) + + timesteps = self._get_timestep( + min_timestep=0, + max_timestep=self.ar_noise_scheduler.num_train_timesteps, + batch_size=batch_size, + num_frame=num_latent_t, + num_frame_per_block=self.num_frame_per_block, + uniform_timestep=False, + ) + + noise_generator = getattr(self, "noise_gen_cuda", None) + noise = torch.randn( + latents_btchw.shape, + generator=noise_generator, + device=latents_btchw.device, + dtype=latents_btchw.dtype, + ) + + noisy_latents = self.ar_noise_scheduler.add_noise( + latents_btchw.flatten(0, 1), + noise.flatten(0, 1), + timesteps.flatten(0, 1), + ).unflatten(0, (batch_size, num_latent_t)) + + noisy_model_input = noisy_latents.permute(0, 2, 1, 3, 4) + + assert isinstance(training_batch.image_latents, torch.Tensor) + image_latents = training_batch.image_latents.to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + + temporal_compression_ratio = self._get_temporal_compression_ratio() + num_frames = (num_latent_t - 1) * temporal_compression_ratio + 1 + _, _, _, latent_height, latent_width = image_latents.shape + mask_lat_size = torch.ones( + batch_size, 1, num_frames, latent_height, latent_width + ) + mask_lat_size[:, :, 1:] = 0 + + first_frame_mask = mask_lat_size[:, :, :1] + first_frame_mask = torch.repeat_interleave( + first_frame_mask, + dim=2, + repeats=temporal_compression_ratio, + ) + mask_lat_size = torch.cat( + [first_frame_mask, mask_lat_size[:, :, 1:]], dim=2 + ) + mask_lat_size = mask_lat_size.view( + batch_size, + -1, + temporal_compression_ratio, + latent_height, + latent_width, + ) + mask_lat_size = mask_lat_size.transpose(1, 2) + mask_lat_size = mask_lat_size.to( + image_latents.device, dtype=torch.bfloat16 + ) + + noisy_model_input = torch.cat( + [noisy_model_input, mask_lat_size, image_latents], dim=1 + ) + + training_target = self.ar_noise_scheduler.training_target( + latents_btchw.flatten(0, 1), + noise.flatten(0, 1), + timesteps.flatten(0, 1), + ).unflatten(0, (batch_size, num_latent_t)) + + training_batch.noisy_model_input = noisy_model_input + training_batch.timesteps = timesteps + training_batch.noise = noise.permute(0, 2, 1, 3, 4) + training_batch.raw_latent_shape = latents.shape + training_batch._ar_training_target = training_target + + return training_batch + + def _build_input_kwargs( + self, training_batch: TrainingBatch + ) -> TrainingBatch: + image_embeds = training_batch.image_embeds + assert isinstance(image_embeds, torch.Tensor) + assert torch.isnan(image_embeds).sum() == 0 + image_embeds = image_embeds.to( + get_local_torch_device(), dtype=torch.bfloat16 + ) + + timesteps = training_batch.timesteps + assert isinstance(timesteps, torch.Tensor) + assert timesteps.ndim == 2, ( + f"Expected per-frame timesteps [B, T], got shape {timesteps.shape}" + ) + + training_batch.input_kwargs = { + "hidden_states": training_batch.noisy_model_input, + "encoder_hidden_states": None, + "timestep": timesteps.to(get_local_torch_device()), + "encoder_hidden_states_image": image_embeds, + "mouse_cond": training_batch.mouse_cond, + "keyboard_cond": training_batch.keyboard_cond, + "num_frame_per_block": self.num_frame_per_block, + "return_dict": False, + } + return training_batch + + def _transformer_forward_and_compute_loss( + self, training_batch: TrainingBatch + ) -> TrainingBatch: + """ + Run transformer forward pass and compute diffusion-forcing loss. + """ + input_kwargs = training_batch.input_kwargs + + with set_forward_context( + current_timestep=training_batch.timesteps, + attn_metadata=None, + forward_batch=None, + ): + model_pred = self.transformer(**input_kwargs) + model_pred_btchw = model_pred.permute(0, 2, 1, 3, 4) + + training_target = training_batch._ar_training_target.to( + model_pred_btchw.device, dtype=model_pred_btchw.dtype + ) + + loss = F.mse_loss( + model_pred_btchw.float(), + training_target.float(), + reduction="mean", + ) + loss = loss / self.training_args.gradient_accumulation_steps + loss.backward() + + avg_loss = loss.detach().clone() + training_batch.total_loss += avg_loss.item() + + return training_batch + + def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch: + self.transformer.train() + self.optimizer.zero_grad() + training_batch.total_loss = 0.0 + args = cast(TrainingArgs, self.training_args) + + for _ in range(args.gradient_accumulation_steps): + training_batch = self._get_next_batch(training_batch) + training_batch = self._normalize_dit_input(training_batch) + training_batch = self._prepare_dit_inputs(training_batch) + training_batch = self._build_input_kwargs(training_batch) + training_batch = self._transformer_forward_and_compute_loss( + training_batch + ) + + grad_norm = clip_grad_norm_while_handling_failing_dtensor_cases( + [p for p in self.transformer.parameters() if p.requires_grad], + args.max_grad_norm if args.max_grad_norm is not None else 0.0, + ) + + self.optimizer.step() + self.lr_scheduler.step() + + if grad_norm is None: + grad_value = 0.0 + else: + try: + if isinstance(grad_norm, torch.Tensor): + grad_value = float(grad_norm.detach().float().item()) + else: + grad_value = float(grad_norm) + except Exception: + grad_value = 0.0 + training_batch.grad_norm = grad_value + assert training_batch.latents is not None + training_batch.raw_latent_shape = training_batch.latents.shape + return training_batch + + def _prepare_validation_batch( + self, + sampling_param: SamplingParam, + training_args: TrainingArgs, + validation_batch: dict[str, Any], + num_inference_steps: int, + ) -> ForwardBatch: + sampling_param.prompt = validation_batch["prompt"] + sampling_param.height = training_args.num_height + sampling_param.width = training_args.num_width + sampling_param.image_path = validation_batch.get( + "image_path" + ) or validation_batch.get("video_path") + sampling_param.num_inference_steps = num_inference_steps + sampling_param.data_type = "video" + assert self.seed is not None + sampling_param.seed = self.seed + + 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] + temporal_compression_factor = ( + training_args.pipeline_config.vae_config.arch_config + .temporal_compression_ratio + ) + num_frames = ( + training_args.num_latent_t - 1 + ) * temporal_compression_factor + 1 + sampling_param.num_frames = num_frames + batch = ForwardBatch( + **shallow_asdict(sampling_param), + latents=None, + generator=torch.Generator(device="cpu").manual_seed(self.seed), + n_tokens=n_tokens, + eta=0.0, + VSA_sparsity=training_args.VSA_sparsity, + ) + if "image" in validation_batch and validation_batch["image"] is not None: + batch.pil_image = validation_batch["image"] + + if ( + "keyboard_cond" in validation_batch + and validation_batch["keyboard_cond"] is not None + ): + keyboard_cond = validation_batch["keyboard_cond"] + keyboard_cond = torch.tensor(keyboard_cond, dtype=torch.bfloat16) + keyboard_cond = keyboard_cond.unsqueeze(0) + batch.keyboard_cond = keyboard_cond + + if ( + "mouse_cond" in validation_batch + and validation_batch["mouse_cond"] is not None + ): + mouse_cond = validation_batch["mouse_cond"] + mouse_cond = torch.tensor(mouse_cond, dtype=torch.bfloat16) + mouse_cond = mouse_cond.unsqueeze(0) + batch.mouse_cond = mouse_cond + + return batch + + +def main(args) -> None: + logger.info("Starting MatrixGame AR diffusion training pipeline...") + + pipeline = MatrixGameARDiffusionPipeline.from_pretrained( + args.pretrained_model_name_or_path, args=args + ) + args = pipeline.training_args + pipeline.train() + logger.info("MatrixGame AR diffusion training pipeline done") + + +if __name__ == "__main__": + argv = sys.argv + from fastvideo.fastvideo_args import TrainingArgs + from fastvideo.utils import FlexibleArgumentParser + + parser = FlexibleArgumentParser() + parser = TrainingArgs.add_cli_args(parser) + parser = FastVideoArgs.add_cli_args(parser) + args = parser.parse_args() + args.dit_cpu_offload = False + main(args)