Functorch support: RuntimeError: In order to use an autograd.Function with functorch transforms

[ifiaposto]

Hi,

Are you planning to support functorch? I am currently getting:

[rank2]: File "/home/iapostol/miniconda3/envs/scaff/lib/python3.12/site-packages/deepspeed/runtime/zero/linear.py", line 109, in zero3_linear_wrap
[rank2]: return LinearFunctionForZeroStage3.apply(input, weight)
[rank2]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank2]: File "/home/iapostol/miniconda3/envs/scaff/lib/python3.12/site-packages/torch/autograd/function.py", line 577, in apply
[rank2]: raise RuntimeError(
[rank2]: RuntimeError: In order to use an autograd.Function with functorch transforms (vmap, grad, jvp, jacrev, ...), it must override the setup_context staticmethod. For more details, please see https://pytorch.org/docs/main/notes/extending.func.html

[Flink-ddd]

Hi @ifiaposto

Thanks for raising this important issue. The error message clearly points out that DeepSpeed's internal autograd.Function needs to be updated to support functorch (now torch.func) transformations by implementing the setup_context method.

This is a great suggestion for improving compatibility with modern PyTorch features. While I'm not a core developer, I can see this being very valuable for users who leverage functional programming models.

To help the developers better understand the use case and prioritize this, could you perhaps provide a minimal, self-contained code snippet that reproduces the RuntimeError? This would be extremely helpful for debugging and testing a potential fix.

In the meantime, I'll also try to look into the relevant parts of the DeepSpeed source code to see how complex a patch might be.

Best