LoRA/QLoRA — 오픈소스 모델 저비용 파인튜닝

pip install transformers peft bitsandbytes datasets accelerate trl

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer
from datasets import load_dataset

# 1. 4bit 양자화 설정 (QLoRA)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

# 2. 모델 로드
model_name = "mistralai/Mistral-7B-Instruct-v0.2"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# 3. LoRA 설정
lora_config = LoraConfig(
    r=16,                # Rank (낮을수록 파라미터 적음, 보통 4~64)
    lora_alpha=32,       # 스케일링 (보통 r의 2배)
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable params: 6,815,744 || all params: 3,758,489,600 || trainable%: 0.18

# 4. 학습 파라미터
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,   # 유효 배치: 4×4=16
    learning_rate=2e-4,
    fp16=True,
    logging_steps=10,
    save_steps=100,
    eval_steps=100,
    warmup_ratio=0.05,
    lr_scheduler_type="cosine",
    report_to="none",  # "wandb"로 설정 시 실험 추적
)

# 5. 데이터셋 준비
dataset = load_dataset("json", data_files="training.jsonl")

def format_chat(sample):
    """ChatML 형식으로 텍스트 변환"""
    messages = sample["messages"]
    text = ""
    for msg in messages:
        if msg["role"] == "system":
            text += f"<|system|>\n{msg['content']}\n"
        elif msg["role"] == "user":
            text += f"<|user|>\n{msg['content']}\n"
        elif msg["role"] == "assistant":
            text += f"<|assistant|>\n{msg['content']}\n"
    return {"text": text}

formatted = dataset["train"].map(format_chat)

# 6. SFT 학습
trainer = SFTTrainer(
    model=model,
    train_dataset=formatted,
    args=training_args,
    tokenizer=tokenizer,
    max_seq_length=2048,
    dataset_text_field="text",
)

trainer.train()

# LoRA 어댑터만 저장 (작은 파일)
model.save_pretrained("./my-lora-adapter")
tokenizer.save_pretrained("./my-lora-adapter")

# 원본 모델에 LoRA 병합 (추론 최적화)
from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto",
)

merged_model = PeftModel.from_pretrained(base_model, "./my-lora-adapter")
merged_model = merged_model.merge_and_unload()  # 가중치 병합

merged_model.save_pretrained("./my-merged-model")

from peft import PeftModel

# 방법 1: 어댑터 로드
base_model = AutoModelForCausalLM.from_pretrained(model_name, ...)
model = PeftModel.from_pretrained(base_model, "./my-lora-adapter")

# 방법 2: 병합된 모델 직접 로드
model = AutoModelForCausalLM.from_pretrained("./my-merged-model", ...)

def generate(prompt: str, max_new_tokens: int = 512) -> str:
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)