FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Add keyword argument support to activation checkpointing by zupengwang · Pull Request #8182 · deepspeedai/DeepSpeed · GitHub

Add keyword argument support to activation checkpointing - #8182

Merged
tohtana merged 1 commit into
deepspeedai:masterfrom
zupengwang:enhancement/checkpoint-kwargs
Aug 2, 2026
Merged

Add keyword argument support to activation checkpointing#8182
tohtana merged 1 commit into
deepspeedai:masterfrom
zupengwang:enhancement/checkpoint-kwargs

Conversation

Copy link
Copy Markdown
Contributor

Summary

DeepSpeed's activation checkpointing wrapper currently accepts positional
arguments only, while torch.utils.checkpoint.checkpoint also supports keyword
arguments. This change allows callers to pass keyword arguments through
deepspeed.checkpointing.checkpoint.

Keyword names and non-Tensor values are retained for reconstruction during the
forward and recompute passes. Tensor keyword values are flattened into the
inputs passed to CheckpointFunction, so autograd tracks them and returns their
gradients correctly.

The activation checkpointing documentation now describes keyword argument
support, and the regression test covers both Tensor and non-Tensor keyword
arguments as well as gradient propagation.

Validation

  • Activation checkpointing unit tests: 27 passed
  • Pre-commit checks for all changed files: passed
  • 1-GPU CUDA correctness smoke: direct and checkpointed execution matched, with
    zero maximum gradient error
  • 2-GPU distributed CUDA correctness smoke: each rank matched direct execution,
    with zero maximum gradient error
  • DCO sign-off is included in the commit

Fixes #7038

Copy link
Copy Markdown
Collaborator

@zupengwang thanks for this PR. Looking forward to when it is ready for review.

zupengwang marked this pull request as ready for review August 2, 2026 02:18

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.

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 47b2a888a6

ℹ️ About Codex in GitHub

Your team has set up Codex to 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 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".


@compiler.disable # WA from Pytorch repo for compile + zero 3 accuracy issue
def checkpoint(function, *args):
def checkpoint(function, *args, **kwargs):

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

Add the required commit sign-off

The reviewed commit's raw message has no Signed-off-by: trailer, despite claiming that the DCO sign-off is included. This violates the repository's mandatory commit requirement and can prevent the change from passing DCO/CI validation; recreate the commit with --signoff.

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

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Choose a reason Spam Abuse Off Topic Outdated Duplicate Resolved Low Quality

Thanks for flagging this. The reviewed commit 47b2a88 includes the Signed-off-by: Wang Zupeng zupenwang@gmail.com trailer in its raw commit message, and its DCO check passed. The current rebased head 8d56205 retains the same trailer and also passes DCO, so no additional commit rewrite is needed.

Signed-off-by: Wang Zupeng <zupenwang@gmail.com>
zupengwang force-pushed the enhancement/checkpoint-kwargs branch from 47b2a88 to 8d56205 Compare August 2, 2026 02:23

Copy link
Copy Markdown
Contributor Author

@sfc-gh-truwase Thanks! The PR is ready for review now. DCO and test collection have passed, and the remaining CI job is in progress.

tohtana added this pull request to the merge queue Aug 2, 2026
Merged via the queue into deepspeedai:master with commit ebf1531 Aug 2, 2026
13 checks passed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
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.

[REQUEST] activation checkpoint API should have parity with Pytorch, keywords arguments not supported

3 participants


Back | FazBrowse Home | New Git URL