add training_seq_len to eagleconfig; remove deprecated unittest

Signed-off-by: h-guo18 <[email protected]>

Author: h-guo18

examples/speculative_decoding/main.py

@@ -185,8 +185,12 @@ def train():
             }[training_args.mode]["config"]
 
             # overwrite config with custom config
-            if use_offline_training:
-                config["eagle_offline"] = True
+            config.update(
+                {
+                    "eagle_offline": use_offline_training,
+                    "eagle_training_seq_len": training_args.training_seq_len,
+                }
+            )
 
             if eagle_args.eagle_config:
                 with open(eagle_args.eagle_config) as f:
@@ -203,8 +207,6 @@ def train():
                     "draft_vocab_size": custom_config["draft_vocab_size"]
                     if eagle_args.eagle_config and "draft_vocab_size" in custom_config
                     else model.config.vocab_size,
-                    # pass in the seq length for flex attention mask compilation
-                    "training_seq_len": training_args.training_seq_len,
                 }
             )
 

modelopt/torch/speculative/config.py

@@ -95,6 +95,10 @@ class EagleConfig(ModeloptBaseConfig):
         default=0.9, description=("The decay factor for multiple eagle_loss.")
     )
 
+    eagle_training_seq_len: int = ModeloptField(
+        default=1024, description=("The training sequence length.")
+    )
+
     eagle_architecture_config: dict = ModeloptField(
         default={}, description=("The config for eagle module architecture.")
     )

modelopt/torch/speculative/eagle/conversion.py

@@ -47,6 +47,7 @@ def convert_to_eagle_model(model: nn.Module, config: EagleConfig) -> ConvertRetu
         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_training_seq_len=config.eagle_training_seq_len,
         eagle_architecture_config=config.eagle_architecture_config,
     )
 

modelopt/torch/speculative/eagle/eagle_model.py

@@ -35,6 +35,7 @@ def modify(
         eagle_report_acc,
         eagle_reuse_base_decoder,
         eagle_loss_decay_factor,
+        eagle_training_seq_len,
         eagle_architecture_config,
     ):
         """Base Eagle Model modify function. Child class should implement the details."""
