Skip to content

Added Sanity Automation Test for SFT - Takes 5 mins to run #177

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chethan-rfai merged 2 commits into from
Feb 12, 2026
+759 −0
Merged

Added Sanity Automation Test for SFT - Takes 5 mins to run #177

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
522f53c
Added Sanity Automation Test for SFT - Takes 5 mins to run
chethan-rfai Feb 11, 2026
5c65f59
Addressed Review Comments
chethan-rfai Feb 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Empty file added tests/sanity/logs/.gitkeep
View file
Open in desktop
Empty file.
357 changes: 357 additions & 0 deletions tests/sanity/scripts/rf-tutorial-sft-chatqa-sanity-1.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,357 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<div align=\"center\">\n",
"<a href=\"https://rapidfire.ai/\"><img src=\"https://raw.githubusercontent.com/RapidFireAI/rapidfireai/main/docs/images/RapidFire - Blue bug -white text.svg\" width=\"115\"></a>\n",
"<a href=\"https://discord.gg/6vSTtncKNN\"><img src=\"https://raw.githubusercontent.com/RapidFireAI/rapidfireai/main/docs/images/discord-button.svg\" width=\"145\"></a>\n",
"<a href=\"https://oss-docs.rapidfire.ai/\"><img src=\"https://raw.githubusercontent.com/RapidFireAI/rapidfireai/main/docs/images/documentation-button.svg\" width=\"125\"></a>\n",
"<br/>\n",
"Join Discord if you need help + ⭐ <i>Star us on <a href=\"https://github.com/RapidFireAI/rapidfireai\">GitHub</a></i> ⭐\n",
"<br/>\n",
"To install RapidFire AI on your own machine, see the <a href=\"https://oss-docs.rapidfire.ai/en/latest/walkthrough.html\">Install and Get Started</a> guide in our docs.\n",
"</div>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### RapidFire AI Tutorial Use Case: SFT for Customer Support Q&A Chatbot"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from rapidfireai import Experiment\n",
"from rapidfireai.automl import List, RFGridSearch, RFModelConfig, RFLoraConfig, RFSFTConfig"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Load Dataset and Specify Train and Eval Partitions"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from datasets import load_dataset\n",
"\n",
"dataset=load_dataset(\"bitext/Bitext-customer-support-llm-chatbot-training-dataset\")\n",
"\n",
"# Select a subset of the dataset for demo purposes\n",
"train_dataset=dataset[\"train\"].select(range(32))\n",
"eval_dataset=dataset[\"train\"].select(range(32,40))\n",
"train_dataset=train_dataset.shuffle(seed=42)\n",
"eval_dataset=eval_dataset.shuffle(seed=42)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Define Data Processing Function"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def sample_formatting_function(row):\n",
" \"\"\"Function to preprocess each example from dataset\"\"\"\n",
" # Special tokens for formatting\n",
" SYSTEM_PROMPT = \"You are a helpful and friendly customer support assistant. Please answer the user's query to the best of your ability.\"\n",
" return {\n",
" \"prompt\": [\n",
" {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
" {\"role\": \"user\", \"content\": row[\"instruction\"]},\n",
" \n",
" ],\n",
" \"completion\": [\n",
" {\"role\": \"assistant\", \"content\": row[\"response\"]}\n",
" ]\n",
" }"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Initialize Experiment"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Every experiment instance must be uniquely named\n",
"experiment = Experiment(experiment_name=\"exp1-chatqa-sanity-1\", mode=\"fit\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Define Custom Eval Metrics Function"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def sample_compute_metrics(eval_preds): \n",
" \"\"\"Optional function to compute eval metrics based on predictions and labels\"\"\"\n",
" predictions, labels = eval_preds\n",
"\n",
" # Standard text-based eval metrics: Rouge and BLEU\n",
" import evaluate\n",
" rouge = evaluate.load(\"rouge\")\n",
" bleu = evaluate.load(\"bleu\")\n",
"\n",
" rouge_output = rouge.compute(predictions=predictions, references=labels, use_stemmer=True)\n",
" rouge_l = rouge_output[\"rougeL\"]\n",
" bleu_output = bleu.compute(predictions=predictions, references=labels)\n",
" bleu_score = bleu_output[\"bleu\"]\n",
"\n",
" return {\n",
" \"rougeL\": round(rouge_l, 4),\n",
" \"bleu\": round(bleu_score, 4),\n",
" }"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Define Multi-Config Knobs for Model, LoRA, and SFT Trainer using RapidFire AI Wrapper APIs"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# 2 LoRA PEFT configs lite with different adapter capacities\n",
"peft_configs_lite = List([\n",
" RFLoraConfig(\n",
" r=8,\n",
" lora_alpha=16,\n",
" lora_dropout=0.1,\n",
" target_modules=[\"q_proj\", \"v_proj\"], # Standard transformer naming\n",
" bias=\"none\"\n",
" ),\n",
" RFLoraConfig(\n",
" r=32,\n",
" lora_alpha=64,\n",
" lora_dropout=0.1,\n",
" target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\"], # Standard naming\n",
" bias=\"none\"\n",
" )\n",
"])\n",
"\n",
"# 2 base models x 2 peft configs = 4 combinations in total\n",
"config_set_lite = List([\n",
" RFModelConfig(\n",
" model_name=\"TinyLlama/TinyLlama-1.1B-Chat-v1.0\", # 1.1B model\n",
" peft_config=peft_configs_lite,\n",
" training_args=RFSFTConfig(\n",
" learning_rate=1e-3, # Higher LR for very small model\n",
" lr_scheduler_type=\"linear\",\n",
" per_device_train_batch_size=4,\n",
" per_device_eval_batch_size=4,\n",
" max_steps=8,\n",
" gradient_accumulation_steps=1, # No accumulation needed\n",
" logging_steps=2,\n",
" eval_strategy=\"steps\",\n",
" eval_steps=4,\n",
" bf16=True,\n",
" ),\n",
" model_type=\"causal_lm\",\n",
" model_kwargs={\"device_map\": \"auto\", \"torch_dtype\": \"auto\", \"use_cache\": False},\n",
" formatting_func=sample_formatting_function,\n",
" compute_metrics=sample_compute_metrics,\n",
" generation_config={\n",
" \"max_new_tokens\": 256,\n",
" \"temperature\": 0.8, # Higher temp for tiny model\n",
" \"top_p\": 0.9,\n",
" \"top_k\": 30, # Reduced top_k\n",
" \"repetition_penalty\": 1.05,\n",
" }\n",
" ),\n",
" RFModelConfig(\n",
" model_name=\"TinyLlama/TinyLlama-1.1B-Chat-v1.0\", # 1.1B model\n",
" peft_config=peft_configs_lite,\n",
" training_args=RFSFTConfig(\n",
" learning_rate=1e-4, # Higher LR for very small model\n",
" lr_scheduler_type=\"linear\",\n",
" per_device_train_batch_size=4, # Even larger batch size\n",
" per_device_eval_batch_size=4,\n",
" max_steps=8,\n",
" gradient_accumulation_steps=1, # No accumulation needed\n",
" logging_steps=2,\n",
" eval_strategy=\"steps\",\n",
" eval_steps=4,\n",
" bf16=True,\n",
" ),\n",
" model_type=\"causal_lm\",\n",
" model_kwargs={\"device_map\": \"auto\", \"torch_dtype\": \"auto\", \"use_cache\": False},\n",
" formatting_func=sample_formatting_function,\n",
" compute_metrics=sample_compute_metrics,\n",
" generation_config={\n",
" \"max_new_tokens\": 256,\n",
" \"temperature\": 0.8, # Higher temp for tiny model\n",
" \"top_p\": 0.9,\n",
" \"top_k\": 30, # Reduced top_k\n",
" \"repetition_penalty\": 1.05,\n",
" }\n",
" )\n",
"])\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Define Model Creation Function for All Model Types Across Configs"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"\n",
"def sample_create_model(model_config): \n",
" \"\"\"Function to create model object for any given config; must return tuple of (model, tokenizer)\"\"\"\n",
" from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForMaskedLM\n",
"\n",
" model_name = model_config[\"model_name\"]\n",
" model_type = model_config[\"model_type\"]\n",
" model_kwargs = model_config[\"model_kwargs\"]\n",
" \n",
" if model_type == \"causal_lm\":\n",
" model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)\n",
" elif model_type == \"seq2seq_lm\":\n",
" model = AutoModelForSeq2SeqLM.from_pretrained(model_name, **model_kwargs)\n",
" elif model_type == \"masked_lm\":\n",
" model = AutoModelForMaskedLM.from_pretrained(model_name, **model_kwargs)\n",
" elif model_type == \"custom\":\n",
" # Handle custom model loading logic, e.g., loading your own checkpoints\n",
" # model = ... \n",
" pass\n",
" else:\n",
" # Default to causal LM\n",
" model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)\n",
" \n",
" tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
" \n",
" return (model,tokenizer)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Generate Config Group"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Simple grid search across all sets of config knob values = 4 combinations in total\n",
"config_group = RFGridSearch(\n",
" configs=config_set_lite,\n",
" trainer_type=\"SFT\"\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Run Multi-Config Training"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Launch training of all configs in the config_group with swap granularity of 4 chunks\n",
"experiment.run_fit(config_group, sample_create_model, train_dataset, eval_dataset, num_chunks=4, seed=42)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### End Current Experiment"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"experiment.end()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<div align=\"center\">\n",
"<a href=\"https://rapidfire.ai/\"><img src=\"https://raw.githubusercontent.com/RapidFireAI/rapidfireai/main/docs/images/RapidFire - Blue bug -white text.svg\" width=\"115\"></a>\n",
"<a href=\"https://discord.gg/6vSTtncKNN\"><img src=\"https://raw.githubusercontent.com/RapidFireAI/rapidfireai/main/docs/images/discord-button.svg\" width=\"145\"></a>\n",
"<a href=\"https://oss-docs.rapidfire.ai/\"><img src=\"https://raw.githubusercontent.com/RapidFireAI/rapidfireai/main/docs/images/documentation-button.svg\" width=\"125\"></a>\n",
"<br/>\n",
"Thanks for trying RapidFire AI! ⭐ <i>Star us on <a href=\"https://github.com/RapidFireAI/rapidfireai\">GitHub</a></i> ⭐\n",
"</div>"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.7"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
Loading