This is an edge case when using gradient accumulation steps with FSDP2.

If a parameter within a parameter group doesnt have gradients for all but the final backwards pass, the grad is not casted to the reduce dtype resulting in inconsistent dtype between gradients for reduce_scatter.

Specifically, if using a MixedPrecisionPolicy with param_dtype=torch.bfloat16 and reduce_dtype=torch.float32, parameters that had grads for both steps will have gradients of dtype torch.float32, but those that have a gradient only on the final pass will have gradients of dtype torch.bfloat16. This raises:

AssertionError: FSDP reduce-scatter expects uniform gradient dtype but got {torch.bfloat16, torch.float32}

This pr explicitly casts the grad to the correct dtype for the given case.

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta