Error in gradient accumulation when there are unused parameters

I write a minimum reproducible case

import torch
from torch import nn
import deepspeed
import argparse

class Net(nn.Module):

    def __init__(self):
        super(Net, self).__init__()
        self.bias = nn.Parameter(torch.zeros(1))
        self.unused_param = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        return x + self.bias

def parse_args():
    parser = argparse.ArgumentParser("Test DeepSpeed gradient accumulation")
    parser.add_argument('--local_rank', type=int, default=-1)
    parser.add_argument('--train_batch_size', type=int, default=-1)
    parser.add_argument('--gradient_accumulation_steps', type=int, default=-1)
    parser = deepspeed.add_config_arguments(parser)
    args = parser.parse_args()
    return args

args = parse_args()
rank = args.local_rank
batch_size = args.train_batch_size // args.gradient_accumulation_steps
net = Net()
net_engine, optimizer, _, _ = deepspeed.initialize(args=args, model=net, model_parameters=net.parameters())

for i in range(100):
    with torch.cuda.amp.autocast():
        x = torch.rand((2, 3)).to(f'cuda:{rank}')
        loss = torch.mean(net(x))
        net_engine.backward(loss)
        net_engine.step()

with a simple deepspeed config json file

{
  "train_batch_size": 16,
  "gradient_accumulation_steps": 4,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.001,
      "betas": [
        0.9,
        0.999
      ],
      "eps": 1e-8,
      "weight_decay": 0
    }
  },
  "fp16": {
    "enabled": true
  },
  "zero_optimization": true
}

Running deepspeed --include="localhost:0,1,2,3" test_deepspeed_grad_accum.py --deepspeed --deepspeed_config test_deepspeed_grad_accum.json on single machine with multiple GPUs can get the following error:

File "PYTHON_INTERPRETER/site-packages/deepspeed/runtime/zero/stage2.py", line 728, in reduce_independent_p_g_buckets_and_remove_grads
    new_grad_tensor.copy_(param.grad.view(-1))
AttributeError: 'NoneType' object has no attribute 'view'

Unused parameters do not cause error if DeepSpeed is not used.

Start with the minimum reproducible case and the failing path in deepspeed/runtime/zero/stage2.py at line 728, then run the provided multi-GPU DeepSpeed command with the sample configuration. Trace how unused parameters are handled during gradient accumulation. Done means the reproduction no longer raises an AttributeError when a parameter has no gradient.