Stage 1: Pre-training          → "Learn language and world knowledge"
Stage 2: Supervised Fine-tuning → "Learn to follow instructions"
Stage 3: Alignment (RLHF/DPO)  → "Learn human preferences"

Input:  "The capital of France is"
Target: "Paris"

Loss = -log P("Paris" | "The capital of France is")

Loss = -Σᵢ log P(tᵢ | t₁, t₂, ..., tᵢ₋₁)

Input:  "The [MASK] sat on the [MASK]"
Target: "cat", "mat"

Loss = -log P("cat" | context) - log P("mat" | context)

# Llama 2 style configuration
config = {
    "optimizer": "AdamW",
    "learning_rate": 3e-4,          # Peak LR
    "lr_schedule": "cosine_decay",
    "warmup_steps": 2000,
    "weight_decay": 0.1,
    "batch_size": 4_000_000,        # Tokens per batch (4M)
    "max_steps": 500_000,
    "gradient_clipping": 1.0,
    "beta1": 0.9,
    "beta2": 0.95,
    "epsilon": 1e-5,
}

### Instruction:
Explain quantum entanglement in simple terms.

### Response:
Quantum entanglement is when two particles become connected in such a way that
measuring one instantly affects the other, no matter how far apart they are...

# Typical SFT training example
{
    "messages": [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What is photosynthesis?"},
        {"role": "assistant", "content": "Photosynthesis is the process by which plants convert sunlight, water, and CO₂ into glucose and oxygen..."}
    ]
}

# SFT training loop (simplified)
from transformers import AutoModelForCausalLM, Trainer, TrainingArguments

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b")

training_args = TrainingArguments(
    output_dir="./sft-model",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    learning_rate=2e-5,          # Much lower than pre-training
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    bf16=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=sft_dataset,
)
trainer.train()

Prompt: "Explain gravity to a 5-year-old"

Response A: "Gravity is the force of attraction between masses proportional to..."
Response B: "You know how when you throw a ball up, it always comes back down?..."

Human preference: B > A (simpler, age-appropriate)

# Reward model: same architecture as LLM but outputs a scalar
class RewardModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.backbone = base_model
        self.reward_head = nn.Linear(hidden_size, 1)

    def forward(self, input_ids):
        hidden = self.backbone(input_ids).last_hidden_state[:, -1, :]
        return self.reward_head(hidden)  # Single scalar reward

# Training loss: Bradley-Terry model
# P(A > B) = sigmoid(reward(A) - reward(B))
loss = -log(sigmoid(reward_chosen - reward_rejected))

For each prompt:
1. Generate response with current policy
2. Score with reward model
3. Update policy to increase reward
4. Add KL penalty to prevent diverging too far from SFT model

Objective = E[reward(response)] - β · KL(policy || reference_policy)

RLHF: Preferences → Reward Model → RL (PPO) → Better Policy
DPO:  Preferences → Direct Policy Optimization (single step)

# DPO loss function
def dpo_loss(policy_chosen_logps, policy_rejected_logps,
             reference_chosen_logps, reference_rejected_logps, beta=0.1):
    """
    policy_chosen_logps: log P_policy(chosen response)
    reference_chosen_logps: log P_reference(chosen response)
    """
    chosen_rewards = beta * (policy_chosen_logps - reference_chosen_logps)
    rejected_rewards = beta * (policy_rejected_logps - reference_rejected_logps)

    loss = -torch.log(torch.sigmoid(chosen_rewards - rejected_rewards))
    return loss.mean()

Original: y = Wx
LoRA:     y = Wx + BAx

Where:
- W is frozen (original weights, not updated)
- B is (d × r) — small matrix
- A is (r × d) — small matrix
- r << d (rank, typically 8-64)

Full fine-tuning:  Update all of W (d × d parameters)
LoRA:              Only train A and B (2 × d × r parameters)

Example: d=4096, r=16
Full: 16,777,216 parameters per layer
LoRA: 131,072 parameters per layer (0.78%)

# LoRA implementation concept
class LoRALayer(nn.Module):
    def __init__(self, original_layer, rank=16, alpha=32):
        super().__init__()
        self.original = original_layer
        self.original.weight.requires_grad = False  # Freeze

        d_in = original_layer.in_features
        d_out = original_layer.out_features

        self.lora_A = nn.Linear(d_in, rank, bias=False)
        self.lora_B = nn.Linear(rank, d_out, bias=False)
        self.scaling = alpha / rank

        # Initialize A with random, B with zeros (start at original behavior)
        nn.init.kaiming_uniform_(self.lora_A.weight)
        nn.init.zeros_(self.lora_B.weight)

    def forward(self, x):
        original_output = self.original(x)
        lora_output = self.lora_B(self.lora_A(x)) * self.scaling
        return original_output + lora_output

from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(
    r=16,                          # Rank
    lora_alpha=32,                 # Scaling factor
    target_modules=["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(base_model, lora_config)
model.print_trainable_parameters()
# trainable params: 13,631,488 || all params: 6,751,318,016 || trainable%: 0.20%

Standard LoRA: Base model in FP16 (140GB for 70B) + LoRA adapters
QLoRA:         Base model in 4-bit (35GB for 70B) + LoRA adapters in FP16

from transformers import BitsAndBytesConfig

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

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-70b",
    quantization_config=bnb_config,
    device_map="auto",
)
# Now apply LoRA on top of the quantized model

Frozen Layer → [Adapter: down-project → activation → up-project] → Frozen Layer

# Chat format (most common)
def format_chat(example):
    messages = [
        {"role": "system", "content": "You are a helpful coding assistant."},
        {"role": "user", "content": example["question"]},
        {"role": "assistant", "content": example["answer"]},
    ]
    return tokenizer.apply_chat_template(messages, tokenize=False)

# Alpaca format (legacy but common in datasets)
TEMPLATE = """### Instruction:
{instruction}

### Input:
{input}

### Response:
{output}"""

GPU 0: Model copy → Batch 0 → Gradients₀ ─┐
GPU 1: Model copy → Batch 1 → Gradients₁ ─┼→ Average → Update all copies
GPU 2: Model copy → Batch 2 → Gradients₂ ─┘

GPU 0: First half of attention heads
GPU 1: Second half of attention heads
→ All-reduce to combine

GPU 0: Layers 1-8
GPU 1: Layers 9-16
GPU 2: Layers 17-24
GPU 3: Layers 25-32

# PyTorch FSDP
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.bfloat16,
    ),
)

Do you need the model to know NEW FACTS?
├── Yes → Is the knowledge static or changing?
│   ├── Static (rarely updates) → Fine-tune on domain data
│   └── Dynamic (updates frequently) → RAG
│
└── No → Do you need a new BEHAVIOR or FORMAT?
    ├── Yes → Can you describe it in a prompt?
    │   ├── Yes, and it works → Prompt engineering (done!)
    │   └── No, too complex → Fine-tune (SFT)
    │
    └── No → Do you need better QUALITY on specific tasks?
        ├── Marginal improvement needed → Better prompts, few-shot examples
        └── Significant improvement needed → Fine-tune

# Step 1: Load base model with quantization
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer, SFTConfig

model_name = "meta-llama/Llama-3.1-8B"

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

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    attn_implementation="flash_attention_2",
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Step 2: Prepare for training
model = prepare_model_for_kbit_training(model)

lora_config = LoraConfig(
    r=32,
    lora_alpha=64,
    target_modules=["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)

# Step 3: Configure training
training_config = SFTConfig(
    output_dir="./output",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    learning_rate=2e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    max_seq_length=2048,
    bf16=True,
    logging_steps=10,
    save_strategy="epoch",
    evaluation_strategy="epoch",
)

# Step 4: Train
trainer = SFTTrainer(
    model=model,
    args=training_config,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer,
)
trainer.train()

# Step 5: Save adapter
model.save_pretrained("./final-adapter")

# Step 6: Merge adapter into base model for deployment (optional)
from peft import PeftModel
base = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)
merged = PeftModel.from_pretrained(base, "./final-adapter")
merged = merged.merge_and_unload()
merged.save_pretrained("./merged-model")