Fix fp16 ZeRO stage-0 loss scaling in engine.backward()

When fp16 is enabled with ZeRO stage 0, engine.backward() falls through
to loss.backward() directly. FP16_Optimizer.step() always divides
gradients by cur_scale (typically 32768), so gradients that were never
pre-scaled are unscaled by the same factor, making effective parameter
updates ~cur_scale times too small and stalling training.

Add an explicit elif branch for fp16 + ZeRO-0 that routes through
FP16_Optimizer.backward(), which pre-scales the loss by cur_scale before
calling backward(), matching what step() expects.

The existing # TODO comment on line 3325 acknowledges this path was
left incomplete. The ZeRO stages > 0 and AMP paths are handled correctly;
only the fp16 non-ZeRO path was missing.

Tested with a CIFAR-10 MoE benchmark (fp16, ZeRO stage 0, 8 experts):
- Before: accuracy 22-28%, loss stalls at ~1.8
- After:  accuracy 49%, loss converges to ~0.7

Signed-off-by: Roohollah Etemadi <Roohollah.Etemadi@amd.com>

Author: etemadiamd

deepspeed/runtime/engine.py

@@ -3325,6 +3325,16 @@ def backward(self, loss, retain_graph=False, scale_wrt_gas=True):
             # TODO: handle these scaling with direct calls to loss.backward()
             if isinstance(self.optimizer, ZeROOptimizer):
                 loss = self.optimizer.scale_if_loss(loss)
+            elif self.fp16_enabled() and not self.zero_optimization():
+                # fp16 + ZeRO stage 0: FP16_Optimizer.step() always divides gradients
+                # by cur_scale, so loss must be pre-scaled before backward. Calling
+                # loss.backward() directly makes effective gradient updates ~cur_scale
+                # times too small, stalling training. Route through
+                # FP16_Optimizer.backward() which applies the scaling correctly.
+                self.optimizer.backward(loss, **backward_kwargs)
+                self._backward_epilogue()
+                self._running_engine_backward = False
+                return gas_scaled_loss
             elif self.torch_autocast_z0_gradscaler:
                 loss = self.torch_autocast_z0_gradscaler.scale(loss)