pre-backward.
class torch.distributed.fsdp.MixedPrecision(param_dtype=None, reduce_dtype=None, buffer_dtype=None, keep_low_precision_grads=False, cast_forward_inputs=False, cast_root_forward_inputs=True)
This configures FSDP-native mixed precision training.
Variables:
      * param_dtype (torch.dtype) -- This specifies the dtype
        for model parameters, inputs (when "cast_forward_inputs" or
        "cast_root_forward_inputsis set toTrue"), and therefore
        the dtype for computation. However, outside the forward and
        backward passes, parameters are in full precision. Model
        checkpointing always happens in full precision.
  * **reduce_dtype** (*torch.dtype*) -- This specifies the dtype
    for gradient reduction, which is permitted to differ from
    "param_dtype".

  * **buffer_dtype** (*torch.dtype*) -- This specifies the dtype
    for buffers. FSDP does not shard buffers, casts them to
