Added QLoRA support for Decoder transformers with tune_strategy "Normal"

requirements.txt

@@ -1,7 +1,8 @@
 numpy==1.24.2
 datasets==2.10.1
-peft>=0.4.0
-torch>=2.0.0
+tokenizers==0.13.3
+peft==0.4.0
+torch>=2.0.1
 wandb==0.14.0
 deepspeed==0.10.0
 trl==0.5.0
@@ -18,7 +19,6 @@ dill<0.3.5
 bitsandbytes>=0.40.0
 pydantic<=1.10.9
 gradio
-
 accelerate>=0.21.0
 einops>=0.6.1
 scikit-learn==1.2.2

scripts/run_finetune_with_qlora_save_aggregated_weights.sh

@@ -0,0 +1,63 @@
+#!/bin/bash
+# Please run this script under ${project_id} in project directory of
+
+# Parses arguments
+model_name_or_path=meta-llama/Llama-2-13b-hf
+dataset_path=/home/paperspace/LMFlow/alpaca/train
+output_dir=output_models/finetune
+deepspeed_args="--master_port=11000"
+
+while [[ $# -ge 1 ]]; do
+  key="$1"
+  case ${key} in
+    -m|--model_name_or_path)
+      model_name_or_path="$2"
+      shift
+      ;;
+    -d|--dataset_path)
+      dataset_path="$2"
+      shift
+      ;;
+    -o|--output_model_path)
+      output_dir="$2"
+      shift
+      ;;
+    --deepspeed_args)
+      deepspeed_args="$2"
+      shift
+      ;;
+    *)
+      echo "error: unknown option \"${key}\"" 1>&2
+      exit 1
+  esac
+  shift
+done
+
+# Finetune
+exp_id=finetune_with_lora
+project_dir=$(cd "$(dirname $0)"/..; pwd)
+log_dir=${project_dir}/log/${exp_id}
+mkdir -p ${output_dir} ${log_dir}
+
+deepspeed ${deepspeed_args} \
+  examples/finetune.py \
+    --model_name_or_path ${model_name_or_path} \
+    --dataset_path ${dataset_path} \
+    --output_dir ${output_dir} --overwrite_output_dir \
+    --num_train_epochs 0.01 \
+    --learning_rate 1e-4 \
+    --block_size 512 \
+    --per_device_train_batch_size 1 \
+    --use_qlora 1 \
+    --save_aggregated_lora 1\
+    --deepspeed configs/ds_config_zero2.json \
+    --fp16 \
+    --run_name ${exp_id} \
+    --validation_split_percentage 0 \
+    --logging_steps 20 \
+    --do_train \
+    --ddp_timeout 72000 \
+    --save_steps 5000 \
+    --dataloader_num_workers 1 \
+    | tee ${log_dir}/train.log \
+    2> ${log_dir}/train.err

setup.py

@@ -47,6 +47,6 @@
 try:
   gpu_state = subprocess.check_output(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"])
   if b"A100" or b"A40" in gpu_state:
-    subprocess.call(["pip", "install", "flash-attn==2.0.2"])
+    subprocess.call(["pip", "install", "flash-attn==2.0.4"])
 except:
   pass

src/lmflow/args.py

@@ -151,6 +151,14 @@ class ModelArguments:
             )
         },
     )
