LoRA

from peft import LoraConfig, get_peft_model, TaskType
from transformers import AutoModelForCausalLM

# 1. Load base model (fp16 trên GPU)
base_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3-8B",
    torch_dtype="float16",
    device_map="auto",
)

# 2. Cấu hình LoRA
config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=8,                    # Rank — trade-off giữa chất lượng/chi phí
    lora_alpha=16,          # Scaling: W' = W + (16/8)·BA = W + 2·BA
    target_modules=[        # Áp dụng cho attention projections
        "q_proj", "v_proj", # Tối thiểu: Q và V
        # "k_proj", "o_proj"  # Tùy chọn: K và O cho biểu diễn mạnh hơn
    ],
    lora_dropout=0.05,      # Regularization
    bias="none",            # Không huấn luyện bias
    modules_to_save=None,   # Không có module full-trainable
)

# 3. Gắn adapter
model = get_peft_model(base_model, config)
model.print_trainable_parameters()
# → trainable params: 4,194,304 || all params: 8,034,078,720
# → trainable%: 0.0522

# 4. Huấn luyện như bình thường
# trainer = Trainer(model=model, ...)
# trainer.train()

# 5. Lưu adapter (chỉ vài MB, không phải 14GB mô hình gốc!)
model.save_pretrained("./my-lora-adapter")

# 6. Nạp lại khi inference
from peft import PeftModel
base = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3-8B")
model = PeftModel.from_pretrained(base, "./my-lora-adapter")

# 7. (Tùy chọn) Gộp để tăng tốc inference
model = model.merge_and_unload()  # W' = W + BA, xoá adapter

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    Trainer,
)
from peft import (
    LoraConfig,
    get_peft_model,
    prepare_model_for_kbit_training,
    TaskType,
)

# ===== 1. 4-bit quantization config =====
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",          # NormalFloat 4-bit
    bnb_4bit_compute_dtype=torch.bfloat16,  # Tính toán ở bf16
    bnb_4bit_use_double_quant=True,     # Double quantization
)

# ===== 2. Load model 4-bit =====
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3-70B",
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3-70B")

# ===== 3. Chuẩn bị cho k-bit training =====
# Bật gradient checkpointing, cast layer norm về fp32,
# đóng băng mọi tham số non-LoRA
model = prepare_model_for_kbit_training(
    model,
    use_gradient_checkpointing=True,
)

# ===== 4. Gắn LoRA adapter =====
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=16,                           # Rank cao hơn LoRA thường
    lora_alpha=32,                  # α = 2r
    target_modules=[                # Tất cả linear layer trong attention
        "q_proj", "k_proj",
        "v_proj", "o_proj",
        # "gate_proj", "up_proj",   # Cũng có thể thêm MLP
        # "down_proj",
    ],
    lora_dropout=0.05,
    bias="none",
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# → trainable: 40M / 70B = 0.057%

# ===== 5. Training với paged optimizer =====
args = TrainingArguments(
    output_dir="./qlora-llama3-70b",
    num_train_epochs=3,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    optim="paged_adamw_32bit",      # Offload optimizer state
    learning_rate=2e-4,
    bf16=True,
    logging_steps=10,
    save_strategy="epoch",
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
)

# trainer = Trainer(model=model, args=args, ...)
# trainer.train()

# ===== 6. Lưu adapter (~100MB thay vì 140GB!) =====
model.save_pretrained("./qlora-llama3-adapter")