import torch

# ===========================
# LoRA: W_new = W_frozen + B @ A
# ===========================
d = 4096  # hidden dimension (e.g., LLaMA 7B)
r = 16    # LoRA rank (hyperparameter)

# Original weight (FROZEN — no gradient!)
W_frozen = torch.randn(d, d, requires_grad=False)

# LoRA adapters (TRAINABLE — small!)
A = torch.randn(d, r, requires_grad=True)   # "down projection"
B = torch.zeros(r, d, requires_grad=True)   # "up projection" (init zeros!)
# B initialized to ZEROS → at start, B@A = 0 → model unchanged!
# This is important: training STARTS from original model behavior.

# Forward pass
x = torch.randn(1, d)  # input
output = x @ W_frozen + x @ A @ B  # original + LoRA delta
#         ↑ frozen       ↑ trainable (tiny!)

# Parameter count comparison
full_params = d * d                  # 16,777,216
lora_params = d * r + r * d          # 131,072
reduction = full_params / lora_params
print(f"Full:  {full_params:>12,} params per weight matrix")
print(f"LoRA:  {lora_params:>12,} params (r={r})")
print(f"Reduction: {reduction:.0f}× fewer parameters!")
# Full:    16,777,216 params
# LoRA:       131,072 params (r=16)
# Reduction: 128× fewer parameters!

# For entire Mistral-7B (applied to Q, K, V, O projections):
# Full trainable: ~7.24B parameters
# LoRA trainable: ~10M parameters (0.14% of total!)
# Saved model size: full=14.5GB vs LoRA adapter=~40MB

# pip install peft
from peft import LoraConfig, get_peft_model, TaskType
from transformers import AutoModelForCausalLM, AutoTokenizer

# ===========================
# 1. Load base model
# ===========================
model_name = "mistralai/Mistral-7B-v0.3"
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name)

print(f"Base model params: {model.num_parameters():,}")
# Base model params: 7,241,732,096 (7.24B!)

# ===========================
# 2. Configure LoRA
# ===========================
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,  # or SEQ_CLS, TOKEN_CLS, SEQ_2_SEQ_LM
    r=16,                         # rank (8-64, sweet spot: 16)
    lora_alpha=32,                # scaling (usually 2×r)
    lora_dropout=0.05,            # dropout
    target_modules=[              # which layers get LoRA
        "q_proj", "k_proj",        # attention query & key
        "v_proj", "o_proj",        # attention value & output
        # "gate_proj", "up_proj", "down_proj",  # MLP layers (optional, more params)
    ],
    bias="none",                  # don't train biases
)

# ===========================
# 3. Apply LoRA to model
# ===========================
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable params: 13,631,488 || all params: 7,255,363,584 || trainable%: 0.188%
# → Only 13.6M trainable out of 7.24B! (0.188%)
# → Saved adapter size: ~55MB (not 14.5GB!)

# ===========================
# 4. Train with standard Trainer (IDENTICAL to Page 2!)
# ===========================
# trainer = Trainer(model=model, args=args, ...)
# trainer.train()  # ← same API! PEFT is transparent to Trainer.

# ===========================
# 5. Save LoRA adapter (tiny!)
# ===========================
model.save_pretrained("./lora-adapter")
# Saves ONLY the LoRA weights (~55MB)
# NOT the full 14.5GB base model!

# ===========================
# 6. Load LoRA adapter later
# ===========================
from peft import PeftModel
base = AutoModelForCausalLM.from_pretrained(model_name)
model_with_lora = PeftModel.from_pretrained(base, "./lora-adapter")
# → base model + LoRA adapter = fine-tuned model!

# pip install bitsandbytes
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch

# ===========================
# 1. 4-bit quantization config (QLoRA standard!)
# ===========================
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,                    # load model in 4-bit!
    bnb_4bit_quant_type="nf4",            # NormalFloat4 (best for LLMs)
    bnb_4bit_compute_dtype=torch.float16, # compute in FP16 (speed)
    bnb_4bit_use_double_quant=True,       # quantize the quantization constants!
    # → extra 0.4GB savings for 7B model
)

# ===========================
# 2. Load 7B model in 4-bit (~3.6GB instead of 14.5GB!)
# ===========================
model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.3",
    quantization_config=bnb_config,
    device_map="auto",  # auto-place layers across GPU/CPU
)

print(f"Model loaded! Memory: {model.get_memory_footprint()/1e9:.1f} GB")
# Model loaded! Memory: 3.6 GB (bukan 14.5GB!)
# → Muat di Colab T4 (16GB)! Sisa ~12GB untuk training.

# ===========================
# 3. Comparison: precision vs VRAM
# ===========================
# FP32 (32-bit): 7B × 4 bytes = ~29 GB  ← TIDAK muat di A100 40GB
# FP16 (16-bit): 7B × 2 bytes = ~14.5 GB ← Muat di A100, tidak T4
# INT8 (8-bit):  7B × 1 byte  = ~7.2 GB  ← Muat di T4, ketat
# NF4 (4-bit):   7B × 0.5 byte = ~3.6 GB ← Muat NYAMAN di T4! ⭐

#!/usr/bin/env python3
"""
🔥 Fine-Tune Mistral-7B with QLoRA on Google Colab T4 (free!)
From "impossible" (120GB VRAM) to "free" (12GB VRAM) with QLoRA.
"""

# ═══════════════════════════════════════
# STEP 0: INSTALL (run in Colab!)
# ═══════════════════════════════════════
# !pip install -q transformers datasets accelerate peft bitsandbytes trl

import torch
from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
    BitsAndBytesConfig, TrainingArguments
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import load_dataset
from trl import SFTTrainer  # Supervised Fine-Tuning Trainer

# ═══════════════════════════════════════
# STEP 1: LOAD MODEL IN 4-BIT
# ═══════════════════════════════════════
model_name = "mistralai/Mistral-7B-v0.3"
# Alternatives: "meta-llama/Llama-3.2-1B" (smaller, easier)
#               "Qwen/Qwen2.5-7B" (multilingual)

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained(
    model_name, quantization_config=bnb_config,
    device_map="auto", trust_remote_code=True)

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

print(f"Model VRAM: {model.get_memory_footprint()/1e9:.1f} GB")
# Model VRAM: 3.6 GB ← 4-bit! (was 14.5 GB in FP16)

# ═══════════════════════════════════════
# STEP 2: PREPARE MODEL FOR TRAINING
# ═══════════════════════════════════════
model = prepare_model_for_kbit_training(model)
# Enables gradient checkpointing + casts to correct dtypes

# ═══════════════════════════════════════
# STEP 3: ADD LORA ADAPTERS
# ═══════════════════════════════════════
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# trainable params: 13,631,488 || all params: 7,255,363,584 || trainable%: 0.19%

# ═══════════════════════════════════════
# STEP 4: LOAD INSTRUCTION DATASET
# ═══════════════════════════════════════
dataset = load_dataset("databricks/databricks-dolly-15k", split="train")
# 15k instruction-response pairs (open-source!)

def format_instruction(example):
    if example.get("context") and example["context"].strip():
        text = f"""### Instruction:\n{example['instruction']}\n\n### Context:\n{example['context']}\n\n### Response:\n{example['response']}{tokenizer.eos_token}"""
    else:
        text = f"""### Instruction:\n{example['instruction']}\n\n### Response:\n{example['response']}{tokenizer.eos_token}"""
    return {"text": text}

dataset = dataset.map(format_instruction)
print(f"Training examples: {len(dataset)}")
print(f"Sample:\n{dataset[0]['text'][:300]}...")

# ═══════════════════════════════════════
# STEP 5: TRAINING ARGUMENTS
# ═══════════════════════════════════════
args = TrainingArguments(
    output_dir="./qlora-mistral",
    num_train_epochs=1,                  # 1 epoch for 15k examples → enough!
    per_device_train_batch_size=2,       # small batch (4-bit model is large)
    gradient_accumulation_steps=8,       # effective batch = 2 × 8 = 16
    learning_rate=2e-4,                  # higher LR for LoRA (2e-4, not 2e-5!)
    weight_decay=0.01,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    fp16=True,
    logging_steps=25,
    save_strategy="steps",
    save_steps=200,
    save_total_limit=2,
    optim="paged_adamw_8bit",           # 8-bit optimizer → less VRAM!
    gradient_checkpointing=True,        # saves ~30% VRAM (slower training)
    max_grad_norm=0.3,                  # gradient clipping
    report_to="none",
)

# ═══════════════════════════════════════
# STEP 6: TRAIN WITH SFTTrainer (from TRL library)
# ═══════════════════════════════════════
trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=dataset,
    tokenizer=tokenizer,
    dataset_text_field="text",          # column with formatted text
    max_seq_length=512,                 # max token length
    packing=True,                       # pack short examples into one sequence!
    # packing = True → 2-3× faster training (less padding waste)
)

print("🏋️ Training QLoRA...")
print(f"VRAM used: {torch.cuda.memory_allocated()/1e9:.1f} GB")
trainer.train()

# ═══════════════════════════════════════
# STEP 7: SAVE & TEST
# ═══════════════════════════════════════
trainer.save_model("./qlora-mistral-final")

# Test!
prompt = "### Instruction:\nWhat is machine learning?\n\n### Response:\n"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
with torch.no_grad():
    outputs = model.generate(**inputs, max_new_tokens=200,
        temperature=0.7, top_p=0.9, repetition_penalty=1.2)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
print("\\n🏆 QLoRA fine-tuning complete!")

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

# ===========================
# 1. Load base model (FP16, NOT 4-bit!) + LoRA adapter
# ===========================
base_model = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.3",
    torch_dtype=torch.float16,
    device_map="auto",
)
# Load LoRA adapter on top
model = PeftModel.from_pretrained(base_model, "./qlora-mistral-final")

# ===========================
# 2. Merge LoRA into base weights
# ===========================
model = model.merge_and_unload()
# → W_new = W_base + B @ A (permanently merged!)
# → No more LoRA overhead at inference
# → Same speed as original model

# ===========================
# 3. Save merged model
# ===========================
model.save_pretrained("./mistral-7b-finetuned-merged")
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.3")
tokenizer.save_pretrained("./mistral-7b-finetuned-merged")
# → Full 14.5GB model with LoRA baked in
# → Deploy like any normal model (TF Serving, vLLM, TGI, etc.)

# ===========================
# 4. Push to Hub!
# ===========================
# model.push_to_hub("username/mistral-7b-instruction-tuned")
# tokenizer.push_to_hub("username/mistral-7b-instruction-tuned")

# ===========================
# When to merge vs keep separate?
# ===========================
# MERGE: for production deployment (no PEFT dependency needed)
# KEEP SEPARATE: when you have multiple LoRA adapters for different tasks
#   e.g., base_model + lora_sentiment, base_model + lora_translation
#   → swap adapters at runtime without loading model twice!