test
This commit is contained in:
H1yori233
@@ -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[@]}"
|
||||
@@ -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
|
||||
)
|
||||
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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
|
||||
+23
-13
@@ -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)
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user