inplement mix hidden_states for eagle3; deprecate eagle1 (#946)

## What does this PR do?

new feature

**Overview:** 
Enable mix hidden_states in eagle3 training. Deprecate eagle1

## Usage
Add --mix_hidden_states True to launch_train.sh

```python
# Add a code snippet demonstrating how to use this
```

## Testing
<!-- Mention how have you tested your change if applicable. -->

## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->

- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->

## Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added --mix_hidden_states option to enable optional hidden-state
mixing during training.
* Added eagle_ttt_steps setting to control speculative multi-step
iterations.

* **Chores**
* Consolidated speculative decoding to EAGLE3 only; legacy Medusa/EAGLE1
paths removed.
* Unified configuration handling so models and plugins accept a single
config object.

* **Tests**
* Updated and expanded tests for hidden-state mixing and EAGLE3-only
scenarios.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

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

Author: yeyu-nvidia

examples/speculative_decoding/launch_train.sh

@@ -110,6 +110,10 @@ while [ $# -gt 0 ]; do
       if [[ "$1" != *=* ]]; then shift; fi
       HEAD_NODE_IP="${1#*=}"
       ;;
+    --mix_hidden_states*)
+      if [[ "$1" != *=* ]]; then shift; fi
+      MIX_HIDDEN_STATES="${1#*=}"
+      ;;
     *)
       >&2 printf "Error: Invalid argument ${1#*=}\n"
       exit 1
@@ -149,6 +153,7 @@ CP_SIZE=${CP_SIZE:-1}
 DP_SHARD_SIZE=${DP_SHARD_SIZE:-$((TOTAL_GPU/CP_SIZE))}
 LOG_STEPS=${LOG_STEPS:-100}
 DRAFT_VOCAB_CACHE=${DRAFT_VOCAB_CACHE:-""}
+MIX_HIDDEN_STATES=${MIX_HIDDEN_STATES:-"False"}
 
 
 if [[ "$MODE" == "eagle3" ]]; then
@@ -234,6 +239,7 @@ CMD="accelerate launch $MULTI_NODE_ARGS --mixed_precision bf16 ${SCRIPT_DIR}/mai
     --disable_tqdm $DISABLE_TQDM \
     --estimate_ar $ESTIMATE_AR \
     --ar_validate_steps $AR_VALIDATE_STEPS \
+    --mix_hidden_states $MIX_HIDDEN_STATES \
     $DRAFT_VOCAB_CACHE_ARGS \
     $VLM_ARGS \
     $OFFLINE_TRAINING_ARGS \

examples/speculative_decoding/main.py

@@ -43,7 +43,6 @@
     make_eagle_supervised_data_module,
     patch_ring_attention_for_ttt,
 )
-from medusa_utils import make_medusa_supervised_data_module
 from transformers.trainer_utils import get_last_checkpoint
 
 import modelopt.torch.opt as mto
@@ -127,6 +126,10 @@ class EagleArguments:
         default="llama",
         metadata={"help": "The class of eagle decoder to use. Available options: llama, kimik2"},
     )
+    mix_hidden_states: bool = field(
+        default=False,
+        metadata={"help": "Whether to mix hidden states from previous TTT step."},
+    )
 
 
 def train():
@@ -204,6 +207,7 @@ def train():
             config = {
                 "eagle_decoder_type": eagle_args.eagle_decoder_type,
                 "eagle_offline": use_offline_training,
+                "eagle_mix_hidden_states": eagle_args.mix_hidden_states,
                 "eagle_architecture_config": custom_config,
             }
 
@@ -221,9 +225,7 @@ def train():
             raise Exception(f"{training_args.mode} is not supported!")
 
     print_rank_0("Loading dataset...")
-    if training_args.mode == "medusa":
-        data_module = make_medusa_supervised_data_module(tokenizer, data_args)
-    elif training_args.mode == "eagle3":
+    if training_args.mode == "eagle3":
         data_module = make_eagle_supervised_data_module(
             tokenizer, data_args, train_len=training_args.training_seq_len
         )

modelopt/torch/speculative/config.py

@@ -29,12 +29,6 @@
 eagle3_default_config.update({"use_aux_hidden_state": True, "use_last_layernorm": True})
 eagle_mtp_default_config.update({"use_last_layernorm": True, "use_mtp_layernorm": True})
 
-EAGLE1_DEFAULT_CFG = {
-    "algorithm": "eagle",
-    "config": {
-        "eagle_architecture_config": deepcopy(default_eagle_config),
-    },
-}
 
 EAGLE3_DEFAULT_CFG = {
     "algorithm": "eagle",
@@ -105,3 +99,14 @@ class EagleConfig(ModeloptBaseConfig):
         default="llama",
         description=("The class of eagle decoder to use. Available options: llama, kimik2"),
     )
+
+    eagle_ttt_steps: int = ModeloptField(
+        default=4, description=("The number of train-time-test steps in training.")
+    )
+
+    eagle_mix_hidden_states: bool = ModeloptField(
+        default=False,
+        description=(
+            "Whether to mix hidden states of multiple TTT steps. It is a technique to reduce training cost."
+        ),
+    )

modelopt/torch/speculative/eagle/conversion.py

@@ -48,17 +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_model.modify(config)
 
     # no metadata, all specified via config.
     metadata = {}

modelopt/torch/speculative/eagle/eagle_model.py

@@ -26,22 +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,
+        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_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