diff --git a/assets/wechat.jpg b/assets/wechat.jpg index 0181431a..766b95e8 100644 Binary files a/assets/wechat.jpg and b/assets/wechat.jpg differ diff --git a/src/llmtuner/model/loader.py b/src/llmtuner/model/loader.py index c961e9b0..0d671fb4 100644 --- a/src/llmtuner/model/loader.py +++ b/src/llmtuner/model/loader.py @@ -119,16 +119,6 @@ def load_model_and_tokenizer( model_args.rope_scaling, scaling_factor )) - # Set FlashAttention-2 - if model_args.flash_attn: - if not is_flash_attn2_available(): - logger.warning("FlashAttention-2 is not installed.") - elif getattr(config, "model_type", None) == "qwen": - logger.info("Current model automatically enables FlashAttention if installed.") - else: - setattr(config, "attn_implementation", "flash_attention_2") - logger.info("Using FlashAttention-2 for faster training and inference.") - # Set shift short attention (S^2-Attn) if is_trainable and model_args.shift_attn: logger.warning("Shift short attention is temporarily invalid due to breaking changes.") @@ -138,10 +128,19 @@ def load_model_and_tokenizer( # else: # logger.warning("Current model does not support shift short attention.") + # Set FlashAttention-2 + if model_args.flash_attn: + if not is_flash_attn2_available(): + logger.warning("FlashAttention-2 is not installed.") + elif getattr(config, "model_type", None) == "qwen": + logger.info("Current model automatically enables FlashAttention if installed.") + else: + config_kwargs["use_flash_attention_2"] = True + logger.info("Using FlashAttention-2 for faster training and inference.") + # Quantization configurations (using gptq or awq) if getattr(config, "quantization_config", None): - if model_args.quantization_bit is not None: # remove bnb quantization - model_args.quantization_bit = None + model_args.quantization_bit = None # remove bnb quantization config_kwargs["device_map"] = {"": get_current_device()} quantization_config = getattr(config, "quantization_config", None) logger.info("Loading {}-bit pre-quantized model.".format(quantization_config.get("bits", -1))) diff --git a/tests/cal_lr.py b/tests/cal_lr.py index 7261d2be..0cbaa0b9 100644 --- a/tests/cal_lr.py +++ b/tests/cal_lr.py @@ -26,7 +26,7 @@ def calculate_lr( cutoff_len: int, # i.e. maximum input length during training batch_size: int, # total batch size, namely (batch size * gradient accumulation * world size) is_mistral: bool, # mistral model uses a smaller learning rate, - dataset_dir: Optional[str] = "data" + dataset_dir: Optional[str] = "../data" ): model_args, data_args, training_args, finetuning_args, _ = get_train_args(dict( stage="sft", @@ -38,7 +38,7 @@ def calculate_lr( output_dir="dummy_dir" )) trainset = get_dataset(model_args, data_args) - _, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, is_trainable=False, stage="sft") + _, tokenizer = load_model_and_tokenizer(model_args, finetuning_args, is_trainable=False, add_valuehead=False) trainset = preprocess_dataset(trainset, tokenizer, data_args, training_args, stage="sft") data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, label_pad_token_id=IGNORE_INDEX) dataloader = DataLoader( diff --git a/tests/loftq_init.py b/tests/loftq_init.py new file mode 100644 index 00000000..32cb96e0 --- /dev/null +++ b/tests/loftq_init.py @@ -0,0 +1,77 @@ +# coding=utf-8 +# Initializes LoRA weights with LoRA-fine-tuning-aware Quantization (LoftQ) +# Usage: python loftq_init.py --model_name_or_path path_to_model --save_dir output_dir +# Inspired by: https://github.com/huggingface/peft/blob/main/examples/loftq_finetuning/quantize_save_load.py + +import os +import fire +import torch +import torch.nn as nn +from typing import Optional +from transformers import AutoModelForCausalLM, AutoTokenizer +from peft import LoftQConfig, LoraConfig, TaskType, get_peft_model + + +class Shell(nn.Module): + + def __init__(self, weight: torch.Tensor, bias: Optional[torch.Tensor] = None): + super().__init__() + self.weight = nn.Parameter(weight, requires_grad=False) + if bias is not None: + self.bias = nn.Parameter(bias, requires_grad=False) + + +def unwrap_model(model: nn.Module, pattern=".base_layer") -> None: + for name in set([k.split(pattern)[0] for k, _ in model.named_modules() if pattern in k]): + parent_name = ".".join(name.split(".")[:-1]) + child_name = name.split(".")[-1] + parent_module = model.get_submodule(parent_name) + child_module = getattr(parent_module, child_name) + base_layer = getattr(child_module, "base_layer") + weight = getattr(base_layer, "weight", None) + bias = getattr(base_layer, "bias", None) + setattr(parent_module, child_name, Shell(weight, bias)) + + print("Model unwrapped.") + + +def quantize_loftq( + model_name_or_path: str, + save_dir: str, + loftq_bits: Optional[int] = 4, + loftq_iter: Optional[int] = 1, + lora_alpha: Optional[int] = None, + lora_rank: Optional[int] = 16, + lora_target: Optional[str] = "q_proj,v_proj" +): + tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True) + model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True, torch_dtype="auto") + loftq_config = LoftQConfig(loftq_bits=loftq_bits, loftq_iter=loftq_iter) + lora_config = LoraConfig( + task_type=TaskType.CAUSAL_LM, + inference_mode=True, + r=lora_rank, + lora_alpha=lora_alpha if lora_alpha is not None else lora_rank * 2, + lora_dropout=0.1, + target_modules=[name.strip() for name in lora_target.split(",")], + init_lora_weights="loftq", + loftq_config=loftq_config + ) + + # Init LoftQ model + lora_model = get_peft_model(model, lora_config) + base_model = lora_model.get_base_model() + + # Save LoftQ model + setattr(lora_model.base_model.peft_config["default"], "base_model_name_or_path", save_dir) + setattr(lora_model.base_model.peft_config["default"], "init_lora_weights", True) + lora_model.save_pretrained(os.path.join(save_dir, "adapters")) + + # Save base model + unwrap_model(base_model) + base_model.save_pretrained(save_dir) + tokenizer.save_pretrained(save_dir) + + +if __name__ == "__main__": + fire.Fire(quantize_loftq)