deepspeedai / deepspeedai/DeepSpeed
[BUG] Memory leak in loss backward when enabled torch activation checkpoint & zero3 & single-model-multi-branch network
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 43.1k
- Forks
- 5k
- Avg merge
- 4d 15h
- Merged PRs (30d)
- 112
Description
Describe the bug
Huge memory leak in loss backward when enabled torch activation checkpoint & zero3 & single-model-multi-branch network.
To Reproduce
Steps to reproduce the behavior:
Training any single-model-multi-branch network with torch activation checkpoint & zero3 optimizer, e.g training diffsynth's animate-v2-lora model
- git clone https://github.com/modelscope/diffsynth-studio
- update
examples/wanvideo/model_training/lora/Wan2.2-Animate-2-14B.shadd--config_file examples/wanvideo/model_training/full/accelerate_config_zero3.yamlto the launch commad - running: sh examples/wanvideo/model_training/lora/Wan2.2-Animate-2-14B.sh
- and the GPU Memory will increase quickly during the loss backward
Expected behavior
Should release the zero3 all-gathered parameters during loss backward even AC is enabled and recompute happened
Screenshots
System info (please complete the following information):
- OS: [Ubuntu 22.04]
- GPU count and types [single machines with x8 32G RTX5090]
- Python version: 3.10.14
Launcher context
accelerate launch with deepspeed zero3 optimizer
Docker context
No Docker context
Additional context
After some debug and text, I got the memory leak is caused by the wrong ds_grads_remaining reset and count
in deepspeed/runtime/zero/parameter_offload.py's PostBackwardFunctionModule hook,
In AC & zero3 & single-model-multi-branch network env, the PostBackwardFunctionModule would be registered double times but the ds_grads_remaining is reset by _run_after_backward_function in every hook register so post_backward_function only run once, but single-model-multi-branch network should run multiple times depends on the number of branches, so the params gathered by pre-hook would never be released.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with deepspeed/runtime/zero/parameter_offload.py, focusing on the PostBackwardFunctionModule hook, _run_after_backward_function, and ds_grads_remaining. Reproduce the issue with the diffsynth-studio Wan2.2-Animate-2-14B.sh command and the stated zero3 configuration. Done means gathered parameters are released during loss backward across all branches when activation checkpointing is enabled, without continued GPU memory growth.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- distributed-systems, machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100