diff --git a/requirements.txt b/requirements.txt index b39799437..7f0fa95ad 100644 --- a/requirements.txt +++ b/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 \ No newline at end of file diff --git a/scripts/run_finetune_with_qlora_save_aggregated_weights.sh b/scripts/run_finetune_with_qlora_save_aggregated_weights.sh new file mode 100644 index 000000000..1667e02f9 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/setup.py b/setup.py index f44199536..a6b6a7084 100644 --- a/setup.py +++ b/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 diff --git a/src/lmflow/args.py b/src/lmflow/args.py index 79ce90160..e91962a71 100644 --- a/src/lmflow/args.py +++ b/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."}, diff --git a/src/lmflow/models/hf_decoder_model.py b/src/lmflow/models/hf_decoder_model.py index 4e2f594af..09fe3a294 100644 --- a/src/lmflow/models/hf_decoder_model.py +++ b/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): """