Создание совместимого бэкенда модели для инференса

Модели Transformers совместимы с движками инференса, такими как vLLM и SGLang. Используйте одну и ту же модель Transformers в любом месте, избегая необходимости заново реализовывать модель для каждого движка инференса. Это также полезно для моделей, которые движок не реализует нативно, при условии, что модель поддерживает интерфейс бэкенда внимания, описанный ниже.

В этом руководстве показано, как реализовать модель в Transformers, которая работает как backend для любого inference-движка.

Реализация модели

Следуйте руководству по внесению моделей или руководству по внесению пользовательских моделей. В каталоге модели должен находиться корректный файл config.json, а в конфигурации должно быть корректное поле auto_map, указывающее на класс модели.
Используйте класс AttentionInterface для пользовательских и оптимизированных функций внимания. Этот интерфейс обеспечивает доступ к возможностям производительности каждого движка инференса.

Используйте ALL_ATTENTION_FUNCTIONS при определении слоя внимания и передавайте **kwargs из базового класса MyModel в слои внимания. Установите _supports_attention_backend в True в PreTrainedModel.

Расширьте код ниже в качестве примера.
modeling_my_model.py
```
from transformers import PreTrainedModel
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
from torch import nn

class MyAttention(nn.Module):

    def forward(self, hidden_states, **kwargs):
        ...
        attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            **kwargs,
        )
        ...

class MyModel(PreTrainedModel):
    _supports_attention_backend = True
```
Включите необязательное тензорное или конвейерное параллелизм, добавив следующие ключи в PreTrainedConfig.
- base_model_tp_plan обеспечивает тензорное параллелизм путём сопоставления полных имён слоёв (в формате шаблонов) стилям тензорного параллелизма. Поддерживаются только стратегии разделения "colwise" и "rowwise".
- base_model_pp_plan включает параллелизм по конвейеру, сопоставляя имена прямых дочерних слоёв с кортежами списков строк. Первый элемент кортежа содержит имена входных аргументов. Последний элемент содержит имена переменных выходов слоя в коде моделирования.
Расширьте код ниже в качестве примера.
configuration_my_model.py
```
from transformers import PreTrainedConfig

class MyConfig(PreTrainedConfig):
    base_model_tp_plan = {
        "layers.*.self_attn.k_proj": "colwise",
        "layers.*.self_attn.v_proj": "colwise",
        "layers.*.self_attn.o_proj": "rowwise",
        "layers.*.mlp.gate_proj": "colwise",
        "layers.*.mlp.up_proj": "colwise",
        "layers.*.mlp.down_proj": "rowwise",
    }
    base_model_pp_plan = {
        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
        "norm": (["hidden_states"], ["hidden_states"]),
    }
```

Мультимодальные модели

Мультимодальные модели требуют дополнительных изменений сверх правил внесения вкладов в модели «зрение-язык». Эти изменения обеспечивают корректную обработку мультимодальных входных данных.

Класс ProcessorMixin должен включать атрибуты self.image_token и self.image_token_ids. Эти токены-заполнители указывают позиции изображений во входных данных. Один и тот же токен присутствует во входном промпте для изображений и в коде модели для распределения признаков изображений.
Класс ProcessorMixin должен включать метод self._get_num_multimodal_tokens. Этот метод вычисляет количество токенов-заполнителей, необходимых для мультимодальных входных данных заданных размеров. Он возвращает объект MultiModalData. Заполнители между токенами <image>, такие как токены строк или столбцов, не считаются заполнителями изображений. Учитывайте только те токены, которые впоследствии заменяются признаками изображений в коде модели.
Класс ProcessorMixin должен проверять значение return_mm_token_type_ids и возвращать mm_token_type_ids. Это указывает, является ли каждая позиция текстовым токеном (0), токеном-заполнителем изображения (1) или токеном-заполнителем видео (2). Последовательности идентификаторов мультимодальных типов токенов должны быть непрерывными, без разрывов между последовательными токенами. Специальные токены начала, конца, строки и столбца следует рассматривать как заполнители.

Расширьте код ниже в качестве примера.

modeling_my_multimodal_model.py

class MyMultimodalProcessor(ProcessorMixin):

    def __call__(self, images=None, text=None, **kwargs):
        if return_mm_token_type_ids:
            mm_token_type_ids = np.zeros_like(input_ids)
            mm_token_type_ids[input_ids == self.image_token_id] = 1
            text_inputs["mm_token_type_ids"] = mm_token_type_ids.tolist()
        return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors)

    def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs):
        """
        Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.
        Args:
            image_sizes (`list[list[int]]`, *optional*):
                The input sizes formatted as (height, width) per each image.
        Returns:
            `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided
            input modalities, along with other useful data.
        """
        vision_data = {}
        if image_sizes is not None:
            num_image_tokens = [256] * len(image_sizes) # 256 placeholder tokens for each image always
            num_image_patches = [1] * len(image_sizes) # no patching, thus each image is processed as a single base image
            vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
        return MultiModalData(**vision_data)

Ресурсы

Подробнее читайте в блоге Интеграция бэкенда Transformers в vLLM.
Подробнее читайте в блоге Интеграция бэкенда Transformers в SGLang.

Update on GitHub