ros/LLaMA-Factory
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

support pissa

Browse Source 
Former-commit-id: ef8e45f2eaf466c54e9a671512a2974575677b08
This commit is contained in:
hiyouga
2024-06-16 01:08:12 +08:00
parent 05f3a3c944
commit 32f45c9e91
19 changed files with 406 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
src/llamafactory/hparams/finetuning_args.py
View File

@@ -108,6 +108,18 @@ class LoraArguments:
default=False,
metadata={"help": "Whether or not to use the weight-decomposed lora method (DoRA)."},
)
pissa_init: bool = field(
default=False,
metadata={"help": "Whether or not to initialize a PiSSA adapter."},
)
pissa_iter: int = field(
default=4,
metadata={"help": "The number of iteration steps performed by FSVD in PiSSA. Use -1 to disable it."},
)
pissa_convert: bool = field(
default=False,
metadata={"help": "Whether or not to convert the PiSSA adapter to a normal LoRA adapter."},
)
create_new_adapter: bool = field(
default=False,
metadata={"help": "Whether or not to create a new adapter with randomly initialized weight."},
@@ -340,7 +352,7 @@ class FinetuningArguments(FreezeArguments, LoraArguments, RLHFArguments, GaloreA
self.additional_target: Optional[List[str]] = split_arg(self.additional_target)
self.galore_target: List[str] = split_arg(self.galore_target)
self.freeze_vision_tower = self.freeze_vision_tower or self.train_mm_proj_only
self.use_ref_model = self.pref_loss not in ["orpo", "simpo"]
self.use_ref_model = (self.stage == "dpo" and self.pref_loss not in ["orpo", "simpo"])
assert self.finetuning_type in ["lora", "freeze", "full"], "Invalid fine-tuning method."
assert self.ref_model_quantization_bit in [None, 8, 4], "We only accept 4-bit or 8-bit quantization."
@@ -367,5 +379,11 @@ class FinetuningArguments(FreezeArguments, LoraArguments, RLHFArguments, GaloreA
if self.loraplus_lr_ratio is not None and self.finetuning_type != "lora":
raise ValueError("`loraplus_lr_ratio` is only valid for LoRA training.")
if self.pissa_convert and self.finetuning_type != "lora":
raise ValueError("`pissa_convert` is only valid for LoRA training.")
if self.pissa_convert and (self.stage in ["rm", "ppo", "kto"] or self.use_ref_model):
raise ValueError("Cannot use PiSSA for current training stage.")
if self.train_mm_proj_only and self.finetuning_type != "full":
raise ValueError("`train_mm_proj_only` is only valid for full training.")

8
src/llamafactory/hparams/model_args.py
View File

@@ -1,6 +1,6 @@
# Copyright 2024 HuggingFace Inc. and the LlamaFactory team.
#
# This code is inspired by HuggingFace's transformers library.
# This code is inspired by the HuggingFace's transformers library.
# https://github.com/huggingface/transformers/blob/v4.40.0/examples/pytorch/language-modeling/run_clm.py
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -45,6 +45,10 @@ class ModelArguments:
)
},
)
adapter_folder: Optional[str] = field(
default=None,
metadata={"help": "The folder containing the adapter weights to load."},
)
cache_dir: Optional[str] = field(
default=None,
metadata={"help": "Where to store the pre-trained models downloaded from huggingface.co or modelscope.cn."},
@@ -150,7 +154,7 @@ class ModelArguments:
metadata={"help": "Whether or not to disable CUDA graph in the vLLM engine."},
)
vllm_max_lora_rank: int = field(
default=8,
default=32,
metadata={"help": "Maximum rank of all LoRAs in the vLLM engine."},
)
offload_folder: str = field(

5
src/llamafactory/hparams/parser.py
View File

@@ -1,6 +1,6 @@
# Copyright 2024 HuggingFace Inc. and the LlamaFactory team.
#
# This code is inspired by HuggingFace's transformers library.
# This code is inspired by the HuggingFace's transformers library.
# https://github.com/huggingface/transformers/blob/v4.40.0/examples/pytorch/language-modeling/run_clm.py
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -90,6 +90,9 @@ def _verify_model_args(model_args: "ModelArguments", finetuning_args: "Finetunin
if finetuning_args.finetuning_type != "lora":
raise ValueError("Quantization is only compatible with the LoRA method.")
if finetuning_args.use_pissa:
raise ValueError("Please use scripts/pissa_init.py for quantized PiSSA.")
if model_args.resize_vocab:
raise ValueError("Cannot resize embedding layers of a quantized model.")

25
src/llamafactory/model/adapter.py
View File

@@ -179,8 +179,16 @@ def _setup_lora_tuning(
else:
adapter_to_merge = model_args.adapter_name_or_path
init_kwargs = {
"subfolder": model_args.adapter_folder,
"offload_folder": model_args.offload_folder,
"cache_dir": model_args.cache_dir,
"revision": model_args.model_revision,
"token": model_args.hf_hub_token,
}
for adapter in adapter_to_merge:
model: "LoraModel" = PeftModel.from_pretrained(model, adapter, offload_folder=model_args.offload_folder)
model: "LoraModel" = PeftModel.from_pretrained(model, adapter, **init_kwargs)
model = model.merge_and_unload()
if len(adapter_to_merge) > 0:
@@ -190,12 +198,7 @@ def _setup_lora_tuning(
if model_args.use_unsloth:
model = load_unsloth_peft_model(config, model_args, is_trainable=is_trainable)
else:
model = PeftModel.from_pretrained(
model,
adapter_to_resume,
is_trainable=is_trainable,
offload_folder=model_args.offload_folder,
)
model = PeftModel.from_pretrained(model, adapter_to_resume, is_trainable=is_trainable, **init_kwargs)
logger.info("Loaded adapter(s): {}".format(",".join(model_args.adapter_name_or_path)))
@@ -242,6 +245,14 @@ def _setup_lora_tuning(
if model_args.use_unsloth:
model = get_unsloth_peft_model(model, model_args, peft_kwargs)
else:
if finetuning_args.pissa_init:
if finetuning_args.pissa_iter == -1:
logger.info("Using PiSSA initialization.")
peft_kwargs["init_lora_weights"] = "pissa"
else:
logger.info("Using PiSSA initialization with FSVD steps {}.".format(finetuning_args.pissa_iter))
peft_kwargs["init_lora_weights"] = "pissa_niter_{}".format(finetuning_args.pissa_iter)
lora_config = LoraConfig(
task_type=TaskType.CAUSAL_LM,
inference_mode=False,

13
src/llamafactory/train/dpo/trainer.py
View File

@@ -1,6 +1,6 @@
# Copyright 2024 HuggingFace Inc. and the LlamaFactory team.
#
# This code is inspired by HuggingFace's TRL library.
# This code is inspired by the HuggingFace's TRL library.
# https://github.com/huggingface/trl/blob/v0.8.0/trl/trainer/dpo_trainer.py
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -15,6 +15,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import warnings
from collections import defaultdict
from contextlib import nullcontext
@@ -28,7 +29,7 @@ from trl import DPOTrainer
from trl.trainer import disable_dropout_in_model
from ...extras.constants import IGNORE_INDEX
from ..trainer_utils import create_custom_optimzer, create_custom_scheduler, get_batch_logps
from ..trainer_utils import convert_pissa_adapter, create_custom_optimzer, create_custom_scheduler, get_batch_logps
if TYPE_CHECKING:
@@ -91,6 +92,9 @@ class CustomDPOTrainer(DPOTrainer):
self.ref_model = self.accelerator.prepare_model(self.ref_model, evaluation_mode=True)
self.ref_model.eval()
if finetuning_args.pissa_convert:
self.save_model(os.path.join(self.args.output_dir, "pissa_init"))
if finetuning_args.use_badam:
from badam import clip_grad_norm_for_sparse_tensor
@@ -109,8 +113,11 @@ class CustomDPOTrainer(DPOTrainer):
def _save(self, output_dir: Optional[str] = None, state_dict: Optional[Dict[str, "torch.Tensor"]] = None) -> None:
super()._save(output_dir, state_dict)
output_dir = output_dir if output_dir is not None else self.args.output_dir
if self.finetuning_args.pissa_convert:
convert_pissa_adapter(output_dir, state_dict, self.accelerator, self.model, self.args)
if self.processor is not None:
output_dir = output_dir if output_dir is not None else self.args.output_dir
getattr(self.processor, "image_processor").save_pretrained(output_dir)
def odds_ratio_loss(self, chosen_logps: "torch.Tensor", rejected_logps: "torch.Tensor") -> "torch.Tensor":

12
src/llamafactory/train/pt/trainer.py
View File

@@ -12,13 +12,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from types import MethodType
from typing import TYPE_CHECKING, Dict, Optional
from transformers import Trainer
from ...extras.logging import get_logger
from ..trainer_utils import create_custom_optimzer, create_custom_scheduler
from ..trainer_utils import convert_pissa_adapter, create_custom_optimzer, create_custom_scheduler
if TYPE_CHECKING:
@@ -42,6 +43,10 @@ class CustomTrainer(Trainer):
super().__init__(**kwargs)
self.finetuning_args = finetuning_args
self.processor = processor
if finetuning_args.pissa_convert:
self.save_model(os.path.join(self.args.output_dir, "pissa_init"))
if finetuning_args.use_badam:
from badam import clip_grad_norm_for_sparse_tensor
@@ -60,6 +65,9 @@ class CustomTrainer(Trainer):
def _save(self, output_dir: Optional[str] = None, state_dict: Optional[Dict[str, "torch.Tensor"]] = None) -> None:
super()._save(output_dir, state_dict)
output_dir = output_dir if output_dir is not None else self.args.output_dir
if self.finetuning_args.pissa_convert:
convert_pissa_adapter(output_dir, state_dict, self.accelerator, self.model, self.args)
if self.processor is not None:
output_dir = output_dir if output_dir is not None else self.args.output_dir
getattr(self.processor, "image_processor").save_pretrained(output_dir)

13
src/llamafactory/train/sft/trainer.py
View File

@@ -1,6 +1,6 @@
# Copyright 2024 HuggingFace Inc. and the LlamaFactory team.
#
# This code is inspired by HuggingFace's transformers library.
# This code is inspired by the HuggingFace's transformers library.
# https://github.com/huggingface/transformers/blob/v4.40.0/src/transformers/trainer_seq2seq.py
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -26,7 +26,7 @@ from transformers import Seq2SeqTrainer
from ...extras.constants import IGNORE_INDEX
from ...extras.logging import get_logger
from ..trainer_utils import create_custom_optimzer, create_custom_scheduler
from ..trainer_utils import convert_pissa_adapter, create_custom_optimzer, create_custom_scheduler
if TYPE_CHECKING:
@@ -51,6 +51,10 @@ class CustomSeq2SeqTrainer(Seq2SeqTrainer):
super().__init__(**kwargs)
self.finetuning_args = finetuning_args
self.processor = processor
if finetuning_args.pissa_convert:
self.save_model(os.path.join(self.args.output_dir, "pissa_init"))
if finetuning_args.use_badam:
from badam import clip_grad_norm_for_sparse_tensor
@@ -69,8 +73,11 @@ class CustomSeq2SeqTrainer(Seq2SeqTrainer):
def _save(self, output_dir: Optional[str] = None, state_dict: Optional[Dict[str, "torch.Tensor"]] = None) -> None:
super()._save(output_dir, state_dict)
output_dir = output_dir if output_dir is not None else self.args.output_dir
if self.finetuning_args.pissa_convert:
convert_pissa_adapter(output_dir, state_dict, self.accelerator, self.model, self.args)
if self.processor is not None:
output_dir = output_dir if output_dir is not None else self.args.output_dir
getattr(self.processor, "image_processor").save_pretrained(output_dir)
def prediction_step(

54
src/llamafactory/train/trainer_utils.py
View File

@@ -1,9 +1,9 @@
# Copyright 2024 HuggingFace Inc. and the LlamaFactory team.
#
# This code is inspired by the GaLore's implementation: https://github.com/jiaweizzhao/GaLore
# and the LoRA+'s implementation: https://github.com/nikhil-ghosh-berkeley/loraplus
# and the BAdam's implementation: https://github.com/Ledzy/BAdam
# and the TRL's implementation: https://github.com/huggingface/trl
# This code is inspired by the original GaLore's implementation: https://github.com/jiaweizzhao/GaLore
# and the original LoRA+'s implementation: https://github.com/nikhil-ghosh-berkeley/loraplus
# and the original BAdam's implementation: https://github.com/Ledzy/BAdam
# and the HuggingFace's TRL library: https://github.com/huggingface/trl
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
@@ -17,9 +17,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union
import torch
from peft import PeftModel
from transformers import Trainer
from transformers.optimization import get_scheduler
from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
@@ -37,6 +39,7 @@ if is_galore_available():
if TYPE_CHECKING:
from accelerate import Accelerator
from transformers import PreTrainedModel, Seq2SeqTrainingArguments
from trl import AutoModelForCausalLMWithValueHead
@@ -171,6 +174,49 @@ def create_reward_model(
return reward_model
def convert_pissa_adapter(
output_dir: str,
state_dict: Dict[str, "torch.Tensor"],
accelerator: "Accelerator",
model: "PreTrainedModel",
training_args: "Seq2SeqTrainingArguments",
) -> None:
r"""
Converts the PiSSA adapter to a LoRA adapter.
"""
pissa_init_dir = os.path.join(training_args.output_dir, "pissa_init")
pissa_backup_dir = os.path.join(output_dir, "pissa_backup")
if output_dir == pissa_init_dir:
logger.info("Initial PiSSA adatper will be saved at: {}.".format(pissa_init_dir))
unwrapped_model = accelerator.unwrap_model(model)
if isinstance(unwrapped_model, PeftModel):
init_lora_weights = getattr(unwrapped_model.peft_config["default"], "init_lora_weights")
setattr(unwrapped_model.peft_config["default"], "init_lora_weights", True)
unwrapped_model.save_pretrained(
output_dir,
state_dict=state_dict,
safe_serialization=training_args.save_safetensors,
)
setattr(unwrapped_model.peft_config["default"], "init_lora_weights", init_lora_weights)
elif output_dir == training_args.output_dir: # at the end of training
logger.info("Converted PiSSA adapter will be saved at: {}.".format(output_dir))
unwrapped_model = accelerator.unwrap_model(model)
if isinstance(unwrapped_model, PeftModel): # backup the pissa adapter for further use
unwrapped_model.save_pretrained(
pissa_backup_dir,
state_dict=state_dict,
safe_serialization=training_args.save_safetensors,
)
unwrapped_model.save_pretrained(
output_dir,
state_dict=state_dict,
safe_serialization=training_args.save_safetensors,
convert_pissa_to_lora=pissa_init_dir,
)
unwrapped_model.load_adapter(pissa_backup_dir, "default", is_trainable=True)
unwrapped_model.set_adapter("default")
def _get_decay_parameter_names(model: "PreTrainedModel") -> List[str]:
r"""
Returns a list of names of parameters with weight decay. (weights in non-layernorm layers)

9
src/llamafactory/webui/components/train.py
View File

@@ -163,10 +163,9 @@ def create_train_tab(engine: "Engine") -> Dict[str, "Component"]:
create_new_adapter = gr.Checkbox()
with gr.Row():
with gr.Column(scale=1):
use_rslora = gr.Checkbox()
use_dora = gr.Checkbox()
use_rslora = gr.Checkbox()
use_dora = gr.Checkbox()
use_pissa = gr.Checkbox()
lora_target = gr.Textbox(scale=2)
additional_target = gr.Textbox(scale=2)
@@ -179,6 +178,7 @@ def create_train_tab(engine: "Engine") -> Dict[str, "Component"]:
create_new_adapter,
use_rslora,
use_dora,
use_pissa,
lora_target,
additional_target,
}
@@ -193,6 +193,7 @@ def create_train_tab(engine: "Engine") -> Dict[str, "Component"]:
create_new_adapter=create_new_adapter,
use_rslora=use_rslora,
use_dora=use_dora,
use_pissa=use_pissa,
lora_target=lora_target,
additional_target=additional_target,
)

14
src/llamafactory/webui/locales.py
View File

@@ -732,6 +732,20 @@ LOCALES = {
"info": "使用权重分解的 LoRA。",
},
},
"use_pissa": {
"en": {
"label": "Use PiSSA",
"info": "Use PiSSA method.",
},
"ru": {
"label": "используйте PiSSA",
"info": "Используйте метод PiSSA.",
},
"zh": {
"label": "使用 PiSSA",
"info": "使用 PiSSA 方法。",
},
},
"lora_target": {
"en": {
"label": "LoRA modules (optional)",

2
src/llamafactory/webui/runner.py
View File

@@ -173,6 +173,8 @@ class Runner:
args["create_new_adapter"] = get("train.create_new_adapter")
args["use_rslora"] = get("train.use_rslora")
args["use_dora"] = get("train.use_dora")
args["pissa_init"] = get("train.use_pissa")
args["pissa_convert"] = get("train.use_pissa")
args["lora_target"] = get("train.lora_target") or "all"
args["additional_target"] = get("train.additional_target") or None