Add condition when dimension is greater than 2 (#4390)

When dimension is greater than 2, grad_bias should be calculated differently.

deepspeed/runtime/zero/linear.py

@@ -94,7 +94,10 @@ def backward(ctx, grad_output):
             #print(f"Computed grad weight grad_weight {grad_weight.shape}")
         if bias is not None and ctx.needs_input_grad[2]:
             #print("Computing grad bias")
-            grad_bias = grad_output.sum(0)
+            if dim > 2:
+                grad_bias = grad_output.sum([i for i in range(dim - 1)])
+            else:
+                grad_bias = grad_output.sum(0)
             #print("Done computing grad bias")
             #print("needs bias")
         #print(f"backward shaped grad_input {grad_input.shape}, grad_weight {grad_weight.shape}, grad_bias {grad_bias.shape if grad_bias is not None else None}")