LLM: add patching function for llm finetuning (#10247)

python/llm/example/GPU/LLM-Finetuning/HF-PEFT/README.md

@@ -0,0 +1,37 @@
+# Finetuning on Intel GPU using Hugging Face PEFT code
+
+This example demonstrates how to easily run LLM finetuning application of PEFT use BigDL-LLM 4bit optimizations using [Intel GPUs](../../../README.md). By applying BigDL-LLM patch, you could run Hugging Face PEFT code on Intel GPUs using BigDL-LLM optimization without modification.
+
+Note, this example is just used for illustrating related usage and don't guarantee convergence of training.
+
+### 0. Requirements
+To run this example with BigDL-LLM on Intel GPUs, we have some recommended requirements for your machine, please refer to [here](../../README.md#requirements) for more information.
+
+### 1. Install
+
+```bash
+conda create -n llm python=3.9
+conda activate llm
+# below command will install intel_extension_for_pytorch==2.1.10+xpu as default
+pip install --pre --upgrade bigdl-llm[xpu] -f https://developer.intel.com/ipex-whl-stable-xpu
+pip install transformers==4.34.0 datasets
+pip install fire peft==0.5.0
+pip install oneccl_bind_pt==2.1.100 -f https://developer.intel.com/ipex-whl-stable-xpu # necessary to run distributed finetuning
+pip install accelerate==0.23.0
+pip install bitsandbytes scipy
+```
+
+### 2. Configures OneAPI environment variables
+```bash
+source /opt/intel/oneapi/setvars.sh
+```
+
+### 3. Finetune
+
+This example shows how to run [Alpaca LoRA Training](https://github.com/tloen/alpaca-lora/tree/main) directly on Intel GPU.
+
+```
+cd alpaca-lora
+python ./finetune.py --base_model "meta-llama/Llama-2-7b-hf" \
+                     --data_path "yahma/alpaca-cleaned"
+```

python/llm/example/GPU/LLM-Finetuning/HF-PEFT/alpaca-lora/finetune.py

@@ -0,0 +1,319 @@
+#
+# Copyright 2016 The BigDL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file is copied from
+# https://github.com/tloen/alpaca-lora/blob/main/finetune.py
+#
+# Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from bigdl.llm import llm_patch
+llm_patch(train=True)
+
+# The following is the original LLM finetuning code using PEFT (without BigDL-LLM)
+import os
+import sys
+from typing import List
+
+import fire
+import torch
+import transformers
+from datasets import load_dataset
+
+"""
+Unused imports:
+import torch.nn as nn
+import bitsandbytes as bnb
+"""
+
+from peft import (
+    LoraConfig,
+    get_peft_model,
+    get_peft_model_state_dict,
+    prepare_model_for_int8_training,
+    set_peft_model_state_dict,
+)
+from transformers import LlamaForCausalLM, LlamaTokenizer
+
+from utils.prompter import Prompter
+
+
+def train(
+    # model/data params
+    base_model: str = "",  # the only required argument
+    data_path: str = "yahma/alpaca-cleaned",
+    output_dir: str = "./lora-alpaca",
+    # training hyperparams
+    batch_size: int = 128,
+    micro_batch_size: int = 4,
+    num_epochs: int = 3,
+    learning_rate: float = 3e-4,
+    cutoff_len: int = 256,
+    val_set_size: int = 2000,
+    # lora hyperparams
+    lora_r: int = 8,
+    lora_alpha: int = 16,
+    lora_dropout: float = 0.05,
+    lora_target_modules: List[str] = [
+        "q_proj",
+        "v_proj",
+    ],
+    # llm hyperparams
+    train_on_inputs: bool = True,  # if False, masks out inputs in loss
+    add_eos_token: bool = False,
+    group_by_length: bool = False,  # faster, but produces an odd training loss curve
+    # wandb params
+    wandb_project: str = "",
+    wandb_run_name: str = "",
+    wandb_watch: str = "",  # options: false | gradients | all
+    wandb_log_model: str = "",  # options: false | true
+    resume_from_checkpoint: str = None,  # either training checkpoint or final adapter
+    prompt_template_name: str = "alpaca",  # The prompt template to use, will default to alpaca.
+):
+    if int(os.environ.get("LOCAL_RANK", 0)) == 0:
+        print(
+            f"Training Alpaca-LoRA model with params:\n"
+            f"base_model: {base_model}\n"
+            f"data_path: {data_path}\n"
+            f"output_dir: {output_dir}\n"
+            f"batch_size: {batch_size}\n"
+            f"micro_batch_size: {micro_batch_size}\n"
+            f"num_epochs: {num_epochs}\n"
+            f"learning_rate: {learning_rate}\n"
+            f"cutoff_len: {cutoff_len}\n"
+            f"val_set_size: {val_set_size}\n"
+            f"lora_r: {lora_r}\n"
+            f"lora_alpha: {lora_alpha}\n"
+            f"lora_dropout: {lora_dropout}\n"
+            f"lora_target_modules: {lora_target_modules}\n"
+            f"train_on_inputs: {train_on_inputs}\n"
+            f"add_eos_token: {add_eos_token}\n"
+            f"group_by_length: {group_by_length}\n"
+            f"wandb_project: {wandb_project}\n"
+            f"wandb_run_name: {wandb_run_name}\n"
+            f"wandb_watch: {wandb_watch}\n"
+            f"wandb_log_model: {wandb_log_model}\n"
+            f"resume_from_checkpoint: {resume_from_checkpoint or False}\n"
+            f"prompt template: {prompt_template_name}\n"
+        )
+    assert (
+        base_model
+    ), "Please specify a --base_model, e.g. --base_model='huggyllama/llama-7b'"
+    gradient_accumulation_steps = batch_size // micro_batch_size
+
+    prompter = Prompter(prompt_template_name)
+
+    device_map = "auto"
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    ddp = world_size != 1
+    if ddp:
+        device_map = {"": int(os.environ.get("LOCAL_RANK") or 0)}
+        gradient_accumulation_steps = gradient_accumulation_steps // world_size
+
+    # Check if parameter passed or if set within environ
+    use_wandb = len(wandb_project) > 0 or (
+        "WANDB_PROJECT" in os.environ and len(os.environ["WANDB_PROJECT"]) > 0
+    )
+    # Only overwrite environ if wandb param passed
+    if len(wandb_project) > 0:
+        os.environ["WANDB_PROJECT"] = wandb_project
+    if len(wandb_watch) > 0:
+        os.environ["WANDB_WATCH"] = wandb_watch
+    if len(wandb_log_model) > 0:
+        os.environ["WANDB_LOG_MODEL"] = wandb_log_model
+
+    model = LlamaForCausalLM.from_pretrained(
+        base_model,
+        load_in_8bit=True,
+        torch_dtype=torch.float16,
+        device_map=device_map,
+    )
+
+    tokenizer = LlamaTokenizer.from_pretrained(base_model)
+
+    tokenizer.pad_token_id = (
+        0  # unk. we want this to be different from the eos token
+    )
+    tokenizer.padding_side = "left"  # Allow batched inference
+
+    def tokenize(prompt, add_eos_token=True):
+        # there's probably a way to do this with the tokenizer settings
+        # but again, gotta move fast
+        result = tokenizer(
+            prompt,
+            truncation=True,
+            max_length=cutoff_len,
+            padding=False,
+            return_tensors=None,
+        )
+        if (
+            result["input_ids"][-1] != tokenizer.eos_token_id
+            and len(result["input_ids"]) < cutoff_len
+            and add_eos_token
+        ):
+            result["input_ids"].append(tokenizer.eos_token_id)
+            result["attention_mask"].append(1)
+
+        result["labels"] = result["input_ids"].copy()
+
+        return result
+
+    def generate_and_tokenize_prompt(data_point):
+        full_prompt = prompter.generate_prompt(
+            data_point["instruction"],
+            data_point["input"],
+            data_point["output"],
+        )
+        tokenized_full_prompt = tokenize(full_prompt)
+        if not train_on_inputs:
+            user_prompt = prompter.generate_prompt(
+                data_point["instruction"], data_point["input"]
+            )
+            tokenized_user_prompt = tokenize(
+                user_prompt, add_eos_token=add_eos_token
+            )
+            user_prompt_len = len(tokenized_user_prompt["input_ids"])
+
+            if add_eos_token:
+                user_prompt_len -= 1
+
+            tokenized_full_prompt["labels"] = [
+                -100
+            ] * user_prompt_len + tokenized_full_prompt["labels"][
+                user_prompt_len:
+            ]  # could be sped up, probably
+        return tokenized_full_prompt
+
+    model = prepare_model_for_int8_training(model)
+
+    config = LoraConfig(
+        r=lora_r,
+        lora_alpha=lora_alpha,
+        target_modules=lora_target_modules,
+        lora_dropout=lora_dropout,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+    model = get_peft_model(model, config)
+
+    if data_path.endswith(".json") or data_path.endswith(".jsonl"):
+        data = load_dataset("json", data_files=data_path)
+    else:
+        data = load_dataset(data_path)
+
+    if resume_from_checkpoint:
+        # Check the available weights and load them
+        checkpoint_name = os.path.join(
+            resume_from_checkpoint, "pytorch_model.bin"
+        )  # Full checkpoint
+        if not os.path.exists(checkpoint_name):
+            checkpoint_name = os.path.join(
+                resume_from_checkpoint, "adapter_model.bin"
+            )  # only LoRA model - LoRA config above has to fit
+            resume_from_checkpoint = (
+                False  # So the trainer won't try loading its state
+            )
+        # The two files above have a different name depending on how they were saved, but are actually the same.
+        if os.path.exists(checkpoint_name):
+            print(f"Restarting from {checkpoint_name}")
+            adapters_weights = torch.load(checkpoint_name)
+            set_peft_model_state_dict(model, adapters_weights)
+        else:
+            print(f"Checkpoint {checkpoint_name} not found")
+
+    model.print_trainable_parameters()  # Be more transparent about the % of trainable params.
+
+    if val_set_size > 0:
+        train_val = data["train"].train_test_split(
+            test_size=val_set_size, shuffle=True, seed=42
+        )
+        train_data = (
+            train_val["train"].shuffle().map(generate_and_tokenize_prompt)
+        )
+        val_data = (
+            train_val["test"].shuffle().map(generate_and_tokenize_prompt)
+        )
+    else:
+        train_data = data["train"].shuffle().map(generate_and_tokenize_prompt)
+        val_data = None
+
+    if not ddp and torch.cuda.device_count() > 1:
+        # keeps Trainer from trying its own DataParallelism when more than 1 gpu is available
+        model.is_parallelizable = True
+        model.model_parallel = True
+
+    trainer = transformers.Trainer(
+        model=model,
+        train_dataset=train_data,
+        eval_dataset=val_data,
+        args=transformers.TrainingArguments(
+            per_device_train_batch_size=micro_batch_size,
+            gradient_accumulation_steps=gradient_accumulation_steps,
+            warmup_steps=100,
+            num_train_epochs=num_epochs,
+            learning_rate=learning_rate,
+            fp16=True,
+            logging_steps=10,
+            optim="adamw_torch",
+            evaluation_strategy="steps" if val_set_size > 0 else "no",
+            save_strategy="steps",
+            eval_steps=200 if val_set_size > 0 else None,
+            save_steps=200,
+            output_dir=output_dir,
+            save_total_limit=3,
+            load_best_model_at_end=True if val_set_size > 0 else False,
+            ddp_find_unused_parameters=False if ddp else None,
+            group_by_length=group_by_length,
+            report_to="wandb" if use_wandb else None,
+            run_name=wandb_run_name if use_wandb else None,
+        ),
+        data_collator=transformers.DataCollatorForSeq2Seq(
+            tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True
+        ),
+    )
+    model.config.use_cache = False
+
+    old_state_dict = model.state_dict
+    model.state_dict = (
+        lambda self, *_, **__: get_peft_model_state_dict(
+            self, old_state_dict()
+        )
+    ).__get__(model, type(model))
+
+    if torch.__version__ >= "2" and sys.platform != "win32":
+        model = torch.compile(model)
+
+    trainer.train(resume_from_checkpoint=resume_from_checkpoint)
+
+    model.save_pretrained(output_dir)
+
+    print(
+        "\n If there's a warning about missing keys above, please disregard :)"
+    )
+
+
+if __name__ == "__main__":
+    fire.Fire(train)

python/llm/example/GPU/LLM-Finetuning/HF-PEFT/alpaca-lora/templates/alpaca.json

@@ -0,0 +1,7 @@
+{
+    "//": "This file is copied from https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca.json",
+    "description": "Template used by Alpaca-LoRA.",
+    "prompt_input": "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n",
+    "prompt_no_input": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:\n",
+    "response_split": "### Response:"    
+}

python/llm/example/GPU/LLM-Finetuning/HF-PEFT/alpaca-lora/templates/alpaca_legacy.json

@@ -0,0 +1,7 @@
+{
+    "//": "This file is copied from https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca_legacy.json",
+    "description": "Legacy template, used by Original Alpaca repository.",
+    "prompt_input": "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:",
+    "prompt_no_input": "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:",
+    "response_split": "### Response:"    
+}

python/llm/example/GPU/LLM-Finetuning/HF-PEFT/alpaca-lora/templates/alpaca_short.json

@@ -0,0 +1,7 @@
+{
+    "//": "This file is copied from https://github.com/tloen/alpaca-lora/blob/main/templates/alpaca_short.json",
+    "description": "A shorter template to experiment with.",
+    "prompt_input": "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n",
+    "prompt_no_input": "### Instruction:\n{instruction}\n\n### Response:\n",
+    "response_split": "### Response:"    
+}

python/llm/example/GPU/LLM-Finetuning/HF-PEFT/alpaca-lora/templates/vigogne.json

@@ -0,0 +1,7 @@
+{
+    "//": "This file is copied from https://github.com/tloen/alpaca-lora/blob/main/templates/vigogne.json",
+    "description": "French template, used by Vigogne for finetuning.",
+    "prompt_input": "Ci-dessous se trouve une instruction qui décrit une tâche, associée à une entrée qui fournit un contexte supplémentaire. Écrivez une réponse qui complète correctement la demande.\n\n### Instruction:\n{instruction}\n\n### Entrée:\n{input}\n\n### Réponse:\n",
+    "prompt_no_input": "Ci-dessous se trouve une instruction qui décrit une tâche. Écrivez une réponse qui complète correctement la demande.\n\n### Instruction:\n{instruction}\n\n### Réponse:\n",
+    "response_split": "### Réponse:"
+}

python/llm/example/GPU/LLM-Finetuning/HF-PEFT/alpaca-lora/utils/prompter.py

@@ -0,0 +1,83 @@
+#
+# Copyright 2016 The BigDL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# This file is copied from
+# https://github.com/tloen/alpaca-lora/blob/main/utils/prompter.py
+#
+# Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+A dedicated helper to manage templates and prompt building.
+"""
+
+import json
+import os.path as osp
+from typing import Union
+
+
+class Prompter(object):
+    __slots__ = ("template", "_verbose")
+
+    def __init__(self, template_name: str = "", verbose: bool = False):
+        self._verbose = verbose
+        if not template_name:
+            # Enforce the default here, so the constructor can be called with '' and will not break.
+            template_name = "alpaca"
+        file_name = osp.join("templates", f"{template_name}.json")
+        if not osp.exists(file_name):
+            raise ValueError(f"Can't read {file_name}")
+        with open(file_name) as fp:
+            self.template = json.load(fp)
+        if self._verbose:
+            print(
+                f"Using prompt template {template_name}: {self.template['description']}"
+            )
+
+    def generate_prompt(
+        self,
+        instruction: str,
+        input: Union[None, str] = None,
+        label: Union[None, str] = None,
+    ) -> str:
+        # returns the full prompt from instruction and optional input
+        # if a label (=response, =output) is provided, it's also appended.
+        if input:
+            res = self.template["prompt_input"].format(
+                instruction=instruction, input=input
+            )
+        else:
+            res = self.template["prompt_no_input"].format(
+                instruction=instruction
+            )
+        if label:
+            res = f"{res}{label}"
+        if self._verbose:
+            print(res)
+        return res
+
+    def get_response(self, output: str) -> str:
+        return output.split(self.template["response_split"])[1].strip()

python/llm/example/GPU/LLM-Finetuning/README.md

@@ -8,6 +8,7 @@ This folder contains examples of running different training mode with BigDL-LLM
 - [ReLora](ReLora): examples of running ReLora finetuning
 - [DPO](DPO): examples of running DPO finetuning
 - [common](common): common templates and utility classes in finetuning examples
+- [HF-PEFT](HF-PEFT): run finetuning on Intel GPU using Hugging Face PEFT code without modification
 
 
 ## Troubleshooting

python/llm/src/bigdl/llm/__init__.py

@@ -22,6 +22,7 @@
 from .convert_model import llm_convert
 from .optimize import optimize_model
 import os
+from .llm_patching import llm_patch, llm_unpatch
 
 # Default is false, set to true to auto importing Intel Extension for PyTorch.
 BIGDL_IMPORT_IPEX = os.getenv("BIGDL_IMPORT_IPEX", 'True').lower() in ('true', '1', 't')

python/llm/src/bigdl/llm/llm_patching.py

@@ -0,0 +1,79 @@
+#
+# Copyright 2016 The BigDL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import transformers
+import importlib
+import sys
+from bigdl.llm.utils.common import invalidInputError
+from enum import Enum
+
+bigdl_patched = None  # None or 'Train' or 'Inference'
+attrs = []
+
+
+def replace_attr(obj, name: str, value):
+    original_attr = getattr(obj, name)
+    setattr(obj, name, value)
+    attrs.append((obj, name, original_attr))
+
+
+def llm_patch(train=False):
+    '''
+    llm_patch is used to make users' LLM application benefit from BigDL-LLM optimization
+    with only one-line code patch.
+
+    :param train: Whether to apply bigdl-llm patch for training code, default to be `False`.
+    '''
+    global bigdl_patched
+    if bigdl_patched:
+        return
+
+    # Initial version of patch for llm finetuning, inference support TBD
+    if train:
+        from bigdl.llm.transformers import AutoModelForCausalLM, AutoModel
+        replace_attr(transformers, "AutoModelForCausalLM", AutoModelForCausalLM)
+        replace_attr(transformers, "LlamaForCausalLM", AutoModelForCausalLM)
+        replace_attr(transformers, "AutoModel", AutoModel)
+
+        import_peft_check = 'peft' in sys.modules or 'peft.utils' in sys.modules or \
+            'peft.tuners' in sys.modules or 'peft.mapping' in sys.modules
+        invalidInputError(not import_peft_check,
+                          'llm_patch() should be called at the beginning of your code.')
+        import peft
+        from bigdl.llm.transformers.qlora import get_peft_model, prepare_model_for_kbit_training,\
+            LoraConfig, TrainingArguments
+        replace_attr(transformers, "TrainingArguments", TrainingArguments)
+        get_peft_model_original = getattr(peft, "get_peft_model")
+        replace_attr(peft, "get_peft_model", get_peft_model)
+        setattr(peft, "get_peft_model_original", get_peft_model_original)
+        replace_attr(peft, "prepare_model_for_kbit_training", prepare_model_for_kbit_training)
+        replace_attr(peft, "prepare_model_for_int8_training", prepare_model_for_kbit_training)
+        replace_attr(peft, "LoraConfig", LoraConfig)
+        bigdl_patched = 'Train'
+
+
+def llm_unpatch():
+    '''
+    llm_unpatch is an reverse function to llm_patch.
+    '''
+    global bigdl_patched
+
+    if bigdl_patched is None:
+        return
+
+    for obj, name, torch_attr in attrs:
+        setattr(obj, name, torch_attr)
+    bigdl_patched = None

python/llm/src/bigdl/llm/transformers/model.py

@@ -50,6 +50,8 @@ import warnings
 import copy
 from .utils import logger
 
+patched_training_mode = None
+
 
 def save_low_bit(self, *args, **kwargs):
     invalidInputError(self.config.to_dict().get("bigdl_transformers_low_bit", False),
@@ -215,6 +217,20 @@ class _BaseAutoModelClass:
             optimize_model = False
             kwargs["modules_to_not_convert"] = ["lm_head"]
 
+        load_in_8bit = kwargs.pop("load_in_8bit", False)
+        from bigdl.llm.llm_patching import bigdl_patched
+        if bigdl_patched == 'Train':
+            global patched_training_mode
+            if load_in_low_bit == "nf4" or load_in_low_bit == "sym_int4" or load_in_4bit:
+                # qlora
+                patched_training_mode = 'qlora'
+            else:
+                # lora
+                patched_training_mode = 'lora'
+                load_in_low_bit = "bf16"
+            optimize_model = False
+            kwargs["modules_to_not_convert"] = ["lm_head"]
+
         if load_in_4bit or load_in_low_bit:
 
             if config_dict.get("quantization_config", None) is not None:
@@ -413,6 +429,8 @@ class _BaseAutoModelClass:
         else:
             _load_pre()
             try:
+                # To handle the input CUDA setting (such as 'device_map={"":0}'), ignore it
+                kwargs.pop('device_map', None)
                 model = cls.HF_Model.from_pretrained(*args, **kwargs)
             except NotImplementedError:
                 logger.info("Failed to load models with `low_cpu_mem_usage` specified, "

python/llm/src/bigdl/llm/transformers/qlora.py

@@ -234,6 +234,8 @@ def _create_new_module(create_new_module_func, lora_config, adapter_name, target
 
 from peft.tuners.lora import LoraModel
 from peft.tuners.lora import LoraConfig as LoraConfigBase
+from transformers import TrainingArguments as TrainingArgumentsBase
+from transformers.training_args import OptimizerNames
 from dataclasses import dataclass, field
 
 
@@ -241,6 +243,30 @@ from dataclasses import dataclass, field
 class LoraConfig(LoraConfigBase):
     training_mode: str = field(default="qlora", metadata={"help": "determine training mode"})
 
+    def __init__(self, *args, **kwargs):
+        self.training_mode = kwargs.pop("training_mode", "qlora")
+        super().__init__(*args, **kwargs)
+        from bigdl.llm.llm_patching import bigdl_patched
+        if bigdl_patched == 'Train':
+            from .model import patched_training_mode
+            self.training_mode = patched_training_mode
+
+
+supported_optim = ["adamw_hf", "adamw_torch", "adafactor", "sgd", "adagrad", "rmsprop"]
+
+
+@dataclass
+class TrainingArguments(TrainingArgumentsBase):
+    def __init__(self, *args, **kwargs):
+        kwargs["fp16"] = False
+        kwargs["bf16"] = True
+        for optim in supported_optim.copy():
+            supported_optim.append(OptimizerNames(optim))
+        if kwargs["optim"] not in supported_optim:
+            LOG.info(f"{self.optim} is not supported yet and adamw_torch optimizer is used.")
+            kwargs["optim"] = "adamw_torch"
+        super().__init__(*args, **kwargs)
+
 
 def get_peft_model(*args, **kwargs):
     old_create_new_module = LoraModel._create_new_module
@@ -248,6 +274,10 @@ def get_peft_model(*args, **kwargs):
                                                                   old_create_new_module))
 
     try:
+        from bigdl.llm.llm_patching import bigdl_patched
+        if bigdl_patched == 'Train':
+            from peft import get_peft_model_original
+        else:
             from peft import get_peft_model as get_peft_model_original
         model = get_peft_model_original(*args, **kwargs)
     finally: