Fix Llama-70B-FSDP model loading issue (huggingface#752)

Fix Llama-70B-FSDP model loading issue (#63)
HabanaAI · Mar 11, 2024 · e6ccd26 · e6ccd26
1 parent 2f18086
commit e6ccd26
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/examples/language-modeling/run_lora_clm.py b/examples/language-modeling/run_lora_clm.py
@@ -42,7 +42,6 @@
 from transformers.trainer_utils import is_main_process
 
 from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments
-from optimum.habana.peft.layer import GaudiLoraLayerLinearForward
 from optimum.habana.utils import set_seed
 
 
@@ -675,7 +674,10 @@ def compute_metrics(eval_preds):
         )
         if training_args.gradient_checkpointing:
             model.enable_input_require_grads()
-        tuners.lora.layer.Linear.forward = GaudiLoraLayerLinearForward
+        if training_args.torch_compile:
+            from optimum.habana.peft.layer import GaudiLoraLayerLinearForward
+
+            tuners.lora.layer.Linear.forward = GaudiLoraLayerLinearForward
         lora_model = get_peft_model(model, peft_config)
         if training_args.bf16:
             lora_model = lora_model.to(torch.bfloat16)

diff --git a/optimum/habana/transformers/training_args.py b/optimum/habana/transformers/training_args.py
@@ -568,6 +568,7 @@ def __post_init__(self):
         # accelerate integration for FSDP
         if len(self.fsdp) > 0 and not self.fsdp_config["xla"]:
             os.environ["ACCELERATE_USE_FSDP"] = "true"
+            os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = "true"
             from accelerate.utils.constants import (
                 FSDP_AUTO_WRAP_POLICY,
                 FSDP_SHARDING_STRATEGY,