The Silent Shape Mismatch: A Single Edit That Uncovered a Hidden Evaluation Bug

Message 9109: [edit] /data/dflash/scripts/eval_drafter.py — "Edit applied successfully."

At first glance, message 9109 appears to be one of the most mundane entries in the entire conversation: a three-word status report confirming that a file edit was applied. But this terse message — [assistant] [edit] /data/dflash/scripts/eval_drafter.py\nEdit applied successfully. — represents the culmination of a painstaking debugging chain that had been building across dozens of messages. It is the second of two edits to the evaluation harness, and together they fixed a silent, catastrophic bug: the evaluation code had been loading a v4 checkpoint's weights into a model with the wrong architecture, producing garbage outputs that masqueraded as valid evaluation results.

The Context: A 4x Performance Gap

To understand why this edit mattered, we must trace the debugging arc that led to it. In the preceding messages ([msg 9100] through [msg 9108]), the assistant had been building an evaluation infrastructure to compare the DFlash drafter's training progress against the z-lab reference model. The initial comparison was stark: at step 20k, the assistant's model achieved a DDTree-8 acceptance rate (τ) of approximately 3.0 on fresh coding prompts, while the z-lab model achieved τ≈12.4 — a 4x gap. This prompted a deep investigation into what was going wrong.

The investigation had already uncovered several critical issues. In [msg 9101], the assistant discovered that causal-conv1d was not installed on the training machine (CT200), meaning all training had been using PyTorch's fallback implementation for the Qwen3.5 linear attention layers, while evaluation on CT129 used the fla library with causal-conv1d for fast-path execution. The assistant recognized this as a fundamental mismatch: "We were training and evaluating with different hidden states."

Then, in [msg 9105], the assistant ran the evaluation harness on the v4 step 4k checkpoint and discovered something even more alarming. The log showed:

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

The fc.weight tensor — the projection layer that maps hidden states from the target model's layers into the drafter's dimension — had never been loaded. The evaluation harness was hardcoded to expect a 4-layer fc projection (4 × 5120 = 20480 input dimension), but the v4 checkpoint used a 5-layer fc (5 × 5120 = 25600 input dimension). The checkpoint loader silently skipped the mismatched weight, leaving the fc layer with randomly initialized parameters. Every evaluation result produced up to this point was essentially random.

The Two Edits: Auto-Detection and Hidden State Reconstruction

Message 9106 applied the first edit, which added logic to auto-detect the fc dimension from the checkpoint's actual weight shape rather than relying on a hardcoded constant. The assistant's reasoning in that message laid out the diagnosis clearly: "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!"

But the first edit alone was insufficient. As the assistant noted in [msg 9107]: "Also need to handle the hidden states dimension — for 5-layer fc, we need all 5 layers concatenated." The evaluation harness extracts hidden states from the target model at specific layer indices (typically layers 1, 16, 31, 46, and 61 of the 64-layer Qwen3.6-27B model). With a 4-layer fc, it would concatenate only the first four layers' hidden states (dimension 20480). With a 5-layer fc, it needed to concatenate all five (dimension 25600). The second edit — message 9109 — implemented this reconstruction logic.

Why This Bug Was So Dangerous

The shape mismatch was a "silent failure" in the truest sense. PyTorch's load_state_dict function, when called with strict=False (as is common when loading partial checkpoints), simply skips mismatched keys with a log message. The evaluation harness printed this log message, but it was buried among dozens of other loading messages. The assistant initially saw garbled output from the evaluation and attributed it to other causes — the missing causal-conv1d, incorrect context masking, or fundamental architectural problems with the drafter itself.

The danger of silent shape mismatches in ML evaluation pipelines cannot be overstated. When a model's weights fail to load, the resulting parameters are whatever was left from initialization — typically random. The model produces outputs that look like valid predictions (tokens, probabilities, etc.) but are in fact completely unmoored from the training process. A researcher could spend days or weeks debugging training hyperparameters, data quality, or architectural decisions, all while the evaluation harness itself was fundamentally broken.

The Assumptions That Failed

The original evaluation harness made a critical assumption: that the fc projection would always use the same number of layers. This assumption was baked into the code as a constant (NUM_AUX_LAYERS = 4), derived from the DFlash paper's architecture where the fc projection uses N-1 layers (reserving the last layer exclusively for target logit computation). However, the v4 training run had expanded this to all 5 layers — a deliberate architectural choice that the evaluation harness was never updated to reflect.

This reveals a common pattern in ML development: training code and evaluation code evolve on parallel tracks, and architectural changes in one are not always propagated to the other. The assistant's decision to auto-detect the fc dimension from the checkpoint — rather than hardcoding it — was a direct response to this brittleness. It transformed the evaluation harness from a single-architecture tool into one that could adapt to whatever checkpoint it was given.

Input and Output Knowledge

To understand this message, one needs substantial background knowledge: the DFlash architecture's use of fc projections to compress hidden states from multiple target model layers into the drafter's dimension; the concept of "verifier" layers (the last target layer reserved for computing target logits); the mechanics of PyTorch's load_state_dict with strict=False; and the specific layer structure of Qwen3.6-27B (64 layers with 5 extraction points at layers 1, 16, 31, 46, 61).

The output knowledge created by this edit was a corrected evaluation harness that could properly load v4 checkpoints. More importantly, it enabled the assistant to obtain accurate evaluation results, which in turn fed into the next phase of debugging: the discovery of three critical training bugs (noise corrupting target logits, fc including the target layer, and loss function mismatch) that would ultimately lead to the v5 training run.

The Broader Lesson

Message 9109 is a reminder that in complex ML pipelines, the evaluation infrastructure is not a neutral observer — it is an active participant in the research process. A bug in the evaluation harness does not merely produce wrong numbers; it actively misdirects the researcher's attention. The assistant spent dozens of messages investigating hidden state mismatches, attention implementation differences, and training data quality, all while the real problem was a shape mismatch in a single weight tensor. The fix was a single edit, but the debugging chain that led to it — tracing from garbled outputs to the fc.weight shape mismatch to the hardcoded layer count — required deep understanding of the entire system architecture.

This message, for all its brevity, marks the turning point where the evaluation pipeline was finally aligned with the training pipeline, clearing the way for the genuine architectural fixes that followed.