The Silent Shape Mismatch: How a Hardcoded Dimension Invalidated a Week of DFlash Drafter Evaluation

In the high-stakes world of speculative decoding research, evaluation is the compass that guides every decision. When that compass is broken, you don't just lose your way—you may not even realize you're lost. This is the story of a single message in an opencode coding session ([msg 9106]) where an assistant discovered that its entire evaluation pipeline for a DFlash drafter model had been silently producing garbage, because a single hardcoded dimension in the evaluation harness didn't match the checkpoint being loaded.

The Context: A Drafter Training Investigation

The message sits at a critical juncture in a multi-day investigation. The assistant has been training a DFlash drafter—a small "draft" model that predicts multiple tokens in parallel to accelerate inference of a large target model (Qwen3.6-27B). After multiple training runs (v3, v4), the drafter's performance has been plateauing far below the z-lab reference model. The z-lab model achieves DDTree-8 τ values of 8–12, while the assistant's best model struggles at τ≈3.24 at step 5000.

The preceding messages show a systematic investigation. The assistant has:

The Discovery: What the Eval Log Revealed

The subject message begins with the assistant's reasoning, which immediately identifies two critical issues spotted in the evaluation output:

1. fc.weight shape mismatch: The eval DFlashDrafterEval uses NUM_AUX_LAYERS * H = 4 * 5120 = 20480 as the fc input dim. But v4 checkpoint has fc.weight of [5120, 25600]. The eval harness's default fc_input_layers=4 doesn't match v4's fc_input_layers=5!

This is the moment of discovery. The assistant sees the log line:

skip: fc.weight: shape mismatch torch.Size([5120, 20480]) vs torch.Size([5120, 25600])

PyTorch's load_state_dict function is lenient: when it encounters a shape mismatch between the checkpoint's state dict and the model's state dict, it simply skips loading that parameter and leaves the model's weight randomly initialized. No error is raised. No warning is emphasized. The parameter is silently discarded.

This means the evaluation harness ran the entire inference pipeline—loading hidden states, running the drafter through 15 autoregressive blocks, computing acceptance rates—using a randomly initialized fc projection layer. The fc layer is the critical component that projects the concatenated hidden states from multiple target model layers down to the drafter's hidden dimension. Without the correct weights, the drafter receives meaningless input at every single step.

The second issue is a minor structural mismatch:

2. verifier_norm.weight: not in eval model — this is fine for inference but the v4 checkpoint includes it and the eval model doesn't have it.

The v4 training run included a verifier head (used during training to compute a separate verification loss), but the evaluation model doesn't include this component. The verifier_norm.weight parameter is present in the checkpoint but absent in the eval model, which is expected and harmless—it simply gets skipped during loading.

Why This Matters: The Silent Failure Mode

This discovery is devastating because of its invisibility. The evaluation harness produced output that looked like real results: it printed acceptance rates, streak lengths, and DDTree scores. The numbers were plausible (around τ=2–3, similar to previous runs). Nothing crashed. No assertion failed. The only clue was a single line buried in the output: "skip: fc.weight: shape mismatch."

The assistant had been comparing these eval results against the z-lab reference model, trying to understand why the drafter was underperforming. They had hypothesized about hidden state alignment issues, loss function mismatches, and training data problems. But the real answer was simpler and more insidious: the evaluation was never testing the trained model at all. Every conclusion drawn from those eval runs was potentially invalid.

This is a classic failure mode in ML engineering: the "silent skip" during checkpoint loading. PyTorch's load_state_dict with strict=False (or even strict=True in some configurations) can silently skip mismatched keys, leaving the model partially initialized. The developer sees the model load "successfully" and proceeds to evaluation, unaware that critical weights were never transferred.

The Reasoning Process: From Symptom to Root Cause

The assistant's reasoning shows a tight diagnostic loop. The evaluation output (visible in [msg 9105]) shows the harness starting up, loading cached hidden states, and beginning inference. The assistant must have seen the "skip: fc.weight: shape mismatch" line in the full output and immediately recognized its significance.

