From 4ffc94ea188b7d08fc676f5fc83311d3b11ec291 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Tue, 8 Sep 2026 14:44:32 +0800 Subject: [PATCH 1/2] Say when ZenFlow will never re-select its important columns ZenFlow re-picks the top-k 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 the number of steps in an epoch. DeepSpeed reads that from engine.training_dataloader, which is only set when the caller hands its data to deepspeed.initialize(). A caller driving its own dataloader falls into `engine.select_interval = 0`, and is_zenflow_select_boundary() treats 0 as "never again": (micro_step - full_warm_up_rounds) == 0 or (self.select_interval != 0 and self.micro_step % self.select_interval == 0) so the columns chosen at the first step are the ones used for the whole run. No error, no warning, and the default configuration takes this path. Resolving the intervals directly, over 40 micro-steps: select_strategy select_interval epoch length known resolved selections auto (default) auto no 0 1 auto (default) auto yes 10 4 epoch 1 no 0 1 epoch 1 yes 10 4 step 10 n/a 10 4 The tests in tests/unit/runtime/zenflow/test_zf.py build their dataloader after deepspeed.initialize() and never pass it in, so every ZenFlow test -- including the auto and epoch parametrizations -- runs on select_interval 0. They assert that training does not crash, which it does not. Rather than refuse the configuration, which would force callers onto DeepSpeed's own dataloader, this keeps the run going and: - takes steps_per_epoch from the zenflow config when the user sets it, so a caller with its own dataloader can say what an epoch is; - warns, naming all three ways out, when neither source provides a length; - documents steps_per_epoch as user-settable rather than "should not be set by users". Also replaces three `raise Warning(...)` calls with logger.warning. Warning is an Exception, so select_strategy 'auto' with an integer select_interval raised instead of warning, and the `select_interval = 1` line after it was dead. tests/unit/runtime/zenflow/test_zf_select_interval.py covers all five rows above plus the warning; five of its eight cases fail on master. Signed-off-by: alanhuangyoo --- deepspeed/runtime/zenflow/engine.py | 23 +++- deepspeed/runtime/zenflow/engine_stage3.py | 11 +- deepspeed/runtime/zenflow/zenflow_config.py | 6 +- .../runtime/zenflow/zenflow_stage_1_and_2.py | 14 ++- .../zenflow/test_zf_select_interval.py | 110 ++++++++++++++++++ 5 files changed, 151 insertions(+), 13 deletions(-) create mode 100644 tests/unit/runtime/zenflow/test_zf_select_interval.py diff --git a/deepspeed/runtime/zenflow/engine.py b/deepspeed/runtime/zenflow/engine.py index 2236d097169b..eef144d07935 100644 --- a/deepspeed/runtime/zenflow/engine.py +++ b/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,22 @@ 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: + # `steps_per_epoch` may already be set by the user; otherwise it can only + # come from a dataloader DeepSpeed owns. + if zenflow_config.steps_per_epoch is None and 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) + if zenflow_config.steps_per_epoch is not None: + 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: diff --git a/deepspeed/runtime/zenflow/engine_stage3.py b/deepspeed/runtime/zenflow/engine_stage3.py index 1de39a618964..2bf3a34e7849 100644 --- a/deepspeed/runtime/zenflow/engine_stage3.py +++ b/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): @@ -51,6 +51,11 @@ def configure_zenflow(optimizer_z3, zenflow_config): if zenflow_config.steps_per_epoch is not None: 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: diff --git a/deepspeed/runtime/zenflow/zenflow_config.py b/deepspeed/runtime/zenflow/zenflow_config.py index 2d522e35d779..96f397327eaa 100644 --- a/deepspeed/runtime/zenflow/zenflow_config.py +++ b/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") diff --git a/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py b/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py index 29c702d6813d..18032201915f 100644 --- a/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py +++ b/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): @@ -139,6 +140,13 @@ def _configure_zenflow(self, zenflow_config): if zenflow_config.steps_per_epoch is not None: 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: diff --git a/tests/unit/runtime/zenflow/test_zf_select_interval.py b/tests/unit/runtime/zenflow/test_zf_select_interval.py new file mode 100644 index 000000000000..656b700e8e58 --- /dev/null +++ b/tests/unit/runtime/zenflow/test_zf_select_interval.py @@ -0,0 +1,110 @@ +# 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) + + +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 From ffb0860f0d02de5341a7290257fc787ec9064c81 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Wed, 9 Sep 2026 17:28:40 +0800 Subject: [PATCH 2/2] Refuse a steps_per_epoch that reaches the state the warning exists to announce Reported by @ebarkhordar on #8456. This PR documents steps_per_epoch as something a caller sets, and the guard read `is not None`, so 0 multiplied straight through to select_interval = 0 -- the 'selected once and never re-selected' state the warning is for, reached silently through the knob that documents it, because the warning only fires in the None branch. A negative value was worse: select_interval stayed negative and micro_step % -5 fired on a schedule nobody asked for. steps_per_epoch resolved selections/40 warning None 0 1 YES 0 0 1 NO -5 -5 8 NO 10 10 4 NO validate_fields already rejects an update_interval below 1 one field away; the same check on steps_per_epoch keeps 0 and negatives out, and the three consumers guard on truth rather than None so nothing else can carry a zero through. Adding the validator surfaced a second path to the same 0, which the report did not cover: configure_zenflow assigns len(training_dataloader), the config model validates on assignment, so an empty dataloader turned into a ValidationError raised from inside DeepSpeed. A crash is the wrong answer for a degenerate but legal dataloader, so the length is only assigned when positive and an empty one falls through to the warning. tests/unit/runtime/zenflow/: 105 passed on 2xH20, against a 56 passed / 38 skipped / 0 failed baseline on master -- the skips are the multi-GPU cases this run reached, plus the 3 tests added here. Signed-off-by: alanhuangyoo --- deepspeed/runtime/zenflow/engine.py | 11 ++++-- deepspeed/runtime/zenflow/engine_stage3.py | 2 +- deepspeed/runtime/zenflow/zenflow_config.py | 9 +++++ .../runtime/zenflow/zenflow_stage_1_and_2.py | 2 +- .../zenflow/test_zf_select_interval.py | 35 +++++++++++++++++++ 5 files changed, 54 insertions(+), 5 deletions(-) diff --git a/deepspeed/runtime/zenflow/engine.py b/deepspeed/runtime/zenflow/engine.py index eef144d07935..9c2ab57cd4e1 100644 --- a/deepspeed/runtime/zenflow/engine.py +++ b/deepspeed/runtime/zenflow/engine.py @@ -56,9 +56,14 @@ def configure_zenflow(engine: "DeepSpeedEngine") -> None: if select_strategy == 'epoch': # `steps_per_epoch` may already be set by the user; otherwise it can only # come from a dataloader DeepSpeed owns. - if zenflow_config.steps_per_epoch is None and engine.training_dataloader is not None: - zenflow_config.steps_per_epoch = len(engine.training_dataloader) - if zenflow_config.steps_per_epoch is not None: + 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 diff --git a/deepspeed/runtime/zenflow/engine_stage3.py b/deepspeed/runtime/zenflow/engine_stage3.py index 2bf3a34e7849..1a877bd6cb49 100644 --- a/deepspeed/runtime/zenflow/engine_stage3.py +++ b/deepspeed/runtime/zenflow/engine_stage3.py @@ -48,7 +48,7 @@ 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 " diff --git a/deepspeed/runtime/zenflow/zenflow_config.py b/deepspeed/runtime/zenflow/zenflow_config.py index 96f397327eaa..e6ebbb7b481d 100644 --- a/deepspeed/runtime/zenflow/zenflow_config.py +++ b/deepspeed/runtime/zenflow/zenflow_config.py @@ -65,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') diff --git a/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py b/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py index 18032201915f..bb2bff754f0c 100644 --- a/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py +++ b/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py @@ -137,7 +137,7 @@ 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 diff --git a/tests/unit/runtime/zenflow/test_zf_select_interval.py b/tests/unit/runtime/zenflow/test_zf_select_interval.py index 656b700e8e58..f5622964ab49 100644 --- a/tests/unit/runtime/zenflow/test_zf_select_interval.py +++ b/tests/unit/runtime/zenflow/test_zf_select_interval.py @@ -89,6 +89,41 @@ def test_no_epoch_length_selects_once_and_says_so(select_strategy, select_interv 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)