diff --git a/examples/training/finetune/MatrixGame2.0/finetune_i2v.slurm b/examples/training/finetune/MatrixGame2.0/finetune_i2v.slurm deleted file mode 100644 index 96ebfe555..000000000 --- a/examples/training/finetune/MatrixGame2.0/finetune_i2v.slurm +++ /dev/null @@ -1,131 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=mg-solaris -#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_i2v_output/i2v_%j.out -#SBATCH --error=finetune_i2v_output/i2v_%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/Solaris-6K" -MODEL_PATH="../mg_models/Solaris-30K" -# DATA_DIR="/mnt/weka/home/hao.zhang/kaiqin/solaris/datasets/vpt/vpt/train_81f/preprocessed_0306" -DATA_DIR="/mnt/weka/home/hao.zhang/kaiqin/solaris/datasets/train_81f/preprocessed_0308" -VALIDATION_DATASET_FILE="examples/training/finetune/MatrixGame2.0/validation.json" - -# Training arguments -training_args=( - --tracker_project_name "mg_finetune_solaris" - --output_dir "checkpoints/matrixgame_finetune_${RUN_NAME}" - --wandb_run_name "${RUN_NAME}_bs64_continue" - # --override_transformer_cls_name "MatrixGameWanModel" - --max_train_steps 40000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 2 - --num_latent_t 21 - --num_height 352 - --num_width 640 - --num_frames 81 - # --enable_gradient_checkpointing_type "full" - # --reinit-action-module True - # --override-keyboard-dim 23 -) - -# 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 -) - -# 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 - --multi_phased_distill_schedule "4000-1" - --not_apply_cfg_solver - --dit_precision "fp32" - --num_euler_timesteps 50 - --ema_start_step 0 -) - -mkdir -p finetune_i2v_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_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/MatrixGame2.0/launch_preprocess_slurm.sh b/examples/training/finetune/MatrixGame2.0/launch_preprocess_slurm.sh deleted file mode 100644 index e9ac59439..000000000 --- a/examples/training/finetune/MatrixGame2.0/launch_preprocess_slurm.sh +++ /dev/null @@ -1,20 +0,0 @@ -#!/bin/bash - -# Create output directory if it doesn't exist -mkdir -p preprocess_output - -# Launch 8 jobs, one for each node (Total 64 GPUs) -# Each node processes 8 consecutive files (64 total files / 8 nodes = 8 files per node) -for node_id in {0..3}; do - # Calculate the starting file number for this node - start_file=$((node_id * 8)) - - echo "Launching node $node_id with files merge_${start_file}.txt to merge_$((start_file + 7)).txt" - - sbatch --job-name=mg-vpt-pre-${node_id} \ - --output=preprocess_output/mg-node-${node_id}.out \ - --error=preprocess_output/mg-node-${node_id}.err \ - $(pwd)/examples/training/finetune/MatrixGame2.0/preprocess_worker.slurm $start_file $node_id -done - -echo "All 4 nodes (32 GPUs) launched successfully!" diff --git a/examples/training/finetune/MatrixGame2.0/preprocess_worker.slurm b/examples/training/finetune/MatrixGame2.0/preprocess_worker.slurm deleted file mode 100644 index 991197752..000000000 --- a/examples/training/finetune/MatrixGame2.0/preprocess_worker.slurm +++ /dev/null @@ -1,60 +0,0 @@ -#!/bin/bash -#SBATCH --partition=main -#SBATCH --nodes=1 -#SBATCH --ntasks-per-node=8 -#SBATCH --gres=gpu:8 -#SBATCH --cpus-per-task=16 -#SBATCH --mem=960G -#SBATCH --exclusive -#SBATCH --time=72:00:00 - -# conda init -source ~/conda/miniconda/bin/activate /mnt/weka/home/hao.zhang/conda/miniconda/envs/alexfv -export PYTHONPATH="/mnt/weka/home/hao.zhang/kaiqin/FastVideo_clean:$PYTHONPATH" - -# Accept parameters from launch script -START_FILE=${1:-1} # Starting file number for this node -NODE_ID=${2:-0} # Node identifier (0-7) - -MODEL_PATH="../mg_models/Matrix-Game-2.0-Base-Diffusers" -OUTPUT_BASE="../solaris/datasets/train_81f/preprocessed_0308" - -# Port range calculation -base_port=$((29500 + NODE_ID * 100)) -gpu_ids=(0 1 2 3 4 5 6 7) - -for i in {1..8}; do - port=$((base_port + i)) - gpu=${gpu_ids[((i-1))]} - file_num=$((START_FILE + i - 1)) - - DATA_MERGE_PATH="../solaris/datasets/train_81f/merge_${file_num}.txt" - OUTPUT_DIR="${OUTPUT_BASE}/gpu_${gpu}_file_${file_num}" - echo "DATA_MERGE_PATH: $DATA_MERGE_PATH" - echo "OUTPUT_DIR: $OUTPUT_DIR" - - # CPU binding (optional, kept from syn.slurm logic) - start_cpu=$(( (i-1)*2 )) - end_cpu=$(( start_cpu+1 )) - - echo "Starting GPU $gpu processing file merge_${file_num}.txt on port $port" - - CUDA_VISIBLE_DEVICES=$gpu taskset -c ${start_cpu}-${end_cpu} torchrun --nnodes=1 --nproc_per_node=1 --master_port $port \ - fastvideo/pipelines/preprocess/v1_preprocess.py \ - --model_path $MODEL_PATH \ - --data_merge_path $DATA_MERGE_PATH \ - --preprocess_video_batch_size 1 \ - --seed 42 \ - --max_height 352 \ - --max_width 640 \ - --num_frames 81 \ - --dataloader_num_workers 0 \ - --output_dir=$OUTPUT_DIR \ - --samples_per_file 8 \ - --train_fps 20 \ - --flush_frequency 8 \ - --preprocess_task matrixgame & -done - -wait -echo "Node $NODE_ID processing blocks completed!" diff --git a/fastvideo/training/matrixgame_ar_diffusion_pipeline.py b/fastvideo/training/matrixgame_ar_diffusion_pipeline.py index 3bc321b78..8b642ca4d 100644 --- a/fastvideo/training/matrixgame_ar_diffusion_pipeline.py +++ b/fastvideo/training/matrixgame_ar_diffusion_pipeline.py @@ -550,19 +550,6 @@ class MatrixGameARDiffusionPipeline(TrainingPipeline): return batch - def _post_process_validation_frames( - self, frames: list[np.ndarray], - batch: ForwardBatch) -> list[np.ndarray]: - from fastvideo.models.dits.matrixgame.utils import ( - overlay_validation_actions_on_frames, - ) - - return overlay_validation_actions_on_frames( - frames, - keyboard_cond=getattr(batch, "keyboard_cond", None), - mouse_cond=getattr(batch, "mouse_cond", None), - ) - def main(args) -> None: logger.info("Starting MatrixGame AR diffusion training pipeline...") diff --git a/fastvideo/training/matrixgame_ode_causal_pipeline.py b/fastvideo/training/matrixgame_ode_causal_pipeline.py index a4b3fd37f..71b235b77 100644 --- a/fastvideo/training/matrixgame_ode_causal_pipeline.py +++ b/fastvideo/training/matrixgame_ode_causal_pipeline.py @@ -515,56 +515,6 @@ class MatrixGameODEInitTrainingPipeline(TrainingPipeline): return batch - def _post_process_validation_frames( - self, frames: list[np.ndarray], - batch: ForwardBatch) -> list[np.ndarray]: - """Apply action overlay to validation frames for WanGame. - - Draws keyboard (WASD) and mouse (pitch/yaw) indicators on the video frames. - """ - # Check if action data is available - keyboard_cond = getattr(batch, 'keyboard_cond', None) - mouse_cond = getattr(batch, 'mouse_cond', None) - - if keyboard_cond is None and mouse_cond is None: - return frames - - # Import overlay functions - from fastvideo.models.dits.matrixgame.utils import (draw_keys_on_frame, - draw_mouse_on_frame) - - # Convert tensors to numpy if needed (bfloat16 -> float32 -> numpy) - if keyboard_cond is not None: - keyboard_cond = keyboard_cond.squeeze( - 0).cpu().float().numpy() # (T, 6) - if mouse_cond is not None: - mouse_cond = mouse_cond.squeeze(0).cpu().float().numpy() # (T, 2) - - # MatrixGame convention: keyboard [W, S, A, D, left, right], mouse [Pitch, Yaw] - key_names = ["W", "S", "A", "D", "left", "right"] - - processed_frames = [] - for frame_idx, frame in enumerate(frames): - frame = np.ascontiguousarray(frame.copy()) - - # Draw keyboard overlay - if keyboard_cond is not None and frame_idx < len(keyboard_cond): - keys = { - key_names[i]: bool(keyboard_cond[frame_idx, i]) - for i in range(min(len(key_names), keyboard_cond.shape[1])) - } - draw_keys_on_frame(frame, keys, mode='universal') - - # Draw mouse overlay - if mouse_cond is not None and frame_idx < len(mouse_cond): - pitch = float(mouse_cond[frame_idx, 0]) - yaw = float(mouse_cond[frame_idx, 1]) - draw_mouse_on_frame(frame, pitch, yaw) - - processed_frames.append(frame) - - return processed_frames - def visualize_intermediate_latents(self, training_batch: TrainingBatch, training_args: TrainingArgs, step: int): tracker_loss_dict: dict[str, Any] = {}