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 2424from ..config import EagleConfig
2525
2626EagleDMRegistry = _DMRegistryCls (prefix = "Eagle" ) # global instance for the registry
27- OfflineEagleDMRegistry = _DMRegistryCls (prefix = "DetachedEagle" ) # global instance for the registry
2827
2928
3029def convert_to_eagle_model (model : nn .Module , config : EagleConfig ) -> ConvertReturnType :
3130 """Convert the model to a eagle model as per `config`."""
3231 # initialize the true module if necessary
3332 model = model .init_modellike () if isinstance (model , ModelLikeModule ) else model
3433
35- registry = OfflineEagleDMRegistry if config .eagle_offline else EagleDMRegistry
36-
3734 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 :
4037 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 ])
4239 break
4340
44- eagle_model = registry .convert (model )
41+ eagle_model = EagleDMRegistry .convert (model )
4542 eagle_model .modify (
4643 eagle_offline = config .eagle_offline ,
4744 eagle_hidden_state_distillation = config .eagle_hidden_state_distillation ,
You can’t perform that action at this time.
0 commit comments