From bfcda9ba30430995ed7b5233b51c9c524f5c7f33 Mon Sep 17 00:00:00 2001
From: marksverdhei <marksverdhei@hotmail.com>
Date: Sun, 29 Mar 2026 01:35:25 +0000
Subject: [PATCH] feat: add performance.yaml example and enable
 gradient_checkpointing in borealis_1b (issue #6 items 6-7)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Addresses issue #6 items 6 (gradient checkpointing) and 7 (8-bit Adam):

- New `examples/performance.yaml`: demonstrates all high-impact perf
  knobs in one place — gradient_checkpointing, adamw_bnb_8bit,
  flash_attention_2, larger batch with accumulation, and all-linear
  LoRA targets. Includes inline comments explaining each tradeoff.

- `examples/borealis_1b.yaml`: enables gradient_checkpointing (was
  false). For a 1B model on a 24GB GPU the VRAM savings allow batch_size
  to increase, more than recovering the ~20-30% wall-clock overhead.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 examples/borealis_1b.yaml |  2 +-
 examples/performance.yaml | 65 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)
 create mode 100644 examples/performance.yaml

diff --git a/examples/borealis_1b.yaml b/examples/borealis_1b.yaml
index 9e2dbec..48831fe 100644
--- a/examples/borealis_1b.yaml
+++ b/examples/borealis_1b.yaml
@@ -14,7 +14,7 @@ per_device_train_batch_size: 1
 max_grad_norm: 1.0
 logging_steps: 5
 save_strategy: "epoch"
-gradient_checkpointing: false
+gradient_checkpointing: true          # saves VRAM on 1B+ models
 optim: "adamw_torch"
 bf16: true
 report_to: "none"
diff --git a/examples/performance.yaml b/examples/performance.yaml
new file mode 100644
index 0000000..6ed36bb
--- /dev/null
+++ b/examples/performance.yaml
@@ -0,0 +1,65 @@
+# Performance-optimized prompt baking (RTX 3090 / 24GB VRAM)
+#
+# Demonstrates all high-impact performance knobs from issue #6:
+#   1. gradient_checkpointing — trades ~20-30% wall-clock for VRAM savings,
+#      enabling larger batch sizes that more than compensate on 1B+ models.
+#   2. adamw_bnb_8bit (bitsandbytes) — cuts optimizer state by ~75%;
+#      requires `pip install bitsandbytes` or `bakery[qlora]`.
+#   3. attn_implementation: "flash_attention_2" — 2-4x attention speedup on
+#      long sequences (falls back to "sdpa" if FA2 not installed).
+#   4. per_device_train_batch_size: 4 — larger batches amortize kernel overhead;
+#      only feasible because gradient_checkpointing frees VRAM.
+#   5. gradient_accumulation_steps: 2 — effective batch size 8 without OOM.
+#
+# Usage:
+#   pip install bitsandbytes           # for 8-bit Adam (optional)
+#   pip install flash-attn             # for Flash Attention 2 (optional)
+#   bakery --config examples/performance.yaml
+
+# --- Standard TrainingArguments ---
+output_dir: "./outputs/performance"
+num_train_epochs: 3
+learning_rate: 1e-4
+per_device_train_batch_size: 4
+gradient_accumulation_steps: 2
+max_grad_norm: 1.0
+logging_steps: 5
+save_strategy: "epoch"
+gradient_checkpointing: true          # saves VRAM → allows batch_size 4
+optim: "adamw_bnb_8bit"               # 8-bit Adam; requires: pip install bitsandbytes
+bf16: true
+report_to: "none"
+seed: 42
+lr_scheduler_type: "cosine"
+warmup_ratio: 0.1
+
+# --- Prompt baking ---
+system_prompt: "You are a helpful, harmless, and honest assistant."
+num_trajectories: 4
+trajectory_length: 128
+temperature: 1.0
+sampling_temperature: 0.8
+
+# --- Model ---
+model_name_or_path: "Qwen/Qwen3-1.7B"
+torch_dtype: "bfloat16"
+# Flash Attention 2 gives 2-4x speedup on long sequences.
+# Falls back gracefully to SDPA if flash-attn is not installed:
+#   attn_implementation: "sdpa"
+attn_implementation: "flash_attention_2"
+
+# --- LoRA ---
+r: 64
+lora_alpha: 128
+# "all-linear" targets every linear layer — maximises model capacity.
+# Slower than listing specific modules but often more expressive.
+target_modules: all-linear
+lora_dropout: 0.05
+
+# --- Data ---
+training_prompts:
+  - "What is the capital of France?"
+  - "Explain photosynthesis in simple terms."
+  - "Write a haiku about the ocean."
+  - "What are the benefits of exercise?"
+  - "How does a computer work?"