fix: add setup_context for torch.func compatibility #7916
tohtana
merged
28 commits into
deepspeedai:master
from
roycho96:fix/support-func-torch
fix: fix LinearFunctionForZeroStage3 to support torch.func transforms
33db7c43
fix: always pass bias arg in zero3_linear_wrap to avoid setup_context…
39b1755a
zhangj1an
commented
on 2026-03-21
fix: remove @autocast_custom_fwd from forward, move autocast state to…
6df37afb
fix(zero3): replace custom_bwd with explicit autocast for functorch-s…
c0b9694c
fix(zero): use setup_context for offload pre/post backward Functions
5e83d056
Merge branch 'master' into fix/support-func-torch
74837016
run pre-commit checks
a1e798d5
update unit tests to reproduce main branch error
8762d00f
add reproduce scripts
dd037da2
Merge branch 'master' into fix/support-func-torch
01ee5a64
update reproduce script
f69c1f1a
update reproduce script to skip repeated env setup
e58ac18c
update reproduce script to remove duplicated code
3121a7f9
update reproduce script to print test env
60d20da7
chatgpt-codex-connector
commented
on 2026-03-25
drop PyTorch < 2.0 support and fix autocast backward in ZeRO linear
bb245b2e
change PyTorch version in README
04c456f8
resolve conflict with master
703aad38
Merge pull request #1 from tohtana/tohtana/pr7916-merge-master-resolve
8468149b
remove repro scripts
e309a6f0
update unit test
e4255697
drop support for pytorch<2.0
39f7e3ce
Merge branch 'master' into fix/support-func-torch
39c9a737
Merge branch 'master' into fix/support-func-torch
a5aa09a0
Merge branch 'master' into fix/support-func-torch
10906aea
Merge branch 'master' into fix/support-func-torch
dcd13baf
Merge branch 'master' into fix/support-func-torch
9599f271
Merge branch 'master' into fix/support-func-torch
3b30b39a
Merge branch 'master' into fix/support-func-torch
b6ef66c5
tohtana
approved these changes
on 2026-05-21
Assignees
No one assigned
Labels
None yet
Milestone
No milestone
Login to write a write a comment.
Login via GitHub