[BugFix][SpecDecode] Fix dspark parallel_drafting_token_id init bug (#51602)
Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
This commit is contained in:
co-authored by
Benjamin Chislett
parent
3a749ce816
commit
355a338b8f
@@ -357,6 +357,10 @@ class SpecDecodeBaseProposer:
|
||||
dflash_config = getattr(model_hf_config, "dflash_config", None)
|
||||
if dflash_config and "mask_token_id" in dflash_config:
|
||||
self.parallel_drafting_token_id = dflash_config["mask_token_id"]
|
||||
elif getattr(model_hf_config, "mask_token_id", None) is not None:
|
||||
self.parallel_drafting_token_id = model_hf_config.mask_token_id
|
||||
elif hasattr(model_hf_config, "dspark_noise_token_id"):
|
||||
self.parallel_drafting_token_id = model_hf_config.dspark_noise_token_id
|
||||
elif hasattr(model_hf_config, "pard_token"):
|
||||
self.parallel_drafting_token_id = model_hf_config.pard_token
|
||||
elif hasattr(model_hf_config, "ptd_token_id"):
|
||||
@@ -364,8 +368,9 @@ class SpecDecodeBaseProposer:
|
||||
else:
|
||||
raise ValueError(
|
||||
"For parallel drafting, the draft model config must have "
|
||||
"`pard_token`, `ptd_token_id`, or "
|
||||
"`dflash_config.mask_token_id` specified in its config.json."
|
||||
"`dflash_config.mask_token_id`, `mask_token_id`, "
|
||||
"`dspark_noise_token_id`, `pard_token`, or `ptd_token_id` "
|
||||
"specified in its config.json."
|
||||
)
|
||||
|
||||
if self.pass_hidden_states_to_model:
|
||||
|
||||
Reference in New Issue
Block a user