fix: add setup_context for torch.func compatibility by roycho96 · Pull Request #7916 · deepspeedai/DeepSpeed

roycho96 · 2026-03-21T09:22:42Z

LinearFunctionForZeroStage3 uses the legacy forward(ctx, ...) pattern which is incompatible with torch.func transforms (torch.func.grad, torch.func.grad_and_value, vmap, etc.):

RuntimeError: In order to use an autograd.Function with functorch transforms
(vmap, grad, jvp, jacrev, ...), it must override the setup_context staticmethod.

This affects any library that uses torch.func internally on a ZeRO-3 model.

Fix

Fixes #7913

Note

As pointed out by @zhangj1an in #7913, PostBackwardFunctionModule and PreBackwardFunctionForModule in parameter_offload.py have the same issue. Those will be addressed in a follow-up commit within this PR.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: eed37042bc

ℹ️ 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".

Signed-off-by: Sung Hyun Cho <[email protected]>

… unpack error Signed-off-by: Sung Hyun Cho <[email protected]>

roycho96 · 2026-03-22T03:17:12Z

Thanks for the work! I implemented the same fix, so it looks good to me. To reduce reviewer's effort, I had 2 minor comments. This should lead to linear.py to only have 17 insertions (+) and 7 deletions (-).

Hi @zhangj1an, I've sent you a collaborator invite to my fork. Feel free to push your fix directly to the branch. Thanks for the suggestion!

… setup_context Co-authored-by: zhangj1an <[email protected]> Signed-off-by: Sung Hyun Cho <[email protected]>

…afe linear Avoid asymmetric custom_bwd without custom_fwd on the setup_context forward path; mirror forward AMP in backward via torch.amp.autocast. Signed-off-by: Zhang <[email protected]>

PyTorch versions that expose autograd.Function.setup_context need the modern forward + setup_context shape for torch.func / functorch. Signed-off-by: Zhang <[email protected]>

Signed-off-by: Zhang <[email protected]>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 60d20da79f

ℹ️ 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".

Signed-off-by: Zhang Jian <[email protected]>

roycho96 · 2026-04-04T17:12:38Z

@tohtana
Thank you for the review! Your PR #7920 is merged!
Seems like all the conflicts are solved.

roycho96 · 2026-05-17T09:04:27Z

@tohtana
Gentle ping, happy to address anything if needed, thanks!

tohtana · 2026-05-20T19:18:50Z

Hi @roycho96,
Sorry I have missed the update for a long time, and thank you for pinging me.
I'm happy to merge this PR now, but it seems the current update doesn't cover vmap. Do you want to update this PR before we merge or merge it now leaving it as follow-up work.

roycho96 · 2026-05-21T02:37:17Z

Hi @roycho96, Sorry I have missed the update for a long time, and thank you for pinging me. I'm happy to merge this PR now, but it seems the current update doesn't cover vmap. Do you want to update this PR before we merge or merge it now leaving it as follow-up work.

Thanks @tohtana! Sounds good to merge this one and I'll send a follow-up for vmap. It'll need generate_vmap_rule on LinearFunctionForZeroStage3 and the offload PreBackwardFunctionForModule / PostBackwardFunctionModule.

tohtana

Okay let's leave it for now and merge this PR now.

Follow-up to #7916. Adds `generate_vmap_rule = True` to `LinearFunctionForZeroStage3` so `torch.func.vmap` works on the Function directly. The previous PR covered `grad` / `jacrev` via `setup_context` but not `vmap`. Test: `pytest tests/unit/v1/zero/test_zero_functorch_linear.py::TestLinearFunctionVmap` Signed-off-by: Sung Hyun Cho <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]>

) Follow-up to #7916 and #8023. Makes `torch.func.grad / grad_and_value / jacrev` and `vmap(grad)` work when called directly on a DeepSpeed engine for ZeRO 0/1/2. | API | ZeRO-0 | ZeRO-1 | ZeRO-2 | ZeRO-3 | |---|---|---|---|---| | `torch.func.grad(lambda x: engine(x))(x)` | ✓ | ✓ | ✓ | not yet | | `torch.func.grad_and_value(lambda x: engine(x))(x)` | ✓ | ✓ | ✓ | not yet | | `torch.func.jacrev(lambda x: engine(x))(x)` | ✓ | ✓ | ✓ | not yet | | `torch.func.vmap(torch.func.grad(...))(x_batch)` | ✓ | ✓ | ✓ | not yet | | `torch.func.vmap(lambda x: engine(x))(x_batch)` | already ✓ | already ✓ | already ✓ | not yet | | `engine.backward(loss)` (regression) | ✓ | ✓ | ✓ | ✓ | `vmap` alone runs only the forward graph so it never hit the broken backward hooks and already worked before this PR; included in the table for completeness. Usage: ```python engine, _, _, _ = deepspeed.initialize(model=model, ...) # input gradient g = torch.func.grad(lambda xi: engine(xi))(x) # gradient and value in one pass g, v = torch.func.grad_and_value(lambda xi: engine(xi))(x) # Jacobian of output w.r.t. input J = torch.func.jacrev(lambda xi: engine(xi))(x) # per-sample input gradients (batched) per_sample_g = torch.func.vmap(torch.func.grad(lambda xi: engine(xi)))(x_batch) ``` ZeRO-3 hits a separate SIGSEGV from the same APIs and is tracked separately. Test: `pytest tests/unit/v1/zero/test_zero_torch_func.py` --------- Signed-off-by: Sung Hyun Cho <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]>

