Xây dựng backend mô hình tương thích cho suy luận

Các mô hình Transformers tương thích với các bộ máy suy luận như vLLM và SGLang. Bạn có thể sử dụng cùng một mô hình Transformers ở bất kỳ đâu và tránh việc phải cài đặt lại mô hình từ đầu cho từng bộ máy suy luận. Điều này cũng hữu ích cho các mô hình mà bộ máy không cài đặt lại một cách bản địa, miễn là mô hình hỗ trợ giao diện backend chú ý được mô tả bên dưới.

Hướng dẫn này cho bạn cách triển khai một mô hình trong Transformers hoạt động như một backend cho bất kỳ engine suy luận nào.

Thực hiện mô hình

Tuân theo hướng dẫn đóng góp mô hình hoặc hướng dẫn đóng góp mô hình tùy chỉnh. Mô hình phải có config.json hợp lệ trong thư mục của nó và trường auto_map hợp lệ trỏ đến lớp mô hình trong cấu hình.

Sử dụng lớp AttentionInterface cho các hàm chú ý tùy chỉnh và được tối ưu hóa. Giao diện này mở khóa các tính năng hiệu năng của từng bộ máy suy luận.

Sử dụng ALL_ATTENTION_FUNCTIONS khi định nghĩa lớp attention và truyền **kwargs từ lớp cơ sở MyModel đến các lớp attention. Đặt _supports_attention_backend thành True trong PreTrainedModel.

Mở rộng đoạn mã bên dưới để có một ví dụ.

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

Bật song song tensor hoặc pipeline tùy chọn bằng cách thêm các khóa sau vào PreTrainedConfig.
- base_model_tp_plan cho phép tensor parallelism bằng cách ánh xạ các mẫu tên lớp đầy đủ sang các kiểu tensor parallel. Chỉ hỗ trợ các chiến lược phân vùng "colwise" và "rowwise".
- base_model_pp_plan cho phép pipeline parallelism bằng cách ánh xạ tên của các lớp con trực tiếp sang các tuple chứa các danh sách chuỗi. Phần tử đầu tiên của tuple chứa tên của các đối số đầu vào. Phần tử cuối cùng chứa tên biến của các đầu ra lớp trong mã modeling.
Mở rộng đoạn mã bên dưới để có một ví dụ.
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"]),
    }
```

Mô hình đa phương thức

Các mô hình đa phương thức yêu cầu những thay đổi bổ sung ngoài hướng dẫn đóng góp cho mô hình ngôn ngữ thị giác. Những thay đổi này đảm bảo các đầu vào đa phương thức được xử lý đúng cách.

Lớp ProcessorMixin phải bao gồm các thuộc tính self.image_token và self.image_token_ids. Các token thay thế này chỉ định vị trí của hình ảnh trong đầu vào. Cùng một token này xuất hiện trong prompt đầu vào cho hình ảnh và trong mã mô hình để phân tán các đặc trưng hình ảnh.
Lớp ProcessorMixin phải bao gồm phương thức self._get_num_multimodal_tokens. Phương thức này tính toán số lượng token placeholder cần thiết cho các đầu vào đa phương tiện với kích thước đã cho. Nó trả về một đối tượng MultiModalData. Các placeholder nằm giữa các token <image>, chẳng hạn như token hàng hoặc cột, không được tính là placeholder hình ảnh. Chỉ đếm các token sẽ được thay thế bởi các đặc trưng hình ảnh sau này trong mã mô hình hóa.
Lớp ProcessorMixin phải kiểm tra giá trị của return_mm_token_type_ids và trả về mm_token_type_ids. Điều này cho biết mỗi vị trí là token văn bản (0), token placeholder hình ảnh (1), hay token placeholder video (2). Các chuỗi id loại token đa phương thức phải liên tục, không có khoảng trống giữa các token liên tiếp. Xử lý các token đặc biệt cho token bắt đầu, kết thúc, hàng và cột như các placeholder.

Mở rộng đoạn mã bên dưới để có một ví dụ.

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)

Tài nguyên

Đọc bài đăng trên blog Tích hợp backend Transformers trong vLLM để biết thêm chi tiết.
Đọc bài đăng trên blog Tích hợp backend Transformers trong SGLang để biết thêm chi tiết.

Update on GitHub