Say when ZenFlow will never re-select its important columns

deepspeed/runtime/zenflow/engine.py

@@ -6,6 +6,7 @@
 from deepspeed import comm as dist
 from typing import TYPE_CHECKING
 from deepspeed.utils.torch import required_torch_version
+from deepspeed.utils import logger
 
 if TYPE_CHECKING:
     from deepspeed.runtime.engine import DeepSpeedEngine
@@ -36,9 +37,9 @@ def configure_zenflow(engine: "DeepSpeedEngine") -> None:
     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."
-            )
+            logger.warning(
+                "ZenFlow: select_strategy is 'auto', so select_interval is one epoch and the "
+                "configured value %s is ignored.", zenflow_config.select_interval)
         engine.select_interval = 1
     else:
         if isinstance(zenflow_config.select_interval, str):
@@ -53,10 +54,27 @@ def configure_zenflow(engine: "DeepSpeedEngine") -> None:
         engine.update_interval = int(zenflow_config.update_interval)
 
     if select_strategy == 'epoch':
-        if engine.training_dataloader is not None:
-            zenflow_config.steps_per_epoch = len(engine.training_dataloader)
-            engine.select_interval = engine.select_interval * len(engine.training_dataloader)
+        # `steps_per_epoch` may already be set by the user; otherwise it can only
+        # come from a dataloader DeepSpeed owns.
+        if not zenflow_config.steps_per_epoch and engine.training_dataloader is not None:
+            # An empty dataloader would assign 0, which the config validator now
+            # rejects on assignment -- and a crash is the wrong answer for a
+            # degenerate-but-legal dataloader. Fall through to the warning.
+            epoch_steps = len(engine.training_dataloader)
+            if epoch_steps > 0:
+                zenflow_config.steps_per_epoch = epoch_steps
+        if zenflow_config.steps_per_epoch:
+            engine.select_interval = engine.select_interval * zenflow_config.steps_per_epoch
         else:
+            # is_zenflow_select_boundary() treats 0 as "never again", so this
+            # leaves the columns chosen at the first step in place for the whole
+            # run. Say so rather than degrading in silence.
+            logger.warning("ZenFlow: select_strategy resolves to 'epoch', but the number of steps in an epoch "
+                           "is unknown -- DeepSpeed reads it from a dataloader it owns, and none was given to "
+                           "deepspeed.initialize(). Important columns will be selected once and never "
+                           "re-selected. Set \"steps_per_epoch\" in the zenflow config, or pass training_data= "
+                           "to deepspeed.initialize(), or use \"select_strategy\": \"step\" with an explicit "
+                           "\"select_interval\".")
             engine.select_interval = 0
 
     if not engine.auto_update and engine.select_interval != 0 and engine.select_interval < engine.update_interval:

deepspeed/runtime/zenflow/engine_stage3.py

@@ -31,9 +31,9 @@ def configure_zenflow(optimizer_z3, zenflow_config):
     if optimizer_z3.select_strategy == 'auto':
         optimizer_z3.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."
-            )
+            logger.warning(
+                "ZenFlow: select_strategy is 'auto', so select_interval is one epoch and the "
+                "configured value %s is ignored.", zenflow_config.select_interval)
         optimizer_z3.select_interval = 1
     else:
         if isinstance(zenflow_config.select_interval, str):
@@ -48,9 +48,14 @@ def configure_zenflow(optimizer_z3, zenflow_config):
         optimizer_z3.update_interval = int(zenflow_config.update_interval)
 
     if optimizer_z3.select_strategy == 'epoch':
-        if zenflow_config.steps_per_epoch is not None:
+        if zenflow_config.steps_per_epoch:
             optimizer_z3.select_interval = optimizer_z3.select_interval * zenflow_config.steps_per_epoch
         else:
+            logger.warning("ZenFlow: select_strategy resolves to 'epoch', but the number of steps in an "
+                           "epoch is unknown. Important columns will be selected once and never "
+                           "re-selected. Set \"steps_per_epoch\" in the zenflow config, or pass "
+                           "training_data= to deepspeed.initialize(), or use \"select_strategy\": "
+                           "\"step\" with an explicit \"select_interval\".")
             optimizer_z3.select_interval = 0
 
     if not optimizer_z3.auto_update and optimizer_z3.select_interval != 0 and optimizer_z3.select_interval < optimizer_z3.update_interval:

deepspeed/runtime/zenflow/zenflow_config.py

@@ -45,8 +45,10 @@ class ZenFlowConfig(DeepSpeedConfigModel):
 
     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.",
