Refactor ZenFlow integration in DeepSpeedEngine

- Move `_configure_zenflow` logic to a standalone `configure_zenflow()` function in `zenflow_utils.py`
- Refactor ZenFlow place to decouple it from ZeRO internals

Signed-off-by: Tingfeng Lan <[email protected]>

Author: Antlera

deepspeed/runtime/engine.py

@@ -27,11 +27,12 @@
 from deepspeed.runtime.utils import see_memory_usage, DummyOptim
 from .zero.offload_config import OffloadDeviceEnum, OffloadStateTypeEnum
 from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer
-from deepspeed.runtime.zero.zenflow.zenflow_stage_1_and_2 import ZenFlowZeroOptimizer
+from deepspeed.runtime.zenflow.zenflow_stage_1_and_2 import ZenFlowZeroOptimizer
 from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
 from deepspeed.runtime.zero.utils import is_zero_supported_optimizer, ZeRORuntimeException
 from deepspeed.runtime.zero.parameter_offload import DeepSpeedZeRoOffload
 from deepspeed.runtime.zero.config import ZERO_OPTIMIZATION
+from deepspeed.runtime.zenflow.zenflow_utils import configure_zenflow
 
 from deepspeed.runtime.fp16.fused_optimizer import FP16_Optimizer
 from deepspeed.runtime.fp16.unfused_optimizer import FP16_UnfusedOptimizer
