Cccccz/HY

1#!/bin/bash2 3source ~/.bashrc4 5export PYTHONUNBUFFERED=16 7if [ $# -ge 1 ]; then8    RANK=$19else10    RANK=011fi12 13if [ $# -ge 2 ]; then14    NODES=$215else16    NODES=117fi18 19NUM_GPUS=8                  # Number of GPUs per node20NODES=$NODES                # Number of nodes21 22export MASTER_ADDR=""       # [TODO] Set master address23export WORLD_SIZE=$NODES24export RANK=$RANK           # Node rank, ranged from 0 to N-125export MASTER_PORT=2785826export NUM_NODES=$WORLD_SIZE27 28CACHE_DIR=""              # [TODO] Path to the overall checkpoint cache directory29HUNYUAN_CHECKPOINT=""     # [TODO] Path to the HunYuan checkpoint directory30WORLDPLAY_CHECKPOINT=""   # [TODO] Path to the WorldPlay checkpoint directory31 32TRAIN_LATENTS_DIR=""      # [TODO] Path to the train latents directory33EVAL_LATENTS_DIR=""       # [TODO] Path to the eval latents directory34POSE_PATH=""              # [TODO] Path to the custom action pose json35OUTPUT_DIR=""             # [TODO] Path to the output directory36 37exp_name="WorldCompass"38 39# Optimizer arguments40optimizer_args=(41  --learning_rate 1e-542  --mixed_precision "bf16"43  --checkpointing_steps 644  --weight_decay 1e-445  --max_grad_norm 2.046)47 48# Model arguments49model_args=(50  --cls_name "HunyuanTransformer3DARActionModel"51  --load_from_dir "${HUNYUAN_CHECKPOINT}/transformer/480p_i2v/"52  --ar_action_load_from_dir "${WORLDPLAY_CHECKPOINT}/ar_model/diffusion_pytorch_model.safetensors"53  --eval_only False54)55 56# Training arguments57training_args=(58  # Path related arguments59  --json_path "${TRAIN_LATENTS_DIR}/latents.json"                            # Path to train latents json file60  --eval_json_path "${EVAL_LATENTS_DIR}/latents.json"                        # Path to eval latents json file61  --random_pose_path "${POSE_PATH}"                                          # Path to random pose json file62  --vae_path "${HUNYUAN_CHECKPOINT}/vae/"                                    # Path to VAE model63  --cache_dir "${CACHE_DIR}"                                                 # Path to cache directory64 65  --output_dir "${OUTPUT_DIR}/${exp_name}"                                   # Path to save checkpoints and logs66  --generated_videos_dir "${OUTPUT_DIR}/generated_videos/${exp_name}"        # Path to save generated videos67 68  # Reward model argument69  --camera_estimator "dav3"                                                  # Camera estimator: "dav3" or "worldmirror"70 71  # wandb arguments72  --wandb_key ""                                                             # [TODO] Set wandb key73  --wandb_entity ""                                                          # [TODO] Set wandb entity74  --tracker_project_name "worldcompass"75 76  # Basic training parameters77  --max_train_steps 100                                                      # Number of training iterations78  --train_batch_size 1                                                       # Training batch size79  --train_sp_batch_size 1                                                    # Training special batch size80  --window_frames 64                                                         # Video window length for training (longest training video)81  --single_chunk_size 4                                                      # Length of single chunk for training82 83  # Training84  --gradient_accumulation_steps 285  --enable_gradient_checkpointing_type "full"86 87  # GRPO sampling and training parameters (for reinforcement learning)88  --sampling_steps 40                                                        # Number of rollout steps89  --grpo_generation_num 12                                                   # Number of rollout samples generated per input90  --train_timestep_fraction 0.5                                              # Fraction of total timesteps used for training91  --bestofn 6                                                                # Select best n samples from all rollouts for training92  --sampling_batch_size 1                                                    # Batch size for rollout sampling93 94  # Sampling and training chunk selection strategy95  --chunk_selection_strategy "min2max"                                       # Chunk selection strategy: "min2max" or "max2min"96  --min_chunk_id 1                                                           # Minimum chunk ID (usually 1)97  --max_chunk_id 16                                                          # Maximum chunk ID (usually window_frames // single_chunk_size)98 99  # Reward arguments100  --action_reward_weight 2.0101  --hpsv3_reward_weight 0.0102  --hpsv3_quality_reward_weight 0.0103  --hpsv3_quality_drift_reward_weight 1.0104  --action_reward_type 'fine_action'105  --adv_clip_max 2.0106  --std_type "global"107 108  # EMA arguments109  --ema_min_decay 0.2110  --ema_max_decay 0.9111  --ema_step_decay 0.01112  --ema_ckpt_decay 0.9113 114  # Model task type arguments115  --causal116  --action117  --i2v_rate 1.0118)119 120# Parallelism arguments121parallel_args=(122  --num_gpus $(( NUM_GPUS * NODES ))    # Total GPU count used (NUM_GPUS per node * NODES)123  --sp_size 1124  --tp_size 1125  --gpu_para 1126  --hsdp_replicate_dim $NODES127  --hsdp_shard_dim $NUM_GPUS128)129 130# Dataset arguments131dataset_args=(132  --neg_prompt_path ''                  # Not used133  --neg_byt5_prompt_path ''             # Not used134  --data_path ''                        # Not used135  --model_path ''                       # Not used136  --pretrained_model_name_or_path ''    # Not used137  --dataloader_num_workers 4138)139 140# Miscellaneous arguments141miscellaneous_args=(142  --training_cfg_rate 0.0143)144 145echo "Begin training!"146python -u -m torch.distributed.run \147    --master_addr=$MASTER_ADDR \148    --master_port=$MASTER_PORT \149    --nproc_per_node=$NUM_GPUS \150    --nnodes=$NUM_NODES \151    --node_rank=$RANK \152    fastvideo/training/world_compass_train_pipeline.py \153    "${parallel_args[@]}" \154    "${model_args[@]}" \155    "${dataset_args[@]}" \156    "${training_args[@]}" \157    "${optimizer_args[@]}" \158    "${miscellaneous_args[@]}"159