+        description="Number of steps in one epoch, used by the 'epoch' select strategy (which "
+        "'auto' resolves to) to turn select_interval into a number of steps. DeepSpeed fills this "
+        "in from the training dataloader when it owns one; set it explicitly when you drive your "
+        "own dataloader, otherwise the important columns are selected once and never re-selected.",
         exclude=True)
 
     @model_validator(mode="after")
@@ -63,6 +65,15 @@ def validate_fields(self):
         if isinstance(self.update_interval, int) and self.update_interval < 1:
             raise ValueError("If update_interval is a number, it must be at least 1")
 
+        # 0 would multiply select_interval straight to 0, which
+        # is_zenflow_select_boundary reads as "never re-select" -- the state the
+        # warning in configure_zenflow exists to announce, reached silently
+        # through the knob that documents it. A negative value is worse: it
+        # leaves select_interval negative, and micro_step % -n fires on a
+        # schedule nobody asked for.
+        if self.steps_per_epoch is not None and self.steps_per_epoch < 1:
+            raise ValueError("If steps_per_epoch is set, it must be at least 1")
+
         if not isinstance(self.full_warm_up_rounds, int):
             raise ValueError('full_warm_up_rounds must be an integer')
 

deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py

@@ -14,6 +14,7 @@
 from deepspeed.moe.utils import is_moe_param
 
 from deepspeed.accelerator import get_accelerator
+from deepspeed.utils import logger
 
 from deepspeed.runtime.utils import all_gather_dp_groups
 
@@ -119,9 +120,9 @@ def _configure_zenflow(self, zenflow_config):
         if self.select_strategy == 'auto':
             self.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."
-                )
+                logger.warning(
+                    "ZenFlow: select_strategy is 'auto', so select_interval is one epoch and the "
+                    "configured value %s is ignored.", zenflow_config.select_interval)
             self.select_interval = 1
         else:
             if isinstance(zenflow_config.select_interval, str):
@@ -136,9 +137,16 @@ def _configure_zenflow(self, zenflow_config):
             self.update_interval = int(zenflow_config.update_interval)
 
         if self.select_strategy == 'epoch':
-            if zenflow_config.steps_per_epoch is not None:
+            if zenflow_config.steps_per_epoch:
                 self.select_interval = self.select_interval * zenflow_config.steps_per_epoch
             else:
+                # 0 makes is_zenflow_select_boundary() true exactly once, so the
+                # columns picked at the first step are used for the whole run.
+                logger.warning("ZenFlow: select_strategy resolves to 'epoch', but the number of steps in an "
+                               "epoch is unknown. Important columns will be selected once and never "
+                               "re-selected. Set \"steps_per_epoch\" in the zenflow config, or pass "
+                               "training_data= to deepspeed.initialize(), or use \"select_strategy\": "
+                               "\"step\" with an explicit \"select_interval\".")
                 self.select_interval = 0
 
         if not self.auto_update and self.select_interval != 0 and self.select_interval < self.update_interval:

tests/unit/runtime/zenflow/test_zf_select_interval.py

