wehub/huggingface--peft
1
0
Fork 0
Code Issues Pull Requests Actions 13 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
caf324b09d1dd9b19281b4c3746a4934607dc8bd
huggingface--peft/tests/test_gpu_examples.py
T
wehub-resource-sync caf324b09d
tests / check_code_quality (push) Waiting to run
Details
tests / tests (ubuntu-latest, 3.10) (push) Blocked by required conditions
Details
tests / tests (ubuntu-latest, 3.11) (push) Blocked by required conditions
Details
Deploy "method_comparison" Gradio to Spaces / deploy (push) Waiting to run
Details
Deploy "PEFT shop" Gradio app to Spaces / deploy (push) Waiting to run
Details
tests on transformers main / tests (push) Waiting to run
Details
tests / tests (ubuntu-latest, 3.12) (push) Blocked by required conditions
Details
tests / tests (ubuntu-latest, 3.13) (push) Blocked by required conditions
Details
tests / tests (windows-latest, 3.10) (push) Blocked by required conditions
Details
tests / tests (windows-latest, 3.11) (push) Blocked by required conditions
Details
tests / tests (windows-latest, 3.12) (push) Blocked by required conditions
Details
tests / tests (windows-latest, 3.13) (push) Blocked by required conditions
Details
Secret Leaks / trufflehog (push) Waiting to run
Details
CI security linting / zizmor latest via Cargo (push) Waiting to run
Details
Build documentation / build (push) Failing after 0s
Details
chore: import upstream snapshot with attribution
2026-07-13 13:24:42 +08:00

6990 lines
275 KiB
Python
Raw Blame History