The reasoning connects two facts:

  1. The eval harness hardcodes NUM_AUX_LAYERS * H = 4 * 5120 = 20480 as the fc input dimension
  2. The v4 checkpoint was trained with 5 target layers (5 * 5120 = 25600) This architectural mismatch exists because the v4 training run changed the fc layer to use 5 target layers (matching the z-lab architecture), but the evaluation harness was never updated to reflect this change. The eval code still assumed the old 4-layer architecture. The assistant's diagnosis is precise: "The eval DFlashDrafterEval uses NUM_AUX_LAYERS * H = 4 * 5120 = 20480 as the fc input dim. But v4 checkpoint has fc.weight of [5120, 25600]." This shows a clear understanding of both the codebase and the checkpoint structure.

The Fix: Auto-Detection Instead of Hardcoding

The assistant's solution is elegant and prevents this class of bug from recurring:

Let me fix the eval to automatically detect the fc dimension from the checkpoint instead of hardcoding it.

Instead of manually updating the hardcoded value to 5 (which would break again if the architecture changed in v5), the assistant modifies the evaluation harness to inspect the checkpoint's fc.weight shape and derive the fc input dimension dynamically. This makes the eval code robust to architectural changes between training runs.

The edit is applied immediately: "[edit] /data/dflash/scripts/eval_drafter.py. Edit applied successfully." The fix is concise but significant—it changes the eval from a brittle, assumption-laden design to one that adapts to whatever checkpoint it receives.

The Deeper Implications

This message reveals several important lessons about the practice of ML research engineering:

First, evaluation infrastructure is as important as training infrastructure. The assistant had invested enormous effort in building the training pipeline, debugging loss functions, and tuning hyperparameters. But the evaluation harness—the tool that measures whether those efforts are working—had a silent bug that invalidated every measurement.

Second, silent failures are the most dangerous. The shape mismatch didn't crash the program. It produced plausible-looking numbers. Without carefully inspecting the loading logs, the assistant would have continued drawing incorrect conclusions about model performance.

Third, architectural consistency across training and evaluation is non-trivial. When the training code changes the model architecture (adding a layer to the fc projection), every downstream component that interacts with that checkpoint must be updated. This includes the evaluation harness, any visualization tools, and any downstream consumers of the model's outputs.

Fourth, the "skip" behavior in PyTorch's state dict loading is a design trade-off that can mislead. While skipping mismatched keys enables useful patterns like partial checkpoint loading or model surgery, it also creates a failure mode where developers unknowingly evaluate with random weights. A more defensive approach—logging every skipped key prominently, or requiring explicit opt-in for non-strict loading—would catch these issues earlier.

What This Enabled

By fixing the eval harness, the assistant could finally get accurate measurements of the v4 checkpoint's true performance. This directly enabled the subsequent discoveries detailed in the next chunk ([chunk 52.1]): the noise corrupting target logits, the fc shortcut including the target layer, and the loss function mismatch. Without this fix, those investigations would have been comparing apples to oranges—evaluating with random weights against a reference model with correct weights.

The message also demonstrates a key cognitive skill: the ability to recognize when a measurement tool is broken rather than assuming the model is underperforming. The assistant could have continued investigating training hyperparameters, data quality, or architectural choices. Instead, they questioned the evaluation itself, tracing the symptom (bad eval numbers) through the loading process to find the root cause.

Conclusion

Message [msg 9106] captures a pivotal debugging moment in a complex ML research project. It shows how a single hardcoded constant—NUM_AUX_LAYERS * H = 4 * 5120—can silently invalidate an entire evaluation pipeline, and how systematic reasoning can trace a performance symptom back to its root cause in the checkpoint loading code. The fix—auto-detecting the fc dimension from the checkpoint—not only resolves the immediate bug but makes the evaluation infrastructure more robust against future architectural changes. In the broader narrative of this coding session, this message marks the transition from investigating phantom performance issues to discovering the real bugs that were limiting the DFlash drafter's training.