@zhangj1an

`LinearFunctionForZeroStage3` uses the legacy `forward(ctx, ...)` pattern which is incompatible with `torch.func` transforms (`torch.func.grad`, `torch.func.grad_and_value`, `vmap`, etc.): ``` RuntimeError: In order to use an autograd.Function with functorch transforms (vmap, grad, jvp, jacrev, ...), it must override the setup_context staticmethod. ``` This affects any library that uses `torch.func` internally on a ZeRO-3 model. ## Fix Fixes deepspeedai#7913 ## Note As pointed out by @zhangj1an in deepspeedai#7913, `PostBackwardFunctionModule` and `PreBackwardFunctionForModule` in `parameter_offload.py` have the same issue. Those will be addressed in a follow-up commit within this PR. --------- Signed-off-by: Sung Hyun Cho <[email protected]> Signed-off-by: Zhang <[email protected]> Signed-off-by: Masahiro Tanaka <[email protected]> Signed-off-by: Zhang Jian <[email protected]> Co-authored-by: zhangj1an <[email protected]> Co-authored-by: Zhang Jian <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]>

Follow-up to deepspeedai#7916. Adds `generate_vmap_rule = True` to `LinearFunctionForZeroStage3` so `torch.func.vmap` works on the Function directly. The previous PR covered `grad` / `jacrev` via `setup_context` but not `vmap`. Test: `pytest tests/unit/v1/zero/test_zero_functorch_linear.py::TestLinearFunctionVmap` Signed-off-by: Sung Hyun Cho <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]>

…epspeedai#8026) Follow-up to deepspeedai#7916 and deepspeedai#8023. Makes `torch.func.grad / grad_and_value / jacrev` and `vmap(grad)` work when called directly on a DeepSpeed engine for ZeRO 0/1/2. | API | ZeRO-0 | ZeRO-1 | ZeRO-2 | ZeRO-3 | |---|---|---|---|---| | `torch.func.grad(lambda x: engine(x))(x)` | ✓ | ✓ | ✓ | not yet | | `torch.func.grad_and_value(lambda x: engine(x))(x)` | ✓ | ✓ | ✓ | not yet | | `torch.func.jacrev(lambda x: engine(x))(x)` | ✓ | ✓ | ✓ | not yet | | `torch.func.vmap(torch.func.grad(...))(x_batch)` | ✓ | ✓ | ✓ | not yet | | `torch.func.vmap(lambda x: engine(x))(x_batch)` | already ✓ | already ✓ | already ✓ | not yet | | `engine.backward(loss)` (regression) | ✓ | ✓ | ✓ | ✓ | `vmap` alone runs only the forward graph so it never hit the broken backward hooks and already worked before this PR; included in the table for completeness. Usage: ```python engine, _, _, _ = deepspeed.initialize(model=model, ...) # input gradient g = torch.func.grad(lambda xi: engine(xi))(x) # gradient and value in one pass g, v = torch.func.grad_and_value(lambda xi: engine(xi))(x) # Jacobian of output w.r.t. input J = torch.func.jacrev(lambda xi: engine(xi))(x) # per-sample input gradients (batched) per_sample_g = torch.func.vmap(torch.func.grad(lambda xi: engine(xi)))(x_batch) ``` ZeRO-3 hits a separate SIGSEGV from the same APIs and is tracked separately. Test: `pytest tests/unit/v1/zero/test_zero_torch_func.py` --------- Signed-off-by: Sung Hyun Cho <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]>

@zhangj1an

