I tried fine-tuning Nemotron 3.5 Lightning with LoRA post-training to make it serve as a judge for an LLM router
This page has been translated by machine translation. View original
Introduction
Hello, I'm Morishige from Classmethod's Manufacturing Business Technology Department.
Some of you may be wondering how to integrate the Nemotron 3.5 Lightning we measured in the first-touch article into actual work. Beyond its impressive specs of a 1M context 30B MoE, this model is strongly designed with the intent of being used as a base for post-training (taking a pre-trained model and applying additional training to tailor it for specific use cases). This time, to verify that claim through actual measurement, I fine-tuned it using LoRA SFT as the "capability judge" in the LLM router I use daily.
The capability judge is a model that reads requests from a coding agent and determines each time whether "this task is sufficient for a cheap model (weak) or requires a high-end model (strong)." I've been running this configuration with a router called NeMo Switchyard, using deepseek-v4-flash-0731 running on 2 DGX Sparks as the judge (the 2-node configuration is covered in the article below, written as of 2026-08). However, this model has 284B parameters. Keeping such a heavyweight requiring 2 machines just for per-request judgments is clearly too resource-intensive for that role. I had long thought about replacing it with a lighter model dedicated to judging.
To state the conclusion upfront: in an evaluation of 87 real routing cases, it surpassed the distillation source flash-0731 (89.7% → 94.3%), and we were able to replace the judgment role that 284B had been handling with a 21GB 30B model. It has also partially gained the ability to route difficult problems that the teacher model was structurally unable to judge. However, along the way, there was a puzzling incident where the trained model went haywire only in the production path. The cause was JSON key order. I believe this story is the most shareable learning from this entire experience.
Since this is a continuation of previous articles, please refer to the first-touch article for Lightning's own characteristics, and the Switchyard article for the router mechanism (both are articles as of 2026-08).
This article introduces the entire flow from training data design, LoRA SFT, NVFP4 quantization, to real-machine routing evaluation, along with the pitfalls encountered along the way. I hope it resonates with people thinking about the next step beyond "measuring" a local LLM — that is, "tailoring" it for business workloads.
Vanilla Lightning Cannot Be Used As-Is as a Judge
First, here is the overall configuration. It's sufficient to just grasp where the judge sits within the router.
The judge returns a 4-field JSON with "the probability that a weak model will complete this task in one attempt (p_solve)" for each request, and the router compares it against a threshold to determine the destination. Since judgment runs on every request, speed, cost, and format robustness are all required — a deceptively demanding role.
What happens if you put vanilla Lightning directly into this role? The results below compare it against the current judge, flash-0731, on a holdout set of 237 items set aside for validation (not used for training).
| Metric | flash-0731 (current) | Vanilla Lightning |
|---|---|---|
| JSON validity check pass | 237/237 | 212/237 |
| Judgment category match | — (this is the ground truth) | 17.3% |
| weak/strong judgment match | — | 72.2% |
What was more serious than the numbers was the content: 234 out of 237 cases were routed to the "unmatched (unable to judge)" category, and for 10 real operational requests, all were sent to strong. If everything flows to the expensive model, the router has no purpose. It can mostly comply with the format, but it's completely missing the vocabulary and probability intuition of a judge. This is where post-training comes in.
Training Data Is Built from Measured Accuracy
Training data refers to pairs of inputs and ideal outputs that you want the model to imitate. This time I prepared 2 streams. The core of the design is to directly use the measured accuracy of the weak model as the p_solve values (labels) on the ideal output side.
When I previously retrained the NVIDIA LLM Router v3, I had a painful experience with bugs and biases in labels judged by an LLM (article from 2026-07). Learning from that, this time I had the weak model actually solve 195 coding benchmark problems twice each, and used the pass frequency as labels. The judgment vocabulary (judgment categories and a single sentence of reasoning) is copied from the current judge flash-0731's judgments as a model (this is what's called distillation), while the critical probabilities are replaced with actual measurements.
| Stream | Count | Source of Labels |
|---|---|---|
| Coding benchmark measured | 195 problems (excluding 1 with broken teacher response: 167 training + 27 validation) | Weak model's measured accuracy + difficulty band smoothing |
| Synthetic tasks | 2,085 problems (1,875 training + 210 validation) | Distilled from flash-0731's judgments |
The synthetic side was generated by prompting 8 archetypes such as "brainstorming consultation," "ambiguous formatting requests," and "image-dependent requests," filling in the bias in judgment category distribution that would result in all cases being "supported" with only the measured data. Generation and distillation calls totaled about 4,000, all fed directly to the current judge flash-0731.
One completed training data item looks like this. The input contains the task instruction and a capability card summarizing the weak model's strengths and weaknesses, and the ideal output provides a 4-field JSON.
{
"crux": "Sort each row descending, then for each column take the maximum value across rows and sum those maxima.",
"primary_rule": "SUP-3",
"capability_boundary": "supported",
"p_solve": 0.88
}
crux is a summary of the solution approach, primary_rule is the rule ID of the judgment basis, capability_boundary is the judgment category, and p_solve is the probability of the weak model completing the task. The model will internalize the judge's vocabulary and probability intuition by imitating thousands of these examples. And the order of these 4 fields becomes the foreshadowing of the incident in the latter half.
Another subtly important point is freezing the evaluation set. The evaluation data, including "discrimination 14 problems" where weak failed twice in a row and strong passed, was excluded from training data without a single item remaining, verified by hash comparison. Since these are the exam questions for judging "did we surpass the distillation source" later, contaminating them would be pointless.
Learning Twice with the Official Cookbook's LoRA SFT
Training was based on the official cookbook from the Nemotron repository (the Megatron-Bridge LoRA recipe), with only the dataset portion replaced for the judge. LoRA is a method that adds only small differential updates for learning instead of rewriting all 30B weights. And SFT refers to supervised learning that directly imitates the example pairs from the previous chapter. This combination allows even large models to be tailored quickly and at low cost. The hardware was 2 H100s rented via NVIDIA Brev, with settings of rank 8 / alpha 32 / 2 epochs / sequence length 3072, taking 46 minutes of training + 2 minutes to merge the differential into the original weights per run. GPU cost comes to around $10 per run.
With the first training run (v0), the formatting issues were fixed in one shot. JSON validity checks went from 212/237 to 237/237, and judgment category matching jumped from 17.3% to 73.4%. However, of the benchmark problems set aside for validation, 11 difficult problems that should have been routed to strong were not routed to strong at all. Digging into the cause, the teacher flash-0731 itself had a blind spot of optimistically judging all problems of this type as "solvable by weak," and the distillation faithfully copied that blind spot. Despite preparing measured labels, the optimistic distillation labels from the synthetic data side won out numerically.
So for the second run (v1), I adjusted the mix. I replicated the benchmark problems with measured labels 6 times and excluded one synthetic category with strong optimistic bias. The results are as follows.
| Metric | Vanilla Lightning | v0 (initial mix) | v1 (adjusted mix) |
|---|---|---|---|
| JSON validity check pass | 212/237 | 237/237 | 237/237 |
| weak/strong judgment match | 72.2% | 78.1% | 80.2% |
| p_solve mean absolute error | 0.225 | 0.151 | 0.139 |
| Difficult validation problems routed to strong | 0/11 | 0/11 | 2/11 |
Just by changing the data mix, the ability to identify difficult problems began to emerge. The training recipe was not touched at all. This was a direct experience of the oft-heard wisdom that the success or failure of post-training depends more on data than on the recipe.
Story of walls encountered in the training environment (2 walls of the nemo container)
Wall 1: The container specified in the cookbook cannot be pulled from outside
The nvcr.io/nvidian/nemo:26.08 specified by the cookbook is an NVIDIA internal registry image and cannot be pulled from outside. I switched to the public nvcr.io/nvidia/nemo:26.06, then checked out 4 training recipe files for Lightning from Megatron-Bridge's unmerged branch and inserted them. Internal storage paths hardcoded in the recipes also need to be rewritten to local checkpoint paths. The lesson is that even with an official cookbook, "it works as-is" is not guaranteed — checking the source of the container and recipe comes first.
Wall 2: API differences between 26.06 and the cookbook
The cookbook's train script assumes the 26.08-series API, with about 3 places where dataset configuration differs in 26.06. I made a small modification to skip preprocessing by using a dummy dataset specification to feed the pre-built training.jsonl.
JSON Key Order Was Causing the Judge to Go Haywire
Here comes the main topic. When v1 was actually run through the real router, about 18% of real operational requests had broken responses, even though direct calls were fine. The failure mode was the same every time: after outputting all fields of the JSON, it would endlessly emit only whitespace characters up to the maximum 4096 tokens. The judgment was treated as a failure, and as a failsafe the request was sent to strong, so the actual harm was quiet — but if one judgment takes 52 seconds, there is no point in moving the judgment role.
It only happens when going through the router. So I diffed the requests arriving at vLLM between router-mediated and direct calls, and the only difference was the key order in the JSON schema within response_format. The router (Rust implementation) sorts object keys alphabetically when re-serializing the configuration. Meanwhile, xgrammar, which handles structured output (the mechanism that forces output to conform to a specified schema JSON) on the vLLM side, enforces the order of properties in the schema as the generation order. And all 2,610 training data items for v1 had label JSON in the exact same key order.
In other words, the model had memorized the judgment JSON as a token sequence "starting with crux and ending with p_solve." When the grammar engine enforces alphabetical order (an arrangement where p_solve appears in the middle), the model loses its known exit and starts emitting whitespace. Reproducing the alphabetical schema directly caused 4 out of 10 runs to go haywire, while the natural order matching training time was clean for 53 consecutive runs.
I implemented countermeasures on both fronts. For the router side, I submitted a PR upstream with a one-line fix to preserve key order. For the model side, I retrained with only the key order of the label JSON in the training data shuffled per sample (v2). This is a data augmentation that changes not a single bit of meaningful content — just thoroughly covering all 24 orderings of the 4 fields. The retraining cost was the same ~$10 per run as v1.
The effect was clear: zero rampages in a probe of 200 calls with 5 variations of key order, and zero rampages in 100 real machine calls directly connected to the router. Judgment quality maintained the same level as v1 (weak/strong match 80.2% → 81.4%).
| Path | v1 | v2 (key order shuffle training) |
|---|---|---|
| Alphabetical schema direct call | 4/10 rampages | 0/40 |
| Real machine calls via router | ~18% rampages | 0/100 |
In fine-tuning that produces JSON output with structured output in mind, key order is not "an implementation detail that should have no meaning" — it is part of the contract. Moreover, OpenAI-compatible proxies and routers often re-serialize JSON, so key order can change at any time. As long as you train with a fixed key order in labels, this landmine can be triggered in any configuration, so if you're training to output JSON, it's safest to include key order shuffling from the start.
Quantizing to NVFP4 and Running on DGX Spark
The trained model is 62GB in BF16. While it fits within DGX Spark's 128GB memory, when judging long-context real operational inputs, the prefill (the pre-processing phase of reading the input) alone takes 6.8 to 89.3 seconds, making it unusable as a per-request judge. So I proceeded to quantization — replacing weights with coarser numerical representations to compress size and computation. Using PTQ (post-training quantization) that doesn't require retraining, I used Lightning's official NVFP4 recipe (4/6 method) included with Model-Optimizer. On a single GB10, it took just under an hour, reducing 62GB to 21GB.
The concern was whether quantization would break the judge's probability intuition, but on the holdout 237 items, the mean absolute error of p_solve was 0.131, which did not worsen from BF16's 0.137, and judgment latency settled at a median of 1.27 seconds. Even for long-context real operational inputs, it was 1.8 to 2.6 seconds. Since the current judge flash-0731 takes just over 2 seconds, entrusting judgment to the 30B model is equally fast or faster.
One point: I hit a snag where vLLM could not load the quantized model. An internal suffix had been mixed into all exported tensor names, and clearing it by rewriting the safetensors header resolved the issue. If you follow the same procedure (Model-Optimizer main branch + hf_ptq.py) and get a load error, suspect the tensor names.
How Far Did It Surpass the Distillation Source flash-0731?
Finally, I compared the three — current judge, vanilla Lightning, and post-training (v2 NVFP4) — under identical conditions. The evaluation consisted of two parts: a probe of 47 labeled problems and a real machine evaluation running 87 items through the actual router. The expected judgments here are not the teacher's output, but ground truth determined independently from actual measurements of whether weak can solve the problem and case design. Since the teacher flash-0731 also takes the same test, we can measure whether the student surpasses the distillation source on this axis.
| Metric | flash-0731 (current) | Vanilla Lightning | Post-training v2-NVFP4 |
|---|---|---|---|
| Match on 47 labeled problems | 38/47 | 32/47 | 45/47 |
| Expected judgment match on 87 real cases | 78/87 (89.7%) | — (not deployable) | 82/87 (94.3%) |
| Difficult 14 problems routed to strong | 0/14 | — | 5/14 |
| Judgment latency (median) | 2.2 sec | — | 1.27 sec (long text 1.8–2.6 sec) |
| Resources for judgment role | 284B · 2 DGX Sparks | — | 21GB · fits on 1 unit |
We confirmed surpassing the distillation source with actual measurements. In particular, the 14 discrimination problems are a group that the teacher flash-0731 is structurally unable to route even one to strong. The effect of mixing measured labels shows here, and I feel the difference between 0 and 5 is larger than it looks. The threshold was also adjusted to 0.80 from 0.75 using the post-training probability distribution (80/87 → 82/87 on 87 real cases).
This judgment role is a component that plugs into the environment our team operates — two open-weight models plus automatic routing (described in the article below). The judgment that required a 284B model on 2 machines now fits on NVFP4's 21GB on a single machine, and the freed-up 2 machines can be dedicated to running inference. The configuration of bringing both judgment and execution closer to local has advanced one more step.
And the first real operational deployment has begun not on the coding agent side, but on the resident agent side. I set up another instance of the NVFP4 judge and Switchyard on DGX Spark and switched the NemoHermes requests that have been running continuously to automatic routing. In the first 16 hours of operation, there were zero rampages, and the judgments correctly split into 92 weak and 24 strong cases. It's subtly satisfying as a demonstration of post-trained component reusability that a judge tailored for coding agents plugs directly into a different harness just by looking at the request content. NemoHermes itself is introduced in the following article (article as of 2026-06).
Two Lineages: Predictive Routers and Generative Judges
The judge this time is a model that generates judgment text, but the NVIDIA LLM Router v3 I tried previously was a predictive type that outputs a score directly from the internal state at the time of reading the input, without generating text. Having retrained both myself, the contrast has become much more three-dimensional, so let me summarize it briefly.
| Aspect | Predictive (LLM Router v3) | Generative (this judge) |
|---|---|---|
| Judgment output | Classification head score | JSON with reasoning |
| Adaptation cost | 480 problems · $11 · 5 min training | 2,610 items · $10 · 50 min training |
| Critical point in label design | Bias in LLM-judged labels | Distillation labels copy the teacher's blind spots |
| Critical point in operation | Retrain every time pool changes | Fixing output contract (down to key order) |
Both are the same in that quality is determined by "the source of labels" and "the inference cost of the judgment model itself" becomes a fixed cost. Personally, the generative type is easier to debug since the reasoning is preserved in JSON, while it carries generation-specific failure modes like rampages — that was my hands-on impression. In the long term, there are suggestions that the judgment role will be folded into the prefill of the weak execution layer rather than remaining as an independent model, and if that happens, these two lineages may converge.
Summary
I took Nemotron 3.5 Lightning, measured in the first-touch article, all the way through training data design → LoRA SFT → NVFP4 quantization → real machine evaluation, and was able to tailor it as a router judge that surpasses the distillation source. This judge is already routing real NemoHermes operational traffic as the judgment role for Switchyard on DGX Spark. The training itself was convenient at around $10 per run and just under an hour based on the official cookbook, and what made the difference was the data. Using measured accuracy as labels, the fact that distillation copies the teacher's blind spots, and the fact that even JSON key order is part of the training data contract — these three are the takeaways from this experience.
Things that are not yet done: the 14 discrimination problems are still halfway at 5/14, and there is one case remaining after quantization where the model deterministically rampages on a specific input (accepted for operations since it falls to the safe side with a failsafe). Judgment category vocabulary matching also dropped slightly as the price of stripping the rote memorization of key order.
Next, I'd like to try QLoRA training on DGX Spark alone using the same dataset. Whether training can be completed entirely on local hardware without renting an H100, and how close the quality can get, is planned to be covered in a follow-up.
Reference Links
- NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 (Hugging Face)
- NVIDIA-NeMo/Nemotron — Repository containing the LoRA SFT usage-cookbook used this time
- NVIDIA-NeMo/Switchyard
- NVIDIA-NeMo/Megatron-Bridge
- NVIDIA/Model-Optimizer — The NVFP4 4/6 quantization recipe is included here
- vLLM v0.27.1 Release (used for serving the post-training model)
