Feature/zero multi loss separate backward

[nathon-lee]

Summary

Fix ZeRO-3 so two separate backward passes on the same forward graph work correctly when retain_graph=True is used on the first backward.

What changed

• Propagate retain_graph through the ZeRO backward path.
• Delay ZeRO-3 parameter release during retained backward so saved tensors remain valid for the second backward.
• Clear the retained-backward flag safely in finally to avoid state leakage.
• Re-enable ZeRO-3 coverage for the two-loss separate-backward regression test.

Why

This fixes the ZeRO-3 failure where the second backward on the same forward graph hit a tensor size mismatch after zero_grad(). The regression is tracked by issue #7352.

Validation

Test command

CUDA_VISIBLE_DEVICES=0,1 DS_DISABLE_REUSE_DIST_ENV=1 NCCL_IB_DISABLE=1 NCCL_P2P_DISABLE=1 TORCH_NCCL_ASYNC_ERROR_HANDLING=1 TORCH_DISTRIBUTED_DEBUG=DETAIL DS_UNITTEST_TIMEOUT=120 pytest test_zero_user_backward.py -k "test_two_losses_separate_backward_gas1" -vv -s -rs

Result

3 passed
0 skipped
55 deselected

root@964c299dd1b5:/workspace/DeepSpeed_woo/tests# CUDA_VISIBLE_DEVICES=0,1 DS_DISABLE_REUSE_DIST_ENV=1 NCCL_IB_DISABLE=1 NCCL_P2P_DISABLE=1 TORCH_NCCL_ASYNC_ERROR_HANDLING=1 TORCH_DISTRIBUTED_DEBUG=DETAIL DS_UNITTEST_TIMEOUT=120 pytest unit/v1/zero/test_zero_user_backward.py -k "test_two_losses_separate_manual_backward_gas1" -vv -s -rs
================================================================== test session starts ===================================================================
platform linux -- Python 3.12.3, pytest-9.0.3, pluggy-1.6.0 -- /usr/bin/python3.12
cachedir: .pytest_cache
rootdir: /workspace/DeepSpeed_woo/tests
configfile: pytest.ini
plugins: anyio-4.12.0
collected 61 items / 58 deselected / 3 selected                                                                                                          

unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_manual_backward_gas1[1] [Gloo] Rank 0 is connected to 1 peer ranks. Expected number of connected peer ranks is : 1
[Gloo] Rank 1 is connected to 1 peer ranks. Expected number of connected peer ranks is : 1
[rank0]:[W603 07:29:17.335227904 ProcessGroupNCCL.cpp:5072] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
[Gloo] Rank [Gloo] Rank 10 is connected to  is connected to 11 peer ranks.  peer ranks. Expected number of connected peer ranks is : Expected number of connected peer ranks is : 11

[W603 07:29:18.304862233 ProcessGroupNCCL.cpp:1524] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
PASSED
unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_manual_backward_gas1[2] [Gloo] Rank 1 is connected to 1 peer ranks. Expected number of connected peer ranks is : 1
[Gloo] Rank 0 is connected to 1 peer ranks. Expected number of connected peer ranks is : 1
[rank0]:[W603 07:29:30.820676734 ProcessGroupNCCL.cpp:5072] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
[Gloo] Rank 0 is connected to 1 peer ranks. Expected number of connected peer ranks is : 1
[Gloo] Rank 1 is connected to 1 peer ranks. Expected number of connected peer ranks is : 1
[W603 07:29:31.829038701 ProcessGroupNCCL.cpp:1524] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
PASSED
unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_manual_backward_gas1[3] [Gloo] Rank 1 is connected to 1 peer ranks. Expected number of connected peer ranks is : 1
[Gloo] Rank 0 is connected to 1 peer ranks. Expected number of connected peer ranks is : 1
[rank0]:[W603 07:29:42.066717923 ProcessGroupNCCL.cpp:5072] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
[Gloo] Rank 0 is connected to 1 peer ranks. Expected number of connected peer ranks is : 1
[Gloo] Rank 1 is connected to 1 peer ranks. Expected number of connected peer ranks is : 1
PASSED

==================================================================== warnings summary ====================================================================
<string>:8
  <string>:8: PytestDeprecationWarning: A private pytest class or function was used.

unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_manual_backward_gas1[1]
  /workspace/DeepSpeed_woo/tests/conftest.py:47: UserWarning: Running test without verifying torch version, please provide an expected torch version with --torch_ver
    warnings.warn(

unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_manual_backward_gas1[1]
  /workspace/DeepSpeed_woo/tests/conftest.py:54: UserWarning: Running test without verifying cuda version, please provide an expected cuda version with --cuda_ver
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
=================================================================== slowest durations ====================================================================
12.53s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_manual_backward_gas1[2]
12.37s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_manual_backward_gas1[1]
12.00s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_manual_backward_gas1[3]

(6 durations < 1s hidden.)
===================================================== 3 passed, 58 deselected, 3 warnings in 46.71s ======================================================
root@964c299dd1b5:/workspace/DeepSpeed_woo/tests# 
root@964c299dd1b5:/workspace/DeepSpeed_woo/tests# 
root@964c299dd1b5:/workspace/DeepSpeed_woo/tests# CUDA_VISIBLE_DEVICES=0,1 DS_DISABLE_REUSE_DIST_ENV=1 NCCL_IB_DISABLE=1 NCCL_P2P_DISABLE=1 TORCH_NCCL_ASYNC_ERROR_HANDLING=1 TORCH_DISTRIBUTED_DEBUG=DETAIL DS_UNITTEST_TIMEOUT=120 pytest unit/v1/zero/test_zero_user_backward.py -k^C-vv -s -rs
root@964c299dd1b5:/workspace/DeepSpeed_woo/tests# CUDA_VISIBLE_DEVICES=0,1 DS_DISABLE_REUSE_DIST_ENV=1 NCCL_IB_DISABLE=1 NCCL_P2P_DISABLE=1 TORCH_NCCL_ASYNC_ERROR_HANDLING=1 TORCH_DISTRIBUTED_DEBUG=DETAIL DS_UNITTEST_TIMEOUT=120 pytest unit/v1/zero/test_zero_user_backward.py -k "test_two_losses_separate_backward_gas1" -vv -s -rs
================================================================== test session starts ===================================================================
platform linux -- Python 3.12.3, pytest-9.0.3, pluggy-1.6.0 -- /usr/bin/python3.12
cachedir: .pytest_cache
rootdir: /workspace/DeepSpeed_woo/tests
configfile: pytest.ini
plugins: anyio-4.12.0
collected 61 items / 58 deselected / 3 selected                                                                                                          

unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_backward_gas1[1] [Gloo] Rank [Gloo] Rank 1 is connected to 01 is connected to  peer ranks. 1Expected number of connected peer ranks is :  peer ranks. 1Expected number of connected peer ranks is : 
1
[rank0]:[W603 07:30:49.080985243 ProcessGroupNCCL.cpp:5072] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
[Gloo] Rank [Gloo] Rank 10 is connected to  is connected to 11 peer ranks.  peer ranks. Expected number of connected peer ranks is : Expected number of connected peer ranks is : 11

[W603 07:30:50.032584977 ProcessGroupNCCL.cpp:1524] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
PASSED
unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_backward_gas1[2] [Gloo] Rank [Gloo] Rank 01 is connected to  is connected to 11 peer ranks.  peer ranks. Expected number of connected peer ranks is : Expected number of connected peer ranks is : 11

[rank0]:[W603 07:31:01.285483828 ProcessGroupNCCL.cpp:5072] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
[Gloo] Rank [Gloo] Rank 10 is connected to  is connected to 11 peer ranks.  peer ranks. Expected number of connected peer ranks is : Expected number of connected peer ranks is : 11

[W603 07:31:02.257472436 ProcessGroupNCCL.cpp:1524] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
PASSED
unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_backward_gas1[3] [Gloo] Rank 1 is connected to 1 peer ranks. Expected number of connected peer ranks is : 1
[Gloo] Rank 0 is connected to 1 peer ranks. Expected number of connected peer ranks is : 1
[rank0]:[W603 07:31:14.971748647 ProcessGroupNCCL.cpp:5072] Guessing device ID based on global rank. This can cause a hang if rank to GPU mapping is heterogeneous. You can specify device_id in init_process_group()
[Gloo] Rank 1[Gloo] Rank  is connected to 01 is connected to  peer ranks. 1Expected number of connected peer ranks is :  peer ranks. 1Expected number of connected peer ranks is : 
1
PASSED

==================================================================== warnings summary ====================================================================
<string>:8
  <string>:8: PytestDeprecationWarning: A private pytest class or function was used.

unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_backward_gas1[1]
  /workspace/DeepSpeed_woo/tests/conftest.py:47: UserWarning: Running test without verifying torch version, please provide an expected torch version with --torch_ver
    warnings.warn(

unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_backward_gas1[1]
  /workspace/DeepSpeed_woo/tests/conftest.py:54: UserWarning: Running test without verifying cuda version, please provide an expected cuda version with --cuda_ver
    warnings.warn(

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
=================================================================== slowest durations ====================================================================
12.47s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_backward_gas1[3]
12.23s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_backward_gas1[2]
12.09s call     unit/v1/zero/test_zero_user_backward.py::TestZeroUserBackwardSeparateLoss::test_two_losses_separate_backward_gas1[1]

(6 durations < 1s hidden.)
===================================================== 3 passed, 58 deselected, 3 warnings in 44.97s ======================================================
root@964c299dd1b5:/workspace/DeepSpeed_woo/tests# 

Notes

The test previously reproduced a ZeRO-3 RuntimeError on the second backward.
The current run confirms the stage 1/2/3 regression coverage is passing.

cc @tohtana @sfc-gh-truwase

[chatgpt-codex-connector]

💡 Codex Review

DeepSpeed/deepspeed/runtime/engine.py

Lines 4348 to 4350 in f873f69

	if rank == 0:
	self.checkpoint_engine.makedirs(save_dir, exist_ok=True)
	dist.barrier()

Validate checkpoint directory before the barrier

After removing the all-rank save_dir validation, a misconfigured save_checkpoint(save_dir=None) or empty string now fails only on rank 0 inside makedirs, while the other ranks proceed to dist.barrier() and can hang indefinitely. The previous check failed synchronously on every rank before this rank-gated filesystem call, so distributed checkpoint misconfiguration no longer fails cleanly.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Use block eigenvalues when logging events

When monitoring and eigenvalue logging are enabled at a gradient-accumulation boundary, this block now reads self.ev_values[i][0], but self.ev_values is never initialized anywhere on DeepSpeedEngine (the local value is ev_values = self.block_eigenvalue.values()). That makes engine.step() raise AttributeError instead of writing the eigenvalue summaries in that configuration.

Useful? React with 👍 / 👎.