Skip to content

Fixed memory leak in loss backward with AC & zero3 & single-model-multi-branch network - #8482

Open
supermeng wants to merge 2 commits into
deepspeedai:masterfrom
supermeng:fix/zero3_ac_mul_branck_mem_leak
Open

supermeng wants to merge 2 commits into
deepspeedai:masterfrom
supermeng:fix/zero3_ac_mul_branck_mem_leak

Conversation

@supermeng

Copy link
Copy Markdown

To fix #8481

image

…kpoint & zero3 & single-model-multi-branch network

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 96e8807dad

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

(output_in, ) = inputs
ctx.module = module
if output_in.requires_grad:
if output_in.requires_grad and torch._C._current_graph_task_id() == -1: # hooks only in really forward phase

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Preserve the callback for reentrant checkpoint recomputation

When use_reentrant=True and the checkpoint input requires gradients, the initial no-grad pass increments this counter with graph task ID -1, but the recomputation runs inside backward and fails this new condition. Its ctx therefore never receives post_backward_function; nevertheless, PostBackwardFunctionModule.backward decrements the earlier counter to zero and accesses that missing attribute, aborting the ZeRO-3 backward. Assign the callback independently of the guarded increment and cover the reentrant training loop on an actual device.

AGENTS.md reference: AGENTS.md:L35-L36

Useful? React with 👍 / 👎.

Comment on lines +515 to +516
if "ds_grads_remaining" not in module.__dict__:
module.ds_grads_remaining = 0

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Add the required Signed-off-by trailer

This is a non-merge commit, but its commit message has no Signed-off-by trailer, so it violates the repository's mandatory commit requirement and may fail DCO enforcement. Recreate the commit with --signoff using the configured Git identity.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.

@supermeng
supermeng force-pushed the fix/zero3_ac_mul_branck_mem_leak branch from 96e8807 to 9b4ffb6 Compare September 11, 2026 10:51

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant