bug: require dflash_config.target_layer_ids instead of guessing it - #157
bug: require dflash_config.target_layer_ids instead of guessing it#157aqilmarwan wants to merge 1 commit into
Conversation
f35babf to
8f38fe9
Compare
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
8f38fe9 to
fc7f36b
Compare
|
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 One data point on ecosystem impact: this helper was copied into downstream engines. SGLang's LGTM. |
|
gentle ping - @jianc99 |
Makes
dflash_config.target_layer_idsrequired instead of guessing it, and deletesbuild_target_layer_ids. Rebased onto 07ebd93. Details and evidence in #156.Why:
fc's input width, so it surfaces as a weight shape mismatch rather than a missing config key.dict.getevaluates its default eagerly, the guess ran on every load, making the unreadnum_target_layersmandatory in every config. That's why the MLXDFlashConfigfield 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_sizeandmask_token_idagainst each. All load, same asmain. No released checkpoint changes behaviour.