refactor(speculative): simplify EAGLE3 trainer and data structure

angelslim/compressor/speculative/train/data/dataset_builder/__init__.py

@@ -13,10 +13,8 @@
 # limitations under the License.
 
 from .dataset_builder_factory import DatasetBuilderFactory
-from .offline_llm_dataset_builder import OfflineLLMDatasetBuilder
-from .offline_vlm_dataset_builder import OfflineVLMDatasetBuilder
-from .online_llm_dataset_builder import OnlineLLMDatasetBuilder
-from .online_vlm_dataset_builder import OnlineVLMDatasetBuilder
+from .offline_dataset_builder import OfflineLLMDatasetBuilder, OfflineVLMDatasetBuilder
+from .online_dataset_builder import OnlineLLMDatasetBuilder, OnlineVLMDatasetBuilder
 
 __all__ = [
     "OnlineLLMDatasetBuilder",

angelslim/compressor/speculative/train/data/dataset_builder/base_dataset_builder.py

@@ -209,20 +209,38 @@ def _process_single_conversation(
                 add_generation_prompt=False,
             )
 
-            # Tokenize conversation
-            encoding = self.tokenizer(
-                conversation,
-                return_offsets_mapping=True,
-                max_length=self.max_length,
-                truncation=True,
-                padding=False,
+            # Check if tokenizer supports offset_mapping
+            is_fast_tokenizer = (
+                hasattr(self.tokenizer, "is_fast") and self.tokenizer.is_fast
             )
 
-            input_ids = encoding.input_ids
-            offsets = encoding.offset_mapping
+            # Tokenize conversation
+            if is_fast_tokenizer:
+                encoding = self.tokenizer(
+                    conversation,
+                    return_offsets_mapping=True,
+                    max_length=self.max_length,
+                    truncation=True,
+                    padding=False,
+                )
+                input_ids = encoding.input_ids
+                offsets = encoding.offset_mapping
+                # Create loss mask for assistant responses
+                loss_mask = self._create_loss_mask_from_offsets(conversation, offsets)
+            else:
+                # For Python tokenizers, use alternative approach
+                encoding = self.tokenizer(
+                    conversation,
+                    max_length=self.max_length,
+                    truncation=True,
+                    padding=False,
+                )
+                input_ids = torch.tensor(encoding.input_ids)
+                # Create loss mask without offsets (alternative implementation needed)
+                loss_mask = self._create_loss_mask_without_offsets(
+                    conversation, input_ids
+                )
 
-            # Create loss mask for assistant responses
-            loss_mask = self._create_loss_mask_from_offsets(conversation, offsets)
             input_ids = torch.tensor(input_ids)
             attention_mask = torch.ones_like(input_ids)
 
@@ -241,6 +259,37 @@ def _process_single_conversation(
             rank0_print(f"Error processing conversation: {e}")
             return None
 
+    def _create_loss_mask_without_offsets(self, conversation, input_ids):
+        # Implement alternative loss mask creation logic for Python tokenizers
+        loss_mask = torch.ones_like(input_ids)
+
+        turns = conversation.split(self.user_header)
+        if len(turns) == 1:
+            # Handle single-turn conversations
+            parts = turns[0].split(self.assistant_header)
+            instruction_part = parts[0] + self.assistant_header
+            instruction_len = len(self.tokenizer(instruction_part).input_ids)
+            loss_mask[:instruction_len] = 0
+        else:
+            # Handle multi-turn conversations
+            cur_len = 0
+            user_header_len = len((self.tokenizer(self.user_header)).input_ids)
+
+            for _, turn in enumerate(turns):
+                parts = turn.split(self.assistant_header)
+                instruction_part = parts[0] + self.assistant_header
+
+                instruction_len = len(self.tokenizer(instruction_part).input_ids)
+                loss_mask[cur_len : cur_len + instruction_len] = 0
+
+                turn_len = len(self.tokenizer(turn).input_ids)
+                cur_len += turn_len
+                cur_len += user_header_len
+
+                loss_mask[cur_len - user_header_len : cur_len] = 0
+
+        return loss_mask
+
     # Copied from https://github.com/NickL77/BaldEagle/blob/master/generate_data/generate_data.py # noqa: E501
     def _create_loss_mask_from_offsets(
         self, conversation: str, offsets: torch.Tensor

angelslim/compressor/speculative/train/data/dataset_builder/offline_llm_dataset_builder.py → angelslim/compressor/speculative/train/data/dataset_builder/offline_dataset_builder.py

@@ -21,7 +21,7 @@
 
 from angelslim.utils import rank0_print
 
-from ..data_utils import DataCollatorWithPadding
+from ..data_utils import DataCollatorWithPadding, VLMDataCollatorWithPadding
 from .base_dataset_builder import DatasetBuilder
 from .dataset_builder_factory import DatasetBuilderFactory
 
@@ -200,6 +200,66 @@ def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
             )
 
 
+class OfflineVLMEagle3Dataset(OfflineEagle3Dataset):
+    def _load_ckpt(self, idx: int) -> Optional[Dict[str, torch.Tensor]]:
+        """
+        Load a checkpoint file.
+
+        Args:
+            idx: Index of the checkpoint file
+
+        Returns:
+            Dictionary containing input_ids, target_hiddens,
+                hidden_states, and loss_mask, or None if loading fails
+        """
+        ckpt_path = self.ckpt_files[idx]
+
+        try:
+            data = torch.load(ckpt_path, map_location="cpu")
+        except Exception as e:
+            warnings.warn(
+                f"Failed to load checkpoint {ckpt_path}: {e}. Skipping this file.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+            return None
+
+        # Validate required keys
+        required_keys = [
+            "input_ids",  # B, N
+            "target_hiddens",  # B, N, D
+            "hidden_states",  # B, N, 3*D
+            "loss_mask",  # B, N
+            "inputs_embeds",  # B, N, D
+            "position_ids",  # 3, B, N
+        ]
+        missing_keys = [key for key in required_keys if key not in data]
+
+        if missing_keys:
+            warnings.warn(
+                f"Checkpoint {ckpt_path} is missing required keys: {missing_keys}. "
+                f"Skipping this file.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+            return None
+
+        # Validate tensor types
+        for key in required_keys:
+            if not isinstance(data[key], torch.Tensor):
+                warnings.warn(
+                    f"Value for key '{key}' in {ckpt_path} is not a torch.Tensor. "
+                    f"Skipping this file.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+                return None
+
+        attention_mask = torch.ones_like(data["input_ids"])
+        data["attention_mask"] = attention_mask  # B, N
+        return data
+
+
 @DatasetBuilderFactory.register("offline", "LLM")
 class OfflineLLMDatasetBuilder(DatasetBuilder):
     def __init__(
@@ -220,3 +280,25 @@ def build_dataset(self, datapath: str, **kwargs: Any) -> Dataset:
 
     def get_data_collator(self) -> Any:
         return DataCollatorWithPadding()
+
+
+@DatasetBuilderFactory.register("offline", "VLM")
+class OfflineVLMDatasetBuilder(DatasetBuilder):
+    def __init__(
+        self, file_pattern: str = "*.ckpt", cache_in_memory: bool = False, **kwargs: Any
+    ):
+        self.file_pattern = file_pattern
+        self.cache_in_memory = cache_in_memory
+
+    def build_dataset(self, datapath: str, **kwargs: Any) -> Dataset:
+        """
+        Create offline datasets from pre-computed .ckpt files.
+        """
+        return OfflineVLMEagle3Dataset(
+            data_dir=datapath,
+            file_pattern=self.file_pattern,
+            cache_in_memory=self.cache_in_memory,
+        )
+
+    def get_data_collator(self) -> Any:
+        return VLMDataCollatorWithPadding()