Fix the dead grad_accum branch in the ZenFlow gradient copy

[vineethsaivs]

ZenFlowZeroOptimizerParallel.async_inplace_copy_grad_to_fp32_buffer_from_gpu overrides the ZeRO-1/2 copy and diverges from it in two places.

1. The dead grad_accum is None branch

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 arm meant to handle a 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 declares.

DeepSpeedZeroOptimizer.async_inplace_copy_grad_to_fp32_buffer_from_gpu carried the identical four lines until Coverity flagged them, and 1a8ad24 (#7431, 2025-08-02) replaced them with assert grad_accum is not None. zenflow_stage_1_and_2.py landed thirteen days later in #7391, written against the pre-#7431 base, so it reintroduced the branch.

2. Clearing the wrong gradient attribute

The base ends with self.clear_grad_attribute(param); the override ends with param.grad = None. Both read through get_param_gradient_attribute, so under use_grad_accum_attribute (ZeRO-1 with a grad_accum_dtype that differs from the parameter dtype) the copy consumes param.grad_accum and then clears param.grad, which _fill_param_grad_accum_attribute has already set to None. The consumed gradient stays live, and the next micro step's param.grad_accum.add_(...) accumulates onto it.

Verification

ZenFlow needs an accelerator, so the two methods were driven out of the source files with a stub:

	master	this PR
gradient attribute missing	AttributeError: 'NoneType' object has no attribute 'view'	AssertionError, as the base declares
param.grad_accum after the copy	tensor([1., 1., 1., 1.])	None

The copied buffer is [1.0, 1.0, 1.0, 1.0] either way, so the copy itself is unchanged.

Test

Two cases in tests/unit/runtime/zenflow/test_zf.py, plain functions next to the existing test_num_selected_columns_has_nonzero_floor, neither needing an accelerator: one asserts the missing-gradient contract, one asserts the consumed attribute is cleared. Both use the real get_param_gradient_attribute and clear_grad_attribute off DeepSpeedZeroOptimizer so the override is checked against the base behaviour rather than a copy of it.

yapf --diff and flake8 with the repo's own .style.yapf and .flake8 are clean on both files.

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

I have not run a full ZenFlow ZeRO-1 job with a separate grad_accum_dtype end to end, so part 2 rests on the base-class contract rather than on a measured loss curve. Happy to split it out if you would rather review the two separately.

[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. I traced it statically too (no GPU here), and I read None as unreachable on that path, so the assert holds.

copy_grads_in_partition is only called from the ipg-bucket drain at stage_1_and_2.py:1761 and :1764, so a param only gets there after it produced a gradient this backward. _fill_param_grad_accum_attribute (:1133) then sets grad_accum for every param whose param.grad is not None.

The stronger signal is inside the same block: async_accumulate_grad_in_cpu_via_gpu already dereferences grad_accum with no guard at :1547 and :1559, and it runs before the copy in every case except micro_step_id == 0 at the accumulation boundary. If None were reachable there, that would raise first.

So I read :1589s fallback and :1509s guard as defensive leftovers rather than live cases, which matches what you found. Leaving the diff as is.

[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.

[vineethsaivs]

Rebased on master and widened: the same override also clears the wrong gradient attribute.

The base ends with self.clear_grad_attribute(param); this override ends with param.grad = None. Both read through get_param_gradient_attribute, so under use_grad_accum_attribute (ZeRO-1 with a grad_accum_dtype that differs from the parameter dtype) the copy consumes param.grad_accum and then clears param.grad, which _fill_param_grad_accum_attribute has already set to None. The consumed gradient stays live, and the next micro step's param.grad_accum.add_(...) accumulates onto it.

Driving the two methods out of the source files with a stub, since ZenFlow itself needs an accelerator:

	master	this PR
gradient attribute missing	AttributeError: 'NoneType' object has no attribute 'view'	AssertionError, as the base declares
param.grad_accum after the copy	tensor([1., 1., 1., 1.])	None

The copied buffer is [1.0, 1.0, 1.0, 1.0] either way, so the copy itself is unchanged.

tests/unit/runtime/zenflow/test_zf.py now carries a case for each, both plain functions next to test_num_selected_columns_has_nonzero_floor and neither needing an accelerator. yapf --diff and flake8 with the repo's own .style.yapf and .flake8 are clean on both files.

I have not been able to run a full ZenFlow ZeRO-1 job with a separate grad_accum_dtype end to end, so the second half rests on the base-class contract rather than on a measured loss curve. Happy to drop it into its own PR if you would rather review the two separately.