add sh

Author: susanbao

run_wan_on_vm.sh

@@ -0,0 +1,68 @@
+export PYTHONPATH="/home/sanbao_google_com/maxdiffusion/src:$PYTHONPATH"
+
+RUN_NAME=sanbao-v5p-test-${RANDOM}
+OUTPUT_DIR=gs://sanbao-bucket/wan/sanbao-v5p-test
+DATASET_DIR=gs://sanbao-bucket/wan_tfr_dataset_pusa_v1/train/
+EVAL_DATA_DIR=gs://sanbao-bucket/wan_tfr_dataset_pusa_v1/eval_timesteps/
+SAVE_DATASET_DIR=gs://sanbao-bucket/wan_tfr_dataset_pusa_v1/save/
+RANDOM=123456789
+CKPT_PATH=gs://sanbao-bucket/wan_ckp
+
+export LIBTPU_INIT_ARGS='--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true \
+--xla_tpu_megacore_fusion_allow_ags=false \
+--xla_enable_async_collective_permute=true \
+--xla_tpu_enable_ag_backward_pipelining=true \
+--xla_tpu_enable_data_parallel_all_reduce_opt=true \
+--xla_tpu_data_parallel_opt_different_sized_ops=true \
+--xla_tpu_enable_async_collective_fusion=true \
+--xla_tpu_enable_async_collective_fusion_multiple_steps=true \
+--xla_tpu_overlap_compute_collective_tc=true \
+--xla_enable_async_all_gather=true \
+--xla_tpu_scoped_vmem_limit_kib=65536 \
+--xla_tpu_enable_async_all_to_all=true \
+--xla_tpu_enable_all_experimental_scheduler_features=true \
+--xla_tpu_enable_scheduler_memory_pressure_tracking=true \
+--xla_tpu_host_transfer_overlap_limit=24 \
+--xla_tpu_aggressive_opt_barrier_removal=ENABLED \
+--xla_lhs_prioritize_async_depth_over_stall=ENABLED \
+--xla_should_allow_loop_variant_parameter_in_chain=ENABLED \
+--xla_should_add_loop_invariant_op_in_chain=ENABLED \
+--xla_max_concurrent_host_send_recv=100 \
+--xla_tpu_scheduler_percent_shared_memory_limit=100 \
+--xla_latency_hiding_scheduler_rerun=2 \
+--xla_tpu_use_minor_sharding_for_major_trivial_input=true \
+--xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 \
+--xla_tpu_assign_all_reduce_scatter_layout=true'
+
+HF_HUB_CACHE=/dev/shm/ python3 -m src.maxdiffusion.train_wan \
+src/maxdiffusion/configs/base_wan_14b.yml \
+attention='flash' \
+weights_dtype=bfloat16 \
+activations_dtype=bfloat16 \
+guidance_scale=5.0 \
+flow_shift=5.0 \
+fps=16 \
+skip_jax_distributed_system=False \
+run_name=${RUN_NAME} \
+output_dir=${OUTPUT_DIR} \
+train_data_dir=${DATASET_DIR} \
+load_tfrecord_cached=True \
+height=1280 \
+width=720 \
+num_frames=81 \
+num_inference_steps=50 \
+jax_cache_dir=${OUTPUT_DIR}/jax_cache/ \
+max_train_steps=20 \
+enable_profiler=True \
+dataset_save_location=${SAVE_DATASET_DIR} \
+remat_policy='HIDDEN_STATE_WITH_OFFLOAD' \
+flash_min_seq_length=0 \
+seed=$RANDOM \
+skip_first_n_steps_for_profiler=3 \
+profiler_steps=3 \
+per_device_batch_size=0.5 \
+ici_data_parallelism=2 \
+ici_fsdp_parallelism=2 \
+ici_tensor_parallelism=1 \
+enable_ssim=False \
+flash_block_sizes='{\"block_q\":2048,\"block_kv_compute\":512,\"block_kv\":2048,\"block_q_dkv\":2048,\"block_kv_dkv\":2048,\"block_kv_dkv_compute\":512,\"use_fused_bwd_kernel\":true}'