Çekirdekler yazma

Bu kılavuz, durumdan bağımsız bir forward yerinin ötesine geçen çekirdeklerin nasıl yazılacağını açıklar. Genişletilmiş KernelConfig API'sinin desteklediği iki yeteneği kapsar:

1. Parametre dönüşümü: çekirdek, orijinal modelden farklı bir yerleşimde ağırlıklar bekler (örneğin, yeniden adlandırılmış veya birleştirilmiş parametreler).
2. Modül füzyonu: çekirdek, birden fazla bitişik modülü tek bir füzyonlu uygulama ile değiştirir.

Temel çekirdekler (parametre değişikliği olmayan, durum içermeyen forward yerinde kullanımları) için kernels kütüphanesi belgelerine bakın.

İki sınıflı desen

Kendi parametrelerini taşıyan herhangi bir çekirdek, iki sınıflı bir deseni izler.

• KernelName: yalnızca forward geçişini içerir. kernels kütüphanesi, durumlu kernel sınıflarına izin vermediği için modeli kernelize etmek bu sınıfı kullanır.
• KernelNameLayout: parametreleri barındıran ve checkpoint yüklenmeden önce orijinal modülü monkey-patch'leyen bir nn.Module. Çalışma zamanında kernelize, forward metodunu KernelName'in forward metoduyla değiştirir. forward tanımlamanıza gerek yoktur. Transformers, KernelName.forward ile aynı imzaya sahip olan bir metodunu otomatik olarak enjekte eder.

Adlandırma kuralı katıdır. Yerleşim sınıfı {KernelName}Layout olarak adlandırılmalı ve KernelName ile aynı modülde tanımlanmalıdır.

Parametre dönüşümü

Çekirdeğin, orijinal model kontrol noktasındaki ağırlıklardan farklı adlar veya farklı bir biçim altında ağırlıklar beklediği durumlarda bu deseni kullanın.

KernelNameLayout sınıfı, yerini aldığı modülle aynı __init__ imzasına sahiptir ve Transformers'ın checkpoint anahtarlarını yeni parametre adlarına nasıl yeniden eşleyeceğini bildiren bir conversion_mapping sınıf özniteliği tanımlar (daha fazla ayrıntı için Dynamic weight loading bölümüne bakın).

# 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 kütüphanesi, çekirdek sınıflarını bir layers modülü (ya da layers.py ya da layers/__init__.py) üzerinden keşfeder. Hem çekirdek hem de yerleşim sınıflarını doğrudan layers.py içinde tanımlayın ve layers modülünü paketin __init__.py dosyasından dışa aktarın. Transformers, CustomRMSNormLayout sınıfını CustomRMSNorm ile aynı modülde bulur.

layers.py içindeki Transformers'a özgü içe aktarmaları (WeightConverter, Concatenate, WeightRenaming) koruyun, böylece çekirdek, Transformers'ın kurulu olmadığı ortamlarda da içe aktarılabilir.

Bu çekirdeği, depo ve sınıf adını KernelConfig nesnesine ileterek yükleyin. Anahtar, modeldeki orijinal modül sınıfı adını belirtir. Değer ise depodaki KernelName sınıfına (değilse Layout sınıfına değil) işaret eder.

Değer, "owner/repo:ClassName" biçiminde bir dize veya (repo_string, metadata) biçiminde bir demettir. Demet biçiminde, version veya revision (en az biri) ve isteğe bağlı olarak trust_remote_code parametrelerini iletin. Ana çekirdek sürümü için version, etiket, dal veya commit için revision kullanın. Güvenilir kernels-community kuruluşunun dışından yüklemek için trust_remote_code=True olarak ayarlayın. Sadece bir dize verilirse varsayılan olarak version=1 ve trust_remode_code=False kullanılır.

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",
)

Model yüklendiğinde, Transformers:

1. CustomRMSNorm bileşenini repodan yükler ve aynı modülde CustomRMSNormLayout bileşenini arar.
2. Modeldeki her RMSNorm öğesini CustomRMSNormLayout ile monkey-patch'ler.
3. Yeni parametre adlarına yüklenmeleri için conversion_mapping kullanılarak checkpoint ağırlıklarını yeniden eşler.
4. kernelize çağrılır; bu, CustomRMSNormLayout.forward yerine CustomRMSNorm.forward kullanılır.

Modül füzyonu

Bir çekirdeğin, birden fazla bitişik modülü tek bir füzyonlu uygulama ile değiştirdiği durumlarda bu deseni kullanın. Füzyonlu modül, birden fazla orijinal modülün parametrelerini birleştirdiği için KernelNameLayout.__init__, bu modüllerin yapıcı argümanları yerine örneklenmiş alt modülleri alır.

# 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

Modülleri birleştirmek için, düz bir dize yerine KernelConfig anahtarı olarak (class_name, path_pattern) çiftlerinden oluşan bir tuple geçirin. Tüm desenler aynı üst modülü paylaşmalıdır (Transformers, bu üst modüldeki alt öğeleri birleştirir). * joker karakteri, herhangi bir tek yol segmentiyle eşleşir.

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",
)

Model yüklendiğinde, Transformers:

1. RMSNormMLP öğesini repodan yükler ve aynı modülde RMSNormMLPLayout öğesini bulur.
2. model.layers.* içindeki her decoder katmanını eşleştirir ve __init__ metodu RMSNormMLPLayout(post_attention_layernorm, mlp) çağrısını yapan birleşik bir üst sınıf oluşturur.
3. Kalan alt modülü (mlp) üst modülün arayüzünü korumak için nn.Identity() ile değiştirir.
4. conversion_mapping kullanarak checkpoint ağırlıklarını yeniden eşler.
5. kernelize çağrılır; bu, RMSNormMLPLayout.forward yerine RMSNormMLP.forward kullanılır.

Füzyon demetindeki çiftlerin sırası, KernelNameLayout.__init__ fonksiyonuna geçirilen argümanların sırasını belirler.