`LinearFunctionForZeroStage3` uses the legacy `forward(ctx, ...)` pattern which is incompatible with `torch.func` transforms (`torch.func.grad`, `torch.func.grad_and_value`, `vmap`, etc.): ``` RuntimeError: In order to use an autograd.Function with functorch transforms (vmap, grad, jvp, jacrev, ...), it must override the setup_context staticmethod. ``` This affects any library that uses `torch.func` internally on a ZeRO-3 model. ## Fix Fixes deepspeedai#7913 ## Note As pointed out by @zhangj1an in deepspeedai#7913, `PostBackwardFunctionModule` and `PreBackwardFunctionForModule` in `parameter_offload.py` have the same issue. Those will be addressed in a follow-up commit within this PR. --------- Signed-off-by: Sung Hyun Cho <[email protected]> Signed-off-by: Zhang <[email protected]> Signed-off-by: Masahiro Tanaka <[email protected]> Signed-off-by: Zhang Jian <[email protected]> Co-authored-by: zhangj1an <[email protected]> Co-authored-by: Zhang Jian <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]> Signed-off-by: nathon-lee <[email protected]>

Follow-up to deepspeedai#7916. Adds `generate_vmap_rule = True` to `LinearFunctionForZeroStage3` so `torch.func.vmap` works on the Function directly. The previous PR covered `grad` / `jacrev` via `setup_context` but not `vmap`. Test: `pytest tests/unit/v1/zero/test_zero_functorch_linear.py::TestLinearFunctionVmap` Signed-off-by: Sung Hyun Cho <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]> Signed-off-by: nathon-lee <[email protected]>

…epspeedai#8026) Follow-up to deepspeedai#7916 and deepspeedai#8023. Makes `torch.func.grad / grad_and_value / jacrev` and `vmap(grad)` work when called directly on a DeepSpeed engine for ZeRO 0/1/2. | API | ZeRO-0 | ZeRO-1 | ZeRO-2 | ZeRO-3 | |---|---|---|---|---| | `torch.func.grad(lambda x: engine(x))(x)` | ✓ | ✓ | ✓ | not yet | | `torch.func.grad_and_value(lambda x: engine(x))(x)` | ✓ | ✓ | ✓ | not yet | | `torch.func.jacrev(lambda x: engine(x))(x)` | ✓ | ✓ | ✓ | not yet | | `torch.func.vmap(torch.func.grad(...))(x_batch)` | ✓ | ✓ | ✓ | not yet | | `torch.func.vmap(lambda x: engine(x))(x_batch)` | already ✓ | already ✓ | already ✓ | not yet | | `engine.backward(loss)` (regression) | ✓ | ✓ | ✓ | ✓ | `vmap` alone runs only the forward graph so it never hit the broken backward hooks and already worked before this PR; included in the table for completeness. Usage: ```python engine, _, _, _ = deepspeed.initialize(model=model, ...) # input gradient g = torch.func.grad(lambda xi: engine(xi))(x) # gradient and value in one pass g, v = torch.func.grad_and_value(lambda xi: engine(xi))(x) # Jacobian of output w.r.t. input J = torch.func.jacrev(lambda xi: engine(xi))(x) # per-sample input gradients (batched) per_sample_g = torch.func.vmap(torch.func.grad(lambda xi: engine(xi)))(x_batch) ``` ZeRO-3 hits a separate SIGSEGV from the same APIs and is tracked separately. Test: `pytest tests/unit/v1/zero/test_zero_torch_func.py` --------- Signed-off-by: Sung Hyun Cho <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]> Signed-off-by: nathon-lee <[email protected]>

@zhangj1an

`LinearFunctionForZeroStage3` uses the legacy `forward(ctx, ...)` pattern which is incompatible with `torch.func` transforms (`torch.func.grad`, `torch.func.grad_and_value`, `vmap`, etc.): ``` RuntimeError: In order to use an autograd.Function with functorch transforms (vmap, grad, jvp, jacrev, ...), it must override the setup_context staticmethod. ``` This affects any library that uses `torch.func` internally on a ZeRO-3 model. ## Fix Fixes deepspeedai#7913 ## Note As pointed out by @zhangj1an in deepspeedai#7913, `PostBackwardFunctionModule` and `PreBackwardFunctionForModule` in `parameter_offload.py` have the same issue. Those will be addressed in a follow-up commit within this PR. --------- Signed-off-by: Sung Hyun Cho <[email protected]> Signed-off-by: Zhang <[email protected]> Signed-off-by: Masahiro Tanaka <[email protected]> Signed-off-by: Zhang Jian <[email protected]> Co-authored-by: zhangj1an <[email protected]> Co-authored-by: Zhang Jian <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]>

