QLoRA

# pip install "transformers>=4.41" "peft>=0.11" "bitsandbytes>=0.43" accelerate datasets
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

MODEL_ID = "meta-llama/Llama-3-70B"

# Bước 1: Cấu hình QLoRA — NF4 + Double Quantization + BF16 compute
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",          # NormalFloat 4-bit (không phải "fp4"!)
    bnb_4bit_use_double_quant=True,     # tiết kiệm thêm ~0.37 bit/param
    bnb_4bit_compute_dtype=torch.bfloat16,  # tính toán bằng BF16 sau dequant
)

# Bước 2: Load base model ở 4-bit
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    trust_remote_code=False,
)
# 70B × 4.125 bit = ~36GB + overhead ≈ 38GB VRAM cho weights

# Bước 3: Chuẩn bị cho k-bit training (cast LN sang FP32, bật gradient checkpointing)
model = prepare_model_for_kbit_training(
    model, use_gradient_checkpointing=True
)

# Bước 4: Gắn LoRA adapter ở BF16 — chỉ phần này có gradient
lora_config = LoraConfig(
    r=16,                    # rank; 8-64 là phổ biến
    lora_alpha=32,           # thường = 2 × r
    target_modules=[         # Llama/Mistral: thường đủ 7 proj
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
    ],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable: ~30M (0.04% của 70B) → gradient + optimizer rất nhẹ

# Bước 5: Training args — dùng paged optimizer
training_args = TrainingArguments(
    output_dir="./qlora-llama3-70b",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,
    learning_rate=2e-4,              # LR cao hơn full FT vì adapter nhỏ
    max_steps=1000,
    optim="paged_adamw_8bit",        # PAGED — trick tránh OOM
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    bf16=True,                       # tính toán BF16
    logging_steps=10,
    save_steps=200,
    gradient_checkpointing=True,
)

# Bước 6: Train
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
dataset = load_dataset("timdettmers/openassistant-guanaco", split="train")

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    args=training_args,
    dataset_text_field="text",
    max_seq_length=2048,
    peft_config=lora_config,
)
trainer.train()

# Bước 7: Lưu — chỉ adapter ~100MB
model.save_pretrained("./qlora-adapter")
# Inference sau này: load base model (có thể 4-bit) + merge adapter

import torch
import bitsandbytes as bnb
from bitsandbytes.functional import quantize_nf4, dequantize_nf4

# Giả lập 1 layer linear weight 4096x4096 (giống Llama hidden size)
W_fp16 = torch.randn(4096, 4096, dtype=torch.float16, device="cuda") * 0.02

# Quantize sang NF4 — block size 64
W_nf4, quant_state = quantize_nf4(W_fp16, blocksize=64)

print(f"FP16 size: {W_fp16.nelement() * 2 / 1e6:.2f} MB")
print(f"NF4 size:  {W_nf4.nelement() / 2 / 1e6:.2f} MB (chia 2 vì pack 2 val/byte)")
print(f"→ Tỉ lệ nén: {2 * W_fp16.nelement() / W_nf4.nelement():.2f}x")

# Dequantize lại — chỉ dùng khi cần tính toán
W_recon = dequantize_nf4(W_nf4, quant_state)
err = (W_fp16 - W_recon).abs().mean().item()
print(f"Sai số trung bình |W - W_recon|: {err:.5f}")
print(f"Tương đối: {err / W_fp16.abs().mean().item() * 100:.2f}%")

# Forward giả: x @ W — cần dequant trước
x = torch.randn(1, 4096, dtype=torch.float16, device="cuda")
y_fp16 = x @ W_fp16
y_nf4  = x @ W_recon
print(f"|y_fp16 - y_nf4| mean: {(y_fp16 - y_nf4).abs().mean():.5f}")

# Trong thực tế, bnb.nn.Linear4bit tự làm bước dequant trong CUDA kernel
layer = bnb.nn.Linear4bit(
    4096, 4096,
    bias=False,
    compute_dtype=torch.bfloat16,
    quant_type="nf4",
)
layer.weight.data = W_fp16  # trước khi .cuda() được tự quantize