File size: 1,104 Bytes
2dd5f57 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | from speculators.losses import kl_div_loss
from speculators.models.dfly import DFlyDraftModel
from speculators.train.config import TrainConfig
def test_dfly_cli_resolves_algorithm_specific_settings():
cfg = TrainConfig.resolve(
[
"--verifier-name-or-path",
"dummy",
"--speculator-type",
"dfly",
"--no-enable-hidden-correction",
"--hidden-correction-intermediate-size",
"96",
]
)
args = cfg.flatten()
assert args["speculator_type"] == "dfly"
assert args["draft_arch"] == "qwen3"
assert args["enable_hidden_correction"] is False
assert args["hidden_correction_intermediate_size"] == 96
def test_dfly_reuses_dflash_training_loss_arguments():
args = TrainConfig(
speculator_type="dfly",
).flatten()
train_kwargs, validation_kwargs = DFlyDraftModel.get_trainer_kwargs(**args)
assert train_kwargs["loss_config"]["kl_div"][0] is kl_div_loss
assert validation_kwargs["loss_config"]["kl_div"][0] is kl_div_loss
assert train_kwargs["gamma"] == 4.0
|