diff --git a/deepspeed/runtime/zenflow/engine.py b/deepspeed/runtime/zenflow/engine.py index 2236d097169b..9c2ab57cd4e1 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,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: diff --git a/deepspeed/runtime/zenflow/engine_stage3.py b/deepspeed/runtime/zenflow/engine_stage3.py index 1de39a618964..1a877bd6cb49 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): @@ -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: diff --git a/deepspeed/runtime/zenflow/zenflow_config.py b/deepspeed/runtime/zenflow/zenflow_config.py index 2d522e35d779..e6ebbb7b481d 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") @@ -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') diff --git a/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py b/deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py index 29c702d6813d..bb2bff754f0c 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): @@ -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: 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..f5622964ab49 --- /dev/null +++ b/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