# Copyright 2023-present the HuggingFace Inc. team.
# 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 datetime
import gc
import importlib
import itertools
import os
import re
import socket
import tempfile
import unittest
import warnings
from collections import Counter, defaultdict
from copy import deepcopy
from dataclasses import dataclass
from functools import partial
from pathlib import Path
from typing import Any, Union
import numpy as np
import packaging
import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import transformers
from accelerate import infer_auto_device_map
from accelerate.test_utils.testing import get_backend, run_command
from accelerate.utils import patch_environment
from accelerate.utils.imports import is_bf16_available
from accelerate.utils.memory import clear_device_cache
from datasets import Audio, Dataset, DatasetDict, load_dataset
from packaging import version
from parameterized import parameterized
from safetensors.torch import load_file, save_file
from torch.distributed import init_process_group
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.utils.data import DataLoader
from transformers import (
AutoModelForCausalLM,
AutoModelForSeq2SeqLM,
AutoTokenizer,
BitsAndBytesConfig,
DataCollatorForLanguageModeling,
FineGrainedFP8Config,
Seq2SeqTrainer,
Seq2SeqTrainingArguments,
Trainer,
TrainerCallback,
TrainingArguments,
WhisperFeatureExtractor,
WhisperForConditionalGeneration,
WhisperProcessor,
WhisperTokenizer,
)
from transformers.pytorch_utils import Conv1D
from peft import (
AdaLoraConfig,
ArrowConfig,
EvaConfig,
FrodConfig,
HiraConfig,
LoftQConfig,
LoraConfig,
PeftModel,
PrefixTuningConfig,
PromptEncoderConfig,
PveraConfig,
RandLoraConfig,
RoadConfig,
TaskType,
VeraConfig,
create_arrow_model,
get_eva_state_dict,
get_peft_model,
get_peft_model_state_dict,
initialize_lora_eva_weights,
inject_adapter_in_model,
prepare_model_for_kbit_training,
replace_lora_weights_loftq,
set_peft_model_state_dict,
)
from peft.import_utils import (
is_diffusers_available,
is_te_available,
is_transformers_ge_v5,
is_xpu_available,
)
from peft.tuners import boft
from peft.tuners.lora import LoraLayer
from peft.tuners.tuners_utils import BaseTunerLayer
from peft.utils import SAFETENSORS_WEIGHTS_NAME, infer_device
from peft.utils.hotswap import hotswap_adapter, prepare_model_for_compiled_hotswap
from peft.utils.loftq_utils import NFQuantizer
from peft.utils.other import fsdp_auto_wrap_policy
from tests.testing_utils import hub_online_once
from .testing_utils import (
DEVICE_MAP_MAP,
device_count,
load_dataset_english_quotes,
require_aqlm,
require_bitsandbytes,
require_deterministic_for_xpu,
require_eetq,
require_gptqmodel,
require_hqq,
require_non_cpu,
require_non_xpu,
require_optimum,
require_torch_gpu,
require_torch_multi_accelerator,
require_torch_multi_gpu,
require_torchao,
torch_device,
)
device, _, _ = get_backend()
if device == "cpu":
pytest.skip(allow_module_level=True, reason="GPU tests require hardware accelerator, got CPU only")
# A full testing suite that tests all the necessary features on GPU. The tests should
# rely on the example scripts to test the features.
class FrodRuntimeOffloadMLP(torch.nn.Module):
def __init__(self):
super().__init__()
self.relu = torch.nn.ReLU()
self.lin0 = torch.nn.Linear(10, 20)
self.lin1 = torch.nn.Linear(20, 20)
self.lin2 = torch.nn.Linear(20, 20)
self.lin3 = torch.nn.Linear(20, 2)
def forward(self, inputs):
hidden = self.lin0(inputs)
hidden = self.relu(hidden)
hidden = self.lin1(hidden)
hidden = self.relu(hidden)
hidden = self.lin2(hidden)
hidden = self.relu(hidden)
return self.lin3(hidden)
@pytest.mark.single_gpu_tests
def test_frod_runtime_offload_keeps_base_weight_on_cpu_after_accelerator_move():
config = FrodConfig(target_modules=["lin1", "lin2"], runtime_offload_base_weight=True)
peft_model = get_peft_model(FrodRuntimeOffloadMLP(), config).to(torch_device)
lin1 = peft_model.base_model.model.lin1
assert lin1.get_base_layer().weight.device.type == "cpu"
assert lin1.frod_U["default"].device.type == torch_device
assert lin1.frod_lambda_l["default"].device.type == torch_device
inputs = torch.randn(5, 10, device=torch_device)
output = peft_model(inputs)
assert output.device.type == torch_device
assert lin1.get_base_layer().weight.device.type == "cpu"
@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
r"""
Directly copied from:
https://github.com/huggingface/peft/blob/main/examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb
"""
processor: Any
def __call__(self, features: list[dict[str, Union[list[int], torch.Tensor]]]) -> dict[str, torch.Tensor]:
# split inputs and labels since they have to be of different lengths and need different padding methods
# first treat the audio inputs by simply returning torch tensors
input_features = [{"input_features": feature["input_features"]} for feature in features]
batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
# get the tokenized label sequences
label_features = [{"input_ids": feature["labels"]} for feature in features]
# pad the labels to max length
labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
# replace padding with -100 to ignore loss correctly
labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
# if bos token is appended in previous tokenization step,
# cut bos token here as it's append later anyways
if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
labels = labels[:, 1:]
batch["labels"] = labels
return batch
@require_non_cpu
@require_bitsandbytes
class PeftBnbGPUExampleTests(unittest.TestCase):
r"""
A single GPU int8 + fp4 test suite, this will test if training fits correctly on a single GPU device (1x NVIDIA T4
16GB) using bitsandbytes.
The tests are the following:
- Seq2Seq model training based on:
https://github.com/huggingface/peft/blob/main/examples/int8_training/Finetune_flan_t5_large_bnb_peft.ipynb
- Causal LM model training based on:
https://github.com/huggingface/peft/blob/main/examples/int8_training/Finetune_opt_bnb_peft.ipynb
- Audio model training based on:
https://github.com/huggingface/peft/blob/main/examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb
"""
def setUp(self):
self.seq2seq_model_id = "google/flan-t5-base"
self.causal_lm_model_id = "facebook/opt-6.7b"
self.tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
self.audio_model_id = "openai/whisper-large"
def tearDown(self):
r"""
Efficient mechanism to free GPU memory after each test. Based on
https://github.com/huggingface/transformers/issues/21094
"""
clear_device_cache(garbage_collection=True)
def _check_inference_finite(self, model, batch):
# try inference without Trainer class
training = model.training
model.eval()
output = model(**batch.to(model.device))
assert torch.isfinite(output.logits).all()
model.train(training)
@pytest.mark.single_gpu_tests
def test_causal_lm_training(self):
r"""
Test the CausalLM training on a single GPU device. This test is a converted version of
https://github.com/huggingface/peft/blob/main/examples/int8_training/Finetune_opt_bnb_peft.ipynb where we train
`opt-6.7b` on `english_quotes` dataset in few steps. The test would simply fail if the adapters are not set
correctly.
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_4bit(self):
r"""
Test the CausalLM training on a single GPU device. This test is a converted version of
https://github.com/huggingface/peft/blob/main/examples/int8_training/Finetune_opt_bnb_peft.ipynb where we train
`opt-6.7b` on `english_quotes` dataset in few steps using 4bit base model. The test would simply fail if the
adapters are not set correctly.
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
def test_causal_lm_training_multi_gpu_4bit(self):
r"""
Test the CausalLM training on a multi-GPU device with 4bit base model. The test would simply fail if the
adapters are not set correctly.
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=DEVICE_MAP_MAP[self.causal_lm_model_id],
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
@require_non_cpu
def test_4bit_adalora_causalLM(self):
r"""
Tests the 4bit training with adalora
"""
model_id = "facebook/opt-350m"
# for >3 GPUs, might need: device_map={"": "cuda:0"}
model = AutoModelForCausalLM.from_pretrained(
model_id, quantization_config=BitsAndBytesConfig(load_in_4bit=True)
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)
peft_config = AdaLoraConfig(
init_r=6,
target_r=4,
tinit=2,
tfinal=2,
total_step=6,
deltaT=5,
beta1=0.3,
beta2=0.3,
orth_reg_weight=0.2,
lora_alpha=32,
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, peft_config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
batch = tokenizer(data["train"][:3]["quote"], return_tensors="pt", padding=True)
self._check_inference_finite(model, batch)
class OptimizerStepCallback(TrainerCallback):
def on_optimizer_step(self, args, state, control, **kwargs):
model.update_and_allocate(state.global_step)
step_callback = OptimizerStepCallback()
with tempfile.TemporaryDirectory() as tmp_dir:
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=6,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.add_callback(step_callback)
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
@require_non_cpu
def test_8bit_adalora_causalLM(self):
r"""
Tests the 8bit training with adalora
"""
model_id = "facebook/opt-350m"
model = AutoModelForCausalLM.from_pretrained(
model_id, quantization_config=BitsAndBytesConfig(load_in_8bit=True)
)
tokenizer = AutoTokenizer.from_pretrained(model_id)
model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)
peft_config = AdaLoraConfig(
init_r=6,
target_r=4,
tinit=2,
tfinal=2,
total_step=6,
deltaT=5,
beta1=0.3,
beta2=0.3,
orth_reg_weight=0.2,
lora_alpha=32,
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, peft_config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
batch = tokenizer(data["train"][:3]["quote"], return_tensors="pt", padding=True)
self._check_inference_finite(model, batch)
class OptimizerStepCallback(TrainerCallback):
def on_optimizer_step(self, args, state, control, **kwargs):
model.update_and_allocate(state.global_step)
step_callback = OptimizerStepCallback()
with tempfile.TemporaryDirectory() as tmp_dir:
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=6,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.add_callback(step_callback)
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
@require_torch_multi_accelerator
def test_causal_lm_training_multi_gpu(self):
r"""
Test the CausalLM training on a multi-GPU device. This test is a converted version of
https://github.com/huggingface/peft/blob/main/examples/int8_training/Finetune_opt_bnb_peft.ipynb where we train
`opt-6.7b` on `english_quotes` dataset in few steps. The test would simply fail if the adapters are not set
correctly.
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
device_map="auto",
)
print(f"device map: {model.hf_device_map}")
assert set(model.hf_device_map.values()) == set(range(device_count))
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_seq2seq_lm_training_single_gpu(self):
r"""
Test the Seq2SeqLM training on a single GPU device. This test is a converted version of
https://github.com/huggingface/peft/blob/main/examples/int8_training/Finetune_opt_bnb_peft.ipynb where we train
`flan-large` on `english_quotes` dataset in few steps. The test would simply fail if the adapters are not set
correctly.
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForSeq2SeqLM.from_pretrained(
self.seq2seq_model_id,
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
device_map={"": 0},
)
# note: transformers v5 doesn't set the device map if there's only one device
assert not hasattr(model, "hf_device_map") or set(model.hf_device_map.values()) == {0}
tokenizer = AutoTokenizer.from_pretrained(self.seq2seq_model_id)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q", "v"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
@require_torch_multi_accelerator
def test_seq2seq_lm_training_multi_gpu(self):
r"""
Test the Seq2SeqLM training on a multi-GPU device. This test is a converted version of
https://github.com/huggingface/peft/blob/main/examples/int8_training/Finetune_opt_bnb_peft.ipynb where we train
`flan-large` on `english_quotes` dataset in few steps. The test would simply fail if the adapters are not set
correctly.
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForSeq2SeqLM.from_pretrained(
self.seq2seq_model_id,
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
device_map=DEVICE_MAP_MAP[self.seq2seq_model_id],
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
tokenizer = AutoTokenizer.from_pretrained(self.seq2seq_model_id)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q", "v"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir="outputs",
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
# TODO skipping to see if this leads to single GPU tests passing
@pytest.mark.skip
@pytest.mark.single_gpu_tests
def test_audio_model_training(self):
r"""
Test the audio model training on a single GPU device. This test is a converted version of
https://github.com/huggingface/peft/blob/main/examples/int8_training/peft_bnb_whisper_large_v2_training.ipynb
"""
with tempfile.TemporaryDirectory() as tmp_dir:
dataset_name = "ybelkada/common_voice_mr_11_0_copy"
task = "transcribe"
language = "Marathi"
common_voice = DatasetDict()
common_voice["train"] = load_dataset(dataset_name, split="train+validation")
common_voice = common_voice.remove_columns(
["accent", "age", "client_id", "down_votes", "gender", "locale", "path", "segment", "up_votes"]
)
feature_extractor = WhisperFeatureExtractor.from_pretrained(self.audio_model_id)
tokenizer = WhisperTokenizer.from_pretrained(self.audio_model_id, language=language, task=task)
processor = WhisperProcessor.from_pretrained(self.audio_model_id, language=language, task=task)
common_voice = common_voice.cast_column("audio", Audio(sampling_rate=16000))
def prepare_dataset(batch):
# load and resample audio data from 48 to 16kHz
audio = batch["audio"]
# compute log-Mel input features from input audio array
batch["input_features"] = feature_extractor(
audio["array"], sampling_rate=audio["sampling_rate"]
).input_features[0]
# encode target text to label ids
batch["labels"] = tokenizer(batch["sentence"]).input_ids
return batch
common_voice = common_voice.map(
prepare_dataset, remove_columns=common_voice.column_names["train"], num_proc=2
)
data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
model = WhisperForConditionalGeneration.from_pretrained(
self.audio_model_id, quantization_config=BitsAndBytesConfig(load_in_8bit=True), device_map="auto"
)
model.config.forced_decoder_ids = None
model.config.suppress_tokens = []
model = prepare_model_for_kbit_training(model)
# as Whisper model uses Conv layer in encoder, checkpointing disables grad computation
# to avoid this, make the inputs trainable
def make_inputs_require_grad(module, input, output):
output.requires_grad_(True)
model.model.encoder.conv1.register_forward_hook(make_inputs_require_grad)
config = LoraConfig(
r=32, lora_alpha=64, target_modules=["q_proj", "v_proj"], lora_dropout=0.05, bias="none"
)
model = get_peft_model(model, config)
model.print_trainable_parameters()
training_args = Seq2SeqTrainingArguments(
output_dir=tmp_dir, # change to a repo name of your choice
per_device_train_batch_size=8,
gradient_accumulation_steps=1, # increase by 2x for every 2x decrease in batch size
learning_rate=1e-3,
warmup_steps=2,
max_steps=3,
fp16=True,
per_device_eval_batch_size=8,
generation_max_length=128,
logging_steps=25,
remove_unused_columns=False, # required as the PeftModel forward doesn't have the signature of the wrapped model's forward
label_names=["labels"], # same reason as above
)
trainer = Seq2SeqTrainer(
args=training_args,
model=model,
train_dataset=common_voice["train"],
data_collator=data_collator,
tokenizer=processor.feature_extractor,
)
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_4bit_non_default_adapter_name(self):
# See PR 1294
config = LoraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
bias="none",
task_type="CAUSAL_LM",
)
# default adapter name
model = AutoModelForCausalLM.from_pretrained(
"peft-internal-testing/opt-125m",
device_map="auto",
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
)
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, config)
n_trainable_default, n_total_default = model.get_nb_trainable_parameters()
# other adapter name
model = AutoModelForCausalLM.from_pretrained(
"peft-internal-testing/opt-125m",
device_map="auto",
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
)
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, config, adapter_name="other")
n_trainable_other, n_total_other = model.get_nb_trainable_parameters()
assert n_trainable_other > 0
# sanity check
assert n_trainable_default == n_trainable_other
assert n_total_default == n_total_other
@pytest.mark.single_gpu_tests
def test_8bit_non_default_adapter_name(self):
# See PR 1294
config = LoraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
bias="none",
task_type="CAUSAL_LM",
)
# default adapter name
model = AutoModelForCausalLM.from_pretrained(
"peft-internal-testing/opt-125m",
device_map="auto",
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
)
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, config)
n_trainable_default, n_total_default = model.get_nb_trainable_parameters()
# other adapter name
model = AutoModelForCausalLM.from_pretrained(
"peft-internal-testing/opt-125m",
device_map="auto",
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
)
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, config, adapter_name="other")
n_trainable_other, n_total_other = model.get_nb_trainable_parameters()
assert n_trainable_other > 0
# sanity check
assert n_trainable_default == n_trainable_other
assert n_total_default == n_total_other
@pytest.mark.single_gpu_tests
def test_causal_lm_training_4bit_dora(self):
r"""
Same as test_causal_lm_training_4bit but with DoRA
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
use_dora=True,
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
def test_causal_lm_training_multi_gpu_4bit_dora(self):
r"""
Same as test_causal_lm_training_multi_gpu_4bit but with DoRA
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=DEVICE_MAP_MAP[self.causal_lm_model_id],
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
use_dora=True,
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_8bit_dora(self):
r"""
Same as test_causal_lm_training_4bit_dora but with 8bit
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
use_dora=True,
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
def test_causal_lm_training_multi_gpu_8bit_dora(self):
r"""
Same as test_causal_lm_training_multi_gpu_4bit_dora but with 8bit
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=DEVICE_MAP_MAP[self.causal_lm_model_id],
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
use_dora=True,
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_gpt2_dora(self):
r"""
Same as test_causal_lm_training_4bit but with DoRA
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained("gpt2", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
use_dora=True,
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@parameterized.expand(["4bit", "8bit"])
def test_initialize_dora_with_bnb_on_cpu(self, kbit):
# 1674
# The issue is that to initialize DoRA, we need to dequantize the weights. That only works on GPU for bnb.
# Therefore, initializing DoRA with bnb on CPU used to fail.
model_id = "peft-internal-testing/opt-125m"
if kbit == "4bit":
bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4")
elif kbit == "8bit":
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
else:
raise ValueError("Only 4bit and 8bit bnb allowed")
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config)
model = model.cpu() # ensure that we're on CPU
# sanity check that all weights are on CPU
weights_not_cpu = [name for name, p in model.named_parameters() if p.device != torch.device("cpu")]
assert not weights_not_cpu
lora_config = LoraConfig(use_dora=True)
# should not raise
peft_model = get_peft_model(model, lora_config)
# check that the weights are still on CPU
weights_not_cpu = [name for name, p in peft_model.named_parameters() if p.device != torch.device("cpu")]
assert not weights_not_cpu
@pytest.mark.single_gpu_tests
def test_causal_lm_training_vera(self):
r"""
Same as test_causal_lm_training but with VeRA
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = VeraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
vera_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_pvera(self):
r"""
Same as test_causal_lm_training but with PVeRA
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = PveraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
pvera_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_4bit_vera(self):
r"""
Same as test_causal_lm_training_4bit but with VeRA
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = VeraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
vera_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_4bit_pvera(self):
r"""
Same as test_causal_lm_training_4bit but with PVeRA
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = PveraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
pvera_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
def test_causal_lm_training_multi_gpu_vera(self):
r"""
Same as test_causal_lm_training_multi_gpu but with VeRA
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=DEVICE_MAP_MAP[self.causal_lm_model_id],
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = VeraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
vera_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
def test_causal_lm_training_multi_gpu_pvera(self):
r"""
Same as test_causal_lm_training_multi_gpu but with PVeRA
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=DEVICE_MAP_MAP[self.causal_lm_model_id],
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = PveraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
vera_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
def test_causal_lm_training_multi_gpu_4bit_vera(self):
r"""
Same as test_causal_lm_training_multi_gpu_4bit but with VeRA
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=DEVICE_MAP_MAP[self.causal_lm_model_id],
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = VeraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
vera_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
def test_causal_lm_training_multi_gpu_4bit_pvera(self):
r"""
Same as test_causal_lm_training_multi_gpu_4bit but with PVeRA
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=DEVICE_MAP_MAP[self.causal_lm_model_id],
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = PveraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
pvera_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_8bit_randlora(self):
r"""
Same as test_causal_lm_training but with RandLora
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = RandLoraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
randlora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset("ybelkada/english_quotes_copy")
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_4bit_randlora(self):
r"""
Same as test_causal_lm_training_4bit but with RandLora
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = RandLoraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
randlora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset("ybelkada/english_quotes_copy")
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
def test_causal_lm_training_multi_gpu_8bit_randlora(self):
r"""
Same as test_causal_lm_training_multi_gpu but with RandLoRA
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=DEVICE_MAP_MAP[self.causal_lm_model_id],
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = RandLoraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
randlora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset("Abirate/english_quotes")
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
def test_causal_lm_training_multi_gpu_4bit_randlora(self):
r"""
Same as test_causal_lm_training_multi_gpu_4bit but with RandLora
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=DEVICE_MAP_MAP[self.causal_lm_model_id],
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = RandLoraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
randlora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset("Abirate/english_quotes")
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_8bit_road(self):
r"""
Same as test_causal_lm_training but with RoAd
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = RoadConfig(
variant="road_1",
target_modules=["q_proj", "v_proj"],
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset("ybelkada/english_quotes_copy")
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=1e-3,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_4bit_road(self):
r"""
Same as test_causal_lm_training_4bit but with RoAd
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = RoadConfig(
variant="road_1",
target_modules=["q_proj", "v_proj"],
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset("ybelkada/english_quotes_copy")
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=1e-3,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
def test_causal_lm_training_multi_gpu_8bit_road(self):
r"""
Same as test_causal_lm_training_multi_gpu but with RoAd
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=DEVICE_MAP_MAP[self.causal_lm_model_id],
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = RoadConfig(
variant="road_1",
target_modules=["q_proj", "v_proj"],
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset("Abirate/english_quotes")
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=1e-3,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
def test_causal_lm_training_multi_gpu_4bit_road(self):
r"""
Same as test_causal_lm_training_multi_gpu_4bit but with RoAd
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=DEVICE_MAP_MAP[self.causal_lm_model_id],
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = RoadConfig(
variant="road_1",
target_modules=["q_proj", "v_proj"],
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset("Abirate/english_quotes")
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=1e-3,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_lora_resize_embeddings_trainable_tokens(self):
r"""
Test LoRA with trainable tokens on a resized embedding matrix
"""
with tempfile.TemporaryDirectory() as tmp_dir:
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_quant_storage=torch.float16,
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=bnb_config,
device_map="auto",
)
# add 2 new tokens
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
new_tokens = ["<think>", "</think>"]
tokenizer.add_special_tokens({"additional_special_tokens": new_tokens})
trainable_token_indices = [tokenizer.vocab[token] for token in new_tokens]
cur_emb_size = model.model.decoder.embed_tokens.weight.shape[0]
model.resize_token_embeddings(max(tokenizer.vocab_size, cur_emb_size))
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
trainable_token_indices={"embed_tokens": trainable_token_indices},
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
def tokenize(samples):
# add new tokens to samples
samples = [f"<think>{row}</think>" for row in samples["quote"]]
return tokenizer(samples)
data = data.map(tokenize, batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
# higher learning rate, as embeddings are a bit slow to update
learning_rate=1e-3,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
# ensure that the new trainable tokens have been updated
embedding = model.base_model.model.model.decoder.embed_tokens
tol = 1e-4
assert not torch.allclose(
embedding.token_adapter.trainable_tokens_delta["default"],
embedding.original_module.weight[trainable_token_indices],
atol=tol,
rtol=tol,
)
# check size of the checkpoint, should be small since the embedding matrix does not need to be stored
stat = os.stat(os.path.join(tmp_dir, SAFETENSORS_WEIGHTS_NAME))
embed_params = model.base_model.model.model.decoder.embed_tokens.original_module.weight.numel()
# fp32 -> 4x
emb_file_size = 4 * embed_params
assert stat.st_size < emb_file_size
# sanity check: assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_hira(self):
r"""
Test the CausalLM training on a single GPU device. This test is a converted version of
https://github.com/huggingface/peft/blob/main/examples/int8_training/Finetune_opt_bnb_peft.ipynb where we train
`opt-6.7b` on `english_quotes` dataset in few steps. The test would simply fail if the adapters are not set
correctly.
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = HiraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
hira_dropout=0.05,
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_4bit_hira(self):
r"""
Test the CausalLM training on a single GPU device. This test is a converted version of
https://github.com/huggingface/peft/blob/main/examples/int8_training/Finetune_opt_bnb_peft.ipynb where we train
`opt-6.7b` on `english_quotes` dataset in few steps using 4bit base model. The test would simply fail if the
adapters are not set correctly.
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_4bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
config = HiraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
hira_dropout=0.05,
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@require_non_cpu
@require_gptqmodel
@require_optimum
class PeftGPTQGPUTests(unittest.TestCase):
r"""
GPT-QModel + PEFT tests
"""
def setUp(self):
from transformers import GPTQConfig
from transformers.utils.quantization_config import AwqBackend
self.causal_lm_model_id = "marcsun13/opt-350m-gptq-4bit"
self.quantization_config = GPTQConfig(bits=4, backend=AwqBackend.AUTO_TRAINABLE)
self.tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
def tearDown(self):
r"""
Efficient mechanism to free GPU memory after each test. Based on
https://github.com/huggingface/transformers/issues/21094
"""
clear_device_cache(garbage_collection=True)
def _check_inference_finite(self, model, batch):
# try inference without Trainer class
training = model.training
model.eval()
output = model(**batch.to(model.device))
assert torch.isfinite(output.logits).all()
model.train(training)
@pytest.mark.single_gpu_tests
def test_causal_lm_training(self):
r"""
Test the CausalLM training on a single GPU device. The test would simply fail if the adapters are not set
correctly.
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
dtype=torch.float16,
device_map="auto",
quantization_config=self.quantization_config,
)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_adalora_causalLM(self):
r"""
Tests the gptq training with adalora
"""
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
dtype=torch.float16,
device_map="auto",
quantization_config=self.quantization_config,
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
model = prepare_model_for_kbit_training(model)
peft_config = AdaLoraConfig(
init_r=6,
target_r=4,
tinit=2,
tfinal=2,
total_step=6,
deltaT=5,
beta1=0.3,
beta2=0.3,
orth_reg_weight=0.2,
lora_alpha=32,
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, peft_config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
batch = tokenizer(data["train"][:3]["quote"], return_tensors="pt", padding=True)
self._check_inference_finite(model, batch)
class OptimizerStepCallback(TrainerCallback):
def on_optimizer_step(self, args, state, control, **kwargs):
model.update_and_allocate(state.global_step)
step_callback = OptimizerStepCallback()
with tempfile.TemporaryDirectory() as tmp_dir:
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=6,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
trainer.add_callback(step_callback)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_gptq_qalora(self):
"""
Test QALoRA with GPTQ quantization. The test would simply fail if the adapters are not set correctly.
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
dtype=torch.float16,
device_map="auto",
quantization_config=self.quantization_config,
)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
use_qalora=True,
qalora_group_size=32,
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
@require_torch_multi_accelerator
def test_causal_lm_training_multi_gpu(self):
r"""
Test the CausalLM training on a multi-GPU device. The test would simply fail if the adapters are not set
correctly.
"""
device_map = {
"model.decoder.embed_tokens": 0,
"lm_head": 0,
"model.decoder.embed_positions": 0,
"model.decoder.project_out": 0,
"model.decoder.project_in": 0,
"model.decoder.layers.0": 0,
"model.decoder.layers.1": 0,
"model.decoder.layers.2": 0,
"model.decoder.layers.3": 0,
"model.decoder.layers.4": 0,
"model.decoder.layers.5": 0,
"model.decoder.layers.6": 1,
"model.decoder.layers.7": 1,
"model.decoder.layers.8": 1,
"model.decoder.layers.9": 1,
"model.decoder.layers.10": 1,
"model.decoder.layers.11": 1,
"model.decoder.final_layer_norm": 1,
}
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
dtype=torch.float16,
device_map=device_map,
quantization_config=self.quantization_config,
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_non_default_adapter_name(self):
# See issue 1346
config = LoraConfig(
r=16,
target_modules=["q_proj", "v_proj"],
task_type="CAUSAL_LM",
)
# default adapter name
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
dtype=torch.float16,
device_map="auto",
quantization_config=self.quantization_config,
)
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, config)
n_trainable_default, n_total_default = model.get_nb_trainable_parameters()
# other adapter name
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
dtype=torch.float16,
device_map="auto",
quantization_config=self.quantization_config,
)
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, config, adapter_name="other")
n_trainable_other, n_total_other = model.get_nb_trainable_parameters()
assert n_trainable_other > 0
# sanity check
assert n_trainable_default == n_trainable_other
assert n_total_default == n_total_other
@require_non_cpu
class TestOffloadSave:
causal_lm_model_id = "gpt2"
@pytest.fixture(scope="class")
def tear_down(self):
r"""
Efficient mechanism to free GPU memory after each test. Based on
https://github.com/huggingface/transformers/issues/21094
"""
yield
clear_device_cache(garbage_collection=True)
def test_offload_load(self, tmp_path):
r"""
Test the loading of a LoRA model with CPU- and disk-offloaded modules
"""
torch.manual_seed(0)
model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
memory_limits = {"cpu": "0.4GIB"} # no "disk" for PeftModel.from_pretrained() compatibility
# offload around half of all transformer modules to the disk
device_map = infer_auto_device_map(model, max_memory=memory_limits)
assert "cpu" in device_map.values()
assert "disk" in device_map.values()
config = LoraConfig(task_type="CAUSAL_LM", init_lora_weights=False, target_modules=["c_attn"])
model = get_peft_model(model, config)
model.save_pretrained(tmp_path)
del model
model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id, device_map="cpu")
lora_model = PeftModel.from_pretrained(model, tmp_path).eval()
input_tokens = tokenizer.encode("Four score and seven years ago", return_tensors="pt")
output = lora_model(input_tokens)[0]
# load the model with device_map
offloaded_model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id, device_map=device_map, offload_folder=tmp_path
)
assert len({p.device for p in offloaded_model.parameters()}) == 2 # 'cpu' and 'meta'
offloaded_lora_model = PeftModel.from_pretrained(
offloaded_model, tmp_path, max_memory=memory_limits, offload_folder=tmp_path
).eval()
offloaded_output = offloaded_lora_model(input_tokens)[0]
assert torch.allclose(output, offloaded_output, atol=1e-5)
@pytest.mark.single_gpu_tests
def test_offload_merge(self, tmp_path):
r"""
Test merging, unmerging, and unloading of a model with CPU- and disk- offloaded modules.
"""
torch.manual_seed(0)
model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
memory_limits = {0: "0.2GIB", "cpu": "0.2GIB"} # no "disk" for PeftModel.from_pretrained() compatibility
# offloads around half of all transformer modules
device_map = infer_auto_device_map(model, max_memory=memory_limits)
assert 0 in device_map.values()
assert "cpu" in device_map.values()
assert "disk" in device_map.values()
config = LoraConfig(task_type="CAUSAL_LM", init_lora_weights=False, target_modules=["c_attn"])
model = get_peft_model(model, config)
model.save_pretrained(tmp_path)
del model
# load the model with device_map
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id, device_map=device_map, offload_folder=tmp_path
).eval()
assert len({p.device for p in model.parameters()}) == 2
model = PeftModel.from_pretrained(model, tmp_path, max_memory=memory_limits, offload_folder=tmp_path)
input_tokens = tokenizer.encode("Four score and seven years ago", return_tensors="pt")
model.eval()
# test peft model adapter merge
pre_merge_olayer = model(input_tokens)[0]
model.merge_adapter()
post_merge_olayer = model(input_tokens)[0]
assert torch.allclose(post_merge_olayer, pre_merge_olayer)
# test peft model adapter unmerge
model.unmerge_adapter()
post_unmerge_olayer = model(input_tokens)[0]
assert torch.allclose(post_unmerge_olayer, pre_merge_olayer)
# test LoRA merge and unload
model = model.merge_and_unload()
post_unload_merge_olayer = model(input_tokens)[0]
assert torch.allclose(post_unload_merge_olayer, pre_merge_olayer)
@pytest.mark.skipif(not (torch.cuda.is_available() or is_xpu_available()), reason="test requires a GPU or XPU")
@pytest.mark.single_gpu_tests
class TestPiSSA:
r"""
Tests for PiSSA to ensure that it reduces the quantization error compared to normal LoRA quantization.
"""
# The error factor indicates by how much the quantization error should be decreased when using PiSSA compared to
# quantization without PiSSA. Thus 1.03 means that the error should be decreased by 3% at least. This is a very
# conservative value to prevent flakiness, in practice most gains are > 1.5
error_factor = 1.03
def quantize_model(self, model, num_bits=4, device="cuda"):
# Quantize the `weight.data` of the linear layer in the model to `num_bits` and store it with full precision.
quantizer = NFQuantizer(num_bits=num_bits, device=device, method="normal", block_size=64)
for name, module in model.named_modules():
if isinstance(module, (torch.nn.Linear, Conv1D)) and "lm_head" not in name:
quantized_weight, max_abs, shape = quantizer.quantize_block(module.weight.data.to(device))
module.weight.data = quantizer.dequantize_block(quantized_weight, max_abs, shape)
return model
def nuclear_norm(self, base_model, quantized_model):
# Calculate the nuclear norm (sum of singular values) of the error matrices between the `quantized_model` and the `base_model`.
error_list = []
for name, module in base_model.named_modules():
if isinstance(module, (torch.nn.Linear, Conv1D)) and "lm_head" not in name:
quant_module = quantized_model.get_submodule(name)
error_list.append(torch.linalg.svdvals(module.weight.data - quant_module.weight.data).sum())
return torch.Tensor(error_list).sum()
def get_errors(
self,
tmp_path,
bits=4,
device="cuda",
model_id="peft-internal-testing/tiny-random-BloomForCausalLM",
):
# Comparing the quantized LoRA model to the base model, vs the PiSSA quantized model to the base model.
# We expect the PiSSA quantized model to have less error than the normal LoRA quantized model.
cls = AutoModelForSeq2SeqLM if "t5" in str(model_id) else AutoModelForCausalLM
base_model = cls.from_pretrained(model_id).eval().to(device)
task_type = TaskType.SEQ_2_SEQ_LM if base_model.config.is_encoder_decoder else TaskType.CAUSAL_LM
# logits from the normal quantized LoRA model
target_modules = "all-linear" if task_type != TaskType.SEQ_2_SEQ_LM else ["o", "k", "wi", "q", "v"]
lora_config = LoraConfig(task_type=task_type, target_modules=target_modules)
qlora_model = self.quantize_model(cls.from_pretrained(model_id).eval().to(device), bits, device)
qlora_model = get_peft_model(
qlora_model,
lora_config,
)
qlora_model = qlora_model.merge_and_unload()
qlora_error = self.nuclear_norm(base_model, qlora_model)
del qlora_model
clear_device_cache(garbage_collection=True)
# logits from quantized LoRA model using PiSSA
lora_config = LoraConfig(
task_type=task_type,
init_lora_weights="pissa",
target_modules=target_modules,
)
pissa_model = cls.from_pretrained(model_id).eval().to(device)
pissa_model = get_peft_model(pissa_model, lora_config)
# save LoRA weights, they should be initialized such that they minimize the quantization error
pissa_model.base_model.peft_config["default"].init_lora_weights = True
pissa_model.save_pretrained(tmp_path / "pissa_model")
pissa_model = pissa_model.unload()
pissa_model.save_pretrained(tmp_path / "residual_model")
del pissa_model
clear_device_cache(garbage_collection=True)
# now load quantized model and apply PiSSA-initialized weights on top
qpissa_model = self.quantize_model(
cls.from_pretrained(tmp_path / "residual_model").eval().to(device), bits, device
)
qpissa_model = PeftModel.from_pretrained(qpissa_model, tmp_path / "pissa_model")
qpissa_model = qpissa_model.merge_and_unload()
qpissa_error = self.nuclear_norm(base_model, qpissa_model)
del qpissa_model
clear_device_cache(garbage_collection=True)
assert qlora_error > 0.0
assert qpissa_error > 0.0
# next, check that PiSSA quantization errors are smaller than LoRA errors by a certain margin
assert qpissa_error < (qlora_error / self.error_factor)
@pytest.mark.parametrize("device", [torch_device, "cpu"])
def test_bloomz_pissa_4bit(self, device, tmp_path):
# In this test, we compare the logits of the base model, the quantized LoRA model, and the quantized model
# using PiSSA. When quantizing, we expect a certain level of error. However, we expect the PiSSA quantized
# model to have less error than the normal LoRA quantized model. Note that when using normal LoRA, the
# quantization error is simply the error from quantization without LoRA, as LoRA is a no-op before training.
# We still apply LoRA for the test for consistency.
self.get_errors(bits=4, device=device, tmp_path=tmp_path)
@pytest.mark.parametrize("device", [torch_device, "cpu"])
def test_bloomz_pissa_8bit(self, device, tmp_path):
# Same test as test_bloomz_pissa_4bit but with 8 bits.
self.get_errors(bits=8, device=device, tmp_path=tmp_path)
@pytest.mark.parametrize("device", [torch_device, "cpu"])
def test_t5_pissa_4bit(self, device, tmp_path):
self.get_errors(bits=4, device=device, model_id="t5-small", tmp_path=tmp_path)
@pytest.mark.parametrize("device", [torch_device, "cpu"])
def test_t5_pissa_8bit(self, device, tmp_path):
self.get_errors(bits=8, device=device, model_id="t5-small", tmp_path=tmp_path)
@pytest.mark.parametrize("device", [torch_device, "cpu"])
def test_gpt2_pissa_4bit(self, device, tmp_path):
# see 2104
self.get_errors(bits=4, device=device, model_id="gpt2", tmp_path=tmp_path)
@pytest.mark.parametrize("device", [torch_device, "cpu"])
def test_gpt2_pissa_8bit(self, device, tmp_path):
# see 2104
self.get_errors(bits=8, device=device, model_id="gpt2", tmp_path=tmp_path)
@require_bitsandbytes
def test_lora_pissa_conversion_same_output_after_loading_with_quantization(self, tmp_path):
# A copy of the test `test_lora_pissa_conversion_same_output_after_loading` in peft/tests/test_initialization.py,
# that would fail if bitsandbytes quantization is used because Quant(W_res) + AB !=Quant(W) + \Delta(AB).
import bitsandbytes as bnb
torch.manual_seed(0)
data = torch.rand(10, 1000).to(torch_device)
class MyModule(torch.nn.Module):
def __init__(self):
super().__init__()
# choose a large weight so that averages are close to expected values
self.linear = torch.nn.Linear(1000, 1000)
self.embed = torch.nn.Embedding(1000, 1000)
self.conv2d = torch.nn.Conv2d(100, 100, 3)
def forward(self, x):
x_int = (100 * x).int()
x_4d = x.flatten().reshape(1, 100, 10, 10)
return self.linear(x), self.embed(x_int), self.conv2d(x_4d)
model = MyModule().to(torch_device)
output_base = model(data)[0]
config = LoraConfig(init_lora_weights="pissa", target_modules=["linear"], r=8)
peft_model = get_peft_model(deepcopy(model), config)
# save the initial model
peft_model.peft_config["default"].init_lora_weights = True
peft_model.save_pretrained(tmp_path / "init-model")
peft_model = peft_model.unload()
torch.save(peft_model.state_dict(), tmp_path / "residual-model")
del peft_model
# create 4bit base model
base_model = deepcopy(model)
base_model.load_state_dict(torch.load(tmp_path / "residual-model"))
# sanity check: the base model weights were indeed changed
tol = 1e-06
assert not torch.allclose(model.linear.weight, base_model.linear.weight, atol=tol, rtol=tol)
# quantize the linear layer
linear4bit = bnb.nn.Linear4bit(base_model.linear.in_features, base_model.linear.out_features)
linear4bit.load_state_dict(base_model.linear.state_dict())
linear4bit.to(0)
base_model.linear = linear4bit
peft_model = PeftModel.from_pretrained(deepcopy(base_model), tmp_path / "init-model")
output_quantized_pissa = peft_model(data)[0]
# sanity check
tol = 1e-06
assert not torch.allclose(output_base, output_quantized_pissa, atol=tol, rtol=tol)
# modify the weights, or else the adapter performs an identity transformation
peft_model.base_model.linear.lora_B["default"].weight.data *= 2.0
output_finetuned_pissa = peft_model(data)[0]
# sanity check
tol = 1e-06
assert not torch.allclose(output_quantized_pissa, output_finetuned_pissa, atol=tol, rtol=tol)
# save the model normally
peft_model.save_pretrained(tmp_path / "pissa-model")
model_loaded = PeftModel.from_pretrained(deepcopy(base_model), tmp_path / "pissa-model")
output_loaded = model_loaded(data)[0]
assert torch.allclose(output_finetuned_pissa, output_loaded, atol=tol, rtol=tol)
# sanity check: ranks should still be 8 as initially
assert model_loaded.peft_config["default"].r == 8
assert model_loaded.base_model.model.linear.lora_A["default"].weight.shape[0] == 8
# save the model with conversion
peft_model.save_pretrained(
tmp_path / "pissa-model-converted", path_initial_model_for_weight_conversion=tmp_path / "init-model"
)
model_converted = PeftModel.from_pretrained(deepcopy(model), tmp_path / "pissa-model-converted")
output_converted = model_converted(data)[0]
# rank should be double of what it was initially
assert model_converted.peft_config["default"].r == 16
assert model_converted.base_model.model.linear.lora_A["default"].weight.shape[0] == 16
# base model weights should be the same as the initial model
assert torch.allclose(
model.linear.weight, model_converted.base_model.model.linear.base_layer.weight, atol=tol, rtol=tol
)
# This check is expected to fail when using bnb
assert not torch.allclose(output_finetuned_pissa, output_converted, atol=tol, rtol=tol)
@pytest.mark.skipif(not (torch.cuda.is_available() or is_xpu_available()), reason="test requires a GPU or XPU")
@pytest.mark.single_gpu_tests
class TestOLoRA:
r"""
Tests for OLoRA to ensure that it reduces the quantization error compared to normal LoRA quantization.
"""
# The error factor indicates by how much the quantization error should be decreased when using OLoRA compared to
# quantization without OLoRA. Thus 1.03 means that the error should be decreased by 3% at least. This is a very
# conservative value to prevent flakiness, in practice most gains are > 1.5
error_factor = 1.2
def quantize_model(self, model, num_bits=4, device="cuda"):
# Quantize the `weight.data` of the linear layer in the model to `num_bits` and store it with full precision.
quantizer = NFQuantizer(num_bits=num_bits, device=device, method="normal", block_size=64)
for name, module in model.named_modules():
if isinstance(module, torch.nn.Linear) and "lm_head" not in name:
quantized_weight, max_abs, shape = quantizer.quantize_block(module.weight.data.to(device))
module.weight.data = quantizer.dequantize_block(quantized_weight, max_abs, shape)
return model
def nuclear_norm(self, base_model, quantized_model):
# Calculate the nuclear norm (sum of singular values) of the error matrices between the `quantized_model` and the `base_model`.
error_list = []
for name, module in base_model.named_modules():
if isinstance(module, torch.nn.Linear) and "lm_head" not in name:
quant_module = quantized_model.get_submodule(name)
error_list.append(torch.linalg.svdvals(module.weight.data - quant_module.weight.data).sum())
return torch.Tensor(error_list).sum()
def get_errors(
self,
tmp_path,
bits=4,
device="cuda",
model_id="peft-internal-testing/tiny-random-BloomForCausalLM",
):
# Comparing the quantized LoRA model to the base model, vs the OLoRA quantized model to the base model.
# We expect the OLoRA quantized model to have less error than the normal LoRA quantized model.
cls = AutoModelForSeq2SeqLM if "t5" in str(model_id) else AutoModelForCausalLM
base_model = cls.from_pretrained(model_id).eval().to(device)
task_type = TaskType.SEQ_2_SEQ_LM if base_model.config.is_encoder_decoder else TaskType.CAUSAL_LM
# logits from the normal quantized LoRA model
target_modules = "all-linear" if task_type != TaskType.SEQ_2_SEQ_LM else ["o", "k", "wi", "q", "v"]
lora_config = LoraConfig(task_type=task_type, target_modules=target_modules)
qlora_model = self.quantize_model(cls.from_pretrained(model_id).eval().to(device), bits, device)
qlora_model = get_peft_model(
qlora_model,
lora_config,
)
qlora_model = qlora_model.merge_and_unload()
qlora_error = self.nuclear_norm(base_model, qlora_model)
del qlora_model
clear_device_cache(garbage_collection=True)
# logits from quantized LoRA model using OLoRA
lora_config = LoraConfig(
task_type=task_type,
init_lora_weights="olora",
target_modules=target_modules,
)
olora_model = cls.from_pretrained(model_id).eval().to(device)
olora_model = get_peft_model(olora_model, lora_config)
# save LoRA weights, they should be initialized such that they minimize the quantization error
olora_model.base_model.peft_config["default"].init_lora_weights = True
olora_model.save_pretrained(tmp_path / "olora_model")
olora_model = olora_model.unload()
olora_model.save_pretrained(tmp_path / "residual_model")
del olora_model
clear_device_cache(garbage_collection=True)
# now load quantized model and apply OLoRA-initialized weights on top
qolora_model = self.quantize_model(
cls.from_pretrained(tmp_path / "residual_model").eval().to(device), bits, device
)
qolora_model = PeftModel.from_pretrained(qolora_model, tmp_path / "olora_model")
qolora_model = qolora_model.merge_and_unload()
qolora_error = self.nuclear_norm(base_model, qolora_model)
del qolora_model
clear_device_cache(garbage_collection=True)
assert qlora_error > 0.0
assert qolora_error > 0.0
# next, check that OLoRA quantization errors are smaller than LoRA errors by a certain margin
assert qolora_error < (qlora_error / self.error_factor)
@pytest.mark.parametrize("device", [torch_device, "cpu"])
def test_bloomz_olora_4bit(self, device, tmp_path):
# In this test, we compare the logits of the base model, the quantized LoRA model, and the quantized model
# using OLoRA. When quantizing, we expect a certain level of error. However, we expect the OLoRA quantized
# model to have less error than the normal LoRA quantized model. Note that when using normal LoRA, the
# quantization error is simply the error from quantization without LoRA, as LoRA is a no-op before training.
# We still apply LoRA for the test for consistency.
self.get_errors(bits=4, device=device, tmp_path=tmp_path)
@pytest.mark.parametrize("device", [torch_device, "cpu"])
def test_bloomz_olora_8bit(self, device, tmp_path):
# Same test as test_bloomz_olora_4bit but with 8 bits.
self.get_errors(bits=8, device=device, tmp_path=tmp_path)
@pytest.mark.parametrize("bits", [4, 8])
def test_olora_with_quantized_model(self, bits):
import bitsandbytes as bnb
# issue 1999
model_id = "peft-internal-testing/tiny-random-OPTForCausalLM"
if bits == 4:
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_quant_storage=torch.float16,
bnb_4bit_use_double_quant=True,
)
elif bits == 8:
bnb_config = BitsAndBytesConfig(load_in_8bit=True)
else:
raise ValueError("bits must be 4 or 8")
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(init_lora_weights="olora")
model = get_peft_model(model, config)
# check that the correct type is used for the weights
base_layer = model.base_model.model.model.decoder.layers[0].self_attn.v_proj.base_layer.weight
if bits == 4:
assert isinstance(base_layer, bnb.nn.modules.Params4bit)
else:
assert isinstance(base_layer, bnb.nn.modules.Int8Params)
inputs = torch.arange(10).unsqueeze(0).to(model.device)
logits = model(inputs).logits # does not raise
assert torch.isfinite(logits).all()
@pytest.mark.skipif(
not (torch.cuda.is_available() or is_xpu_available()), reason="test requires a hardware accelerator"
)
@pytest.mark.single_gpu_tests
@require_bitsandbytes
class TestLoftQ:
r"""
Tests for LoftQ to ensure that it reduces the quantization error compared to normal LoRA quantization.
"""
OPT_MODEL_ID = "facebook/opt-125m"
def get_error_factor(self, device, n_bits):
# The error factor indicates by how much the quantization error should be decreased when using LoftQ compared to
# quantization without LoftQ. Thus 1.03 means that the error should be decreased by 3% at least. This is a very
# conservative value to prevent flakiness, in practice most gains are > 1.5
return 1.03
def get_input(self, model_id, device):
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer(
[
"All I want is",
"All I need",
"Forever yours truly: ",
"Translate French to German: Tu l'as lu?",
"Translate German to French: Last du es?",
],
padding=True,
return_tensors="pt",
)
inputs = inputs.to(device)
return inputs
def get_base_model(self, model_id, device, **kwargs):
cls = AutoModelForSeq2SeqLM if "t5" in str(model_id) else AutoModelForCausalLM
with hub_online_once(model_id):
model = cls.from_pretrained(model_id, device_map=device, **kwargs).eval()
return model
def get_logits(self, model, inputs):
if model.config.is_encoder_decoder:
input_ids = inputs["input_ids"]
input_ids = model._shift_right(input_ids)
return model(input_ids=input_ids, decoder_input_ids=input_ids).logits
return model(**inputs).logits
def error_mask(self, error, attention_mask=None):
# Make sure to ignore padding values in the error term
if attention_mask is not None:
# attention_mask shape: [batch_size, seq_len]
# error shape: [batch_size, seq_len, vocab_size]
# apply the mask (zeros out the squared error for padding tokens)
mask = attention_mask.unsqueeze(-1).expand_as(error)
masked_error = error * mask
return masked_error.sum() / mask.sum()
return error.mean()
def mse(self, a, b, attention_mask=None):
error = torch.pow(a - b, 2)
return self.error_mask(error, attention_mask=attention_mask)
def mae(self, a, b, attention_mask=None):
error = torch.abs(a - b)
return self.error_mask(error, attention_mask=attention_mask)
def get_errors(
self,
tmp_path,
bits=4,
loftq_iter=1,
device="cuda",
model_id=None,
use_dora=False,
):
# Helper function that returns the quantization errors (MAE and MSE) when comparing the quantized LoRA model
# to the base model, vs the LoftQ quantized model to the base model. We expect the LoftQ quantized model to
# have less error than the normal LoRA quantized model. Since we compare logits, the observed error is
# already somewhat dampened because of the softmax.
torch.manual_seed(0)
model = self.get_base_model(model_id, device)
task_type = TaskType.SEQ_2_SEQ_LM if model.config.is_encoder_decoder else TaskType.CAUSAL_LM
inputs = self.get_input(model_id, device)
# the base logits are the reference, we try to match those as closely as possible
logits_base = self.get_logits(model, inputs)
# clean up
del model
clear_device_cache(garbage_collection=True)
# logits from the normal quantized LoRA model
#
# t5 has `_keep_in_fp32=["wo"]` which is why we target all linear except for "wo" - if we'd target all
# layers including wo we'd introduce quantization error that's not present by applying the mitigation
target_modules = "all-linear" if task_type != TaskType.SEQ_2_SEQ_LM else ["o", "k", "wi", "q", "v"]
lora_config = LoraConfig(task_type=task_type, use_dora=use_dora, target_modules=target_modules)
kwargs = {}
if bits == 4:
kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4")
elif bits == 8:
# threshold > 0 will introduce errors uncorrectable by static methods like LoftQ
kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True, llm_int8_threshold=0.0)
else:
raise ValueError("bits must be 4 or 8")
quantized_base_model = self.get_base_model(model_id, device, **kwargs, dtype=torch.float32)
if target_modules != "all-linear":
# make sure that the manual targeting catches all layers and doesn't target too many
layer_name = "Linear8bitLt" if bits == 8 else "Linear4bit"
quantized_suffixes = {
n.split(".")[-1] for n, p in quantized_base_model.named_modules() if p.__class__.__name__ == layer_name
}
assert set(target_modules) == quantized_suffixes
quantized_model = get_peft_model(
quantized_base_model,
lora_config,
)
torch.manual_seed(0)
logits_quantized = self.get_logits(quantized_model, inputs)
del quantized_model
del quantized_base_model
clear_device_cache(garbage_collection=True)
# logits from quantized LoRA model using LoftQ
loftq_config = LoftQConfig(loftq_bits=bits, loftq_iter=loftq_iter)
lora_config = LoraConfig(
r=64,
lora_alpha=32,
task_type=task_type,
init_lora_weights="loftq",
loftq_config=loftq_config,
use_dora=use_dora,
target_modules=target_modules,
)
model = self.get_base_model(model_id, device)
if device != "cpu":
model = model.to(device)
loftq_model = get_peft_model(model, lora_config)
if device != "cpu":
loftq_model = loftq_model.to(device)
# save LoRA weights, they should be initialized such that they minimize the quantization error
loftq_model.base_model.peft_config["default"].init_lora_weights = True
loftq_model.save_pretrained(tmp_path / "loftq_model")
loftq_model = loftq_model.unload()
loftq_model.save_pretrained(tmp_path / "base_model")
del loftq_model
clear_device_cache(garbage_collection=True)
# now load quantized model and apply LoftQ-initialized weights on top
base_model = self.get_base_model(tmp_path / "base_model", device=device, **kwargs, dtype=torch.float32)
loftq_model = PeftModel.from_pretrained(base_model, tmp_path / "loftq_model", is_trainable=True)
# TODO sanity check: model is quantized
torch.manual_seed(0)
logits_loftq = self.get_logits(loftq_model, inputs)
del loftq_model
clear_device_cache(garbage_collection=True)
mae_quantized = self.mae(logits_base, logits_quantized, attention_mask=inputs["attention_mask"])
mse_quantized = self.mse(logits_base, logits_quantized, attention_mask=inputs["attention_mask"])
mae_loftq = self.mae(logits_base, logits_loftq, attention_mask=inputs["attention_mask"])
mse_loftq = self.mse(logits_base, logits_loftq, attention_mask=inputs["attention_mask"])
return mae_quantized, mse_quantized, mae_loftq, mse_loftq
@pytest.mark.parametrize("device", [torch_device, "cpu"])
@pytest.mark.parametrize("n_bits", [4, 8], ids=["4bit", "8bit"])
@pytest.mark.parametrize("n_iter", [1, 3], ids=["1iter", "3iter"])
def test_opt_loftq(self, device, n_bits, n_iter, tmp_path):
# In this test, we compare the logits of the base model, the quantized LoRA model, and the quantized model
# using LoftQ. When quantizing, we expect a certain level of error. However, we expect the LoftQ quantized
# model to have less error than the normal LoRA quantized model. Note that when using normal LoRA, the
# quantization error is simply the error from quantization without LoRA, as LoRA is a no-op before training.
# We still apply LoRA for the test for consistency.
#
# With 5 iterations we should expect the error to be even smaller with more
# iterations, but in practice the difference is not that large, at least not for this small base model.
mae_quantized, mse_quantized, mae_loftq, mse_loftq = self.get_errors(
bits=n_bits,
loftq_iter=n_iter,
device=device,
model_id=self.OPT_MODEL_ID,
tmp_path=tmp_path,
)
# first, sanity check that all errors are > 0.0
assert mae_quantized > 0.0
assert mse_quantized > 0.0
assert mae_loftq > 0.0
assert mse_loftq > 0.0
if n_bits == 4:
# next, check that LoftQ quantization errors are smaller than LoRA errors by a certain margin
error_factor = self.get_error_factor(device, n_bits)
assert mse_loftq < (mse_quantized / error_factor)
assert mae_loftq < (mae_quantized / error_factor)
elif n_bits == 8:
# for int8 compensation we're happy to achieve parity but without additional measures we will
# not be better since int8 also quantizes the layer input which introduces error loftq can't
# compensate
assert torch.allclose(mse_loftq, mse_quantized, atol=0.06)
assert torch.allclose(mae_loftq, mae_quantized, atol=0.06)
@pytest.mark.parametrize("device", [torch_device, "cpu"])
@pytest.mark.parametrize("n_bits", [4, 8], ids=["4bit", "8bit"])
@pytest.mark.parametrize("n_iter", [1, 3], ids=["1iter", "3iter"])
def test_t5_loftq(self, device, n_bits, n_iter, tmp_path):
if n_bits == 8 and n_iter > 1:
pytest.xfail("n_iter > 1 produces strictly worse results for a yet unknown reason")
mae_quantized, mse_quantized, mae_loftq, mse_loftq = self.get_errors(
bits=n_bits, loftq_iter=n_iter, device=device, model_id="t5-small", tmp_path=tmp_path
)
# first, sanity check that all errors are > 0.0
assert mae_quantized > 0.0
assert mse_quantized > 0.0
assert mae_loftq > 0.0
assert mse_loftq > 0.0
if n_bits == 4:
# next, check that LoftQ quantization errors are smaller than LoRA errors by a certain margin
error_factor = self.get_error_factor(device, n_bits)
assert mse_loftq < (mse_quantized / error_factor)
assert mae_loftq < (mae_quantized / error_factor)
elif n_bits == 8:
# for int8 compensation we're happy to achieve parity but without additional measures we will
# not be better since int8 also quantizes the layer input which introduces error loftq can't
# compensate
assert torch.allclose(mse_loftq, mse_quantized, atol=0.06)
assert torch.allclose(mae_loftq, mae_quantized, atol=0.06)
@pytest.mark.xfail # failing for now, but having DoRA pass is only a nice-to-have, not a must, so we're good
@pytest.mark.parametrize("device", [torch_device, "cpu"])
def test_opt_loftq_4bit_dora(self, device, tmp_path):
# same as test_opt_loftq_4bit but with DoRA
mae_quantized, mse_quantized, mae_loftq, mse_loftq = self.get_errors(
bits=4,
device=device,
use_dora=True,
model_id=self.OPT_MODEL_ID,
tmp_path=tmp_path,
)
# first, sanity check that all errors are > 0.0
assert mae_quantized > 0.0
assert mse_quantized > 0.0
assert mae_loftq > 0.0
assert mse_loftq > 0.0
# next, check that LoftQ quantization errors are smaller than LoRA errors by a certain margin
factor = 3
assert mae_loftq < (mae_quantized / factor)
assert mse_loftq < (mse_quantized / factor)
@pytest.mark.parametrize("device", [torch_device, "cpu"])
def test_opt_loftq_8bit_dora(self, device, tmp_path):
# same as test_opt_loftq_8bit but with DoRA
mae_quantized, mse_quantized, mae_loftq, mse_loftq = self.get_errors(
bits=8,
device=device,
use_dora=True,
model_id=self.OPT_MODEL_ID,
tmp_path=tmp_path,
)
# first, sanity check that all errors are > 0.0
assert mae_quantized > 0.0
assert mse_quantized > 0.0
assert mae_loftq > 0.0
assert mse_loftq > 0.0
# again, parity is the best we can get here due to int8 activation quant.
assert torch.allclose(mse_loftq, mse_quantized, atol=0.03)
assert torch.allclose(mae_loftq, mae_quantized, atol=0.03)
def test_replace_lora_weights_with_loftq_using_callable(self):
"""
Test replacing LoRa weights with LoFTQ using a callable.
Using the replace_lora_weights_loftq function, we replace the LoRa weights of a bnb-quantized model with LoRA
weights initialized by LoftQ on the fly. We use a callable to decide whether to replace the weights or not.
This callable checks, for each weight, if replacing it would actually result in logits that are closer to the
original logits of the non-quantized model.
"""
torch.manual_seed(0)
model_id = "bigscience/bloomz-560m"
device = torch_device
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer("The dog was", padding=True, return_tensors="pt").to(device)
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
logits_base = model(**inputs).logits
model.save_pretrained(tmp_dir)
# load in 4bit
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config)
model = get_peft_model(model, LoraConfig(task_type="CAUSAL_LM", target_modules="all-linear"))
logits_lora = model(**inputs).logits
current_mse = float("inf")
logs = []
def my_callback(model, module_name):
"""Callable to replace weights with LoFTQ if the mse is lower than the current best one."""
nonlocal current_mse
logits = model(**inputs).logits
mse = ((logits_base - logits) ** 2).mean()
if mse < current_mse:
current_mse = mse
logs.append(True)
return True
logs.append(False)
return False
replace_lora_weights_loftq(model, model_path=tmp_dir, callback=my_callback)
logits_loftq = model(**inputs).logits
mae_lora = (logits_base - logits_lora).abs().mean()
mae_loftq = (logits_base - logits_loftq).abs().mean()
mse_lora = ((logits_base - logits_lora) ** 2).mean()
mse_loftq = ((logits_base - logits_loftq) ** 2).mean()
# check that the error was reduced by a certain margin
assert mae_loftq * 1.5 < mae_lora
assert mse_loftq * 2.5 < mse_lora
# check that the callback has returned some True and some False values
assert any(logs)
assert not all(logs)
del model
clear_device_cache(garbage_collection=True)
def test_replace_lora_weights_with_local_model(self):
# see issue 2020
torch.manual_seed(0)
model_id = "peft-internal-testing/tiny-random-OPTForCausalLM"
device = torch_device
with tempfile.TemporaryDirectory() as tmp_dir:
# save base model locally
model = AutoModelForCausalLM.from_pretrained(model_id).to(device)
model.save_pretrained(tmp_dir)
del model
# load in 4bit
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
)
# load the base model from local directory
model = AutoModelForCausalLM.from_pretrained(tmp_dir, quantization_config=bnb_config)
model = get_peft_model(model, LoraConfig())
# passing the local path directly works
replace_lora_weights_loftq(model, model_path=tmp_dir)
del model
# load the base model from local directory
model = AutoModelForCausalLM.from_pretrained(tmp_dir, quantization_config=bnb_config)
model = get_peft_model(model, LoraConfig())
# when not passing, ensure that users are made aware of the `model_path` argument
with pytest.raises(ValueError, match="model_path"):
replace_lora_weights_loftq(model)
del model
clear_device_cache(garbage_collection=True)
def test_config_no_loftq_init(self):
with pytest.warns(
UserWarning,
match="`loftq_config` specified but will be ignored when `init_lora_weights` is not 'loftq'.",
):
LoraConfig(loftq_config=LoftQConfig())
def test_config_no_loftq_config(self):
with pytest.raises(ValueError, match="`loftq_config` must be specified when `init_lora_weights` is 'loftq'."):
LoraConfig(init_lora_weights="loftq")
@require_bitsandbytes
@require_non_cpu
class MultiprocessTester(unittest.TestCase):
def test_notebook_launcher(self):
script_path = os.path.join("scripts", "launch_notebook_mp.py")
cmd = ["python", script_path]
with patch_environment(omp_num_threads=1):
run_command(cmd, env=os.environ.copy())
@require_non_cpu
class MixedPrecisionTests(unittest.TestCase):
def setUp(self):
self.causal_lm_model_id = "peft-internal-testing/opt-125m"
self.tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
self.config = LoraConfig(
r=16,
lora_alpha=32,
task_type="CAUSAL_LM",
)
data = load_dataset_english_quotes()
self.data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
def tearDown(self):
r"""
Efficient mechanism to free GPU memory after each test. Based on
https://github.com/huggingface/transformers/issues/21094
"""
clear_device_cache(garbage_collection=True)
gc.collect()
@pytest.mark.single_gpu_tests
def test_model_using_float16_with_amp_raises(self):
# This test shows the issue with using a model in fp16 and then trying to use it with mixed precision training,
# which should not use fp16.
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
dtype=torch.float16,
)
model = get_peft_model(model, self.config, autocast_adapter_dtype=False)
with tempfile.TemporaryDirectory() as tmp_dir:
trainer = Trainer(
model=model,
train_dataset=self.data["train"],
args=TrainingArguments(
fp16=True, # <= this is required for the error to be raised
output_dir=tmp_dir,
max_steps=3,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
with pytest.raises(ValueError, match="Attempting to unscale FP16 gradients."):
trainer.train()
@pytest.mark.single_gpu_tests
def test_model_using_float16_autocast_dtype(self):
# Here we use autocast_adapter_dtype=True (the default) to automatically promote the adapter weights to float32.
# No exception should be raised.
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
dtype=torch.float16,
)
model = get_peft_model(model, self.config, autocast_adapter_dtype=True)
with tempfile.TemporaryDirectory() as tmp_dir:
trainer = Trainer(
model=model,
train_dataset=self.data["train"],
args=TrainingArguments(
fp16=True, # <= this is required for the error to be raised
output_dir=tmp_dir,
max_steps=3,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
trainer.train() # does not raise
@pytest.mark.single_gpu_tests
def test_model_using_float16_explicit_cast(self):
# Same test as above but containing the fix to make it work
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
dtype=torch.float16,
)
model = get_peft_model(model, self.config, autocast_adapter_dtype=False)
# here we manually promote the adapter weights to float32
for param in model.parameters():
if param.requires_grad:
param.data = param.data.float()
dtype_counts_before = Counter(p.dtype for p in model.parameters())
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
dtype=torch.float16,
)
model = get_peft_model(model, self.config, autocast_adapter_dtype=True)
dtype_counts_after = Counter(p.dtype for p in model.parameters())
assert dtype_counts_before == dtype_counts_after
with tempfile.TemporaryDirectory() as tmp_dir:
trainer = Trainer(
model=model,
train_dataset=self.data["train"],
args=TrainingArguments(
fp16=True, # <= this is required for the error to be raised
max_steps=3,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
trainer.train() # does not raise
@pytest.mark.single_gpu_tests
def test_load_model_using_float16_with_amp_raises(self):
# Same as previous tests, but loading the adapter with PeftModel.from_pretrained instead
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
dtype=torch.float16,
)
model = get_peft_model(model, self.config, autocast_adapter_dtype=False)
with tempfile.TemporaryDirectory() as tmp_dir:
model.save_pretrained(tmp_dir)
model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id, dtype=torch.float16)
model = PeftModel.from_pretrained(model, tmp_dir, autocast_adapter_dtype=False, is_trainable=True)
trainer = Trainer(
model=model,
train_dataset=self.data["train"],
args=TrainingArguments(
fp16=True, # <= this is required for the error to be raised
output_dir=tmp_dir,
max_steps=3,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
with pytest.raises(ValueError, match="Attempting to unscale FP16 gradients."):
trainer.train()
@pytest.mark.single_gpu_tests
def test_load_model_using_float16_autocast_dtype(self):
# Same as previous tests, but loading the adapter with PeftModel.from_pretrained instead
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
dtype=torch.float16,
)
# Below, we purposefully set autocast_adapter_dtype=False so that the saved adapter uses float16. We still want
# the loaded adapter to use float32 when we load it with autocast_adapter_dtype=True.
model = get_peft_model(model, self.config, autocast_adapter_dtype=False)
# sanity check: this should have float16 adapter weights:
assert (
model.base_model.model.model.decoder.layers[0].self_attn.v_proj.lora_A["default"].weight.dtype
== torch.float16
)
with tempfile.TemporaryDirectory() as tmp_dir:
model.save_pretrained(tmp_dir)
model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id, dtype=torch.float16)
model = PeftModel.from_pretrained(model, tmp_dir, autocast_adapter_dtype=True, is_trainable=True)
# sanity check: this should NOT have float16 adapter weights:
assert (
model.base_model.model.model.decoder.layers[0].self_attn.v_proj.lora_A["default"].weight.dtype
== torch.float32
)
trainer = Trainer(
model=model,
train_dataset=self.data["train"],
args=TrainingArguments(
fp16=True, # <= this is required for the error to be raised
output_dir=tmp_dir,
max_steps=3,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
trainer.train() # does not raise
@pytest.mark.single_gpu_tests
def test_load_adapter_using_float16_autocast_dtype(self):
# Here we test the load_adapter method with autocast_adapter_dtype. We show that autocasting is prevented when
# calling load_model(..., autocast_adapter_dtype=False) and that it is enabled when calling
# load_model(..., autocast_adapter_dtype=True) (the default).
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
dtype=torch.float16,
)
# Below, we purposefully set autocast_adapter_dtype=False so that the saved adapter uses float16. We still want
# the loaded adapter to use float32 when we load it with autocast_adapter_dtype=True.
model = get_peft_model(model, self.config, autocast_adapter_dtype=False)
# sanity check: this should have float16 adapter weights:
assert (
model.base_model.model.model.decoder.layers[0].self_attn.v_proj.lora_A["default"].weight.dtype
== torch.float16
)
with tempfile.TemporaryDirectory() as tmp_dir:
model.save_pretrained(tmp_dir)
model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id, dtype=torch.float16)
# the default adapter is now in float16
model = get_peft_model(model, self.config, autocast_adapter_dtype=False)
# sanity check: this should NOT have float16 adapter weights:
assert (
model.base_model.model.model.decoder.layers[0].self_attn.v_proj.lora_A["default"].weight.dtype
== torch.float16
)
# now load the first adapter in float16 using the adapter name "loaded16"
model.load_adapter(tmp_dir, "loaded16", autocast_adapter_dtype=False)
assert (
model.base_model.model.model.decoder.layers[0].self_attn.v_proj.lora_A["loaded16"].weight.dtype
== torch.float16
)
# now load the first adapter in float32 using the adapter name "loaded32"
model.load_adapter(tmp_dir, "loaded32", autocast_adapter_dtype=True)
assert (
model.base_model.model.model.decoder.layers[0].self_attn.v_proj.lora_A["loaded32"].weight.dtype
== torch.float32
)
# training with the default adapter, which is in float16, should raise
model.set_adapter("default")
trainer = Trainer(
model=model,
train_dataset=self.data["train"],
args=TrainingArguments(
fp16=True, # <= this is required for the error to be raised
output_dir=tmp_dir,
max_steps=3,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
with pytest.raises(ValueError, match="Attempting to unscale FP16 gradients."):
trainer.train()
# training the model with the adapter "loaded16", which is in float16, should also raise
model.set_adapter("loaded16")
trainer = Trainer(
model=model,
train_dataset=self.data["train"],
args=TrainingArguments(
fp16=True, # <= this is required for the error to be raised
output_dir=tmp_dir,
max_steps=3,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
with pytest.raises(ValueError, match="Attempting to unscale FP16 gradients."):
trainer.train()
# training the model with the adapter "loaded32", which is in float32, should not raise
model.set_adapter("loaded32")
trainer = Trainer(
model=model,
train_dataset=self.data["train"],
args=TrainingArguments(
fp16=True, # <= this is required for the error to be raised
output_dir=tmp_dir,
max_steps=3,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
trainer.train() # does not raise
@require_non_xpu
@require_torch_gpu
@require_aqlm
@pytest.mark.skipif(
not version.parse(importlib.metadata.version("transformers")) >= version.parse("4.38.0"),
reason="test requires `transformers>=4.38.0`",
)
class PeftAqlmGPUTests(unittest.TestCase):
r"""
AQLM + peft tests
"""
def setUp(self):
self.causal_lm_model_id = "BlackSamorez/TinyLlama-1_1B-Chat-v1_0-AQLM-2Bit-1x16-hf"
self.tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
def tearDown(self):
r"""
Efficient mechanism to free GPU memory after each test. Based on
https://github.com/huggingface/transformers/issues/21094
"""
clear_device_cache(garbage_collection=True)
def _check_inference_finite(self, model, batch):
# try inference without Trainer class
training = model.training
model.eval()
output = model(**batch.to(model.device))
assert torch.isfinite(output.logits).all()
model.train(training)
@pytest.mark.single_gpu_tests
def test_causal_lm_training_aqlm(self):
r"""
Test the CausalLM training on a single GPU device. The test would simply fail if the adapters are not set
correctly.
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map="cuda",
dtype="auto",
)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
logging_steps=1,
output_dir=tmp_dir,
fp16=True,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@require_non_xpu
@require_torch_gpu
@require_hqq
@pytest.mark.skipif(
not version.parse(importlib.metadata.version("transformers")) >= version.parse("4.36.1"),
reason="test requires `transformers>=4.36.1`",
)
class PeftHqqGPUTests(unittest.TestCase):
r"""
HQQ + peft tests
"""
def setUp(self):
self.causal_lm_model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
self.tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
def tearDown(self):
r"""
Efficient mechanism to free GPU memory after each test. Based on
https://github.com/huggingface/transformers/issues/21094
"""
clear_device_cache(garbage_collection=True)
@pytest.mark.xfail(
reason="HQQ is not yet supported by Transformers v5",
condition=is_transformers_ge_v5,
strict=True,
raises=NotImplementedError,
)
@pytest.mark.single_gpu_tests
@parameterized.expand([False, True])
def test_causal_lm_training_hqq(self, use_dora):
r"""
Test the CausalLM training on a single GPU device. The test would simply fail if the adapters are not set
correctly.
"""
from transformers import HqqConfig
with tempfile.TemporaryDirectory() as tmp_dir:
device = "cuda"
compute_dtype = torch.float16
quant_config = HqqConfig(nbits=4, group_size=64)
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=device,
dtype=compute_dtype,
quantization_config=quant_config,
)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
use_dora=use_dora,
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
logging_steps=1,
output_dir=tmp_dir,
fp16=True,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.xfail(
reason="HQQ is not yet supported by Transformers v5",
condition=is_transformers_ge_v5,
strict=True,
raises=NotImplementedError,
)
@pytest.mark.single_gpu_tests
def test_hqq_lora_model_outputs(self):
# check that the outputs generated by HQQ with LoRA are similar to those without HQQ
from transformers import HqqConfig
device = "cuda"
compute_dtype = torch.float16
min_correlation = 0.96
# first load the model without HQQ
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=device,
dtype=compute_dtype,
)
config = LoraConfig(
target_modules=["q_proj", "v_proj"],
task_type="CAUSAL_LM",
init_lora_weights=False,
)
torch.manual_seed(0)
model = get_peft_model(model, config).eval()
inputs = self.tokenizer("The meaning of unit tests is", return_tensors="pt").to(model.device)
with torch.inference_mode():
output_normal = model(**inputs).logits
assert torch.isfinite(output_normal).all()
del model
clear_device_cache(garbage_collection=True)
# now load with HQQ
quant_config = HqqConfig(nbits=4, group_size=64)
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=device,
dtype=compute_dtype,
quantization_config=quant_config,
)
torch.manual_seed(0)
model = get_peft_model(model, config).eval()
with torch.inference_mode():
output_hqq = model(**inputs).logits
# check that outputs of HQQ are highly correlated; there are outliers, so don't check for equality
cc_matrix = torch.corrcoef(torch.stack((output_normal.float().flatten(), output_hqq.float().flatten())))
assert cc_matrix.min() > min_correlation
# check that outputs are the same after merging
cc_matrix = torch.corrcoef(torch.stack((output_normal.float().flatten(), output_hqq.float().flatten())))
assert cc_matrix.min() > min_correlation
# check outputs are the same after unmerging
model.unmerge_adapter()
with torch.inference_mode():
output_unmerged = model(**inputs).logits
cc_matrix = torch.corrcoef(torch.stack((output_normal.float().flatten(), output_unmerged.float().flatten())))
assert cc_matrix.min() > min_correlation
# check that the results are the same after saving and loading
with tempfile.TemporaryDirectory() as tmp_dir:
model.save_pretrained(tmp_dir)
del model
clear_device_cache(garbage_collection=True)
quant_config = HqqConfig(nbits=4, group_size=64)
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=device,
dtype=compute_dtype,
quantization_config=quant_config,
)
model = PeftModel.from_pretrained(model, tmp_dir)
with torch.inference_mode():
output_loaded = model(**inputs).logits
# for loading, we expect high precision, so check for equality and not just correlation
atol, rtol = 1e-6, 1e-6
assert torch.allclose(output_hqq, output_loaded, atol=atol, rtol=rtol)
# check that outputs are the same after merge_and_unload
model = model.merge_and_unload()
with torch.inference_mode():
output_merged_unloaded = model(**inputs).logits
cc_matrix = torch.corrcoef(
torch.stack((output_normal.float().flatten(), output_merged_unloaded.float().flatten()))
)
assert cc_matrix.min() > min_correlation
@require_non_cpu
@require_gptqmodel
class PeftAwqGPUTests(unittest.TestCase):
r"""
Awq + peft tests
"""
def setUp(self):
from transformers import AwqConfig
from transformers.utils.quantization_config import AwqBackend
self.causal_lm_model_id = "peft-internal-testing/opt-125m-awq"
self.quantization_config = AwqConfig(backend=AwqBackend.AUTO_TRAINABLE)
self.tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
def tearDown(self):
r"""
Efficient mechanism to free accelerator memory after each test. Based on
https://github.com/huggingface/transformers/issues/21094
"""
clear_device_cache(garbage_collection=True)
def _check_inference_finite(self, model, batch):
# try inference without Trainer class
training = model.training
model.eval()
output = model(**batch.to(model.device))
assert torch.isfinite(output.logits).all()
model.train(training)
@pytest.mark.single_gpu_tests
def test_causal_lm_training_awq(self):
r"""
Test the CausalLM training on a single accelerator. The test would simply fail if the adapters are not set
correctly.
"""
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map="auto",
quantization_config=self.quantization_config,
)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
# TODO: deal correctly with this case in transformers
model._is_quantized_training_enabled = True
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
logging_steps=1,
output_dir=tmp_dir,
fp16=True,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
@require_torch_multi_accelerator
def test_causal_lm_training_multi_accelerator(self):
r"""
Test the CausalLM training on a multi-accelerator device. The test would simply fail if the adapters are not
set correctly.
"""
device_map = {
"model.decoder.embed_tokens": 0,
"lm_head": 0,
"model.decoder.embed_positions": 0,
"model.decoder.project_out": 0,
"model.decoder.project_in": 0,
"model.decoder.layers.0": 0,
"model.decoder.layers.1": 0,
"model.decoder.layers.2": 0,
"model.decoder.layers.3": 0,
"model.decoder.layers.4": 0,
"model.decoder.layers.5": 0,
"model.decoder.layers.6": 1,
"model.decoder.layers.7": 1,
"model.decoder.layers.8": 1,
"model.decoder.layers.9": 1,
"model.decoder.layers.10": 1,
"model.decoder.layers.11": 1,
"model.decoder.final_layer_norm": 1,
}
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=device_map,
quantization_config=self.quantization_config,
)
assert set(model.hf_device_map.values()) == {0, 1}
assert {p.device.index for p in model.parameters()} == {0, 1}
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@require_non_xpu
@require_torch_gpu
@require_eetq
class PeftEetqGPUTests(unittest.TestCase):
r"""
EETQ + peft tests
"""
def setUp(self):
self.causal_lm_model_id = "peft-internal-testing/opt-125m"
self.tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
def tearDown(self):
r"""
Efficient mechanism to free GPU memory after each test. Based on
https://github.com/huggingface/transformers/issues/21094
"""
clear_device_cache(garbage_collection=True)
def _check_inference_finite(self, model, batch):
# try inference without Trainer class
training = model.training
model.eval()
output = model(**batch.to(model.device))
assert torch.isfinite(output.logits).all()
model.train(training)
@pytest.mark.single_gpu_tests
def test_causal_lm_training_eetq(self):
r"""
Test the CausalLM training on a single GPU device. The test would simply fail if the adapters are not set
correctly.
"""
from transformers import EetqConfig
with tempfile.TemporaryDirectory() as tmp_dir:
quantization_config = EetqConfig("int8")
try:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id, device_map="auto", quantization_config=quantization_config
)
except FileNotFoundError:
pytest.skip("There is no kernel for EETQ on this architecture, skipping this test.")
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
@require_torch_multi_gpu
def test_causal_lm_training_multi_gpu_eetq(self):
r"""
Test the CausalLM training on a multi-GPU device. The test would simply fail if the adapters are not set
correctly.
"""
from transformers import EetqConfig
with tempfile.TemporaryDirectory() as tmp_dir:
quantization_config = EetqConfig("int8")
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=DEVICE_MAP_MAP[self.causal_lm_model_id],
quantization_config=quantization_config,
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: self.tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(self.tokenizer, mlm=False),
)
model.config.use_cache = False
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@require_non_cpu
@require_torchao
class TestPeftTorchao:
causal_lm_model_id = "peft-internal-testing/opt-125m"
supported_quant_types = [
"int8_weight_only",
"int8_dynamic_activation_int8_weight",
# int4_weight_only raises an error:
# RuntimeError: We encountered some issues during automatic conversion of the weights
# "int4_weight_only",
]
@pytest.fixture(scope="class")
def tokenizer(self):
return AutoTokenizer.from_pretrained(self.causal_lm_model_id)
@pytest.fixture(scope="class", autouse=True)
def setup_teardown(self):
# Efficient mechanism to free GPU memory after each test. Based on
# https://github.com/huggingface/transformers/issues/21094
yield
clear_device_cache(garbage_collection=True)
@staticmethod
def get_quant_type(quant_type: str):
from torchao.quantization import (
Int4WeightOnlyConfig,
Int8DynamicActivationInt8WeightConfig,
Int8WeightOnlyConfig,
)
return {
"int4_weight_only": Int4WeightOnlyConfig(),
"int8_weight_only": Int8WeightOnlyConfig(),
"int8_dynamic_activation_int8_weight": Int8DynamicActivationInt8WeightConfig(),
}[quant_type]
@pytest.mark.parametrize("quant_type", supported_quant_types)
@pytest.mark.single_gpu_tests
def test_causal_lm_training_single_gpu_torchao(self, quant_type, tokenizer):
from transformers import TorchAoConfig
device = 0
with tempfile.TemporaryDirectory() as tmp_dir:
quantization_config = TorchAoConfig(quant_type=self.get_quant_type(quant_type))
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id, device_map=device, quantization_config=quantization_config
)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.model.config.use_cache = False
trainer.train()
model.save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_single_gpu_torchao_dora_int8_weight_only(self, tokenizer):
from transformers import TorchAoConfig
device = 0
with tempfile.TemporaryDirectory() as tmp_dir:
quantization_config = TorchAoConfig(quant_type=self.get_quant_type("int8_weight_only"))
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id, device_map=device, quantization_config=quantization_config
)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
use_dora=True,
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.model.config.use_cache = False
trainer.train()
model.save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
def test_causal_lm_training_single_gpu_torchao_dora_int8_dynamic_activation_int8_weight_raises(self):
from transformers import TorchAoConfig
device = 0
quantization_config = TorchAoConfig(quant_type=self.get_quant_type("int8_dynamic_activation_int8_weight"))
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id, device_map=device, quantization_config=quantization_config
)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
use_dora=True,
)
with pytest.raises(NotImplementedError):
get_peft_model(model, config)
@pytest.mark.single_gpu_tests
@pytest.mark.xfail(
reason="int4_weight_only still has issues",
raises=(RuntimeError, ValueError),
)
def test_causal_lm_training_single_gpu_torchao_int4_raises(self):
# TODO: Once proper torchao support for int4 is added, remove this test and add int4 to supported_quant_types
from transformers import TorchAoConfig
device = 0
quantization_config = TorchAoConfig(quant_type=self.get_quant_type("int4_weight_only"))
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id, device_map=device, quantization_config=quantization_config
)
model = prepare_model_for_kbit_training(model)
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
inputs = torch.arange(10).view(1, -1).to(device)
# this raises:
# > RuntimeError: cutlass cannot initialize
# tested in multiple matchines
model(inputs)
@pytest.mark.parametrize("quant_type", supported_quant_types)
@pytest.mark.multi_gpu_tests
@require_torch_multi_accelerator
def test_causal_lm_training_multi_accelerator_torchao(self, quant_type, tokenizer):
from transformers import TorchAoConfig
device_map = {
"model.decoder.embed_tokens": 0,
"lm_head": 0,
"model.decoder.embed_positions": 0,
"model.decoder.project_out": 0,
"model.decoder.project_in": 0,
"model.decoder.layers.0": 0,
"model.decoder.layers.1": 0,
"model.decoder.layers.2": 0,
"model.decoder.layers.3": 0,
"model.decoder.layers.4": 0,
"model.decoder.layers.5": 0,
"model.decoder.layers.6": 1,
"model.decoder.layers.7": 1,
"model.decoder.layers.8": 1,
"model.decoder.layers.9": 1,
"model.decoder.layers.10": 1,
"model.decoder.layers.11": 1,
"model.decoder.final_layer_norm": 1,
}
with tempfile.TemporaryDirectory() as tmp_dir:
quantization_config = TorchAoConfig(quant_type=self.get_quant_type(quant_type))
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=device_map,
quantization_config=quantization_config,
dtype=torch.bfloat16,
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.model.config.use_cache = False
trainer.train()
model.save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.multi_gpu_tests
@require_torch_multi_accelerator
def test_causal_lm_training_multi_accelerator_torchao_int4_raises(self):
# int4_weight_only raises an error:
# RuntimeError: derivative for aten::_weight_int4pack_mm is not implemented
# TODO: Once proper torchao support for int4 is added, remove this test and add int4 to supported_quant_types
from transformers import TorchAoConfig
device_map = {
"model.decoder.embed_tokens": 0,
"lm_head": 0,
"model.decoder.embed_positions": 0,
"model.decoder.project_out": 0,
"model.decoder.project_in": 0,
"model.decoder.layers.0": 0,
"model.decoder.layers.1": 0,
"model.decoder.layers.2": 0,
"model.decoder.layers.3": 0,
"model.decoder.layers.4": 0,
"model.decoder.layers.5": 0,
"model.decoder.layers.6": 1,
"model.decoder.layers.7": 1,
"model.decoder.layers.8": 1,
"model.decoder.layers.9": 1,
"model.decoder.layers.10": 1,
"model.decoder.layers.11": 1,
"model.decoder.final_layer_norm": 1,
}
quantization_config = TorchAoConfig(self.get_quant_type(quant_type="int4_weight_only"))
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
device_map=device_map,
quantization_config=quantization_config,
dtype=torch.bfloat16,
)
assert set(model.hf_device_map.values()) == set(range(device_count))
assert {p.device.index for p in model.parameters()} == set(range(device_count))
model = prepare_model_for_kbit_training(model)
model.model_parallel = True
model.is_parallelizable = True
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
msg = re.escape("TorchaoLoraLinear only supports int8 weights for now")
with pytest.raises(ValueError, match=msg):
get_peft_model(model, config)
@pytest.mark.single_gpu_tests
def test_torchao_merge_layers_int8_weight_only(self):
from torchao.utils import TorchAOBaseTensor
from transformers import TorchAoConfig
quant_type = "int8_weight_only"
torch.manual_seed(0)
device = 0
dummy_input = torch.arange(10).view(-1, 1).to(device)
quantization_config = TorchAoConfig(self.get_quant_type(quant_type=quant_type))
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id, device_map=device, quantization_config=quantization_config
).eval()
logits_base = model(dummy_input)[0]
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
init_lora_weights=False,
)
model = get_peft_model(model, config)
model.eval()
logits = model(dummy_input)[0]
# sanity check: outputs changed
# precision is quite low, so we need to use high atol and rtol
atol, rtol = 1e-1, 1e-1
assert not torch.allclose(logits, logits_base, atol=atol, rtol=rtol)
model.merge_adapter()
logits_merged = model(dummy_input)[0]
for name, module in model.named_modules():
if "base_layer" in name:
assert isinstance(module.weight, TorchAOBaseTensor)
model.unmerge_adapter()
logits_unmerged = model(dummy_input)[0]
for name, module in model.named_modules():
if "base_layer" in name:
assert isinstance(module.weight, TorchAOBaseTensor)
model = model.merge_and_unload()
logits_merged_unloaded = model(dummy_input)[0]
assert torch.allclose(logits, logits_merged, atol=atol, rtol=rtol)
assert torch.allclose(logits, logits_unmerged, atol=atol, rtol=rtol)
assert torch.allclose(logits, logits_merged_unloaded, atol=atol, rtol=rtol)
@pytest.mark.single_gpu_tests
def test_torchao_merge_layers_int8_dynamic_activation_int8_weight_raises(self):
# int8_dynamic_activation_int8_weight does not support dequantize, thus merging does not work
from transformers import TorchAoConfig
quant_type = "int8_dynamic_activation_int8_weight"
torch.manual_seed(0)
device = 0
quantization_config = TorchAoConfig(quant_type=self.get_quant_type(quant_type))
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id, device_map=device, quantization_config=quantization_config
).eval()
config = LoraConfig(
r=16,
lora_alpha=32,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
init_lora_weights=False,
)
model = get_peft_model(model, config)
msg = re.escape(
"Weights of type LinearActivationQuantizedTensor do not support dequantization (yet), which is needed to "
"support merging."
)
with pytest.raises(NotImplementedError, match=msg):
model.merge_adapter()
@pytest.mark.single_gpu_tests
def test_torchao_lora_warns_when_base_not_quantized_via_transformers(self):
# Manually quantizing the base model with torchao.quantize_ leaves PEFT without
# `get_apply_tensor_subclass`, so the LoRA torchao linear emits a warning at init
# to flag that merge()/unmerge() will not work. When the base model is loaded via
# `TorchAoConfig`, PEFT recovers the requantization subclass and no warning fires.
from torchao.quantization import Int8WeightOnlyConfig, quantize_
from transformers import TorchAoConfig
device = 0
lora_config = LoraConfig(
r=8,
lora_alpha=16,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM",
)
# Path 1: manual quantize_ -> warning expected, and merge() must raise.
model = AutoModelForCausalLM.from_pretrained(self.causal_lm_model_id, device_map=device)
quantize_(model, Int8WeightOnlyConfig())
with pytest.warns(UserWarning, match="get_apply_tensor_subclass"):
peft_model = get_peft_model(model, lora_config)
merge_msg = re.escape(
"was instantiated without `get_apply_tensor_subclass`, which is "
"required to re-quantize the base layer after merging."
)
with pytest.raises(ValueError, match=merge_msg):
peft_model.merge_adapter()
del peft_model, model
clear_device_cache(garbage_collection=True)
# Path 2: TorchAoConfig -> no such warning.
quantization_config = TorchAoConfig(quant_type=Int8WeightOnlyConfig())
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id, device_map=device, quantization_config=quantization_config
)
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
get_peft_model(model, lora_config)
assert not any("get_apply_tensor_subclass" in str(w.message) for w in caught)
PRECISIONS = [(torch.float32), (torch.float16), (torch.bfloat16)]
LORA_PARAMS = {
"r": 8,
"lora_alpha": 16,
"lora_dropout": 0.05,
}
class SimpleModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.embedding_layer = torch.nn.Embedding(1000, 768)
self.layer_norm = torch.nn.LayerNorm(768)
self.linear_transform = torch.nn.Linear(768, 256)
def forward(self, input_ids):
embedded_output = self.embedding_layer(input_ids)
norm_output = self.layer_norm(embedded_output)
linear_output = self.linear_transform(norm_output)
return linear_output
class SimpleConv2DModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.embedding_layer = torch.nn.Embedding(1000, 768)
self.layer_norm = torch.nn.LayerNorm(768)
self.conv2d_transform = torch.nn.Conv2d(1, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
def forward(self, input_ids):
# Additional layers for your custom model
embedded_output = self.embedding_layer(input_ids)
norm_output = self.layer_norm(embedded_output)
# Reshape for Conv2d input (add batch size dimension)
norm_output = norm_output.unsqueeze(1)
conv_output = self.conv2d_transform(norm_output)
# Remove batch size dimension
conv_output = conv_output.squeeze(1)
return conv_output
@require_non_cpu
class TestAutoCast(unittest.TestCase):
device = infer_device()
# This test makes sure, that Lora dtypes are consistent with the types
# inferred by torch.autocast under tested PRECISIONS
@parameterized.expand(PRECISIONS)
def test_simple_model(self, *args, **kwargs):
self._test_model(SimpleModel(), *args, **kwargs)
@parameterized.expand(PRECISIONS)
def test_simple_lora_linear_model(self, *args, **kwargs):
simple_model = SimpleModel()
config = LoraConfig(
**LORA_PARAMS,
target_modules=["linear_transform"],
)
lora_model = get_peft_model(simple_model, config)
self._test_model(lora_model, *args, **kwargs)
@parameterized.expand(PRECISIONS)
def test_simple_lora_embedding_model(self, *args, **kwargs):
simple_model = SimpleModel()
config = LoraConfig(
**LORA_PARAMS,
target_modules=["embedding_layer"],
)
lora_model = get_peft_model(simple_model, config)
self._test_model(lora_model, *args, **kwargs)
@parameterized.expand(PRECISIONS)
def test_simple_conv2d_model(self, *args, **kwargs):
self._test_model(SimpleConv2DModel(), *args, **kwargs)
@parameterized.expand(PRECISIONS)
def test_simple_lora_conv2d_model(self, *args, **kwargs):
simple_model = SimpleConv2DModel()
config = LoraConfig(
**LORA_PARAMS,
target_modules=["conv2d_transform"],
)
lora_model = get_peft_model(simple_model, config)
self._test_model(lora_model, *args, **kwargs)
def _test_model(self, model, precision):
# Move model to GPU
model = model.to(self.device)
# Prepare dummy inputs
input_ids = torch.randint(0, 1000, (2, 10)).to(self.device)
if precision == torch.bfloat16:
if not is_bf16_available():
pytest.skip("Bfloat16 not supported on this device")
# Forward pass with test precision
with torch.autocast(enabled=True, dtype=precision, device_type=self.device):
outputs = model(input_ids)
assert outputs.dtype == precision
class TestFSDPWrap:
"""
Test that we can successfully initialize an FSDP instance of the module.
This is a very simple test, as it does not perform actual FSDP training. Here we just ensure that the FSDP instance
can be created. This can fail for several reasons, e.g. int dtype from BNB or inconsistent requires_grad settings
due to the auto wrap policy.
"""
@pytest.mark.single_gpu_tests
@require_bitsandbytes
def test_bnb_4bit_wrap_fsdp(self):
quant_config = BitsAndBytesConfig(
load_in_4bit=True,
# float32 must be used, or else FSDP will complain about mixed int and float dtypes
bnb_4bit_compute_dtype=torch.float32,
bnb_4bit_quant_storage=torch.float32,
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
"peft-internal-testing/opt-125m",
quantization_config=quant_config,
dtype=torch.float32,
)
# model = prepare_model_for_kbit_training(model)
config = LoraConfig(
target_modules=["q_proj", "v_proj"],
task_type="CAUSAL_LM",
use_dora=True,
)
model = get_peft_model(model, config)
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "29501"
init_process_group(world_size=1, rank=0)
# check that this does not raise:
FSDP(model, auto_wrap_policy=fsdp_auto_wrap_policy(model), use_orig_params=False, sync_module_states=True)
def test_fsdp_auto_wrap_policy_does_not_raise_on_custom_model(self):
# See #2167
# Avoid raising on custom models since Trainer uses fsdp_auto_wrap_policy automatically for PEFT + FSDP
fsdp_auto_wrap_policy(SimpleModel()) # does not raise
class TestBOFT:
"""
Test that we can correctly use half-precision models with BOFT.
"""
device = infer_device()
@require_non_cpu
@pytest.mark.single_gpu_tests
def test_boft_half_linear(self):
# Check that we can use BoFT with model loaded in half precision
config = boft.config.BOFTConfig(boft_n_butterfly_factor=2)
layer = torch.nn.Linear(160, 160).to(self.device)
layer = boft.layer.Linear(layer, "layer", config=config).to(dtype=torch.bfloat16)
x = torch.randn(160, 160, device=self.device, dtype=torch.bfloat16)
layer(x) # does not raise
@require_non_cpu
@pytest.mark.single_gpu_tests
def test_boft_half_conv(self):
conv = torch.nn.Conv2d(1, 1, 4).to(self.device)
config = boft.config.BOFTConfig(boft_n_butterfly_factor=2)
conv = boft.layer.Conv2d(conv, "conv", config=config).to(dtype=torch.bfloat16)
x = torch.randn(1, 160, 160, device=self.device, dtype=torch.bfloat16)
conv(x) # does not raise
class TestPTuningReproducibility:
device = infer_device()
causal_lm_model_id = "peft-internal-testing/opt-125m"
@require_non_cpu
@require_deterministic_for_xpu
def test_p_tuning_exactly_reproducible_after_loading(self, tmp_path):
# See: https://github.com/huggingface/peft/issues/2043#issuecomment-2321522577
# Ensure that after loading a p-tuning checkpoint, results are exactly reproducible (before the patch, they were
# only _almost_ identical).
# The model must be sufficiently large for the effect to be measurable, which is why this test requires is not
# run on CPU.
model_id = "peft-internal-testing/opt-125m"
inputs = torch.arange(10).view(-1, 1).to(self.device)
torch.manual_seed(0)
model = AutoModelForCausalLM.from_pretrained(model_id).to(self.device)
peft_config = PromptEncoderConfig(task_type="CAUSAL_LM", num_virtual_tokens=20, encoder_hidden_size=128)
model = get_peft_model(model, peft_config).eval()
with torch.inference_mode():
output_peft = model(inputs).logits
gen_peft = model.generate(inputs, min_new_tokens=10, max_new_tokens=10)
model.save_pretrained(tmp_path)
del model
clear_device_cache(garbage_collection=True)
model = AutoModelForCausalLM.from_pretrained(model_id).to(self.device)
model = PeftModel.from_pretrained(model, tmp_path)
with torch.inference_mode():
output_loaded = model(inputs).logits
gen_loaded = model.generate(inputs, min_new_tokens=10, max_new_tokens=10)
torch.testing.assert_close(output_loaded, output_peft)
torch.testing.assert_close(gen_loaded, gen_peft)
@require_bitsandbytes
@pytest.mark.single_gpu_tests
def test_p_tuning_causal_lm_training_8bit_bnb(self):
# test is analog to PeftBnbGPUExampleTests.test_causal_lm_training
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
config = PromptEncoderConfig(task_type="CAUSAL_LM", num_virtual_tokens=20, encoder_hidden_size=128)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.single_gpu_tests
class TestLowCpuMemUsageDifferentDevices:
"""Test for the low CPU memory usage option for loading PEFT models.
There are already tests for low_cpu_mem_usage=True in test_initialization.py but here we want to run tests that
require a GPU.
"""
model_id = "peft-internal-testing/tiny-random-OPTForCausalLM"
device = infer_device()
@require_non_cpu
@pytest.mark.parametrize("device_model, device_sd", [("cpu", infer_device()), (infer_device(), "cpu")])
def test_low_cpu_mem_usage_model_model_on_gpu_state_dict_on_cpu_works(self, device_model, device_sd):
# specifically test diverging devices for the model and state_dict
inputs = {"input_ids": torch.randint(0, 100, (1, 10)), "attention_mask": torch.ones(1, 10)}
inputs = {k: v.to(device_model) for k, v in inputs.items()}
model = AutoModelForCausalLM.from_pretrained(self.model_id).to(device_model)
lora_config = LoraConfig(init_lora_weights=False, target_modules="all-linear")
model = get_peft_model(model, lora_config)
model.eval()
logits_not_low_cpu_mem = model(**inputs).logits
state_dict = get_peft_model_state_dict(model)
peft_model_state_dict = {}
# remap the state dict so that it can be correctly loaded, and move weights to the other device
prefix = "base_model.model."
for k, v in state_dict.items():
k = k[len(prefix) :]
peft_model_state_dict[k] = v.to(device_sd)
del model
model = AutoModelForCausalLM.from_pretrained(self.model_id).to(device_model)
model.eval()
inject_adapter_in_model(lora_config, model, low_cpu_mem_usage=True)
load_result = set_peft_model_state_dict(model, peft_model_state_dict, low_cpu_mem_usage=True)
# sanity check: all lora keys are matched
assert not any("lora" in k for k in load_result.missing_keys)
assert not any("lora" in k for k in load_result.unexpected_keys)
logits_low_cpu_mem = model(**inputs).logits
assert torch.allclose(logits_low_cpu_mem, logits_not_low_cpu_mem)
assert {p.device.type for p in model.parameters()} == {device_model}
@require_bitsandbytes
@pytest.mark.parametrize("quantization_method", ["bnb-4bit", "bnb-8bit"])
def test_low_cpu_mem_usage_with_quantization(self, quantization_method):
# Ensure that low_cpu_mem_usage works with quantization
# See also https://github.com/huggingface/diffusers/issues/10550
if quantization_method == "bnb-4bit":
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float32,
bnb_4bit_quant_storage=torch.float32,
bnb_4bit_use_double_quant=True,
)
elif quantization_method == "bnb-8bit":
quantization_config = BitsAndBytesConfig(load_in_8bit=True)
else:
raise ValueError(f"Unknown quantization method {quantization_method}")
model = AutoModelForCausalLM.from_pretrained(self.model_id, quantization_config=quantization_config)
if model.device.type != self.device:
# calling model.to("cuda") with 8 bit bnb raises an error, thus guard against it
model = model.to(self.device)
lora_config = LoraConfig(init_lora_weights=False, target_modules="all-linear")
# We use get_peft_model with low_cpu_mem_usage=True here. This is not typically done in practice (the option is
# mostly interesting for loading trained adapters), but it does the job for testing purposes.
model = get_peft_model(model, lora_config, low_cpu_mem_usage=True) # this should not raise
assert {p.device.type for p in model.parameters()} == {self.device, "meta"}
@pytest.mark.single_gpu_tests
@require_non_cpu
class TestEvaInitializationGPU:
"""GPU tests for the Eva initialization method.
This test suite verifies:
1. Consistency of initialization across different seeds
2. Proper error handling for invalid inputs
3. Compatibility with different model architectures
4. Reproducibility of results
5. Proper handling of edge cases
"""
# Constants for test configuration
COSINE_SIMILARITY_THRESHOLD = 0.75
NUM_SEEDS = 3
BATCH_SIZE = 4
MAX_LENGTH = 256
LORA_DIM = 8
LORA_ALPHA = 1
DEVICE = infer_device()
@pytest.fixture
def tokenizer(self):
tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
tokenizer.pad_token = tokenizer.eos_token
return tokenizer
@pytest.fixture
def dataset(self, tokenizer):
# concatenate examples
examples = []
example = ""
ds = load_dataset_english_quotes()["train"] # cached
for data in ds:
if len(example) >= self.MAX_LENGTH:
examples.append(example)
example = ""
example = example + " " + data["quote"]
dataset = Dataset.from_dict({"text": examples})
# tokenize
dataset = dataset.map(
lambda x: tokenizer(x["text"], padding="max_length", truncation=True, max_length=self.MAX_LENGTH),
batched=True,
remove_columns=dataset.column_names,
)
dataset.set_format(type="torch")
return dataset
@pytest.fixture
def model(self):
model_id = "openai-community/gpt2"
with hub_online_once(model_id):
model = AutoModelForCausalLM.from_pretrained(model_id)
model.transformer.h = model.transformer.h[:2] # truncate to 2 layers
return model.to(self.DEVICE)
@pytest.fixture
def model_bnb(self):
model_id = "openai-community/gpt2"
with hub_online_once(model_id):
bnb_config = BitsAndBytesConfig(load_in_4bit=True)
model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config=bnb_config,
attn_implementation="eager", # gpt2 doesn't support flash attention
)
model.transformer.h = model.transformer.h[:2] # truncate to 2 layers
model = prepare_model_for_kbit_training(model)
return model
@pytest.fixture
def model_fixture(self, request):
return request.getfixturevalue(request.param)
@pytest.fixture
def peft_config(self):
return LoraConfig(
r=self.LORA_DIM,
lora_alpha=self.LORA_ALPHA,
target_modules=["c_attn"],
init_lora_weights="eva",
eva_config=EvaConfig(rho=2),
)
def is_bnb_model(self, model):
return hasattr(model.config, "quantization_config")
@staticmethod
def collate_fn(examples):
return {k: torch.stack([v[k] for v in examples], dim=0) for k in examples[0].keys()}
def get_dataloader(self, dataset):
return DataLoader(
dataset,
batch_size=self.BATCH_SIZE,
collate_fn=self.collate_fn,
shuffle=False,
)
@require_bitsandbytes
@pytest.mark.parametrize(
"eva_config",
[
EvaConfig(rho=2, tau=0.99),
EvaConfig(rho=1, tau=0.9),
EvaConfig(rho=1, whiten=True, tau=0.9),
EvaConfig(rho=1.0001, tau=0.9),
],
)
@pytest.mark.parametrize("model_fixture", ["model", "model_bnb"], indirect=True)
def test_eva_initialization_consistency(self, model_fixture, dataset, peft_config, eva_config):
"""Test that the state dict returned by get_eva_state_dict loaded correctly and is consistent across different
seeds based on the cosine similarity of the svd components.
"""
peft_config = deepcopy(peft_config)
peft_config.eva_config = eva_config
state_dicts = []
for seed in range(self.NUM_SEEDS):
shuffled_dataset = dataset.shuffle(seed=seed)
dataloader = self.get_dataloader(shuffled_dataset)
sd = get_eva_state_dict(deepcopy(model_fixture), dataloader, peft_config, show_progress_bar=False)
state_dicts.append(sd)
cos_sims = defaultdict(list)
for i, j in itertools.combinations(range(self.NUM_SEEDS), 2):
for k, v1 in state_dicts[i].items():
v2 = state_dicts[j][k]
min_size = min(v1.size(0), v2.size(0))
cos_sims[k].extend(
torch.cosine_similarity(v1[:min_size].abs(), v2[:min_size].abs(), dim=1).abs().tolist()
)
mean_cosine_similarities = {k: torch.tensor(v).mean() for k, v in cos_sims.items()}
for mean_cosine_similarity in mean_cosine_similarities.values():
assert mean_cosine_similarity > self.COSINE_SIMILARITY_THRESHOLD, (
f"Mean absolute cosine similarity {mean_cosine_similarity:.4f} "
f"is not greater than {self.COSINE_SIMILARITY_THRESHOLD}"
)
@pytest.mark.parametrize(
"prepare_layer_inputs_keys, expected_outcome",
[
(None, "success"),
(["transformer.h.0.attn.c_attn"], "success"),
(
["transformer.h.0.attn.c_attn", "transformer.h.1.attn.c_attn", "transformer.h.2.attn.c_attn"],
"value_error",
),
],
)
def test_eva_state_dict_prepare_inputs_mapping(
self, model, dataset, peft_config, prepare_layer_inputs_keys, expected_outcome
):
"""
Tests for cases where prepare_layer_inputs_fn is a mapping. Checks that if not all target modules are present,
the prepare_layer_inputs_fn for the remaining modules is set to None. Also checks that if more keys than target
modules are present, a ValueError is raised.
"""
def fn(x, *args):
return x[0].view(-1, x[0].size(-1))
if prepare_layer_inputs_keys is None:
prepare_layer_inputs_fn = fn
else:
prepare_layer_inputs_fn = {k: fn for k in prepare_layer_inputs_keys}
shuffled_dataset = dataset.shuffle(seed=0)
dataloader = self.get_dataloader(shuffled_dataset)
modified_peft_config = deepcopy(peft_config)
modified_peft_config.eva_config.tau = 0 # converge immediately
if expected_outcome == "success":
sd = get_eva_state_dict(
model,
dataloader,
modified_peft_config,
prepare_model_inputs_fn=None,
prepare_layer_inputs_fn=prepare_layer_inputs_fn,
)
assert len(sd) == 2
assert "transformer.h.0.attn.c_attn" in sd
assert "transformer.h.1.attn.c_attn" in sd
else:
with pytest.raises(
ValueError, match="prepare_layer_inputs_fn is a mapping but the following module names were not found"
):
get_eva_state_dict(
model,
dataloader,
modified_peft_config,
prepare_model_inputs_fn=None,
prepare_layer_inputs_fn=prepare_layer_inputs_fn,
)
@pytest.mark.parametrize(
"eva_config",
[EvaConfig(rho=2, adjust_scaling_factors=True)],
)
def test_eva_state_dict_adjust_scaling_factors(self, model, dataset, peft_config, eva_config):
"""
Tests that the scaling factors are adjusted so that all LoRA gradients have the same scale regardless of their
rank.
"""
modified_peft_config = deepcopy(peft_config)
modified_peft_config.eva_config = eva_config
dataloader = self.get_dataloader(dataset)
peft_model = get_peft_model(deepcopy(model), modified_peft_config)
scaling_factors_before = {}
for n, m in peft_model.named_modules():
if isinstance(m, LoraLayer):
scaling_factors_before[n] = m.scaling["default"]
initialize_lora_eva_weights(peft_model, dataloader)
for n, m in peft_model.named_modules():
if isinstance(m, LoraLayer):
assert m.scaling["default"] == scaling_factors_before[n]
@pytest.mark.parametrize("has_rank_zero", [True, False])
def test_load_eva_state_dict(self, model, dataset, peft_config, tmp_path, has_rank_zero):
"""
Tests that the `eva_state_dict` argument in `initialize_lora_eva_weights` can be used to initialize a model
with EVA weights and that the initialized model can be saved and loaded correctly.
"""
dataloader = self.get_dataloader(dataset)
peft_model = get_peft_model(deepcopy(model), peft_config)
sd = get_eva_state_dict(peft_model, dataloader)
if has_rank_zero:
k = "base_model.model.transformer.h.0.attn.c_attn"
sd[k] = sd[k][:0]
initialize_lora_eva_weights(peft_model, eva_state_dict=sd)
if has_rank_zero:
assert not isinstance(peft_model.model.transformer.h[0].attn.c_attn, LoraLayer)
else:
assert isinstance(peft_model.model.transformer.h[0].attn.c_attn, LoraLayer)
peft_model.save_pretrained(tmp_path)
peft_model = PeftModel.from_pretrained(model, tmp_path, torch_device=self.DEVICE, low_cpu_mem_usage=True)
peft_model(**{k: v.to(self.DEVICE) for k, v in next(iter(dataloader)).items()})
def test_missing_eva_inits(self, model, dataset, peft_config):
"""
Tests that a warning is raised when some adapter modules were not initialized with EVA weights.
"""
modified_peft_config = deepcopy(peft_config)
modified_peft_config.target_modules = ["wte"]
dataloader = self.get_dataloader(dataset)
peft_model = get_peft_model(deepcopy(model), modified_peft_config)
msg = (
"the following layers were initialized with init_lora_weights=True because they were not found in the eva "
"state_dict:*"
)
with pytest.warns(UserWarning, match=msg):
initialize_lora_eva_weights(peft_model, dataloader)
def test_load_eva_model(self, model, dataset, peft_config, tmp_path):
"""
Tests that a model initialized with EVA weights can be loaded correctly.
"""
dataloader = self.get_dataloader(dataset)
peft_model = get_peft_model(deepcopy(model), peft_config)
initialize_lora_eva_weights(peft_model, dataloader)
peft_model.save_pretrained(tmp_path)
peft_model = PeftModel.from_pretrained(model, tmp_path, torch_device=self.DEVICE, low_cpu_mem_usage=True)
peft_model(**{k: v.to(self.DEVICE) for k, v in next(iter(dataloader)).items()})
@pytest.mark.parametrize("use_label_mask", [True, False])
def test_eva_label_mask(self, model, dataset, peft_config, use_label_mask):
"""
Tests that label masking works correctly in get_eva_state_dict (see PR #3234).
"""
def add_labels(x):
return {
"labels": torch.where(
x["attention_mask"].bool(),
torch.ones_like(x["attention_mask"]),
-100 * torch.ones_like(x["attention_mask"]),
)
}
dataset_with_labels = dataset.map(add_labels)
dataset_with_labels.set_format(type="torch")
dataloader = self.get_dataloader(dataset_with_labels)
modified_peft_config = deepcopy(peft_config)
modified_peft_config.eva_config = EvaConfig(rho=2, use_label_mask=use_label_mask)
sd = get_eva_state_dict(deepcopy(model), dataloader, modified_peft_config, show_progress_bar=False)
assert len(sd) > 0
def test_eva_initialization_with_invalid_dataloader(self, model, peft_config):
"""Test that appropriate error is raised when dataloader is empty."""
empty_dataset = Dataset.from_dict({"text": []})
dataloader = self.get_dataloader(empty_dataset)
with pytest.raises(ValueError, match="dataloader is empty"):
get_eva_state_dict(model, dataloader, peft_config)
class TestALoRAInferenceGPU:
"""GPU inference for Activated LoRA."""
# Constants for test configuration
NUM_SEEDS = 3
LORA_DIM = 8
LORA_ALPHA = 1
DEVICE = infer_device()
@pytest.fixture
def tokenizer(self):
tokenizer = AutoTokenizer.from_pretrained("peft-internal-testing/opt-125m")
tokenizer.pad_token = tokenizer.eos_token
return tokenizer
@pytest.fixture
def model(self):
model = AutoModelForCausalLM.from_pretrained("peft-internal-testing/opt-125m")
model.model.decoder.layers = model.model.decoder.layers[:2] # truncate to 2 layers
return model.to(self.DEVICE)
@pytest.fixture
def model_bnb(self):
bnb_config = BitsAndBytesConfig(load_in_4bit=True)
model = AutoModelForCausalLM.from_pretrained(
"peft-internal-testing/opt-125m",
quantization_config=bnb_config,
)
model.model.decoder.layers = model.model.decoder.layers[:2] # truncate to 2 layers
model = prepare_model_for_kbit_training(model)
return model
@pytest.fixture
def peft_config(self):
return LoraConfig(
r=self.LORA_DIM,
task_type="CAUSAL_LM",
lora_alpha=self.LORA_ALPHA,
target_modules=["q_proj"],
alora_invocation_tokens=[2], # id for </s>
init_lora_weights=False,
)
@require_non_cpu
@require_bitsandbytes
@pytest.mark.single_gpu_tests
def test_alora_forward_consistency(self, model, model_bnb, peft_config):
"""Test that the forwards of the model with adapter are similar across quantizations."""
for seed in range(self.NUM_SEEDS):
torch.manual_seed(seed)
# random.seed(seed)
np.random.seed(seed)
peft_model = get_peft_model(deepcopy(model), peft_config)
torch.manual_seed(seed)
# random.seed(seed)
np.random.seed(seed)
peft_model_bnb = get_peft_model(deepcopy(model_bnb), peft_config)
peft_model.eval()
peft_model_bnb.eval()
input_ids = torch.tensor([[0, 1, 2, 3]]).to(self.DEVICE)
with torch.no_grad():
peft_out = peft_model(input_ids=input_ids, return_dict=True, output_hidden_states=True)
peft_out_bnb = peft_model_bnb(input_ids=input_ids, return_dict=True, output_hidden_states=True)
h_fp = peft_out.hidden_states[-1]
h_4b = peft_out_bnb.hidden_states[-1]
a = h_fp.detach().to(torch.float32).cpu()
b = h_4b.detach().to(torch.float32).cpu()
import torch.nn.functional as F
cos = F.cosine_similarity(a.flatten(), b.flatten(), dim=0).item()
assert cos > 0.9
@pytest.mark.multi_gpu_tests
class TestPrefixTuning:
device = infer_device()
causal_lm_model_id = "peft-internal-testing/opt-125m"
@require_torch_multi_accelerator
def test_prefix_tuning_multiple_devices_decoder_model(self):
# See issue 2134
model_id = "hf-internal-testing/tiny-random-MistralForCausalLM"
tokenizer = AutoTokenizer.from_pretrained(model_id, padding="left")
inputs = tokenizer(["A list of colors: red, blue"], return_tensors="pt").to(self.device)
device_map = {
"model.embed_tokens": 0,
"model.layers.0": 0,
"model.layers.1": 1,
"model.norm": 1,
"model.rotary_emb": 1,
"lm_head": 1,
}
model = AutoModelForCausalLM.from_pretrained(model_id, device_map=device_map)
# sanity check, as the test passes trivially for a single device
assert len({p.device for p in model.parameters()}) > 1
# sanity check: this should work without peft
model.generate(**inputs) # does not raise
peft_config = PrefixTuningConfig(num_virtual_tokens=10, task_type="CAUSAL_LM")
model = get_peft_model(model, peft_config)
model.generate(**inputs) # does not raise
@require_torch_multi_accelerator
def test_prefix_tuning_multiple_devices_encoder_decoder_model(self):
# See issue 2134
model_id = "peft-internal-testing/tiny-random-t5"
tokenizer = AutoTokenizer.from_pretrained(model_id, padding="left")
inputs = tokenizer(["A list of colors: red, blue"], return_tensors="pt").to(self.device)
device_map = {
"shared": 0,
"encoder.embed_tokens": 0,
"encoder.block.0": 0,
"encoder.block.1": 0,
"encoder.block.2": 1,
"encoder.block.3": 1,
"encoder.block.4": 1,
"encoder.final_layer_norm": 1,
"decoder.embed_tokens": 0,
"decoder.block.0": 0,
"decoder.block.1": 0,
"decoder.block.2": 1,
"decoder.block.3": 1,
"decoder.block.4": 1,
"decoder.final_layer_norm": 1,
"lm_head": 0,
}
model = AutoModelForSeq2SeqLM.from_pretrained(model_id, device_map=device_map)
# sanity check, as the test passes trivially for a single device
assert len({p.device for p in model.parameters()}) > 1
# sanity check: this should work without peft
model.generate(**inputs) # does not raise
peft_config = PrefixTuningConfig(num_virtual_tokens=10, task_type="SEQ_2_SEQ_LM")
model = get_peft_model(model, peft_config)
model.generate(**inputs) # does not raise
@require_bitsandbytes
@pytest.mark.single_gpu_tests
def test_prefix_tuning_causal_lm_training_8bit_bnb(self):
# test is analog to PeftBnbGPUExampleTests.test_causal_lm_training
with tempfile.TemporaryDirectory() as tmp_dir:
model = AutoModelForCausalLM.from_pretrained(
self.causal_lm_model_id,
quantization_config=BitsAndBytesConfig(load_in_8bit=True),
device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(self.causal_lm_model_id)
config = PrefixTuningConfig(num_virtual_tokens=10, task_type=TaskType.CAUSAL_LM)
model = get_peft_model(model, config)
data = load_dataset_english_quotes()
data = data.map(lambda samples: tokenizer(samples["quote"]), batched=True)
trainer = Trainer(
model=model,
train_dataset=data["train"],
args=TrainingArguments(
per_device_train_batch_size=4,
warmup_steps=2,
max_steps=3,
learning_rate=2e-4,
fp16=True,
logging_steps=1,
output_dir=tmp_dir,
),
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.train()
model.cpu().save_pretrained(tmp_dir)
assert "adapter_config.json" in os.listdir(tmp_dir)
assert SAFETENSORS_WEIGHTS_NAME in os.listdir(tmp_dir)
# assert loss is not None
assert trainer.state.log_history[-1]["train_loss"] is not None
@pytest.mark.skipif(not (torch.cuda.is_available() or is_xpu_available()), reason="test requires a GPU or XPU")
@pytest.mark.single_gpu_tests
class TestHotSwapping:
"""
Test hotswapping on compiled models.
This test suite is only run on GPU as it is quite slow.
"""
torch_device = infer_device()
@pytest.fixture(scope="class", autouse=True)
def reset_float32_matmul_precision(self):
# Earlier tests may run torchao, which, at the time this was added, sets the float32 matmul precision to 'high'.
# This in turn results in some models producing different outputs when compiled (but only for some seeds).
# Therefore, we need to ensure that the precision is reset to "highest", which is the default.
# TODO: if torchao removes the side effect, this fixture can be deleted.
# https://github.com/pytorch/ao/blob/ffb4350640e76c7e7f449dd1e36d33f19fe384c8/torchao/quantization/utils.py#L589
torch.set_float32_matmul_precision("highest")
@pytest.fixture(autouse=True)
def reset_dynamo_cache(self):
# It is critical that the dynamo cache is reset for each test. Otherwise, if the test re-uses the same model,
# there will be recompilation errors, as torch caches the model when run in the same process.
torch._dynamo.reset()
yield
#######
# LLM #
#######
def check_hotswap(self, do_hotswap, ranks, alpha_scalings):
"""
Test hotswapping with a compiled model.
Passing do_hotswap=False should trigger recompilation. Use the raise_error_on_recompile context manager to
raise an error when recompilation occurs.
"""
torch.manual_seed(0)
inputs = torch.arange(10).view(-1, 1).to(self.torch_device)
model_id = "peft-internal-testing/tiny-random-OPTForCausalLM"
model = AutoModelForCausalLM.from_pretrained(model_id).to(self.torch_device)
rank0, rank1 = ranks
alpha0, alpha1 = alpha_scalings
# note that the 2nd adapter targeting a subset of the 1st adapter is okay, but not the other way round
config0 = LoraConfig(init_lora_weights=False, r=rank0, lora_alpha=alpha0, target_modules=["q_proj", "v_proj"])
config1 = LoraConfig(init_lora_weights=False, r=rank1, lora_alpha=alpha1, target_modules=["q_proj"])
model = get_peft_model(model, config0, adapter_name="adapter0").eval()
with torch.inference_mode():
output0 = model(inputs).logits
model.add_adapter("adapter1", config1)
model.set_adapter("adapter1")
with torch.inference_mode():
output1 = model(inputs).logits
# sanity check:
tol = 1e-4
assert not torch.allclose(output0, output1, atol=tol, rtol=tol)
with tempfile.TemporaryDirectory() as tmp_dirname:
model.save_pretrained(tmp_dirname)
del model
model = AutoModelForCausalLM.from_pretrained(model_id).to(self.torch_device)
model = PeftModel.from_pretrained(model, os.path.join(tmp_dirname, "adapter0")).eval()
if do_hotswap:
prepare_model_for_compiled_hotswap(model, config=model.peft_config, target_rank=max(ranks))
model = torch.compile(model, mode="reduce-overhead")
output_after0 = model(inputs).logits
assert torch.allclose(output0, output_after0, atol=tol, rtol=tol)
# swap and check that we get the output from adapter1
if do_hotswap:
hotswap_adapter(model, os.path.join(tmp_dirname, "adapter1"), adapter_name="default")
else:
model.load_adapter(os.path.join(tmp_dirname, "adapter1"), adapter_name="other")
model.set_adapter("other")
# we need to call forward to potentially trigger recompilation
output_after1 = model(inputs).logits
assert torch.allclose(output1, output_after1, atol=tol, rtol=tol)
# we need to call forward third time since cudagraphs are not recorded in first call.
if do_hotswap:
hotswap_adapter(model, os.path.join(tmp_dirname, "adapter0"), adapter_name="default")
output_after2 = model(inputs).logits
assert torch.allclose(output0, output_after2, atol=tol, rtol=tol)
# it is important to check hotswapping small to large ranks and large to small ranks
@pytest.mark.parametrize("ranks", [(11, 11), (7, 13), (13, 7)])
def test_hotswapping_compiled_model_does_not_trigger_recompilation(self, ranks):
# here we set three configs to ensure no recompilation or cudagraph re-record occurs:
# 1. error_on_recompile: raise an error on recompilation
# 2. inline_inbuilt_nn_modules: needed to raise an error on static input address changes instead of re-recording
# 3. triton.cudagraph_support_input_mutation: same as above
dynamo_config_ctx = torch._dynamo.config.patch(error_on_recompile=True, inline_inbuilt_nn_modules=False)
inductor_config_ctx = torch._inductor.config.patch("triton.cudagraph_support_input_mutation", False)
with dynamo_config_ctx, inductor_config_ctx:
self.check_hotswap(do_hotswap=True, ranks=ranks, alpha_scalings=ranks)
def test_no_hotswapping_compiled_model_triggers_recompilation(self):
# contingency test to ensure that hotswapping is actually needed to prevent recompilation
ranks = 7, 13
with torch._dynamo.config.patch(error_on_recompile=True):
with pytest.raises(torch._dynamo.exc.RecompileError): # raise an error on recompilation
self.check_hotswap(do_hotswap=False, ranks=ranks, alpha_scalings=ranks)
@pytest.mark.parametrize("do_compile", [False, True])
def test_hotswap_lora_target_parameters(self, do_compile, tmp_path):
# Test that hotswapping works with target_parameters. In this test, there is no need to call
# prepare_model_for_compiled_hotswap, as we use the same LoRA shapes. Due to (re-)compilation, the test is
# relatively slow.
atol, rtol = 1e-6, 1e-6
model_id = "trl-internal-testing/tiny-GptOssForCausalLM"
inputs = torch.arange(10).view(1, -1).to(self.torch_device)
def strong_init(model):
# increase the scale of the LoRA weights so that the difference between adapters is more pronounced, making the test more robust
for name, param in model.named_parameters():
if "lora_" in name:
param.data *= 10.0
with hub_online_once(model_id):
model = AutoModelForCausalLM.from_pretrained(model_id).to(self.torch_device)
with torch.inference_mode():
output_base = model(inputs).logits
# create adapter 0
config = LoraConfig(
target_parameters=[
"mlp.experts.down_proj",
"mlp.experts.gate_up_proj",
],
init_lora_weights=False,
)
torch.manual_seed(0)
model = get_peft_model(model, config)
strong_init(model)
# Note: For compilation, use eager backend, as the parameter targeting, which uses nn.utils.parametrize,
# leads to recompiles/graph breaks, which can significantly affect the outputs, even if weights are correctly
# loaded.
model = torch.compile(model, backend="eager")
with torch.inference_mode():
torch.manual_seed(0)
output0 = model(inputs).logits
# sanity check: outputs differ
assert not torch.allclose(output_base, output0, atol=1e-3, rtol=1e-3)
model.save_pretrained(tmp_path / "adapter0")
del model
# create adapter 1
model = AutoModelForCausalLM.from_pretrained(model_id).to(self.torch_device)
torch.manual_seed(1)
model = get_peft_model(model, config)
strong_init(model)
model = torch.compile(model, backend="eager")
model.eval()
with torch.inference_mode():
torch.manual_seed(0)
output1 = model(inputs).logits
model.save_pretrained(tmp_path / "adapter1")
# sanity check: they're not the same
assert not torch.allclose(output0, output1, atol=1e-3, rtol=1e-3)
del model
# load adapter 0
model = AutoModelForCausalLM.from_pretrained(model_id).to(self.torch_device)
model = PeftModel.from_pretrained(model, tmp_path / "adapter0")
model = torch.compile(model, backend="eager")
with torch.inference_mode():
torch.manual_seed(0)
output_loaded0 = model(inputs).logits
# sanity check: same output after loading for adapter 0
assert torch.allclose(output0, output_loaded0, atol=atol, rtol=rtol)
# hotswap with adapter 1
hotswap_adapter(model, tmp_path / "adapter1", adapter_name="default")
with torch.inference_mode():
torch.manual_seed(0)
output_loaded1 = model(inputs).logits
# real check: model now behaves like adapter 1
assert torch.allclose(output1, output_loaded1, atol=atol, rtol=rtol)
# hotswap back to adapter 0
hotswap_adapter(model, tmp_path / "adapter0", adapter_name="default")
with torch.inference_mode():
torch.manual_seed(0)
output_loaded_back0 = model(inputs).logits
# real check: model now behaves again like adapter 0
assert torch.allclose(output0, output_loaded_back0, atol=atol, rtol=rtol)
###################
# DIFFUSION MODEL #
###################
def get_small_unet(self):
# from diffusers UNet2DConditionModelTests
from diffusers import UNet2DConditionModel
torch.manual_seed(0)
init_dict = {
"block_out_channels": (4, 8),
"norm_num_groups": 4,
"down_block_types": ("CrossAttnDownBlock2D", "DownBlock2D"),
"up_block_types": ("UpBlock2D", "CrossAttnUpBlock2D"),
"cross_attention_dim": 8,
"attention_head_dim": 2,
"out_channels": 4,
"in_channels": 4,
"layers_per_block": 1,
"sample_size": 16,
}
model = UNet2DConditionModel(**init_dict)
return model.to(self.torch_device)
def get_unet_lora_config(self, lora_rank, lora_alpha, target_modules):
# from diffusers test_models_unet_2d_condition.py
# note that this only targets linear layers by default
unet_lora_config = LoraConfig(
r=lora_rank,
lora_alpha=lora_alpha,
target_modules=target_modules,
init_lora_weights=False,
use_dora=False,
)
return unet_lora_config
def get_dummy_input(self):
pipeline_inputs = {
"prompt": "A painting of a squirrel eating a burger",
"num_inference_steps": 5,
"guidance_scale": 6.0,
"output_type": "np",
"return_dict": False,
}
return pipeline_inputs
def set_lora_device(self, model, adapter_names, device):
# copied from diffusers LoraBaseMixin.set_lora_device
for module in model.modules():
if isinstance(module, BaseTunerLayer):
for adapter_name in adapter_names:
module.lora_A[adapter_name].to(device)
module.lora_B[adapter_name].to(device)
# this is a param, not a module, so device placement is not in-place -> re-assign
if hasattr(module, "lora_magnitude_vector") and module.lora_magnitude_vector is not None:
if adapter_name in module.lora_magnitude_vector:
module.lora_magnitude_vector[adapter_name] = module.lora_magnitude_vector[adapter_name].to(
device
)
def check_hotswap_diffusion(self, ranks, alpha_scalings, target_modules):
"""
Check that hotswapping works on a pipeline.
This is essentially the same test as:
https://github.com/huggingface/diffusers/blob/d7dd924ece56cddf261cd8b9dd901cbfa594c62c/tests/pipelines/test_pipelines.py#L2264
Steps:
- create 2 LoRA adapters and save them
- load the first adapter
- hotswap the second adapter
- check that the outputs are correct
- optionally compile the model
Note: We set rank == alpha here because save_lora_adapter does not save the alpha scalings, thus the test would
fail if the values are different. Since rank != alpha does not matter for the purpose of this test, this is
fine.
"""
from diffusers import StableDiffusionPipeline
# create 2 adapters with different ranks and alphas
dummy_input = self.get_dummy_input()
pipeline = StableDiffusionPipeline.from_pretrained("hf-internal-testing/tiny-sd-pipe").to(torch_device)
rank0, rank1 = ranks
alpha0, alpha1 = alpha_scalings
max_rank = max([rank0, rank1])
lora_config0 = self.get_unet_lora_config(rank0, alpha0, target_modules)
lora_config1 = self.get_unet_lora_config(rank1, alpha1, target_modules)
torch.manual_seed(0)
pipeline.unet.add_adapter(lora_config0, adapter_name="adapter0")
output0_before = pipeline(**dummy_input, generator=torch.manual_seed(0))[0]
torch.manual_seed(1)
pipeline.unet.add_adapter(lora_config1, adapter_name="adapter1")
pipeline.unet.set_adapter("adapter1")
output1_before = pipeline(**dummy_input, generator=torch.manual_seed(0))[0]
# sanity check
tol = 1e-3
assert not np.allclose(output0_before, output1_before, atol=tol, rtol=tol)
assert not (output0_before == 0).all()
assert not (output1_before == 0).all()
with tempfile.TemporaryDirectory() as tmp_dirname:
# save the adapter checkpoints
sd0 = get_peft_model_state_dict(pipeline.unet, adapter_name="adapter0")
StableDiffusionPipeline.save_lora_weights(
save_directory=os.path.join(tmp_dirname, "adapter0"), safe_serialization=True, unet_lora_layers=sd0
)
sd1 = get_peft_model_state_dict(pipeline.unet, adapter_name="adapter1")
StableDiffusionPipeline.save_lora_weights(
save_directory=os.path.join(tmp_dirname, "adapter1"), safe_serialization=True, unet_lora_layers=sd1
)
del pipeline
# load the first adapter
pipeline = StableDiffusionPipeline.from_pretrained("hf-internal-testing/tiny-sd-pipe").to(torch_device)
# no need to prepare if the model is not compiled or if the ranks are identical
pipeline.enable_lora_hotswap(target_rank=max_rank)
file_name0 = os.path.join(tmp_dirname, "adapter0", "pytorch_lora_weights.safetensors")
file_name1 = os.path.join(tmp_dirname, "adapter1", "pytorch_lora_weights.safetensors")
pipeline.load_lora_weights(file_name0)
pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead")
output0_after = pipeline(**dummy_input, generator=torch.manual_seed(0))[0]
# sanity check: still same result
assert np.allclose(output0_before, output0_after, atol=tol, rtol=tol)
# hotswap the 2nd adapter
pipeline.load_lora_weights(file_name1, hotswap=True, adapter_name="default_0")
output1_after = pipeline(**dummy_input, generator=torch.manual_seed(0))[0]
# sanity check: since it's the same LoRA, the results should be identical
assert np.allclose(output1_before, output1_after, atol=tol, rtol=tol)
# we need to call forward third time since cudagraphs are not recorded in first call.
pipeline.load_lora_weights(file_name0, hotswap=True, adapter_name="default_0")
output2_after = pipeline(**dummy_input, generator=torch.manual_seed(0))[0]
assert np.allclose(output0_before, output2_after, atol=tol, rtol=tol)
@pytest.mark.skipif(not is_diffusers_available(), reason="Test requires diffusers to be installed")
# it is important to check hotswapping small to large ranks and large to small ranks
@pytest.mark.parametrize("ranks", [(11, 11), (7, 13), (13, 7)])
@pytest.mark.parametrize(
"target_modules",
[
["to_q", "to_k", "to_v", "to_out.0"], # Linear layers
["conv", "conv1", "conv2"], # Conv2d layers
["to_q", "conv"], # mix of Linear and Conv2d
],
)
def test_hotswapping_compiled_diffusers_model_does_not_trigger_recompilation(self, ranks, target_modules):
# here we set three configs to ensure no recompilation or cudagraph re-record occurs:
# 1. error_on_recompile: raise an error on recompilation
# 2. inline_inbuilt_nn_modules: needed to raise an error on static input address changes instead of re-recording
# 3. triton.cudagraph_support_input_mutation: same as above
dynamo_config_ctx = torch._dynamo.config.patch(error_on_recompile=True, inline_inbuilt_nn_modules=False)
inductor_config_ctx = torch._inductor.config.patch("triton.cudagraph_support_input_mutation", False)
with dynamo_config_ctx, inductor_config_ctx:
self.check_hotswap_diffusion(ranks=ranks, alpha_scalings=ranks, target_modules=target_modules)
# Test: 4-bit load + Arrow + generate
class TestArrowQuantized:
@pytest.fixture(scope="class")
def workdir(self, tmp_path_factory):
"""Create and return a temp directory path for this class (no chdir)."""
wd = tmp_path_factory.mktemp("arrow_workdir")
return Path(wd)
def _create_and_save_adapter_opt(self, out_dir: Path, rank: int = 4):
"""
Build a randomly initialized LoRA adapter for OPT-125M and save into `out_dir`. We construct a model from
CONFIG (no pretrained weights) to avoid slow downloads here.
"""
model_id = "peft-internal-testing/opt-125m"
# Target all linear layers so the adapter matches whatever we later quantize/load.
lora_cfg = LoraConfig(
r=rank,
target_modules="all-linear",
task_type="CAUSAL_LM",
init_lora_weights=False,
)
# Load the adapter on the model and save it
with hub_online_once(model_id):
model = AutoModelForCausalLM.from_pretrained(model_id)
peft_model = get_peft_model(model, lora_cfg)
peft_model.save_pretrained(out_dir)
@pytest.fixture(scope="class")
def ts_adapters_opt(self, workdir: Path):
"""
Build 3 locally-saved task-specific adapters for OPT-125M and return their absolute paths.
"""
paths = []
for i in range(3):
sub = workdir / f"ts_expert_{i}"
self._create_and_save_adapter_opt(sub)
paths.append(str(sub))
return paths
@require_bitsandbytes
@pytest.mark.single_gpu_tests
def test_arrow_4bit_opt125m_load_and_generate_with_local_adapters(self, ts_adapters_opt):
model_id = "peft-internal-testing/opt-125m"
# Quantization config (nf4, bf16 compute)
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=False,
)
with hub_online_once(model_id):
# Load quantized base model
base_model = AutoModelForCausalLM.from_pretrained(
model_id,
dtype=torch.bfloat16,
device_map="auto",
quantization_config=bnb_config,
)
with hub_online_once(model_id + "tokenizer"):
tok = AutoTokenizer.from_pretrained(model_id, use_fast=True)
# Build Arrow model from the locally created adapters
arrow_cfg = ArrowConfig(top_k=2, router_temperature=1.0, rng_seed=42)
model = create_arrow_model(
base_model=base_model,
task_specific_adapter_paths=ts_adapters_opt, # local dirs (each has adapter_config.json)
arrow_config=arrow_cfg,
).eval()
# Quick generate smoke test
inputs = tok("Hello world", return_tensors="pt")
inputs = {k: v.to(model.device) for k, v in inputs.items()}
with torch.no_grad():
out = model.generate(**inputs, max_new_tokens=8)
assert out is not None
assert out.shape[0] == 1 # batch size 1
@require_non_cpu
@require_bitsandbytes
class TestDtypeAutocastBnb:
"""Ensure that the dtype of the PEFT weights have the expected value, even when using quantized base models.
The autocast argument should be honored.
"""
model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
# no need to check each possible peft type, a selection should be enough
peft_types_to_test = ["lora", "vera", "lora-target-param"]
def check_dtype(self, quant_config, autocast_adapter_dtype, base_dtype, expected_dtype, peft_type, tmp_path=None):
"""helper function that creates the PEFT model and checks that the dtype of the PEFT adapter is as expected.
Checks:
- get_peft_model
- add_adapter
- PeftModel.from_pretrained
- load_adapter
"""
if peft_type == "lora":
peft_config = LoraConfig()
elif peft_type == "vera":
peft_config = VeraConfig()
elif peft_type == "lora-target-param":
peft_config = LoraConfig(target_modules=[], target_parameters=["q_proj.weight", "v_proj.weight"])
else:
raise ValueError("Argument must be one of 'lora' or 'vera'")
with hub_online_once(self.model_id):
model = AutoModelForCausalLM.from_pretrained(
self.model_id,
quantization_config=quant_config,
dtype=base_dtype,
device_map="auto",
)
model = get_peft_model(model, peft_config, autocast_adapter_dtype=autocast_adapter_dtype)
if peft_type != "lora-target-param":
# target_parameters does not allow multiple adapters on the same parameter
model.add_adapter("other", peft_config, autocast_adapter_dtype=autocast_adapter_dtype)
peft_params = [p for n, p in model.named_parameters() if model.prefix in n]
assert all(p.dtype == expected_dtype for p in peft_params)
model.save_pretrained(tmp_path)
del model
model = AutoModelForCausalLM.from_pretrained(
self.model_id,
quantization_config=quant_config,
dtype=base_dtype,
device_map="auto",
)
model = PeftModel.from_pretrained(model, tmp_path, autocast_adapter_dtype=autocast_adapter_dtype)
if peft_type != "lora-target-param":
# target_parameters does not allow multiple adapters on the same parameter
model.load_adapter(
tmp_path / "other", adapter_name="other", autocast_adapter_dtype=autocast_adapter_dtype
)
peft_params = [p for n, p in model.named_parameters() if model.prefix in n]
assert all(p.dtype == expected_dtype for p in peft_params)
@pytest.mark.parametrize("peft_type", peft_types_to_test)
@pytest.mark.parametrize("base_dtype", [torch.float32, torch.float16, torch.bfloat16])
def test_lora_no_quantization_dtype_no_autocast(self, base_dtype, peft_type, tmp_path):
# sanity check that without bnb, everything works as expected
quant_config = None
self.check_dtype(
quant_config,
autocast_adapter_dtype=False,
base_dtype=base_dtype,
expected_dtype=base_dtype,
peft_type=peft_type,
tmp_path=tmp_path,
)
@pytest.mark.parametrize("peft_type", peft_types_to_test)
@pytest.mark.parametrize("base_dtype", [torch.float32, torch.float16, torch.bfloat16])
def test_lora_no_quantization_dtype_autocast(self, base_dtype, peft_type, tmp_path):
# sanity check that without bnb, everything works as expected
quant_config = None
self.check_dtype(
quant_config,
autocast_adapter_dtype=True,
base_dtype=base_dtype,
expected_dtype=torch.float32,
peft_type=peft_type,
tmp_path=tmp_path,
)
@pytest.mark.parametrize("peft_type", peft_types_to_test)
@pytest.mark.parametrize("base_dtype", [torch.float32, torch.float16, torch.bfloat16])
def test_lora_4bit_bnb_dtype_no_autocast(self, base_dtype, peft_type, tmp_path):
# Ensure that the compute dtype of the 4bit weights is honored, see #2889
quant_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=base_dtype,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
)
self.check_dtype(
quant_config,
autocast_adapter_dtype=False,
base_dtype=base_dtype,
expected_dtype=base_dtype,
peft_type=peft_type,
tmp_path=tmp_path,
)
@pytest.mark.parametrize("peft_type", peft_types_to_test)
@pytest.mark.parametrize("base_dtype", [torch.float32, torch.float16, torch.bfloat16])
def test_lora_4bit_bnb_dtype_autocast(self, base_dtype, peft_type, tmp_path):
# With autocast, the adapter weights should always be in float32
quant_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=base_dtype,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
)
self.check_dtype(
quant_config,
autocast_adapter_dtype=True,
base_dtype=base_dtype,
expected_dtype=torch.float32,
peft_type=peft_type,
tmp_path=tmp_path,
)
@pytest.mark.parametrize("peft_type", peft_types_to_test)
def test_lora_4bit_bnb_dtype_no_autocast_compute_dtype_diverges(self, peft_type, tmp_path):
# In this test, the compute dtype of the bnb weights and the dtype of the base model diverge. In this case the
# bnb dtype should 'win'.
quant_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
)
self.check_dtype(
quant_config,
autocast_adapter_dtype=False,
base_dtype=torch.float16,
expected_dtype=torch.bfloat16,
peft_type=peft_type,
tmp_path=tmp_path,
)
@pytest.mark.parametrize("peft_type", peft_types_to_test)
@pytest.mark.parametrize("base_dtype", [torch.float16, torch.bfloat16])
@pytest.mark.xfail(reason="Currently, dtype casting with 8bit bnb does not work", strict=True)
def test_lora_8bit_bnb_dtype_no_autocast(self, base_dtype, peft_type, tmp_path):
# With 8bit bnb, the base layer carries no information about the intended dtype, thus we cannot cast to the same dtype
quant_config = BitsAndBytesConfig(load_in_8bit=True)
self.check_dtype(
quant_config,
autocast_adapter_dtype=False,
base_dtype=base_dtype,
expected_dtype=base_dtype,
peft_type=peft_type,
tmp_path=tmp_path,
)
@pytest.mark.parametrize("peft_type", peft_types_to_test)
def test_lora_8bit_bnb_dtype_no_autocast_float32(self, peft_type, tmp_path):
# for 8bit bnb with float32, everything works as expected
# TODO: once dtype != float32 works, merge this test with the one above
quant_config = BitsAndBytesConfig(load_in_8bit=True)
base_dtype = torch.float32
self.check_dtype(
quant_config,
autocast_adapter_dtype=False,
base_dtype=base_dtype,
expected_dtype=base_dtype,
peft_type=peft_type,
tmp_path=tmp_path,
)
@pytest.mark.parametrize("peft_type", peft_types_to_test)
@pytest.mark.parametrize("base_dtype", [torch.float16, torch.bfloat16])
def test_lora_8bit_bnb_dtype_autocast(self, base_dtype, peft_type, tmp_path):
# With 8bit bnb, the base layer carries no information about the intended dtype, thus we cannot cast to the same dtype
quant_config = BitsAndBytesConfig(load_in_8bit=True)
self.check_dtype(
quant_config,
autocast_adapter_dtype=True,
base_dtype=base_dtype,
expected_dtype=torch.float32,
peft_type=peft_type,
tmp_path=tmp_path,
)
@pytest.mark.parametrize("peft_type", peft_types_to_test)
def test_lora_8bit_bnb_dtype_autocast_float32(self, peft_type, tmp_path):
# for 8bit bnb with float32, everything works as expected
# TODO: once dtype != float32 works, merge this test with the one above
base_dtype = torch.float32
quant_config = BitsAndBytesConfig(load_in_8bit=True)
self.check_dtype(
quant_config,
autocast_adapter_dtype=True,
base_dtype=base_dtype,
expected_dtype=torch.float32,
peft_type=peft_type,
tmp_path=tmp_path,
)
if is_te_available():
import transformer_engine as te
from peft.tuners.lora.te import TeLinear
def _replace_with_te_linear(model, target_module_names):
"""Replace nn.Linear modules whose short name is in target_module_names with te.pytorch.Linear.
Weights and biases are copied from the original module. The replacement is done *in-place* (via ``setattr`` on the
parent module).
"""
replacements = []
for name, module in model.named_modules():
short_name = name.rsplit(".", 1)[-1] if "." in name else name
if short_name in target_module_names and isinstance(module, torch.nn.Linear):
replacements.append((name, short_name, module))
for name, short_name, module in replacements:
parent_name = name.rsplit(".", 1)[0] if "." in name else ""
parent = model.get_submodule(parent_name) if parent_name else model
te_linear = te.pytorch.Linear(
module.in_features,
module.out_features,
bias=module.bias is not None,
params_dtype=module.weight.dtype,
device=module.weight.device,
)
with torch.no_grad():
te_linear.weight.copy_(module.weight)
if module.bias is not None:
te_linear.bias.copy_(module.bias)
setattr(parent, short_name, te_linear)
@pytest.mark.skipif(not is_te_available(), reason="transformer_engine is not available")
class TestTransformerEngine:
"""Tests for LoRA with TransformerEngine layers.
Uses a standard OPT model with selected nn.Linear layers manually replaced by te.pytorch.Linear.
"""
model_id = "facebook/opt-125m"
te_target_modules = ["q_proj", "v_proj"]
@pytest.fixture
def model_with_te_layers(self):
model = AutoModelForCausalLM.from_pretrained(self.model_id, torch_dtype=torch.bfloat16)
_replace_with_te_linear(model, self.te_target_modules)
return model
@pytest.fixture
def tokenizer(self):
return AutoTokenizer.from_pretrained(self.model_id)
@pytest.fixture
def tokenized_inputs(self, tokenizer):
return tokenizer("Hello world", return_tensors="pt")
@require_torch_gpu
@pytest.mark.single_gpu_tests
def test_te_lora_wraps_te_linear_and_keeps_forward_working(self, model_with_te_layers, tokenized_inputs):
cfg = LoraConfig(target_modules=self.te_target_modules, r=2, lora_alpha=8)
lora_model = get_peft_model(model_with_te_layers, cfg).to("cuda")
wrapped_q_proj = lora_model.base_model.model.model.decoder.layers[0].self_attn.q_proj
assert isinstance(wrapped_q_proj, TeLinear)
assert "default" in wrapped_q_proj.lora_A and "default" in wrapped_q_proj.lora_B
assert wrapped_q_proj.get_base_layer().weight.requires_grad is False
assert wrapped_q_proj.lora_A["default"].weight.requires_grad is True
assert wrapped_q_proj.lora_B["default"].weight.requires_grad is True
inputs = {k: v.to("cuda") for k, v in tokenized_inputs.items()}
out = lora_model(**inputs)
assert out.logits.ndim == 3
@require_torch_gpu
@pytest.mark.single_gpu_tests
def test_te_lora_forward_matches_base_before_backward(self, model_with_te_layers, tokenized_inputs):
base_model = deepcopy(model_with_te_layers).to("cuda")
cfg = LoraConfig(target_modules=self.te_target_modules, r=4, lora_alpha=8)
lora_model = get_peft_model(model_with_te_layers, cfg).to("cuda")
inputs = {k: v.to("cuda") for k, v in tokenized_inputs.items()}
lora_model.eval()
base_model.eval()
with torch.no_grad():
lora_result = lora_model(**inputs).logits
base_result = base_model(**inputs).logits
assert torch.allclose(lora_result, base_result, rtol=1e-3, atol=1e-3)
@require_torch_gpu
@pytest.mark.single_gpu_tests
def test_te_lora_backward(self, model_with_te_layers, tokenized_inputs):
cfg = LoraConfig(target_modules=self.te_target_modules, r=4, lora_alpha=8)
lora_model = get_peft_model(model_with_te_layers, cfg).to("cuda")
optimizer = torch.optim.AdamW(lora_model.parameters())
loss_fn = torch.nn.CrossEntropyLoss()
inputs = {k: v.to("cuda") for k, v in tokenized_inputs.items()}
logits = lora_model(**inputs).logits
shift_logits = logits[:, :-1, :].contiguous()
shift_labels = inputs["input_ids"][:, 1:].contiguous()
loss = loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
assert torch.isfinite(loss), f"Loss is not finite: {loss.item()}"
loss.backward()
optimizer.step()
@pytest.mark.skipif(not hasattr(torch, "float8_e4m3fn"), reason="Platform does not support torch.float8_e4m3fn")
@require_torch_gpu
@pytest.mark.single_gpu_tests
class TestDtypeFp8:
"""Tests that float8 models work.
Note that at this time, these lower dtypes require a GPU, so these tests cannot be added to the standard CPU test
suite.
"""
@pytest.fixture(scope="class", autouse=True)
def setup_cleanup(self):
yield
clear_device_cache(garbage_collection=True)
@pytest.fixture
def model_high_prec(self):
model_id = "facebook/opt-125m"
model = AutoModelForCausalLM.from_pretrained(model_id, device_map=0)
return model
@pytest.fixture
def model_low_prec(self):
model_id = "facebook/opt-125m"
# only convert q_proj to fp8, otherwise we get nan results
modules_not_to_convert = ["embed_tokens", "lm_head", "v_proj", "k_proj", "out_proj", "fc1", "fc2"]
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map=0,
quantization_config=FineGrainedFP8Config(
modules_to_not_convert=modules_not_to_convert,
),
)
# sanity check
assert model.model.decoder.layers[0].self_attn.q_proj.weight.dtype == torch.float8_e4m3fn
return model
@pytest.mark.parametrize(
"config",
[
LoraConfig(target_modules=["q_proj", "v_proj"], init_lora_weights=False),
VeraConfig(target_modules=["q_proj", "v_proj"], init_weights=False),
RoadConfig(target_modules=["q_proj", "v_proj"], init_weights=False),
],
ids=lambda c: c.__class__.__name__,
)
def test_target_modules_float8_e4m3fn(self, model_high_prec, model_low_prec, config):
# Test should work with all adapters, but only testing a few here to save time and resources.
inputs = torch.arange(10).view(1, -1).to(model_low_prec.device)
# high precision
torch.manual_seed(0)
model_high_prec = get_peft_model(model_high_prec, config)
with torch.inference_mode():
# check that there are no errors
output_high_prec = model_high_prec(inputs).logits
# low precision
torch.manual_seed(0)
model_low_prec = get_peft_model(model_low_prec, config)
with torch.inference_mode():
# check that there are no errors
output_low_prec = model_low_prec(inputs).logits
# sanity check
assert torch.isfinite(output_low_prec).all()
# use relatively high tolerances because of low precision dtype
mse = ((output_low_prec - output_high_prec) ** 2).mean()
assert mse < 0.05
@pytest.mark.parametrize(
"config",
[
LoraConfig(target_modules=["q_proj", "v_proj"], init_lora_weights=False),
VeraConfig(target_modules=["q_proj", "v_proj"], init_weights=False),
RoadConfig(target_modules=["q_proj", "v_proj"], init_weights=False),
],
ids=lambda c: c.__class__.__name__,
)
@pytest.mark.xfail(reason="Merging with float8 not supported (yet)", strict=True)
def test_merge_with_float8_e4m3fn(self, model_high_prec, model_low_prec, config):
# Test should work with all adapters, but only testing a few here to save time and resources.
inputs = torch.arange(10).view(1, -1).to(model_low_prec.device)
# high precision
torch.manual_seed(0)
model_high_prec = get_peft_model(model_high_prec, config).merge_and_unload()
with torch.inference_mode():
# check that there are no errors
output_high_prec = model_high_prec(inputs).logits
# low precision
torch.manual_seed(0)
model_low_prec = get_peft_model(model_low_prec, config).merge_and_unload()
with torch.inference_mode():
# check that there are no errors
output_low_prec = model_low_prec(inputs).logits
# sanity check
assert torch.isfinite(output_low_prec).all()
# use relatively high tolerances because of low precision dtype
mse = ((output_low_prec - output_high_prec) ** 2).mean()
assert mse < 0.05
def test_lora_target_parameters_float8_e4m3fn(self, model_high_prec, model_low_prec):
inputs = torch.arange(10).view(1, -1).to(model_low_prec.device)
# high precision
torch.manual_seed(0)
config = LoraConfig(
target_modules=["k_proj", "v_proj"], target_parameters=["q_proj.weight"], init_lora_weights=False
)
model_high_prec = get_peft_model(model_high_prec, config)
with torch.inference_mode():
# check that there are no errors
output_high_prec = model_high_prec(inputs).logits
# low precision
torch.manual_seed(0)
model_low_prec = get_peft_model(model_low_prec, config)
with torch.inference_mode():
# check that there are no errors
output_low_prec = model_low_prec(inputs).logits
# sanity check
assert torch.isfinite(output_low_prec).all()
# use relatively high tolerances because of low precision dtype
mse = ((output_low_prec - output_high_prec) ** 2).mean()
assert mse < 0.05
def test_target_modules_no_autocast_prevserves_e4m3fn(self, model_low_prec):
# ensure that users can choose to keep the adapter weights in the same dtype as the original weights by passing
# autocast_adapter_dtype=False, even though the resulting model is not usable (no inference or training
# possible)
config = LoraConfig(target_modules=["q_proj", "v_proj"], init_lora_weights=False)
model = get_peft_model(model_low_prec, config, autocast_adapter_dtype=False)
q_proj = model.base_model.model.model.decoder.layers[0].self_attn.q_proj
assert q_proj.lora_A.default.weight.dtype == torch.float8_e4m3fn
assert q_proj.lora_B.default.weight.dtype == torch.float8_e4m3fn
### LoRA and Tensor Parallelism tests ###
WORLD_SIZE = 2
TINY_MODEL_ID = "peft-internal-testing/zephyr-smol_llama-100m-sft-full"
TARGET_MODULES = ["embed_tokens", "q_proj", "k_proj", "v_proj", "o_proj"]
TIMEOUT_BARRIER = datetime.timedelta(seconds=30)
TP_PLAN = {
"model.embed_tokens": "embedding_rowwise",
"model.layers.*.self_attn.q_proj": "colwise",
"model.layers.*.self_attn.k_proj": "colwise",
"model.layers.*.self_attn.v_proj": "colwise",
"model.layers.*.self_attn.o_proj": "rowwise",
"model.layers.*.mlp.gate_proj": "colwise",
"model.layers.*.mlp.up_proj": "colwise",
"model.layers.*.mlp.down_proj": "rowwise",
}
def _find_free_port():
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("", 0))
return s.getsockname()[1]
_BASE_PORT = _find_free_port()
def _setup_dist(rank, world_size, port):
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = str(port)
os.environ["LOCAL_RANK"] = str(rank)
os.environ["RANK"] = str(rank)
os.environ["WORLD_SIZE"] = str(rank)
dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)
def _teardown_dist():
dist.destroy_process_group()
def _test_function_wrapper(fn, rank, world_size, port, *extra_args):
try:
_setup_dist(rank, world_size, port)
fn(rank, world_size, port, *extra_args)
finally:
_teardown_dist()
def _test_lora_weight_synchronization(rank, world_size, port):
"""
Test that non-sharded LoRA weights are identical across ranks after training step.
"""
model = AutoModelForCausalLM.from_pretrained(TINY_MODEL_ID, tp_plan=TP_PLAN)
lora_config = LoraConfig(r=4, target_modules=TARGET_MODULES, init_lora_weights=True)
model = get_peft_model(model, lora_config)
torch.cuda.set_device(rank)
device = torch.device("cuda", rank)
model.to(device)
tokenizer = AutoTokenizer.from_pretrained(TINY_MODEL_ID)
inputs = tokenizer("Paris is the most beautiful city in the world.", return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
model.train()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
# Test that loss is finite and decreases over multiple steps
for _ in range(3):
outputs = model(**inputs, labels=inputs["input_ids"])
loss = outputs.loss
assert torch.isfinite(loss), f"Loss is not finite: {loss}"
loss.backward()
optimizer.step()
optimizer.zero_grad()
# Test that non-sharded LoRA weights are identical across ranks after training step
for name, module in model.named_modules():
if not isinstance(module, LoraLayer):
continue
base_layer = module.get_base_layer()
tp_plan = getattr(base_layer, "_hf_tp_plan", None)
if tp_plan == "colwise":
weight = module.lora_A["default"].weight.data.contiguous()
gathered = [torch.zeros_like(weight) for _ in range(world_size)]
dist.all_gather(gathered, weight)
for i, g in enumerate(gathered):
assert torch.allclose(weight, g), f"{name}.lora_A differs between rank {rank} and rank {i}"
elif tp_plan == "rowwise":
weight = module.lora_B["default"].weight.data.contiguous()
gathered = [torch.zeros_like(weight) for _ in range(world_size)]
dist.all_gather(gathered, weight)
for i, g in enumerate(gathered):
assert torch.allclose(weight, g), f"{name}.lora_B differs between rank {rank} and rank {i}"
elif tp_plan == "embedding_rowwise":
weight = module.lora_embedding_B["default"].data.contiguous()
gathered = [torch.zeros_like(weight) for _ in range(world_size)]
dist.all_gather(gathered, weight)
for i, g in enumerate(gathered):
assert torch.allclose(weight, g), f"{name}.lora_embedding_B differs between rank {rank} and rank {i}"
def _test_load_from_checkpoint(rank, world_size, port, tmp_dir):
"""
Test that loading from a checkpoint correctly handles the sharding of LoRA weights according to the TP plan.
"""
if rank == 0:
plain_model = AutoModelForCausalLM.from_pretrained(TINY_MODEL_ID)
lora_config = LoraConfig(r=4, target_modules=TARGET_MODULES, init_lora_weights=True)
plain_model = get_peft_model(plain_model, lora_config)
plain_model.save_pretrained(tmp_dir)
dist.monitored_barrier(timeout=TIMEOUT_BARRIER, wait_all_ranks=True)
torch.cuda.set_device(rank)
device = torch.device("cuda", rank)
tp_base = AutoModelForCausalLM.from_pretrained(TINY_MODEL_ID, tp_plan=TP_PLAN)
tp_base.to(device)
tp_model = PeftModel.from_pretrained(tp_base, tmp_dir)
for name, module in tp_model.named_modules():
if not isinstance(module, LoraLayer):
continue
base_layer = module.get_base_layer()
tp_plan = getattr(base_layer, "_hf_tp_plan", None)
if tp_plan == "colwise":
# lora_B output dim must match base layer output dim
lora_b_out = module.lora_B["default"].weight.shape[0]
base_layer_out = base_layer.weight.shape[0]
assert lora_b_out == base_layer_out, (
f"{name}: lora_B out_dim {lora_b_out} != local base out_dim {base_layer_out}"
)
elif tp_plan == "rowwise":
# lora_A input dim must match base layer input dim
lora_a_in = module.lora_A["default"].weight.shape[1]
base_layer_in = base_layer.weight.shape[1]
assert lora_a_in == base_layer_in, (
f"{name}: lora_A in_dim {lora_a_in} != local base in_dim {base_layer_in}"
)
elif tp_plan == "embedding_rowwise":
# lora_embedding_A vocab size must match base layer vocab size
lora_emb_vocab = module.lora_embedding_A["default"].shape[1] # Lora embedding weights are (r, vocab_size)
base_emb_vocab = base_layer.weight.shape[0]
assert lora_emb_vocab == base_emb_vocab, (
f"{name}: lora_embedding_A vocab {lora_emb_vocab} != local base vocab {base_emb_vocab}"
)
tokenizer = AutoTokenizer.from_pretrained(TINY_MODEL_ID)
inputs = tokenizer("The capital of France is Paris.", return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
tp_model.eval()
with torch.no_grad():
outputs = tp_model(**inputs, labels=inputs["input_ids"])
assert torch.isfinite(outputs.loss), f"Loss not finite after checkpoint load: {outputs.loss}"
def _test_save_unsharded_weights(rank, world_size, port, tmp_dir_reference, tmp_dir_tp):
"""
Test that saving a TP PEFT model produces fully unsharded weights identical to the original non-TP weights.
Flow:
1. Rank 0: create a plain (non-TP) PEFT model, save it as the reference.
2. All ranks: load the TP base model, load the reference PEFT weights (shards on load), then save again.
3. Rank 0: compare the re-saved state dict against the original reference, they must match exactly.
"""
if rank == 0:
plain_model = AutoModelForCausalLM.from_pretrained(TINY_MODEL_ID)
lora_config = LoraConfig(r=4, target_modules=TARGET_MODULES, init_lora_weights=True)
plain_model = get_peft_model(plain_model, lora_config)
plain_model.save_pretrained(tmp_dir_reference)
dist.monitored_barrier(timeout=TIMEOUT_BARRIER, wait_all_ranks=True)
torch.cuda.set_device(rank)
device = torch.device("cuda", rank)
tp_base = AutoModelForCausalLM.from_pretrained(TINY_MODEL_ID, tp_plan=TP_PLAN)
tp_base.to(device)
tp_model = PeftModel.from_pretrained(tp_base, tmp_dir_reference)
tp_model.save_pretrained(tmp_dir_tp)
dist.monitored_barrier(timeout=TIMEOUT_BARRIER, wait_all_ranks=True)
if rank == 0:
reference_sd = load_file(f"{tmp_dir_reference}/adapter_model.safetensors")
saved_sd = load_file(f"{tmp_dir_tp}/adapter_model.safetensors")
assert set(reference_sd.keys()) == set(saved_sd.keys()), (
f"State dict keys differ.\n reference: {sorted(reference_sd.keys())}\n saved: {sorted(saved_sd.keys())}"
)
for key in reference_sd:
assert torch.allclose(reference_sd[key], saved_sd[key], atol=1e-6), (
f"Weight mismatch for '{key}': max diff = {(reference_sd[key] - saved_sd[key]).abs().max()}"
)
def _test_multiple_adapters(rank, world_size, port):
"""Two LoRA adapters coexist on a TP model and can be switched between."""
torch.cuda.set_device(rank)
device = torch.device("cuda", rank)
model = AutoModelForCausalLM.from_pretrained(TINY_MODEL_ID, tp_plan=TP_PLAN)
model.to(device)
for adapter_name in ["adapter_a", "adapter_b"]:
lora_config = LoraConfig(r=4, target_modules=TARGET_MODULES, init_lora_weights=True)
model = get_peft_model(model, lora_config, adapter_name=adapter_name)
tokenizer = AutoTokenizer.from_pretrained(TINY_MODEL_ID)
inputs = tokenizer("What is the capital of France?", return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
model.eval()
with torch.no_grad():
for adapter_name in ["adapter_a", "adapter_b"]:
model.set_adapter(adapter_name)
outputs = model(**inputs, labels=inputs["input_ids"])
assert torch.isfinite(outputs.loss), f"Loss not finite with adapter '{adapter_name}': {outputs.loss}"
def _test_load_adapter_forward(rank, world_size, port, tmp_dir_reference):
"""
Test that load_adapter (with a peft_config) works with a TP base model and the forward pass produces the same loss
on every rank and that it is finite.
This exercises the low-level API path where no PeftModel/tuner is created, so TP info must be stored on the lora
modules themselves (via _tp_info) rather than on the tuner.
"""
torch.cuda.set_device(rank)
device = torch.device("cuda", rank)
if rank == 0:
plain_model = AutoModelForCausalLM.from_pretrained(TINY_MODEL_ID)
lora_config = LoraConfig(r=4, target_modules=TARGET_MODULES, init_lora_weights=True)
plain_model = get_peft_model(plain_model, lora_config)
plain_model.save_pretrained(tmp_dir_reference)
dist.monitored_barrier(timeout=TIMEOUT_BARRIER, wait_all_ranks=True)
model = AutoModelForCausalLM.from_pretrained(TINY_MODEL_ID, tp_plan=TP_PLAN)
model.load_adapter(tmp_dir_reference)
model.to(device)
for mod in model.modules():
if isinstance(mod, LoraLayer):
assert hasattr(mod, "_tp_info"), "load_adapter did not store TP info on the LoRA module"
tokenizer = AutoTokenizer.from_pretrained(TINY_MODEL_ID)
inputs = tokenizer("Paris is the capital of France.", return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
model.eval()
with torch.no_grad():
outputs = model(**inputs, labels=inputs["input_ids"])
all_losses = [torch.empty_like(outputs.loss) for _ in range(world_size)]
dist.all_gather(all_losses, outputs.loss)
assert all(loss.item() == all_losses[0].item() for loss in all_losses), (
f"Losses differ across ranks: {[loss.item() for loss in all_losses]}"
)
assert torch.isfinite(outputs.loss), f"Loss is not finite: {outputs.loss}"
def _test_load_adapter_save(rank, world_size, port, tmp_dir_reference, tmp_dir_tp):
"""
Test that get_peft_model_state_dict correctly gathers unsharded TP weights when using load_adapter with a
peft_config.
Flow:
1. Rank 0: create a plain (non-TP) model, load adapter from config, save state dict as reference.
2. All ranks: load TP base model, load adapter from config, save state dict via get_peft_model_state_dict.
3. Rank 0: compare re-saved state dict against the reference, keys and values must match.
"""
if rank == 0:
plain_model = AutoModelForCausalLM.from_pretrained(TINY_MODEL_ID)
lora_config = LoraConfig(r=4, target_modules=TARGET_MODULES, init_lora_weights=True)
plain_model = get_peft_model(plain_model, lora_config)
plain_model.save_pretrained(tmp_dir_reference)
dist.monitored_barrier(timeout=TIMEOUT_BARRIER, wait_all_ranks=True)
torch.cuda.set_device(rank)
device = torch.device("cuda", rank)
tp_base = AutoModelForCausalLM.from_pretrained(TINY_MODEL_ID, tp_plan=TP_PLAN)
tp_base.load_adapter(tmp_dir_reference)
tp_base.to(device)
tp_sd = get_peft_model_state_dict(tp_base)
if rank == 0:
tmp_dir_tp.mkdir(exist_ok=True)
save_file(tp_sd, f"{tmp_dir_tp}/adapter_model.safetensors")
dist.monitored_barrier(timeout=TIMEOUT_BARRIER, wait_all_ranks=True)
if rank == 0:
reference_sd = load_file(f"{tmp_dir_reference}/adapter_model.safetensors")
saved_sd = load_file(f"{tmp_dir_tp}/adapter_model.safetensors")
prefix = "base_model.model."
reference_sd = {k[len(prefix) :]: v for k, v in reference_sd.items() if k.startswith(prefix)}
assert set(reference_sd.keys()) == set(saved_sd.keys()), (
f"State dict keys differ.\n reference: {sorted(reference_sd.keys())}\n saved: {sorted(saved_sd.keys())}"
)
for key in reference_sd:
if reference_sd[key].shape != saved_sd[key].shape:
raise AssertionError(
f"Shape mismatch for '{key}': reference shape {reference_sd[key].shape} vs saved shape "
f"{saved_sd[key].shape}"
)
torch.testing.assert_close(
reference_sd[key].to(torch.bfloat16),
saved_sd[key].to(torch.bfloat16),
rtol=1e-5,
atol=1e-6,
msg=f"Weight mismatch for '{key}'",
)
# transformers >= 5.4.0 is required for TP integration
# transformers >= 5.6.0 is required for PreTrainedModel.load_adapter with TP models integration
is_transformers_ge_v5_6_0 = packaging.version.parse(transformers.__version__) >= packaging.version.parse("5.6.0")
@require_torch_gpu
@pytest.mark.skipif(
not is_transformers_ge_v5_6_0, reason="transformers TP integration supported for transformers >= 5.6.0"
)
@pytest.mark.multi_gpu_tests
class TestLoraTensorParallel:
def _spawn(self, fn, *extra_args, port_offset=0):
port = _BASE_PORT + port_offset
wrapped_fn = partial(_test_function_wrapper, fn)
mp.spawn(wrapped_fn, args=(WORLD_SIZE, port) + extra_args, nprocs=WORLD_SIZE, join=True)
def test_lora_weight_synchronization(self):
self._spawn(_test_lora_weight_synchronization, port_offset=0)
def test_from_checkpoint(self, tmp_path):
self._spawn(_test_load_from_checkpoint, tmp_path, port_offset=1)
def test_save_unsharded_weights(self, tmp_path):
tmp_dir_reference = tmp_path / "reference"
tmp_dir_tp = tmp_path / "tp"
self._spawn(_test_save_unsharded_weights, tmp_dir_reference, tmp_dir_tp, port_offset=3)
def test_multiple_adapters(self):
self._spawn(_test_multiple_adapters, port_offset=4)
def test_load_adapter_forward(self, tmp_path):
self._spawn(_test_load_adapter_forward, tmp_path, port_offset=5)
def test_load_adapter_save(self, tmp_path):
tmp_dir_reference = tmp_path / "reference"
tmp_dir_tp = tmp_path / "tp"
self._spawn(_test_load_adapter_save, tmp_dir_reference, tmp_dir_tp, port_offset=6)
@pytest.mark.single_gpu_tests
@require_bitsandbytes
def test_kappatune_with_4bit_model():
"""Test that KappaTune works with 4-bit quantized models on GPU."""
import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from peft.helpers import find_kappa_target_modules
# Use a very small model for faster testing
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_quant_type="nf4",
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
"hf-internal-testing/tiny-random-LlamaForCausalLM",
quantization_config=quantization_config,
device_map="cuda",
torch_dtype=torch.float16,
)
# Run KappaTune
targets = find_kappa_target_modules(model, top_p=0.3)
# Basic assertions
assert isinstance(targets, dict)
assert "target_modules" in targets
assert isinstance(targets["target_modules"], list)
assert len(targets["target_modules"]) > 0, "Should return at least some target modules"
Reference in New Issue View Git Blame Copy Permalink