Add `MtpModel` helpers for frameworks that own the drafting loop

@@ -44,6 +44,7 @@ if TYPE_CHECKING:     from .cache_utils import MtpCache     from .configuration_utils import PreTrainedConfig+    from .core_model_loading import WeightTransform     from .generation.logits_process import LogitsProcessorList  @@ -502,27 +503,10 @@ def forward(         drafted_logits = []         drafted_tokens = []         loss = None-        for i, mtp_layer in enumerate(self.layers):-            # We need to recompute those every layer since they change-            inputs_embeds = self.embed_tokens(input_ids).to(last_hidden_states.device)-            position_embeddings = (-                self.rotary_emb(inputs_embeds, position_ids=position_ids) if self.rotary_emb is not None else None-            )--            # In full generality, we may need to recompute masks for every layer due to the position offset of each layer-            masks = self.create_masks_for_mtp_layer(i, inputs_embeds, mtp_cache, position_ids)--            last_hidden_states = mtp_layer(-                inputs_embeds,-                last_hidden_states,-                position_embeddings=position_embeddings,-                position_ids=position_ids,-                past_key_values=mtp_cache,-                **masks,-                **kwargs,+        for i in range(len(self.layers)):+            last_hidden_states = self.forward_layer(+                i, last_hidden_states, position_ids, mtp_cache, input_ids=input_ids, **kwargs             )-            if self.use_shared_post_norm:-                last_hidden_states = self.shared_post_norm(last_hidden_states)              # If we are not computing the loss, only compute logits for the next drafted token to save memory             slice_indices = slice(-1, None) if labels is None else slice(None, None)@@ -563,13 +547,53 @@ def forward(         candidate_logits = torch.cat(drafted_logits, dim=1)         return new_candidate_ids, candidate_logits, loss -    @classmethod-    def from_pretrained(cls, main_model: PreTrainedModel, device_map=None, **kwargs) -> MtpModel:-        pretrained_model_name_or_path = main_model.config.name_or_path+    def forward_layer(+        self,+        layer_idx: int,+        last_hidden_states: torch.Tensor,+        position_ids: torch.Tensor | None,+        mtp_cache: MtpCache | None,+        input_ids: torch.Tensor | None = None,+        inputs_embeds: torch.Tensor | None = None,+        **kwargs,+    ) -> torch.Tensor:+        """+        Run the mtp layer `layer_idx` once, i.e. a single step of the drafting loop in `forward`. Exactly one of `input_ids`+        or `inputs_embeds` should be passed. This lets frameworks which own the drafting loop (e.g. vLLM) reuse the per-layer+        logic.+        """+        if (input_ids is None) == (inputs_embeds is None):+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")+        # We need to recompute those every layer since they change+        if inputs_embeds is None:+            inputs_embeds = self.embed_tokens(input_ids).to(last_hidden_states.device)++        # In full generality, we may need to recompute masks for every layer due to the position offset of each layer+        masks = self.create_masks_for_mtp_layer(layer_idx, inputs_embeds, mtp_cache, position_ids)++        position_embeddings = (+            self.rotary_emb(inputs_embeds, position_ids=position_ids) if self.rotary_emb is not None else None+        )+        last_hidden_states = self.layers[layer_idx](+            inputs_embeds,+            last_hidden_states,+            position_embeddings=position_embeddings,+            position_ids=position_ids,+            past_key_values=mtp_cache,+            **masks,+            **kwargs,+        )+        if self.use_shared_post_norm:+            last_hidden_states = self.shared_post_norm(last_hidden_states)+        return last_hidden_states++    @staticmethod+    def get_mtp_weights_regex(main_model: PreTrainedModel) -> re.Pattern:+        """Regex matching the checkpoint keys of the mtp weights of `main_model`."""         num_hidden_layers = main_model.config.get_text_config().num_hidden_layers         # Heuristic: the main model should have the mtp layer patterns under `_keys_to_ignore_on_load_unexpected` to avoid         # loading them by default, so use it to later load the correct keys from the checkpoints-        mtp_patterns = main_model._keys_to_ignore_on_load_unexpected.copy()  # type: ignore+        mtp_patterns = main_model._keys_to_ignore_on_load_unexpected         # Due to different released checkpoints, only keep the ones with layer number >= num_hidden_layers - otherwise         # mtp layers in a smaller checkpoints could be wrongly added as a 2nd mtp layer of a bigger checkpoint         final_mtp_patterns = []@@ -580,7 +604,33 @@ def from_pretrained(cls, main_model: PreTrainedModel, device_map=None, **kwargs)             final_mtp_patterns.append(pattern)         if len(final_mtp_patterns) == 0:             raise ValueError(f"{main_model.__class__.__name__} does not seem to register any known MTP layer patterns")-        mtp_regex = re.compile("|".join(rf"({pattern})" for pattern in final_mtp_patterns))+        return re.compile("|".join(rf"({pattern})" for pattern in final_mtp_patterns))++    def get_weight_conversions(self, main_model: PreTrainedModel) -> list[WeightTransform]:+        """The weight conversions which map the mtp weights of `main_model`'s checkpoint onto this model."""+        num_hidden_layers = main_model.config.get_text_config().num_hidden_layers+        # For the correct conversions, we need first the mtp-specific renamings, then the main_model conversions+        # Note that since the layer numbers are dynamic, we cannot register those conversions - we also add the `mtp_block`+        # part for all weights since we cannot distinguish easily those that are under the main model's block or not. It will+        # be removed after for the few that should not have it+        weight_conversions = [+            WeightRenaming(+                source_patterns=f"layers.{N}.", target_patterns=f"layers.{N - num_hidden_layers}.mtp_block."+            )+            for N in range(num_hidden_layers, num_hidden_layers + self.num_mtp_layers)+        ]+        weight_conversions.extend(get_model_conversion_mapping(self, add_legacy=False))+        # `_weight_conversions` is only attached to models loaded with `from_pretrained`+        main_conversions = getattr(main_model, "_weight_conversions", None)+        if main_conversions is None:+            main_conversions = get_model_conversion_mapping(main_model)+        weight_conversions.extend(main_conversions)+        return weight_conversions++    @classmethod+    def from_pretrained(cls, main_model: PreTrainedModel, device_map=None, **kwargs) -> MtpModel:+        pretrained_model_name_or_path = main_model.config.name_or_path+        mtp_regex = cls.get_mtp_weights_regex(main_model)          # Get the number of layers in the checkpoint         num_mtp_layers = main_model.config.get_text_config().num_mtp_layers@@ -623,18 +673,7 @@ def from_pretrained(cls, main_model: PreTrainedModel, device_map=None, **kwargs)                 ):                     mtp_state_dict[k] = file_pointer.get_slice(k)  # don't materialize yet -        # For the correct conversions, we need first the mtp-specific renamings, then the main_model conversions-        # Note that since the layer numbers are dynamic, we cannot register those conversions - we also add the `mtp_block`-        # part for all weights since we cannot distinguish easily those that are under the main model's block or not. It will-        # be removed after for the few that should not have it-        weight_conversions = [-            WeightRenaming(-                source_patterns=f"layers.{N}.", target_patterns=f"layers.{N - num_hidden_layers}.mtp_block."-            )-            for N in range(num_hidden_layers, num_hidden_layers + num_mtp_layers)-        ]-        weight_conversions.extend(get_model_conversion_mapping(mtp_model, add_legacy=False))-        weight_conversions.extend(main_model._weight_conversions)+        weight_conversions = mtp_model.get_weight_conversions(main_model)          # Load the weights         loading_info, _ = convert_and_load_state_dict_in_model(