Add Zenflow code for Stage 1 & 2

[Antlera]

This PR adds ZenFlow, a importance-aware offloaded training framework for DeepSpeed ZeRO. ZenFlow enables multi-step overlap between computation and communication during offloaded training, improving GPU utilization and reducing stalls.

Highlights:

• New ZenFlow optimizers (ZenFlowCPUAdam, ZenFlowSelectiveAdamW)
• ZenFlowZeroOptimizer for ZeRO Stage 1/2 integration
• Configurable via ZenFlowConfig, integrated with DeepSpeedZeroConfig
• Unit tests and documentation included

Note: This PR focuses on Stage 1 and 2 integration. Stage 3 support will be introduced in a follow-up PR.

[tohtana]

Hi @Antlera, thank you for submitting a great PR!

I added some comments. Overall, I think we need to separate ZenFlow code and minimize changes for ZenFlow in existing code

[Antlera]

Hi @Antlera, thank you for submitting a great PR!

I added some comments. Overall, I think we need to separate ZenFlow code and minimize changes for ZenFlow in existing code

Hi @tohtana, thank you for the thoughtful review and suggestions!

I tried my best to avoid adding ZenFlow logic directly into engine and zero optimizer. But for some shared functions like average_tensor, fully separating it would mean rewriting a large function with mostly duplicated code, which might make future maintenance harder when upstream changes.

I’m happy to improve this further if this is considered a better practice — I’m just not entirely sure if full separation is the right trade-off here.

[tohtana]

Hi @Antlera,
I agree that it wouldn't be a good idea to try full separation. The engine and optimizer are not currently designed for the flexible extensions.
Can you first try to separate the parts I mentioned first? Then we can discuss if we have a chance to do more. If you have any concern, please share it here.

[Antlera]

@sfc-gh-truwase All copyright issues have been fixed.

[sfc-gh-truwase]

@delock do you have additional concerns or can we merge this? Thanks

[Antlera]

The cpu-torch-latest CI is failing because PyTorch has released version 2.8, while the workflow’s pytest invocation still expects --torch_ver="2.7". See HF_HOME=/tmp/hf_home/ pytest $PYTEST_OPTS --forked -n 4 unit/ --torch_ver="2.7".

Since the workflow installs the latest CPU build by default, it pulled 2.8.0+cpu, which caused the version check in tests/conftest.py to fail and exit.

=================================== FAILURES ===================================
______ TestMultipleModels.test_zero_optimizer[True-False-False-2-False-3] ______
[gw3] linux -- Python 3.12.3 /home/runner/work/DeepSpeed/DeepSpeed/unit-test-venv/bin/python
:-1: running the test CRASHED with signal 0
------------------------------- captured stderr --------------------------------
/home/runner/work/DeepSpeed/DeepSpeed/tests/conftest.py:50: _pytest.outcomes.Exit: expected torch version 2.7 did not match found torch version 2.8.0+cpu
___________________ TestNoSyncCtxt.test_zero_stage[0-dtype2] ___________________
[gw0] linux -- Python 3.12.3 /home/runner/work/DeepSpeed/DeepSpeed/unit-test-venv/bin/python
:-1: running the test CRASHED with signal 0
------------------------------- captured stderr --------------------------------
/home/runner/work/DeepSpeed/DeepSpeed/tests/conftest.py:50: _pytest.outcomes.Exit: expected torch version 2.7 did not match found torch version 2.8.0+cpu
______ TestMultipleModels.test_zero_optimizer[True-False-False-2-False-2] ______
[gw1] linux -- Python 3.12.3 /home/runner/work/DeepSpeed/DeepSpeed/unit-test-venv/bin/python
:-1: running the test CRASHED with signal 0
------------------------------- captured stderr --------------------------------
/home/runner/work/DeepSpeed/DeepSpeed/tests/conftest.py:50: _pytest.outcomes.Exit: expected torch version 2.7 did not match found torch version 2.8.0+cpu

[sfc-gh-truwase]

@Antlera should be fixed by #7481

[Antlera]

@sfc-gh-truwase Merged for checking the new ut.

[Antlera]

@sfc-gh-truwase Not sure about what happen to the new up-coming errors in the CIs.

It says

Run modal run -m ci.torch_latest
╭─ Error ──────────────────────────────────────────────────────────────────────╮
│ Token missing. Could not authenticate client. If you have token credentials, │
│ see modal.com/docs/reference/modal.config for setup help. If you are a new   │
│ user, register an account at modal.com, then run `modal token new`.          │
╰──────────────────────────────────────────────────────────────────────────────╯
Error: Process completed with exit code 1.

[Antlera]

This might be relevant to #7289. Possible problem: The CI failures on forked PRs are due to Modal authentication.
modal run -m ci.torch_latest requires a token stored in the repository’s secrets, but forked PRs cannot access these secrets, resulting in a “Token missing” error.

[Antlera]

Merged for checking the new CI. Maybe re-run it will solve the problem. I assume this will make this branch up-to-date.