+    trust_remote_code : bool = field(
+        default=False,
+        metadata={
+            "help": (
+                "Whether to trust remote code when loading model."
+            )
+        },
+    )
     torch_dtype: Optional[str] = field(
         default=None,
         metadata={
@@ -165,6 +173,24 @@ class ModelArguments:
         default=False,
         metadata={"help": "Whether to lora."},
     )
+    use_qlora: bool = field(
+        default=False,
+        metadata={"help": "Whether to use qlora."},
+    )
+    bits: int = field(
+        default=4,
+        metadata={"help": "The number of bits for quantization.",
+                  "choices": [4, 8],},
+    )
+    quant_type: str = field(
+        default='nf4',
+        metadata={"help": "The quantization type for quantization.",
+                  "choices": ["nf4", "fp4"],},
+    )
+    double_quant: bool = field(
+        default=True,
+        metadata={"help": "Whether to use double quantization."},
+    )
     lora_r: int = field(
         default=8,
         metadata={"help": "the rank of the lora parameters. The smaller lora_r is , the fewer parameters lora has."},

src/lmflow/models/hf_decoder_model.py

@@ -21,15 +21,18 @@
 import hashlib
 import logging
 from typing import List, Union
-
+import os, shutil
 import deepspeed
 
+from pathlib import Path
+
 from peft import (
     LoraConfig,
     PeftModel,
     TaskType,
     get_peft_config,
     get_peft_model,
+    prepare_model_for_kbit_training
 )
 
 import torch
@@ -38,6 +41,9 @@
 
 from transformers.testing_utils import CaptureLogger
 
+from transformers import BitsAndBytesConfig
+import bitsandbytes
+
 from transformers import (
     CONFIG_MAPPING,
     AutoConfig,
@@ -253,15 +259,38 @@ def __init__(
 
         if tune_strategy == 'normal':
             if model_args.model_name_or_path:
+                compute_dtype = torch_dtype
+                device_map = "auto"
+                if os.environ.get('LOCAL_RANK') is not None:
+                    local_rank = int(os.environ.get('LOCAL_RANK','0'))
+                    device_map = {'': local_rank}
+
+                if model_args.use_qlora:
+                    model_args.use_lora = True
+                    quant_config = BitsAndBytesConfig(
+                        load_in_4bit=model_args.bits == 4,
+                        load_in_8bit=model_args.bits == 8,
+                        llm_int8_threshold=6.0,
+                        llm_int8_has_fp16_weight=False,
+                        bnb_4bit_compute_dtype=compute_dtype,
+                        bnb_4bit_use_double_quant=model_args.double_quant,
+                        bnb_4bit_quant_type=model_args.quant_type,
+                    )
                 model = AutoModelForCausalLM.from_pretrained(
                     model_args.model_name_or_path,
                     from_tf=bool(".ckpt" in model_args.model_name_or_path),
                     config=config,
+                    quantization_config=quant_config if model_args.use_qlora else None,
                     cache_dir=model_args.cache_dir,
                     revision=model_args.model_revision,
                     use_auth_token=True if model_args.use_auth_token else None,
                     torch_dtype=torch_dtype,
+                    device_map=device_map,
+                    trust_remote_code = model_args.trust_remote_code,
                 )
+                if model_args.use_qlora:
+                    model.gradient_checkpointing_enable()
+                    model = prepare_model_for_kbit_training(model)
             else:
                 model = AutoModelForCausalLM.from_config(config)
                 n_params = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters()).values())
@@ -632,11 +661,54 @@ def inference(self, inputs, use_accelerator=False, *args, **kwargs):
 
 
     def merge_lora_weights(self):
-        if self.model_args.use_lora:
+        if self.model_args.use_lora and not self.model_args.use_qlora:
+            self.get_backend_model().merge_and_unload()
+        elif self.model_args.use_qlora:
+            logger.warning("Reloading base model in 16-bit precision to merge adapter weights. NOTE: Your device must have"
+                           "sufficient memory to reload the model in half-precision without quantization.")
+            self.get_peft_without_qlora()
             self.get_backend_model().merge_and_unload()
         else:
             logger.warning("LoRA training is NOT enabled. Merging LoRA weights is not applicable.")
 
+    def get_peft_without_qlora(self):
+        import tempfile
+
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            print('created temporary directory', tmpdirname)
+
+
+            self.get_backend_model().save_pretrained(tmpdirname)
+
+            torch_dtype = (
+                self.model_args.torch_dtype
+                if self.model_args.torch_dtype in ["auto", None]
+                else getattr(torch, self.model_args.torch_dtype)
+            )
+            config_kwargs = {
+                "cache_dir": self.model_args.cache_dir,
+                "revision": self.model_args.model_revision,
+                "use_auth_token": True if self.model_args.use_auth_token else None,
+            }
+            config = AutoConfig.from_pretrained(self.model_args.model_name_or_path, **config_kwargs)
+            device_map = "auto"
+            if os.environ.get('LOCAL_RANK') is not None:
+                local_rank = int(os.environ.get('LOCAL_RANK','0'))
+                device_map = {'': local_rank}
+
+            self.backend_model_full = AutoModelForCausalLM.from_pretrained(
+                self.model_args.model_name_or_path,
+                from_tf=bool(".ckpt" in self.model_args.model_name_or_path),
+                config=config,
+                cache_dir=self.model_args.cache_dir,
+                revision=self.model_args.model_revision,
+                use_auth_token=True if self.model_args.use_auth_token else None,
+                torch_dtype=torch_dtype,
+                device_map=device_map,
+                trust_remote_code = self.model_args.trust_remote_code,
+            )
+
+            self.backend_model = PeftModel.from_pretrained(self.backend_model_full, tmpdirname)
 
     def save(self, dir, save_full_model=False, *args, **kwargs):
         """