Exception with optimization_stage_3

Using CUDA 11.1, pytorch 1.8.1 and Deepspeed 0.3.14.

Model trains with FP16 and optimization_stage 2, but fails with optimization_stage 3 with the following exception:

b' model_engine.backward(loss)\n'
b' File "/opt/miniconda/lib/python3.7/site-packages/deepspeed/runtime/engine.py", line 997, in backward\n'
b' self.optimizer.backward(loss)\n'
b' File "/opt/miniconda/lib/python3.7/site-packages/deepspeed/runtime/zero/stage3.py", line 2555, in backward\n'
b' self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)\n'
b' File "/opt/miniconda/lib/python3.7/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 53, in backward\n'
b' scaled_loss.backward(retain_graph=retain_graph)\n'
b' File "/opt/miniconda/lib/python3.7/site-packages/torch/tensor.py", line 245, in backward\n'
b' torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)\n'
b' File "/opt/miniconda/lib/python3.7/site-packages/torch/autograd/init.py", line 147, in backward\n'
b' allow_unreachable=True, accumulate_grad=True) # allow_unreachable flag\n'
b' File "/opt/miniconda/lib/python3.7/site-packages/torch/autograd/function.py", line 89, in apply\n'
b' return self._forward_cls.backward(self, *args) # type: ignore\n'
b' File "/opt/miniconda/lib/python3.7/site-packages/deepspeed/runtime/zero/linear.py", line 85, in backward\n'
b' grad_weight = grad_output.t().matmul(input)\n'
b'RuntimeError: expected scalar type Half but found Float\n'

Start with the failing operation at deepspeed/runtime/zero/linear.py:85, then trace the stage-3 backward path through stage3.py and fp16/loss_scaler.py using the reported CUDA, PyTorch, and DeepSpeed versions. Reproduce the Half-versus-Float failure and verify that optimization stage 3 completes backward without the dtype mismatch.