@@ -334,6 +335,8 @@ def __init__(self,
         if self.torch_autocast_enabled():
             init_autocast_params(self, self.torch_autocast_dtype(), self.torch_autocast_lower_precision_safe_modules())
 
+        configure_zenflow(self)
+
         if has_optimizer:
             self._configure_optimizer(optimizer, model_parameters)
             self._configure_lr_scheduler()

deepspeed/runtime/zero/zenflow/__init__.py renamed to deepspeed/runtime/zenflow/__init__.py

@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+from pydantic import Field, model_validator
+from typing import Optional, Union
+
+from deepspeed.runtime.config_utils import DeepSpeedConfigModel
+
+
+class ZenFlowConfig(DeepSpeedConfigModel):
+    """Configuration options for ZenFlow optimization module."""
+
+    topk_ratio: float = Field(0.1, ge=0.0, le=1.0)
+    """Ratio of top-k important gradient columns to retain (range: 0.0 to 1.0)."""
+
+    select_strategy: str = "auto"
+    """Strategy for selecting important gradient indices.
+    Options: "auto", "step", or "epoch"."""
+
+    select_interval: Union[str, int] = "auto"
+    """Interval at which to reselect important gradient indices.
+    Can be "auto" or a fixed integer step/epoch interval."""
+
+    update_interval: Union[str, int] = "auto"
+    """Interval for applying accumulated unimportant gradients to model parameters.
+    Can be "auto" or a fixed integer step interval."""
+
+    overlap_step: bool = False
+    """Whether to overlap CPU-side optimizer steps with forward/backward computation."""
+
+    offload: bool = False
+    """Whether to offload selective optimizer states to CPU to save memory."""
+
+    auto_ratio: float = Field(0.99, ge=0.0, le=1.0)
+    """Threshold used in the "auto" strategy to determine update_interval."""
+
+    full_warm_up_rounds: int = 0
+    """Number of initial rounds during which all gradients are fully updated (no selection)."""
+
+    steps_per_epoch: Optional[int] = Field(
+        default=None,
+        description=
+        "Number of steps per epoch. This field is initialized during execution and should not be set by users.",
+        exclude=True)
+
+    @model_validator(mode="after")
+    def validate_fields(self):
+        if self.select_strategy not in ["auto", "step", "epoch"]:
+            raise ValueError('select_strategy must be one of "auto", "step", or "epoch"')
+
+        if isinstance(self.select_interval, str) and self.select_interval != "auto":
+            raise ValueError('If select_interval is a string, it must be "auto"')
+
+        if isinstance(self.update_interval, str) and self.update_interval != "auto":
+            raise ValueError('If update_interval is a string, it must be "auto"')
+
+        if not isinstance(self.full_warm_up_rounds, int):
+            raise ValueError('full_warm_up_rounds must be an integer')
+
+        return self

deepspeed/runtime/zenflow/zenflow_config.py

@@ -0,0 +1,88 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from deepspeed.runtime.engine import DeepSpeedEngine
+
+
+def _flatten_dense_tensors(tensors):
+    """Flatten dense tensors into a contiguous 1D buffer. Assume tensors are of
+    same dense type.
+
+    Since inputs are dense, the resulting tensor will be a concatenated 1D
+    buffer. Element-wise operation on this buffer will be equivalent to
+    operating individually.
+
+    Args:
+        tensors (Iterable[Tensor]): dense tensors to flatten.
+
+    Returns:
+        A contiguous 1D buffer containing input tensors.
+    """
+    transposed_tensors = [t.transpose(0, 1).contiguous() if t.dim() == 2 else t for t in tensors]
+    return torch._C._nn.flatten_dense_tensors(transposed_tensors)
+
+
+def _unflatten_dense_tensors(flat, tensors):
+    """View a flat buffer using the sizes of tensors. Assume that tensors are of
+    same dense type, and that flat is given by _flatten_dense_tensors.
+
+    Args:
+        flat (Tensor): flattened dense tensors to unflatten.
+        tensors (Iterable[Tensor]): dense tensors whose sizes will be used to
+          unflatten flat.
+
+    Returns:
+        Unflattened dense tensors with sizes same as tensors and values from
+        flat.
+    """
+    transposed_tensors = [t.transpose(0, 1) if t.dim() == 2 else t for t in tensors]
+    unflat = torch._C._nn.unflatten_dense_tensors(flat, transposed_tensors)
+    return [t.transpose(0, 1) if t.dim() == 2 else t for t in unflat]
+
+
+def configure_zenflow(engine: "DeepSpeedEngine") -> None:
+    zenflow_config = engine.zenflow_config()
+    if zenflow_config == None:
+        engine.zenflow = False
+        return
+
+    engine.zenflow = True
+    select_strategy = zenflow_config.select_strategy
+
+    if select_strategy == 'auto':
+        select_strategy = "epoch"
+        if isinstance(zenflow_config.select_interval, int):
+            raise Warning(
+                "If use auto select strategy, select_interval will be set to 1 and select_strategy will be set to epoch, thus select_interval would be overwritten."
+            )
+        engine.select_interval = 1
+    else:
+        if isinstance(zenflow_config.select_interval, str):
+            raise ValueError("If don't use auto select strategy, select_interval must be a number.")
+        engine.select_interval = zenflow_config.select_interval
+
+    if isinstance(zenflow_config.update_interval, str):
+        engine.auto_update = True
+        engine.update_interval = 0
+    else:
+        engine.auto_update = False
+        engine.update_interval = int(zenflow_config.update_interval)
+
+    if select_strategy == 'epoch':
+        zenflow_config.steps_per_epoch = len(engine.training_dataloader)
+        engine.select_interval = engine.select_interval * len(engine.training_dataloader)
+
+    if not engine.auto_update and engine.select_interval != 0 and engine.select_interval < engine.update_interval:
+        raise ValueError("Select interval must be greater or equal to update interval")
+
+    engine.overlap_step = zenflow_config.overlap_step
+
+    engine.full_warm_up_rounds = zenflow_config.full_warm_up_rounds
+
+    engine._config.gradient_accumulation_steps = engine.update_interval

deepspeed/runtime/zero/zenflow/zenflow_stage_1_and_2.py renamed to deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py

@@ -9,7 +9,8 @@
 from pydantic import Field, model_validator
 from deepspeed.runtime.config_utils import get_scalar_param, pp_int, DeepSpeedConfigModel
 from deepspeed.utils import logger
-from .offload_config import DeepSpeedZeroOffloadParamConfig, DeepSpeedZeroOffloadOptimizerConfig, OffloadDeviceEnum, ZenFlowConfig
+from .offload_config import DeepSpeedZeroOffloadParamConfig, DeepSpeedZeroOffloadOptimizerConfig, OffloadDeviceEnum
+from deepspeed.runtime.zenflow.zenflow_config import ZenFlowConfig
 
 # ZeRO optimization. By default, this optimization is not enabled.
 # Users have to configure the desired optimization (0 means disabled) in params.json as below example:

deepspeed/runtime/zenflow/zenflow_utils.py

@@ -6,7 +6,7 @@
 from enum import Enum
 from pathlib import Path
 from pydantic import Field, model_validator
-from typing import Optional, Union
+from typing import Optional
 
 from deepspeed.runtime.config_utils import DeepSpeedConfigModel, pp_int
 
@@ -100,59 +100,6 @@ def set_pipeline(self):
         return self
 
 
-class ZenFlowConfig(DeepSpeedConfigModel):
-    """Configuration options for ZenFlow optimization module."""
-
-    topk_ratio: float = Field(0.1, ge=0.0, le=1.0)
-    """Ratio of top-k important gradient columns to retain (range: 0.0 to 1.0)."""
-
-    select_strategy: str = "auto"
-    """Strategy for selecting important gradient indices.
-    Options: "auto", "step", or "epoch"."""
-
-    select_interval: Union[str, int] = "auto"
-    """Interval at which to reselect important gradient indices.
-    Can be "auto" or a fixed integer step/epoch interval."""
-
-    update_interval: Union[str, int] = "auto"
-    """Interval for applying accumulated unimportant gradients to model parameters.
-    Can be "auto" or a fixed integer step interval."""
-
-    overlap_step: bool = False
-    """Whether to overlap CPU-side optimizer steps with forward/backward computation."""
-
-    offload: bool = False
-    """Whether to offload selective optimizer states to CPU to save memory."""
-
-    auto_ratio: float = Field(0.99, ge=0.0, le=1.0)
-    """Threshold used in the "auto" strategy to determine update_interval."""
-
-    full_warm_up_rounds: int = 0
-    """Number of initial rounds during which all gradients are fully updated (no selection)."""
-
-    steps_per_epoch: Optional[int] = Field(
-        default=None,
-        description=
-        "Number of steps per epoch. This field is initialized during execution and should not be set by users.",
-        exclude=True)
-
-    @model_validator(mode="after")
-    def validate_fields(self):
-        if self.select_strategy not in ["auto", "step", "epoch"]:
-            raise ValueError('select_strategy must be one of "auto", "step", or "epoch"')
-
-        if isinstance(self.select_interval, str) and self.select_interval != "auto":
-            raise ValueError('If select_interval is a string, it must be "auto"')
-
-        if isinstance(self.update_interval, str) and self.update_interval != "auto":
-            raise ValueError('If update_interval is a string, it must be "auto"')
-
-        if not isinstance(self.full_warm_up_rounds, int):
-            raise ValueError('full_warm_up_rounds must be an integer')
-
-        return self
-
-
 class OffloadStateTypeEnum(str, Enum):
     """ Enum for internal buffer types """
     optim_states = "optim_states"

deepspeed/runtime/zero/config.py

@@ -11,7 +11,7 @@
 from typing import List, Dict
 
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
-from deepspeed.runtime.zero.zenflow import zenflow_utils
+from deepspeed.runtime.zenflow import zenflow_utils
 
 from deepspeed.runtime.base_optimizer import ZeROOptimizer
 from deepspeed.runtime.fp16.loss_scaler import CreateLossScaler
@@ -1933,6 +1933,8 @@ def _optimizer_step(self, group_no):
         if self.torch_autocast_gradscaler:
             self.torch_autocast_gradscaler.step(self.optimizer)
             self.torch_autocast_gradscaler.update()
+        elif self.zenflow:
+            self.zenflow_cpu_optimizer_step(group_no)
         else:
             self.optimizer.step()
         self.optimizer.param_groups = original_param_groups

deepspeed/runtime/zero/offload_config.py

@@ -7,7 +7,8 @@
 from pydantic import ValidationError
 
 from deepspeed.runtime.zero.config import DeepSpeedZeroConfig, ZeroStageEnum
-from deepspeed.runtime.zero.offload_config import DeepSpeedZeroOffloadOptimizerConfig, ZenFlowConfig
+from deepspeed.runtime.zenflow.zenflow_config import ZenFlowConfig
+from deepspeed.runtime.zero.offload_config import DeepSpeedZeroOffloadOptimizerConfig
 
 
 def test_stage_enum_accepts_int_and_enum():