Igor Kasianenko 2 mesi fa
parent
commit
b3992f326d
1 ha cambiato i file con 2 aggiunte e 0 eliminazioni
  1. 2 0
      src/llama_cookbook/utils/fsdp_utils.py

+ 2 - 0
src/llama_cookbook/utils/fsdp_utils.py

@@ -7,6 +7,8 @@ import torch.distributed as dist
 
 
 from llama_cookbook.policies import fpSixteen,bfSixteen, get_llama_wrapper
 from llama_cookbook.policies import fpSixteen,bfSixteen, get_llama_wrapper
 from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.device_mesh import init_device_mesh
+from accelerate.utils import is_xpu_available
+
 
 
 def fsdp_auto_wrap_policy(model, transformer_layer_names):
 def fsdp_auto_wrap_policy(model, transformer_layer_names):
     import functools
     import functools