Buckets:
| set -euo pipefail | |
| SOURCE_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) | |
| PROJECT_ROOT=${SPATIAL_MMCOT_PROJECT_ROOT:-$(cd "$SOURCE_ROOT/../.." && pwd)} | |
| OUTPUT_DIR=${OUTPUT_DIR:?set OUTPUT_DIR to one prepared run directory} | |
| MODEL_PATH=${MODEL_PATH:-$PROJECT_ROOT/models/Bagel-7B-MoT} | |
| DATASET_CONFIG=${DATASET_CONFIG:-$SOURCE_ROOT/data/configs/thinkmorph_paper.yaml} | |
| NPROC_PER_NODE=${NPROC_PER_NODE:-4} | |
| TOTAL_STEPS=${TOTAL_STEPS:-8000} | |
| GRADIENT_ACCUMULATION_STEPS=${GRADIENT_ACCUMULATION_STEPS:-4} | |
| EXPECTED_NUM_TOKENS=${EXPECTED_NUM_TOKENS:-32768} | |
| MAX_NUM_TOKENS=${MAX_NUM_TOKENS:-32768} | |
| MAX_NUM_TOKENS_PER_SAMPLE=${MAX_NUM_TOKENS_PER_SAMPLE:-16384} | |
| NUM_WORKERS=${NUM_WORKERS:-4} | |
| SAVE_EVERY=${SAVE_EVERY:-2000} | |
| BACKWARD_PREFETCH=${BACKWARD_PREFETCH:-BACKWARD_PRE} | |
| EXTRA_ARGS=() | |
| if [[ -n ${STOP_AFTER_STEPS:-} ]]; then | |
| EXTRA_ARGS+=(--stop_after_steps "$STOP_AFTER_STEPS") | |
| fi | |
| mkdir -p "$OUTPUT_DIR/logs" | |
| export WANDB_DIR="$OUTPUT_DIR/logs" | |
| export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" | |
| export PYTHONPATH="$SOURCE_ROOT${PYTHONPATH:+:$PYTHONPATH}" | |
| cd "$SOURCE_ROOT" | |
| exec torchrun --standalone --nproc-per-node="$NPROC_PER_NODE" \ | |
| train/pretrain_unified_navit.py \ | |
| --dataset_config_file "$DATASET_CONFIG" \ | |
| --model_path "$MODEL_PATH" \ | |
| --layer_module Qwen2MoTDecoderLayer \ | |
| --llm_qk_norm True \ | |
| --tie_word_embeddings False \ | |
| --latent_patch_size 2 \ | |
| --vit_patch_size 14 \ | |
| --vit_max_num_patch_per_side 70 \ | |
| --vit_select_layer -2 \ | |
| --vit_rope False \ | |
| --interpolate_pos False \ | |
| --connector_act gelu_pytorch_tanh \ | |
| --finetune_from_hf True \ | |
| --auto_resume True \ | |
| --resume_model_only True \ | |
| --finetune_from_ema True \ | |
| --resume_from "$MODEL_PATH" \ | |
| --results_dir "$OUTPUT_DIR/logs" \ | |
| --checkpoint_dir "$OUTPUT_DIR/checkpoints" \ | |
| --wandb_project thinkmorph-paper-reproduction \ | |
| --wandb_name "$(basename "$OUTPUT_DIR")" \ | |
| --wandb_offline True \ | |
| --log_every 1 \ | |
| --save_every "$SAVE_EVERY" \ | |
| --lr 1e-5 \ | |
| --lr_scheduler cosine \ | |
| --warmup_steps 2000 \ | |
| --min_lr 1e-7 \ | |
| --beta1 0.9 \ | |
| --beta2 0.95 \ | |
| --eps 1e-15 \ | |
| --total_steps "$TOTAL_STEPS" \ | |
| --gradient_accumulation_steps "$GRADIENT_ACCUMULATION_STEPS" \ | |
| --num_workers "$NUM_WORKERS" \ | |
| --prefetch_factor 2 \ | |
| --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 16384 \ | |
| --max_buffer_size 50 \ | |
| --max_latent_size 64 \ | |
| --visual_gen True \ | |
| --visual_und True \ | |
| --freeze_llm False \ | |
| --freeze_vit False \ | |
| --freeze_vae True \ | |
| --freeze_und False \ | |
| --copy_init_moe True \ | |
| --text_cond_dropout_prob 0.1 \ | |
| --vit_cond_dropout_prob 0.3 \ | |
| --vae_cond_dropout_prob 0.3 \ | |
| --mse_weight 1 \ | |
| --ce_weight 1 \ | |
| --max_grad_norm 1 \ | |
| --ema 0.9999 \ | |
| --timestep_shift 1 \ | |
| --global_seed 4396 \ | |
| --data_seed 42 \ | |
| --ce_loss_reweighting False \ | |
| --sharding_strategy FULL_SHARD \ | |
| --backward_prefetch "$BACKWARD_PREFETCH" \ | |
| --cpu_offload False \ | |
| --num_replicate 1 \ | |
| --num_shard 4 \ | |
| --use_flex False \ | |
| "${EXTRA_ARGS[@]}" | |
Xet Storage Details
- Size:
- 3.13 kB
- Xet hash:
- dcdaf360bab4096d3ae58fca6a2d2284a125129b3c216391d76bb6f7e4b894b2
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.