Train Your Own Request Router: Fine-Tuning ModernBERT to Pick the Lane in 45 ms

Every request to an AI system starts with a quiet decision: what kind of answer does this need? Words? A picture? A search? A voice message? A video?

In my lab, a 9-billion-parameter model makes that call on every turn: an internal GPTQ build of a Qwen3.5-9B derivative, served by vLLM on an Intel Arc GPU with thinking disabled. It works, but it means a full language-model call before any real work can start. So we trained a small classifier to take over that one decision, running on an NVIDIA Jetson Orin Nano. It is built into our pipeline but currently switched off, pending testing on live traffic; this post covers how it was built and tested.

On a sealed blind test (our second; the first is below), run on the same cases for both routers:

That is one sealed set of 60 cases, each run three times, not a production soak. The classifier is deterministic, so the fair reading is 60 of 60 distinct cases. The 9B's five wrong calls came from three cases: two voice-note requests each went to chat once, and one video request went to image on all three repeats.

The latency rows are measured from a client to each router as deployed, on different hardware: the classifier on the Jetson, the 9B on the Intel Arc server. They compare routing calls in our setup, not model speed on equal hardware. The test's limits are covered at the end.

This post shows how to build one yourself: define the lanes, build the data, train, test it honestly, and put it on edge hardware. We published the model, the training data and the code, all under Apache-2.0:

• Dataset: ericmey/five-lane-router: 2,334 hand-labelled training rows and 19 development rows.
• Model: ericmey/five-lane-router-modernbert-large: the exact weights we tested. The model card carries the full results.
• Code, in the model repo's code/ folder: train_modernbert.py (the exact script that produced the published weights), train_router.py (a portable version used in this post), evaluate_router.py and export_onnx.py.

The sealed test set is not published. It stays private so it can gate future versions; a test set anyone can train on stops being blind.

Try the published model before training anything:

from transformers import pipeline

router = pipeline("text-classification", model="ericmey/five-lane-router-modernbert-large", top_k=None)
router.tokenizer.truncation_side = "left"      # keep the END of long conversations

scores = router("user: can you draw me a lighthouse at dusk?", truncation=True, max_length=512)[0]  # all five lanes, best first
print(scores[0]["label"], scores[0]["score"])   # e.g. image 0.99...

The input is the last few conversation turns as user: / assistant: lines, with the request to route on the last line.

Why a classifier, not a bigger prompt

Routing is a narrow decision. It needs to recognize what the person asked for, quickly and reliably, and then get out of the way. It does not need to write the reply.

A generative model can do it, but you pay for it on every turn: prefill, decoding, JSON parsing, and a timeout budget for when it hangs. An encoder classifier reads the conversation once and emits five numbers. There is no token-by-token decoding and no JSON to parse.

The design we kept: the classifier picks the lane, and a larger model only runs where its richer output is actually needed. Ordinary chat turns skip the big routing call entirely.

Step 1: Write the lanes down before you label anything

This is the step that decides whether the rest works. Our five lanes:

Then write the rules for the ambiguous middle, because that is where every error lives. Ours, in short:

• Act on the best usable reading. "Send me a gif of a cat" → video (create it). If nothing names a medium ("make me one", "send it"), it's chat, where asking "one what?" is a normal reply.
• Read the operation, not the adjective. "Find the real photo from the news" → search. "Make a realistic photo of my cat" → image.
• Talking about media is chat. "Any tips for drawing hands?" is not a drawing request.
• audio means the assistant's own voice. Thunder, a heartbeat or a door knock is not a voice note.

We learned this the slow way. The cost was not compute; it was time and decisions. Our first rubric had a sixth lane, clarify (ask before acting). We ran six careful training rounds against it before anyone confirmed that "ask first" was the behavior I actually wanted. It wasn't. Put your lane rules in front of the person who owns the product before you build a test around them.

Step 2: Build the data (and read every row)

We had no real traffic to label, so we generated candidates with an open-weight LLM and labelled every row by hand. The generator's label was only a suggestion. Three prompt shapes did most of the work:

• Minimal pairs: two requests that differ in exactly one thing ("describe the sunset in words" / "draw the sunset"). These teach the boundary, not the topic.
• Direct asks: one-line requests per lane, in varied tone, length and typos.
• Short conversations: where the last line only makes sense with context ("…yeah, send it").

Each row is one JSON object. text is what the router sees: the last few turns as role: text lines, with the request to route on the last line.

