DeepSpeed

bffaf457 - Fix no-grad grad-fn lookup in ZeRO hook counting on PyTorch 2.3 (#7830) (#7841)

Fix no-grad grad-fn lookup in ZeRO hook counting on PyTorch 2.3 (#7830) (#7841) Fixes #7830 In torch==2.3, `_get_grad_fn_or_grad_acc` can fail in backward-hook no-grad context (NoneType.next_functions), which breaks `count_used_parameters_in_backward`. This PR wraps that lookup with `torch.enable_grad()` (matching newer torch behavior) and updates the unit test accordingly. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>