refactor conversion API

Signed-off-by: Ye Yu <yeyu@nvidia.com>

Author: yeyu-nvidia

modelopt/torch/speculative/eagle/conversion.py

@@ -48,19 +48,7 @@ def convert_to_eagle_model(model: nn.Module, config: EagleConfig) -> ConvertRetu
     config.eagle_architecture_config = {**default_arch_config, **custom_config}
 
     eagle_model = EagleDMRegistry.convert(model)
-    eagle_model.modify(
-        eagle_offline=config.eagle_offline,
-        eagle_hidden_state_distillation=config.eagle_hidden_state_distillation,
-        eagle_self_logit_distillation=config.eagle_self_logit_distillation,
-        eagle_freeze_base_model=config.eagle_freeze_base_model,
-        eagle_report_acc=config.eagle_report_acc,
-        eagle_reuse_base_decoder=config.eagle_reuse_base_decoder,
-        eagle_loss_decay_factor=config.eagle_loss_decay_factor,
-        eagle_architecture_config=config.eagle_architecture_config,
-        eagle_decoder_type=config.eagle_decoder_type,
-        eagle_ttt_steps=config.eagle_ttt_steps,
-        eagle_mix_hidden_states=config.eagle_mix_hidden_states,
-    )
+    eagle_model.modify(config)
 
     # no metadata, all specified via config.
     metadata = {}

modelopt/torch/speculative/eagle/eagle_model.py

@@ -26,26 +26,16 @@ def _setup(self):
 
     def modify(
         self,
-        eagle_offline,
-        eagle_hidden_state_distillation,
-        eagle_self_logit_distillation,
-        eagle_freeze_base_model,
-        eagle_report_acc,
-        eagle_reuse_base_decoder,
-        eagle_loss_decay_factor,
-        eagle_architecture_config,
-        eagle_decoder_type,
-        eagle_ttt_steps,
-        eagle_mix_hidden_states,
+        config,
     ):
         """Base Eagle Model modify function. Child class should implement the details."""
-        self.eagle_offline = eagle_offline
-        self.eagle_hidden_state_distillation = eagle_hidden_state_distillation
-        self.eagle_self_logit_distillation = eagle_self_logit_distillation
-        self.eagle_freeze_base_model = eagle_freeze_base_model
-        self.eagle_report_acc = eagle_report_acc
-        self.eagle_reuse_base_decoder = eagle_reuse_base_decoder
-        self.eagle_loss_decay_factor = eagle_loss_decay_factor
-        self.eagle_decoder_type = eagle_decoder_type
-        self.eagle_ttt_steps = eagle_ttt_steps
-        self.eagle_mix_hidden_states = eagle_mix_hidden_states
+        self.eagle_offline = config.eagle_offline
+        self.eagle_hidden_state_distillation = config.eagle_hidden_state_distillation
+        self.eagle_self_logit_distillation = config.eagle_self_logit_distillation
+        self.eagle_freeze_base_model = config.eagle_freeze_base_model
+        self.eagle_report_acc = config.eagle_report_acc
+        self.eagle_reuse_base_decoder = config.eagle_reuse_base_decoder
+        self.eagle_loss_decay_factor = config.eagle_loss_decay_factor
+        self.eagle_decoder_type = config.eagle_decoder_type
+        self.eagle_ttt_steps = config.eagle_ttt_steps
+        self.eagle_mix_hidden_states = config.eagle_mix_hidden_states

modelopt/torch/speculative/plugins/megatron_eagle.py

@@ -682,17 +682,7 @@ def _setup(self):
 
     def modify(
         self,
-        eagle_offline,
-        eagle_hidden_state_distillation,
-        eagle_self_logit_distillation,
-        eagle_freeze_base_model,
-        eagle_report_acc,
-        eagle_reuse_base_decoder,
-        eagle_loss_decay_factor,
-        eagle_architecture_config,
-        eagle_decoder_type,
-        eagle_ttt_steps,
-        eagle_mix_hidden_states,
+        config,
     ):
         if self.config.pipeline_model_parallel_size > 1:
             warnings.warn(
@@ -705,26 +695,14 @@ def modify(
         if hasattr(self.config, "hetereogenous_dist_checkpoint"):
             self.config.hetereogenous_dist_checkpoint = True
 
-        super().modify(
-            eagle_offline=eagle_offline,
-            eagle_hidden_state_distillation=eagle_hidden_state_distillation,
-            eagle_self_logit_distillation=eagle_self_logit_distillation,
-            eagle_freeze_base_model=eagle_freeze_base_model,
-            eagle_report_acc=eagle_report_acc,
-            eagle_reuse_base_decoder=eagle_reuse_base_decoder,
-            eagle_loss_decay_factor=eagle_loss_decay_factor,
-            eagle_architecture_config=eagle_architecture_config,
-            eagle_decoder_type=eagle_decoder_type,
-            eagle_ttt_steps=eagle_ttt_steps,
-            eagle_mix_hidden_states=eagle_mix_hidden_states,
-        )
+        super().modify(config)
 
         # sequence_parallel is not used in offline eagle
         if self.eagle_offline:
             self.config.sequence_parallel = False
 
         self.eagle_config = dict_to_config(
-            eagle_architecture_config,
+            config.eagle_architecture_config,
             self.config.use_cpu_initialization,
             self.config.fp16,
             self.config.bf16,
@@ -740,7 +718,7 @@ def modify(
         )
 
         if self.eagle_config.draft_vocab_size != self.eagle_config.vocab_size:
-            assert eagle_self_logit_distillation, (
+            assert self.eagle_self_logit_distillation, (
                 "Only logit distillation is supported when draft_vocab_size != vocab_size!"
             )
 

modelopt/torch/speculative/plugins/transformers.py

@@ -532,45 +532,23 @@ def _get_eagle_device(self):
 
     def modify(
         self,
-        eagle_offline,
-        eagle_hidden_state_distillation,
-        eagle_self_logit_distillation,
-        eagle_freeze_base_model,
-        eagle_report_acc,
-        eagle_reuse_base_decoder,
-        eagle_loss_decay_factor,
-        eagle_architecture_config,
-        eagle_decoder_type,
-        eagle_ttt_steps,
-        eagle_mix_hidden_states,
+        config,
     ):
         """Constructor.
 
         Args:
             config: The config for eagle decoder layers.
         """
-        super().modify(
-            eagle_offline=eagle_offline,
-            eagle_hidden_state_distillation=eagle_hidden_state_distillation,
-            eagle_self_logit_distillation=eagle_self_logit_distillation,
-            eagle_freeze_base_model=eagle_freeze_base_model,
-            eagle_report_acc=eagle_report_acc,
-            eagle_reuse_base_decoder=eagle_reuse_base_decoder,
-            eagle_loss_decay_factor=eagle_loss_decay_factor,
-            eagle_architecture_config=eagle_architecture_config,
-            eagle_decoder_type=eagle_decoder_type,
-            eagle_ttt_steps=eagle_ttt_steps,
-            eagle_mix_hidden_states=eagle_mix_hidden_states,
-        )
+        super().modify(config)
 
-        if eagle_decoder_type == "llama":
+        if self.eagle_decoder_type == "llama":
             # Use default eagle config
             decoder_cls = LlamaDecoderLayer
-        elif eagle_decoder_type == "kimik2":
+        elif self.eagle_decoder_type == "kimik2":
             decoder_cls = _setup_kimi_k2_decoder()
 
-        self.eagle_config = PretrainedConfig.from_dict(eagle_architecture_config)
-        self.eagle_config.eagle_decoder_type = eagle_decoder_type
+        self.eagle_config = PretrainedConfig.from_dict(config.eagle_architecture_config)
+        self.eagle_config.eagle_decoder_type = self.eagle_decoder_type
         # Hidden size and vocab size must match base model
         self.eagle_config.hidden_size = self._base_llm_config.hidden_size
         self.eagle_config.vocab_size = self._base_llm_config.vocab_size
@@ -609,14 +587,14 @@ def modify(
         self.eagle_module.to(self._base_model.dtype).to(self._get_eagle_device())
 
         # EAGLE-3 auxiliary hidden_states
-        if (not eagle_offline) and self.eagle_config.use_aux_hidden_state:
+        if (not self.eagle_offline) and self.eagle_config.use_aux_hidden_state:
             self._aux_hidden_states = []
             for layer_idx, layer in enumerate(self._base_model.layers):
                 if layer_idx in self.eagle_config.eagle_aux_hidden_state_layer_ids:
                     layer.register_forward_hook(self._collect_aux_hidden_states_forward_hook)
 
         # delete base model layers for offline training
-        if eagle_offline:
+        if self.eagle_offline:
             self._base_model._modules.pop("layers")
 
         # NOTE: this is a temporary hack to bypass hf trainer check: