Performance: training pipeline optimizations

[marksverdhei]

Performance Optimization Opportunities

Research and codebase analysis of training pipeline efficiency. Targeting single-GPU setups (RTX 3090, 24GB VRAM).

---

High Impact

1. Unsloth integration for model loading and LoRA

Unsloth provides custom CUDA kernels for LoRA training that report 2x speedup and 60-70% VRAM reduction across Llama, Qwen, Mistral, and Gemma models. It works with standard HF Trainer (not just SFTTrainer), which is what bakery uses.

Integration path:

• Replace AutoModelForCausalLM.from_pretrained + get_peft_model with FastLanguageModel.from_pretrained + FastLanguageModel.get_peft_model
• The custom compute_loss() and adapter toggling in PromptBakingTrainer should still work since unsloth patches at the model level, not the trainer level
• Need to verify disable_adapter_layers() / enable_adapter_layers() compatibility

Concern: Unsloth modifies internal model structures. The teacher/student adapter toggling pattern is non-standard and needs testing.

2. Enable Flash Attention / SDPA (cli.py:135)

Model loading doesn't request Flash Attention. Adding attn_implementation="flash_attention_2" (or "sdpa" as fallback) to load_kwargs gives 2-4x speedup on attention with no code changes. Especially impactful for long sequences (system prompt + conversation).

load_kwargs["attn_implementation"] = "flash_attention_2"  # or "sdpa"

Should be configurable via DataConfig.

3. Vectorize per-sample logit slicing loop (trainer.py:205-228)

The loss computation loops over each sample individually in Python:

for i in range(len(pairs)):
    t_logits = teacher_outputs.logits[i:i+1, t_start-1:-1, :]
    ...
    loss = compute_kl_divergence(...)
    losses.append(loss)

This could be vectorized with pre-computed offset tensors and a single batched compute_kl_divergence call, eliminating N Python loop iterations and N separate KL computations.

4. Batch trajectory generation (trainer.py:259-264)

On-the-fly trajectory generation runs sequentially — one prompt at a time, num_trajectories times each. Batching all prompts into a single model.generate() call would reduce GPU-CPU round trips from B * T to 1.

---

Medium Impact

5. Pre-compute prompt lengths outside training loop (trainer.py:157-172)

Each compute_loss() call tokenizes prompts individually just to measure their token length. These lengths are deterministic per (system_prompt, user_message) pair and could be cached in the dataset or collator.

6. Gradient checkpointing defaults

Currently gradient_checkpointing: false in examples. For the 1B+ model regime, enabling it trades ~20-30% wall-clock time for significant VRAM savings, allowing larger batch sizes that more than compensate.

7. 8-bit optimizer (adamw_torch -> adamw_bnb_8bit)

AdamW stores 2 state tensors per parameter. 8-bit Adam (via bitsandbytes) cuts optimizer memory by ~75%. With LoRA's small parameter count this matters less, but it's free performance for larger adapter ranks.

8. Move tokenization to data collator (data.py:28-43)

Currently the collator just passes raw strings through, and all tokenization happens inside compute_loss(). Moving formatting + tokenization to the collator enables DataLoader workers to overlap data prep with GPU compute.

---

Lower Impact / Experimental

9. torch.compile

Has known issues with PEFT/LoRA — module names gain _orig_mod. prefix breaking target matching. The dynamic adapter toggling pattern (enable/disable per step) likely causes graph breaks. Not recommended until PEFT compatibility improves.

10. DeepSpeed ZeRO on single GPU

ZeRO-Offload can offload optimizer states to CPU, but with LoRA the optimizer memory is already small. Overhead of CPU offloading likely outweighs benefit for 1B models. More relevant for 7B+ without quantization.

---

Suggested priority

1. Flash Attention (trivial, high payoff)
2. Unsloth model loading (moderate effort, big speedup if compatible)
3. Vectorize loss loop (moderate effort, removes Python overhead)
4. Gradient checkpointing + 8-bit optimizer in examples (config-only)
5. Batch trajectory generation (moderate effort)
6. Pre-compute prompt lengths / move tokenization to collator (refactor)