{"id": "t5-000123", "text": "user: hey, what should i wear to that party on saturday? i dont know what to pick.\nassistant: Since it's a party, comfort is key unless it's super formal. Maybe that blue shirt you liked? It looks great on you.\nuser: yeah i guess. show me", "label": "image"}

"Show me" is arguable on its own. Here the conversation has already named something concrete to depict (the blue shirt), so the best usable reading is a picture: image. Without that context, "show me" alone would be chat.

Two habits kept the data honest:

1. Leak-check against your test set before training. Hash every test prompt and refuse any training row that matches exactly. For near matches, we used a conservative difflib similarity threshold of 0.70 for this dataset; it is our choice, not a general standard, so pick one that fits your data. A model that has seen the test cannot be graded on it.
2. Keep a changelog per row. Our labels live in a tab-separated sidecar file (draft file and line, label, note), separate from the generated text. When a rule changed, every affected row's note recorded it, like was clarify; rule: empty request → chat. When a label is questioned later, you can see why it is what it is. The published dataset keeps only id, text and label.

Load ours to start:

from datasets import load_dataset
ds = load_dataset("ericmey/five-lane-router")
print(ds)                  # train: 2334 rows, validation: 19 rows
print(ds["train"][0])

The 19 validation rows are a plumbing and export-parity check (do PyTorch, ONNX and the Jetson engine agree?), not a benchmark. The training recipe was fixed before the run, and no checkpoint was chosen by it.

Step 3: Set up

A CUDA GPU is the comfortable way to train the large model. Apple Silicon (MPS) and CPU work for the base model and for smoke tests. ModernBERT needs a recent Transformers.

python -m venv .venv && source .venv/bin/activate
pip install torch "transformers>=4.48" datasets onnx onnxruntime onnxscript

The published weights were trained with torch 2.14.0 (CUDA 13.0) and transformers 5.17.0. Pin those if you want to match our environment exactly.

Before a long run, do a 64-row smoke test. It catches plumbing mistakes before they cost you a full run:

python train_router.py --out smoke --limit 64 --epochs 1 --max-length 128 --accumulation 1

Step 4: Train

The full script is in the model repo. The parts that matter:

Truncate from the left. The request being routed is the last line. Default right-truncation throws it away on long conversations.

tokenizer.truncation_side = "left"
enc = tokenizer(rows["text"], truncation=True, max_length=512)

Weight the lanes. Our data has about three chat rows for every search row. Inverse-square-root class weights lift the small lanes without letting them dominate:

counts = Counter(ds["train"]["label"])
raw = [1 / math.sqrt(counts[label]) for label in LABELS]
weights = torch.tensor([w / (sum(raw) / len(raw)) for w in raw], device=device)
loss = torch.nn.functional.cross_entropy(logits.float(), labels, weight=weights)

Save a manifest with the model. Base model, label order, max length, truncation side, hyperparameters, row counts and the training file's hash. When a result is questioned, you can prove what produced it.

The published weights came from ModernBERT-large (about 395M parameters) on a single RTX 5090, trained by train_modernbert.py in bf16 mixed precision. The run took about 24 seconds from start to saved weights, including loading the model and tokenizing the data (Python startup not counted). The portable reproduction of that recipe with train_router.py is:

python train_router.py --base answerdotai/ModernBERT-large --out router-large \
  --epochs 3 --batch-size 8 --accumulation 4 --lr 3e-5 --seed 42 --max-length 512

It uses the same data, hyperparameters and seed, but keeps full precision so it also runs on MPS and CPU, so expect close results rather than byte-identical weights. ModernBERT-base is a good first run: faster, smaller, and often enough for a handful of lanes.

Why 512 tokens when ModernBERT handles 8,192? Our inputs are at most eight short turns, and a fixed, small cap keeps latency on the Jetson predictable. Raise it if your inputs are long.

Step 5: Test it like you mean it

This is where most fine-tuning posts stop at "validation accuracy: 97%." Don't. A router's failures are not equal, and one overall number hides the ones your users feel.

Before you train, write down:

• a sealed test set you will never train on, tune on or peek at (we used 60 cases, each run three times);
• the baseline: how many cases the router you are replacing gets right on those exact cases;
• the lines that must hold, and their limits.

Ours were:

1. at least as many correct routes as the 9B on the same cases,
2. zero chat messages sent to an expensive media lane,
3. zero clear media or search requests dropped into chat,
4. every empty request ("make me one") routed to chat,
5. p95 latency under 100 ms.

evaluate_router.py checks rules 1, 2, 3 and 5 and exits non-zero if any fails, so it can gate a CI job. Rule 4 needs your empty-request cases tagged, so check those separately:

python evaluate_router.py --model router-large --test my-sealed-test.jsonl \
  --repeats 3 --baseline-correct 175 --max-p95-ms 100

It prints overall accuracy, a confusion matrix (gold lanes against predicted lanes), latency, and a PASS/FAIL line per rule. The latency it reports is local PyTorch inference on whatever machine runs it, so set --max-p95-ms for the hardware you will actually deploy on. Don't compare it to our 45.6 ms, which is a TensorRT engine on the Jetson measured from a client.

Our first blind test failed. On that first sealed set, the model scored 174/180, beating the 9B's 172, but it failed rule 2: one chat message (someone asking for advice that mentioned drawing) went to image, on all three repeats: one wrong case, counted three times. The overall score looked like a win. The line that mattered said no.

What we did next is the part I'd copy:

1. Accepted the fail. No re-scoring, no "it's only one case."
2. Froze a fresh sealed test before opening the failed case. Only then did we inspect the miss and revise the training data, so the new test could not be built around the observed failure. The person who owns the tests described the failure shape to whoever wrote the new data, rather than passing on the prompt itself.
3. Fixed the data, not the threshold. Working from those descriptions, we added 253 new contrast rows (included in the 2,334) for the failure shapes: talking about pictures vs asking for one, and explicit motion verbs ("animate", "make it flicker") vs stills.
4. Retrained once and scored once on the new test: 180/180, with zero surprise media routes.

If your model is deterministic too, repeats do not add evidence. State your limits in cases, not calls.

Step 6: Put it on edge hardware

Export to ONNX and verify the export matches PyTorch before trusting it. export_onnx.py refuses an export whose prediction differs or whose logits drift more than 0.01:

python export_onnx.py --model router-large --output router-large/model.onnx
# {"max_abs_logit_error": 3.8e-06, "torch_pred": 1, "ort_pred": 1}   (from our smoke model)

On the Jetson, we built a TensorRT FP16 engine with a dynamic sequence length. This is the exact command from our run; the engine is specific to the Jetson's TensorRT and CUDA versions:

/usr/bin/trtexec --onnx=model.onnx --fp16 \
  --minShapes=input_ids:1x8,attention_mask:1x8 \
  --optShapes=input_ids:1x128,attention_mask:1x128 \
  --maxShapes=input_ids:1x512,attention_mask:1x512 \
  --memPoolSize=workspace:2048 --saveEngine=model-fp16.engine --skipInference

We serve the engine behind a tiny HTTP endpoint that takes {"text": ...} and returns the lane and the five probabilities. The 45.6 ms p95 above is measured end to end from a client on the same network, including tokenization on the Jetson and the HTTP round trip, not just the GPU kernel time.

If you don't have a Jetson, the ONNX model on a CPU or any small GPU is the same idea. Measure your own p95.

Step 7: Wire it in carefully

A few things we'd insist on before sending real traffic through a new router:

• Fallback. If the classifier errors, times out or returns a lane you don't recognize, route that turn the old way. Never drop it into chat by default.
• Keep the brief. Our classifier returns a lane only. Downstream stages that need a short description of the request ("a blue mug on a desk") still get it from the larger model, on media turns only.
• Trace what it did. Log the lane, the confidence and the model version on every turn, so a bad route can be found and fed back as data.
• Soak before you switch. A sealed test is evidence for that test. Shadow live traffic before you trust it with users.

Limits

• One sealed set of 60 cases is strong evidence for that set, not a live-traffic error rate. For intuition: if 60 cases were independent random draws from real traffic, zero errors would still be consistent with an error rate of up to about 5%. Ours were written to probe the boundaries, not sampled from traffic, so even that figure is an illustration, not an estimate.
• The data is synthetic, English, and shaped for a companion-style assistant. Your requests will differ. Fork the labels, don't reuse them blindly.
• Lane boundaries are product decisions. Another team could reasonably label "show me my cat" differently.

Build your own

Start with our dataset to get the pipeline working end to end, then replace the lanes with yours: a support-ticket router, a tool picker for an agent, a document-type sorter. The recipe doesn't change:

1. Write the lanes and the ambiguous-case rules down, and get the owner to agree.
2. Generate candidates, and label every row by reading it.
3. Seal a test set and write your pass lines before training.
4. Train small, verify the export, and measure latency where it will run.
5. When the test says no, fix the data and test again on a fresh set.

The question I'd leave you with: which decision does your system make on every request that doesn't need your biggest model?