Fine-Tuning Qwen3-0.6B for Tool Calling with XYZ-Aquila-SFT

Supervised fine-tuning LLMs for reliable tool-calling remains one of the more finicky corners of applied ML — existing chat templates silently mutate training data, and naive JSON extraction breaks on nested argument objects. A new end-to-end pipeline walks through every failure mode explicitly, using the XYZAILab/XYZ-Aquila-SFT dataset and Qwen/Qwen3-0.6B as concrete targets. The result is a reproducible workflow a practitioner can clone, adapt, and scale. Given that agentic deployment patterns are reshaping how models interact with external systems, getting tool-calling supervision right at the data-preparation stage has compounding downstream value.

Dataset Streaming and Trajectory Parsing

The pipeline streams 400 rows from XYZAILab/XYZ-Aquila-SFT (English split) using Hugging Face datasets>=3.0.0 with streaming=True, avoiding a full download. Each row exposes a question, answer, number of tool calls integer, and a trajectory list of role-tagged messages.

A critical implementation decision is the JSON extractor. The pipeline explicitly rejects re pattern matching — citing breakage on nested arguments objects present in every real tool call — in favour of a nesting-safe scanner that walks the raw string character by character using json.JSONDecoder.raw_decode. The parse_row function assembles a Trajectory dataclass, splitting the system message on a # Tools header, extracting embedded tool schemas with a <tools>...</tools> regex, harvesting all <tool_call> blocks from assistant turns, and separately counting <think> reasoning blocks and <tool_response> observations. Parser agreement with the dataset's declared number of tool calls field is logged at inference time as a direct validation check.

ChatML Rendering and Loss Masking

Rather than calling tok.apply_chat_template(), the pipeline manually renders every trajectory into ChatML using <|im_start|>, <|im_end|>, and newline delimiters. The rationale is explicit: Qwen3's built-in template deletes <think>...</think> blocks from every assistant turn except the last, silently destroying the reasoning supervision signal for multi-step trajectories.

The manual renderer applies label = -100 to all header and non-assistant tokens, preserving loss only on assistant-generated body and tail tokens. Sequences exceeding MAX_SEQ_LEN = 2048 are truncated by default (LENGTH_POLICY = "truncate"), with a drop option available. Examples where all labels remain -100 after truncation are filtered out. The supervised-token ratio (mean, p10, p90) is reported across the encoded corpus as a dataset-health diagnostic — a low ratio indicates most sequence budget is consumed by prompt context rather than learnable output.

LoRA Configuration and Training Loop

The model is loaded as Qwen/Qwen3-0.6B, initialised with attn_implementation="sdpa" and BF16 dtype where the GPU supports it. Gradient checkpointing is enabled before LoRA attachment. The LoRA config sets r=16, lora_alpha=32 (2×r), lora_dropout=0.05, bias="none", and task_type="CAUSAL_LM". Training runs for MAX_STEPS=30 with batch size 1 and GRAD_ACCUM=8, yielding an effective 8 trajectories per step. The optimiser is AdamW with lr=1e-4, weight_decay=0.0, and betas=(0.9, 0.95). A cosine schedule with 5 warmup steps governs the learning rate. Loss, EMA loss, learning rate, and perplexity are logged every 5 steps. The tutorial frames this 30-step run explicitly as a smoke test on roughly 350 trajectories, not a publishable result, and recommends scaling N_STREAM and MAX_STEPS for production use.

Evaluation Protocol

Post-training evaluation uses teacher-forced probes: the trajectory is cut immediately before an assistant turn containing a tool call, the prefix is fed to the model with max_new_tokens=160 and greedy decoding (do_sample=False), and the generation is parsed for tool calls. Three metrics are reported: parseable rate, tool-name accuracy, and argument-key F1. The pipeline runs evaluation both before and after LoRA fine-tuning to produce a delta, using up to 24 probes (N_EVAL_PROBES=24) constructed from the held-out split.

The most consequential decisions in this pipeline — preserving <think> blocks by bypassing apply_chat_template, using nesting-safe JSON extraction, and verifying parser agreement against ground-truth counts — are exactly the steps that get skipped when teams move fast. As tool-calling models become standard infrastructure for AI agents interacting with external data sources, these low-level correctness guarantees become load-bearing. The XYZ-Aquila-SFT + Qwen3-0.6B pairing is Colab-accessible; the architecture of the data pipeline is what scales to larger models and datasets.