File tree Expand file tree Collapse file tree 3 files changed +70
-446
lines changed
modelopt/torch/speculative Expand file tree Collapse file tree 3 files changed +70
-446
lines changed Original file line number Diff line number Diff line change 24
24
from ..config import EagleConfig
25
25
26
26
EagleDMRegistry = _DMRegistryCls (prefix = "Eagle" ) # global instance for the registry
27
- OfflineEagleDMRegistry = _DMRegistryCls (prefix = "DetachedEagle" ) # global instance for the registry
28
27
29
28
30
29
def convert_to_eagle_model (model : nn .Module , config : EagleConfig ) -> ConvertReturnType :
31
30
"""Convert the model to a eagle model as per `config`."""
32
31
# initialize the true module if necessary
33
32
model = model .init_modellike () if isinstance (model , ModelLikeModule ) else model
34
33
35
- registry = OfflineEagleDMRegistry if config .eagle_offline else EagleDMRegistry
36
-
37
34
original_cls = type (model )
38
- if original_cls not in registry :
39
- for cls in registry ._registry :
35
+ if original_cls not in EagleDMRegistry :
36
+ for cls in EagleDMRegistry ._registry :
40
37
if issubclass (original_cls , cls ):
41
- registry .register ({original_cls : "base_model_class" })(registry [cls ])
38
+ EagleDMRegistry .register ({original_cls : "base_model_class" })(EagleDMRegistry [cls ])
42
39
break
43
40
44
- eagle_model = registry .convert (model )
41
+ eagle_model = EagleDMRegistry .convert (model )
45
42
eagle_model .modify (
46
43
eagle_offline = config .eagle_offline ,
47
44
eagle_hidden_state_distillation = config .eagle_hidden_state_distillation ,
You can’t perform that action at this time.
0 commit comments