Skip to content

bug: require dflash_config.target_layer_ids instead of guessing it - #157

Open
aqilmarwan wants to merge 1 commit into
z-lab:mainfrom
aqilmarwan:fix/require-target-layer-ids
Open

bug: require dflash_config.target_layer_ids instead of guessing it#157
aqilmarwan wants to merge 1 commit into
z-lab:mainfrom
aqilmarwan:fix/require-target-layer-ids

Conversation

@aqilmarwan

@aqilmarwan aqilmarwan commented Aug 17, 2026

Copy link
Copy Markdown

Makes dflash_config.target_layer_ids required instead of guessing it, and deletes build_target_layer_ids. Rebased onto 07ebd93. Details and evidence in #156.

self.target_layer_ids = _draft_value(config, "target_layer_ids")
if self.target_layer_ids is None:
    raise ValueError("Draft config must define dflash_config.target_layer_ids.")

Why:

  • The guess never runs. All 20 checkpoints I could download set the key.
  • If it did run it would usually be wrong. It passes the draft's layer count where the formula needs the number of target layers read, which matches only 6 of 20 checkpoints (18 of 20 with the right count). A wrong guess sets fc's input width, so it surfaces as a weight shape mismatch rather than a missing config key.
  • Because dict.get evaluates its default eagerly, the guess ran on every load, making the unread num_target_layers mandatory in every config. That's why the MLX DFlashConfig field goes too.

MLX already required the key with no fallback; torch now matches.

Verified by building the draft model from all 20 reachable configs, including both DFlash 2 ones, and checking target_layer_ids, fc.in_features, block_size and mask_token_id against each. All load, same as main. No released checkpoint changes behaviour.

@aqilmarwan
aqilmarwan force-pushed the fix/require-target-layer-ids branch from f35babf to 8f38fe9 Compare August 17, 2026 05:13
@aqilmarwan aqilmarwan changed the title Require dflash_config.target_layer_ids instead of guessing it Error: Require dflash_config.target_layer_ids instead of guessing it Aug 17, 2026
build_target_layer_ids was only ever the default argument of
dflash_config.get("target_layer_ids", ...), and every released
z-lab/*-DFlash checkpoint sets that key, so it never supplied a value.

Had it fired it would have been wrong. It derived the tap count from the
draft's num_hidden_layers, but the number of target taps is an independent
hyperparameter: the two differ in 13 of the 18 released checkpoints, and
the formula only reproduces the shipped ids when given len(target_layer_ids)
rather than the draft depth. Since len(target_layer_ids) sets fc's input
width, a wrong guess surfaces as a checkpoint shape mismatch that points at
the weights instead of the missing config key.

Because dict.get evaluates its default eagerly, the dead call also ran on
every load, making num_target_layers a hard requirement of every draft
config even though nothing reads it. A config that sets target_layer_ids
but omits num_target_layers previously raised AttributeError.

Drop the helper, require the key with an error that names it, and remove
the unused num_target_layers field from the MLX DraftConfig. The MLX
loader already read target_layer_ids with no fallback; the torch path now
matches.

Refs: z-lab#156
@aqilmarwan
aqilmarwan force-pushed the fix/require-target-layer-ids branch from 8f38fe9 to fc7f36b Compare August 28, 2026 04:22
@aqilmarwan aqilmarwan changed the title Error: Require dflash_config.target_layer_ids instead of guessing it Require dflash_config.target_layer_ids instead of guessing it Aug 28, 2026
@harshal-96

Copy link
Copy Markdown

I hit this same helper while porting DFlash to LMDeploy and reviewing the TensorRT-LLM and LMDeploy implementations, so I read this change closely.

Verified against the evidence in #156 and my own ports: every published checkpoint sets dflash_config.target_layer_ids, the taps are a training-time choice baked into fc's input width, and DFlash 2 (Qwen3.8-27B) proves matching counts no longer imply matching ids, so no load-time formula can be correct. Requiring the key and turning a shape mismatch into a named config error is the right call, and dropping the unused num_target_layers field from the MLX config matches.

One data point on ecosystem impact: this helper was copied into downstream engines. SGLang's srt/speculative/dflash_utils.py still carries it as a live fallback in resolve_target_layer_ids, where the count-collision case (right tap count, wrong layers) would load cleanly and silently read the wrong target layers. I am filing that separately against SGLang, referencing this issue. LMDeploy's in-flight DFlash PR carries the same fallback and I will flag it there too. Landing this in the reference repo gives those fixes something to point at.

LGTM.

@aqilmarwan aqilmarwan changed the title Require dflash_config.target_layer_ids instead of guessing it bug: require dflash_config.target_layer_ids instead of guessing it Sep 2, 2026
@aqilmarwan

Copy link
Copy Markdown
Author

gentle ping - @jianc99

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants