Say when ZenFlow will never re-select its important columns

[alanhuangyoo]

Problem

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 it from engine.training_dataloader, which is only set when the caller hands its data to deepspeed.initialize():

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)
    else:
        engine.select_interval = 0

and is_zenflow_select_boundary() reads 0 as "never again":

return self.zenflow and (self.micro_step - self.full_warm_up_rounds) >= 0 and (
    (self.micro_step - self.full_warm_up_rounds) == 0 or
    (self.select_interval != 0 and self.micro_step % self.select_interval == 0))

So a caller that drives its own dataloader — the common case — gets the columns chosen at the first step 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 select_interval	selections in 40 steps
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

Two things I want to flag rather than bury:

• The existing tests run on this path. tests/unit/runtime/zenflow/test_zf.py::run_training_distributed builds its dataloader after deepspeed.initialize() and never passes it in, so every ZenFlow test — including the ("auto", "auto", "auto") and ("epoch", 1, 4) parametrizations — resolves select_interval to 0. They assert that training does not crash, which it does not, so the suite cannot see this.
• raise Warning(...) raises. Warning is a subclass of Exception, so select_strategy: "auto" together with an integer select_interval did not warn — it aborted, with a message that describes a fallback ("would be overwritten"), and the select_interval = 1 line right after it was unreachable. Three copies, in engine.py, engine_stage3.py and zenflow_stage_1_and_2.py.

Solution

My first attempt raised a ValueError for this configuration. That was wrong, and this repo's own tests said so: forcing callers onto training_data= also hands them DeepSpeed's dataloader, which builds 2 * device_count worker processes — 36 of the 56 ZenFlow tests then died on daemonic processes are not allowed to have children. Imposing a dataloader is not a bug fix.

So the run keeps going, and instead:

• steps_per_epoch is taken from the zenflow config when the user sets it, so a caller with its own dataloader can say what an epoch is. The field already existed and was already settable; it was documented as "initialized during execution and should not be set by users", which is exactly the gap.
• When neither source provides a length, logger.warning says the columns will be selected once and never re-selected, and names all three ways out.
• Three raise Warning(...) become logger.warning.

Nothing changes for a caller that already passes training_data=.

Verification

Full ZenFlow suite, 1×H20 (DS_SKIP_CUDA_CHECK=1 for the CPU-Adam JIT build):

master   : 56 passed, 38 skipped, 0 failed
this PR  : 64 passed, 38 skipped, 0 failed

64 − 56 = the 8 tests added here; no existing test changes state.

tests/unit/runtime/zenflow/test_zf_select_interval.py drives the real configure_zenflow against a stub engine, so it needs no GPU and no CPU-Adam build. It covers every row of the table, the user-supplied steps_per_epoch path, and the warning. Five of its eight cases fail on master.

pre-commit run passes on all five files.

Note

I have not touched the is_zenflow_select_boundary predicate or the meaning of select_interval == 0. If you would rather the epoch strategy refuse to start without a length, that is a smaller change than this one — I went with the non-breaking side because the current default lands there and refusing would break every existing run.

[ebarkhordar]

I ran this against 4ffc94ea in a clean container (python:3.12-slim, CPU torch), driving the real configure_zenflow the way test_zf_select_interval.py does. Your 8 new cases pass there.

This PR makes steps_per_epoch something a caller is now told to set, and the guard on the next line reads is not None, so steps_per_epoch: 0 multiplies straight through to select_interval = 0. That is the same "selected once and never re-selected" state this PR exists to announce, reached silently through the knob it documents, because the new warning only fires in the None branch.

$ docker run --rm -v $PWD:/src:ro -w /src -e PYTHONPATH=/src ds-hunt:cpu python zf_probe.py
deepspeed from: /src/deepspeed/__init__.py
steps_per_epoch unset (None)     -> select_interval=0     selections_in_40_steps=1  warning=YES
steps_per_epoch=0                -> select_interval=0     selections_in_40_steps=1  warning=NO
steps_per_epoch=-5               -> select_interval=-5    selections_in_40_steps=8  warning=NO
steps_per_epoch=10               -> select_interval=10    selections_in_40_steps=4  warning=NO

All three copies guard the same way, so a zero passes all of them: engine.py:61, engine_stage3.py:51, zenflow_stage_1_and_2.py:140.

validate_fields already rejects an update_interval below 1 one field away. The same check on steps_per_epoch, or gt=0 on this Field, keeps the zero case in the branch that warns and costs one line.

I only exercised configure_zenflow, so I have not run the rest of the ZenFlow suite here.

[alanhuangyoo]

Fixed in ffb0860 — sorry for the silence, I pushed it the same day and never replied here.

validate_fields now rejects it one field away from the update_interval check you pointed at:

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")

and all three consumers moved from is not None to a truth guard, so None and a rejected value cannot take different paths through them.

Your probe against the current head:

steps_per_epoch          result
  unset (None)           accepted, steps_per_epoch=None      -> the branch that warns
  0                      ValidationError at config validation
  -5                     ValidationError at config validation
  10                     accepted, steps_per_epoch=10

The negative case was the worse of the two and I would not have looked for it: select_interval stayed negative and micro_step % -5 fired on a schedule nobody asked for — 8 selections in 40 steps rather than the 4 that 10 gives. Your table is in the commit message.

Full suite on 2×H20 after the change: tests/unit/runtime/zenflow/ 105 passed, against a 56 passed / 38 skipped baseline on one GPU.

[alanhuangyoo]

You were right, and I reproduced your table before changing anything — same four rows, same numbers.

Fixed in ffb0860: validate_fields now rejects a steps_per_epoch below 1, next to the update_interval check you pointed at, and the three consumers guard on truth rather than is not None so nothing else can carry a zero through.

Adding the validator surfaced a second path to the same 0 that your report did not cover, and I would rather say so than let it pass: configure_zenflow assigns len(engine.training_dataloader), and this 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. That case has its own test.

tests/unit/runtime/zenflow/: 105 passed on 2×H20, against a 56 passed / 38 skipped / 0 failed baseline on master — the difference is the multi-GPU cases this run reached plus the 3 tests added here.

Worth naming the pattern, since this is the second time you have caught the same shape in my work: on #8435 the guard was per tensor and the invariant I wrote was global, and here I documented a knob and only guarded one of its values. Both times the fix was verified against the path that motivated it and described in wider language than that path supports. Your habit of sweeping the domain and printing the table is the thing that catches it, and it is cheap — I am adopting it for anything I add or document rather than waiting for review to do it.

[ebarkhordar]

Thanks. I read the diff at ffb0860 rather than re-running it: the validator rejects 0 and negatives one field away from the update_interval check, and with all three consumers on a truth guard there is no path left that carries a zero into select_interval.

The empty-dataloader case was not in my report. Assigning len(dataloader) into a model that validates on assignment turns a legal-if-degenerate input into a ValidationError raised from inside DeepSpeed, and falling through to the warning is the right answer there.