TinyToolLM — End-to-End SLM Training Stack

A fully custom pipeline for pretraining and fine-tuning a 40M-parameter transformer on smart-home tool-use tasks. Covers everything from raw dataset preparation and distributed pretraining, through synthetic data distillation and supervised fine-tuning, to automated evaluation and interactive inference.

---

Results

Metric	Score
Intent / Tool Accuracy	97.7%
Slot Extraction Accuracy	69.5%
JSON Validity	99.0%
Negative / Refusal Recall	91.5%

Evaluated on an 800-example unseen holdout set generated by the same distillation pipeline.

---

Architecture

Component	Detail
Parameters	~40M
Layers	8 Transformer blocks
Hidden Dimension	1024
Attention Heads	8
Normalization	Fused Triton RMSNorm
Vocabulary	GPT-2 BPE (50,304 tokens)
Max Sequence Length	1,024 (pretraining) / 256 (SFT)

---

Repository Structure

.
├── custom_kernels.py               # Fused Triton RMSNorm kernel + autograd wrapper
├── prepare_pretraining_dataset.py  # Download & tokenize TinyStories → binary file
├── pretraining.py                  # Distributed FSDP pretraining on TinyStories
├── synthetic_data_distillation.py  # Qwen2.5-7B teacher → 5K SFT samples
├── finetuning.py                   # Distributed FSDP supervised fine-tuning
├── evaluate.py                     # Automated evaluation harness (holdout set)
└── inference.py                    # Interactive chat interface

---

Pipeline Overview

TinyStories (500M tokens)
        │
        ▼
prepare_pretraining_dataset.py  →  /data/training_dataset.bin
        │
        ▼
pretraining.py  →  /model/pretrained_model.pt
        │
        ▼
synthetic_data_distillation.py  →  /data/sft_training_dataset.pt
                                    /data/sft_test_dataset.pt
        │
        ▼
finetuning.py  →  /model/finetuned_model.pt
        │
        ├──▶  evaluate.py   (automated metrics)
        └──▶  inference.py  (interactive chat)

---

Key Components

1. Custom Triton RMSNorm (custom_kernels.py)

A fused GPU kernel for RMSNorm written in Triton, integrated as a drop-in nn.Module via a custom torch.autograd.Function.

• Forward pass — computes x / RMS(x) * w in a single kernel launch, storing inv_rms for the backward pass.
• Backward pass — recomputes normalized activations and accumulates weight gradients atomically, avoiding a second read of the input.
• All intermediate math is promoted to float32 for numerical stability while storing outputs in the original dtype.

2. Dataset Preparation (prepare_pretraining_dataset.py)

Downloads the TinyStories dataset from Hugging Face and tokenizes it using the GPT-2 BPE tokenizer across all available CPU cores. The full token stream is written to a flat uint16 binary file for memory-mapped random access during training.

3. Distributed Pretraining (pretraining.py)

Trains the model on ~500M tokens from TinyStories using PyTorch Fully Sharded Data Parallel (FSDP) with per-block wrapping, bfloat16 mixed precision, and AdamW. Token throughput and loss are tracked in real time via Weights & Biases.

4. Synthetic Data Distillation (synthetic_data_distillation.py)

Uses Qwen2.5-7B-Instruct as a teacher model to generate 5,000 high-quality SFT samples across four categories:

Category	Tool	Slot
lights	smart_lights	Room name
weather	get_weather	City name
music	media_player	Genre
negative	(refusal)	—

Prompt diversity is achieved through persona injection — each generation pass conditions the teacher on one of seven distinct user speaking styles (e.g. frustrated teenager, elderly formal speaker, keyword-command style). The pipeline validates slot annotations, deduplicates samples, and serializes train/test tensors with SFT-style label masking (-100 on prompt tokens).

5. Supervised Fine-Tuning (finetuning.py)

Resumes from pretrained weights and fine-tunes for 5 epochs using FSDP, bfloat16, and a cross-entropy loss that ignores prompt tokens. A 5% validation split is evaluated after each epoch, with the global validation loss reduced across all ranks via dist.all_reduce.

6. Evaluation (evaluate.py)

Runs greedy-ish decoding (temperature 0.3, top-10 sampling) over the 800-example holdout set and reports four metrics:

• Valid JSON Syntax — strict json.loads parse of <tool_call> content
• Intent / Tool Accuracy — exact match on the name field
• Slot Extraction Accuracy — case-insensitive match on room / location / genre
• Negative / Refusal Recall — no <tool_call> or { in output for out-of-scope queries

7. Inference (inference.py)

Interactive terminal chat loop. Formats prompts in the same <system> / <user> / <assistant> template used during SFT and runs autoregressive decoding with repetition penalty until EOS or the 50-token budget is reached.

---

Getting Started

Prerequisites

pip install torch transformers triton datasets wandb tqdm numpy
# For faster HuggingFace downloads
pip install hf_transfer
export HF_TOKEN="your_HF_token"

Step 1 — Prepare Pretraining Data

python prepare_pretraining_dataset.py
# Output: /data/training_dataset.bin

Step 2 — Pretrain

torchrun --nproc_per_node=<NUM_GPUS> pretraining.py
# Output: /model/pretrained_model.pt

Step 3 — Generate SFT Dataset

python synthetic_data_distillation.py
# Output: /data/sft_training_dataset.pt
#         /data/sft_test_dataset.pt

Step 4 — Fine-Tune

torchrun --nproc_per_node=<NUM_GPUS> finetuning.py
# Output: /model/finetuned_model.pt

Step 5 — Evaluate

python evaluate.py

Step 6 — Chat

python inference.py

---

Example Outputs

You: Turn on the lights in the kitchen
Assistant: <tool_call>{"name": "smart_lights", "room": "kitchen", "action": "turn_on"}</tool_call>

You: What's the weather like in Paris?
Assistant: <tool_call>{"name": "get_weather", "location": "paris"}</tool_call>

You: Play some jazz music
Assistant: <tool_call>{"name": "media_player", "genre": "jazz", "action": "play"}</tool_call>

You: What is the capital of France?
Assistant: I am a smart home assistant. I can control lights, weather, and music, but I cannot help with that.

---

Training Configuration

Setting	Pretraining	Fine-Tuning
Optimizer	AdamW	AdamW
Learning Rate	4e-4	5e-5
Weight Decay	0.1	0.1
Batch Size	16	32
Sequence Length	1,024	256
Precision	bfloat16	bfloat16
Gradient Clip	—	1.0
Steps / Epochs	14,466 steps (2 epochs)	5 epochs
Parallelism	FSDP	FSDP
Monitoring	W&B	W&B

---

Supported Tools

Tool Name	Trigger	Parameters
smart_lights	Lighting requests	room, action
get_weather	Weather queries	location
media_player	Music playback	genre, action

All other queries produce a natural-language refusal response.