modelscope / modelscope/DiffSynth-Studio

8 * 4090 GPUs using LoRA to fine-tune the Wan2.1-T2V-14B model results in OOM (Out of Memory)

Open
#581 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
13.1k
Forks
1.3k
Avg merge
13h 12m
Merged PRs (30d)
45

Description

run:

python train_wan_t2v.py \
  --task data_process \
  --dataset_path data \
  --text_encoder_path "/model/Wan2.1-T2V-14B/models_t5_umt5-xxl-enc-bf16.pth" \
  --vae_path "/model/Wan2.1-T2V-14B/Wan2.1_VAE.pth" \
  --tiled \
  --num_frames 81 \
  --height 864 \
  --width 480
CUDA_VISIBLE_DEVICES="0,1,2,3,4,5,6,7" \
python train_wan_t2v.py \
    --task train \
    --train_architecture lora \
    --dataset_path ./data/video_train \
    --output_path ./lora/ \
    --dit_path "/model/Wan2.1-T2V-14B/diffusion_pytorch_model-00001-of-00006.safetensors,\
/model/Wan2.1-T2V-14B/diffusion_pytorch_model-00002-of-00006.safetensors,\
/model/Wan2.1-T2V-14B/diffusion_pytorch_model-00003-of-00006.safetensors,\
/model/Wan2.1-T2V-14B/diffusion_pytorch_model-00004-of-00006.safetensors,\
/model/Wan2.1-T2V-14B/diffusion_pytorch_model-00005-of-00006.safetensors,\
/model/Wan2.1-T2V-14B/diffusion_pytorch_model-00006-of-00006.safetensors" \
    --steps_per_epoch 500 \
    --max_epochs 10 \
    --learning_rate 1e-4 \
    --lora_rank 4 \
    --lora_alpha 4 \
    --lora_target_modules "q" \
    --accumulate_grad_batches 1 \
    --use_gradient_checkpointing_offload \
    --use_gradient_checkpointing \
    --training_strategy "deepspeed_stage_3"

log:

[INFO] [real_accelerator.py:239:get_accelerator] Setting ds_accelerator to cuda (auto detect)
initializing deepspeed distributed: GLOBAL_RANK: 1, MEMBER: 2/8
initializing deepspeed distributed: GLOBAL_RANK: 4, MEMBER: 5/8
initializing deepspeed distributed: GLOBAL_RANK: 2, MEMBER: 3/8
initializing deepspeed distributed: GLOBAL_RANK: 7, MEMBER: 8/8
initializing deepspeed distributed: GLOBAL_RANK: 6, MEMBER: 7/8
initializing deepspeed distributed: GLOBAL_RANK: 5, MEMBER: 6/8
LOCAL_RANK: 1 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
LOCAL_RANK: 6 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
LOCAL_RANK: 4 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
LOCAL_RANK: 2 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
Enabling DeepSpeed BF16. Model parameters and inputs will be cast to `bfloat16`.
LOCAL_RANK: 5 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
LOCAL_RANK: 7 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
LOCAL_RANK: 3 - CUDA_VISIBLE_DEVICES: [0,1,2,3,4,5,6,7]
[rank1]: Traceback (most recent call last):
...
[rank1]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 50.00 MiB. GPU 1 has a total capacity of 23.53 GiB of which 30.50 MiB is free. Including non-PyTorch memory, this process has 23.49 GiB memory in use. Of the allocated memory 22.90 GiB is allocated by PyTorch, and 19.48 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
....
[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 50.00 MiB. GPU 0 has a total capacity of 23.53 GiB of which 30.50 MiB is free. Including non-PyTorch memory, this process has 23.49 GiB memory in use. Of the allocated memory 22.90 GiB is allocated by PyTorch, and 19.48 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
...

I've tried adjusting the following parameters, always resulted in OOM:

CUDA_VISIBLE_DEVICES to only "0,1,2,3"
--lora_rank 8(16) --lora_alpha 8(16)
--accumulate_grad_batches 8(16)
--training_strategy "deepspeed_stage_3"(deepspeed_stage_2/deepspeed_stage_1)

Is it that the 8*4090 GPUs cannot handle fine-tuning the Wan2.1-T2V-14B model?

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with train_wan_t2v.py and the supplied training command, then inspect the logged per-GPU allocation under the listed DeepSpeed and gradient-checkpointing settings. Done means establishing whether the configuration can run on 24-GB RTX 4090 cards and recording a confirmed supported configuration or explicit hardware limitation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
ai, machine-learning
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.