Support compute dtype in spmd fsdp (#22)

AlibabaPAI · Sep 20, 2024 · 27eccca · 27eccca
1 parent de882ac
commit 27eccca
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/torchacc/dist/spmd_fsdp.py b/torchacc/dist/spmd_fsdp.py
@@ -49,6 +49,12 @@ def fsdp(self, model: torch.nn.Module, config: Config):
             transformer_layer_cls=layer_cls,
         )
 
+        dtype = torch.float32
+        if config.compute.fp16:
+            dtype = torch.float16
+        if config.compute.bf16:
+            dtype = torch.bfloat16
+
         auto_wrapper_callable = None
         if config.memory.gc and (config.memory.gc_cls
                                  == config.dist.fsdp.wrap_layer_cls):
@@ -69,6 +75,7 @@ def auto_wrapper_callable(m, *args, **kwargs):
             model,
             mesh,
             shard_output=self.shard_output_callable,
+            compute_dtype=dtype,
             auto_wrap_policy=auto_wrap_policy,
             auto_wrapper_callable=auto_wrapper_callable)
         return model