@@ -45,7 +46,7 @@ def modify(
         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_training_seq_len = eagle_training_seq_len
         if eagle_architecture_config.get("parallel_draft_step", 1) > 1:
             for i in range(eagle_architecture_config.get("parallel_draft_step") - 1):
                 self.register_buffer(f"mask_token_{i}", torch.tensor(-1))

modelopt/torch/speculative/plugins/megatron_eagle.py

@@ -721,6 +721,7 @@ def modify(
         eagle_report_acc,
         eagle_reuse_base_decoder,
         eagle_loss_decay_factor,
+        eagle_training_seq_len,
         eagle_architecture_config,
     ):
         if self.config.pipeline_model_parallel_size > 1:
@@ -742,6 +743,7 @@ def modify(
             eagle_report_acc=eagle_report_acc,
             eagle_reuse_base_decoder=eagle_reuse_base_decoder,
             eagle_loss_decay_factor=eagle_loss_decay_factor,
+            eagle_training_seq_len=eagle_training_seq_len,
             eagle_architecture_config=eagle_architecture_config,
         )
 

modelopt/torch/speculative/plugins/transformers.py

@@ -382,6 +382,7 @@ def modify(
         eagle_report_acc,
         eagle_reuse_base_decoder,
         eagle_loss_decay_factor,
+        eagle_training_seq_len,
         eagle_architecture_config,
     ):
         """Constructor.
@@ -397,6 +398,7 @@ def modify(
             eagle_report_acc=eagle_report_acc,
             eagle_reuse_base_decoder=eagle_reuse_base_decoder,
             eagle_loss_decay_factor=eagle_loss_decay_factor,
+            eagle_training_seq_len=eagle_training_seq_len,
             eagle_architecture_config=eagle_architecture_config,
         )
         self.eagle_config = PretrainedConfig.from_dict(eagle_architecture_config)
@@ -454,7 +456,7 @@ def modify(
         self.num_ttt_steps = 3  # NOTE: (hg) hardcoded for now. Might add to config later.
         # compile and cach flex attention masks
         self.cached_attn_blk_masks = [
-            self._compile_ttt_block_mask(self.eagle_config.training_seq_len, i)
+            self._compile_ttt_block_mask(eagle_training_seq_len, i)
             for i in range(self.num_ttt_steps)
         ]
 

tests/unit/torch/speculative/plugins/test_hf_speculative.py

@@ -17,7 +17,6 @@
 from copy import deepcopy
 
 import pytest
-import torch
 from _test_utils.torch_model.transformers_models import (
     create_tiny_llama_dir,
     get_tiny_llama,
@@ -69,122 +68,3 @@ def test_eagle_model_convert_save_and_restore(tmp_path, eagle_config):
     model_test = AutoModelForCausalLM.from_pretrained(tmp_path / "modelopt_model")
     assert isinstance(model_test, mtsp.plugins.HFEagleModel)
     tf_modelopt_state_and_output_tester(model_ref, model_test)
-
-
-# fmt: off
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-def test_eagle_model_prepare_eagle_inputs(dtype):
-    dummy_model = get_tiny_llama(num_hidden_layers=4)
-
-    config = EAGLE3_DEFAULT_CFG["config"]
-    config["eagle_architecture_config"].update({
-        "draft_vocab_size": dummy_model.config.vocab_size,
-        "hidden_size": dummy_model.config.hidden_size,
-    })
-    mtsp.convert(dummy_model, mode=[("eagle", config)])
-
-    eagle_input_ids_0 = torch.tensor([[10, 20, 30, 40]], dtype=torch.long)
-    position_ids_0 = torch.tensor([[0, 1, 2, 3]], dtype=torch.long)
-
-
-    #This is concatenated from 3 intermediate base model layers
-    cat_aux_hidden_states = torch.randn(1, 4, 32, dtype=dtype)
-
-    #This is eagle output from previous eagle forward pass
-    dummy_eagle_output_hidden_states = torch.randn(1, 4, 32, dtype=dtype)
-
-    #This is the causal mask for the 0th eagle step
-    m = torch.finfo(dtype).min
-    attention_mask_0 = torch.tensor([[0, m, m, m], #  input tok 10-> predicting token 20
-                                     [0, 0, m, m], #  20 -> 30
-                                     [0, 0, 0, m], #  30 -> 40
-                                     [0, 0, 0, 0]] #  40 -> tok after 40
-
-                                    , dtype=dtype).view(1, 1, 4, 4)
-
-    # 2nd eagle step
-    eagle_input_h_1, eagle_input_ids_1, attention_mask_1, position_ids_1 = dummy_model._concat_eagle_inputs(
-        eagle_input_ids_0,
-        cat_aux_hidden_states,
-        attention_mask_0,
-        position_ids_0,
-        dummy_eagle_output_hidden_states,
-    )
-
-    assert eagle_input_ids_1.equal(torch.tensor([[10, 20, 30, 40, 10, 20, 30, 40]], dtype=torch.long))
-    assert position_ids_1.equal(torch.tensor([[0, 1, 2, 3, 0, 1, 2, 3]], dtype=torch.long))
-
-    assert attention_mask_1.equal(torch.tensor([[0, m, m, m,  m, m, m, m], # (x) output discarded
-                                                [0, 0, m, m,  m, m, m, m], # (x)
-                                                [0, 0, 0, m,  m, m, m, m], # (x)
-                                                [0, 0, 0, 0,  m, m, m, m], # (x)
-
-                                                [m, m, m, m,  m, m, m, m], # (x) input tok 10-> predicting token 20
-                                                [0, m, m, m,  m, 0, m, m], # 20 -> 30
-                                                [0, 0, m, m,  m, m, 0, m], # 30 -> 40
-                                                [0, 0, 0, 0,  m, m, m, m], # (x) 40 -> tok after 40
-                                                ], dtype=dtype).view(1, 1, 8, 8))
-
-    # 3rd eagle step
-    eagle_input_hidden_states_2, eagle_input_ids_2, attention_mask_2, position_ids_2 = dummy_model._concat_eagle_inputs(
-        eagle_input_ids_0,
-        cat_aux_hidden_states,
-        attention_mask_0,
-        position_ids_0,
-        torch.cat([dummy_eagle_output_hidden_states, dummy_eagle_output_hidden_states], dim=1),
-    )
-    assert eagle_input_ids_2.equal(torch.tensor([[10, 20, 30, 40,  10, 20, 30, 40,  10, 20, 30, 40]], dtype=torch.long))
-    assert position_ids_2.equal(torch.tensor([[0, 1, 2, 3,  0, 1, 2, 3,  0, 1, 2, 3]], dtype=torch.long))
-
-    assert attention_mask_2.equal(torch.tensor([[0, m, m, m,  m, m, m, m,  m, m, m, m], # (x)
-                                                [0, 0, m, m,  m, m, m, m,  m, m, m, m], # (x)
-                                                [0, 0, 0, m,  m, m, m, m,  m, m, m, m], # (x)
-                                                [0, 0, 0, 0,  m, m, m, m,  m, m, m, m], # (x)
-
-                                                [m, m, m, m,  m, m, m, m,  m, m, m, m], # (x)
-                                                [0, m, m, m,  m, 0, m, m,  m, m, m, m], # (x)
-                                                [0, 0, m, m,  m, m, 0, m,  m, m, m, m], # (x)
-                                                [0, 0, 0, 0,  m, m, m, m,  m, m, m, m], # (x)
-
-                                                [m, m, m, m,  m, m, m, m,  m, m, m, m], # (x)10 -> 20
-                                                [m, m, m, m,  m, m, m, m,  m, m, m, m], # (x)20 -> 30
-                                                [0, m, m, m,  m, 0, m, m,  m, m, 0, m], # 30 -> 40
-                                                [0, 0, 0, 0,  m, m, m, m,  m, m, m, m], # (x) 40 -> tok after 40
-
-                                                ], dtype=dtype).view(1, 1, 12, 12))
-
-    # 4th eagle step
-    eagle_input_hidden_states_3, eagle_input_ids_3, attention_mask_3, position_ids_3 = dummy_model._concat_eagle_inputs(
-        eagle_input_ids_0,
-        cat_aux_hidden_states,
-        attention_mask_0,
-        position_ids_0,
-        torch.cat([dummy_eagle_output_hidden_states, dummy_eagle_output_hidden_states,
-                   dummy_eagle_output_hidden_states],dim=1),
-    )
-
-    assert eagle_input_ids_3.equal(torch.tensor([[10, 20, 30, 40,  10, 20, 30, 40,
-                                                  10, 20, 30, 40,  10, 20, 30, 40]], dtype=torch.long))
-    assert position_ids_3.equal(torch.tensor([[0, 1, 2, 3,  0, 1, 2, 3,  0, 1, 2, 3,  0, 1, 2, 3]], dtype=torch.long))
-
-    assert attention_mask_3.equal(torch.tensor([[0, m, m, m,  m, m, m, m,  m, m, m, m,  m, m, m, m], # (x)
-                                                [0, 0, m, m,  m, m, m, m,  m, m, m, m,  m, m, m, m], # (x)
-                                                [0, 0, 0, m,  m, m, m, m,  m, m, m, m,  m, m, m, m], # (x)
-                                                [0, 0, 0, 0,  m, m, m, m,  m, m, m, m,  m, m, m, m], # (x)
-
-                                                [m, m, m, m,  m, m, m, m,  m, m, m, m,   m, m, m, m], # (x)
-                                                [0, m, m, m,  m, 0, m, m,  m, m, m, m,   m, m, m, m], # (x)
-                                                [0, 0, m, m,  m, m, 0, m,  m, m, m, m,   m, m, m, m], # (x)
-                                                [0, 0, 0, 0,  m, m, m, m,  m, m, m, m,   m, m, m, m], # (x)
-
-                                                [m, m, m, m,  m, m, m, m,  m, m, m, m,   m, m, m, m], # (x)
-                                                [m, m, m, m,  m, m, m, m,  m, m, m, m,   m, m, m, m], # (x)
-                                                [0, m, m, m,  m, 0, m, m,  m, m, 0, m,   m, m, m, m], # (x)
-                                                [0, 0, 0, 0,  m, m, m, m,  m, m, m, m,   m, m, m, m], # (x)
-
-                                                [m, m, m, m,  m, m, m, m,  m, m, m, m,   m, m, m, m], # (x)10 -> 20
-                                                [m, m, m, m,  m, m, m, m,  m, m, m, m,   m, m, m, m], # (x)20 -> 30
-                                                [m, m, m, m,  m, m, m, m,  m, m, m, m,   m, m, m, m], # (x)
-                                                [0, 0, 0, 0,  m, m, m, m,  m, m, m, m,   m, m, m, m], # (x)
-
-                                                ], dtype=dtype).view(1, 1, 16, 16))