disable FSDP mixed precision for model buffers

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/585

Disable FSDP mixed precision for model buffers. Buffers are usually small in size so there's very limited performance gain for enabling mixed precision. Plus, applications like BatchNorm layers and diffusion models are very sensitive to the precision of buffers. Thus, we stick to full precision for buffers in FSDP.

Reviewed By: wat3rBro

Differential Revision: D46951673

fbshipit-source-id: 12bb1a47fbd8b3dd85c7f781bab707206044af15

d2go/trainer/fsdp.py

@@ -297,7 +297,7 @@ class FSDPModelingHook(ModelingHook):
             use_backward_prefetch=self.cfg.FSDP.BACKWARD_PREFETCH,
             param_dtype=precision_dtype,
             reduce_dtype=precision_dtype,
-            buffer_dtype=precision_dtype,
+            buffer_dtype=None,
             amp_autocast_dtype=precision_dtype,
             use_local_state_dict=self.cfg.FSDP.USE_LOCAL_STATE_DICT,
             load_local_state_dict=self.cfg.FSDP.USE_LOCAL_STATE_DICT,