[BUG] set_norm_for_param_grad_in_gpu's grad_accum fallback is dead in both accumulation modes

[vineethsaivs]

Describe the bug

DeepSpeedZeroOptimizer.set_norm_for_param_grad_in_gpu (deepspeed/runtime/zero/stage_1_and_2.py:1585) falls back to param.grad when the gradient attribute is None:

def set_norm_for_param_grad_in_gpu(self, param):
    param_id = self.get_param_id(param)
    grad_accum = self.get_param_gradient_attribute(param)
    if grad_accum is None:
        accumulated_grad = param.grad
    else:
        accumulated_grad = grad_accum
    ...
    accumulated_grad = accumulated_grad.view(-1).narrow(0, start, num_elements)

That fallback cannot supply a usable tensor in either mode, so instead of rescuing the None case it defers it by one line, where .view(-1) raises AttributeError: 'NoneType' object has no attribute 'view' rather than a diagnosable error.

Why the fallback is dead

use_grad_accum_attribute is True only for ZeRO stage 1 with a separate accumulation dtype, and get_param_gradient_attribute is:

def get_param_gradient_attribute(self, param):
    return param.grad_accum if self.use_grad_accum_attribute else param.grad

• use_grad_accum_attribute=False. get_param_gradient_attribute is param.grad, so grad_accum is None is exactly param.grad is None. The fallback re-assigns the same None. It is a tautology.

• use_grad_accum_attribute=True. It is worse than dead. _fill_param_grad_accum_attribute (:1133) moves the gradient into grad_accum and then clears the source:

  def _fill_param_grad_accum_attribute(self, param):
      if param.grad is not None:
          ...
          param.grad = None

  So the one state the fallback is written for, grad_accum unset while param.grad holds the gradient, does not survive the fill. After it runs, param.grad is None by construction, and the fallback reaches for what was just cleared.

To Reproduce

No accelerator needed. Extracting the three methods with ast and driving them against a stub parameter reproduces it directly:

mode A (use_grad_accum_attribute=True), grad_accum unset and param.grad set,
      i.e. exactly what the fallback is written to rescue:
   before fill: get_param_gradient_attribute -> None
   after  fill: param.grad = None | grad_accum = T(real)

mode A, grad_accum None at the call (fill did not run or found no grad):
   set_norm_for_param_grad_in_gpu -> AttributeError: 'NoneType' object has no attribute 'view'

mode B (use_grad_accum_attribute=False), param.grad None:
   get_param_gradient_attribute IS param.grad -> True
   set_norm_for_param_grad_in_gpu -> AttributeError: 'NoneType' object has no attribute 'view'

Expected behavior

Either the branch goes, since grad_accum is the only source in both modes:

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

or, if a None gradient is meant to be tolerated here, the method should return early rather than fall through to .view(-1).

The sibling immediately above it, set_norm_for_param_grad at :1572, already takes the first shape and has no branch at all, which is what makes this look like a leftover rather than a deliberate difference.

Context

This came up while reviewing #8360, which fixes the same grad_accum is None shape in ZenFlow's override of async_inplace_copy_grad_to_fp32_buffer_from_gpu. That one is a copy of the branch commit 1a8ad24f (#7431) removed from the base class after Coverity flagged it. set_norm_for_param_grad_in_gpu is a different method on the base class and nothing about it needs that PR, so filing separately at the reviewer's suggestion rather than widening #8360.

Whether it is reachable in practice is a separate question from whether the branch does anything: on the three routes into copy_grads_in_partition the gradient is non-None either because process_gradients fills it first or because the caller guards with get_gradient_for_reduction(param) is None. So this is a latent defect and a misleading error rather than a live crash. Happy to send a PR for whichever shape maintainers prefer.

System info

• Read against main at 183c7f9. Static analysis plus the extracted-method reproduction above, so no ds_report or GPU configuration applies.