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 <hope5487@gmail.com>

… unpack error Signed-off-by: Sung Hyun Cho <hope5487@gmail.com>

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 <jianmusings@gmail.com> Signed-off-by: Sung Hyun Cho <hope5487@gmail.com>

…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 <jianmusings@gmail.com>

PyTorch versions that expose autograd.Function.setup_context need the modern forward + setup_context shape for torch.func / functorch. Signed-off-by: Zhang <jianmusings@gmail.com>

Signed-off-by: Zhang <jianmusings@gmail.com>

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 <jianmusings@gmail.com>

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 <hope5487@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>

) 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 <hope5487@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>

@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 <hope5487@gmail.com> Signed-off-by: Zhang <jianmusings@gmail.com> Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: Zhang Jian <jianmusings@gmail.com> Co-authored-by: zhangj1an <jianmusings@gmail.com> Co-authored-by: Zhang Jian <zhang.jian@u.nus.edu> Co-authored-by: Masahiro Tanaka <mtanaka@anyscale.com>

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 <hope5487@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>

…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 <hope5487@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>

@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 <hope5487@gmail.com> Signed-off-by: Zhang <jianmusings@gmail.com> Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: Zhang Jian <jianmusings@gmail.com> Co-authored-by: zhangj1an <jianmusings@gmail.com> Co-authored-by: Zhang Jian <zhang.jian@u.nus.edu> Co-authored-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: nathon-lee <leejianwoo@gmail.com>

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 <hope5487@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> Signed-off-by: nathon-lee <leejianwoo@gmail.com>

…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 <hope5487@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> Signed-off-by: nathon-lee <leejianwoo@gmail.com>

@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 <hope5487@gmail.com> Signed-off-by: Zhang <jianmusings@gmail.com> Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: Zhang Jian <jianmusings@gmail.com> Co-authored-by: zhangj1an <jianmusings@gmail.com> Co-authored-by: Zhang Jian <zhang.jian@u.nus.edu> Co-authored-by: Masahiro Tanaka <mtanaka@anyscale.com>

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 <hope5487@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>

…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 <hope5487@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>

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 Outdated

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

roycho96 added 2 commits March 21, 2026 19:27

fix: fix LinearFunctionForZeroStage3 to support torch.func transforms

33db7c4

Signed-off-by: Sung Hyun Cho <hope5487@gmail.com>

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

39b1755

… unpack error Signed-off-by: Sung Hyun Cho <hope5487@gmail.com>

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 <jianmusings@gmail.com> Signed-off-by: Sung Hyun Cho <hope5487@gmail.com>

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 <jianmusings@gmail.com>

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 <jianmusings@gmail.com>

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 <jianmusings@gmail.com>

update unit tests to reproduce main branch error

8762d00

Signed-off-by: Zhang <jianmusings@gmail.com>

add reproduce scripts

dd037da

Signed-off-by: Zhang <jianmusings@gmail.com>

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

01ee5a6

update reproduce script

f69c1f1

Signed-off-by: Zhang <jianmusings@gmail.com>

update reproduce script to skip repeated env setup

e58ac18

Signed-off-by: Zhang <jianmusings@gmail.com>

update reproduce script to remove duplicated code

3121a7f

Signed-off-by: Zhang <jianmusings@gmail.com>

update reproduce script to print test env

60d20da

Signed-off-by: Zhang <jianmusings@gmail.com>

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 Outdated

zhangj1an and others added 5 commits March 30, 2026 02:29

remove repro scripts

e309a6f

Signed-off-by: Zhang Jian <jianmusings@gmail.com>

update unit test

e425569

Signed-off-by: Zhang Jian <jianmusings@gmail.com>

drop support for pytorch<2.0

39f7e3c

Signed-off-by: Zhang Jian <jianmusings@gmail.com>

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

Uh oh!

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

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants