|
| 1 | +#!/bin/bash |
| 2 | +# End-to-end Wan2.2 27B T2V generation using the new ulysses_ring_custom |
| 3 | +# (Ulysses + Ring 2D context parallel) attention kernel. |
| 4 | +set -uo pipefail |
| 5 | +cd /mnt/data_500g/maxdiffusion |
| 6 | + |
| 7 | +export GLOG_minloglevel=3 |
| 8 | +export JAX_CPP_MIN_LOG_LEVEL=3 |
| 9 | +export TPU_LOGS="/mnt/data_500g/sagar_tmp/tpu_logs/" |
| 10 | +mkdir -p "$TPU_LOGS" |
| 11 | + |
| 12 | +# Complete local Wan2.2-T2V-A14B HF cache (weights) on this device. |
| 13 | +export HF_HUB_CACHE="/mnt/data_500g/maxdiffusion_runtime/wan_hf_cache/" |
| 14 | +export HF_HUB_ENABLE_HF_TRANSFER=1 |
| 15 | + |
| 16 | +export JAX_COMPILATION_CACHE_DIR="/mnt/data_500g/sagar_jax_cache_wan22" |
| 17 | +export JAX_PERSISTENT_CACHE_MIN_ENTRY_SIZE_BYTES=-1 |
| 18 | +export JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS=0 |
| 19 | +export XLA_PYTHON_CLIENT_MEM_FRACTION=0.95 |
| 20 | +export JAX_DEFAULT_MATMUL_PRECISION=bfloat16 |
| 21 | + |
| 22 | +export LIBTPU_INIT_ARGS='--xla_tpu_dvfs_p_state=7 \ |
| 23 | +--xla_tpu_spmd_rng_bit_generator_unsafe=true \ |
| 24 | +--xla_tpu_enable_dot_strength_reduction=true \ |
| 25 | +--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true \ |
| 26 | +--xla_enable_async_collective_permute=true \ |
| 27 | +--xla_tpu_enable_data_parallel_all_reduce_opt=true \ |
| 28 | +--xla_tpu_data_parallel_opt_different_sized_ops=true \ |
| 29 | +--xla_tpu_enable_async_collective_fusion=true \ |
| 30 | +--xla_tpu_enable_async_collective_fusion_multiple_steps=true \ |
| 31 | +--xla_tpu_overlap_compute_collective_tc=true \ |
| 32 | +--xla_enable_async_all_gather=true \ |
| 33 | +--xla_tpu_scoped_vmem_limit_kib=8192 \ |
| 34 | +--xla_tpu_enable_async_all_to_all=true \ |
| 35 | +--xla_tpu_enable_all_experimental_scheduler_features=true \ |
| 36 | +--xla_tpu_enable_scheduler_memory_pressure_tracking=true \ |
| 37 | +--xla_tpu_host_transfer_overlap_limit=24 \ |
| 38 | +--xla_max_concurrent_host_send_recv=100 \ |
| 39 | +--xla_tpu_scheduler_percent_shared_memory_limit=100 \ |
| 40 | +--xla_latency_hiding_scheduler_rerun=2 \ |
| 41 | +--xla_tpu_use_minor_sharding_for_major_trivial_input=true \ |
| 42 | +--xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 \ |
| 43 | +--xla_tpu_enable_latency_hiding_scheduler=true \ |
| 44 | +--xla_tpu_memory_bound_loop_optimizer_options=enabled:true \ |
| 45 | +--xla_tpu_use_single_sparse_core_for_all_gather_offload=true \ |
| 46 | +--xla_tpu_sparse_core_all_gather_latency_multiplier=1 \ |
| 47 | +--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \ |
| 48 | +--xla_tpu_enable_sparse_core_collective_aggregator=true \ |
| 49 | +--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \ |
| 50 | +--xla_tpu_enable_sparse_core_reduce_scatter_v2=true \ |
| 51 | +--xla_tpu_enable_sparse_core_collective_offload_all_gather=true \ |
| 52 | +--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \ |
| 53 | +--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \ |
| 54 | +--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \ |
| 55 | +--xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \ |
| 56 | +--xla_tpu_enable_concurrent_sparse_core_offloading=true \ |
| 57 | +--xla_tpu_assign_all_reduce_scatter_layout=true' |
| 58 | + |
| 59 | +export PYTHONPATH=src |
| 60 | + |
| 61 | +.venv/bin/python src/maxdiffusion/generate_wan.py \ |
| 62 | + src/maxdiffusion/configs/base_wan_27b.yml \ |
| 63 | + run_name=wan-ulysses-ring-custom \ |
| 64 | + pretrained_orbax_dir=/mnt/data_500g/sagar_orbax_22_27b_v2 \ |
| 65 | + seed=12345 \ |
| 66 | + attention="ulysses_ring_custom" \ |
| 67 | + ulysses_shards=2 \ |
| 68 | + flash_min_seq_length=0 \ |
| 69 | + num_inference_steps=40 num_frames=81 width=1280 height=720 \ |
| 70 | + per_device_batch_size=.125 \ |
| 71 | + vae_spatial=4 vae_decode_chunk=-1 vae_weights_dtype='bfloat16' vae_dtype='bfloat16' \ |
| 72 | + text_encoder_dtype='bfloat16' compile_text_encoder=true \ |
| 73 | + ici_data_parallelism=2 ici_context_parallelism=4 fps=16 \ |
| 74 | + use_kv_cache=true use_base2_exp=true use_experimental_scheduler=true use_batched_text_encoder=true \ |
| 75 | + flash_block_sizes='{"block_q" : 6400, "block_kv" : 1024, "block_kv_compute" : 1024, "block_kv_compute_in" : 1024, "heads_per_tile" : 1, "vmem_limit_bytes" : 67108864}' |
0 commit comments