File size: 5,940 Bytes
469a471 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | #!/bin/bash
source ~/.bashrc
export PYTHONUNBUFFERED=1
if [ $# -ge 1 ]; then
RANK=$1
else
RANK=0
fi
if [ $# -ge 2 ]; then
NODES=$2
else
NODES=1
fi
NUM_GPUS=8 # Number of GPUs per node
NODES=$NODES # Number of nodes
export MASTER_ADDR="" # [TODO] Set master address
export WORLD_SIZE=$NODES
export RANK=$RANK # Node rank, ranged from 0 to N-1
export MASTER_PORT=27858
export NUM_NODES=$WORLD_SIZE
CACHE_DIR="" # [TODO] Path to the overall checkpoint cache directory
HUNYUAN_CHECKPOINT="" # [TODO] Path to the HunYuan checkpoint directory
WORLDPLAY_CHECKPOINT="" # [TODO] Path to the WorldPlay checkpoint directory
TRAIN_LATENTS_DIR="" # [TODO] Path to the train latents directory
EVAL_LATENTS_DIR="" # [TODO] Path to the eval latents directory
POSE_PATH="" # [TODO] Path to the custom action pose json
OUTPUT_DIR="" # [TODO] Path to the output directory
exp_name="WorldCompass"
# Optimizer arguments
optimizer_args=(
--learning_rate 1e-5
--mixed_precision "bf16"
--checkpointing_steps 6
--weight_decay 1e-4
--max_grad_norm 2.0
)
# Model arguments
model_args=(
--cls_name "HunyuanTransformer3DARActionModel"
--load_from_dir "${HUNYUAN_CHECKPOINT}/transformer/480p_i2v/"
--ar_action_load_from_dir "${WORLDPLAY_CHECKPOINT}/ar_model/diffusion_pytorch_model.safetensors"
--eval_only False
)
# Training arguments
training_args=(
# Path related arguments
--json_path "${TRAIN_LATENTS_DIR}/latents.json" # Path to train latents json file
--eval_json_path "${EVAL_LATENTS_DIR}/latents.json" # Path to eval latents json file
--random_pose_path "${POSE_PATH}" # Path to random pose json file
--vae_path "${HUNYUAN_CHECKPOINT}/vae/" # Path to VAE model
--cache_dir "${CACHE_DIR}" # Path to cache directory
--output_dir "${OUTPUT_DIR}/${exp_name}" # Path to save checkpoints and logs
--generated_videos_dir "${OUTPUT_DIR}/generated_videos/${exp_name}" # Path to save generated videos
# Reward model argument
--camera_estimator "dav3" # Camera estimator: "dav3" or "worldmirror"
# wandb arguments
--wandb_key "" # [TODO] Set wandb key
--wandb_entity "" # [TODO] Set wandb entity
--tracker_project_name "worldcompass"
# Basic training parameters
--max_train_steps 100 # Number of training iterations
--train_batch_size 1 # Training batch size
--train_sp_batch_size 1 # Training special batch size
--window_frames 64 # Video window length for training (longest training video)
--single_chunk_size 4 # Length of single chunk for training
# Training
--gradient_accumulation_steps 2
--enable_gradient_checkpointing_type "full"
# GRPO sampling and training parameters (for reinforcement learning)
--sampling_steps 40 # Number of rollout steps
--grpo_generation_num 12 # Number of rollout samples generated per input
--train_timestep_fraction 0.5 # Fraction of total timesteps used for training
--bestofn 6 # Select best n samples from all rollouts for training
--sampling_batch_size 1 # Batch size for rollout sampling
# Sampling and training chunk selection strategy
--chunk_selection_strategy "min2max" # Chunk selection strategy: "min2max" or "max2min"
--min_chunk_id 1 # Minimum chunk ID (usually 1)
--max_chunk_id 16 # Maximum chunk ID (usually window_frames // single_chunk_size)
# Reward arguments
--action_reward_weight 2.0
--hpsv3_reward_weight 0.0
--hpsv3_quality_reward_weight 0.0
--hpsv3_quality_drift_reward_weight 1.0
--action_reward_type 'fine_action'
--adv_clip_max 2.0
--std_type "global"
# EMA arguments
--ema_min_decay 0.2
--ema_max_decay 0.9
--ema_step_decay 0.01
--ema_ckpt_decay 0.9
# Model task type arguments
--causal
--action
--i2v_rate 1.0
)
# Parallelism arguments
parallel_args=(
--num_gpus $(( NUM_GPUS * NODES )) # Total GPU count used (NUM_GPUS per node * NODES)
--sp_size 1
--tp_size 1
--gpu_para 1
--hsdp_replicate_dim $NODES
--hsdp_shard_dim $NUM_GPUS
)
# Dataset arguments
dataset_args=(
--neg_prompt_path '' # Not used
--neg_byt5_prompt_path '' # Not used
--data_path '' # Not used
--model_path '' # Not used
--pretrained_model_name_or_path '' # Not used
--dataloader_num_workers 4
)
# Miscellaneous arguments
miscellaneous_args=(
--training_cfg_rate 0.0
)
echo "Begin training!"
python -u -m torch.distributed.run \
--master_addr=$MASTER_ADDR \
--master_port=$MASTER_PORT \
--nproc_per_node=$NUM_GPUS \
--nnodes=$NUM_NODES \
--node_rank=$RANK \
fastvideo/training/world_compass_train_pipeline.py \
"${parallel_args[@]}" \
"${model_args[@]}" \
"${dataset_args[@]}" \
"${training_args[@]}" \
"${optimizer_args[@]}" \
"${miscellaneous_args[@]}"
|