|
| 1 | +# Part 1 — Text-to-SQL (full / LoRA / QLoRA) |
| 2 | + |
| 3 | +Fine-tune a small causal LM on text-to-SQL with **full fine-tuning**, **LoRA**, or **QLoRA**. Adapted from Union AI's [llm-fine-tuning-lora-qlora](https://github.com/unionai/workshops/tree/main/tutorials/llm-fine-tuning-lora-qlora) workshop. Default model: [SmolLM2-135M](https://huggingface.co/HuggingFaceTB/SmolLM2-135M). |
| 4 | + |
| 5 | +**Why Flyte 2:** a one-off `Trainer` script is enough for a single laptop run. This tutorial is two machines and more than one kind of work. Follow along on this **MacBook Pro** (M4 Max, 128 GB, 16-core 12P+4E, MPS) with a smoke test. Union-scale samples and QLoRA run on a remote **NVIDIA RTX PRO 6000 Blackwell (96 GB)**. Part 3 swaps in an 8B model on the same DAG. Without Flyte that is three scripts. With it, one `workflow.py`: CPU `prepare_data`, GPU `train` / `evaluate`, optional serve. `flyte run --local` here; `flyte run` on the GPU. Change `--method` or `--model_name`, not the pipeline. |
| 6 | + |
| 7 | +## SDK note (already Flyte 2) |
| 8 | + |
| 9 | +This workshop does **not** need a Flyte 1 → 2 migration. The upstream code already uses the Flyte 2 SDK: |
| 10 | + |
| 11 | +| Flyte 1 | This tutorial (Flyte 2) | |
| 12 | +|---------|-------------------------| |
| 13 | +| `import flytekit` | `import flyte` | |
| 14 | +| `@task` / `@workflow` | `TaskEnvironment` + `@cpu_env.task` / `@gpu_env.task` | |
| 15 | +| `FlyteDirectory` | `flyte.io.Dir` | |
| 16 | +| `pyflyte run` | `flyte run` | |
| 17 | +| Decks | `report=True` | |
| 18 | + |
| 19 | +`requirements.txt` pins `flyte[tui]>=2.0`. See Union's [Flyte 1 → 2 guide](https://www.union.ai/docs/v2/flyte/user-guide/migration/flyte-2/) if you later port other Flyte 1 code. |
| 20 | + |
| 21 | +## What's here |
| 22 | + |
| 23 | +| File | Role | |
| 24 | +|------|------| |
| 25 | +| `device_utils.py` / `run_profile.py` | CUDA vs MPS vs CPU; laptop vs GPU run sizes | |
| 26 | +| `config.py` | Flyte environments — CPU for data prep, GPU for training (`gpu=0` on a Mac so `--local` works) | |
| 27 | +| `workflow.py` | prepare data → train → evaluate base vs fine-tuned | |
| 28 | +| `report_helpers.py` | HTML/SVG reports for the Flyte UI | |
| 29 | +| `serve.py` | FastAPI serving (optional, needs a Flyte cluster) | |
| 30 | +| `app_gradio.py` | Gradio UI in front of `serve.py` (optional) | |
| 31 | +| `llm-fine-tune-tutorial.ipynb` | Notebook: pipeline diagram first, then train; placeholders, not a personal GPU endpoint | |
| 32 | +| `images/flyte_pipeline.png` | High-level DAG: two machines, CPU vs GPU tasks | |
| 33 | + |
| 34 | +## Setup |
| 35 | + |
| 36 | +```bash |
| 37 | +cd NLP/Fine_Tuning/01_text_to_sql |
| 38 | + |
| 39 | +uv venv .venv --python 3.11 |
| 40 | +source .venv/bin/activate |
| 41 | +uv pip install -r requirements.txt |
| 42 | +``` |
| 43 | + |
| 44 | +Optional Hugging Face token for gated models (never commit `.env`): |
| 45 | + |
| 46 | +```bash |
| 47 | +cp .env.example .env |
| 48 | +# then edit .env |
| 49 | +``` |
| 50 | + |
| 51 | +Remote Flyte/Union config is **not** committed. If you use a cluster, generate it locally: |
| 52 | + |
| 53 | +```bash |
| 54 | +flyte create config \ |
| 55 | + --endpoint YOUR_CLUSTER.hosted.unionai.cloud \ |
| 56 | + --project YOUR_PROJECT \ |
| 57 | + --domain development \ |
| 58 | + --builder remote |
| 59 | +``` |
| 60 | + |
| 61 | +That writes `.flyte/config.yaml`, which is gitignored. A dummy layout is in `config.yaml.example`. |
| 62 | + |
| 63 | +## How LoRA and QLoRA differ from full fine-tuning |
| 64 | + |
| 65 | +**Full fine-tuning** updates every weight. Effective, but expensive to train, store, and deploy. |
| 66 | + |
| 67 | +**LoRA** freezes the base weights `W` and trains a low-rank update `A × B`, scaled by `alpha / r`: |
| 68 | + |
| 69 | +``` |
| 70 | + ┌─────────────────────────┐ |
| 71 | + │ Original Weight W │ |
| 72 | +input ────────────→ │ (frozen) │──→ main output |
| 73 | + │ └─────────────────────────┘ │ |
| 74 | + │ ┌───────────┐ ┌───────────┐ │ |
| 75 | + └─────────────→ │ A (d × r) │→│ B (r × d) │→ × α/r ──→ + ──→ combined |
| 76 | + └───────────┘ └───────────┘ |
| 77 | + (trainable adapters) |
| 78 | +``` |
| 79 | + |
| 80 | +- `r` — adapter rank (capacity vs size) |
| 81 | +- `alpha` — scale; common default is `alpha = 2 * r`. Adds no extra parameters. |
| 82 | + |
| 83 | +**QLoRA** keeps those adapters in higher precision but stores the frozen base in 4-bit (NF4). That 4-bit path is [bitsandbytes](https://github.com/bitsandbytes-foundation/bitsandbytes), which **requires NVIDIA CUDA** — not Apple MPS or CPU. `requirements.txt` therefore installs `bitsandbytes` only when `sys_platform != "darwin"`. On a Mac, use `lora` or `full`; `method=qlora` raises a clear error. Run QLoRA on the remote **NVIDIA RTX PRO 6000 Blackwell (96 GB)** (same pipeline). |
| 84 | + |
| 85 | +On SmolLM2-135M, QLoRA is overkill (quality often drops) and is here so the same `method` flag works when you swap in a larger model. |
| 86 | + |
| 87 | +LoRA targets in `workflow.py` are LLaMA-style names (`q_proj`, `k_proj`, `v_proj`, `o_proj`, `gate_proj`, `up_proj`, `down_proj`). RoBERTa uses different names — that is handled in part 2. |
| 88 | + |
| 89 | +## Run |
| 90 | + |
| 91 | +Start small and **local**. The notebook (`llm-fine-tune-tutorial.ipynb`) picks MPS vs CUDA and a smoke-test size automatically. |
| 92 | + |
| 93 | +```bash |
| 94 | +# Mac laptop (MPS) or any machine without NVIDIA — LoRA smoke test |
| 95 | +flyte run --local workflow.py pipeline \ |
| 96 | + --method lora --epochs 1 --batch_size 2 \ |
| 97 | + --max_train_samples 100 --max_eval_samples 20 --num_eval_examples 10 |
| 98 | + |
| 99 | +# Remote NVIDIA RTX PRO 6000 Blackwell (96 GB) — Union-scale LoRA |
| 100 | +flyte run --local workflow.py pipeline --method lora --epochs 3 |
| 101 | + |
| 102 | +# QLoRA — NVIDIA CUDA only (bitsandbytes 4-bit; skipped on Mac) |
| 103 | +flyte run --local workflow.py pipeline --method qlora |
| 104 | +``` |
| 105 | + |
| 106 | +On a Mac, `recommend_jobs()` omits QLoRA. Treat CUDA as the QLoRA path. |
| 107 | + |
| 108 | +On a Flyte 2 cluster you already configured (config file stays local): |
| 109 | + |
| 110 | +```bash |
| 111 | +FLYTE_GPUS=1 flyte run workflow.py pipeline --method lora --epochs 3 |
| 112 | +FLYTE_GPUS=4 flyte run workflow.py pipeline --method lora --epochs 3 |
| 113 | +``` |
| 114 | + |
| 115 | +`FLYTE_GPUS` is how many NVIDIA GPUs the train task requests (default 1). Same `workflow.py` for one GPU or four on one node. Keep `FLYTE_GPUS=1` unless a single task should own the node. |
| 116 | + |
| 117 | +### Useful flags |
| 118 | + |
| 119 | +| Flag | Default | Meaning | |
| 120 | +|------|---------|---------| |
| 121 | +| `--model_name` | `HuggingFaceTB/SmolLM2-135M` | Hugging Face model | |
| 122 | +| `--dataset_name` | `b-mc2/sql-create-context` | Hugging Face dataset | |
| 123 | +| `--method` | `lora` | `full`, `lora`, or `qlora` | |
| 124 | +| `--epochs` | `3` | Training epochs | |
| 125 | +| `--lr` | `2e-4` | Learning rate | |
| 126 | +| `--batch_size` | `4` | Per-device batch size | |
| 127 | +| `--max_train_samples` | `5000` | Cap on train size | |
| 128 | +| `--max_eval_samples` | `500` | Cap on eval size | |
| 129 | +| `--num_eval_examples` | `50` | Examples in the before/after report | |
| 130 | +| `--lora_r` | `16` | LoRA rank | |
| 131 | +| `--lora_alpha` | `32` | LoRA alpha | |
| 132 | + |
| 133 | +## Evaluation |
| 134 | + |
| 135 | +The evaluate task runs the same prompts on the base model and the fine-tuned model: |
| 136 | + |
| 137 | +- Exact-match accuracy on normalized SQL |
| 138 | +- Side-by-side raw outputs (base models often ramble; fine-tuned models usually stop after the query) |
| 139 | + |
| 140 | +**Results live in the [notebook](llm-fine-tune-tutorial.ipynb) (section 3).** That table is full vs LoRA vs QLoRA (QLoRA only when CUDA/bitsandbytes is available). Re-run the results cell anytime — it reads `outputs/run_<method>.json`. |
| 141 | + |
| 142 | +```bash |
| 143 | +python run_grid.py # smoke: full + LoRA (QLoRA if NVIDIA) |
| 144 | +python run_grid.py --no-smoke # Union-scale sample counts (RTX PRO 6000 Blackwell) |
| 145 | +``` |
| 146 | + |
| 147 | +## Serve (optional) |
| 148 | + |
| 149 | +Only after a successful training run, and only if you have a Flyte cluster: |
| 150 | + |
| 151 | +```bash |
| 152 | +python serve.py |
| 153 | +# python serve.py --run-name YOUR_RUN_NAME |
| 154 | +``` |
| 155 | + |
| 156 | +Point curl at **your** deployed URL, not a sample host: |
| 157 | + |
| 158 | +```bash |
| 159 | +curl -X POST https://YOUR_APP_URL/generate \ |
| 160 | + -H "Content-Type: application/json" \ |
| 161 | + -d '{ |
| 162 | + "schema": "CREATE TABLE employees (id INT, name VARCHAR, department VARCHAR, salary INT)", |
| 163 | + "question": "What is the average salary by department?" |
| 164 | + }' |
| 165 | +``` |
| 166 | + |
| 167 | +Gradio UI: |
| 168 | + |
| 169 | +```bash |
| 170 | +SERVER_URL=https://YOUR_APP_URL python app_gradio.py |
| 171 | +``` |
| 172 | + |
| 173 | +## Next |
| 174 | + |
| 175 | +[Part 2 — RoBERTa on FOMC hawkish–dovish](../02_fomc_roberta/) uses the same three methods on a classification task from [Trillion Dollar Words](https://arxiv.org/abs/2305.07972). |
0 commit comments