JoshWoo2003 7cb1b88ec4 Add ZenFlow code for Stage 3 (#7516) ...

This PR completes the ZenFlow integration for DeepSpeed ZeRO Stage 3. 

Highlights:

- ZenFlowSelectiveAdamW_stage3: Optimizer with importance-aware
selective parameter updates for ZeRO Stage 3.
- ZenFlowZeroOptimizer_Stage3: Full Stage 3 optimizer integration with
partitioned parameters and CPU offload.
- Configurable via ZenFlowConfig, fully integrated with
DeepSpeedZeroConfig for Stage 3.
- Unit tests for Stage 3 cases ensuring correctness and compatibility.

Note: Intergration with ZeRO Stage 1&2 was introduced in #7391

---------

Signed-off-by: Yusen Wu <xrn4ub@virginia.edu>
Co-authored-by: Ma, Guokai <guokai.ma@intel.com>
Co-authored-by: Olatunji Ruwase <tunji.ruwase@snowflake.com>
Co-authored-by: Tingfeng Lan <erc8gx@virginia.edu>

2025-10-13 12:19:18 -04:00

..

autotuning

fix typo s/1014 /1024 (#7528)

2025-09-01 01:12:40 +00:00

checkpoint

Fix the universal checkpoint issue for stage3 when there are multiple subgroups. (#7585)

2025-09-27 17:39:43 +00:00

comm

avoid setting device_id to init_process_group (#7542)

2025-09-05 06:06:26 +00:00

compile

DeepCompile: Use min_cut_rematerialization for partitioning joint graphs (#7609)

2025-10-03 03:39:38 +00:00

compression

Fix invalid f-strings (#7457)

2025-08-16 18:22:19 +00:00

elasticity

Update yapf version (#6721)

2024-11-06 18:57:12 +00:00

inference

[MoE] Fix misuse of num_experts as expert parallel group size (ep_size) (#7551)

2025-09-09 22:31:44 -07:00

io

DeepNVMe update (#7215)

2025-06-06 18:49:41 -04:00

launcher

[logging] less startup noise (#7526)

2025-09-02 19:14:57 +00:00

linear

Fix invalid f-strings (#7457)

2025-08-16 18:22:19 +00:00

model_implementations

Fix invalid f-strings (#7457)

2025-08-16 18:22:19 +00:00

module_inject

Fix invalid f-strings (#7457)

2025-08-16 18:22:19 +00:00

moe

Fix invalid f-strings (#7457)

2025-08-16 18:22:19 +00:00

monitor

fix wandb.log() call by removing sync kwarg (#7383)

2025-06-23 16:07:35 +00:00

nebula

fix typo with deepspeed/ (#3547)

2023-06-02 00:47:14 +00:00

nvme

Fix invalid f-strings (#7457)

2025-08-16 18:22:19 +00:00

ops

Add ZenFlow code for Stage 3 (#7516)

2025-10-13 12:19:18 -04:00

pipe

Update DeepSpeed copyright license to Apache 2.0 (#3111)

2023-03-30 17:14:38 -07:00

profiling

Flops profiler support for F.interpolate (#7353)

2025-06-20 15:02:56 +00:00

runtime

Add ZenFlow code for Stage 3 (#7516)

2025-10-13 12:19:18 -04:00

sequence

fix issues raised by Coverity scans (#7431)

2025-08-02 12:16:10 -04:00

utils

add print_dist util (#7621)

2025-10-03 19:30:26 -07:00

__init__.py

Make Muon optimizer easier to enable (#7555)

2025-09-17 09:52:11 -04:00

accelerator

Abstract accelerator (step 2) (#2560)

2023-01-06 23:40:58 -05:00

constants.py

DeepNVMe update (#7215)

2025-06-06 18:49:41 -04:00

env_report.py

Fix invalid f-strings (#7457)

2025-08-16 18:22:19 +00:00

git_version_info.py

Make op builder detection adapt to accelerator change (#5206)

2024-03-12 20:48:29 +00:00