deepspeedai / deepspeedai/DeepSpeed

[BUG] Memory leak in loss backward when enabled torch activation checkpoint & zero3 & single-model-multi-branch network

Open
#8,481 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug training
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

  1. git clone https://github.com/modelscope/diffsynth-studio
  2. update examples/wanvideo/model_training/lora/Wan2.2-Animate-2-14B.sh add --config_file examples/wanvideo/model_training/full/accelerate_config_zero3.yaml to the launch commad
  3. running: sh examples/wanvideo/model_training/lora/Wan2.2-Animate-2-14B.sh
  4. 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
Image

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

Open the contributing guide

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 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.