[Feature] support offline eagle3 training (#142)

Author: liusong1222

angelslim/compressor/speculative/__init__.py

@@ -17,7 +17,9 @@
     DataCollatorWithPadding,
     DatasetManager,
     DraftModelConfig,
+    OfflineEagle3Trainer,
     OnlineEagle3Trainer,
+    TargetHead,
     convert_sharegpt_data,
     convert_ultrachat_data,
     create_draft_model,
@@ -34,10 +36,12 @@
     "DraftModelConfig",
     "create_target_model",
     "OnlineEagle3Trainer",
+    "OfflineEagle3Trainer",
     "data_generation_work_flow",
     "DataCollatorWithPadding",
     "convert_sharegpt_data",
     "convert_ultrachat_data",
     "DatasetManager",
     "get_supported_chat_template_type_strings",
+    "TargetHead",
 ]

angelslim/compressor/speculative/train/__init__.py

@@ -6,18 +6,25 @@
     data_generation_work_flow,
     get_supported_chat_template_type_strings,
 )
-from .models import DraftModelConfig, create_draft_model, create_target_model
-from .trainer import OnlineEagle3Trainer
+from .models import (
+    DraftModelConfig,
+    TargetHead,
+    create_draft_model,
+    create_target_model,
+)
+from .trainer import OfflineEagle3Trainer, OnlineEagle3Trainer
 
 __all__ = [
     "create_draft_model",
     "DraftModelConfig",
     "create_target_model",
     "OnlineEagle3Trainer",
+    "OfflineEagle3Trainer",
     "data_generation_work_flow",
     "DataCollatorWithPadding",
     "convert_sharegpt_data",
     "convert_ultrachat_data",
     "DatasetManager",
     "get_supported_chat_template_type_strings",
+    "TargetHead",
 ]

angelslim/compressor/speculative/train/data/data_utils.py

@@ -133,5 +133,24 @@ def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
             "input_ids": batch_input_ids,
             "attention_mask": batch_attention_mask,
             "loss_mask": batch_loss_mask,
+            "hidden_states": None,
+            "target_hiddens": None,
         }
+
+        # Check if both hidden_states and target_hiddens exist in all features
+        if all(
+            "hidden_states" in item and "target_hiddens" in item for item in features
+        ):
+            batch["hidden_states"] = torch.cat(
+                [
+                    self.paddingtensor(item["hidden_states"], max_length)
+                    for item in features
+                ]
+            )
+            batch["target_hiddens"] = torch.cat(
+                [
+                    self.paddingtensor(item["target_hiddens"], max_length)
+                    for item in features
+                ]
+            )
         return batch