A small, fine-tuned Qwen2.5-1.5B-Instruct model that replaces two
LLM calls in a LangGraph-based multi-tool investment research agent
(classifierNode + extractFiltersFromQuestions) with a single
structured-output call — trained through a full CPT → SFT → GRPO
pipeline.
Why this exists
The original agent used two separate LLM calls per research question,
both served by a 120B-class model:
classifierNode — decomposes the question into sub-questions and
assigns each a tool (RAG, WEB, or STOCK).
extractFiltersFromQuestions — for each RAG sub-question, a
second LLM call extracts companies, years, and a cleaned search
query. WEB/STOCK sub-questions got no structured extraction at
all — the raw question string was passed straight to Tavily / Yahoo
Finance.
This also hardcoded the RAG-eligible company list to exactly three
companies (NVIDIA, AMD, Microsoft), baked directly into both prompts.
The router model collapses this into one model call that outputs
tool routing and structured extraction (companies, years, search
query) for all three tools uniformly — including WEB/STOCK,
which the original system never structured at all. It also drops the
hardcoded company enum: extraction generalizes to any named company,
not just the original three.
Pipeline
Base model (Qwen2.5-1.5B-Instruct)
|
Continued Pretraining (CPT) -- domain vocabulary/register adaptation
|
Supervised Fine-Tuning (SFT) -- teaches the routing/extraction task
|
Reinforcement Learning (GRPO) -- closes schema/edge-case gaps SFT couldn't
|
grpo_router_final_merged <- final deployed router
1. Base model
Qwen2.5-1.5B-Instruct — chosen for its small footprint (deployable
on modest hardware, fast single-call inference) relative to the
120B-class model used for the original two-call baseline. The
Instruct checkpoint was kept throughout the entire pipeline
(never switched to the raw base model) — a deliberate choice: the SFT
dataset was sized assuming the model already knows how to follow
instructions and only needs to specialize, not learn instruction-
following from scratch. CPT was correspondingly kept light-touch,
specifically because it perturbs an already-aligned model.
2. Continued Pretraining (CPT)
Goal: light domain-vocabulary adaptation to SEC-filing register
(financial/legal terminology, disclosure structure) before task-specific
SFT — not deep domain rewiring, and not meant to teach the model to
answer financial questions (that's SFT's job).
Corpus: TheFinAI/SEC_2025 only — real FY2025 SEC filing text
(8-K item disclosures + 10-Q MD&A excerpts). A second candidate corpus
(a scraped tech-news dataset) was evaluated and deliberately dropped in
favor of a single-source, controlled experiment — avoids confounding
CPT's effect with an unrelated data-quality/licensing variable. A
mixed-source news+filings CPT experiment remains a documented future
option, not something ruled out permanently.
Final corpus: 14 deduplicated documents, ~6,000 tokens (Qwen2.5
tokenizer), capped at 600 tokens/doc. Corpus construction went through
two iterations — the first pull (10 docs) surfaced zero mentions of the
target priority companies (NVIDIA, AMD, Microsoft, Apple, etc.) despite
the source dataset containing them; the builder was hardened with
explicit stream shuffling, per-row progress logging, and a scan cap
before the second pull, which surfaced 3 genuine filings by priority
companies (AMD's own executive-appointment 8-K, Salesforce's and
Amazon's own earnings 8-Ks) among the 14. The corpus skews toward Item
5.02 personnel-filing boilerplate (8 of 14 docs) over Item 2.02/MD&A
financial-results language (6 of 14) — a known, accepted trade-off
given the light-touch CPT scope, not a blocking issue.
Training config (Unsloth + LoRA, QLoRA/4-bit):
r=16, lora_alpha=32 (2:1 ratio — CPT convention, more headroom than
SFT's later 1:1 to shift the embedding-space distribution)
target_modules: attention + FFN plus embed_tokens and lm_head
— the key CPT-specific inclusion, since vocabulary/output-distribution
adaptation is the whole point of this stage; SFT's LoRA later
deliberately excludes these two modules
learning_rate=2e-4, embedding_learning_rate=2e-5 (10x smaller —
embeddings are the model's foundation, updated gently)
3 epochs (more than a typical CPT pass would need for a larger corpus
— needed here specifically because the corpus is tiny; repeating a
small corpus more times is standard practice and doesn't meaningfully
raise catastrophic-forgetting risk given the corpus size stays small)
No 80/20 domain/general-text data mixing — judged unnecessary at this
scale (tiny corpus, few epochs, LoRA-only, frozen base weights); a
direct instruction-following sanity check was used instead of mixing
as the catastrophic-forgetting safeguard (see below)
Evaluation — two checks, not just perplexity, since CPT here
perturbs an aligned Instruct model and instruction-following damage
was the specific risk being guarded against:
The perplexity drop is intentionally modest (Qwen2.5 already has
substantial financial text in its pretraining mix, so there wasn't much
"surprise" left to remove — a large drop would have suggested
overfitting the tiny corpus, not healthy adaptation). The sanity-
generation check passing cleanly confirmed the higher-risk part of this
stage — embed_tokens/lm_head perturbation on an aligned model —
didn't degrade instruction-following. Held-out set size (4 docs) means
the perplexity number is directional, not precise.
Output: cpt_router_adapter (LoRA adapter, not merged) → later
merged into a standalone checkpoint, cpt_merged_model, which SFT was
built on top of.
3. Supervised Fine-Tuning (SFT)
3.1 Dataset construction
Synthetic dataset generated by a custom pipeline (generate_dataset.py
builders.py + assemble.py), targeting a fixed JSON output schema:
Key design principles (see DATASET_CARD.md for full rationale):
No hardcoded company/year enum — companies are extracted as a
general skill, not memorized from a fixed list.
years: [] for unresolved relative time ("recent", "latest",
"last year") — resolving relative time to concrete years is left to
the app layer, not baked into the model.
Minimal splitting — multiple companies/metrics for the same
tool stay in ONE sub-question; only a genuine tool change forces a
split.
Canonical-name extraction — tickers (NVDA) and known aliases
(Facebook) resolve to the canonical company name (NVIDIA, Meta)
in the output, never the raw form as written.
3.2 Dataset scaling
Version
Total examples
Train / Val
Companies
Structural signatures
v1
~760
--
--
--
v2
3,010
2,709 / 301
32
10
v3 (final SFT set)
8,513
7,662 / 851
64
25
v3 additions over v2: expanded entity vocab (finance, retail, energy,
healthcare, telecom, airlines, industrials, international tech —
64 companies total), a ticker-to-canonical-name map (40 tickers),
register diversity (informal abbreviations, Slack-terse phrasing,
formal analyst-memo phrasing), 7 boundary-trap builders (RAG/WEB/STOCK
boundary questions designed to be hard), multi-hop dependent-phrasing
builders, and new structural shapes (reversed tool orderings, mixed
company-count-bucket compounds) — raising distinct structural
signature coverage from 10 to 25, specifically to avoid scaling by
permuting entities within the same handful of shapes.
3.3 Post-hoc dataset patch (v3.1) — attempted, then deliberately not used
After the initial SFT run, out-of-domain evaluation surfaced 4 real
gaps: alias-name canonicalization not generalizing (tickers worked,
plain-English aliases like "Facebook" didn't), a schema leak
(searchQuery appearing on STOCK sub-questions with explicit
historical years), relative-time phrasing beyond "recent/latest/current"
not resolving to years: [], and a generator bug in one compound-
question builder where the target year/period wasn't stated in the
visible question text.
A small 374-example patch dataset was built and used for a short
continued-SFT pass. Two separate patch attempts (4 epochs, then 2
epochs on a clean-restarted model) both introduced new regressions
on the original held-out val set — schema-field hallucinations,
requiredTools-ordering errors — while only partially closing the
original gaps. This pointed to instability from training on too
narrow/isolated a correction set rather than a fixable epoch-count
issue.
Decision: the patch was discarded. The final SFT checkpoint used
downstream is the original v3, unpatched, 8,513-example SFT run.
The 4 known gaps were deliberately deferred to the GRPO stage instead,
where they could be addressed via reward-function checks on live
rollouts rather than more small-scale SFT correction passes.
3.4 SFT training run
Framework: Unsloth + TRL SFTTrainer, QLoRA-style LoRA adapters,
starting from cpt_merged_model (the CPT-adapted checkpoint, not raw
base) — target_modules restricted to attention + FFN only this time
(no embed_tokens/lm_head — that was CPT's job, already merged in)
r=16, lora_alpha=16 (1:1 — standard SFT convention, behavior
shaping through existing representations rather than pushing the
embedding space)
Trainable parameters: 18,464,768 of 1,562,179,072 (1.18%)
3 epochs, effective batch size 16 (batch 4 x grad accumulation 4,
later switched to packing=True mid-run for throughput), response-only
loss masking (train_on_responses_only — loss computed only on the
assistant's JSON output, not the repeated system prompt/question)
Final training loss: 0.000626, final validation loss: 0.001455 —
monotonically improving through the full run, no overfitting reversal
3.5 SFT evaluation
On a 120-question held-out sample from the v3 val set:
Metric
Result
JSON parse rate
100.0%
Schema validity rate
100.0%
Tool-sequence exact match
100.0%
Full exact match
99.2% (119/120)
The single mismatch was later confirmed to be a dataset artifact (a
compound-question builder not stating its target year/period in the
visible question), not a model error — see §3.3. Out-of-domain checks
(generic non-router prompts) confirmed no mode collapse into
always-emit-JSON behavior.
Output: sft_router_merged — the checkpoint GRPO starts from.
4. Reinforcement Learning (GRPO)
Starting checkpoint: sft_router_merged (the original, unpatched
v3 SFT model — not the abandoned v3.1 patch, see §3.3).
Why GRPO, and why these specific gaps: the 4 gaps deferred from SFT
(§3.3) are all programmatically verifiable from the question text or
against a gold label — exactly GRPO's ideal use case. A reward function
built on rule-based checks can directly shape these behaviors via live
rollouts, rather than depending on a small, narrow correction dataset
(which SFT patching had already shown to be unstable).
4.1 Reward function
router_reward.py is fully rule-based — no separate reward model. It
combines a hard schema-validity gate (malformed JSON or a schema
violation like searchQuery on a non-RAG sub-question caps the reward
at a fixed penalty, regardless of anything else) with weighted
component checks:
Component
Weight
Checks
required_tools_correct
1.0
requiredTools matches the sub-questions' tools, in canonical order
minimal_splitting
1.0
No tool appears more than once across adjacent sub-questions
An explicit year in the question -> STOCK/WEB sub-questions populate years, still omit searchQuery
alias_ticker_canonicalization
1.5
Tickers/aliases in the question resolve to the canonical company name in companies[]
gold_tool_sequence_match
2.0
Tool sequence matches gold, where gold is available
gold_company_overlap
1.0
Company extraction overlaps gold
gold_year_match
1.0
Year extraction overlaps gold
gold_exact_match_bonus
1.5
Full exact match against gold
Most components are self-checkable from the question alone (no gold
label needed) — the GRPO prompt pool for those checks can be larger
than the labeled SFT set. The reward function's ticker/alias list was
deliberately kept simple: single-letter tickers ("V" for Visa, "T" for
AT&T) and tickers that double as ordinary English words (NOW, SHOP,
COST, SNOW) were excluded rather than special-cased, since real
usage essentially never writes a bare single letter or common word as a
ticker reference. Word-boundary matching (not naive substring matching)
and an explicit-year regex covering both "in 2022" and "FY2024"
phrasing were fixed after an initial version produced false positives
on real model output — verified via a 15-case self-test and manual
inspection against live generations before being trusted for training.
4.2 Prompt pool
2,198 questions total, three sources:
Source
Count
Gold label?
Purpose
gold_sft
1,801
Yes
Stratified sample of the 8,513 SFT questions, preserving all 11 tool-sequence signature proportions — grounds the core routing skill
gap_targeted
352
No
New questions specifically exercising the 4 known gaps (alias/ticker mentions, explicit-year STOCK/WEB, relative-time phrasing, minimal-splitting temptation)
realistic_noise
45
No
Hand-written messy real-world phrasing (typos, run-ons, ALL CAPS, informal celebrity-based company references like "the zuck company") — a diagnostic stress test, not a training target in itself
4.3 Training attempts
Three training attempts were run; the middle one produced the final
checkpoint.
Attempt 1 — failed, no meaningful learning. Initial config
(num_generations=6, default temperature, learning_rate=5e-6,
400 planned steps) was first found to be too slow to be practical on a
free-tier T4 (~0.14 it/s); an attempt to enable Unsloth's vLLM-backed
fast-generation path failed due to a CUDA version mismatch (vLLM built
for CUDA 13, Colab's runtime on CUDA 12.8) and was abandoned in favor of
trimming the standard-generation config instead
(num_generations=4, max_completion_length=160, max_steps=200).
That trimmed run completed, with misleadingly high aggregate reward
(mostly 2.0-3.0 out of 3.0) — but reward_std was 0.000000 on over
half of logged steps, and kl divergence never meaningfully exceeded
~0.002 across the entire run. Diagnosis: the SFT-converged policy was
confident enough (SFT final loss ~0.0006) that default sampling
produced near-identical completions across the group on most prompts —
no differentiated signal for GRPO to learn from. Post-training
evaluation confirmed all 4 target gaps were completely unchanged
from the pre-GRPO baseline, including on the exact reward components
they were meant to fix. Key lesson: a high aggregate reward number
can fully mask zero progress on the specific failures a reward function
was designed to catch — always check reward_std/kl, and evaluate
against the specific target cases, not just the mean reward.
Attempt 2 — the final checkpoint. Restarted from sft_router_merged
fresh (attempt 1's near-zero policy movement meant nothing was lost by
not continuing from it). Two changes targeted the diagnosed cause
directly: temperature=1.3, top_p=0.95 (forces real sampling
diversity) and learning_rate=1e-5 (up from 5e-6 — headroom to move
when signal exists). Prompt pool reweighted toward the harder gap/noise
subsets (300 gold + all gap-targeted x2 + all noise x2 = 1,094
questions) rather than the full 2,198, concentrating training budget on
what needed fixing. kl divergence reached as high as 0.0248 at
individual steps (vs. attempt 1's ~0.002 ceiling) and reward_std
showed real, sometimes large, spread (e.g. 0.948 at step 30) — genuine
differentiated learning signal this time. Trained to step 175 (of a
200-step plan; the runtime session ended before completion, and the
checkpoint-175 save was used as-is rather than continuing).
Evaluation at checkpoint-175:
Gap
Status
Relative-time rule
Fully fixed — generalized correctly to novel relative-time phrasings not seen in training
STOCK/WEB schema leak
Partially fixed — the schema-crash (illegal searchQuery on STOCK/WEB) is eliminated; the underlying year-extraction on STOCK/WEB with an explicit year is still incomplete (years stays [] instead of populating)
Alias canonicalization
Not fixed — "Facebook" and "Alphabet" still don't resolve to "Meta"/"Google"; ticker resolution (already working pre-GRPO) is unaffected
requiredTools ordering
Not specifically regressed; not separately stress-tested in this eval batch
Regression check against the original 8,513-question SFT benchmark:
parse 100/100, schema 100/100, tool-sequence 100/100, exact match
99/100 — statistically identical to the pre-GRPO SFT baseline. Net
result: real, generalizing improvement on one gap, a safety-relevant
partial fix on a second, zero cost to the base task.
Attempt 3 — a targeted top-up, discarded. Continued training
directly from the attempt-2 checkpoint (not a fresh LoRA) on a pool
concentrated almost entirely on the one remaining clean-failure gap
(alias/ticker canonicalization weighted 3x, explicit-year 2x, only a
small 120-question gold anchor for regression protection — 637
questions total, 80 steps). Result: no improvement on alias
canonicalization (still 0.0 on every tested case) and a net
regression — a previously-correct compound RAG question now failed
schema validation (searchQuery missing), and the regression-suite
exact-match score dropped from 99/100 to 91/100. Discarded; attempt 2's
checkpoint-175 remains the final model. Alias/name-based
canonicalization has now failed to improve across four separate
correction attempts (two SFT patches, one general GRPO run, one
targeted GRPO top-up) — treated as a documented, deliberate open
limitation rather than a target for further correction attempts within
this project's scope (see §7).
4.4 Final merge
sft_router_merged (base) + attempt 2's checkpoint-175 adapter ->
merged -> grpo_router_final_merged, the final deployed router.
Post-merge sanity check (this project's own canonical example question:
"Compare NVIDIA vs AMD's data center revenue, recent news sentiment,
and stock performance") produced fully correct, schema-valid,
correctly-decomposed output on the first try.
5. Benchmark: Router vs. Original Baseline
Baseline = original classifierNode + extractFiltersFromQuestions
prompts, reproduced verbatim, served via Groq (openai/gpt-oss-120b),
JSON mode forced. Router = grpo_router_final_merged, run unbatched on
a free-tier Colab T4.
Evaluated on two held-out subsets, split deliberately to keep the
comparison honest:
In-scope: questions whose companies are within the baseline's
original hardcoded enum (NVIDIA / AMD / Microsoft) — a genuine
apples-to-apples routing comparison.
Out-of-scope: questions naming companies outside that enum — the
baseline cannot serve these by design (the enum is hardcoded into
its prompt), so this subset demonstrates a capability gain, not a
routing-accuracy contest.
*Latency caveat: baseline runs on Groq's hosted inference
infrastructure; the router ran unbatched on a free-tier T4 with no
serving optimization. This gap reflects infrastructure, not model
quality — it is not presented as a head-to-head speed claim. With
proper serving (e.g. vLLM, batching, or hosted inference), the router's
1.5B parameter count and single-call design would be expected to be
substantially faster than the 120B two-call baseline in practice.
Reading these numbers honestly
In-scope accuracy gap (61.3% -> 96.8%): the val set specifically
includes RAG/WEB/STOCK "boundary trap" questions (e.g. a filing
metric using STOCK-suggestive wording like "shares outstanding")
designed to be hard. The original baseline prompt has no rules for
these boundary cases; the router's training data was built
specifically to cover them.
In-scope RAG field accuracy gap (41.2% -> 94.1%): partly a
genuine design-philosophy difference, not purely "worse extraction."
The baseline's filter-extraction prompt explicitly resolves relative
time ("recent", "latest") to concrete years ([2024, 2025]); the
router's schema requires years: [] for unresolved relative time.
Every relative-time RAG question is a near-automatic baseline "miss"
under the router's schema by construction — still a fair comparison
since the router's schema is what the downstream system actually
consumes, but worth stating plainly rather than implying the baseline
extracts sloppily across the board.
Out-of-scope numbers: not a baseline failure so much as a scope
mismatch — the original prompt hardcodes exactly 3 companies. The
router's 87.5-100% here demonstrates genuine generalization
(extraction works on any company, no code changes, no retraining),
which is the headline capability claim of this project.
Sample sizes (n=~30 per subset) are enough to be directionally
solid but small enough that any single percentage point (e.g. the
out-of-scope subset slightly outperforming in-scope) shouldn't be
over-read — treat these as illustrative, not a rigorous statistical
claim, unless re-run at a larger n.
6. Repo structure
generate_dataset.py # system prompt, vocab lists, ticker/alias maps
builders.py # dataset builder functions + hand-crafted edge cases
assemble.py # weighted sampling, dedup, train/val split, stats
router_sft_train.jsonl # 7,662 SFT examples
router_sft_val.jsonl # 851 SFT examples
DATASET_CARD.md # full schema + design rationale + composition
build_cpt_corpus_colab.py # CPT corpus builder (TheFinAI/SEC_2025, stream-shuffled, deduplicated)
cpt_notebook.py # CPT training notebook (Qwen2.5-1.5B-Instruct, QLoRA)
sft_notebook.py # SFT training notebook (chat-template formatting, response-only loss masking)
router_reward.py # GRPO reward function (schema + rule-based checks)
grpo_prompt_pool.jsonl # 2,198-question GRPO prompt pool (gold + gap-targeted + noise)
grpo_router_final_merged/ # final deployed model
7. Known limitations
Alias/name-based company canonicalization does not work.
"Facebook" does not resolve to "Meta," "Alphabet" does not resolve to
"Google," and similar plain-English aliases are extracted verbatim
rather than canonicalized. Ticker-based resolution (NVDA ->
NVIDIA, MSFT -> Microsoft) works correctly and generalizes well.
This gap was targeted across 4 separate correction attempts (2 SFT
patch passes, 1 general GRPO run, 1 targeted GRPO top-up) — none
succeeded, and the final targeted attempt actively regressed the base
task instead. Documented here as a deliberate, accepted limitation
rather than continuing to spend training budget on diminishing
returns; a dedicated, isolated SFT pass with much heavier alias-
example density is the most likely path to actually closing this,
left as future work.
Explicit-year extraction on STOCK/WEB sub-questions is incomplete.
The dangerous half of this gap (an illegal searchQuery field
appearing on a STOCK/WEB sub-question, which would break downstream
JSON.parse()-based consumption) is fixed. The correctness half
(populating years with the stated year, e.g. "Tesla's stock price in 2022" -> years: [2022]) is not — the model currently returns
years: [] in these cases instead.
Latency benchmark reflects unoptimized single-request inference on a
free-tier GPU, not a production serving setup — see §5 caveat.
Benchmark sample sizes (~30/subset) are illustrative; a larger-n
re-run is recommended before quoting these percentages in a more
formal context.
CPT's held-out perplexity evaluation used only 4 documents —
directionally useful, not a precise measurement.
8. Acknowledgments
Built with Unsloth, TRL, and Groq. Base model: Qwen2.5-1.5B-Instruct
(Alibaba Cloud / Qwen team). CPT domain corpus: TheFinAI/SEC_2025.