Follow-up to deepspeedai#7916. Adds `generate_vmap_rule = True` to `LinearFunctionForZeroStage3` so `torch.func.vmap` works on the Function directly. The previous PR covered `grad` / `jacrev` via `setup_context` but not `vmap`. Test: `pytest tests/unit/v1/zero/test_zero_functorch_linear.py::TestLinearFunctionVmap` Signed-off-by: Sung Hyun Cho <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]>

…epspeedai#8026) Follow-up to deepspeedai#7916 and deepspeedai#8023. Makes `torch.func.grad / grad_and_value / jacrev` and `vmap(grad)` work when called directly on a DeepSpeed engine for ZeRO 0/1/2. | API | ZeRO-0 | ZeRO-1 | ZeRO-2 | ZeRO-3 | |---|---|---|---|---| | `torch.func.grad(lambda x: engine(x))(x)` | ✓ | ✓ | ✓ | not yet | | `torch.func.grad_and_value(lambda x: engine(x))(x)` | ✓ | ✓ | ✓ | not yet | | `torch.func.jacrev(lambda x: engine(x))(x)` | ✓ | ✓ | ✓ | not yet | | `torch.func.vmap(torch.func.grad(...))(x_batch)` | ✓ | ✓ | ✓ | not yet | | `torch.func.vmap(lambda x: engine(x))(x_batch)` | already ✓ | already ✓ | already ✓ | not yet | | `engine.backward(loss)` (regression) | ✓ | ✓ | ✓ | ✓ | `vmap` alone runs only the forward graph so it never hit the broken backward hooks and already worked before this PR; included in the table for completeness. Usage: ```python engine, _, _, _ = deepspeed.initialize(model=model, ...) # input gradient g = torch.func.grad(lambda xi: engine(xi))(x) # gradient and value in one pass g, v = torch.func.grad_and_value(lambda xi: engine(xi))(x) # Jacobian of output w.r.t. input J = torch.func.jacrev(lambda xi: engine(xi))(x) # per-sample input gradients (batched) per_sample_g = torch.func.vmap(torch.func.grad(lambda xi: engine(xi)))(x_batch) ``` ZeRO-3 hits a separate SIGSEGV from the same APIs and is tracked separately. Test: `pytest tests/unit/v1/zero/test_zero_torch_func.py` --------- Signed-off-by: Sung Hyun Cho <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]>

Xref deepspeedai/DeepSpeed#7916

* updated v0.19.5 * MNT: Re-rendered with conda-smithy 2026.8.9 and conda-forge-pinning 2026.08.13.08.26.4 * Bump minimum version of pytorch to 2.0.0 Xref deepspeedai/DeepSpeed#7916 --------- Co-authored-by: Wei Ji <[email protected]>

roycho96 requested review from tjruwase and tohtana as code owners March 21, 2026 09:22

roycho96 marked this pull request as draft March 21, 2026 09:22

roycho96 mentioned this pull request Mar 21, 2026

[BUG] LinearFunctionForZeroStage3 crashes with torch.func transforms (missing setup_context) #7913

Closed

chatgpt-codex-connector Bot reviewed Mar 21, 2026

View reviewed changes

Comment thread deepspeed/runtime/zero/linear.py Resolved

Comment thread deepspeed/runtime/zero/linear.py Resolved

roycho96 added 2 commits March 21, 2026 19:27

fix: fix LinearFunctionForZeroStage3 to support torch.func transforms

33db7c4

Signed-off-by: Sung Hyun Cho <[email protected]>

fix: always pass bias arg in zero3_linear_wrap to avoid setup_context…

39b1755

… unpack error Signed-off-by: Sung Hyun Cho <[email protected]>

roycho96 force-pushed the fix/support-func-torch branch from 252aea1 to 39b1755 Compare March 21, 2026 10:28

This comment was marked as resolved.

Sign in to view

fix: remove @autocast_custom_fwd from forward, move autocast state to…

6df37af

… setup_context Co-authored-by: zhangj1an <[email protected]> Signed-off-by: Sung Hyun Cho <[email protected]>

zhangj1an force-pushed the fix/support-func-torch branch from 444122c to 6df37af Compare March 22, 2026 08:45

zhangj1an added 2 commits March 22, 2026 08:52

fix(zero3): replace custom_bwd with explicit autocast for functorch-s…

c0b9694

…afe linear Avoid asymmetric custom_bwd without custom_fwd on the setup_context forward path; mirror forward AMP in backward via torch.amp.autocast. Signed-off-by: Zhang <[email protected]>

