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,