Homework

Task 1: Rank Ablation Study

Goal: Develop intuition for how LoRA rank affects fine-tuning quality and find the sweet spot for a given dataset.

Instructions

Using the same QLoRA setup from the lab (TinyLlama + dolly-15k subset), train four separate adapters with different rank values:

Run	Rank (r)	Alpha	Target Modules
A	4	8	q_proj, v_proj
B	8	16	q_proj, v_proj
C	16	32	q_proj, v_proj
D	32	64	q_proj, v_proj

Keep all other hyperparameters identical (learning rate, batch size, epochs, dataset). For each run, record:

1. Final training loss
2. Trainable parameter count and percentage
3. Peak GPU memory usage (from torch.cuda.max_memory_allocated())
4. Total training time

After training all four, run the same 5 test prompts through each adapter and qualitatively rate the outputs.

Deliverable

A comparison table with all four runs and a 2–3 sentence conclusion about which rank provided the best trade-off for this dataset and task.

Self-check

• ✓ You trained all four rank variants with identical hyperparameters (except r and alpha).
• ✓ Your table includes loss, params, memory, and time for each run.
• ✓ You tested inference on the same prompts across all variants.
• ✓ Your conclusion explains why a particular rank was the sweet spot (not just which one had the lowest loss).

Task 2: Memory Profiling — LoRA vs QLoRA

Goal: Quantify the exact memory savings of QLoRA compared to standard LoRA, and understand where the savings come from.

Instructions

Write a script that loads the same model in three configurations and reports memory usage for each:

1. fp16 + LoRA: Load model in fp16, apply LoRA (r=16, q_proj + v_proj)
2. 4-bit + LoRA (QLoRA): Load model in 4-bit NF4, apply same LoRA config
3. 4-bit + LoRA + double quant: Same as (2) but with bnb_4bit_use_double_quant=True

For each configuration, record:

• torch.cuda.memory_allocated() — memory used by tensors
• torch.cuda.memory_reserved() — memory reserved by the caching allocator
• Model size on disk (save and check file sizes)

Deliverable

A markdown table or printed comparison showing memory usage across all three configurations, plus a short paragraph explaining where the memory savings come from in each step.

Self-check

• ✓ You measured all three configurations on the same model.
• ✓ Each measurement was taken in a clean process (no leftover allocations).
• ✓ You can explain why 4-bit saves ~4x memory over fp16 for the base model weights.
• ✓ You can explain what double quantization saves and roughly how much.

Task 3: Target Module Analysis

Goal: Understand which layers in a transformer are most impactful for LoRA adaptation.

Instructions

This task combines code exploration with written analysis. No GPU training is required — you will inspect the model architecture and reason about layer selection.

1. Load any Llama-family model (even in CPU mode) and print all named modules:

   for name, module in model.named_modules():
       if isinstance(module, torch.nn.Linear):
           print(f"{name}: {module.in_features} x {module.out_features}")

2. Categorize all linear layers into groups: attention projections (q, k, v, o), MLP layers (gate, up, down), and other (lm_head, embeddings).
3. For each group, calculate:
   • Total parameters in that group
   • Percentage of total model parameters
   • LoRA trainable params if you applied r=16 to that group
4. Write a ~200-word analysis explaining:
   • Why attention projections are the default LoRA target
   • When adding MLP layers would be worthwhile
   • Why lm_head is typically not adapted with LoRA

Deliverable

A table of layer groups with parameter counts, plus a 200-word written analysis.

Self-check

• ✓ You identified all linear layers in the model and grouped them correctly.
• ✓ Your parameter counts are computed from the actual model, not estimated.
• ✓ Your analysis addresses all three questions (attention default, MLP trade-off, lm_head exclusion).
• ✓ You can explain the trade-off between adding more target modules and overfitting risk on small datasets.