Reject invalid ZenFlow ratio and update interval boundaries

[tandede]

Summary

• require zenflow.topk_ratio to be strictly between 0 and 1
• require numeric zenflow.update_interval values to be at least 1
• add regression coverage for the invalid ratio endpoints and non-positive update intervals

Why

ZenFlow's automatic update path normalizes selected gradients by topk_ratio and unselected gradients by 1 - topk_ratio. The explicit update path selects its overlap buffer using micro_step // update_interval.

The previous validation accepted ratio endpoints as well as zero and negative numeric update intervals. These values could therefore trigger division by zero or invalid scheduling after training had already initialized. Rejecting them during configuration parsing provides an immediate validation error while preserving all usable configurations.

Testing

• python -m pytest -q tests/unit/runtime/zero/test_zero_config.py tests/unit/runtime/zenflow/test_zf_config.py tests/unit/runtime/zenflow/test_zf.py::test_split_affinity — 21 passed
• pre-commit run --files deepspeed/runtime/zenflow/zenflow_config.py tests/unit/runtime/zenflow/test_zf_config.py

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[sfc-gh-truwase]

@Antlera can you please help review?

[ebarkhordar]

I ran the config at head 09641d5 in a clean python:3.11-slim container. Both endpoints are rejected now. The other divisor in that same expression still has no bound, and it is in this config class:

topk_ratio=0.0: rejected (ValidationError)
topk_ratio=1.0: rejected (ValidationError)
update_interval=0: ACCEPTED, value=0
update_interval=-5: ACCEPTED, value=-5

An explicit integer sets auto_update = False (zenflow_stage_1_and_2.py:129-130), and that is the branch which divides: (self.micro_step // self.update_interval) & 1, at zenflow_stage_1_and_2.py:658 and engine_stage3.py:451. full_warm_up_rounds defaults to 0, so get_overlap_step_state reaches it on the first optimizer step. The topk_ratio divisions you fixed sit inside if self.auto_update:, so yours is the auto path and this is the explicit one.

Field bounds are awkward on a Union[str, int], but validate_fields already checks this field and one clause covers it:

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

I only constructed configs. No GPU on this box, so I did not run a training step and have not watched the error happen. Happy for it to be out of scope here.

[tandede]

Thanks for catching this. The explicit numeric update path does divide by update_interval, so zero and negative values are invalid for the same reason.

I added configuration validation requiring numeric update_interval values to be at least 1, together with regression coverage for 0 and -5. The related configuration and ZenFlow tests now report 21 passed, and the pre-commit checks pass.

[ebarkhordar]

Confirmed at 2f74e7e in a clean container: update_interval now rejects 0, -5, "0" and 0.0, while 1 and "auto" still construct. Thanks for the quick turnaround.

[tandede]

Confirmed at 2f74e7e in a clean container: update_interval now rejects 0, -5, "0" and 0.0, while 1 and "auto" still construct. Thanks for the quick turnaround.

Thanks for re-testing and confirming the fix! Really appreciate the thorough review.

[Antlera]

Hi @tandede. Thanks for this nice PR. I think most of the code is LGTM.

One more related comment, if you are willing to take it in this PR.

Rejecting topk_ratio in {0, 1} fixes the auto-path divisors, but a legal ratio can still collapse to zero selected columns on ZeRO-1/2:

num_select = int(self.topk_ratio * (end_column - start_column))

int() truncates toward zero, so e.g. topk_ratio=0.01 and 50 columns
yields 0. That later hits torch.topk(..., k=0). ZeRO-3 already guards
this with max(1, int(...)) in engine_stage3.py.

Please apply the same floor in zenflow_stage_1_and_2.py, and keep the
buffer sizes in lockstep — do not change only the two num_select lines:

• update_selected_channels: num_select
• _process_selected_fp32_groups_grad: num_select
• average_tensor: curr_column_size (index_buffer length)
• average_tensor: curr_selected_reduce_size (grad_buffer length)

Selection is per partition, so sizing with int(num_col * topk_ratio) on the full parameter can still under-allocate after a per-partition max(1, ...). Use the same per-partition formula for allocation and
selection. Skip the floor when the column/element count is already 0.

Happy for this to stay out of scope if you prefer a follow-up.

[tandede]

Hi @tandede. Thanks for this nice PR. I think most of the code is LGTM.

One more related comment, if you are willing to take it in this PR.

Rejecting topk_ratio in {0, 1} fixes the auto-path divisors, but a legal ratio can still collapse to zero selected columns on ZeRO-1/2:

num_select = int(self.topk_ratio * (end_column - start_column))

int() truncates toward zero, so e.g. topk_ratio=0.01 and 50 columns yields 0. That later hits torch.topk(..., k=0). ZeRO-3 already guards this with max(1, int(...)) in engine_stage3.py.

Please apply the same floor in zenflow_stage_1_and_2.py, and keep the buffer sizes in lockstep — do not change only the two num_select lines:

• update_selected_channels: num_select
• _process_selected_fp32_groups_grad: num_select
• average_tensor: curr_column_size (index_buffer length)
• average_tensor: curr_selected_reduce_size (grad_buffer length)

Selection is per partition, so sizing with int(num_col * topk_ratio) on the full parameter can still under-allocate after a per-partition max(1, ...). Use the same per-partition formula for allocation and selection. Skip the floor when the column/element count is already 0.

Happy for this to stay out of scope if you prefer a follow-up.

Thanks, this is a good point. I addressed it in 5b5beb6.

I added a shared per-partition selection-count helper so that every non-empty partition uses max(1, int(num_columns * topk_ratio)), while a partition with zero complete columns still selects zero.

Both num_select paths now use this calculation, and the index-buffer and gradient-buffer sizes use the same per-partition formula, keeping their allocations in lockstep with selection.

I also added regression coverage for the topk_ratio=0.01 and 50-column case, the zero-column case, and a case selecting more than one column. The changed-file pre-commit checks pass, and the targeted ZenFlow/config tests report 18 passed.

Thanks for calling out the buffer-sizing requirement as well.

[tohtana]

Thank you for your contribution, @tandede!
The code looks good to me. I will approve after @Antlera gives the greenlight.

[Antlera]

Thank you for your contribution, @tandede! The code looks good to me. I will approve after @Antlera gives the greenlight.

Hi @tohtana, thanks for following up on this, will approve the changes once the positive case as I commented is added.

[tandede]

Hi @tohtana, thanks for following up on this, will approve the changes once the positive case as I commented is added.

Thanks for the clarification. I added an end-to-end positive regression test in 1bdca9f.

The test runs the actual ZeRO-1/2 distributed training path with two ranks, a 50-column parameter, and topk_ratio=0.01. This exercises the per-partition nonzero selection floor together with the corresponding index and gradient buffer allocations.

The focused local tests report 13 passed and 2 GPU-dependent cases skipped. The GPU cases are now queued in CI.

[Antlera]

Hi @tohtana, thanks for following up on this, will approve the changes once the positive case as I commented is added.

Thanks for the clarification. I added an end-to-end positive regression test in 1bdca9f.

The test runs the actual ZeRO-1/2 distributed training path with two ranks, a 50-column parameter, and topk_ratio=0.01. This exercises the per-partition nonzero selection floor together with the corresponding index and gradient buffer allocations.

The focused local tests report 13 passed and 2 GPU-dependent cases skipped. The GPU cases are now queued in CI.

Thanks. LGTM.

[Antlera]

LGTM

[Antlera]

Thank you for your contribution, @tandede! The code looks good to me. I will approve after @Antlera gives the greenlight.

Hi @tohtana just approved after the latest commit.