@@ -0,0 +1,145 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+"""ZenFlow's select interval, and what it takes for re-selection to happen.
+
+ZenFlow re-picks the important gradient columns every ``select_interval``
+micro-steps. The 'epoch' strategy -- which 'auto', the default, resolves to --
+expresses that interval in epochs, so it needs to know how many steps an epoch
+is. DeepSpeed can only read that from a dataloader it owns, so a caller driving
+its own dataloader has to say. These pin what each combination produces.
+"""
+
+import types
+
+import pytest
+
+from deepspeed.runtime.zenflow.engine import configure_zenflow
+from deepspeed.runtime.zenflow.zenflow_config import ZenFlowConfig
+
+
+class _StubEngine:
+    """The attributes ``configure_zenflow`` reads, and nothing else."""
+
+    def __init__(self, config, training_dataloader=None):
+        self._zenflow_config = config
+        self.training_dataloader = training_dataloader
+        self._config = types.SimpleNamespace(gradient_accumulation_steps=1)
+
+    def zenflow_config(self):
+        return self._zenflow_config
+
+
+class _Loader:
+
+    def __init__(self, steps):
+        self.steps = steps
+
+    def __len__(self):
+        return self.steps
+
+
+def _select_boundaries(select_interval, steps, full_warm_up_rounds=0):
+    """Count the boundaries ``is_zenflow_select_boundary`` would report."""
+    return sum(1 for micro_step in range(steps) if (micro_step - full_warm_up_rounds) >= 0 and (
+        (micro_step - full_warm_up_rounds) == 0 or (select_interval != 0 and micro_step % select_interval == 0)))
+
+
+@pytest.mark.parametrize("select_strategy,select_interval", [("auto", "auto"), ("epoch", 1)])
+def test_epoch_length_from_a_dataloader_deepspeed_owns(select_strategy, select_interval):
+    config = ZenFlowConfig(select_strategy=select_strategy, select_interval=select_interval, update_interval="auto")
+    engine = _StubEngine(config, training_dataloader=_Loader(10))
+
+    configure_zenflow(engine)
+
+    assert config.steps_per_epoch == 10
+    assert engine.select_interval == 10
+    assert _select_boundaries(engine.select_interval, steps=40) == 4
+
+
+@pytest.mark.parametrize("select_strategy,select_interval", [("auto", "auto"), ("epoch", 1)])
+def test_epoch_length_given_by_the_user_needs_no_dataloader(select_strategy, select_interval):
+    # The path for a caller that drives its own dataloader.
+    config = ZenFlowConfig(select_strategy=select_strategy,
+                           select_interval=select_interval,
+                           update_interval="auto",
+                           steps_per_epoch=10)
+    engine = _StubEngine(config, training_dataloader=None)
+
+    configure_zenflow(engine)
+
+    assert engine.select_interval == 10
+    assert _select_boundaries(engine.select_interval, steps=40) == 4
+
+
+@pytest.mark.parametrize("select_strategy,select_interval", [("auto", "auto"), ("epoch", 1)])
+def test_no_epoch_length_selects_once_and_says_so(select_strategy, select_interval, caplog):
+    # select_interval 0 makes is_zenflow_select_boundary true exactly once, so
+    # the columns chosen at the first step are used for the whole run. That is
+    # the shape of the bug; the point of the test is that it is now announced.
+    config = ZenFlowConfig(select_strategy=select_strategy, select_interval=select_interval, update_interval="auto")
+    engine = _StubEngine(config, training_dataloader=None)
+
+    with caplog.at_level("WARNING"):
+        configure_zenflow(engine)
+
+    assert engine.select_interval == 0
+    assert _select_boundaries(engine.select_interval, steps=40) == 1
+    assert any("never re-selected" in record.getMessage() for record in caplog.records)
+
+
+@pytest.mark.parametrize("steps_per_epoch", [0, -5])
+def test_a_non_positive_epoch_length_is_refused(steps_per_epoch):
+    """Reported by @ebarkhordar on #8456.
+
+    0 multiplies select_interval straight to 0, which is_zenflow_select_boundary
+    reads as "never re-select" -- the state the warning exists to announce, reached
+    silently through the knob that documents it, because the warning only fires in
+    the None branch. -5 is worse: select_interval stays negative and
+    `micro_step % -5` fires on a schedule nobody asked for, 8 times in 40 steps.
+
+    validate_fields already rejects an update_interval below 1 one field away.
+    """
+    with pytest.raises(ValueError, match="steps_per_epoch"):
+        ZenFlowConfig(select_strategy="auto",
+                      select_interval="auto",
+                      update_interval="auto",
+                      steps_per_epoch=steps_per_epoch)
+
+
+def test_an_empty_dataloader_warns_rather_than_selecting_once_in_silence(caplog):
+    """The programmatic path to the same 0: len(dataloader) == 0.
+
+    configure_zenflow fills steps_per_epoch from the dataloader it owns, and the
+    config validator never sees that assignment.
+    """
+    config = ZenFlowConfig(select_strategy="auto", select_interval="auto", update_interval="auto")
+    engine = _StubEngine(config, training_dataloader=_Loader(0))
+
+    with caplog.at_level("WARNING"):
+        configure_zenflow(engine)
+
+    assert engine.select_interval == 0
+    assert any("never re-selected" in record.getMessage() for record in caplog.records)
+
+
+def test_step_strategy_needs_no_epoch_length():
+    config = ZenFlowConfig(select_strategy="step", select_interval=10, update_interval="auto")
+    engine = _StubEngine(config, training_dataloader=None)
+
+    configure_zenflow(engine)
+
+    assert engine.select_interval == 10
+    assert _select_boundaries(engine.select_interval, steps=40) == 4
+
+
+def test_auto_strategy_warns_rather_than_raising_on_an_ignored_interval():
+    # `raise Warning(...)` raises: Warning is an Exception. This combination
+    # could not run at all, and the `select_interval = 1` line after it was dead.
+    config = ZenFlowConfig(select_strategy="auto", select_interval=5, update_interval="auto")
+    engine = _StubEngine(config, training_dataloader=_Loader(10))
+
+    configure_zenflow(engine)
+
+    assert engine.select_interval == 10  # one epoch, not the ignored 5