fix(zero): use setup_context for offload pre/post backward Functions

5e83d05

PyTorch versions that expose autograd.Function.setup_context need the modern forward + setup_context shape for torch.func / functorch. Signed-off-by: Zhang <[email protected]>

zhangj1an force-pushed the fix/support-func-torch branch from 0a66444 to 5e83d05 Compare March 22, 2026 09:34

This comment was marked as resolved.

Sign in to view

zhangj1an and others added 9 commits March 24, 2026 22:21

Merge branch 'master' into fix/support-func-torch

7483701

run pre-commit checks

a1e798d

Signed-off-by: Zhang <[email protected]>

update unit tests to reproduce main branch error

8762d00

Signed-off-by: Zhang <[email protected]>

add reproduce scripts

dd037da

Signed-off-by: Zhang <[email protected]>

Merge branch 'master' into fix/support-func-torch

01ee5a6

update reproduce script

f69c1f1

Signed-off-by: Zhang <[email protected]>

update reproduce script to skip repeated env setup

e58ac18

Signed-off-by: Zhang <[email protected]>

update reproduce script to remove duplicated code

3121a7f

Signed-off-by: Zhang <[email protected]>

update reproduce script to print test env

60d20da

Signed-off-by: Zhang <[email protected]>

This comment was marked as resolved.

Sign in to view

roycho96 marked this pull request as ready for review March 25, 2026 14:35

roycho96 requested a review from loadams as a code owner March 25, 2026 14:35

chatgpt-codex-connector Bot reviewed Mar 25, 2026

View reviewed changes

Comment thread deepspeed/runtime/zero/linear.py Resolved

zhangj1an and others added 3 commits March 30, 2026 02:36

drop support for pytorch<2.0

39f7e3c

Signed-off-by: Zhang Jian <[email protected]>

Merge branch 'master' into fix/support-func-torch

39c9a73

Merge branch 'master' into fix/support-func-torch

a5aa09a

roycho96 added 3 commits April 20, 2026 20:07

Merge branch 'master' into fix/support-func-torch

10906ae

Merge branch 'master' into fix/support-func-torch

dcd13ba

Merge branch 'master' into fix/support-func-torch

9599f27

Merge branch 'master' into fix/support-func-torch

3b30b39

Merge branch 'master' into fix/support-func-torch

b6ef66c

tohtana approved these changes May 21, 2026

View reviewed changes

tohtana merged commit 67b9e21 into deepspeedai:master May 21, 2026
10 of 11 checks passed

roycho96 deleted the fix/support-func-torch branch May 21, 2026 05:27

roycho96 mentioned this pull request May 22, 2026

fix(zero): enable vmap on LinearFunctionForZeroStage3 #8023

Merged

roycho96 mentioned this pull request May 25, 2026

feat(zero): enable torch.func transforms on engine for ZeRO 0/1/2 #8026

Merged

weiji14 added a commit to regro-cf-autotick-bot/deepspeed-feedstock that referenced this pull request Aug 13, 2026

Bump minimum version of pytorch to 2.0.0

b05b939

Xref deepspeedai/DeepSpeed#7916

Navigation Menu

Repository navigation

fix: add setup_context for torch.func compatibility - #7916

fix: add setup_context for torch.func compatibility#7916
tohtana merged 28 commits into
deepspeedai:masterfrom
roycho96:fix/support-func-torch

roycho96 commented Mar 21, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

Uh oh!

Uh oh!

This comment was marked as resolved.

Uh oh!

This comment was marked as resolved.

roycho96 commented Mar 22, 2026

Uh oh!

This comment was marked as resolved.

This comment was marked as resolved.

chatgpt-codex-connector Bot left a comment

Uh oh!

Uh oh!

roycho96 commented Apr 4, 2026

Uh oh!

roycho96 commented May 17, 2026

Uh oh!

tohtana commented May 20, 2026

Uh oh!

roycho96 commented May 21, 2026

Uh oh!

tohtana left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

roycho96 commented Mar 21, 2026

Fix

Note

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

This comment was marked as resolved.

Uh oh!

This comment was marked as resolved.

roycho96 commented Mar 22, 2026

Uh oh!

This comment was marked as resolved.

This comment was marked as resolved.

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

roycho96 commented Apr 4, 2026

Uh oh!

roycho96 commented May 17, 2026

Uh oh!

tohtana commented May 20, 2026

Uh oh!

roycho96 commented May 21, 2026

Uh oh!

tohtana left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants