Aluode/PerceptionLabPortable
0
1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with5# the License. You may obtain a copy of the License at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on10# an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the11# specific language governing permissions and limitations under the License.12 13import logging14from typing import Callable, Optional15 16import torch17 18from ..cache_utils import (19 DynamicCache,20 DynamicLayer,21 DynamicSlidingWindowLayer,22 EncoderDecoderCache,23 StaticCache,24)25from ..generation.configuration_utils import GenerationConfig26from ..masking_utils import (27 ALL_MASK_ATTENTION_FUNCTIONS,28 _ignore_causal_mask_sdpa,29 _is_torch_greater_or_equal_than_2_5,30 prepare_padding_mask,31)32from ..modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel33from ..pytorch_utils import (34 is_torch_greater_or_equal,35 is_torch_greater_or_equal_than_2_3,36 is_torch_greater_or_equal_than_2_6,37)38 39 40class TorchExportableModuleForVLM:41 """42 A wrapper class for exporting Vision-Language Models (VLMs) like SmolVLM2 for ExecuTorch.43 44 This class handles the export of three main components:45 1. Vision encoder (processes images to visual features)46 2. Connector/projector (maps visual features to text embedding space)47 3. Text decoder (generates text from combined visual and text tokens)48 """49 50 def __init__(self, model, max_batch_size: int = 1, max_cache_len: int = 1024):51 """52 Initialize the exportable VLM module.53 54 Args:55 model: The VLM (e.g. SmolVLM) model instance56 max_batch_size: Maximum batch size. Always 1 for ExecuTorch57 max_cache_len: Maximum cache length for text generation58 """59 self.model = model60 self.max_batch_size = max_batch_size61 self.max_cache_len = max_cache_len62 self.config = model.config63 64 # Extract individual components65 self.vision_encoder = model.model.vision_model66 self.connector = model.model.connector67 self.text_decoder = model.model.text_model68 69 # Store exported programs70 self.exported_vision_encoder = None71 self.exported_connector = None72 self.exported_text_decoder = None73 74 def export_vision_encoder(self):75 """Export the vision encoder component."""76 self.vision_encoder.eval()77 78 # Create example input79 pixel_values = torch.randn(1, 3, 384, 384, dtype=torch.float32)80 81 # Define dynamic shapes82 dynamic_shapes = {83 "pixel_values": {84 2: torch.export.Dim.AUTO,85 3: torch.export.Dim.AUTO,86 }87 }88 89 self.exported_vision_encoder = torch.export.export(90 self.vision_encoder,91 args=(pixel_values,),92 dynamic_shapes=dynamic_shapes,93 strict=False,94 )95 96 return self.exported_vision_encoder97 98 def export_connector(self):99 """Export the connector component."""100 self.connector.eval()101 102 # Vision encoder output shape: [batch_size, num_patches, vision_hidden_size]103 vision_hidden_size = self.config.vision_config.hidden_size104 image_size = self.config.vision_config.image_size105 patch_size = self.config.vision_config.patch_size106 patches_per_dim = image_size // patch_size107 num_patches = patches_per_dim * patches_per_dim108 image_hidden_states = torch.randn(1, num_patches, vision_hidden_size, dtype=torch.float32)109 110 # Define dynamic shapes - static batch_size=1, dynamic num_patches111 dynamic_shapes = {"image_hidden_states": {1: torch.export.Dim.AUTO}}112 113 # Export the connector using torch.export114 self.exported_connector = torch.export.export(115 self.connector,116 args=(image_hidden_states,),117 dynamic_shapes=dynamic_shapes,118 strict=False,119 )120 121 return self.exported_connector122 123 def export_text_decoder(self):124 """Export the text decoder component."""125 126 # Create text decoder exportable wrapper127 self.exportable_text_decoder = TorchExportableModuleForDecoderOnlyLM(model=self.text_decoder)128 129 # Use the existing text decoder exportable wrapper130 seq_length = 3131 input_ids = torch.zeros((1, seq_length), dtype=torch.long)132 cache_position = torch.arange(seq_length, dtype=torch.long)133 max_seq_length = min(self.max_cache_len, self.config.text_config.max_position_embeddings)134 seq_len_dim = torch.export.Dim("seq_length_dim", max=max_seq_length - 1)135 136 dynamic_shapes = {137 "input_ids": {1: seq_len_dim},138 "cache_position": {0: seq_len_dim},139 }140 141 self.exported_text_decoder = self.exportable_text_decoder.export(142 input_ids=input_ids,143 cache_position=cache_position,144 dynamic_shapes=dynamic_shapes,145 strict=False,146 )147 148 return self.exported_text_decoder149 150 def export(self, **kwargs):151 """Export all components of the VLM model."""152 self.export_vision_encoder(**kwargs)153 self.export_connector(**kwargs)154 self.export_text_decoder(**kwargs)155 return {156 "vision_encoder": self.exported_vision_encoder,157 "connector": self.exported_connector,158 "text_decoder": self.exported_text_decoder,159 }160 161 def forward(self, pixel_values, input_ids, cache_position):162 """163 Simplified forward pass for inference with guaranteed non-null input_ids and cache_position.164 165 Args:166 pixel_values: Input images [1, channels, height, width] (optional)167 input_ids: Text token IDs [1, seq_len] (required - won't be None)168 cache_position: Cache positions [seq_len] (required - won't be None)169 170 Returns:171 Output with logits for text generation172 """173 pass174 175 def generate(176 self, pixel_values=None, input_ids=None, max_new_tokens=50, do_sample=False, temperature=1.0, **kwargs177 ):178 """179 Simplified generate method with guaranteed non-null input_ids.180 181 Args:182 pixel_values: Input images [1, channels, height, width] (optional)183 input_ids: Initial text tokens [1, seq_len] (required - won't be None)184 max_new_tokens: Maximum number of tokens to generate185 do_sample: Whether to use sampling or greedy decoding186 temperature: Temperature for sampling187 188 Returns:189 Generated sequences190 """191 pass192 193 194class TorchExportableModuleForDecoderOnlyLM(torch.nn.Module):195 """196 A recipe module designed to make a `PreTrainedModel` exportable with `torch.export`,197 specifically for decoder-only LM with cache. This module ensures that the198 exported model is compatible with further lowering and execution in `ExecuTorch`.199 """200 201 def __init__(202 self,203 model: PreTrainedModel,204 batch_size: Optional[int] = None,205 max_cache_len: Optional[int] = None,206 device: Optional[torch.device] = None,207 ) -> None:208 """209 Initializes the exportable module.210 211 Args:212 model (`PreTrainedModel`): The pretrained model to wrap.213 214 Raises:215 ValueError: If the model is configured with a unsupported cache implementation.216 """217 super().__init__()218 219 config = model.config.get_text_config()220 221 if not hasattr(config, "use_cache") or config.use_cache is False:222 raise ValueError("The model must have caching enabled to be performant.")223 224 if hasattr(config, "layer_types") and getattr(config, "sliding_window", None) is not None:225 self.model = TorchExportableModuleWithHybridCache(model, batch_size, max_cache_len, device)226 else:227 # If `layer_types` is not specified explicitly in the config or `sliding_window` is null,228 # there is only 1 type of layers, so export will use `StaticCache` by default.229 logging.info(230 "Using `StaticCache` for export as `layer_types` is not specified or `sliding_window` is `null` in the config."231 )232 self.model = TorchExportableModuleWithStaticCache(model, batch_size, max_cache_len, device)233 # This is the same as sdpa, but mask creation does not use `vmap` which is not exportable234 ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa_without_vmap", sdpa_mask_without_vmap)235 ALL_ATTENTION_FUNCTIONS.register("sdpa_without_vmap", ALL_ATTENTION_FUNCTIONS["sdpa"])236 self.model.model.config._attn_implementation = "sdpa_without_vmap"237 238 def forward(239 self,240 input_ids: Optional[torch.Tensor] = None,241 inputs_embeds: Optional[torch.Tensor] = None,242 cache_position: Optional[torch.Tensor] = None,243 ) -> torch.Tensor:244 """245 Forward pass of the module, which is compatible with the ExecuTorch llm runner.246 247 Args:248 input_ids (`torch.Tensor`): Tensor representing current input token id to the module.249 inputs_embeds (`torch.Tensor`): Tensor representing current input embeddings to the module.250 cache_position (`torch.Tensor`): Tensor representing current input position in the cache.251 252 Returns:253 torch.Tensor: Logits output from the model.254 """255 return self.model.forward(256 input_ids=input_ids,257 inputs_embeds=inputs_embeds,258 cache_position=cache_position,259 )260 261 def export(262 self,263 input_ids: Optional[torch.Tensor] = None,264 inputs_embeds: Optional[torch.Tensor] = None,265 cache_position: Optional[torch.Tensor] = None,266 dynamic_shapes: Optional[dict] = None,267 strict: Optional[bool] = None,268 ) -> torch.export.ExportedProgram:269 """270 Export the wrapped module using `torch.export`.271 272 Args:273 input_ids (`Optional[torch.Tensor]`):274 Tensor representing current input token id to the module. Must specify either this or inputs_embeds.275 inputs_embeds (`Optional[torch.Tensor]`):276 Tensor representing current input embeddings to the module. Must specify either this or input_ids.277 cache_position (`Optional[torch.Tensor]`):278 Tensor representing current input position in the cache. If not provided, a default tensor will be used.279 dynamic_shapes (`Optional[dict]`):280 Dynamic shapes to use for export if specified.281 strict(`Optional[bool]`):282 Flag to instruct `torch.export` to use `torchdynamo`.283 284 Returns:285 torch.export.ExportedProgram: The exported program that can be used for inference.286 287 Examples:288 Export with input_ids:289 ```python290 # Prepare inputs291 input_ids = torch.tensor([[1, 2, 3]], dtype=torch.long, device=model.device)292 cache_position = torch.arange(input_ids.shape[-1], dtype=torch.long, device=model.device)293 294 # Export295 exported = exportable_module.export(296 input_ids=input_ids,297 cache_position=cache_position298 )299 ```300 301 Export with inputs_embeds:302 ```python303 # Prepare embeddings304 inputs_embeds = torch.randn(1, 3, 768, device=model.device) # batch_size=1, seq_len=3, hidden_size=768305 cache_position = torch.arange(inputs_embeds.shape[1], dtype=torch.long, device=model.device)306 307 # Export308 exported = exportable_module.export(309 inputs_embeds=inputs_embeds,310 cache_position=cache_position311 )312 ```313 """314 if not (input_ids is None) ^ (inputs_embeds is None):315 raise ValueError("Need to specify either input_ids or inputs_embeds.")316 317 if hasattr(self.model, "base_model_prefix"):318 base = getattr(self.model, self.model.base_model_prefix, self.model)319 model_device = base.device320 elif hasattr(self.model, "model"):321 model_device = self.model.model.device322 else:323 model_device = "cpu"324 logging.warning(325 "TorchExportableModuleForDecoderOnlyLM.export Can't infer device from the model. Set to CPU by default."326 )327 328 if input_ids is not None:329 input_kwargs = {330 "input_ids": input_ids,331 "cache_position": cache_position332 if cache_position is not None333 else torch.arange(input_ids.shape[-1], dtype=torch.long, device=model_device),334 }335 else: # inputs_embeds336 input_kwargs = {337 "inputs_embeds": inputs_embeds,338 "cache_position": cache_position339 if cache_position is not None340 else torch.arange(inputs_embeds.shape[1], dtype=torch.long, device=model_device),341 }342 343 exported_program = torch.export.export(344 self.model,345 args=(),346 kwargs=input_kwargs,347 dynamic_shapes=dynamic_shapes,348 strict=strict if strict is not None else True,349 )350 351 return exported_program352 353 @staticmethod354 def generate(355 exported_program: torch.export.ExportedProgram,356 tokenizer,357 prompt: str,358 max_new_tokens: int = 20,359 do_sample: bool = False,360 temperature: float = 1.0,361 top_k: int = 50,362 top_p: float = 1.0,363 device: str = "cpu",364 ) -> str:365 """366 Generate a sequence of tokens using an exported program.367 368 Args:369 exported_program (`torch.export.ExportedProgram`): The exported model being used for generate.370 tokenizer: The tokenizer to use.371 prompt (str): The input prompt.372 max_new_tokens (int): Maximum number of new tokens to generate.373 do_sample (bool): Whether to use sampling or greedy decoding.374 temperature (float): The temperature for sampling.375 top_k (int): The number of highest probability tokens to keep for top-k sampling.376 top_p (float): The cumulative probability for nucleus sampling.377 device (str): The device to use.378 379 Returns:380 str: The generated text.381 """382 # Get the module from the exported program383 exported_module = exported_program.module()384 385 # Tokenize the prompt386 input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)387 388 # Initialize with the prompt389 generated_ids = input_ids.clone()390 391 # Process the prompt tokens first392 curr_position = 0393 for i in range(input_ids.shape[1]):394 # Process one token at a time395 curr_input_ids = input_ids[:, i : i + 1]396 curr_cache_position = torch.tensor([curr_position], dtype=torch.long, device=device)397 398 # Forward pass399 _ = exported_module(input_ids=curr_input_ids, cache_position=curr_cache_position)400 curr_position += 1401 402 # Generate new tokens403 for _ in range(max_new_tokens):404 # Get the last token as input405 curr_input_ids = generated_ids[:, -1:]406 curr_cache_position = torch.tensor([curr_position], dtype=torch.long, device=device)407 408 # Forward pass to get next token logits409 outputs = exported_module(input_ids=curr_input_ids, cache_position=curr_cache_position)410 411 # Get the next token ID412 if do_sample:413 # Apply temperature414 if temperature > 0:415 logits = outputs / temperature416 else:417 logits = outputs418 419 # Apply top-k filtering420 if top_k > 0:421 indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]422 logits[indices_to_remove] = float("-inf")423 424 # Apply top-p (nucleus) filtering425 if top_p < 1.0:426 sorted_logits, sorted_indices = torch.sort(logits, descending=True)427 cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)428 429 # Remove tokens with cumulative probability above the threshold430 sorted_indices_to_remove = cumulative_probs > top_p431 # Shift the indices to the right to keep also the first token above the threshold432 sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()433 sorted_indices_to_remove[..., 0] = 0434 435 # Scatter sorted tensors to original indexing436 indices_to_remove = sorted_indices_to_remove.scatter(-1, sorted_indices, sorted_indices_to_remove)437 logits[indices_to_remove] = float("-inf")438 439 # Sample from the filtered distribution440 probs = torch.softmax(logits, dim=-1)441 next_token_id = torch.multinomial(probs, num_samples=1)442 else:443 # Greedy decoding444 next_token_id = outputs.argmax(dim=-1, keepdim=True)445 446 # Ensure next_token_id has the right shape before concatenation447 if next_token_id.dim() > 2:448 next_token_id = next_token_id.squeeze(-1)449 450 # Append to the generated sequence451 generated_ids = torch.cat([generated_ids, next_token_id], dim=-1)452 curr_position += 1453 454 # Stop if we generate an EOS token455 if next_token_id.item() == tokenizer.eos_token_id:456 break457 458 # Decode the generated text459 return tokenizer.decode(generated_ids[0], skip_special_tokens=True)460 461 462class TorchExportableModuleWithStaticCache(torch.nn.Module):463 """464 A recipe module designed to make a `PreTrainedModel` exportable with `torch.export`,465 specifically for decoder-only LM to `StaticCache`. This module ensures that the466 exported model is compatible with further lowering and execution in `ExecuTorch`.467 468 Note:469 This class is specifically designed to support export process using `torch.export`470 in a way that ensures the model can be further lowered and run efficiently in `ExecuTorch`.471 """472 473 def __init__(474 self,475 model: PreTrainedModel,476 batch_size: Optional[int] = None,477 max_cache_len: Optional[int] = None,478 device: Optional[torch.device] = None,479 ) -> None:480 """481 Initializes the wrapper module with the pretrained model.482 483 Args:484 model (`PreTrainedModel`): The pretrained model to wrap. The model must have caching485 enabled and use a 'static' caching implementation.486 batch_size (`Optional[int]`): The batch size of the model. If not provided, we check if a value can be found487 in `generation_config.cache_config` and otherwise we raise a ValueError.488 max_cache_len (`Optional[int]`): The maximum cache length for generation. Same mechanism as `batch_size` if489 not provided.490 device (`Optional[torch.device]`): The device to use. If not provided, we check if a value can be found491 in `generation_config.cache_config` and otherwise we use `model.device` (no error is raised).492 493 Raises:494 AssertionError: If the pretrained model does not have caching enabled or if it does495 not use a 'static' caching implementation in `model.generation_config`.496 ValueError: If `batch_size` or `max_cache_len` is not provided, either as an argument or in `cache_config`.497 """498 super().__init__()499 500 config = model.config.get_text_config()501 generation_config = model.generation_config502 503 # Sanity checks504 if generation_config is None:505 raise AssertionError(506 "The model must have a generation config to be exported with static caching. "507 "Please set `generation_config` in `model`."508 )509 if not generation_config.use_cache:510 raise AssertionError(511 "The model must have caching enabled to be exported with static caching. "512 "Please set `generation_config.use_cache=True`."513 )514 if generation_config.cache_implementation != "static":515 raise AssertionError(516 "The model must use a 'static' caching implementation to be exported with static caching. "517 "Please set `generation_config.cache_implementation='static'`."518 )519 520 cache_config = {} if generation_config.cache_config is None else generation_config.cache_config521 522 # Ensure batch_size and max_cache_len are set523 if batch_size is None:524 batch_size = cache_config.get("batch_size", None)525 if batch_size is None:526 raise ValueError("batch_size must be provided, either as an argument or in cache_config.")527 if max_cache_len is None:528 max_cache_len = cache_config.get("max_cache_len", None)529 if max_cache_len is None:530 raise ValueError("max_cache_len must be provided, either as an argument or in cache_config.")531 # Infer device if not provided532 if device is None:533 device = cache_config.get("device", model.device)534 535 # Initialize the static cache536 self.model = model537 self.static_cache = StaticCache(max_cache_len=max_cache_len, config=config)538 head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)539 num_heads = getattr(config, "num_key_value_heads", config.num_attention_heads)540 dtype = self.model.dtype541 # We need this call to initialize all the layers (otherwise it's done lazily, which is not exportable)542 self.static_cache.early_initialization(batch_size, num_heads, head_dim, dtype, device)543 for i in range(len(self.static_cache)):544 self.register_buffer(f"key_cache_{i}", self.static_cache.layers[i].keys, persistent=False)545 self.register_buffer(f"value_cache_{i}", self.static_cache.layers[i].values, persistent=False)546 547 def forward(548 self,549 input_ids: Optional[torch.LongTensor] = None,550 inputs_embeds: Optional[torch.Tensor] = None,551 cache_position: Optional[torch.Tensor] = None,552 ):553 """554 Forward pass of the module, which is compatible with the ExecuTorch runtime.555 556 Args:557 input_ids (`torch.Tensor`): Tensor representing current input token id to the module.558 inputs_embeds (`torch.Tensor`): Tensor representing current input embeddings to the module.559 cache_position (`torch.Tensor`): Tensor representing current input position in the cache.560 561 Returns:562 torch.Tensor: Logits output from the model.563 564 This forward adapter serves two primary purposes:565 566 1. **Making the Model `torch.export`-Compatible**:567 The adapter hides unsupported objects, such as the `Cache`, from the graph inputs and outputs,568 enabling the model to be exportable using `torch.export` without encountering issues.569 570 2. **Ensuring Compatibility with `ExecuTorch` runtime**:571 The adapter matches the model's forward signature with that in `executorch/extension/llm/runner`,572 ensuring that the exported model can be executed in `ExecuTorch` out-of-the-box.573 """574 past_key_values = self.static_cache575 576 outs = self.model(577 input_ids=input_ids,578 inputs_embeds=inputs_embeds,579 cache_position=cache_position,580 attention_mask=None,581 past_key_values=past_key_values,582 use_cache=True,583 )584 if hasattr(outs, "logits"):585 # Returned outputs is `CausalLMOutputWithPast`586 return outs.logits587 else:588 # Returned the `last_hidden_state` from `BaseModelOutputWithPast`589 return outs.last_hidden_state590 591 @staticmethod592 def generate(593 exported_program: torch.export.ExportedProgram,594 prompt_token_ids: torch.Tensor,595 max_new_tokens: int,596 ) -> torch.Tensor:597 """598 Generate a sequence of tokens using an exported program.599 600 This util function is designed to test exported models by simulating the generation process.601 It processes the input prompt tokens sequentially (no parallel prefill).602 This generate function is not intended to replace the original `generate` method, and the support603 for leveraging the original `generate` is potentially planned!604 605 Args:606 exported_program (`torch.export.ExportedProgram`): The exported program generated via `torch.export`.607 prompt_token_ids (`torch.Tensor`): Tensor representing the input prompt token IDs.608 max_new_tokens (`int`): Maximum number of new tokens to generate. Note that the total generation609 length is limited by both `max_new_tokens` and the model's cache size.610 611 Returns:612 torch.Tensor: A tensor containing the generated sequence of token IDs, including the original prompt tokens.613 """614 device = prompt_token_ids.device615 prompt_token_len = prompt_token_ids.shape[-1]616 max_generation_length = prompt_token_len + max_new_tokens617 for buffer_name, buffer in exported_program.named_buffers():618 if buffer_name.startswith("key_cache"):619 max_cache_len = buffer.shape[2]620 max_generation_length = min(max_generation_length, max_cache_len)621 break622 623 response_tokens = []624 for input_pos in range(min(max_generation_length, prompt_token_len)):625 result = exported_program.module().forward(626 input_ids=prompt_token_ids[:, input_pos : input_pos + 1],627 cache_position=torch.tensor([input_pos], dtype=torch.long, device=device),628 )629 response_tokens.append(prompt_token_ids[0][input_pos].item())630 631 current_token = torch.argmax(result[:, -1, :], dim=-1).item()632 response_tokens.append(current_token)633 634 while len(response_tokens) < max_generation_length:635 result = exported_program.module().forward(636 input_ids=torch.tensor([[current_token]], dtype=torch.long, device=device),637 cache_position=torch.tensor([len(response_tokens)], dtype=torch.long, device=device),638 )639 current_token = torch.argmax(result[:, -1, :], dim=-1).item()640 response_tokens.append(current_token)641 642 return torch.tensor([response_tokens], dtype=torch.long, device=device)643 644 645class TorchExportableModuleWithHybridCache(torch.nn.Module):646 """647 A recipe module designed to make a `PreTrainedModel` exportable with `torch.export`,648 specifically for decoder-only LM to hybrid `StaticCache`. This module ensures that the649 exported model is compatible with further lowering and execution in `ExecuTorch`.650 """651 652 def __init__(653 self,654 model: PreTrainedModel,655 batch_size: Optional[int] = None,656 max_cache_len: Optional[int] = None,657 device: Optional[torch.device] = None,658 ) -> None:659 """660 Initializes the exportable module.661 662 Args:663 model (`PreTrainedModel`): The pretrained model to wrap.664 batch_size (`Optional[int]`): The batch size of the model. If not provided, we check if a value can be found665 in `generation_config.cache_config` and otherwise we raise a ValueError.666 max_cache_len (`Optional[int]`): The maximum cache length for generation. Same mechanism as `batch_size` if667 not provided.668 device (`Optional[torch.device]`): The device to use. If not provided, we check if a value can be found669 in `generation_config.cache_config` and otherwise we use `model.device` (no error is raised).670 Raises:671 AssertionError: If the model doesn't have the expected configuration for hybrid StaticCache.672 ValueError: If `batch_size` or `max_cache_len` is not provided, either as an argument or in `cache_config`.673 """674 super().__init__()675 self.model = model676 config = model.config.get_text_config()677 generation_config = model.generation_config678 679 # Sanity checks680 if generation_config is None:681 raise AssertionError(682 "The model must have a generation config to be exported with static caching. "683 "Please set `generation_config` in `model`."684 )685 if not config.use_cache:686 raise AssertionError("Model must have caching enabled.")687 688 cache_config = {} if generation_config.cache_config is None else generation_config.cache_config689 # Ensure batch_size and max_cache_len are set690 if batch_size is None:691 batch_size = cache_config.get("batch_size", None)692 if batch_size is None:693 raise ValueError("batch_size must be provided, either as an argument or in cache_config.")694 if max_cache_len is None:695 max_cache_len = cache_config.get("max_cache_len", None)696 if max_cache_len is None:697 raise ValueError("max_cache_len must be provided, either as an argument or in cache_config.")698 # Infer device if not provided699 if device is None:700 device = cache_config.get("device", model.device)701 702 # Initialize the cache703 self.cache = StaticCache(config=config, max_cache_len=max_cache_len)704 head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)705 num_heads = getattr(config, "num_key_value_heads", config.num_attention_heads)706 dtype = self.model.dtype707 # We need this call to initialize all the layers (otherwise it's done lazily, which is not exportable)708 self.cache.early_initialization(batch_size, num_heads, head_dim, dtype, device)709 710 # Register all key and value cache tensors as buffers711 for i in range(len(self.cache)):712 self.register_buffer(f"key_cache_{i}", self.cache.layers[i].keys, persistent=False)713 self.register_buffer(f"value_cache_{i}", self.cache.layers[i].values, persistent=False)714 715 def forward(716 self,717 input_ids: Optional[torch.LongTensor] = None,718 inputs_embeds: Optional[torch.Tensor] = None,719 cache_position: Optional[torch.Tensor] = None,720 ) -> torch.Tensor:721 """722 Forward pass of the module, which is compatible with the ExecuTorch llm runner.723 724 Args:725 input_ids (`torch.Tensor`): Tensor representing current input token id to the module.726 inputs_embeds (`Optional[torch.Tensor]`): Tensor representing current input embeddings to the module.727 cache_position (`torch.Tensor`): Tensor representing current input position in the cache.728 729 Returns:730 torch.Tensor: Logits output from the model.731 """732 # Forward pass with the model733 outputs = self.model(734 input_ids=input_ids,735 inputs_embeds=inputs_embeds,736 cache_position=cache_position,737 attention_mask=None,738 past_key_values=self.cache,739 use_cache=True,740 )741 742 # Return only the logits to simplify the export743 return outputs.logits744 745 746def convert_and_export_with_cache(747 model: PreTrainedModel,748 example_input_ids: Optional[torch.Tensor] = None,749 example_cache_position: Optional[torch.Tensor] = None,750 dynamic_shapes: Optional[dict] = None,751 strict: Optional[bool] = None,752):753 """754 Convert a `PreTrainedModel` into an exportable module and export it using `torch.export`,755 ensuring the exported model is compatible with `ExecuTorch`.756 757 Args:758 model (`PreTrainedModel`): The pretrained model to be exported.759 example_input_ids (`Optional[torch.Tensor]`): Example input token id used by `torch.export`.760 example_cache_position (`Optional[torch.Tensor]`): Example current cache position used by `torch.export`.761 dynamic_shapes(`Optional[dict]`): Dynamic shapes used by `torch.export`.762 strict(`Optional[bool]`): Flag to instruct `torch.export` to use `torchdynamo`.763 764 Returns:765 Exported program (`torch.export.ExportedProgram`): The exported program generated via `torch.export`.766 """767 if not is_torch_greater_or_equal_than_2_3:768 raise ImportError("torch >= 2.3 is required.")769 770 import torch.export._trace771 772 # This is the same as sdpa, but mask creation does not use `vmap` which is not exportable773 ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa_without_vmap", sdpa_mask_without_vmap)774 ALL_ATTENTION_FUNCTIONS.register("sdpa_without_vmap", ALL_ATTENTION_FUNCTIONS["sdpa"])775 model.config._attn_implementation = "sdpa_without_vmap"776 777 with torch.no_grad():778 # TODO: The default inputs only work for text models. We need to add support for vision/audio models.779 example_input_ids = (780 example_input_ids781 if example_input_ids is not None782 else torch.tensor([[1]], dtype=torch.long, device=model.device)783 )784 example_cache_position = (785 example_cache_position786 if example_cache_position is not None787 else torch.tensor([0], dtype=torch.long, device=model.device)788 )789 790 if is_torch_greater_or_equal("2.6.0"):791 exported_program = torch.export.export(792 TorchExportableModuleWithStaticCache(model),793 args=(),794 kwargs={"input_ids": example_input_ids, "cache_position": example_cache_position},795 dynamic_shapes=dynamic_shapes,796 strict=strict if strict is not None else True,797 )798 else:799 if dynamic_shapes is not None:800 logging.warning(801 "Dynamic shapes spec will be ignored by convert_and_export_with_cache for torch < 2.6.0."802 )803 if strict is not None:804 logging.warning("The strict flag will be ignored by convert_and_export_with_cache for torch < 2.6.0.")805 # We have to keep this path for BC.806 #807 # Due to issue https://github.com/pytorch/pytorch/issues/128394, we need to switch to use an internal808 # export API and pre_dispatch=False. Switch to use the public API once the issue is included in 2.5 release.809 exported_program = torch.export._trace._export(810 TorchExportableModuleWithStaticCache(model),811 args=(),812 kwargs={"input_ids": example_input_ids, "cache_position": example_cache_position},813 pre_dispatch=False,814 strict=True,815 )816 return exported_program817 818 819class Seq2SeqLMEncoderExportableModule(torch.nn.Module):820 """821 A wrapper module designed to make a Seq2Seq LM encoder exportable with `torch.export`.822 This module ensures that the exported encoder model is compatible with ExecuTorch.823 """824 825 def __init__(self, encoder_model):826 super().__init__()827 self.encoder = encoder_model828 829 def forward(self, input_ids):830 return self.encoder(input_ids=input_ids).last_hidden_state831 832 833class Seq2SeqLMDecoderExportableModuleWithStaticCache(torch.nn.Module):834 """835 A wrapper module designed to make a Seq2Seq LM decoder exportable with `torch.export`,836 specifically for use with static caching. This module ensures the exported decoder837 is compatible with ExecuTorch.838 """839 840 def __init__(self, model, max_static_cache_length, batch_size):841 super().__init__()842 843 # Get the decoder component844 self.decoder = model.get_decoder()845 self.lm_head = model.lm_head846 self.config = model.config847 848 # Detect the device of the exported models by checking a parameter849 # We'll use the model's device as the target device850 model_device = next(model.parameters()).device851 852 # Initialize static cache for decoder and DynamicCache for encoder853 self.static_cache = StaticCache(config=self.config, max_cache_len=max_static_cache_length)854 head_dim = getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads)855 num_heads = getattr(self.config, "num_key_value_heads", self.config.num_attention_heads)856 self.static_cache.early_initialization(batch_size, num_heads, head_dim, torch.float32, model_device)857 self.cache = EncoderDecoderCache(self.static_cache, DynamicCache(config=self.config))858 859 register_dynamic_cache_export_support()860 861 # Register cache buffers to make them exportable862 for i in range(len(self.static_cache)):863 self.register_buffer(f"key_cache_{i}", self.static_cache.layers[i].keys, persistent=False)864 self.register_buffer(f"value_cache_{i}", self.static_cache.layers[i].values, persistent=False)865 866 def forward(self, decoder_input_ids, encoder_hidden_states, cache_position):867 # Get outputs from decoder868 outputs = self.decoder(869 input_ids=decoder_input_ids,870 encoder_hidden_states=encoder_hidden_states,871 past_key_values=self.cache,872 use_cache=True,873 cache_position=cache_position,874 )875 876 # Apply language model head877 lm_logits = self.lm_head(outputs[0])878 879 return lm_logits880 881 882class Seq2SeqLMExportableModule(torch.nn.Module):883 def __init__(884 self, model, batch_size=1, max_hidden_seq_length=4096, cache_implementation="static", max_cache_length=1024885 ):886 super().__init__()887 888 self.full_model = model889 self.encoder = model.get_encoder()890 self.config = model.config891 self.max_hidden_seq_length = max_hidden_seq_length892 self.generation_config = GenerationConfig(893 use_cache=True,894 max_length=max_cache_length,895 cache_implementation=cache_implementation,896 cache_config={897 "batch_size": batch_size,898 "max_cache_len": max_cache_length,899 },900 )901 self.exported_encoder = None902 self.exported_decoder = None903 904 def _export_encoder(self, encoder_input_ids):905 wrapped_encoder = Seq2SeqLMEncoderExportableModule(self.encoder).to(self.full_model.device).eval()906 907 # Define dynamic sequence length for encoder908 seq_len_dim = torch.export.Dim("encoder_seq_length", max=self.max_hidden_seq_length)909 910 # Export the encoder911 with torch.no_grad():912 exported_encoder = torch.export.export(913 wrapped_encoder, (encoder_input_ids,), dynamic_shapes={"input_ids": {1: seq_len_dim}}, strict=True914 )915 916 return exported_encoder917 918 def _export_decoder(self, decoder_input_ids, encoder_hidden_states, cache_position):919 target_device = self.full_model.device920 wrapped_decoder = (921 Seq2SeqLMDecoderExportableModuleWithStaticCache(922 model=self.full_model,923 max_static_cache_length=self.generation_config.cache_config.get("max_cache_len"),924 batch_size=self.generation_config.cache_config.get("batch_size"),925 )926 .to(target_device)927 .eval()928 )929 930 # Move input tensors to the same device as the wrapped decoder931 decoder_input_ids = decoder_input_ids.to(target_device)932 encoder_hidden_states = encoder_hidden_states.to(target_device)933 cache_position = cache_position.to(target_device)934 935 # Define dynamic dimension for encoder output sequence length936 encoder_seq_len_dim = torch.export.Dim("encoder_hidden_seq_length", max=self.max_hidden_seq_length)937 938 # Export the decoder939 with torch.no_grad():940 exported_decoder = torch.export.export(941 wrapped_decoder,942 (decoder_input_ids, encoder_hidden_states, cache_position),943 dynamic_shapes={944 "decoder_input_ids": None,945 "encoder_hidden_states": {1: encoder_seq_len_dim},946 "cache_position": None,947 },948 strict=True,949 )950 951 return exported_decoder952 953 def export(self, encoder_input_ids=None, decoder_input_ids=None, encoder_hidden_states=None, cache_position=None):954 device = self.full_model.device955 example_encoder_input_ids = (956 encoder_input_ids957 if encoder_input_ids is not None958 else torch.ones((1, 10), dtype=torch.long, device=device)959 )960 example_decoder_input_ids = (961 decoder_input_ids962 if decoder_input_ids is not None963 else torch.tensor([[0]], dtype=torch.long, device=device)964 ) # Start token965 example_cache_position = (966 cache_position if cache_position is not None else torch.tensor([0], dtype=torch.long, device=device)967 )968 example_encoder_hidden_states = (969 encoder_hidden_states970 if encoder_hidden_states is not None971 else torch.zeros(972 (self.generation_config.cache_config.get("batch_size"), 10, self.config.d_model),973 dtype=torch.float32,974 device=device,975 )976 )977 self.exported_encoder = self._export_encoder(example_encoder_input_ids)978 self.exported_decoder = self._export_decoder(979 example_decoder_input_ids, example_encoder_hidden_states, example_cache_position980 )981 982 # Return self to allow chaining983 return self984 985 def generate(self, prompt_token_ids, max_new_tokens):986 with torch.no_grad():987 model_device = self.full_model.device988 989 # Move input to the model's device if it's on a different device990 if prompt_token_ids.device != model_device:991 prompt_token_ids = prompt_token_ids.to(model_device)992 993 # Run encoder994 encoder_output = self.exported_encoder.module()(prompt_token_ids)995 996 # Initialize with start token (0 for T5) on the correct device997 decoder_input_ids = torch.tensor([[0]], dtype=torch.long, device=model_device)998 generated_ids = [0]999 1000 # Generate tokens one by one1001 for i in range(max_new_tokens - 1):1002 # Run decoder for next token prediction1003 logits = self.exported_decoder.module()(1004 decoder_input_ids, encoder_output, torch.tensor([i], dtype=torch.long, device=model_device)1005 )1006 1007 # Get next token1008 next_token = torch.argmax(logits[:, -1, :], dim=-1).item()1009 generated_ids.append(next_token)1010 1011 # Update input for next iteration on the correct device1012 decoder_input_ids = torch.tensor([[next_token]], dtype=torch.long, device=model_device)1013 1014 # Check if EOS token1015 if next_token == self.config.eos_token_id:1016 break1017 1018 return generated_ids1019 1020 1021def export_with_dynamic_cache(1022 model: PreTrainedModel,1023 example_input_ids: Optional[torch.Tensor] = None,1024 example_attention_mask: Optional[torch.Tensor] = None,1025):1026 """1027 Export a model with DynamicCache using `torch.export`, ensuring the exported model is compatible with `ExecuTorch`.1028 1029 Args:1030 model (`PreTrainedModel`): The pretrained model to be exported.1031 example_input_ids (`Optional[torch.Tensor]`): Example input token id used by `torch.export`.1032 example_attention_mask (`Optional[torch.Tensor]`): Example attention mask used by `torch.export`.1033 1034 Returns:1035 Exported program (`torch.export.ExportedProgram`): The exported program generated via `torch.export`.1036 """1037 if not is_torch_greater_or_equal_than_2_3:1038 raise ImportError("torch >= 2.3 is required.")1039 1040 # This is the same as sdpa, but mask creation does not use `vmap` which is not exportable1041 ALL_MASK_ATTENTION_FUNCTIONS.register("sdpa_without_vmap", sdpa_mask_without_vmap)1042 ALL_ATTENTION_FUNCTIONS.register("sdpa_without_vmap", ALL_ATTENTION_FUNCTIONS["sdpa"])1043 model.config._attn_implementation = "sdpa_without_vmap"1044 1045 register_dynamic_cache_export_support()1046 1047 with torch.no_grad():1048 exported_program = torch.export.export(1049 model,1050 (),1051 {1052 "input_ids": example_input_ids,1053 "attention_mask": example_attention_mask,1054 "past_key_values": DynamicCache(config=model.config),1055 "use_cache": True,1056 },1057 strict=False,1058 )1059 return exported_program1060 1061 1062def register_dynamic_cache_export_support():1063 """1064 Utilities for `DynamicCache` <> torch.export support1065 """1066 1067 try:1068 torch.utils._pytree.register_pytree_node(1069 DynamicCache,1070 lambda dynamic_cache: torch.utils._pytree._dict_flatten(_get_cache_dict(dynamic_cache)),1071 _unflatten_dynamic_cache,1072 serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",1073 flatten_with_keys_fn=lambda dynamic_cache: torch.utils._pytree._dict_flatten_with_keys(1074 _get_cache_dict(dynamic_cache)1075 ),1076 )1077 # TODO (tmanlaibaatar) This won't be needed in torch 2.7.1078 torch.fx._pytree.register_pytree_flatten_spec(1079 DynamicCache,1080 lambda cache, spec: torch.fx._pytree._dict_flatten_spec(_get_cache_dict(cache), spec),1081 )1082 # Catching this in case there are multiple runs for some test runs1083 except ValueError as e:1084 if "already registered as pytree node" not in str(e):1085 raise1086 1087 1088def _get_cache_dict(cache: DynamicCache):1089 """Convert cache to dictionary format for pytree operations."""1090 if any(not isinstance(layer, (DynamicLayer, DynamicSlidingWindowLayer)) for layer in cache.layers):1091 raise RuntimeError("This pytree flattening function should only be applied to DynamicCache")1092 1093 if not is_torch_greater_or_equal_than_2_6:1094 logging.warning("DynamicCache + torch.export is tested on torch 2.6.0+ and may not work on earlier versions.")1095 1096 return {1097 "key_cache": [layer.keys for layer in cache.layers if layer.keys is not None],1098 "value_cache": [layer.values for layer in cache.layers if layer.values is not None],1099 }1100 1101 1102def _unflatten_dynamic_cache(values, context: torch.utils._pytree.Context):1103 dictionary = torch.utils._pytree._dict_unflatten(values, context)1104 cache = DynamicCache()1105 # Reconstruct layers from keys and values lists1106 key_list = dictionary.get("key_cache", [])1107 value_list = dictionary.get("value_cache", [])1108 for idx in range(max(len(key_list), len(value_list))):1109 key = key_list[idx] if idx < len(key_list) else None1110 value = value_list[idx] if idx < len(value_list) else None1111 cache.update(key, value, idx)1112 return cache1113 1114 1115def sdpa_mask_without_vmap(1116 batch_size: int,1117 cache_position: torch.Tensor,1118 kv_length: int,1119 kv_offset: int = 0,1120 mask_function: Optional[Callable] = None,1121 attention_mask: Optional[torch.Tensor] = None,1122 local_size: Optional[int] = None,1123 allow_is_causal_skip: bool = True,1124 allow_torch_fix: bool = True,1125 **kwargs,1126) -> Optional[torch.Tensor]:1127 """1128 Create a 4D boolean mask of shape `(batch_size, 1, query_length, kv_length)` where a value of True indicates that1129 the element should take part in the attention computation, and False that it should not.1130 1131 This is similar to `masking_utils.sdpa_mask` but does not use `vmap` which is incompatible with export.1132 1133 Args:1134 batch_size (`int`):1135 The batch size of the input sequence.1136 cache_position (`torch.Tensor`):1137 A tensor of shape (query_length,) indicating the current indices of the input sequence elements.1138 kv_length (`int`):1139 The size that the key and value states will have during the attention computation.1140 kv_offset (`int`, optional):1141 An optional offset to indicate at which first position the key and values states will refer to.1142 mask_function (`Callable`):1143 The mask factory function describing the mask pattern.1144 attention_mask (`torch.Tensor`, optional):1145 The 2D attention mask corresponding to padded tokens of shape (batch_size, number_of_seen_tokens+q_length)1146 local_size (`int`, optional):1147 The size of the local attention, if we do not use full attention. This is used only if `allow_is_causal_skip=True`1148 to try to skip mask creation if possible.1149 allow_is_causal_skip (`bool`, optional):1150 Whether to allow to return `None` for the mask under conditions where we can use the `is_causal` argument in1151 `torch.sdpa` instead. Default to `True`.1152 allow_torch_fix (`bool`, optional):1153 Whether to update the mask in case a query is not attending to any tokens, to solve a bug in torch's older1154 versions. We need an arg to skip it when using eager. By default `True`.1155 1156 """1157 1158 q_length = cache_position.shape[0]1159 # Potentially pad the 2D mask, and slice it correctly1160 padding_mask = prepare_padding_mask(attention_mask, kv_length, kv_offset)1161 1162 # Under specific conditions, we can avoid materializing the mask, instead relying on the `is_causal` argument1163 if allow_is_causal_skip and _ignore_causal_mask_sdpa(padding_mask, q_length, kv_length, local_size):1164 return None1165 1166 # Similar to `kv_arange = torch.arange(start=kv_offset, end=kv_offset + kv_length, device=cache_position.device)`1167 # but without data-dependent slicing (i.e. torch.compile friendly)1168 kv_arange = torch.arange(kv_length, device=cache_position.device)1169 kv_arange += kv_offset1170 reshaped_cache_position = cache_position.view(-1, 1)1171 1172 # This is a bit hacky to know what pattern we are using, but all mask creation function actually forward1173 # the config through kwargs anyway, so it allows to rely on it1174 # Usually, the `mask_function` is the only entry-point to define the pattern - we could do for loops over it,1175 # but this is more efficient1176 sliding_window = getattr(kwargs["config"], "sliding_window", None)1177 chunk_size = getattr(kwargs["config"], "attention_chunk_size", None)1178 1179 if sliding_window is not None and chunk_size is not None:1180 raise ValueError("Cannot use both `sliding_window` and `attention_chunk_size`")1181 1182 # Simplest and most efficient way to obtain a causal mask1183 causal_mask = kv_arange <= reshaped_cache_position1184 # If using sliding window, add the sliding mask1185 if sliding_window is not None:1186 sliding_mask_overlay = kv_arange > reshaped_cache_position - sliding_window1187 causal_mask *= sliding_mask_overlay1188 # If using chunk attention, add the chunked mask1189 elif chunk_size is not None:1190 chunked_mask_overlay = kv_arange // chunk_size == reshaped_cache_position // chunk_size1191 causal_mask *= chunked_mask_overlay1192 1193 causal_mask = causal_mask[None, None, :, :].expand(batch_size, -1, -1, -1)1194 if padding_mask is not None:1195 causal_mask = causal_mask * padding_mask[:, None, None, :]1196 1197 # Due to a bug in some older torch version, we need to update the mask in case a query is not attending to any1198 # tokens (due to padding). See details in https://github.com/pytorch/pytorch/issues/1102131199 if not _is_torch_greater_or_equal_than_2_5 and allow_torch_fix:1200 causal_mask |= torch.all(~causal_mask, dim=-1, keepdim=True)