HY / worldcompass /scripts /train_worldcompass.sh

Upload batch 66: 8 files (0.00 GiB)

469a471 verified about 1 month ago

5.94 kB

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

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