Создание ядер

В этом руководстве объясняется, как писать ядра, выходящие за рамки замены без состояния forward. Рассматриваются две возможности, поддерживаемые расширенным API KernelConfig:

1. Преобразование параметров: ядро ожидает веса в другом расположении, чем в исходной модели (например, переименованные или объединённые параметры).
2. Слияние модулей: ядро заменяет несколько соседних модулей единой объединённой реализацией.

Для базовых ядер (бессостоятельных замен forward без изменения параметров) см. документацию библиотеки kernels.

Двухклассовый паттерн

Любое ядро, содержащее собственные параметры, следует двухклассовой схеме.

• KernelName: содержит только проход forward. Библиотека kernels использует этот класс для кернализации модели, так как не допускает классов ядер с состоянием.
• KernelNameLayout: nn.Module, который хранит параметры и подменяет исходный модуль (monkey-patch) до загрузки чекпоинта. Во время выполнения kernelize заменяет его forward на forward из KernelName. Вам не нужно определять forward. Transformers автоматически вставляет его с той же сигнатурой, что и у KernelName.forward.

Именование строго регламентировано. Класс раскладки должен называться {KernelName}Layout и определяться в том же модуле, что и KernelName.

Преобразование параметров

Используйте этот шаблон, если ядро ожидает веса с другими именами или в другой форме, чем в исходной контрольной точке модели.

Класс KernelNameLayout имеет ту же сигнатуру __init__, что и заменяемый им модуль, и объявляет атрибут класса conversion_mapping, который указывает Transformers, как переназначать ключи контрольной точки на новые имена параметров (см. Динамическая загрузка весов для подробностей).

# layers.py
try:
    from transformers import Concatenate, WeightConverter
    from transformers.conversion_mapping import WeightRenaming
except ImportError:
    Concatenate = WeightConverter = WeightRenaming = None

import torch
import torch.nn as nn


class CustomRMSNormLayout(nn.Module):
    conversion_mapping = [...]  # rules that remap checkpoint keys to the new parameter names

    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps


class CustomRMSNorm(nn.Module):
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.scale * hidden_states.to(input_dtype)

# __init__.py
from . import layers

Библиотека kernels обнаруживает классы ядер через модуль layers (либо layers.py, либо layers/__init__.py). Определите классы ядра и раскладки непосредственно в layers.py и экспортируйте модуль layers из __init__.py пакета. Transformers находит CustomRMSNormLayout в том же модуле, что и CustomRMSNorm.

Защитите импорты, специфичные для Transformers (WeightConverter, Concatenate, WeightRenaming) в layers.py, чтобы ядро можно было импортировать и в средах, где Transformers не установлен.

Загрузите это ядро, передав имя репозитория и имя класса в KernelConfig. Ключом является исходное имя класса модуля из модели. Значение указывает на класс KernelName (а не Layout) в репозитории.

Значением является либо строка "owner/repo:ClassName", либо кортеж (repo_string, metadata). В форме кортежа передайте version или revision (хотя бы один из них), а также, при необходимости, trust_remote_code. Используйте version для основной версии ядра или revision для тега, ветки или коммита. Установите trust_remote_code=True, чтобы загружать код извне доверенной организации kernels-community. Для простой строки по умолчанию используются version=1 и trust_remode_code=False.

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig({"RMSNorm": ("owner/my-kernel:CustomRMSNorm", {"trust_remote_code": True, "version": 1})})
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto",
)

При загрузке модели Transformers:

1. Загружает CustomRMSNorm из репозитория и ищет CustomRMSNormLayout в том же модуле.
2. Заменяет каждый RMSNorm в модели на CustomRMSNormLayout с помощью monkey-patching.
3. Переназначает веса контрольной точки с помощью conversion_mapping, чтобы они загружались под новыми именами параметров.
4. Вызывает kernelize, который заменяет CustomRMSNormLayout.forward на CustomRMSNorm.forward.

Слияние модулей

Используйте этот шаблон, когда ядро заменяет несколько соседних модулей единым слиянным (fused) представлением. Поскольку слиянный модуль объединяет параметры из нескольких исходных модулей, KernelNameLayout.__init__ получает уже созданные дочерние модули, а не аргументы их конструкторов.

# layers.py
try:
    from transformers import Concatenate, WeightConverter
    from transformers.conversion_mapping import WeightRenaming
except ImportError:
    Concatenate = WeightConverter = WeightRenaming = None

import torch
import torch.nn as nn


class RMSNormMLPLayout(nn.Module):
    conversion_mapping = [...]  # rules that remap checkpoint keys to the fused parameter names

    def __init__(self, norm, mlp):
        super().__init__()
        self.variance_epsilon = norm.variance_epsilon
        self.scale = nn.Parameter(torch.empty_like(norm.weight))
        self.gate_up_proj = nn.Linear(
            mlp.gate_proj.in_features,
            mlp.gate_proj.out_features + mlp.up_proj.out_features,
            bias=mlp.gate_proj.bias is not None,
            device=mlp.gate_proj.weight.device,
            dtype=mlp.gate_proj.weight.dtype,
        )
        self.down_proj = nn.Linear(
            mlp.down_proj.in_features,
            mlp.down_proj.out_features,
            bias=mlp.down_proj.bias is not None,
            device=mlp.down_proj.weight.device,
            dtype=mlp.down_proj.weight.dtype,
        )
        self.act_fn = mlp.act_fn


class RMSNormMLP(nn.Module):
    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        hidden_states = self.scale * hidden_states.to(input_dtype)
        gate, up = self.gate_up_proj(hidden_states).chunk(2, dim=-1)
        return self.down_proj(self.act_fn(gate) * up)

# __init__.py
from . import layers

Чтобы объединить модули, передайте кортеж пар (class_name, path_pattern) в качестве ключа в KernelConfig вместо обычного строкового значения. Все паттерны должны иметь общий родительский модуль (Transformers объединяет дочерние элементы этого родителя). Звёздочка * соответствует любому одиночному сегменту пути.

from transformers import AutoModelForCausalLM, KernelConfig

kernel_config = KernelConfig(
    {
        (
            ("RMSNorm", "model.layers.*.post_attention_layernorm"),
            ("MLP",     "model.layers.*.mlp"),
        ): ("owner/my-kernel:RMSNormMLP", {"trust_remote_code": True, "version": 1}),
    }
)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    use_kernels=True,
    kernel_config=kernel_config,
    device_map="auto",
)

При загрузке модели Transformers:

1. Загружает RMSNormMLP из репозитория и находит RMSNormMLPLayout в том же модуле.
2. Совпадает со всеми слоями декодера в model.layers.* и создаёт объединённый родительский класс, в __init__ которого вызывается RMSNormMLPLayout(post_attention_layernorm, mlp).
3. Заменяет оставшийся дочерний модуль (mlp) на nn.Identity(), чтобы сохранить интерфейс родительского модуля.
4. Переназначает веса контрольной точки с помощью conversion_mapping.
5. Вызывает kernelize, который заменяет RMSNormMLPLayout.forward на RMSNormMLP.forward.

Порядок пар в кортеже слияния определяет порядок аргументов, передаваемых в KernelNameLayout.__init__.