fix mp policy forwarding in fsdp2 (#970)

Summary: Pull Request resolved: #970

Reviewed By: galrotem, anshulverma

Differential Revision: D69669442

fbshipit-source-id: fff1e475ab1a31fc3291ee249c21292af1fc0561

Author: JKSenthil

tests/utils/test_prepare_module_gpu.py

@@ -348,7 +348,7 @@ def _test_prepare_fsdp2_shard_all() -> None:
 
         module = SimpleModule()
         device = torch.device("cuda")
-        strategy = FSDP2Strategy(modules_to_shard="all")
+        strategy = FSDP2Strategy(modules_to_shard="all", mp_policy=torch.bfloat16)
         prepare_fsdp2(module, device, strategy)
 
         for submodule in module.modules():

torchtnt/utils/prepare_module.py

@@ -371,9 +371,9 @@ def prepare_fsdp2(
         fsdp_kwargs["offload_policy"] = CPUOffloadPolicy()
     if (mp_policy := strategy.mp_policy) is not None:
         if isinstance(mp_policy, MixedPrecisionPolicy):
-            fsdp_kwargs["mixed_precision"] = mp_policy
+            fsdp_kwargs["mp_policy"] = mp_policy
         else:
-            fsdp_kwargs["mixed_precision"] = MixedPrecisionPolicy(
+            fsdp_kwargs["mp_policy"] = MixedPrecisionPolicy(
                 param_dtype=mp_policy,
                 reduce_dtype=mp_policy,
                 output_dtype=mp_policy,