diff --git a/angelslim/compressor/speculative/train/data/dataset_builder/__init__.py b/angelslim/compressor/speculative/train/data/dataset_builder/__init__.py index 3d47de8c..79cef188 100644 --- a/angelslim/compressor/speculative/train/data/dataset_builder/__init__.py +++ b/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", diff --git a/angelslim/compressor/speculative/train/data/dataset_builder/base_dataset_builder.py b/angelslim/compressor/speculative/train/data/dataset_builder/base_dataset_builder.py index 6b96775e..a2bc4ddb 100644 --- a/angelslim/compressor/speculative/train/data/dataset_builder/base_dataset_builder.py +++ b/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 diff --git a/angelslim/compressor/speculative/train/data/dataset_builder/offline_llm_dataset_builder.py b/angelslim/compressor/speculative/train/data/dataset_builder/offline_dataset_builder.py similarity index 74% rename from angelslim/compressor/speculative/train/data/dataset_builder/offline_llm_dataset_builder.py rename to angelslim/compressor/speculative/train/data/dataset_builder/offline_dataset_builder.py index 4d133c90..f45a8a6a 100644 --- a/angelslim/compressor/speculative/train/data/dataset_builder/offline_llm_dataset_builder.py +++ b/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() diff --git a/angelslim/compressor/speculative/train/data/dataset_builder/offline_vlm_dataset_builder.py b/angelslim/compressor/speculative/train/data/dataset_builder/offline_vlm_dataset_builder.py deleted file mode 100644 index e4d4de50..00000000 --- a/angelslim/compressor/speculative/train/data/dataset_builder/offline_vlm_dataset_builder.py +++ /dev/null @@ -1,221 +0,0 @@ -# Copyright 2025 Tencent Inc. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import warnings -from pathlib import Path -from typing import Any, Dict, List, Optional - -import torch -from torch.utils.data import Dataset - -from angelslim.utils import rank0_print - -from ..data_utils import VLMDataCollatorWithPadding -from .base_dataset_builder import DatasetBuilder -from .dataset_builder_factory import DatasetBuilderFactory - - -class OfflineVLMEagle3Dataset(Dataset): - """ - Offline Dataset for EAGLE3 training. - - Loads pre-computed hidden states, logits, and other data from .ckpt files. - Each .ckpt file contains a dictionary with keys: input_ids, target_logits, - hidden_states, and loss_mask. - """ - - def __init__( - self, data_dir: str, file_pattern: str = "*.ckpt", cache_in_memory: bool = False - ): - """ - Initialize the OfflineVLMEagle3Dataset. - - Args: - data_dir: Directory containing .ckpt files - (will search recursively in subdirectories) - file_pattern: Pattern to match checkpoint files (default: "*.ckpt") - cache_in_memory: Whether to cache all data in memory (default: False) - """ - self.data_dir = Path(data_dir) - self.cache_in_memory = cache_in_memory - - if not self.data_dir.exists(): - raise ValueError(f"Data directory does not exist: {data_dir}") - - # Recursively find all checkpoint files in subdirectories - self.ckpt_files = sorted(list(self.data_dir.rglob(file_pattern))) - - if len(self.ckpt_files) == 0: - raise ValueError( - f"No checkpoint files found in {data_dir} " - f"(including subdirectories) with pattern {file_pattern}" - ) - - rank0_print( - f"Found {len(self.ckpt_files)} checkpoint files " - f"in {data_dir} (including subdirectories)" - ) - - # Track valid indices (files that can be loaded successfully) - self.valid_indices = list(range(len(self.ckpt_files))) - - # Cache data in memory if requested - self.cached_data: Optional[List[Dict[str, torch.Tensor]]] = None - if self.cache_in_memory: - rank0_print("Caching all data in memory...") - self.cached_data = [] - failed_count = 0 - for i in range(len(self.ckpt_files)): - data = self._load_ckpt(i) - if data is not None: - self.cached_data.append(data) - else: - failed_count += 1 - - # Update valid indices based on successful loads - self.valid_indices = list(range(len(self.cached_data))) - - if failed_count > 0: - rank0_print( - f"Data caching completed. " - f"Successfully loaded {len(self.cached_data)} files, " - f"failed to load {failed_count} files" - ) - else: - rank0_print("Data caching completed") - - 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 - - def __len__(self) -> int: - """Return the number of valid samples in the dataset.""" - if self.cached_data is not None: - return len(self.cached_data) - return len(self.valid_indices) - - def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]: - """ - Get a sample from the dataset. - - Args: - idx: Index of the sample - - Returns: - Dictionary containing: - - input_ids: Token IDs (torch.Tensor) - - target_logits: Pre-computed logits from target - model (torch.Tensor) - - hidden_states: Pre-computed hidden states from - target model (torch.Tensor) - - loss_mask: Mask for loss computation (torch.Tensor) - """ - if self.cached_data is not None: - return self.cached_data[idx] - else: - # Try to load the checkpoint, retry with next valid index if fails - max_retries = len(self.valid_indices) - for _attempt in range(max_retries): - actual_idx = self.valid_indices[idx % len(self.valid_indices)] - data = self._load_ckpt(actual_idx) - if data is not None: - return data - - # Remove failed index from valid_indices - self.valid_indices.remove(actual_idx) - if len(self.valid_indices) == 0: - raise RuntimeError( - "All checkpoint files failed to load. " "Cannot continue training." - ) - - # If all retries failed, raise error - raise RuntimeError( - f"Failed to load any valid checkpoint after {max_retries} attempts" - ) - - -@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() diff --git a/angelslim/compressor/speculative/train/data/dataset_builder/online_vlm_dataset_builder.py b/angelslim/compressor/speculative/train/data/dataset_builder/online_dataset_builder.py similarity index 93% rename from angelslim/compressor/speculative/train/data/dataset_builder/online_vlm_dataset_builder.py rename to angelslim/compressor/speculative/train/data/dataset_builder/online_dataset_builder.py index 5a9e9167..f8f9417c 100644 --- a/angelslim/compressor/speculative/train/data/dataset_builder/online_vlm_dataset_builder.py +++ b/angelslim/compressor/speculative/train/data/dataset_builder/online_dataset_builder.py @@ -22,11 +22,34 @@ from angelslim.utils import rank0_print from ..chat_templates import ChatTemplateType -from ..data_utils import VLMDataCollatorWithPadding +from ..data_utils import DataCollatorWithPadding, VLMDataCollatorWithPadding from .base_dataset_builder import OnlineDatasetBuilder from .dataset_builder_factory import DatasetBuilderFactory +@DatasetBuilderFactory.register("online", "LLM") +class OnlineLLMDatasetBuilder(OnlineDatasetBuilder): + def __init__( + self, + tokenizer: Union[AutoTokenizer, AutoProcessor], + max_length: int = 2048, + shuffle_seed: int = 42, + chat_template_type: ChatTemplateType = ChatTemplateType.QWEN3, + display: bool = False, + **kwargs: Any, + ): + super().__init__( + tokenizer, + max_length, + shuffle_seed, + chat_template_type, + display, + ) + + def get_data_collator(self) -> Any: + return DataCollatorWithPadding() + + @DatasetBuilderFactory.register("online", "VLM") class OnlineVLMDatasetBuilder(OnlineDatasetBuilder): def __init__( diff --git a/angelslim/compressor/speculative/train/data/dataset_builder/online_llm_dataset_builder.py b/angelslim/compressor/speculative/train/data/dataset_builder/online_llm_dataset_builder.py deleted file mode 100644 index cd9ecde7..00000000 --- a/angelslim/compressor/speculative/train/data/dataset_builder/online_llm_dataset_builder.py +++ /dev/null @@ -1,170 +0,0 @@ -# Copyright 2025 Tencent Inc. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from typing import Any, Dict, List, Optional, Union - -import torch -from transformers import AutoProcessor, AutoTokenizer - -from angelslim.utils import rank0_print - -from ..chat_templates import ChatTemplateType -from ..data_utils import DataCollatorWithPadding -from .base_dataset_builder import OnlineDatasetBuilder -from .dataset_builder_factory import DatasetBuilderFactory - - -@DatasetBuilderFactory.register("online", "LLM") -class OnlineLLMDatasetBuilder(OnlineDatasetBuilder): - def __init__( - self, - tokenizer: Union[AutoTokenizer, AutoProcessor], - max_length: int = 2048, - shuffle_seed: int = 42, - chat_template_type: ChatTemplateType = ChatTemplateType.QWEN3, - display: bool = False, - **kwargs: Any, - ): - super().__init__( - tokenizer, - max_length, - shuffle_seed, - chat_template_type, - display, - ) - - def get_data_collator(self) -> Any: - return DataCollatorWithPadding() - - def _preprocess_function(self, examples: Dict[str, List]) -> Dict[str, List]: - new_examples = {"input_ids": [], "attention_mask": [], "loss_mask": []} - - for i in range(len(examples["id"])): - try: - processed_example = self._process_single_conversation( - examples["conversations"][i] - ) - - if processed_example is not None: - for key, value in processed_example.items(): - if key in new_examples: - new_examples[key].append(value) - - except Exception as e: - rank0_print(f"Error processing example: {e}") - # Add None placeholders to maintain batch consistency - for key in new_examples: - new_examples[key].append(None) - - return new_examples - - def _process_single_conversation( - self, conversation_data: List[Dict] - ) -> Optional[Dict]: - if not conversation_data or not isinstance(conversation_data, list): - return None - - try: - # Build messages with system prompt - messages = self._build_messages(conversation_data) - if not messages: - return None - - # Apply chat template - conversation = self.tokenizer.apply_chat_template( - messages, - tokenize=False, - add_generation_prompt=False, - ) - - # Check if tokenizer supports offset_mapping - is_fast_tokenizer = ( - hasattr(self.tokenizer, "is_fast") and self.tokenizer.is_fast - ) - - # 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 - ) - - input_ids = torch.tensor(input_ids) - attention_mask = torch.ones_like(input_ids) - - # Visualize loss mask if display mode is enabled - if self.display and self.display_count == 0: - self._visualize_loss_mask(input_ids, loss_mask, conversation) - self.display_count += 1 - - return { - "input_ids": input_ids[None, :], - "attention_mask": attention_mask[None, :], - "loss_mask": loss_mask[None, :], - } - - except Exception as e: - 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 diff --git a/angelslim/compressor/speculative/train/trainer/__init__.py b/angelslim/compressor/speculative/train/trainer/__init__.py index 41085649..b7012159 100644 --- a/angelslim/compressor/speculative/train/trainer/__init__.py +++ b/angelslim/compressor/speculative/train/trainer/__init__.py @@ -12,10 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .offline_eagle3_trainer import OfflineEagle3Trainer -from .offline_vlm_eagle3_trainer import OfflineVLMEagle3Trainer -from .online_eagle3_trainer import OnlineEagle3Trainer -from .online_vlm_eagle3_trainer import OnlineVLMEagle3Trainer +from .offline_eagle3_trainer import OfflineEagle3Trainer, OfflineVLMEagle3Trainer +from .online_eagle3_trainer import OnlineEagle3Trainer, OnlineVLMEagle3Trainer from .trainer_factory import Eagle3TrainerFactory __all__ = [ diff --git a/angelslim/compressor/speculative/train/trainer/eagle3_trainer.py b/angelslim/compressor/speculative/train/trainer/eagle3_trainer.py index 08c29402..64b34ba4 100644 --- a/angelslim/compressor/speculative/train/trainer/eagle3_trainer.py +++ b/angelslim/compressor/speculative/train/trainer/eagle3_trainer.py @@ -73,7 +73,7 @@ def compute_loss( data_for_draft_model = self.prepare_data_for_draft_model(inputs) attention_mask = data_for_draft_model["attention_mask"] # Batch x Seq - position_ids = data_for_draft_model["position_ids"] # Batch x Seq + position_ids = data_for_draft_model["position_ids"] input_ids = data_for_draft_model["input_ids"] # Batch x Seq target_logits = data_for_draft_model["target_logits"] # Batch x Seq x Vocab loss_mask = data_for_draft_model["loss_mask"] # Batch x Seq x 1 @@ -90,6 +90,7 @@ def compute_loss( position_ids, target_logits, loss_mask, + log_prefix="train", ) return loss @@ -290,3 +291,39 @@ def _save_zero3_model(self, output_dir: str, _internal_call: bool = False): # Wait for all processes self.accelerator.wait_for_everyone() + + def prediction_step( + self, + model: nn.Module, + inputs: Dict[str, torch.Tensor], + prediction_loss_only: bool, + ignore_keys: Optional[List[str]] = None, + ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + """ + Perform an evaluation step on `model` using `inputs`. + """ + data_for_draft_model = self.prepare_data_for_draft_model(**inputs) + + attention_mask = data_for_draft_model["attention_mask"] + # inputs_embeds = data_for_draft_model["inputs_embeds"] + position_ids = data_for_draft_model.get("position_ids", None) + input_ids = data_for_draft_model["input_ids"] + target_logits = data_for_draft_model["target_logits"] + loss_mask = data_for_draft_model["loss_mask"] + hidden_states = data_for_draft_model["hidden_states"] + + with torch.no_grad(): + hidden_states = self.down_project_hidden_states(hidden_states) + attention_mask, position_ids = self.prepare_attention_mask_and_position_ids( + hidden_states, attention_mask, position_ids + ) + loss = self.draft_model_training_time_test( + input_ids, + hidden_states, + attention_mask, + position_ids, + target_logits, + loss_mask, + log_prefix="eval", + ) + return (loss, None, None) diff --git a/angelslim/compressor/speculative/train/trainer/offline_eagle3_trainer.py b/angelslim/compressor/speculative/train/trainer/offline_eagle3_trainer.py index 6fb0e653..ad4235f2 100644 --- a/angelslim/compressor/speculative/train/trainer/offline_eagle3_trainer.py +++ b/angelslim/compressor/speculative/train/trainer/offline_eagle3_trainer.py @@ -98,3 +98,85 @@ def prepare_data_for_draft_model( outputs["input_ids"] = input_ids return outputs + + +@Eagle3TrainerFactory.register("offline", "VLM") +class OfflineVLMEagle3Trainer(Eagle3Trainer): + """ + Offline EAGLE3 Trainer for speculative decoding training. + + Uses pre-computed hidden states and logits from offline processing, + avoiding the need for online target model inference. + """ + + def __init__( + self, draft_model: nn.Module, target_head: nn.Module, length: int, **kwargs + ): + """ + Initialize the OnlineEagle3Trainer. + + Args: + draft_model: Draft model for token prediction + length: Number of speculative decoding steps + **kwargs: Additional arguments passed to parent Trainer + """ + super().__init__(draft_model=draft_model, length=length, **kwargs) + self.target_head = target_head + + def prepare_data_for_draft_model( + self, inputs: Dict[str, torch.Tensor] + ) -> Dict[str, torch.Tensor]: + """ + Prepare data for draft model training from offline-generated inputs. + + Args: + inputs: Dictionary containing: + - input_ids: Token IDs + - target_hiddens: Pre-computed last hidden states from target model + - hidden_states: Pre-computed aux hidden states from target model + - attention_mask: Attention mask + - loss_mask: Mask for loss computation + - inputs_embeds: Input embeddings (text and visual) + - position_ids (optional): Position IDs (3D for VLMs mrope) + + Returns: + Dictionary with prepared data for draft model training + """ + inputs_fields = [ + "input_ids", + "target_hiddens", + "hidden_states", + "attention_mask", + "loss_mask", + "inputs_embeds", + "position_ids", + ] + output_fields = [ + "input_ids", + "target_logits", + "hidden_states", + "attention_mask", + "loss_mask", + "inputs_embeds", + "position_ids", + ] + + target_logits = self.target_head(inputs["target_hiddens"]) + loss_mask = inputs["loss_mask"] + input_ids = inputs["input_ids"] + # inputs_embeds = inputs.get("inputs_embeds", None) + position_ids = inputs.get("position_ids", None) + + # Apply right padding and move tensors to correct device + target_logits = padding(target_logits, left=False).to(input_ids.device) + input_ids = padding(input_ids, left=False) + loss_mask = loss_mask[..., None].to(input_ids.device) + + outputs = {k: inputs[k] for k in inputs_fields if k in output_fields} + outputs["target_logits"] = target_logits + outputs["loss_mask"] = loss_mask + outputs["input_ids"] = input_ids + # outputs["inputs_embeds"] = inputs_embeds + outputs["position_ids"] = position_ids + + return outputs diff --git a/angelslim/compressor/speculative/train/trainer/offline_vlm_eagle3_trainer.py b/angelslim/compressor/speculative/train/trainer/offline_vlm_eagle3_trainer.py deleted file mode 100644 index f2ff71a5..00000000 --- a/angelslim/compressor/speculative/train/trainer/offline_vlm_eagle3_trainer.py +++ /dev/null @@ -1,104 +0,0 @@ -# Copyright 2025 Tencent Inc. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from typing import Dict - -import torch -from torch import nn - -from ...utils import padding -from .eagle3_trainer import Eagle3Trainer -from .trainer_factory import Eagle3TrainerFactory - - -@Eagle3TrainerFactory.register("offline", "VLM") -class OfflineVLMEagle3Trainer(Eagle3Trainer): - """ - Offline EAGLE3 Trainer for speculative decoding training. - - Uses pre-computed hidden states and logits from offline processing, - avoiding the need for online target model inference. - """ - - def __init__( - self, draft_model: nn.Module, target_head: nn.Module, length: int, **kwargs - ): - """ - Initialize the OnlineEagle3Trainer. - - Args: - draft_model: Draft model for token prediction - length: Number of speculative decoding steps - **kwargs: Additional arguments passed to parent Trainer - """ - super().__init__(draft_model=draft_model, length=length, **kwargs) - self.target_head = target_head - - def prepare_data_for_draft_model( - self, inputs: Dict[str, torch.Tensor] - ) -> Dict[str, torch.Tensor]: - """ - Prepare data for draft model training from offline-generated inputs. - - Args: - inputs: Dictionary containing: - - input_ids: Token IDs - - target_hiddens: Pre-computed last hidden states from target model - - hidden_states: Pre-computed aux hidden states from target model - - attention_mask: Attention mask - - loss_mask: Mask for loss computation - - inputs_embeds: Input embeddings (text and visual) - - position_ids (optional): Position IDs (3D for VLMs mrope) - - Returns: - Dictionary with prepared data for draft model training - """ - inputs_fields = [ - "input_ids", - "target_hiddens", - "hidden_states", - "attention_mask", - "loss_mask", - "inputs_embeds", - "position_ids", - ] - output_fields = [ - "input_ids", - "target_logits", - "hidden_states", - "attention_mask", - "loss_mask", - "inputs_embeds", - "position_ids", - ] - - target_logits = self.target_head(inputs["target_hiddens"]) - loss_mask = inputs["loss_mask"] - input_ids = inputs["input_ids"] - inputs_embeds = inputs.get("inputs_embeds", None) - position_ids = inputs.get("position_ids", None) - - # Apply right padding and move tensors to correct device - target_logits = padding(target_logits, left=False).to(input_ids.device) - input_ids = padding(input_ids, left=False) - loss_mask = loss_mask[..., None].to(input_ids.device) - - outputs = {k: inputs[k] for k in inputs_fields if k in output_fields} - outputs["target_logits"] = target_logits - outputs["loss_mask"] = loss_mask - outputs["input_ids"] = input_ids - outputs["inputs_embeds"] = inputs_embeds - outputs["position_ids"] = position_ids - - return outputs diff --git a/angelslim/compressor/speculative/train/trainer/online_eagle3_trainer.py b/angelslim/compressor/speculative/train/trainer/online_eagle3_trainer.py index 049186a4..4185a285 100644 --- a/angelslim/compressor/speculative/train/trainer/online_eagle3_trainer.py +++ b/angelslim/compressor/speculative/train/trainer/online_eagle3_trainer.py @@ -81,3 +81,63 @@ def prepare_data_for_draft_model(self, inputs): "position_ids": position_ids, "attention_mask": attention_mask, } + + +@Eagle3TrainerFactory.register("online", "VLM") +class OnlineVLMEagle3Trainer(Eagle3Trainer): + """ + Online EAGLE3 Trainer for speculative decoding training. + Implements training logic for EAGLE3 model using a draft model to predict + tokens based on hidden states from a target model. + """ + + def __init__( + self, + draft_model: nn.Module, + target_model: nn.Module, + length: int, + draft_model_config: Dict[str, Any], + **kwargs, + ): + """ + Initialize the OnlineEagle3Trainer. + Args: + draft_model: Draft model for token prediction + target_model: Target model for generating hidden states + length: Number of speculative decoding steps + draft_model_config: Configuration dictionary for draft model + **kwargs: Additional arguments passed to parent Trainer + """ + super().__init__(draft_model=draft_model, length=length, **kwargs) + self.target_model = target_model + self._aux_hidden_states_layer_ids = getattr( + draft_model_config, "aux_hidden_states_layer_ids", None + ) + + def prepare_data_for_draft_model(self, inputs): + input_ids = inputs["input_ids"] + attention_mask = inputs["attention_mask"] + loss_mask = inputs["loss_mask"] + + # Get hidden states and logits from target model + hidden_states, target_logits, _, position_ids = ( + self.target_model.get_hidden_states_and_logits( + input_ids=input_ids, + attention_mask=attention_mask, + aux_hidden_states_layer_ids=self._aux_hidden_states_layer_ids, + ) + ) + + # Apply right padding and move tensors to correct device + target_logits = padding(target_logits, left=False).to(input_ids.device) + input_ids = padding(input_ids, left=False) + loss_mask = loss_mask[..., None].to(input_ids.device) + + return { + "hidden_states": hidden_states, + "target_logits": target_logits, + "input_ids": input_ids, + "loss_mask": loss_mask, + "position_ids": position_ids, + "attention_mask": attention_mask, + } diff --git a/angelslim/compressor/speculative/train/trainer/online_vlm_eagle3_trainer.py b/angelslim/compressor/speculative/train/trainer/online_vlm_eagle3_trainer.py deleted file mode 100644 index 503e7c76..00000000 --- a/angelslim/compressor/speculative/train/trainer/online_vlm_eagle3_trainer.py +++ /dev/null @@ -1,253 +0,0 @@ -# Copyright 2025 Tencent Inc. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from typing import Any, Dict, List, Optional, Tuple - -import torch -from torch import nn - -from ...utils import padding -from .eagle3_trainer import Eagle3Trainer -from .trainer_factory import Eagle3TrainerFactory - - -@Eagle3TrainerFactory.register("online", "VLM") -class OnlineVLMEagle3Trainer(Eagle3Trainer): - """ - Online EAGLE3 Trainer for speculative decoding training. - Implements training logic for EAGLE3 model using a draft model to predict - tokens based on hidden states from a target model. - """ - - def __init__( - self, - draft_model: nn.Module, - target_model: nn.Module, - length: int, - draft_model_config: Dict[str, Any], - **kwargs, - ): - """ - Initialize the OnlineEagle3Trainer. - Args: - draft_model: Draft model for token prediction - target_model: Target model for generating hidden states - length: Number of speculative decoding steps - draft_model_config: Configuration dictionary for draft model - **kwargs: Additional arguments passed to parent Trainer - """ - super().__init__(draft_model=draft_model, length=length, **kwargs) - self.target_model = target_model - self._aux_hidden_states_layer_ids = getattr( - draft_model_config, "aux_hidden_states_layer_ids", None - ) - - def prepare_data_for_draft_model( - self, input_ids, attention_mask, loss_mask, **kwargs - ): - # Get hidden states and logits from target model - hidden_states, target_logits, inputs_embeds, position_ids = ( - self.target_model.get_hidden_states_and_logits( - input_ids=input_ids, - attention_mask=attention_mask, - aux_hidden_states_layer_ids=self._aux_hidden_states_layer_ids, - **kwargs, - ) - ) - - # Apply right padding and move tensors to correct device - target_logits = padding(target_logits, left=False).to(input_ids.device) - input_ids = padding(input_ids, left=False) - loss_mask = loss_mask[..., None].to(input_ids.device) - - result_dict = {} - result_dict.update(kwargs) - result_dict.update( - { - "hidden_states": hidden_states, - "target_logits": target_logits, - "input_ids": input_ids, - "inputs_embeds": inputs_embeds, - "position_ids": position_ids, - "loss_mask": loss_mask, - "attention_mask": attention_mask, - } - ) - return result_dict - - def compute_loss( - self, - model: nn.Module, - inputs: Dict[str, torch.Tensor], - num_items_in_batch: Optional[int] = None, - return_outputs: bool = False, - ) -> Tuple[List[torch.Tensor], List, List[float]]: - """ - Compute the training loss for the model. - Args: - model: The model for which to compute the loss - inputs: Input data dictionary with input_ids, attention_mask, - loss_mask, position_ids - num_items_in_batch: Number of items in batch (unused) - return_outputs: Whether to return model outputs (unused) - Returns: - Tuple of (prediction_losses, value_losses, accuracies) for each step - """ - data_for_draft_model = self.prepare_data_for_draft_model(**inputs) - - attention_mask = data_for_draft_model["attention_mask"] - inputs_embeds = data_for_draft_model["inputs_embeds"] - position_ids = data_for_draft_model.get("position_ids", None) - input_ids = data_for_draft_model["input_ids"] - target_logits = data_for_draft_model["target_logits"] - loss_mask = data_for_draft_model["loss_mask"] - hidden_states = data_for_draft_model["hidden_states"] - - hidden_states = self.down_project_hidden_states(hidden_states) - attention_mask, position_ids = self.prepare_attention_mask_and_position_ids( - hidden_states, attention_mask, position_ids - ) - loss = self.draft_model_training_time_test( - input_ids, - hidden_states, - attention_mask, - position_ids, - target_logits, - loss_mask, - inputs_embeds, - ) - - return loss - - def draft_model_training_time_test( - self, - input_ids, - hidden_states, - attention_mask, - position_ids, - target_logits, - loss_mask, - inputs_embeds, - log_prefix="", - ): - _, seq_length, _ = hidden_states.shape - - # Initialize containers for losses, accuracies and cache - plosses, acces = [], [] - cache_hidden = [[], []] - - # Iterative speculative decoding training loop - for idx in range(self.length): - # Get input embeddings with gradient tracking - inputs_embeds = self.draft_model.get_input_embeddings(input_ids) - if not inputs_embeds.requires_grad: - inputs_embeds.requires_grad = True - - # Encode through draft model layers - hidden_states, cache_hidden = self.draft_model.encode_layers( - inputs_embeds=inputs_embeds, - hidden_states=hidden_states, - cache_hidden=cache_hidden, - attention_mask=attention_mask, - position_ids=position_ids, - use_cache=True, - ) - - # Compute logits from hidden states - logits = self.draft_model.compute_logits(hidden_states) - - # Compute target distribution and position mask - with torch.no_grad(): - target_max_token = target_logits.argmax(-1) - target_mask = self.draft_model.t2d[target_max_token][..., None].int() - position_mask = target_mask * loss_mask - - target_head = target_logits[..., self.draft_model.t2d].float() - target_p = nn.Softmax(dim=2)(target_head).detach() - - # Compute loss - out_logp = nn.LogSoftmax(dim=2)(logits) - loss = -torch.sum(position_mask * target_p * out_logp, dim=2).mean() - - # Compute accuracy - with torch.no_grad(): - correct = ( - logits.argmax(-1) == target_p.argmax(-1) - ) * position_mask.squeeze(-1) - accuracy = correct.sum().item() / (loss_mask.sum().item() + 1e-6) - - # Store loss and accuracy - plosses.append(loss) - acces.append(accuracy) - - # Update inputs for next iteration (skip on last step) - if idx < self.length - 1: - input_ids = padding(input_ids, left=False) - target_logits = padding(target_logits, left=False) - loss_mask = padding(loss_mask, left=False) - - # Compute weighted loss - ploss_weight = [0.8**i for i in range(len(plosses))] - ploss = sum([ploss_weight[i] * plosses[i] for i in range(len(plosses))]) - - log = { - f"{log_prefix}acc_{i}": round(float(acces[i]), 3) for i in range(len(acces)) - } - log.update( - { - f"{log_prefix}ploss_{i}": round(float(plosses[i].item()), 3) - for i in range(len(plosses)) - } - ) - self.log(log) - - # Return loss - return ploss - - def prediction_step( - self, - model: nn.Module, - inputs: Dict[str, torch.Tensor], - prediction_loss_only: bool, - ignore_keys: Optional[List[str]] = None, - ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: - """ - Perform an evaluation step on `model` using `inputs`. - """ - data_for_draft_model = self.prepare_data_for_draft_model(**inputs) - - attention_mask = data_for_draft_model["attention_mask"] - inputs_embeds = data_for_draft_model["inputs_embeds"] - position_ids = data_for_draft_model.get("position_ids", None) - input_ids = data_for_draft_model["input_ids"] - target_logits = data_for_draft_model["target_logits"] - loss_mask = data_for_draft_model["loss_mask"] - hidden_states = data_for_draft_model["hidden_states"] - - with torch.no_grad(): - hidden_states = self.down_project_hidden_states(hidden_states) - attention_mask, position_ids = self.prepare_attention_mask_and_position_ids( - hidden_states, attention_mask, position_ids - ) - loss = self.draft_model_training_time_test( - input_ids, - hidden_states, - attention_mask, - position_ids, - target_logits, - loss_mask, - inputs_embeds, - log_prefix="eval_", - ) - return (loss, None, None)