[Bug] The trained dflash has an extremely low acceptance rate

[baifanxxx]

Checklist

• 1. I have searched related issues but cannot get the expected help.
• 2. The bug has not been fixed in the latest version.
• 3. Please note that if the bug-related issue you submitted lacks corresponding environment info and a minimal reproducible demo, it will be challenging for us to reproduce and resolve the issue, reducing the likelihood of receiving feedback.
• 4. If the issue you raised is not a bug but a question, please raise a discussion at https://github.com/sgl-project/SpecForge/discussions/new/choose Otherwise, it will be closed.
• 5. Please use English, otherwise it will be closed.

Describe the bug

I cloned the latest SpecForge codebase and noticed that it now supports training with DFlash. Based on this, I launched a DFlash training job using the script below.

During training, both loss and accuracy behaved normally. On the evaluation set, there was only mild overfitting, which did not seem significant. However, when I loaded the trained weights into the official DFlash benchmark script on gsm8k dataset , (https://github.com/z-lab/dflash/blob/main/run_benchmark.sh), I observed an acceptance rate of only 1.29/(1+3), which is extremely low. This suggests that the training has effectively failed.

I would like to ask whether anyone has successfully trained and inferred a DFlash model. Any discussion or help in locating the root cause would be greatly appreciated.

For context, I have previously successfully trained an Eagle3 model, which indicates that my data preprocessing, training pipeline, and evaluation setup should generally be correct.

Below is the training script I used:

#!/bin/bash

SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
ROOT_DIR=$(dirname $SCRIPT_DIR)
export TORCHINDUCTOR_CACHE_DIR=$ROOT_DIR/cache/compiled_kernels
export SPECFORGE_DATA_NUM_PROC=32

NUM_GPUS=16

ATTENTION_BACKEND=sdpa

torchrun \
    --standalone \
    --nproc_per_node $NUM_GPUS \
    $ROOT_DIR/scripts/train_dflash.py \
    --target-model-path /data/weights/qwen3-8b \
    --draft-config-path $ROOT_DIR/configs/qwen3-8b-dflash.json \
    --train-data-path $ROOT_DIR/cache/dataset/perfectblend_train.jsonl \
    --output-dir $ROOT_DIR/outputs/qwen3-8b-dflash-perfectblend-baseline \
    --num-epochs 15 \
    --batch-size 1 \
    --learning-rate 1e-4 \
    --max-length 2048 \
    --chat-template qwen \
    --attention-backend $ATTENTION_BACKEND \
    --log-interval 100 \
    --eval-interval 5000 \
    --save-interval 10000 \
    --eval-data-path $ROOT_DIR/cache/dataset/opc_test.jsonl \
    --cache-dir $ROOT_DIR/cache \
    --report-to tensorboard \
    --target-model-backend sglang \
    --resume

Looking forward to any insights or suggestions. Thanks!

Best regards,
BAI Fan

Reproduction

Below is the training script I used:

#!/bin/bash

SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
ROOT_DIR=$(dirname $SCRIPT_DIR)
export TORCHINDUCTOR_CACHE_DIR=$ROOT_DIR/cache/compiled_kernels
export SPECFORGE_DATA_NUM_PROC=32

NUM_GPUS=16

ATTENTION_BACKEND=sdpa

torchrun \
    --standalone \
    --nproc_per_node $NUM_GPUS \
    $ROOT_DIR/scripts/train_dflash.py \
    --target-model-path /data/weights/qwen3-8b \
    --draft-config-path $ROOT_DIR/configs/qwen3-8b-dflash.json \
    --train-data-path $ROOT_DIR/cache/dataset/perfectblend_train.jsonl \
    --output-dir $ROOT_DIR/outputs/qwen3-8b-dflash-perfectblend-baseline \
    --num-epochs 15 \
    --batch-size 1 \
    --learning-rate 1e-4 \
    --max-length 2048 \
    --chat-template qwen \
    --attention-backend $ATTENTION_BACKEND \
    --log-interval 100 \
    --eval-interval 5000 \
    --save-interval 10000 \
    --eval-data-path $ROOT_DIR/cache/dataset/opc_test.jsonl \
    --cache-dir $ROOT_DIR/cache \
    --report-to tensorboard \
    --target-model-backend sglang \
    --resume

Environment

SpecForge-main

[Walid-Ahmed]

Which dataset you used as
"perfectblend_train.jsonl" ?

[baifanxxx]

Hi,
This is a sufficiently large mixed dataset that can support multiple tasks. The link is as follows:
https://huggingface.co/datasets/mlabonne/open-perfectblend

[baifanxxx]

@xiaomin-D
I hope you can discuss and resolve this issue with me, and I look forward to your reply.

[xiaomin-D]

@xiaomin-D I hope you can discuss and resolve this issue with me, and I look forward to your reply.

Thanks for your experiments~
Since ZLab hasn’t released the official training details yet, we currently have to reverse-engineer the setup from the inference code, so we can’t fully align with the official pipeline for now. We’ll follow up once their paper or code is public.

For the current situation:

1. Please try aligning the chat template. By default, the dflash model disables , while SpecForge enables it. A mismatch between training and inference chat templates can significantly affect acceptance length. The current SpecForge default keeps mainly for better generality across different models and use cases.

2. On the loss side, different choices can matter. We use CE loss by default, but in our experiments, KL loss — as well as other loss variants — can sometimes give better results depending on the setup.

3. It may also help to further increase the dataset size, and make sure the dataset is regenerated , which tends to improve stability and acceptance length.

Although there is still a gap compared to the official model, with aligned chat templates and KL loss, we can reach around above 3 tokens on average acceptance length, for reference.

You can also try different training techniques, and any contributions would be appreciated.

#427 (comment)

[baifanxxx]

Thank you for your reply. In fact, when I trained dflash using only two repeated data samples, the loss could be overfitted to zero, but I still couldn't achieve a satisfactory acceptance rate during dflash inference. This doesn't seem to be an issue with training techniques or the loss function; there must be some bug causing the misalignment.

[xiaomin-D]

Thank you for your reply. In fact, when I trained dflash using only two repeated data samples, the loss could be overfitted to zero, but I still couldn't achieve a satisfactory acceptance rate during dflash inference. This doesn't seem to be an issue with training techniques or the loss function; there must be some bug causing the misalignment.

There must also be some training techniques that are still not well understood yet, but I think you can try following what I suggested first, because the inference and training templates, data regeneration, and different loss functions all have a significant impact on the acceptance rate.

[xiaomin-D]

Thank you for your reply. In fact, when I trained dflash using only two repeated data samples, the loss could be overfitted to zero, but I still couldn't achieve a satisfactory acceptance rate during dflash inference. This doesn't seem to be an issue with training techniques or the loss function; there must be some bug causing the misalignment.

I don’t think this should be considered a bug unless a concrete implementation issue is found. Given the current differences between training and inference losses and the lack of a reference implementation, this could be an implementation or objective mismatch rather than a correctness issue.

Maybe you can wait for the official release of the training details.