Supervised Fine-Tuning (SFT)

# Configure PyTorch training with gradient accumulation and mixed precision
# These settings control memory usage and effective batch size
import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from trl import SFTTrainer, SFTConfig

# 1. Load model and tokenizer
model_name = "meta-llama/Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Ensure pad token is set (required for batched training)
if tokenizer.pad_token is None:
 tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
 model_name,
 torch_dtype=torch.bfloat16,
 device_map="auto",
 attn_implementation="flash_attention_2", # Faster training
)

# 2. Load and prepare dataset (ChatML/messages format)
dataset = load_dataset("json", data_files={
 "train": "data/train.jsonl",
 "validation": "data/val.jsonl"
})

# 3. Configure SFT training
sft_config = SFTConfig(
 output_dir="./checkpoints/llama-sft",

 # Core hyperparameters
 num_train_epochs=3,
 per_device_train_batch_size=4,
 per_device_eval_batch_size=4,
 gradient_accumulation_steps=8, # Effective batch = 4 * 8 = 32
 learning_rate=2e-5,
 weight_decay=0.01,
 warmup_ratio=0.1, # 10% of steps for warmup
 lr_scheduler_type="cosine",

 # Sequence configuration
 max_seq_length=2048,
 packing=True, # Enable sequence packing

 # Precision and optimization
 bf16=True, # Use bfloat16 mixed precision
 gradient_checkpointing=True, # Trade compute for memory
 gradient_checkpointing_kwargs={"use_reentrant": False},

 # Logging and evaluation
 logging_steps=10,
 eval_strategy="steps",
 eval_steps=100,
 save_strategy="steps",
 save_steps=100,
 save_total_limit=3, # Keep only 3 best checkpoints
 load_best_model_at_end=True,
 metric_for_best_model="eval_loss",

 # Monitoring
 report_to="wandb", # or "tensorboard"
 run_name="llama-8b-sft-v1",

 # Reproducibility
 seed=42,
 data_seed=42,
)

# 4. Create trainer and start training
trainer = SFTTrainer(
 model=model,
 args=sft_config,
 train_dataset=dataset["train"],
 eval_dataset=dataset["validation"],
 processing_class=tokenizer,
)

# 5. Train
trainer.train()

# 6. Save the final model
trainer.save_model("./models/llama-sft-final")
tokenizer.save_pretrained("./models/llama-sft-final")

# Calculating effective batch size
def compute_effective_batch_size(
 per_device_batch_size: int,
 gradient_accumulation_steps: int,
 num_gpus: int = 1
) -> dict:
 """Calculate effective batch size and training throughput."""
 effective_batch = per_device_batch_size * gradient_accumulation_steps * num_gpus

 return {
 "per_device_batch_size": per_device_batch_size,
 "gradient_accumulation_steps": gradient_accumulation_steps,
 "num_gpus": num_gpus,
 "effective_batch_size": effective_batch,
 "optimizer_steps_per_epoch": "num_examples / effective_batch_size",
 }

# Common configurations
configs = [
 (2, 16, 1), # Single GPU, small memory
 (4, 8, 1), # Single GPU, moderate memory
 (4, 4, 4), # 4 GPUs, distributed
 (8, 2, 8), # 8 GPUs, large cluster
]

print(f"{'Per-device':>12} {'Grad Accum':>12} {'GPUs':>6} {'Effective BS':>14}")
print("-" * 50)
for pd_bs, ga, gpus in configs:
 result = compute_effective_batch_size(pd_bs, ga, gpus)
 print(f"{pd_bs:>12} {ga:>12} {gpus:>6} {result['effective_batch_size']:>14}")

# Visualizing different schedulers
from transformers import get_scheduler
import torch

def visualize_schedulers(
 total_steps: int = 1000,
 warmup_steps: int = 100,
 learning_rate: float = 2e-5
):
 """Compare learning rate schedules side by side."""
 schedules = {}

 for sched_type in ["cosine", "linear", "constant_with_warmup"]:
 # Create a dummy optimizer
 param = torch.nn.Parameter(torch.zeros(1))
 optimizer = torch.optim.AdamW([param], lr=learning_rate)

 scheduler = get_scheduler(
 name=sched_type,
 optimizer=optimizer,
 num_warmup_steps=warmup_steps,
 num_training_steps=total_steps,
 )

 lrs = []
 for step in range(total_steps):
 lrs.append(optimizer.param_groups[0]["lr"])
 optimizer.step()
 scheduler.step()

 schedules[sched_type] = lrs

 # Print sample points
 checkpoints = [0, 50, 100, 250, 500, 750, 999]
 print(f"{'Step':>6} {'Cosine':>12} {'Linear':>12} {'Constant':>12}")
 print("-" * 44)
 for step in checkpoints:
 print(f"{step:>6} {schedules['cosine'][step]:>12.2e} "
 f"{schedules['linear'][step]:>12.2e} "
 f"{schedules['constant_with_warmup'][step]:>12.2e}")

 return schedules

schedules = visualize_schedulers()

# Sanity check: verify training is working correctly
def run_sanity_check(trainer, tokenizer, dataset, num_samples=3):
 """Quick sanity check before committing to a full training run."""

 print("=" * 60)
 print("SANITY CHECK")
 print("=" * 60)

 # 1. Check a few tokenized examples
 print("\n1. Sample tokenized examples:")
 for i in range(min(num_samples, len(dataset))):
 example = dataset[i]
 messages = example["messages"]
 text = tokenizer.apply_chat_template(messages, tokenize=False)
 tokens = tokenizer(text)["input_ids"]
 print(f" Example {i}: {len(tokens)} tokens")
 # Decode and check it looks reasonable
 decoded = tokenizer.decode(tokens[:50])
 print(f" First 50 tokens: {decoded[:100]}...")

 # 2. Run a few training steps
 print("\n2. Running 10 training steps...")
 trainer.args.max_steps = 10
 trainer.args.logging_steps = 1
 result = trainer.train()

 # 3. Check loss trajectory
 logs = trainer.state.log_history
 losses = [l["loss"] for l in logs if "loss" in l]
 print(f" Loss trajectory: {[f'{l:.4f}' for l in losses]}")

 if len(losses) >= 2:
 if losses[-1] < losses[0]:
 print(" [PASS] Loss is decreasing")
 else:
 print(" [WARN] Loss is not decreasing; check learning rate")

 # 4. Generate a sample response
 print("\n3. Sample generation:")
 model = trainer.model
 model.eval()
 test_messages = [{"role": "user", "content": "Hello, how are you?"}]
 inputs = tokenizer.apply_chat_template(
 test_messages, return_tensors="pt", add_generation_prompt=True
 ).to(model.device)

 with torch.no_grad():
 output = model.generate(inputs, max_new_tokens=50, temperature=0.7)

 response = tokenizer.decode(output[0][inputs.shape[1]:], skip_special_tokens=True)
 print(f" Response: {response[:200]}")

 print("\n" + "=" * 60)
 print("Sanity check complete. Review above before full training.")
 print("=" * 60)

# pip install accelerate
from accelerate import Accelerator
from torch.utils.data import DataLoader
import torch

accelerator = Accelerator(mixed_precision="bf16")

# Standard PyTorch setup (model, optimizer, dataloader)
model = ... # your AutoModelForCausalLM
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
train_loader = DataLoader(dataset, batch_size=4, shuffle=True)

# Accelerate wraps everything for distributed + mixed precision
model, optimizer, train_loader = accelerator.prepare(
 model, optimizer, train_loader
)

# Training loop: identical to single-GPU PyTorch
for batch in train_loader:
 outputs = model(**batch)
 loss = outputs.loss
 accelerator.backward(loss) # replaces loss.backward()
 optimizer.step()
 optimizer.zero_grad()

# Launch with: accelerate launch --num_processes 4 train.py