Skip to content

Commit 2b1510e

Browse files
committed
remove export dependency on transformers
Signed-off-by: h-guo18 <[email protected]>
1 parent 240eac1 commit 2b1510e

File tree

2 files changed

+3
-9
lines changed

2 files changed

+3
-9
lines changed

modelopt/torch/export/plugins/__init__.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,5 +19,5 @@
1919

2020
with import_plugin("megatron_importer"):
2121
from .megatron_importer import *
22-
with import_plugin("transformers"):
23-
from .hf_spec_export import *
22+
23+
from .hf_spec_export import *

modelopt/torch/export/plugins/hf_spec_export.py

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -17,9 +17,6 @@
1717

1818
import torch
1919
import torch.nn as nn
20-
import transformers
21-
22-
from modelopt.torch.speculative.plugins.transformers import HFEagleModel
2320

2421
EAGLE_MODELOPT_TO_OFFICIAL = {
2522
"required": {
@@ -63,7 +60,6 @@ def rename_and_prune_if_spec_decoding(model: nn.Module, post_state_dict: dict):
6360
# if there's other opts, return as is
6461
return post_state_dict
6562

66-
assert isinstance(model, HFEagleModel)
6763
# Check if the state dict keys match
6864
_check_state_dict_keys_match(model.eagle_module, EAGLE_MODELOPT_TO_OFFICIAL["required"])
6965

@@ -90,8 +86,6 @@ def set_config_if_spec_decoding(model: nn.Module, config_data: dict):
9086
# return as is
9187
return config_data
9288

93-
assert isinstance(model, HFEagleModel)
94-
9589
# This is the config keys in official checkpoint.
9690
template_config = {
9791
"architectures": ["LlamaForCausalLMEagle3"],
@@ -110,7 +104,7 @@ def set_config_if_spec_decoding(model: nn.Module, config_data: dict):
110104
"rms_norm_eps": None,
111105
"tie_word_embeddings": False,
112106
"torch_dtype": None,
113-
"transformers_version": transformers.__version__,
107+
"transformers_version": None,
114108
"use_cache": None,
115109
"vocab_size": None,
116110
"draft_vocab_size": None,

0 commit comments

Comments
 (0)