كتابة النوى

يشرح هذا الدليل كيفية كتابة نوى تتجاوز مجرد بديل forward بلا حالة. ويتناول القدرتين اللتين يدعمهما واجهة KernelConfig الموسعة:

1. تحويل المعاملات: يتوقع النواة (kernel) الأوزان بتنسيق مختلف عن النموذج الأصلي (على سبيل المثال، معاملات أُعيدت تسميتها أو دُمجت).
2. دمج الوحدات: يستبدل النواة عدة وحدات متجاورة بتنفيذ مدمج واحد.

بالنسبة إلى النوى الأساسية (استبدالات forward بلا حالة دون تغييرات في المعاملات)، راجع توثيق مكتبة kernels.

نمط الفئتين

أي نواة تحمل معلماتها الخاصة تتبع نمط فئتين.

• KernelName: يحتوي على مسار forward فقط. يستخدم مكتبة kernels هذه الفئة لتحويل النموذج إلى نواة، وذلك لأنها لا تسمح بفئات النواة التي تحتفظ بالحالة.
• KernelNameLayout: وهو nn.Module يحتفظ بالمعاملات ويعدّل الوحدة الأصلية (monkey-patches) قبل تحميل نقطة الحفظ. عند التشغيل، يستبدل kernelize دالة forward الخاصة به بدالة forward المستمدة من KernelName. لا تحتاج إلى تعريف forward، إذ يقوم Transformers بحقن واحدة تلقائيًا تحمل التوقيع نفسه الخاص بـKernelName.forward.

تكون تسمية الأسماء صارمة. يجب أن يُسمّى فئة التخطيط {KernelName}Layout ويُعرَّف في الوحدة نفسها التي تُعرَّف فيها KernelName.

تحويل المعلمات

استخدم هذا النمط عندما يتوقع النواة (kernel) الأوزان بأسماء مختلفة أو بشكل مختلف عن نقطة التفتيش (checkpoint) الأصلية للنموذج.

تتمتع فئة 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 مثبّتًا.

قم بتحميل هذا النواة (kernel) عن طريق تمرير اسم المستودع واسم الفئة إلى 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.
3. يعيد تعيين أوزان نقطة الحفظ باستخدام conversion_mapping بحيث تُحمَّل في أسماء المعاملات الجديدة.
4. يستدعي kernelize، والذي يستبدل CustomRMSNormLayout.forward بـ CustomRMSNorm.forward.

دمج الوحدات

استخدم هذا النمط عندما يستبدل نواة (kernel) عدة وحدات متجاورة بتنفيذ مدمج واحد. نظرًا لأن الوحدة المدمجة تجمع بين المعاملات من عدة وحدات أصلية، فإن 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__.