Skip to content

fix(zenflow): restore assert and remove duplicated identical branch in gradient copy - #8445

Open
simpleqt wants to merge 1 commit into
deepspeedai:masterfrom
simpleqt:fix/zenflow-grad-accum-noner
Open

fix(zenflow): restore assert and remove duplicated identical branch in gradient copy#8445
simpleqt wants to merge 1 commit into
deepspeedai:masterfrom
simpleqt:fix/zenflow-grad-accum-noner

Conversation

@simpleqt

@simpleqt simpleqt commented Sep 6, 2026

Copy link
Copy Markdown

In deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py (ZenFlow overlap path), the fork of async_inplace_copy_grad_to_fp32_buffer_from_gpu mangled the original logic into two byte-identical if/else arms:

grad_accum = self.get_param_gradient_attribute(param)
if grad_accum is None:
    src_tensor = grad_accum.view(-1).narrow(...)  # AttributeError on None!
else:
    src_tensor = grad_accum.view(-1).narrow(...)

If the condition were ever true, the None branch itself calls .view() on None. Restored the upstream shape from runtime/zero/stage_1_and_2.py (assert grad_accum is not None + single assignment). No test covers this method (grep -rn async_inplace_copy tests/ is empty), which is why the mangle survived.

Signed-off-by: simpleqt 89645338+simpleqt@users.noreply.github.com

…n gradient copy

The ZenFlow fork of async_inplace_copy_grad_to_fp32_buffer_from_gpu
mangled the original logic into two byte-identical if/else arms. If
grad_accum were ever None, the None branch would itself call
grad_accum.view(-1) and raise AttributeError, defeating the check.
Restored the upstream assert + single-assignment shape from
runtime/zero/stage_1_and_2.py.

Signed-off-by: simpleqt <89645338+simpleqt@users.noreply.github.com>
Copilot AI lite review requested due to automatic review settings September 6, 2026 17:34

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants