Fine-Tuning Qwen3 with LoRA for Reliable Tool Calls: A Practical XYZ-Aquila-SFT Recipe

Fine‑Tuning Tool‑Calling LLMs: a practical Qwen3 + XYZ‑Aquila‑SFT SFT recipe

This guide shows how to fine‑tune Qwen3 with LoRA so it reliably emits structured tool calls from the XYZ‑Aquila‑SFT trajectories while preserving in‑turn reasoning (<think> blocks). Follow the pipeline and you’ll get a reproducible set of structured JSONL artifacts, a corpus stats report, and a small LoRA adapter you can iterate on (the notebook is Colab/GPU friendly).

What you’ll build, quickly

Stream and inspect XYZAILab/XYZ‑Aquila‑SFT examples (demo: N_STREAM = 400, reserve N_EVAL = 40).
Parse multi‑turn trajectories with embedded <tools> schemas and <tool_call> objects; extract schemas and re‑render them as a Qwen‑compatible system message.
Render manual ChatML and apply assistant‑only loss masking (labels = −100 for non‑assistant tokens) so internal reasoning is preserved.
Create a PyTorch Dataset and collator, attach LoRA adapters via PEFT, and run a mixed‑precision LoRA fine‑tune on Qwen/Qwen3‑0.6B.
Build teacher‑forced probes and evaluate parseable tool‑call fraction, tool‑name accuracy, and argument‑key F1.
Export: structured JSONL, a corpus_stats.json, and a saved LoRA adapter in /content/aquila_out.

Key demo constants and artifacts (as used in the notebook)

Dataset: XYZAILab/XYZ-Aquila-SFT (language = “en”)
Model: Qwen/Qwen3-0.6B
N_STREAM = 400, N_EVAL = 40
MAX_SEQ_LEN = 2048, LENGTH_POLICY = “truncate”
RUN_TRAINING = True, MAX_STEPS = 30 (smoke test)
LR = 1e-4, LORA_R = 16, GRAD_ACCUM = 8
N_EVAL_PROBES = 24 (teacher‑forced probes)
OUT_DIR = “/content/aquila_out” (writes aquila_en_structured_tools.jsonl and corpus_stats.json; LoRA adapter saved to {OUT_DIR}/lora_adapter)
DEV = “cuda” if torch.cuda.is_available() else “cpu”; BF16 detection via torch.cuda.is_bf16_supported()

Why manual ChatML and assistant‑only masking matter

The tutorial highlights a concrete issue: the tokenizer helper can silently remove the very reasoning tokens you want to train on. The notebook calls this out directly:

“Manual ChatML so we control masking token-exactly. WHY NOT apply_chat_template(): Qwen3’s template deletes <think>…</think> from every assistant turn except the last. On this dataset that silently destroys most of the reasoning supervision you are paying to train on.”

Treat that as an empirical observation from the notebook. Tokenizer and template behavior can change across library versions. The practical checklist is short:

Does your dataset include in‑turn reasoning tokens like <think> blocks?
If yes, test whether your tokenizer.apply_chat_template preserves them. If it removes or canonicalizes them, render ChatML yourself.
When you render ChatML manually, mark assistant spans and set labels = −100 for every token that should not contribute to loss, so loss is computed only on assistant tokens you intend to supervise.

Quick validation steps you can do without special tooling:

Construct a short multi‑turn conversation with at least two assistant turns containing <think>…</think>.
Run your tokenizer’s apply_chat_template on that conversation and inspect the string and tokenized output to confirm whether <think> blocks survive in every assistant turn.
If any <think> tokens were removed, switch to manual ChatML and implement token‑level masking as described below.

Manual ChatML + token masking, an explicit example (conceptual)

Render an explicit ChatML sequence of messages using the exact markers the model expects (the notebook uses markers such as <|im_start|> and <|im_end|>). For example, a single trajectory might round‑trip as:

<|im_start|>system
System instructions…
<|im_end|>

<|im_start|>user
What’s the weather in Seoul?
<|im_end|>

<|im_start|>assistant
I will check. <think>look up best source</think>
<tool_call>{“name”:”get_weather”, “arguments”:{“city”:”Seoul”}}</tool_call>
<|im_end|>

Tokenization and labels are produced so that labels[i] == -100 for every token that is not part of the assistant’s supervised output, and labels for the assistant tokens contain the target token ids (standard Transformers/Trainer ignore label == -100). In practice implement an assertion to enforce label and input_id alignment: len(labels) == len(input_ids) and labels contain at least one non‑-100 value for supervised examples.

Supervised‑token ratio (computed in the notebook) is simply:

supervised_token_ratio = (number of tokens with label != -100) / (total tokens)

Monitoring this ratio helps detect examples that are nearly all context (low ratio) or almost entirely supervised (high ratio). The notebook prints mean/p10/p90/max values for calls per trajectory, messages per trajectory, and chars per trajectory to help debug example composition.

