fix usage of acclerate

examples/llama_acc.sh

@@ -2,7 +2,7 @@
 set -ex
 
 # FSDP
-./examples/run.sh --model ./hf_models/config/llama-1b --accelerator acc --gc --mbs 4 --fsdp 4
+./examples/run.sh --model ./hf_models/config/llama-1b --accelerator acc --gc --mbs 4 --fsdp 4 --use_flash_attn
 
 # TP
 # ./examples/run.sh --model ./hf_models/config/llama-1b --accelerator acc --gc --mbs 24 --tp 4

flashmodels/accelerators/acc_baichuan_accelerator.py

@@ -17,7 +17,7 @@ def accelerate_internal(self, model, loader):
                 raise NotImplementedError("resume_from_checkpoint.")
 
             config = self.get_config(model)
-            model = ta.accelerate(model, config)
+            model = ta.accelerate(model, config=config)
             return model, loader
 
     def get_config(self, model):

flashmodels/accelerators/acc_gemma_accelerator.py

@@ -12,7 +12,7 @@ def accelerate(self, model, loader):
 
     def accelerate_internal(self, model, loader):
         config = self.get_config()
-        model = ta.accelerate(model, config)
+        model = ta.accelerate(model, config=config)
         return model, loader
 
     def get_config(self):

flashmodels/accelerators/acc_glm_accelerator.py

@@ -17,7 +17,7 @@ def accelerate_internal(self, model, loader):
                 raise NotImplementedError("resume_from_checkpoint.")
 
             config = self.get_config(model)
-            model = ta.accelerate(model, config)
+            model = ta.accelerate(model, config=config)
             return model, loader
 
     def get_config(self, model):

flashmodels/accelerators/acc_gpt_accelerator.py

@@ -20,7 +20,7 @@ def accelerate_internal(self, model, loader):
                 raise NotImplementedError("resume_from_checkpoint.")
 
             config = self.get_config(model)
-            model = ta.accelerate(model, config)
+            model = ta.accelerate(model, config=config)
             return model, loader
 
         device = lazy_device()

flashmodels/accelerators/acc_llama_accelerator.py

@@ -99,7 +99,7 @@ def accelerate_internal(self, model, loader):
                                                        self.args.sp)
 
         config = self.get_config(model)
-        model = ta.accelerate(model, config)
+        model = ta.accelerate(model, config=config)
 
         if self.args.tp_num > 1 and self.args.pp_num > 1:
             self.parallel_3d(model._get_underlay_model())

flashmodels/accelerators/acc_olmo_accelerator.py

@@ -17,7 +17,7 @@ def accelerate_internal(self, model, loader):
                 raise NotImplementedError("resume_from_checkpoint.")
 
             config = self.get_config(model)
-            model = ta.accelerate(model, config)
+            model = ta.accelerate(model, config=config)
             return model, loader
         else:
             raise NotImplementedError("Currently, only FSDP is supported.")

flashmodels/accelerators/acc_qwen_accelerator.py

@@ -37,7 +37,7 @@ def accelerate_internal(self, model, loader):
                 raise NotImplementedError("resume_from_checkpoint.")
 
             config = self.get_config(model)
-            model = ta.accelerate(model, config)
+            model = ta.accelerate(model, config=config)
             return model, loader
 
     def get_config(self, model):

flashmodels/accelerators/cuda_llama_accelerator.py

@@ -13,6 +13,7 @@
     FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
 from torch.distributed.fsdp.wrap import ModuleWrapPolicy
+from transformers.models.llama.modeling_llama import LlamaDecoderLayer
 
 from flashmodels.accelerators.accelerator import (Accelerator,
                                                   AcceleratorFactory)
@@ -70,7 +71,7 @@ def apply_checkpointing(self, model):
             checkpoint_wrapper,
             checkpoint_impl=CheckpointImpl.NO_REENTRANT,
         )
-        check_fn = lambda submodule: isinstance(submodule, transformers.models.llama.modeling_llama.LlamaDecoderLayer)
+        check_fn = lambda submodule: isinstance(LlamaDecoderLayer)
         apply_activation_checkpointing(
             model,
             checkpoint_wrapper_fn=non_reentrant_wrapper,
@@ -96,9 +97,7 @@ def fsdp(self, model):
                 convert_outputs_to_fp32(model.forward.__func__), model)
 
         # Use auto_wrap_poliy for nested wrapping instead of only a top-level FSDP.
-        auto_wrap_policy = ModuleWrapPolicy({
-            transformers.models.llama.modeling_llama.LlamaDecoderLayer,
-        })
+        auto_wrap_policy = ModuleWrapPolicy({LlamaDecoderLayer, })
 
         mixed_precision_policy = None
         if self.args.fp16 or self.args.bf16: