# LLM fine-tuning with LoRA and QLoRA
> [!NOTE]
> Code available [here](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/llm_fine_tuning_lora_qlora).
This tutorial fine-tunes a language model for SQL generation using three methods in one workflow: **full** fine-tuning, **LoRA** adapters, and **QLoRA** (4-bit quantized base + LoRA). The pipeline prepares an instruction dataset from HuggingFace, trains with [TRL](https://huggingface.co/docs/trl) `SFTTrainer`, evaluates against a base-model baseline, and streams training charts into Flyte reports.
Flyte provides:
- **GPU training** with live loss and learning-rate charts via `report=True`.
- **Method switching** through a single `method` parameter (`full`, `lora`, or `qlora`).
- **Cached dataset preparation** for fast iteration on hyperparameters.
## Define the task environments
The GPU environment declares a HuggingFace token secret for gated models.
```
# /// script
# requires-python = ">=3.12"
# dependencies = [
# "flyte>=2.4.0",
# "torch>=2.1.0",
# "transformers>=4.45.0",
# "peft>=0.13.0",
# "trl>=0.12.0",
# "datasets>=3.0.0",
# "bitsandbytes>=0.44.0",
# "accelerate>=0.34.0",
# ]
# main = "pipeline"
# params = ""
# ///
import asyncio
import json
import logging
import os
import tempfile
import flyte
import flyte.io
import flyte.report
# {{docs-fragment env}}
import os
main_img = flyte.Image.from_uv_script(__file__, name="llm-fine-tuning-lora-qlora", pre=True)
gpu_env = flyte.TaskEnvironment(
name="llm-fine-tuning-lora-qlora-gpu",
image=main_img,
resources=flyte.Resources(cpu=4, memory="24Gi", gpu=1),
secrets=[flyte.Secret(key="huggingface-token", as_env_var="HF_TOKEN")],
)
cpu_env = flyte.TaskEnvironment(
name="llm-fine-tuning-lora-qlora-cpu",
image=main_img,
resources=flyte.Resources(cpu=2, memory="8Gi"),
depends_on=[gpu_env],
)
HF_TOKEN = os.environ.get("HF_TOKEN")
# {{/docs-fragment env}}
from report_helpers import make_bar_chart, make_line_chart, pipeline_step_indicator, wrap_report
logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)
# ------------------------------------------------------------------
# Task 1: Prepare dataset
# ------------------------------------------------------------------
@cpu_env.task(cache="auto")
async def prepare_data(
dataset_name: str = "b-mc2/sql-create-context",
max_train_samples: int = 5000,
max_eval_samples: int = 500,
) -> flyte.io.Dir:
"""Download dataset from HuggingFace and format for instruction fine-tuning."""
from datasets import DatasetDict, load_dataset
log.info(f"Loading dataset: {dataset_name}")
ds = load_dataset(dataset_name, split="train")
def format_example(ex):
return {
"text": (
"### Task: Generate a SQL query to answer the question.\n"
f"### Schema:\n{ex['context']}\n"
f"### Question:\n{ex['question']}\n"
f"### SQL:\n{ex['answer']}\n<|endoftext|>"
)
}
ds = ds.map(format_example)
# Split into train and eval
total = len(ds)
train_end = min(max_train_samples, total - max_eval_samples)
eval_start = train_end
eval_end = min(eval_start + max_eval_samples, total)
processed = DatasetDict({
"train": ds.select(range(train_end)),
"eval": ds.select(range(eval_start, eval_end)),
})
output_dir = os.path.join(tempfile.mkdtemp(), "dataset")
processed.save_to_disk(output_dir)
log.info(f"Dataset ready: {len(processed['train'])} train, {len(processed['eval'])} eval")
return await flyte.io.Dir.from_local(output_dir)
# ------------------------------------------------------------------
# Task 2: Train
# ------------------------------------------------------------------
@gpu_env.task(report=True)
async def train(
model_name: str,
data_dir: flyte.io.Dir,
method: str = "lora",
epochs: int = 3,
lr: float = 2e-4,
batch_size: int = 4,
lora_r: int = 16,
lora_alpha: int = 32,
) -> flyte.io.Dir:
"""Fine-tune a model using full, LoRA, or QLoRA method."""
import torch
from datasets import load_from_disk
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainerCallback
from trl import SFTConfig, SFTTrainer
log.info(f"Training: model={model_name}, method={method}")
# -- Load data --
data_path = await data_dir.download()
dataset = load_from_disk(data_path)
# -- Load tokenizer --
token_kwargs = {"token": HF_TOKEN} if HF_TOKEN else {}
tokenizer = AutoTokenizer.from_pretrained(model_name, **token_kwargs)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# -- Initial report: loading model --
await flyte.report.replace.aio(
wrap_report(
f"Loading Model...
"
f"{model_name}
"
f''
f"
Method: {method.upper()}
"
f"
Dataset: {len(dataset['train']):,} train / {len(dataset['eval']):,} eval
"
f"
Training in Progress...
{model_name}
{len(dataset['train']):,}
Train Examples
{trainable_params / total_params * 100:.1f}%
Trainable
Method: {method_badge} | Total params: {total_params:,} | Trainable: {trainable_params:,}
"""
charts_html = ""
if training_log:
current = training_log[-1]
progress_pct = current["step"] / max_steps * 100 if max_steps else 0
charts_html += f"""
Step {current['step']}/{max_steps}
({progress_pct:.0f}%) |
Epoch {current['epoch']:.2f}/{epochs} |
Loss:
{current['loss']:.4f}
"""
loss_chart = make_line_chart(
data=training_log,
x_key="epoch",
y_keys=["loss"],
title="Training Loss",
x_label="Epoch",
y_label="Loss",
colors=["#5a7db5"],
)
charts_html += f'{loss_chart}
'
if "lr" in training_log[0]:
lr_chart = make_line_chart(
data=training_log,
x_key="epoch",
y_keys=["lr"],
title="Learning Rate Schedule",
x_label="Epoch",
y_label="LR",
colors=["#0f3460"],
)
charts_html += f'{lr_chart}
'
if "grad_norm" in training_log[0]:
grad_chart = make_line_chart(
data=training_log,
x_key="epoch",
y_keys=["grad_norm"],
title="Gradient Norm",
x_label="Epoch",
y_label="Grad Norm",
colors=["#06d6a0"],
)
charts_html += f'{grad_chart}
'
return wrap_report(stats_html + charts_html)
# -- Metrics callback with live report updates --
class MetricsCallback(TrainerCallback):
def on_log(self, args, state, control, logs=None, **kwargs):
if not logs or "loss" not in logs:
return
entry = {
"step": state.global_step,
"epoch": round(logs.get("epoch", 0), 2),
"loss": round(logs["loss"], 4),
}
if "learning_rate" in logs:
entry["lr"] = logs["learning_rate"]
if "grad_norm" in logs:
entry["grad_norm"] = round(float(logs["grad_norm"]), 4)
training_log.append(entry)
log.info(
f"step={state.global_step}/{state.max_steps} "
f"epoch={entry['epoch']:.2f} "
f"loss={entry['loss']:.4f}"
)
asyncio.run_coroutine_threadsafe(
flyte.report.replace.aio(
_build_training_report(state.max_steps),
do_flush=True,
),
loop,
)
# -- Train --
output_dir = os.path.join(tempfile.mkdtemp(), "checkpoints")
training_args = SFTConfig(
output_dir=output_dir,
num_train_epochs=epochs,
per_device_train_batch_size=batch_size,
learning_rate=lr,
logging_steps=10,
save_strategy="epoch",
bf16=use_bf16,
fp16=not use_bf16 and torch.cuda.is_available(),
gradient_accumulation_steps=4,
warmup_steps=10,
report_to="none",
)
trainer = SFTTrainer(
model=model,
args=training_args,
train_dataset=dataset["train"],
eval_dataset=dataset["eval"],
processing_class=tokenizer,
callbacks=[MetricsCallback()],
)
log.info("Starting training...")
await asyncio.to_thread(trainer.train)
log.info("Training complete.")
# -- Merge LoRA weights and save --
save_dir = os.path.join(tempfile.mkdtemp(), "finetuned_model")
if method in ("lora", "qlora"):
log.info("Merging LoRA weights into base model...")
model = model.merge_and_unload()
model.save_pretrained(save_dir)
tokenizer.save_pretrained(save_dir)
log.info(f"Model saved to {save_dir}")
# -- Final training report --
final_loss = training_log[-1]["loss"] if training_log else "N/A"
loss_chart = make_line_chart(
data=training_log,
x_key="epoch",
y_keys=["loss"],
title="Training Loss",
x_label="Epoch",
y_label="Loss",
colors=["#5a7db5"],
) if training_log else ""
lr_chart = ""
if training_log and "lr" in training_log[0]:
lr_chart = make_line_chart(
data=training_log,
x_key="epoch",
y_keys=["lr"],
title="Learning Rate Schedule",
x_label="Epoch",
y_label="LR",
colors=["#0f3460"],
)
await flyte.report.replace.aio(
wrap_report(
f"Training Complete
"
f"{model_name}
"
f''
f'
'
f'
'
f'
'
f'
{total_params:,}
Total Params
'
f'
{trainable_params:,}
Trainable Params
'
f'
{trainable_params / total_params * 100:.1f}%
% Trainable
'
f'
{loss_chart}
'
f'{f"""{lr_chart}
""" if lr_chart else ""}'
),
do_flush=True,
)
return await flyte.io.Dir.from_local(save_dir)
# ------------------------------------------------------------------
# Task 3: Evaluate — before/after comparison
# ------------------------------------------------------------------
@gpu_env.task(report=True)
async def evaluate(
model_name: str,
finetuned_dir: flyte.io.Dir,
data_dir: flyte.io.Dir,
num_examples: int = 50,
) -> str:
"""Compare base model vs fine-tuned model on test examples."""
import torch
from datasets import load_from_disk
from transformers import AutoModelForCausalLM, AutoTokenizer
log.info("Starting evaluation...")
await flyte.report.replace.aio(
wrap_report(
"Evaluation
"
'Loading models and running inference...
Evaluation
"
f''
f'
{len(eval_ds)}
Eval Examples
'
f'
'
f'
Running base model inference...
'
f'
Evaluation
"
f'Running base model inference... {i + 1}/{len(eval_ds)}
'
f'
Evaluation
"
f'Running fine-tuned model inference...
'
f'
Evaluation
"
f'Running fine-tuned model inference... {i + 1}/{len(eval_ds)}
'
f'
Q: {c['question']}
Schema: {c['schema'][:200]}...
| Source | SQL | Result |
|---|
| Expected | {c['expected']} | |
| Base | {c['base'][:200]} | {base_badge} |
| Fine-tuned | {c['finetuned'][:200]} | {ft_badge} |
"""
await flyte.report.replace.aio(
wrap_report(
f"Evaluation Results
"
f''
f'
{base_acc:.1f}%
Base Accuracy
'
f'
{ft_acc:.1f}%
Fine-Tuned Accuracy
'
f'
{improvement:+.1f}pp
Improvement
'
f'
'
f'
{bar_chart}
'
f'Example Comparisons {improvement_badge}
'
f'{examples_html}'
f''
f'Note: Exact match accuracy compares normalized SQL output. '
f'The fine-tuned model may generate semantically correct queries that differ in formatting.'
f'
'
),
do_flush=True,
)
return json.dumps({
"base_accuracy": round(base_acc, 1),
"finetuned_accuracy": round(ft_acc, 1),
"improvement": round(ft_acc - base_acc, 1),
"num_examples": total,
"comparisons": comparisons[:10],
})
# ------------------------------------------------------------------
# Pipeline: orchestrate everything
# ------------------------------------------------------------------
# {{docs-fragment pipeline}}
@cpu_env.task(report=True)
async def pipeline(
model_name: str = "HuggingFaceTB/SmolLM2-135M",
dataset_name: str = "b-mc2/sql-create-context",
method: str = "lora",
epochs: int = 3,
lr: float = 2e-4,
batch_size: int = 4,
max_train_samples: int = 5000,
max_eval_samples: int = 500,
num_eval_examples: int = 50,
lora_r: int = 16,
lora_alpha: int = 32,
) -> flyte.io.Dir:
"""
End-to-end LLM fine-tuning pipeline.
1. Download and format dataset
2. Fine-tune model (full / LoRA / QLoRA)
3. Evaluate: before/after comparison on test set
Returns the fine-tuned model directory so it can be served directly.
"""
log.info(f"Pipeline: {model_name} | method={method} | dataset={dataset_name}")
steps = ["Prepare Data", "Train", "Evaluate"]
method_badge = f'{method.upper()}'
# Step 1: Prepare data
await flyte.report.replace.aio(
wrap_report(
f"LLM Fine-Tuning Pipeline
"
f"{model_name} {method_badge}
"
f'{pipeline_step_indicator(0, steps)}'
f'Downloading and formatting dataset: {dataset_name}...
LLM Fine-Tuning Pipeline
"
f"{model_name} {method_badge}
"
f'{pipeline_step_indicator(1, steps)}'
f'Training in progress... check the train task report for live charts.
LLM Fine-Tuning Pipeline
"
f"{model_name} {method_badge}
"
f'{pipeline_step_indicator(2, steps)}'
f'Evaluating base vs fine-tuned model...
Pipeline Complete
"
f"{model_name} {method_badge}
"
f'{pipeline_step_indicator(3, steps)}'
f''
f'
{metrics["base_accuracy"]}%
Base Accuracy
'
f'
{metrics["finetuned_accuracy"]}%
Fine-Tuned Accuracy
'
f'
{improvement:+.1f}pp
Improvement {improvement_badge}
'
f'
'
f'
'
f'
{metrics["num_examples"]}
Eval Examples
'
f'
'
f'Check the train task report for training loss/LR charts, '
f'and the evaluate task report for detailed example comparisons.'
f'
'
),
do_flush=True,
)
log.info(f"Pipeline complete. Improvement: {metrics['improvement']:+.1f}pp")
return finetuned_dir
# {{/docs-fragment pipeline}}
if __name__ == "__main__":
flyte.init_from_config()
run = flyte.run(pipeline)
print(run.url)
run.wait()
```
*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/llm_fine_tuning_lora_qlora/llm_fine_tuning_lora_qlora.py*
```
# /// script
# requires-python = ">=3.12"
# dependencies = [
# "flyte>=2.4.0",
# "torch>=2.1.0",
# "transformers>=4.45.0",
# "peft>=0.13.0",
# "trl>=0.12.0",
# "bitsandbytes>=0.44.0",
# ...
# ]
# ///
```
## Orchestrate the pipeline
```
# /// script
# requires-python = ">=3.12"
# dependencies = [
# "flyte>=2.4.0",
# "torch>=2.1.0",
# "transformers>=4.45.0",
# "peft>=0.13.0",
# "trl>=0.12.0",
# "datasets>=3.0.0",
# "bitsandbytes>=0.44.0",
# "accelerate>=0.34.0",
# ]
# main = "pipeline"
# params = ""
# ///
import asyncio
import json
import logging
import os
import tempfile
import flyte
import flyte.io
import flyte.report
# {{docs-fragment env}}
import os
main_img = flyte.Image.from_uv_script(__file__, name="llm-fine-tuning-lora-qlora", pre=True)
gpu_env = flyte.TaskEnvironment(
name="llm-fine-tuning-lora-qlora-gpu",
image=main_img,
resources=flyte.Resources(cpu=4, memory="24Gi", gpu=1),
secrets=[flyte.Secret(key="huggingface-token", as_env_var="HF_TOKEN")],
)
cpu_env = flyte.TaskEnvironment(
name="llm-fine-tuning-lora-qlora-cpu",
image=main_img,
resources=flyte.Resources(cpu=2, memory="8Gi"),
depends_on=[gpu_env],
)
HF_TOKEN = os.environ.get("HF_TOKEN")
# {{/docs-fragment env}}
from report_helpers import make_bar_chart, make_line_chart, pipeline_step_indicator, wrap_report
logging.basicConfig(level=logging.WARNING, format="%(message)s", force=True)
log = logging.getLogger(__name__)
log.setLevel(logging.INFO)
# ------------------------------------------------------------------
# Task 1: Prepare dataset
# ------------------------------------------------------------------
@cpu_env.task(cache="auto")
async def prepare_data(
dataset_name: str = "b-mc2/sql-create-context",
max_train_samples: int = 5000,
max_eval_samples: int = 500,
) -> flyte.io.Dir:
"""Download dataset from HuggingFace and format for instruction fine-tuning."""
from datasets import DatasetDict, load_dataset
log.info(f"Loading dataset: {dataset_name}")
ds = load_dataset(dataset_name, split="train")
def format_example(ex):
return {
"text": (
"### Task: Generate a SQL query to answer the question.\n"
f"### Schema:\n{ex['context']}\n"
f"### Question:\n{ex['question']}\n"
f"### SQL:\n{ex['answer']}\n<|endoftext|>"
)
}
ds = ds.map(format_example)
# Split into train and eval
total = len(ds)
train_end = min(max_train_samples, total - max_eval_samples)
eval_start = train_end
eval_end = min(eval_start + max_eval_samples, total)
processed = DatasetDict({
"train": ds.select(range(train_end)),
"eval": ds.select(range(eval_start, eval_end)),
})
output_dir = os.path.join(tempfile.mkdtemp(), "dataset")
processed.save_to_disk(output_dir)
log.info(f"Dataset ready: {len(processed['train'])} train, {len(processed['eval'])} eval")
return await flyte.io.Dir.from_local(output_dir)
# ------------------------------------------------------------------
# Task 2: Train
# ------------------------------------------------------------------
@gpu_env.task(report=True)
async def train(
model_name: str,
data_dir: flyte.io.Dir,
method: str = "lora",
epochs: int = 3,
lr: float = 2e-4,
batch_size: int = 4,
lora_r: int = 16,
lora_alpha: int = 32,
) -> flyte.io.Dir:
"""Fine-tune a model using full, LoRA, or QLoRA method."""
import torch
from datasets import load_from_disk
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainerCallback
from trl import SFTConfig, SFTTrainer
log.info(f"Training: model={model_name}, method={method}")
# -- Load data --
data_path = await data_dir.download()
dataset = load_from_disk(data_path)
# -- Load tokenizer --
token_kwargs = {"token": HF_TOKEN} if HF_TOKEN else {}
tokenizer = AutoTokenizer.from_pretrained(model_name, **token_kwargs)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# -- Initial report: loading model --
await flyte.report.replace.aio(
wrap_report(
f"Loading Model...
"
f"{model_name}
"
f''
f"
Method: {method.upper()}
"
f"
Dataset: {len(dataset['train']):,} train / {len(dataset['eval']):,} eval
"
f"
Training in Progress...
{model_name}
{len(dataset['train']):,}
Train Examples
{trainable_params / total_params * 100:.1f}%
Trainable
Method: {method_badge} | Total params: {total_params:,} | Trainable: {trainable_params:,}
"""
charts_html = ""
if training_log:
current = training_log[-1]
progress_pct = current["step"] / max_steps * 100 if max_steps else 0
charts_html += f"""
Step {current['step']}/{max_steps}
({progress_pct:.0f}%) |
Epoch {current['epoch']:.2f}/{epochs} |
Loss:
{current['loss']:.4f}
"""
loss_chart = make_line_chart(
data=training_log,
x_key="epoch",
y_keys=["loss"],
title="Training Loss",
x_label="Epoch",
y_label="Loss",
colors=["#5a7db5"],
)
charts_html += f'{loss_chart}
'
if "lr" in training_log[0]:
lr_chart = make_line_chart(
data=training_log,
x_key="epoch",
y_keys=["lr"],
title="Learning Rate Schedule",
x_label="Epoch",
y_label="LR",
colors=["#0f3460"],
)
charts_html += f'{lr_chart}
'
if "grad_norm" in training_log[0]:
grad_chart = make_line_chart(
data=training_log,
x_key="epoch",
y_keys=["grad_norm"],
title="Gradient Norm",
x_label="Epoch",
y_label="Grad Norm",
colors=["#06d6a0"],
)
charts_html += f'{grad_chart}
'
return wrap_report(stats_html + charts_html)
# -- Metrics callback with live report updates --
class MetricsCallback(TrainerCallback):
def on_log(self, args, state, control, logs=None, **kwargs):
if not logs or "loss" not in logs:
return
entry = {
"step": state.global_step,
"epoch": round(logs.get("epoch", 0), 2),
"loss": round(logs["loss"], 4),
}
if "learning_rate" in logs:
entry["lr"] = logs["learning_rate"]
if "grad_norm" in logs:
entry["grad_norm"] = round(float(logs["grad_norm"]), 4)
training_log.append(entry)
log.info(
f"step={state.global_step}/{state.max_steps} "
f"epoch={entry['epoch']:.2f} "
f"loss={entry['loss']:.4f}"
)
asyncio.run_coroutine_threadsafe(
flyte.report.replace.aio(
_build_training_report(state.max_steps),
do_flush=True,
),
loop,
)
# -- Train --
output_dir = os.path.join(tempfile.mkdtemp(), "checkpoints")
training_args = SFTConfig(
output_dir=output_dir,
num_train_epochs=epochs,
per_device_train_batch_size=batch_size,
learning_rate=lr,
logging_steps=10,
save_strategy="epoch",
bf16=use_bf16,
fp16=not use_bf16 and torch.cuda.is_available(),
gradient_accumulation_steps=4,
warmup_steps=10,
report_to="none",
)
trainer = SFTTrainer(
model=model,
args=training_args,
train_dataset=dataset["train"],
eval_dataset=dataset["eval"],
processing_class=tokenizer,
callbacks=[MetricsCallback()],
)
log.info("Starting training...")
await asyncio.to_thread(trainer.train)
log.info("Training complete.")
# -- Merge LoRA weights and save --
save_dir = os.path.join(tempfile.mkdtemp(), "finetuned_model")
if method in ("lora", "qlora"):
log.info("Merging LoRA weights into base model...")
model = model.merge_and_unload()
model.save_pretrained(save_dir)
tokenizer.save_pretrained(save_dir)
log.info(f"Model saved to {save_dir}")
# -- Final training report --
final_loss = training_log[-1]["loss"] if training_log else "N/A"
loss_chart = make_line_chart(
data=training_log,
x_key="epoch",
y_keys=["loss"],
title="Training Loss",
x_label="Epoch",
y_label="Loss",
colors=["#5a7db5"],
) if training_log else ""
lr_chart = ""
if training_log and "lr" in training_log[0]:
lr_chart = make_line_chart(
data=training_log,
x_key="epoch",
y_keys=["lr"],
title="Learning Rate Schedule",
x_label="Epoch",
y_label="LR",
colors=["#0f3460"],
)
await flyte.report.replace.aio(
wrap_report(
f"Training Complete
"
f"{model_name}
"
f''
f'
'
f'
'
f'
'
f'
{total_params:,}
Total Params
'
f'
{trainable_params:,}
Trainable Params
'
f'
{trainable_params / total_params * 100:.1f}%
% Trainable
'
f'
{loss_chart}
'
f'{f"""{lr_chart}
""" if lr_chart else ""}'
),
do_flush=True,
)
return await flyte.io.Dir.from_local(save_dir)
# ------------------------------------------------------------------
# Task 3: Evaluate — before/after comparison
# ------------------------------------------------------------------
@gpu_env.task(report=True)
async def evaluate(
model_name: str,
finetuned_dir: flyte.io.Dir,
data_dir: flyte.io.Dir,
num_examples: int = 50,
) -> str:
"""Compare base model vs fine-tuned model on test examples."""
import torch
from datasets import load_from_disk
from transformers import AutoModelForCausalLM, AutoTokenizer
log.info("Starting evaluation...")
await flyte.report.replace.aio(
wrap_report(
"Evaluation
"
'Loading models and running inference...
Evaluation
"
f''
f'
{len(eval_ds)}
Eval Examples
'
f'
'
f'
Running base model inference...
'
f'
Evaluation
"
f'Running base model inference... {i + 1}/{len(eval_ds)}
'
f'
Evaluation
"
f'Running fine-tuned model inference...
'
f'
Evaluation
"
f'Running fine-tuned model inference... {i + 1}/{len(eval_ds)}
'
f'
Q: {c['question']}
Schema: {c['schema'][:200]}...
| Source | SQL | Result |
|---|
| Expected | {c['expected']} | |
| Base | {c['base'][:200]} | {base_badge} |
| Fine-tuned | {c['finetuned'][:200]} | {ft_badge} |
"""
await flyte.report.replace.aio(
wrap_report(
f"Evaluation Results
"
f''
f'
{base_acc:.1f}%
Base Accuracy
'
f'
{ft_acc:.1f}%
Fine-Tuned Accuracy
'
f'
{improvement:+.1f}pp
Improvement
'
f'
'
f'
{bar_chart}
'
f'Example Comparisons {improvement_badge}
'
f'{examples_html}'
f''
f'Note: Exact match accuracy compares normalized SQL output. '
f'The fine-tuned model may generate semantically correct queries that differ in formatting.'
f'
'
),
do_flush=True,
)
return json.dumps({
"base_accuracy": round(base_acc, 1),
"finetuned_accuracy": round(ft_acc, 1),
"improvement": round(ft_acc - base_acc, 1),
"num_examples": total,
"comparisons": comparisons[:10],
})
# ------------------------------------------------------------------
# Pipeline: orchestrate everything
# ------------------------------------------------------------------
# {{docs-fragment pipeline}}
@cpu_env.task(report=True)
async def pipeline(
model_name: str = "HuggingFaceTB/SmolLM2-135M",
dataset_name: str = "b-mc2/sql-create-context",
method: str = "lora",
epochs: int = 3,
lr: float = 2e-4,
batch_size: int = 4,
max_train_samples: int = 5000,
max_eval_samples: int = 500,
num_eval_examples: int = 50,
lora_r: int = 16,
lora_alpha: int = 32,
) -> flyte.io.Dir:
"""
End-to-end LLM fine-tuning pipeline.
1. Download and format dataset
2. Fine-tune model (full / LoRA / QLoRA)
3. Evaluate: before/after comparison on test set
Returns the fine-tuned model directory so it can be served directly.
"""
log.info(f"Pipeline: {model_name} | method={method} | dataset={dataset_name}")
steps = ["Prepare Data", "Train", "Evaluate"]
method_badge = f'{method.upper()}'
# Step 1: Prepare data
await flyte.report.replace.aio(
wrap_report(
f"LLM Fine-Tuning Pipeline
"
f"{model_name} {method_badge}
"
f'{pipeline_step_indicator(0, steps)}'
f'Downloading and formatting dataset: {dataset_name}...
LLM Fine-Tuning Pipeline
"
f"{model_name} {method_badge}
"
f'{pipeline_step_indicator(1, steps)}'
f'Training in progress... check the train task report for live charts.
LLM Fine-Tuning Pipeline
"
f"{model_name} {method_badge}
"
f'{pipeline_step_indicator(2, steps)}'
f'Evaluating base vs fine-tuned model...
Pipeline Complete
"
f"{model_name} {method_badge}
"
f'{pipeline_step_indicator(3, steps)}'
f''
f'
{metrics["base_accuracy"]}%
Base Accuracy
'
f'
{metrics["finetuned_accuracy"]}%
Fine-Tuned Accuracy
'
f'
{improvement:+.1f}pp
Improvement {improvement_badge}
'
f'
'
f'
'
f'
{metrics["num_examples"]}
Eval Examples
'
f'
'
f'Check the train task report for training loss/LR charts, '
f'and the evaluate task report for detailed example comparisons.'
f'
'
),
do_flush=True,
)
log.info(f"Pipeline complete. Improvement: {metrics['improvement']:+.1f}pp")
return finetuned_dir
# {{/docs-fragment pipeline}}
if __name__ == "__main__":
flyte.init_from_config()
run = flyte.run(pipeline)
print(run.url)
run.wait()
```
*Source: https://github.com/unionai/unionai-examples/blob/main/v2/tutorials/llm_fine_tuning_lora_qlora/llm_fine_tuning_lora_qlora.py*
## Run the workflow
Create a HuggingFace token secret if you use a gated base model:
```
flyte create secret huggingface-token
```
From the [example directory](https://github.com/unionai/unionai-examples/tree/main/v2/tutorials/llm_fine_tuning_lora_qlora):
```
cd v2/tutorials/llm_fine_tuning_lora_qlora
uv run --script llm_fine_tuning_lora_qlora.py
```
Try QLoRA on a GPU:
```
flyte run llm_fine_tuning_lora_qlora.py pipeline --method qlora --epochs 3
```
QLoRA requires CUDA; LoRA and full fine-tuning follow the same entry point with different memory requirements.
---
**Source**: https://github.com/unionai/unionai-docs/blob/main/content/tutorials/model-training/llm-fine-tuning-lora-qlora/_index.md
**HTML**: https://www.union.ai/docs/v2/union/tutorials/model-training/llm-fine-tuning-lora-qlora/