Allow FSDP to use with torch.autocast for bfloat16 mixed precision (#…

…2033) * Ignore native_amp when FSDP is used * Rollback condition * Fix mixed precision of bfloat16 for FSDP
huggingface · Oct 6, 2023 · 5ae6111 · 5ae6111
1 parent 230a5f5
commit 5ae6111
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
@@ -1458,6 +1458,7 @@ def get_mixed_precision_context_manager(native_amp: bool = False, autocast_kwarg
             DistributedType.MULTI_GPU,
             DistributedType.MULTI_NPU,
             DistributedType.MULTI_XPU,
+            DistributedType.FSDP,
         ]:
             return torch.autocast(device_type=state.device.type, dtype=torch.bfloat16, **autocast_kwargs)
         else: