wehub-resource-sync
217 lines
9.4 KiB
Bash
217 lines
9.4 KiB
Bash
#!/bin/bash
|
|
|
|
# ==============================Manual launch mode==============================
|
|
MANUAL_START=${MANUAL_START:-True} # True: manual launch; False: automatic launch
|
|
|
|
if [ "$MANUAL_START" = "True" ]; then
|
|
# for one machine
|
|
export ARNOLD_WORKER_NUM=${ARNOLD_WORKER_NUM:-1} # number of machines
|
|
if [ -z "${ARNOLD_WORKER_GPU:-}" ]; then
|
|
if command -v nvidia-smi >/dev/null 2>&1; then
|
|
ARNOLD_WORKER_GPU=$(nvidia-smi -L | wc -l)
|
|
else
|
|
ARNOLD_WORKER_GPU=1
|
|
fi
|
|
export ARNOLD_WORKER_GPU
|
|
fi
|
|
export ARNOLD_ID=${ARNOLD_ID:-0} # machine ID
|
|
MAIN_PROCESS_IP=127.0.0.1
|
|
MAIN_PROCESS_PORT=${ARNOLD_WORKER_0_PORT:-6668}
|
|
MACHINE_RANK=$ARNOLD_ID
|
|
|
|
# # for multiple machines
|
|
# export ARNOLD_WORKER_NUM=1 # 2 # number of machines
|
|
# OFFSET=0
|
|
# eval "MAIN_PROCESS_IP=\$ARNOLD_WORKER_${OFFSET}_HOST"
|
|
# eval "MAIN_PROCESS_PORT=\$(echo '$ARNOLD_WORKER_0_PORT' | cut -d "," -f 1)"
|
|
# MACHINE_RANK=$((ARNOLD_ID - OFFSET))
|
|
else
|
|
MACHINE_RANK=$ARNOLD_ID
|
|
fi
|
|
|
|
# ==============================Distributed configuration==============================
|
|
export NCCL_DEBUG=WARN # Log warnings and errors only
|
|
export PYTHONWARNINGS="${PYTHONWARNINGS:+$PYTHONWARNINGS,}ignore::FutureWarning:deepspeed.runtime.zero.linear"
|
|
TOTAL_RANK=$((ARNOLD_WORKER_NUM * ARNOLD_WORKER_GPU))
|
|
NUM_SHARD=${NUM_SHARD:-$((TOTAL_RANK < 4 ? TOTAL_RANK : 4))}
|
|
if [ "$TOTAL_RANK" -lt 1 ]; then
|
|
echo "TOTAL_RANK must be >= 1, got $TOTAL_RANK"
|
|
exit 1
|
|
fi
|
|
if [ "$NUM_SHARD" -lt 1 ] || [ $((TOTAL_RANK % NUM_SHARD)) -ne 0 ]; then
|
|
echo "NUM_SHARD must divide TOTAL_RANK. Got TOTAL_RANK=$TOTAL_RANK, NUM_SHARD=$NUM_SHARD"
|
|
exit 1
|
|
fi
|
|
NUM_REPLICATE=$((TOTAL_RANK / NUM_SHARD))
|
|
|
|
echo "ARNOLD_WORKER_NUM: $ARNOLD_WORKER_NUM"
|
|
echo "ARNOLD_WORKER_GPU: $ARNOLD_WORKER_GPU"
|
|
echo "TOTAL_RANK: $TOTAL_RANK"
|
|
echo "MACHINE_RANK: $MACHINE_RANK"
|
|
echo "MAIN_PROCESS_IP: $MAIN_PROCESS_IP"
|
|
echo "MAIN_PROCESS_PORT: $MAIN_PROCESS_PORT"
|
|
echo "NUM_REPLICATE: $NUM_REPLICATE"
|
|
echo "NUM_SHARD: $NUM_SHARD"
|
|
|
|
# ==============================Model and initialization configuration==============================
|
|
VIT_PATH=./downloads/Qwen2.5-VL-ViT
|
|
MODEL_PATH=${MODEL_PATH:-./downloads/Lance_3B_Video} # or use ./downloads/Lance_3B for image-related task
|
|
VIT_TYPE="qwen_2_5_vl_original"
|
|
|
|
LOAD_FROM_LANCE_CHECKPOINT=True
|
|
INIT_FROM_VLM_CHECKPOINT=False
|
|
COPY_INIT_MOE=True
|
|
|
|
LLM_QK_NORM=True # Use QK_NORM in LLM_UND and LLM_GEN, only when both LLM_QK_NORM_UND and LLM_QK_NORM_GEN are False
|
|
LLM_QK_NORM_UND=True # Use QK_NORM in LLM_UND
|
|
LLM_QK_NORM_GEN=True # Use QK_NORM in LLM_GEN
|
|
TIE_WORD_EMBEDDINGS=False
|
|
|
|
LATENT_PATCH_SIZE="1 1 1"
|
|
MAX_LATENT_SIZE=64
|
|
MAX_NUM_FRAMES=121
|
|
VAE_MODEL_TYPE="wan"
|
|
VISUAL_UND=True
|
|
VISUAL_GEN=False # Set False for understanding-only tasks
|
|
|
|
# ==============================Freeze configuration==============================
|
|
FREEZE_LLM=False # Freeze all LLM_UND and LLM_GEN parameters
|
|
FREEZE_LLM_EMBED_TOKENS=True # Freeze the LLM embedding layer
|
|
FREEZE_UND_PARAMS=False # Freeze all LLM_UND parameters
|
|
FREEZE_UND=False # Freeze LLM_UND behavior
|
|
|
|
# ==============================Data configuration==============================
|
|
DATASET_CONFIG_FILE=${DATASET_CONFIG_FILE:-config/train_local/multi_und.yaml}
|
|
VAL_DATASET_CONFIG_FILE=${VAL_DATASET_CONFIG_FILE:-config/train_local/multi_und.yaml}
|
|
|
|
NUM_WORKERS=${NUM_WORKERS:-8}
|
|
PREFETCH_FACTOR=${PREFETCH_FACTOR:-2}
|
|
MAX_NUM_TOKENS=65536
|
|
EXPECTED_NUM_TOKENS=65536
|
|
MAX_NUM_TOKENS_PER_SAMPLE=65536
|
|
PREFER_BUFFER_BEFORE=65536
|
|
MAX_BUFFER_SIZE=50
|
|
REQUIRE_UND_GEN=True
|
|
|
|
DATA_SEED=2026
|
|
GLOBAL_SEED=2026
|
|
|
|
TEXT_COND_DROPOUT_PROB=0.1
|
|
VAE_COND_DROPOUT_PROB=0
|
|
VIT_COND_DROPOUT_PROB=0.1
|
|
|
|
# ==============================Optimizer and training configuration==============================
|
|
LR=0.0001
|
|
MIN_LR=0.000001
|
|
WARMUP_STEPS=500
|
|
LR_SCHEDULER="constant" # "cosine" | "constant"
|
|
|
|
CE_WEIGHT=0.25
|
|
MSE_WEIGHT=1.0
|
|
TIMESTEP_SHIFT=4.0
|
|
USE_FLEX=True
|
|
|
|
LOG_EVERY=2
|
|
SAVE_EVERY=10 # 200 # 1000
|
|
CKPT_DEBUG_STEPS=10
|
|
|
|
# ==============================EMA, resume, and output configuration==============================
|
|
USE_EMA=True
|
|
AUTO_RESUME=False
|
|
RESUME_MODEL_ONLY=True
|
|
LOAD_DATA_STATUS=False
|
|
|
|
DATE=v0608
|
|
WANDB_PROJECT=huangmengqi_sft_lance
|
|
WANDB_NAME=${DATE}_"$NUM_SHARD"x"$NUM_REPLICATE"_multi_gen
|
|
WANDB_OFFLINE=True # False
|
|
|
|
OUTPUTS_DIR=./outputs/${WANDB_PROJECT}
|
|
|
|
# ==============================Validation configuration==============================
|
|
VALIDATION_STEP=-1 # 100 | -1 disables validation
|
|
VALIDATION_TYPE="und_gen"
|
|
VALIDATION_NUM_TIMESTEPS=10
|
|
VALIDATION_TIMESTEP_SHIFT=3.0
|
|
VALIDATION_MAX_SAMPLES=$((128 / TOTAL_RANK))
|
|
VALIDATION_LOG_TYPE="direct"
|
|
echo "VALIDATION_MAX_SAMPLES: $VALIDATION_MAX_SAMPLES"
|
|
|
|
# ==============================Runtime switches==============================
|
|
APPLY_QWEN_2_5_VL_POS_EMB=True
|
|
APPLY_CHAT_TEMPLATE=False
|
|
|
|
# ==============================Launch training==============================
|
|
accelerate launch \
|
|
--num_machines $ARNOLD_WORKER_NUM \
|
|
--num_processes $TOTAL_RANK \
|
|
--machine_rank $MACHINE_RANK \
|
|
--main_process_ip $MAIN_PROCESS_IP \
|
|
--main_process_port $MAIN_PROCESS_PORT \
|
|
--mixed_precision bf16 \
|
|
train/unified_train.py \
|
|
--model_path $MODEL_PATH \
|
|
--vit_path $VIT_PATH \
|
|
--vit_type $VIT_TYPE \
|
|
--llm_qk_norm $LLM_QK_NORM \
|
|
--llm_qk_norm_und $LLM_QK_NORM_UND \
|
|
--llm_qk_norm_gen $LLM_QK_NORM_GEN \
|
|
--tie_word_embeddings $TIE_WORD_EMBEDDINGS \
|
|
--max_num_frames $MAX_NUM_FRAMES \
|
|
--max_latent_size $MAX_LATENT_SIZE \
|
|
--latent_patch_size $LATENT_PATCH_SIZE \
|
|
--load_from_lance_checkpoint $LOAD_FROM_LANCE_CHECKPOINT \
|
|
--init_from_vlm_checkpoint $INIT_FROM_VLM_CHECKPOINT \
|
|
--copy_init_moe $COPY_INIT_MOE \
|
|
--visual_und $VISUAL_UND \
|
|
--visual_gen $VISUAL_GEN \
|
|
--vae_model_type $VAE_MODEL_TYPE \
|
|
--freeze_llm $FREEZE_LLM \
|
|
--freeze_llm_embed_tokens $FREEZE_LLM_EMBED_TOKENS \
|
|
--freeze_und_params $FREEZE_UND_PARAMS \
|
|
--freeze_und $FREEZE_UND \
|
|
--dataset_config_file $DATASET_CONFIG_FILE \
|
|
--val_dataset_config_file $VAL_DATASET_CONFIG_FILE \
|
|
--num_workers $NUM_WORKERS \
|
|
--prefetch_factor $PREFETCH_FACTOR \
|
|
--expected_num_tokens $EXPECTED_NUM_TOKENS \
|
|
--max_num_tokens $MAX_NUM_TOKENS \
|
|
--max_num_tokens_per_sample $MAX_NUM_TOKENS_PER_SAMPLE \
|
|
--prefer_buffer_before $PREFER_BUFFER_BEFORE \
|
|
--max_buffer_size $MAX_BUFFER_SIZE \
|
|
--require_und_gen $REQUIRE_UND_GEN \
|
|
--data_seed $DATA_SEED \
|
|
--global_seed $GLOBAL_SEED \
|
|
--text_cond_dropout_prob $TEXT_COND_DROPOUT_PROB \
|
|
--vae_cond_dropout_prob $VAE_COND_DROPOUT_PROB \
|
|
--vit_cond_dropout_prob $VIT_COND_DROPOUT_PROB \
|
|
--warmup_steps $WARMUP_STEPS \
|
|
--lr_scheduler $LR_SCHEDULER \
|
|
--lr $LR \
|
|
--min_lr $MIN_LR \
|
|
--ce_weight $CE_WEIGHT \
|
|
--mse_weight $MSE_WEIGHT \
|
|
--timestep_shift $TIMESTEP_SHIFT \
|
|
--use_flex $USE_FLEX \
|
|
--num_replicate $NUM_REPLICATE \
|
|
--num_shard $NUM_SHARD \
|
|
--use_ema $USE_EMA \
|
|
--auto_resume $AUTO_RESUME \
|
|
--resume_model_only $RESUME_MODEL_ONLY \
|
|
--load_data_status $LOAD_DATA_STATUS \
|
|
--outputs_dir $OUTPUTS_DIR \
|
|
--save_every $SAVE_EVERY \
|
|
--log_every $LOG_EVERY \
|
|
--ckpt_debug_steps $CKPT_DEBUG_STEPS \
|
|
--wandb_project $WANDB_PROJECT \
|
|
--wandb_name $WANDB_NAME \
|
|
--wandb_offline $WANDB_OFFLINE \
|
|
--validation_step $VALIDATION_STEP \
|
|
--validation_type $VALIDATION_TYPE \
|
|
--validation_num_timesteps $VALIDATION_NUM_TIMESTEPS \
|
|
--validation_timestep_shift $VALIDATION_TIMESTEP_SHIFT \
|
|
--validation_max_samples $VALIDATION_MAX_SAMPLES \
|
|
--validation_log_type $VALIDATION_LOG_TYPE \
|
|
--apply_qwen_2_5_vl_pos_emb $APPLY_QWEN_2_5_VL_POS_EMB \
|
|
--apply_chat_template $APPLY_CHAT_TEMPLATE
|