Parsing tool calls robustly

XYZ‑Aquila‑SFT stores tool calls as JSON embedded inside XML‑like markers (<tool_call>…</tool_call>) and tool schemas appear in <tools> blocks. Real outputs are noisy, so the notebook uses a nesting‑safe JSON scanner (iter_json_objects) and a parse_tool_calls routine to extract JSON safely even with nested braces.

Illustrative (not dataset) example of the pattern the parser targets:

Noisy assistant output: “Okay, calling tool now: <tool_call>{ “name”: “get_weather”, “arguments”: {“city”:”Seoul”} }</tool_call>, I’ll return the result.”

iter_json_objects tolerates surrounding text and nested braces to yield the embedded JSON object {“name”:”get_weather”, “arguments”:{“city”:”Seoul”}}. The notebook also renders a Qwen tools template to teach the model the expected wrapper and JSON shape; the template begins with the instructional fragment the notebook uses:

“You are provided with function signatures within <tools></tools> XML tags:
<tools>
{lines}
</tools>

For each function call, return a json object with function name and arguments within <tool_call> XML tags:
<tool_call>
{{\”name\”: <function-name>, \”arguments\”: <args-json-object>}}
</tool_call>”

Concrete robustness advice:

Require the <tool_call> wrapper at generation time when possible, it makes parsing far less brittle.
Keep defensive extractors: nesting‑safe JSON scanning, JSON key validation, and fallback heuristics such as trying to fix missing quotes or balance braces for noisy outputs.
For production, consider constrained decoding or a two‑step pipeline where the model generates and a small post‑processor canonicalizes to strict JSON to reduce parser failures.

Teacher‑forced probes: what they measure and how the notebook builds them

The notebook defines teacher‑forced probes as:

“Teacher-forced probes: cut the trajectory right before an assistant turn that issues a tool call; the gold label is that call.”

This evaluation answers a focused question: given the correct dialogue prefix, can the model produce the precise structured call? The notebook reports three practical signals for each probe:

Parseable fraction: portion of model outputs that decode into JSON tool calls.
Tool‑name accuracy: whether the predicted “name” matches gold.
Argument‑key F1: F1 computed on the set of argument keys (matches in keys only; value mismatches are not penalized by this metric).

Important experimental caveats the notebook calls out explicitly: the demo uses N_EVAL_PROBES = 24 and MAX_STEPS = 30, small numbers suitable for smoke tests. The notebook prints the reminder:

“(30 steps on ~350 trajectories is a smoke test, not a result, expect noise, and scale N_STREAM/MAX_STEPS for anything real.)”

Recommendation: for a stable comparison use hundreds to low thousands of probes (for example, 200-2, 000) depending on variance, and repeat runs with different seeds for statistical confidence.

LoRA + PEFT training specifics and practical tuning guidance

The notebook attaches LoRA adapters to Qwen3 using PEFT (LoraConfig + get_peft_model) and trains only the adapter parameters. Demo hyperparameters in the notebook are:

LR = 1e‑4
LORA_R = 16
GRAD_ACCUM = 8
MAX_STEPS = 30 (smoke test)
Mixed precision and autocast when the device supports it; BF16 detection uses torch.cuda.is_bf16_supported()

Practical tuning ranges and notes:

For small models (≤1B) try LR in [5e‑5, 2e‑4]. For mid and large models start lower and sweep carefully.
Sweep LORA_R (rank) over 4, 8, 16, 32. Higher ranks increase capacity but also memory and potential overfitting.
Use gradient accumulation to emulate larger batch sizes when GPU memory is constrained. Monitor training loss and validation metrics frequently, and save adapters on improvement.
Checkpoint adapters regularly and save tokenizer and config alongside the adapter to ensure reproducibility.

Qwen3 docs also recommend careful sampling in thinking mode (non‑greedy sampling) during generation; follow model family guidance for thinking and non‑thinking sampling hyperparameters when you create teacher‑forced probes.

Practical Colab/GPU checklist, get to a working smoke test

Stream N_STREAM = 400 rows and inspect the prevalence and quality of <tools> and <tool_call> markup. If structured rows are rare, increase N_STREAM before training.
Parse rows into a Trajectory record: question, answer, messages, declared_calls, parsed calls. Save structured tools JSONL at /content/aquila_out/aquila_en_structured_tools.jsonl.
Render manual ChatML and produce labels with non‑assistant tokens masked to −100. Add an assertion: len(input_ids) == len(labels).
Build the PyTorch SFT dataset and collator that pads input_ids and labels consistently; use LENGTH_POLICY = “truncate” or “drop” as appropriate for your data budget.
Attach LoRA via PEFT, set mixed precision and autocast and gradient accumulation, then run the smoke training loop (MAX_STEPS = 30 is a debug run; scale up for real results).
Build teacher‑forced probes (demo N_EVAL_PROBES = 24) and compute parseable fraction, tool‑name accuracy, and argument‑key F1. Complement with at least one end‑to‑end execution test against a sandboxed tool implementation.
Save outputs: aquila_en_structured_tools.jsonl, corpus_stats.json, and the adapter directory at /content/aquila_out/lora_adapter.

Robustness and limitations, what to watch for

Tokenizer/template differences. The notebook observed an apply_chat_template behavior that removed intermediate <think> blocks; validate your tokenizer and transformers/PEFT versions locally and render ChatML manually if needed.
Parsing brittleness. Even nesting‑safe extraction can fail on free‑form model text. Use strict wrappers, constrained decoding, or a canonicalizer post‑processor for production.
Teacher‑forced ≠ free‑running. Success on teacher‑forced probes shows conditional generation ability under a gold prefix; it does not guarantee robust agentic behavior when earlier outputs deviate. Run end‑to‑end execution tests.
Small demo runs are noisy. The notebook’s demo hyperparameters (MAX_STEPS = 30, N_EVAL_PROBES = 24) are explicitly a smoke test, don’t treat them as tuned defaults or proof of generalization.
Versioning and reproducibility. Save the adapter, tokenizer config, and corpus_stats.json. Differences in transformers/PEFT versions can change tokenizer and chat template behavior and results.

Where this approach is a good fit, and where to invest more

LoRA plus assistant‑only masking is an efficient way to teach structural tool‑call behavior given dialogue prefixes. It gives you small checkpoints, fast iteration, and focused supervision. This method suits projects that need compact adapters to map conversation context to structured API calls.

Do not expect this pipeline on its own to produce robust multi‑turn agents or out‑of‑distribution tool generalization. For production agents invest in:

Larger SFT runs and hyperparameter sweeps (scale N_STREAM and MAX_STEPS).
End‑to‑end execution testing with sandboxed or mocked tool endpoints.
Free‑running evaluation, error‑handling behavior, and execution‑guided finetuning (or RLHF) if you need the agent to learn corrective behavior.

Next practical nudges (three quick experiments)

Sanity check the dataset: sample a few hundred rows and measure the fraction that include <tools> or <tool_call> markup. If structured markup is sparse, prioritize data curation.
Validate template behavior: create a mini conversation with multiple assistant turns that include <think> blocks and check whether tokenizer.apply_chat_template preserves them; if not, switch to manual ChatML and token‑masking.
Run a smoke LoRA training (the notebook shows MAX_STEPS = 30) to validate end‑to‑end data flow, then scale data and steps before drawing conclusions from metrics.

Key questions, answered

Does tokenizer.apply_chat_template always strip intermediate <think> blocks?

The tutorial author observed that behavior in their environment and therefore used manual ChatML: “Manual ChatML so we control masking token-exactly. WHY NOT apply_chat_template(): Qwen3’s template deletes <think>…</think> from every assistant turn except the last.” Treat this as an empirical observation, validate with your tokenizer and transformers/PEFT versions before deciding.
Will a 30‑step LoRA run prove the SFT worked?

No. As the notebook cautions: “(30 steps on ~350 trajectories is a smoke test, not a result, expect noise, and scale N_STREAM/MAX_STEPS for anything real.)” Use short runs to debug; scale data and steps for meaningful results.
What do teacher‑forced probes actually measure?

They truncate a trajectory before a tool call and ask the model to predict that call. This measures conditional tool‑call prediction (parseable output, tool‑name accuracy, argument‑key F1) under a gold prefix, not free‑running agent behavior or execution correctness.
Are the demo LoRA hyperparameters production‑ready?

The notebook uses LR=1e‑4, LORA_R=16 and GRAD_ACCUM=8 as demonstration values. They are sensible starting points for a smoke test but not tuned defaults, sweep learning rates and ranks, and scale training steps and data for production experiments.
How should I guard against brittle parsing?

Require strict wrappers (<tool_call>…</tool_call>), use a nesting‑safe JSON extractor, add a post‑generation canonicalizer, consider constrained decoding, and validate predicted calls by executing them against a sandboxed API to detect malformed calls early.

Final practical checklist before you scale

Confirm dataset prevalence and quality of <tools>/<tool_call> markup.
Verify tokenizer/template behavior for <think> blocks; if it’s destructive, use manual ChatML and token‑level masking with labels = −100 for non‑assistant tokens.
Run a smoke LoRA iteration to validate the end‑to‑end flow, then scale N_STREAM and MAX_STEPS and run repeated seeds for robust comparisons.
Complement teacher‑forced probes with at least one end‑to‑end execution test against a sandbox to measure practical utility.
Save adapter, tokenizer config, and corpus_stats.json in /content/aquila_out for reproducibility.

If you want a drop‑in checklist or a short methods‑verification snippet to paste into the notebook (sanity checks for <think> preservation, label alignment assertions, and a minimal iter_json_objects test), I can provide that next, tailored to your transformers and PEFT versions.