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

[roycho96]

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:

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

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 4a4bd2ad5a

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