[Question] Huge vram usage on GPU0 when loading models (train WAN2.2 TI2V 5B)

I noticed that when loading models from local .safetensors files, the vram usage on GPU0 will be huge, which will sometimes cause the OOM error on GPU0.

After doing some observation, I found that when training the models in multi_process status, all the processes will store their models on GPU0 and then transfer to the assigned device, it is this which cause the OOM and huge usage on GPU0 when loading models.

How can I avoid this problem?

Model: Wan2.2 TI2V 5B, full fine-tuning training.
GPU: A100(80G) * 4
num_process: 4
python version: 3.10.12
torch version: 2.8.0+cu128

the scripts below:

accelerate launch --config_file ./$ACCELERATE_CONF ./train_local.py \
  --dataset_base_path "./video" \
  --dataset_metadata_path $DATASET_META_PATH \
  --data_file_keys "video" \
  --num_frames 49 \
  --dataset_repeat 1 \
  --model_id_with_origin_paths "$MODEL_ID:$MODEL_PATH_BASE/diffusion_pytorch_model*.safetensors,$MODEL_ID:$MODEL_PATH_BASE/models_t5_umt5-xxl-enc-bf16.pth,$MODEL_ID:$MODEL_PATH_BASE/Wan2.2_VAE.pth" \
  --learning_rate 1e-5 \
  --num_epochs 2 \
  --remove_prefix_in_ckpt "pipe.dit." \
  --output_path "./train_output" \
  --trainable_models "dit" \
  --save_steps 1000 \
  --use_gradient_checkpointing \
  --use_gradient_checkpointing_offload \
  --gradient_accumulation_steps 1

and $ACCELERATE_CONF yaml files:

compute_environment: LOCAL_MACHINE
debug: false
deepspeed_config:
  gradient_accumulation_steps: 1
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: false
  zero_stage: 2
distributed_type: DEEPSPEED
downcast_bf16: 'no'
enable_cpu_affinity: false
machine_rank: 0
main_training_function: main
mixed_precision: bf16
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

Start with train_local.py and the provided Accelerate/DeepSpeed configuration, tracing how local .safetensors models are loaded and assigned across the four processes. Reproduce the multi-process run with the Wan2.2 TI2V 5B command and verify that model loading no longer causes excessive GPU0 VRAM use or an OOM.