Fix the dead grad_accum branch in the ZenFlow gradient copy

[vineethsaivs]

ZenFlowZeroOptimizerParallel.async_inplace_copy_grad_to_fp32_buffer_from_gpu tests grad_accum is None and then dereferences the None it just tested for:

grad_accum = self.get_param_gradient_attribute(param)
if grad_accum is None:
    src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements)
else:
    src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements)

Both arms are the same expression, so the branch does nothing, and the arm that is supposed to handle the missing gradient is the one that cannot run: it raises AttributeError: 'NoneType' object has no attribute 'view' from inside the copy instead of the assertion the base ZeRO-1/2 method declares.

Where it came from

This is not a new defect, it is a copy of one that was already fixed. DeepSpeedZeroOptimizer.async_inplace_copy_grad_to_fp32_buffer_from_gpu carried the identical four lines until Coverity flagged them, and 1a8ad24 ("fix issues raised by Coverity scans", #7431, 2025-08-02) replaced them with:

grad_accum = self.get_param_gradient_attribute(param)
assert grad_accum is not None

src_tensor = grad_accum.view(-1).narrow(0, source_offset, num_elements)

deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py was added on 2025-08-15 by #7391, thirteen days later, and its override was written against the pre-#7431 base, so it reintroduced the branch. It has been there since. This PR applies #7431's own resolution to the override, so both copies of the method now state the same contract.

Scanning the rest of deepspeed/runtime/zenflow/ turns up no other branch of this shape.

Test

test_async_inplace_copy_grad_requires_a_gradient in tests/unit/runtime/zenflow/test_zf.py, a plain function alongside the existing test_num_selected_columns_has_nonzero_floor. It drives the override with a stub whose gradient attribute is unset, which is the only way to reach the branch, and needs no accelerator.

$ pytest tests/unit/runtime/zenflow/test_zf.py -k "num_selected_columns or requires_a_gradient" -q
4 passed, 79 deselected

# without the source change
E           AttributeError: 'NoneType' object has no attribute 'view'
deepspeed/runtime/zenflow/zenflow_stage_1_and_2.py:689: AttributeError
1 failed, 82 deselected

yapf --diff and flake8 are clean on both files.

[ebarkhordar]

Static review, no GPU here, so I parsed rather than ran it.

One question about assert on the copy_grads_in_partition path. ZenFlowZeroOptimizerParallel overrides neither set_norm_for_param_grad_in_gpu nor update_offload_overflow_tracker_for_param_grad, so stage_1_and_2.py:1672-1676 runs three calls back to back on the same param and all three read get_param_gradient_attribute(param):

• :1589 treats None as reachable and falls back to param.grad
• :1509 guards with if grad is not None
• the method you patch now asserts it is not None

At the other call site the assert is clearly safe: _finalize_cpu_offload_gradient_accumulation calls _restore_cpu_offload_grad_to_gpu first, which assigns the attribute at :1024, and only then runs the same three at :1007-1009. I could not reach that conclusion for copy_grads_in_partition, where nothing between the reduction and :1676 sets it.

So are the fallback at :1589 and the guard at :1509 dead on that path, or can grad_accum still be None there for the stage 1 plus separate accumulation dtype case :363-366 enables? The disagreement predates your PR, so this is not a request to change the diff.

On your closing line: parsing all 658 files under deepspeed/ for if/else nodes whose arms unparse identically leaves only writer_factory.py:22, stage_1_and_2.py:1731 and pipe/module.py:418, none of which dereference what they test.

[vineethsaivs]

Good question, and the answer turns out to be that the three do not disagree about reachability. Two of them assume grad_accum is not None, one loudly and one silently, and the third opts out of the work rather than tolerating it.

The :1589 fallback is not a safety net in either mode. I extracted the four real methods with ast and drove them with a stub self:

A)  use_grad_accum_attribute=True, after _fill_param_grad_accum_attribute
      param.grad=None  grad_accum=set      set_norm ok    tracker ok   async_copy ok

A') same mode, grad_accum forced to None   set_norm AttributeError: 'NoneType' object has no attribute 'view'
                                           tracker ok    async_copy AssertionError

B)  use_grad_accum_attribute=False         get_param_gradient_attribute(param) IS param.grad
    param.grad=None                        set_norm AttributeError: 'NoneType' object has no attribute 'view'
                                           tracker ok    async_copy AssertionError

In mode B the fallback is a tautology: get_param_gradient_attribute returns param.grad, so grad_accum is None is exactly param.grad is None, and accumulated_grad = param.grad re-assigns the same None before .view(-1). In mode A it is worse than dead, because _fill_param_grad_accum_attribute sets param.grad = None on its way out (:1139), so the one thing the fallback reaches for is the thing that was just cleared.

So :1589 cannot save that path. If grad_accum were ever None there it would raise AttributeError two lines later, three lines before the assert you are asking about would have raised AssertionError.

On whether it can be None on the copy_grads_in_partition path. I do not think it can, and the reason is one call earlier than the site you looked at. process_gradients (:1768) runs

if self.use_grad_accum_attribute:
    self._fill_param_grad_accum_attribute(param)

before the reduction that eventually reaches :1761/:1764, and it is called from grad_handling_hook, registered through register_grad_hook (utils/torch.py:40), which is register_post_accumulate_grad_hook on torch >= 2.1. So autograd has just written param.grad when it fires, the fill runs, and param.grad_accum is set. In the other mode grad_accum is param.grad, non-None for the same reason. Both branches arrive at :1672-1676 with a gradient.

Your reading of the _finalize_cpu_offload_gradient_accumulation site is right, and it is the same shape: something assigns the attribute first, _restore_cpu_offload_grad_to_gpu at :1024 there, the fill at :1770 here.

The :1509 guard is the only one that genuinely tolerates None, and it tolerates it by skipping the overflow check entirely. Worth its own look at some point: if it ever did fire, the step would proceed with local_overflow unset for that parameter, which is a wrong answer rather than a crash. I have not chased that.

For this diff specifically, none of it changes the risk. copy_grads_in_partition is inherited by ZenFlowZeroOptimizerParallel unchanged, as are set_norm_for_param_grad_in_gpu, update_offload_overflow_tracker_for_param_grad, process_gradients and _fill_param_grad_accum_attribute, so both classes reach :1676 by the identical route. The base has asserted there since 1a8ad24 (#7431). All this PR does is make ZenFlow raise the base's AssertionError on an input where it currently raises AttributeError from inside the copy, so it cannot introduce a failure the base does not already have.

The :1589 tautology is a real defect and predates both of us, but it is in the base class and on a different method, so I would rather file it separately than widen this PR. Happy to do that, or to leave it if you would rather not have another open item on this file.

On the scan: that matches mine exactly. identical_if.py and dupbranch.py over deepspeed/runtime/zenflow/ returned only this site, and over the whole tree the same three you list, writer_factory.py:22, stage_1_and_2.py:1731 and pipe/module.py:418, none of which dereference what they test. writer_factory.py:22 is additionally unreachable-with-None, since FastCheckpointEngine sets self._writer = None rather than constructing the factory when dp_writer_config is None, so its and self._data_parallel_writer is not None is dead on both arms.

[ebarkhordar]

That answers it, thanks. The fill in process_gradients at :1773 running before the reduction is the piece I had missed, and _fill_param_grad_accum_attribute clearing param.grad at :1139 on its way out does make the :1589 fallback dead the way you describe.

One thing to add, still static, no run on my side. process_gradients is not the only route into :1672-1676. Parsing the call sites across deepspeed/, copy_grads_in_partition has exactly two callers, both inside reduce_ipg_grads, and reduce_ready_partitions_and_remove_grads has six, of which two arrive without the per-param fill:

• reduce_gradients at :889, the not self.overlap_comm path
• engine.py:3236, the coalesced flush that your own early return at :1769 hands the work to

Both are still safe, but by a different mechanism than the one you describe. Each guards with if get_gradient_for_reduction(param) is None: continue, and that reads the same attribute the assert checks, param.grad_accum in mode A and param.grad in mode B. So the conclusion holds on all three routes; on two of them it is the caller's own None check doing the work rather than the fill.

On :1589, a separate issue is the right call. It is a different method on the base class and nothing about it needs this PR.

[vineethsaivs]

Checked your call-site analysis against current main (183c7f9) rather than taking it, and it holds. Counts first:

copy_grads_in_partition has exactly two callers, stage_1_and_2.py:1761 and :1764, both inside reduce_ipg_grads (which spans 1701-1765), so that matches.

reduce_ready_partitions_and_remove_grads has six. Splitting them by which optimizer they reach:

caller	class	fill runs first?
stage_1_and_2.py:889, in reduce_gradients (871-891), the not self.overlap_comm path	ZeRO-1/2	no
stage_1_and_2.py:1775, in process_gradients (1768-1775)	ZeRO-1/2	yes
engine.py:3213, in _flush_coalesced_reduction	ZeRO-1/2	no
stage3.py:1469, stage3.py:1501	ZeRO-3	n/a
engine.py:3224, in _flush_coalesced_reduction_zero3	ZeRO-3	n/a

So three of the six reach ZeroOptimizer.copy_grads_in_partition at all, and your two fill-less routes are the right two. One line-number nit: you cited engine.py:3236; on current main the ZeRO-1/2 flush call is :3213 and its guard is :3211. :3224 is the stage-3 flush, which iterates optimizer.fp16_groups and guards on param.grad is not None, so it never gets here.

The part worth pinning down is why those guards are a precondition for the assert rather than a coincidence. The two accessors bottom out on the same attribute under the same flag:

def get_gradient_for_reduction(self, param):        # what the caller guards read
    if self.use_grad_accum_attribute:
        return param.grad_accum.to(self.dtype) if param.grad_accum is not None else None
    else:
        return param.grad

def get_param_gradient_attribute(self, param):      # what the assert checks
    return param.grad_accum if self.use_grad_accum_attribute else param.grad

.to(self.dtype) only runs on the non-None branch, so it can neither create a None nor consume one. get_gradient_for_reduction(param) is None is therefore true exactly when get_param_gradient_attribute(param) is None, in both modes. That makes the guard the same test as the assert, one frame up.

Worth noting :889 spells it positively, grad_reduc = self.get_gradient_for_reduction(param) then if grad_reduc is not None:, rather than the is None: continue shape the engine uses. Same effect, just not greppable as one idiom, which is probably why it reads as an exception.

So: three routes into :1672-1676, all safe, by the fill on one and by the caller's own None check on the other two. Nothing here changes the diff.

Agreed on :1589 being its own issue. I will open one against the base class and link it here so it does not ride along on this PR.

[vineethsaivs]

Filed as #8371, with the mode-by-mode reproduction and the set_norm_for_param_grad sibling as the reference shape. Left this PR's scope alone.