Skip to content

[RFC]: Add DFlash model converter #616

Description

@shanjiaz

Motivation.

The Speculators library supports finetuning pretrained speculator checkpoints via --from-pretrained arg in train.py. However, this only works for checkpoints already in the Speculators format (DFlashSpeculatorConfig + safetensors). There is no way to convert an externally-trained DFlash checkpoint into the Speculators format.

Proposed Change.

Address the TODO in dflash/model,py, "once conversion is added, need to handle the case where a non speculator config is passed in as a kwarg and auto convert." When from_pretrained loads a checkpoint whose config.json is not a SpeculatorModelConfig, instead of raising TypeError, it will attempt auto-conversion:

  1. In SpeculatorModel.from_pretrained(), when the loaded config is not a SpeculatorModelConfig, identify if it's an external DFlash checkpoint (e.g., by inspecting config fields or weight key patterns).
  2. Invoke a DFlashConverter that:
    - Remaps external weight names to the DFlashDraftModel state dict layout (layers, fc, norm, hidden_norm, etc.)
    - Builds a DFlashSpeculatorConfig from the source config.json + verifier model config, including DFlash-specific fields (block_size, max_anchors, mask_token_id, aux_hidden_state_layer_ids, sliding_window_non_causal)
    - Handles t2d/d2t vocab mapping tensors if present
    - Saves the converted checkpoint to a cache/temp directory
  3. Load converted checkpoint — Continue the normal from_pretrained flow with the converted config and weights, returning a ready-to-use DFlashDraftModel.

Any Other Things.

Optionally update the pathway for EAGLE3 as well. We should unify the behavior for both EAGLE3 and DFlash. Users should be able to use --from-pretrained to finetune EAGLE3 models.

Metadata

Metadata

Assignees

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions