Fine-Tune a 7B LLM with LoRA on a Single 24 GB GPU

python -m venv venv && source venv/bin/activate

pip install \
  torch==2.3.1 \
  transformers>=4.38.0 \
  peft>=0.10.0 \
  bitsandbytes>=0.43.0 \
  datasets>=2.16.0 \
  accelerate>=0.27.0

pip install flash-attn --no-build-isolation

"""
LoRA fine-tuning of tiiuae/falcon-7b-instruct on a single 24 GB GPU.
Requires: transformers>=4.38, peft>=0.10, bitsandbytes>=0.43, datasets, accelerate
"""

import json
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    Trainer,
    DataCollatorForSeq2Seq,
)
from peft import LoraConfig, get_peft_model, TaskType
from datasets import Dataset

# ─────────────────────────────────────────────
# CONFIG
# ─────────────────────────────────────────────
MODEL_NAME   = "tiiuae/falcon-7b-instruct"
DATA_PATH    = "train.jsonl"
OUTPUT_DIR   = "falcon-7b-lora"
SEQ_LENGTH   = 1024      # Falcon context is 2048; 1024 keeps activations in check
BATCH_SIZE   = 1         # per-device
GRAD_ACCUM   = 16        # effective batch = 16
NUM_EPOCHS   = 3
LR           = 2e-4
LORA_R       = 16
LORA_ALPHA   = 32
LORA_DROPOUT = 0.05

# ─────────────────────────────────────────────
# 1. TOKENIZER
# ─────────────────────────────────────────────
# Falcon quirk: no default pad_token. Set it explicitly.
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.pad_token_id = tokenizer.eos_token_id

print(f"Vocab size: {tokenizer.vocab_size}")
print(f"EOS token: {tokenizer.eos_token} (id={tokenizer.eos_token_id})")

# ─────────────────────────────────────────────
# 2. DATASET
# ─────────────────────────────────────────────
def load_and_format(path: str) -> Dataset:
    """Load JSONL and apply Falcon's chat template."""
    records = []
    with open(path) as f:
        for line in f:
            obj = json.loads(line)
            messages = [
                {"role": "user", "content": obj["instruction"]},
                {"role": "assistant", "content": obj["output"]},
            ]
            text = tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=False
            )
            records.append({"text": text})

    ds = Dataset.from_list(records)

    def tokenize(examples):
        return tokenizer(
            examples["text"],
            max_length=SEQ_LENGTH,
            truncation=True,
            padding="max_length",
        )

    return ds.map(tokenize, batched=True, remove_columns=["text"])

print("Loading dataset...")
dataset = load_and_format(DATA_PATH)
print(f"Dataset size: {len(dataset)}")

# ─────────────────────────────────────────────
# 3. MODEL (4-bit quantization)
# ─────────────────────────────────────────────
# Falcon quirk: native training precision is fp16 (NOT bf16 like LLaMA).
# Using bf16 compute dtype can cause numerical issues with Falcon's MQA.
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",              # NormalFloat4 — better than FP4 for 7B
    bnb_4bit_compute_dtype=torch.float16,   # Falcon quirk: use fp16, not bf16
    bnb_4bit_use_double_quant=True,         # NF4 + double quant saves ~0.4 GB
)

# Falcon quirk: attention implementation
#   "flash_attention_2" — most VRAM-efficient, requires flash-attn wheel
#   "sdpa"              — PyTorch native, good fallback
#   "eager"             — default, most VRAM-hungry
try:
    import flash_attn  # noqa: F401
    ATTN_IMPL = "flash_attention_2"
except ImportError:
    ATTN_IMPL = "sdpa"
    print("flash-attn not found; falling back to SDPA attention.")

print(f"Loading model with 4-bit NF4 + {ATTN_IMPL}...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    device_map="auto",
    attn_implementation=ATTN_IMPL,
)

# Falcon quirk: disable KV cache during training to avoid memory bloat
model.config.use_cache = False
# Required for gradient checkpointing + LoRA to work together
model.enable_input_require_grads()

# ─────────────────────────────────────────────
# 4. LoRA CONFIGURATION
# ─────────────────────────────────────────────
# Falcon architecture specifics:
#
#   FalconAttention:
#     query_key_value: Linear(hidden, 3*hidden)  ← FUSED QKV (single linear!)
#     dense:           Linear(hidden, hidden)    ← output projection
#
#   FalconMLP:
#     dense:  Linear(hidden, inter)             ← up-projection
#     output: Linear(inter, hidden)             ← down-projection
#
# Unlike LLaMA (separate q_proj/k_proj/v_proj), Falcon fuses Q, K, V into a
# single "query_key_value" linear. LoRA on this one module adapts all three
# projections simultaneously.
#
# "dense" appears in BOTH attention and MLP, so it targets both.
lora_config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=LORA_R,
    lora_alpha=LORA_ALPHA,
    lora_dropout=LORA_DROPOUT,
    target_modules=["query_key_value", "dense"],
    bias="none",
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
# Expected: "trainable params: ~20.9M || all params: 7,568,000,000 || 0.28%"

# ─────────────────────────────────────────────
# 5. TRAINING
# ─────────────────────────────────────────────
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=NUM_EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRAD_ACCUM,
    learning_rate=LR,
    fp16=True,
    gradient_checkpointing=True,
    logging_steps=10,
    save_strategy="epoch",
    save_total_limit=2,
    load_best_model_at_end=True,
    metric_for_best_model="loss",
    greater_is_better=False,
    report_to="none",
    optim="adamw_torch",
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    data_collator=DataCollatorForSeq2Seq(
        tokenizer=tokenizer,
        padding=True,
        max_length=SEQ_LENGTH,
    ),
)

trainer.train()

# ─────────────────────────────────────────────
# 6. SAVE ADAPTER
# ─────────────────────────────────────────────
adapter_path = f"{OUTPUT_DIR}/final_adapter"
model.save_pretrained(adapter_path)
tokenizer.save_pretrained(adapter_path)
print(f"Adapter saved to {adapter_path}")

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

base_path    = "tiiuae/falcon-7b-instruct"
adapter_path = "falcon-7b-lora/final_adapter"

# Load base model in 4-bit for inference
infer_model = AutoModelForCausalLM.from_pretrained(
    base_path,
    load_in_4bit=True,
    device_map="auto",
)
infer_model = PeftModel.from_pretrained(infer_model, adapter_path)

infer_tokenizer = AutoTokenizer.from_pretrained(base_path)
infer_tokenizer.pad_token = infer_tokenizer.eos_token

test_prompt = (
    "A chat between a curious user and an artificial intelligence assistant. "
    "The assistant gives helpful, detailed, and polite answers.\n"
    "### Instruction: Explain gradient checkpointing in one sentence.\n"
    "### Response:"
)

inputs = infer_tokenizer(test_prompt, return_tensors="pt").to("cuda")

with torch.inference_mode():
    outputs = infer_model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
        repetition_penalty=1.1,   # Falcon-7B has a known repetition bias
    )

print(infer_tokenizer.decode(outputs[0], skip_special_tokens=True))