[v1] refactor registry plugin structure and params (#10641)

examples/v1/train_batching_strategy/train_full_fsdp2_batching_normal.yaml

@@ -1,12 +1,8 @@
 model: Qwen/Qwen3-0.6B
 model_class: llm
 
-template: qwen3_nothink
-
-
 kernel_config:
   name: auto
-  include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null
 
 quant_config: null
 

examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_batching.yaml

@@ -1,11 +1,8 @@
 model: Qwen/Qwen3-0.6B
 model_class: llm
 
-template: qwen3_nothink
-
 kernel_config:
   name: auto
-  include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null
 
 quant_config: null
 

examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml

@@ -1,11 +1,8 @@
 model: Qwen/Qwen3-0.6B
 model_class: llm
 
-template: qwen3_nothink
-
 kernel_config:
   name: auto
-  include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null
 
 quant_config: null
 

examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml

@@ -1,11 +1,8 @@
 model: Qwen/Qwen3-0.6B
 model_class: llm
 
-template: qwen3_nothink
-
 kernel_config:
   name: auto
-  include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null
 
 quant_config: null
 

examples/v1/train_freeze/train_freeze_sft.yaml

@@ -1,7 +1,6 @@
 model: Qwen/Qwen3-4B
 model_class: llm
 
-
 # Freeze Configuration
 peft_config:
   name: freeze
@@ -12,7 +11,6 @@ peft_config:
 # Kernel Config
 kernel_config:
   name: auto
-  include_kernels: auto
 
 # FSDP Config
 dist_config:
@@ -25,7 +23,6 @@ train_dataset: data/v1_sft_demo.yaml
 ### training
 output_dir: ./outputs/test_freeze
 micro_batch_size: 1
-global_batch_size: 4
 cutoff_len: 2048
 learning_rate: 2.0e-5
 max_steps: 10

examples/v1/train_full/train_full_deepspeed.yaml

@@ -1,10 +1,8 @@
 model: Qwen/Qwen3-0.6B
 model_class: llm
 
-
 kernel_config:
     name: auto
-    include_kernels: auto
 
 dist_config:
     name: deepspeed

examples/v1/train_full/train_full_fsdp2.yaml

@@ -1,10 +1,8 @@
 model: Qwen/Qwen3-0.6B
 model_class: llm
 
-
 kernel_config:
   name: auto
-  include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null
 
 quant_config: null
 

examples/v1/train_full/train_full_liger_kernel.yaml

@@ -1,11 +1,8 @@
 model: Qwen/Qwen3-0.6B
 model_class: llm
 
-template: qwen3_nothink
-
 kernel_config:
   name: liger_kernel
-  include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null
 
 quant_config: null
 

examples/v1/train_full/train_full_ulysses_cp.yaml

@@ -8,8 +8,9 @@ flash_attn: flash_attention_2
 dist_config:
   name: fsdp2
   dcp_path: null
-  cp_mode: ulysses
+
-  cp_size: 2
+cp_mode: ulysses
+cp_size: 2
 
 ### data
 train_dataset: data/v1_sft_demo.yaml

examples/v1/train_lora/train_lora_dpo.yaml

@@ -1,8 +1,6 @@
 model: Qwen/Qwen3-4B
 model_class: llm
 
-template: qwen3_nothink
-
 # PEFT Configuration
 peft_config:
   name: lora
@@ -14,7 +12,6 @@ peft_config:
 # Kernel Config
 kernel_config:
   name: auto
-  include_kernels: auto
 
 # FSDP Config
 dist_config:

examples/v1/train_lora/train_lora_sft.yaml

@@ -1,7 +1,6 @@
 model: Qwen/Qwen3-4B
 model_class: llm
 
-
 # PEFT Configuration
 peft_config:
   name: lora
@@ -13,7 +12,6 @@ peft_config:
 # Kernel Config
 kernel_config:
   name: auto
-  include_kernels: auto
 
 # FSDP Config
 dist_config:

examples/v1/train_lora/train_lora_sft_rank0.yaml

@@ -1,7 +1,6 @@
 model: Qwen/Qwen3-4B
 model_class: llm
 
-
 # PEFT Configuration
 peft_config:
   name: lora
@@ -13,7 +12,6 @@ peft_config:
 # Kernel Config
 kernel_config:
   name: auto
-  include_kernels: auto
 
 # FSDP Config
 dist_config:

examples/v1/train_qlora/quantization.yaml

@@ -1,7 +1,6 @@
 model: Qwen/Qwen3-0.6B
 model_class: llm
 
-
 # PEFT Configuration
 peft_config:
   name: lora
@@ -13,7 +12,6 @@ peft_config:
 # Kernel Config
 kernel_config:
   name: auto
-  include_kernels: auto
 
 # FSDP Config
 dist_config:

src/llamafactory/v1/accelerator/interface.py

@@ -29,16 +29,20 @@ And data parallelism types:
 from dataclasses import dataclass
 from datetime import timedelta
 from enum import StrEnum
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 from torch.distributed import barrier, destroy_process_group, init_process_group
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 
 from ..utils import logging
-from ..utils.types import DistributedConfig, ProcessGroup, TensorLike
+from ..utils.types import ProcessGroup, TensorLike
 from . import helper
 
 
+if TYPE_CHECKING:
+    from ..config.training_args import TrainingArguments
+
+
 logger = logging.get_logger(__name__)
 
 
@@ -128,12 +132,13 @@ class DistributedInterface:
 
         return cls._instance
 
-    def __init__(self, config: DistributedConfig | None = None) -> None:
+    def __init__(
+        self,
+        training_args: "TrainingArguments | None" = None,
+    ) -> None:
         if self._initialized:
             return
 
-        self.dist_config = config
-
         helper.set_device_index()
         self._is_distributed = helper.is_distributed()
         self._rank = helper.get_rank()
@@ -143,17 +148,17 @@ class DistributedInterface:
         self.current_device = helper.get_current_device()
         self.device_count = helper.get_device_count()
 
-        if config is None:
+        if training_args is None:
             self.strategy = DistributedStrategy()
             timeout = 18000
         else:
             self.strategy = DistributedStrategy(
-                mp_replicate_size=config.get("mp_replicate_size", 1),
+                mp_replicate_size=training_args.mp_replicate_size,
-                mp_shard_size=config.get("mp_shard_size", None),
+                mp_shard_size=training_args.mp_shard_size,
-                dp_size=config.get("dp_size", None),
+                dp_size=training_args.dp_size,
-                cp_size=config.get("cp_size", 1),
+                cp_size=training_args.cp_size,
             )
-            timeout = config.get("timeout", 18000)
+            timeout = training_args.dist_timeout
 
         if self._is_distributed:
             init_process_group(timeout=timedelta(seconds=timeout), backend=helper.get_process_group_backend())

src/llamafactory/v1/config/training_args.py

@@ -76,7 +76,31 @@ class TrainingArguments:
     )
     dist_config: PluginConfig | None = field(
         default=None,
-        metadata={"help": "Distribution configuration for training."},
+        metadata={"help": "Distributed backend plugin configuration."},
+    )
+    dp_size: int | None = field(
+        default=None,
+        metadata={"help": "Data parallel size, default to world_size // cp_size."},
+    )
+    cp_size: int = field(
+        default=1,
+        metadata={"help": "Context parallel size."},
+    )
+    cp_mode: str = field(
+        default="ulysses",
+        metadata={"help": "Context parallel implementation."},
+    )
+    mp_replicate_size: int = field(
+        default=1,
+        metadata={"help": "Model parallel replicate size."},
+    )
+    mp_shard_size: int | None = field(
+        default=None,
+        metadata={"help": "Model parallel shard size, default to world_size // mp_replicate_size."},
+    )
+    dist_timeout: int = field(
+        default=18000,
+        metadata={"help": "Distributed process group initialization timeout in seconds."},
     )
     optim_config: PluginConfig | None = field(
         default=None,
@@ -149,6 +173,12 @@ class TrainingArguments:
         self.dist_config = get_plugin_config(self.dist_config)
         self.optim_config = get_plugin_config(self.optim_config)
         self.lr_scheduler_config = get_plugin_config(self.lr_scheduler_config)
+        try:
+            from ..plugins.model_plugins.deepspeed_utils import register_deepspeed_dist_config
+
+            register_deepspeed_dist_config(self.dist_config)
+        except ImportError:
+            pass
 
         # The optimizer learning rate has a single source of truth: ``learning_rate``.
         # Propagate it into ``optim_config["lr"]`` so optimizer plugins (e.g. Muon) pick it up

src/llamafactory/v1/core/base_trainer.py

@@ -94,9 +94,12 @@ class BaseTrainer:
         dist_name = self.args.dist_config.name if self.args.dist_config is not None else None
 
         if dist_name == "deepspeed":
-            from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin
+            if self.args.cp_size > 1:
+                raise ValueError("Context parallelism currently requires `dist_config.name: fsdp2`.")
 
-            self._deepspeed_engine = DistributedPlugin("deepspeed")(
+            from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
+
+            self._deepspeed_engine = DistributedPlugin("deepspeed").shard_model(
                 self.model,
                 self.args.dist_config,
                 num_micro_batch=self.train_batch_generator.num_micro_batch,
@@ -139,7 +142,7 @@ class BaseTrainer:
         self.state.global_step = self.global_step
         self.state.epoch = self._resume_epoch
 
-        if self.args.dist_config is not None and self.args.dist_config.get("cp_size", 1) > 1:
+        if self.args.cp_size > 1:
             # qwen3.5 is not supported because of the different attention implementation, which will be supported in the future.
             if model.config.model_type == "qwen3_5":
                 raise RuntimeError(
@@ -152,7 +155,7 @@ class BaseTrainer:
                     "Sequence parallelism requires flash attention. Please set `flash_attn: flash_attention_2`."
                 )
 
-            SequenceParallelModelPlugin(self.args.dist_config.get("cp_mode", "ulysses"))(model, self.args.dist_config)
+            SequenceParallelModelPlugin(self.args.cp_mode)(model, self.args.cp_size)
 
     def _create_batch_generator(self) -> None:
         if (
@@ -183,9 +186,9 @@ class BaseTrainer:
                 device_ids = None if self.device.type == "cpu" else [self.device.index]
                 self.model = DDP(self.model, device_ids=device_ids)
         else:
-            from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin
+            from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
 
-            self.model = DistributedPlugin(self.args.dist_config.name)(
+            self.model = DistributedPlugin(self.args.dist_config.name).shard_model(
                 self.model,
                 self.args.dist_config,
                 bf16=self.args.bf16,
@@ -256,7 +259,7 @@ class BaseTrainer:
                 step_valid_tokens = DistributedInterface().all_reduce(step_valid_tokens, op=ReduceOp.SUM)
                 num_micro = len(micro_batches)
                 for i, micro_batch in enumerate(micro_batches):
-                    if self.args.dist_config and self.args.dist_config.get("cp_size", 1) > 1:
+                    if self.args.cp_size > 1:
                         from ..plugins.model_plugins.parallelization.sequence_parallel import (
                             SequenceParallelLossPlugin,
                         )
@@ -352,7 +355,7 @@ class BaseTrainer:
     def save_model(self) -> None:
         """Save the model."""
         if self.args.dist_config is not None and self.args.dist_config.name in ("deepspeed", "fsdp2"):
-            from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin
+            from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
 
             DistributedPlugin(self.args.dist_config.name).save_model(
                 self.model, self.args.output_dir, self.renderer.processor

src/llamafactory/v1/core/data_engine.py

@@ -137,9 +137,7 @@ class DataEngine(Dataset):
         messages = sample.get("messages")
         if not messages:
             return [None]
-        cuts = [
+        cuts = [i + 1 for i, m in enumerate(messages) if m["role"] == "assistant" and m.get("loss_weight", 1.0) > 1e-6]
-            i + 1 for i, m in enumerate(messages) if m["role"] == "assistant" and m.get("loss_weight", 1.0) > 1e-6
-        ]
         return cuts or [None]
 
     def _convert_data_sample(self, raw_sample: dict[str, Any], dataset_name: str) -> Sample:

src/llamafactory/v1/core/model_engine.py

@@ -69,24 +69,25 @@ class ModelEngine:
         """Model configuration."""
         self.renderer = Renderer(self.processor)
         """Renderer."""
-        self._dist_config = DistributedInterface().dist_config
-        self._deepspeed_zero3_plugin = None
         self._deepspeed_zero3_enabled = False
 
-        if self.is_train and self._dist_config is not None and self._dist_config.get("name") == "deepspeed":
+        try:
             from ..plugins.model_plugins.deepspeed_utils import (
+                is_deepspeed_zero3_enabled,
                 setup_deepspeed_zero3_model_loading,
                 teardown_deepspeed_zero3_model_loading,
             )
 
+            self._deepspeed_zero3_enabled = self.is_train and is_deepspeed_zero3_enabled()
+        except ImportError:
+            pass
+
+        if self._deepspeed_zero3_enabled:
+            plugin = setup_deepspeed_zero3_model_loading()
             try:
-                self._deepspeed_zero3_plugin = setup_deepspeed_zero3_model_loading(self.is_train, self._dist_config)
-                self._deepspeed_zero3_enabled = self._deepspeed_zero3_plugin is not None
                 self.model = self._init_model()
             finally:
-                teardown_deepspeed_zero3_model_loading(self._deepspeed_zero3_plugin)
+                teardown_deepspeed_zero3_model_loading(plugin)
-                self._deepspeed_zero3_plugin = None
-                self._deepspeed_zero3_enabled = False
         else:
             self.model = self._init_model()
 
@@ -142,9 +143,7 @@ class ModelEngine:
 
             init_kwargs = QuantizationPlugin(self.args.quant_config.name)(
                 init_kwargs=init_kwargs,
-                config=self.model_config,
+                quant_config=self.args.quant_config,
-                tokenizer=self.processor,
-                model_args=self.args,
                 is_trainable=self.is_train,
             )
 
@@ -200,17 +199,16 @@ class ModelEngine:
 
             from ..plugins.model_plugins.peft import PeftPlugin
 
-            model = PeftPlugin(self.args.peft_config.name)(model, self.args.peft_config, self.is_train)
+            model = PeftPlugin(self.args.peft_config.name)(
+                model,
+                peft_config=self.args.peft_config,
+                is_train=self.is_train,
+            )
 
         if self.args.kernel_config is not None:
-            from ..plugins.model_plugins.kernels.interface import KernelPlugin
+            from ..plugins.model_plugins.kernels.interface import apply_kernels
 
-            kernel_config = self.args.kernel_config
+            model = apply_kernels(model, self.args.kernel_config, require_logits=self.is_train)
-            kernel_kwargs: dict = {"model": model, "include_kernels": kernel_config.get("include_kernels")}
-            if kernel_config.name == "liger_kernel":
-                # Fused linear CE omits logits; SFT stage needs logits for loss_weights.
-                kernel_kwargs["require_logits"] = self.is_train
-            model = KernelPlugin(kernel_config.name)(**kernel_kwargs)
 
         return model
 

src/llamafactory/v1/core/rendering/format.py

@@ -52,12 +52,8 @@ def _to_hf_messages(messages: list[Message]) -> list[dict]:
                 except json.JSONDecodeError as e:
                     raise ValueError(f"tool_call value is not valid JSON: {content['value']!r}") from e
                 if not isinstance(tc, dict) or "name" not in tc or "arguments" not in tc:
-                    raise ValueError(
+                    raise ValueError(f"tool_call must be a JSON object with 'name' and 'arguments' keys, got {tc!r}")
-                        f"tool_call must be a JSON object with 'name' and 'arguments' keys, got {tc!r}"
+                tool_calls.append({"type": "function", "function": {"name": tc["name"], "arguments": tc["arguments"]}})
-                    )
-                tool_calls.append(
-                    {"type": "function", "function": {"name": tc["name"], "arguments": tc["arguments"]}}
-                )
         hf_msg = {"role": message["role"], "content": text}
 
         if tool_calls:
@@ -67,4 +63,3 @@ def _to_hf_messages(messages: list[Message]) -> list[dict]:
 
         hf_messages.append(hf_msg)
     return hf_messages
-

src/llamafactory/v1/core/rendering/rendering.py

@@ -50,7 +50,6 @@ def _render_messages(
 
     Note: ``position_ids`` are not produced here; ``process_samples`` assigns a 1-based range.
     """
-
     tokenizer = get_tokenizer(processor)
     if not getattr(tokenizer, "chat_template", None):
         tokenizer.chat_template = _FALLBACK_CHATML_JINJA
@@ -73,6 +72,7 @@ def _render_messages(
             tools_parsed = [tools_parsed]
     if not is_generate and hf_messages and hf_messages[-1].get("reasoning_content"):
         kwargs["enable_thinking"] = True
+
     def _encode(msgs: list[dict], add_generation_prompt: bool) -> list[int]:
         text = tokenizer.apply_chat_template(
             msgs, tokenize=False, add_generation_prompt=add_generation_prompt, tools=tools_parsed, **kwargs
@@ -150,13 +150,7 @@ class Renderer:
         Returns:
             ModelInput with input_ids, attention_mask, labels, and loss_weights.
         """
-        return _render_messages(
+        return _render_messages(self.processor, messages, tools, is_generate, **kwargs)
-            self.processor,
-            messages,
-            tools,
-            is_generate,
-            **kwargs
-        )
 
     def process_samples(self, samples: list[Sample]) -> list[ModelInput]:
         """Process samples to model input.

src/llamafactory/v1/core/utils/checkpoint.py

@@ -251,7 +251,7 @@ class TrainingCheckpointCoordinator:
             )
 
         if self._dist_name in ("fsdp2", "deepspeed"):
-            from ...plugins.trainer_plugins.distributed.hub import DistributedPlugin
+            from ...plugins.trainer_plugins.distributed.interface import DistributedPlugin
 
             DistributedPlugin(self._dist_name).save_checkpoint(
                 self._t.model,
@@ -307,7 +307,7 @@ class TrainingCheckpointCoordinator:
         self._t._resume_epoch = metadata["epoch"]
 
         if self._dist_name in ("fsdp2", "deepspeed"):
-            from ...plugins.trainer_plugins.distributed.hub import DistributedPlugin
+            from ...plugins.trainer_plugins.distributed.interface import DistributedPlugin
 
             DistributedPlugin(self._dist_name).load_checkpoint(
                 self._t.model,

src/llamafactory/v1/plugins/model_plugins/deepspeed_utils.py

@@ -14,9 +14,32 @@
 
 import json
 from copy import deepcopy
+from functools import lru_cache
 from typing import Any
 
 
+_registered_dist_config: Any | None = None
+
+
+def register_deepspeed_dist_config(dist_config: Any | None) -> None:
+    """Register backend config before model loading without involving the accelerator."""
+    global _registered_dist_config
+    _registered_dist_config = dist_config
+    is_deepspeed_zero3_enabled.cache_clear()
+
+
+@lru_cache(maxsize=1)
+def is_deepspeed_zero3_enabled() -> bool:
+    dist_config = _registered_dist_config
+    if dist_config is None or getattr(dist_config, "name", None) != "deepspeed":
+        return False
+
+    config_file = dist_config.get("config_file")
+    if not config_file:
+        return False
+    return _load_deepspeed_config(config_file).get("zero_optimization", {}).get("stage") == 3
+
+
 def _normalize_precision_enabled(value: Any) -> bool | str:
     if isinstance(value, str):
         value_lower = value.lower()
@@ -69,18 +92,19 @@ def _load_deepspeed_config(config_file: str) -> dict[str, Any]:
         return json.load(f)
 
 
-def setup_deepspeed_zero3_model_loading(is_train: bool, dist_config: dict[str, Any] | None):
+def setup_deepspeed_zero3_model_loading():
-    """Enable transformers' ZeRO-3-aware model loading for the current thread."""
+    """Enable ZeRO-3-aware model loading for the registered backend config."""
-    config_file = dist_config.get("config_file")
+    dist_config = _registered_dist_config
+    config_file = dist_config.get("config_file") if dist_config is not None else None
     if not config_file:
         raise ValueError("DeepSpeed config_file is required in dist_config")
 
     from accelerate.utils import DeepSpeedPlugin
 
     try:
-        from transformers.integrations import is_deepspeed_zero3_enabled
+        from transformers.integrations import is_deepspeed_zero3_enabled as _hf_is_deepspeed_zero3_enabled
     except ImportError:
-        from transformers.deepspeed import is_deepspeed_zero3_enabled
+        from transformers.deepspeed import is_deepspeed_zero3_enabled as _hf_is_deepspeed_zero3_enabled
 
     # DeepSpeed configs often use "auto" placeholders that only make sense once
     # we know the current runtime batch settings and precision mode.
@@ -109,7 +133,7 @@ def setup_deepspeed_zero3_model_loading(is_train: bool, dist_config: dict[str, A
     plugin.set_mixed_precision(mixed_precision)
     plugin.set_deepspeed_weakref()
 
-    if not is_deepspeed_zero3_enabled():
+    if not _hf_is_deepspeed_zero3_enabled():
         raise RuntimeError(
             "DeepSpeed ZeRO-3 model-loading bootstrap failed: transformers still reports zero3 disabled "
             "after constructing HfDeepSpeedConfig. This usually means the runtime is using a different transformers "

src/llamafactory/v1/plugins/model_plugins/kernels/base.py

@@ -12,76 +12,40 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""The definition of base kernel class.
-
-Init Phase:
-1. Define base kernel class.
-2. Define abstract methods.
-
-"""
-
 from abc import ABC, abstractmethod
-from typing import Any
 
-from ....accelerator.helper import DeviceType, get_current_accelerator
+from ....utils.plugin import BasePlugin, ensure_methods_implemented
 from ....utils.types import HFModel
 
 
+class KernelPlugin(BasePlugin):
+    """Plugin family for model kernel optimization classes."""
+
+
 class BaseKernel(ABC):
-    r"""Base class for all kernel implementations.
+    """Template base for concrete kernel implementations."""
 
-    Subclasses must implement the abstract methods and define the required class attributes.
+    def __init_subclass__(cls, **kwargs) -> None:
-    """
+        super().__init_subclass__(**kwargs)
+        ensure_methods_implemented(cls)
 
-    _kernel_id: Any = ""  # kernel ID, any hashable value to identify a kernel implementation
+    @staticmethod
-    _device: list[DeviceType] = [DeviceType.CPU]  # "cuda", "npu", "cpu", etc.
-
-    @classmethod
-    def get_kernel_id(cls) -> str:
-        """Returns the unique identifier for the kernel."""
-        return cls._kernel_id
-
-    @classmethod
-    def get_device(cls) -> list[DeviceType]:
-        """Returns the device type list associated with the kernel (e.g., ["cuda", "npu", "cpu"])."""
-        return cls._device
-
-    @classmethod
-    def check_deps(cls) -> bool:
-        """Checks if the required dependencies for the kernel are available.
-
-        Returns:
-            bool: ``True`` if dependencies are met, ``False`` otherwise.
-
-        .. note::
-            In explicit mode, if a user specifies an implementation but this check fails,
-            it should raise an error instead of silently switching.
-            Kernels can override this method to implement custom dependency checks.
-        """
-        if get_current_accelerator().type not in cls._device:
-            return False
-        return True
-
-    @classmethod
     @abstractmethod
+    def check_device() -> None: ...
+
+    @staticmethod
+    def check_deps() -> None:
+        pass
+
+    @classmethod
     def apply(cls, **kwargs) -> HFModel:
-        """Applies the kernel optimization to the model.
+        cls.check_device()
+        cls.check_deps()
+        if kwargs.get("model") is None:
+            raise ValueError(f"HFModel instance is required for {cls.__name__}.")
 
-        Args:
+        return cls._apply(**kwargs)
-            **kwargs: Arbitrary keyword arguments, usually containing the model instance and the kernel configuration.
 
-        Returns:
+    @staticmethod
-            HFModel: The model with the kernel applied.
+    @abstractmethod
-
+    def _apply(**kwargs) -> HFModel: ...
-        Raises:
-            RuntimeError: If the kernel dependencies are not met.
-            NotImplementedError: If the method is not implemented by the subclass.
-
-        Example:
-            >>> from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernel
-            >>> model = HFModel(config=config)
-            >>> model = apply_kernel(model=model, kernel_id="npu_fused_moe")
-        """
-        if not cls.check_deps():
-            raise RuntimeError(f"{cls.__name__} is not available but {cls.__name__} kernel was called.")
-        raise NotImplementedError

src/llamafactory/v1/plugins/model_plugins/kernels/interface.py

@@ -12,174 +12,63 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""The definition of kernel interface.
+from typing import Any
 
-Init Phase:
+from ....accelerator.helper import DeviceType, get_current_accelerator
-1. Scan all kernels.
-2. Register default kernels.
-3. Define kernel plugin.
-
-"""
-
-import importlib
-from pathlib import Path
-
-from ....utils import logging
-from ....utils.plugin import BasePlugin
 from ....utils.types import HFModel
-from .registry import Registry
+from .base import KernelPlugin
+
+# Import built-in implementations so their class decorators populate the registry.
+from .liger_kernel_ops import LigerKernel  # noqa: F401
+from .ops.mlp.cuda_fused_moe import CudaFusedMoEKernel  # noqa: F401
+from .ops.mlp.npu_fused_moe import NpuFusedMoEKernel  # noqa: F401
+from .ops.mlp.npu_swiglu import NpuSwiGluKernel  # noqa: F401
+from .ops.rms_norm.npu_rms_norm import NpuRMSNormKernel  # noqa: F401
+from .ops.rope.npu_rope import NpuRoPEKernel  # noqa: F401
 
 
-logger = logging.get_logger(__name__)
+_AUTO_KERNELS = {
+    DeviceType.NPU: ("npu_fused_moe", "npu_fused_rmsnorm", "npu_fused_rope", "npu_fused_swiglu"),
+}
 
 
-def scan_all_kernels():
+def _apply_auto_kernels(model: HFModel, **kwargs) -> HFModel:
-    """Scan all kernels in the ``ops`` directory.
+    device_type = get_current_accelerator().type
-
+    for kernel_name in _AUTO_KERNELS.get(device_type, ()):
-    Scans the ``ops`` directory for all ``.py`` files and attempts to import them.
+        model = KernelPlugin(kernel_name).apply(model=model, **kwargs)
-    Importing triggers the :func:`~registry.register_kernel` decorator, which automatically registers the kernels.
-
-    Returns:
-        dict[str, type[BaseKernel]]: A dictionary of registered kernels.
-
-    .. note::
-        This function assumes that the ``ops`` directory is located in the same directory as this file.
-        It recursively searches for ``.py`` files and constructs the module path for import.
-    """
-    ops_path = Path(__file__).parent / "ops"
-
-    if not ops_path.exists():
-        return
-
-    base_package = __package__
-
-    for file_path in ops_path.rglob("*.py"):
-        if file_path.name == "__init__.py":
-            continue
-
-        # calculate the relative path:
-        # file_path = .../kernels_v2/ops/mlp/npu_swiglu.py
-        # rel_path  = ops/mlp/npu_swiglu.py
-        rel_path = file_path.relative_to(Path(__file__).parent)
-
-        # build module path:
-        module_name = ".".join(rel_path.parts)[:-3]
-        full_module_name = f"{base_package}.{module_name}"
-
-        try:
-            importlib.import_module(full_module_name)
-        except Exception as e:
-            logger.warning(f"[Kernel Registry] Failed to import {full_module_name} when loading kernels: {e}")
-
-    return Registry.get_registered_kernels()
-
-
-default_kernels = scan_all_kernels()
-
-
-def get_default_kernels():
-    """Get a list of default registered kernel IDs.
-
-    Returns:
-        list[str]: List of kernel IDs.
-    """
-    return list(default_kernels.keys())
-
-
-def apply_kernel(kernel_id: str, **kwargs):
-    """Applies a specific kernel to the model.
-
-    Args:
-        kernel_id (str): The ID of the kernel to apply.
-        **kwargs: Keyword arguments passed to the kernel application function.
-                  Typically includes the model instance.
-
-    Returns:
-        HFModel: The model with applied kernel.
-    """
-    kernel = default_kernels.get(kernel_id)
-    if kernel is None:
-        raise ValueError(f"Kernel {kernel_id} not found")
-
-    kernel.apply(**kwargs)
-
-
-class KernelPlugin(BasePlugin):
-    """Plugin for managing kernel optimizations."""
-
-    pass
-
-
-@KernelPlugin("auto").register()
-def apply_default_kernels(model: HFModel, include_kernels: str = None) -> HFModel:
-    """Applies all default registered kernels to the model.
-
-    Args:
-        model (HFModel): The model instance to apply kernels to.
-        include_kernels (str, optional): Comma-separated list of kernel IDs to apply.
-                                         If "auto" or True, applies all default kernels.
-                                         If None or False, no kernels are applied.
-                                         Defaults to None.
-
-    Returns:
-        HFModel: The model with applied kernels.
-    """
-    if not include_kernels:
-        return model
-    elif include_kernels == "auto" or include_kernels is True:
-        use_kernels = default_kernels.keys()
-    else:
-        use_kernels = include_kernels.split(",")  # "kernel_id1,kernel_id2,kernel_id3"
-
-    for kernel in use_kernels:
-        if kernel not in default_kernels:
-            raise ValueError(f"Kernel {kernel} not found")
-
-        apply_kernel(kernel, model=model)
 
     return model
 
 
-@KernelPlugin("liger_kernel").register()
+def apply_kernels(model: HFModel, config: dict[str, Any], require_logits: bool = False) -> HFModel:
-def apply_liger_kernels(
+    """Apply the comma-separated kernel names selected by ``kernel_config.name``."""
-    model: HFModel,
+    kernel_names = config.get("name")
-    include_kernels: str = None,
+    if not isinstance(kernel_names, str):
-    require_logits: bool = False,
+        raise TypeError("kernel_config.name must be a string.")
-) -> HFModel:
-    """Applies Liger kernel to the model.
 
-    Args:
+    names = [name.strip() for name in kernel_names.split(",") if name.strip()]
-        model (HFModel): The model instance to apply kernels to.
+    if not names:
-        include_kernels (str, optional): If ``"auto"`` or ``True``, apply Liger with
+        raise ValueError("kernel_config.name must contain at least one kernel name.")
-                                         library defaults. If a comma-separated list (e.g.
-                                         ``rope,rms_norm``), enable only those ops; names match
-                                         ``apply_liger_kernel_to_*`` kwargs: ``rope``, ``rms_norm``,
-                                         ``swiglu``, ``cross_entropy``, ``fused_linear_cross_entropy``.
-                                         If ``None`` or ``False``, do nothing. Defaults to ``None``.
-        require_logits (bool, optional): When true, disables ``fused_linear_cross_entropy`` in favor
-                                         of non-fused CE so the forward pass returns ``logits``. Needed
-                                         for trainers that compute weighted loss from logits (e.g. v1
-                                         SFT with ``loss_weights``). Defaults to ``False`` (fused CE
-                                         when supported). The v1 ``run_sft`` entrypoint sets
-                                         ``require_logits`` to true for ``liger_kernel`` when the key
-                                         is omitted so SFT weighted loss keeps working.
 
-    Returns:
+    for name in names:
-        HFModel: The model with Liger kernel applied.
+        if name == "auto":
-    """
+            model = _apply_auto_kernels(model=model, config=config, require_logits=require_logits)
-    if not include_kernels:
+        else:
-        return model
+            model = KernelPlugin(name).apply(model=model, config=config, require_logits=require_logits)
-    if include_kernels == "auto" or include_kernels is True:
-        use_kernels = "auto"
-    else:
-        use_kernels = [k.strip() for k in include_kernels.split(",") if k.strip()]
-        if not use_kernels:
-            return model
 
-    try:
+    return model
-        from .liger_kernel_ops import LigerKernel
+
-    except ImportError as e:
+
-        logger.warning_rank0(f"[Kernel] Failed to import liger_kernel ops, skip. Error: {e}")
+def apply_v1_kernels(model: HFModel, use_v1_kernels: bool) -> HFModel:
+    """Apply v1 automatic kernels for the transitional v0 ``use_v1_kernels`` option."""
+    if not use_v1_kernels:
         return model
 
-    return LigerKernel.apply(use_kernels=use_kernels, model=model, require_logits=require_logits)
+    return apply_kernels(model, {"name": "auto"})
+
+
+def apply_kernel(kernel_id: str, **kwargs) -> HFModel:
+    if kernel_id == "auto":
+        return _apply_auto_kernels(**kwargs)
+
+    return KernelPlugin(kernel_id).apply(**kwargs)

src/llamafactory/v1/plugins/model_plugins/kernels/liger_kernel_ops.py

@@ -25,7 +25,7 @@ import inspect
 from ....accelerator.helper import DeviceType, get_current_accelerator
 from ....utils.logging import get_logger
 from ....utils.types import HFModel
-from .base import BaseKernel
+from .base import BaseKernel, KernelPlugin
 
 
 logger = get_logger(__name__)
@@ -41,26 +41,26 @@ _LIGER_FN_BY_MODEL_TYPE: dict[str, str] = {
 }
 
 
+@KernelPlugin("liger_kernel").register()
 class LigerKernel(BaseKernel):
     """Liger Kernel for optimized model training."""
 
-    _device = [DeviceType.CUDA, DeviceType.NPU]
+    @staticmethod
+    def check_device() -> None:
+        current = get_current_accelerator().type
+        if current not in (DeviceType.CUDA, DeviceType.NPU):
+            raise RuntimeError(f"LigerKernel requires CUDA or NPU, current accelerator is {current}.")
 
-    @classmethod
+    @staticmethod
-    def check_deps(cls) -> bool:
+    def check_deps() -> None:
         """Checks if the required dependencies for the kernel are available."""
         try:
             import liger_kernel  # noqa: F401
-
-            return super().check_deps()
         except ImportError:
-            logger.warning_rank0(
+            raise RuntimeError("Liger kernel is not installed.") from None
-                "Liger kernel is not installed, the kernel_config liger_kernel will be ignored. Please install it from https://github.com/linkedin/Liger-Kernel."
-            )
-            return False
 
-    @classmethod
+    @staticmethod
-    def apply(cls, **kwargs) -> "HFModel":
+    def _apply(**kwargs) -> "HFModel":
         """Applies the Liger kernel to the model.
 
         Args:
@@ -78,16 +78,12 @@ class LigerKernel(BaseKernel):
             RuntimeError: If dependencies are not met.
         """
         model = kwargs.get("model")
-        use_kernels = kwargs.get("use_kernels", None)
+        config = kwargs.get("config")
-        if model is None:
+        use_kernels = kwargs.get("use_kernels", "auto")
-            raise ValueError(f"HFModel instance is required for {cls.__name__}.")
-
-        if not cls.check_deps():
-            raise RuntimeError(
-                f"current device is not supported by liger_kernel. Current device is {get_current_accelerator().type}, supported devices are {cls.get_device()}"
-            )
 
         require_logits = kwargs.get("require_logits", False)
+        if config is not None:
+            require_logits = config.get("require_logits", require_logits)
 
         model_type = getattr(model.config, "model_type", None)
 

src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/cuda_fused_moe.py

@@ -27,16 +27,22 @@ import types
 import torch
 import torch.nn.functional as F
 
-from ......accelerator.helper import DeviceType
+from ......accelerator.helper import DeviceType, get_current_accelerator
 from ......utils.types import HFModel
-from ...base import BaseKernel
+from ...base import BaseKernel, KernelPlugin
-from ...registry import register_kernel
+
-from .triton_grouped_gemm import (
+
-    group_gemm_same_mn,
+try:
-    group_gemm_same_nk,
+    from .triton_grouped_gemm import (
-    moe_gather,
+        group_gemm_same_mn,
-    moe_scatter,
+        group_gemm_same_nk,
-)
+        moe_gather,
+        moe_scatter,
+    )
+except ImportError as exc:
+    _TRITON_IMPORT_ERROR = exc
+else:
+    _TRITON_IMPORT_ERROR = None
 
 
 logger = logging.getLogger(__name__)
@@ -351,7 +357,7 @@ _TRITON_MOE_MAPPING: dict[str, dict[str, object]] = {
 # ---------------------------------------------------------------------------
 
 
-@register_kernel
+@KernelPlugin("cuda_fused_moe").register()
 class CudaFusedMoEKernel(BaseKernel):
     """Pure-Triton fused MoE kernel for NVIDIA CUDA GPUs.
 
@@ -362,30 +368,20 @@ class CudaFusedMoEKernel(BaseKernel):
     Requires: CUDA GPU + Triton
     """
 
-    _kernel_id = "cuda_fused_moe"
+    @staticmethod
-    _device = DeviceType.CUDA
+    def check_device() -> None:
+        current = get_current_accelerator().type
+        if current != DeviceType.CUDA:
+            raise RuntimeError(f"CudaFusedMoEKernel requires CUDA, current accelerator is {current}.")
 
-    @classmethod
+    @staticmethod
-    def check_deps(cls) -> bool:
+    def check_deps() -> None:
-        if not super().check_deps():
+        if _TRITON_IMPORT_ERROR is not None:
-            return False
+            raise RuntimeError("cuda_fused_moe requires Triton.") from _TRITON_IMPORT_ERROR
-        try:
-            import triton  # noqa: F401
 
-            return True
+    @staticmethod
-        except ImportError:
+    def _apply(**kwargs) -> HFModel:
-            logger.info("cuda_fused_moe: Triton not available, kernel disabled.")
-            return False
-
-    @classmethod
-    def apply(cls, **kwargs) -> HFModel:
         model = kwargs.get("model")
-        if model is None:
-            raise ValueError(f"HFModel instance is required for {cls.__name__}.")
-
-        if not cls.check_deps():
-            logger.warning("cuda_fused_moe: Dependencies not met. Skipping kernel application.")
-            return model
 
         archs = getattr(model.config, "architectures", None) or []
         target_mapping = None

src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_fused_moe.py

@@ -32,11 +32,10 @@ try:
 except ImportError:
     pass
 
-from ......accelerator.helper import DeviceType
+from ......accelerator.helper import DeviceType, get_current_accelerator
 from ......utils.packages import is_transformers_version_greater_than
 from ......utils.types import HFModel
-from ...base import BaseKernel
+from ...base import BaseKernel, KernelPlugin
-from ...registry import register_kernel
 
 
 class GmmFunction(torch.autograd.Function):
@@ -334,15 +333,18 @@ else:
     }
 
 
-@register_kernel
+@KernelPlugin("npu_fused_moe").register()
 class NpuFusedMoEKernel(BaseKernel):
     """NPU Fused MoE Kernel implementation."""
 
-    _kernel_id = "npu_fused_moe"
+    @staticmethod
-    _device = DeviceType.NPU
+    def check_device() -> None:
+        current = get_current_accelerator().type
+        if current != DeviceType.NPU:
+            raise RuntimeError(f"NpuFusedMoEKernel requires NPU, current accelerator is {current}.")
 
-    @classmethod
+    @staticmethod
-    def apply(cls, **kwargs) -> HFModel:
+    def _apply(**kwargs) -> HFModel:
         """Applies the NPU fused MoE kernel to the model.
 
         Args:
@@ -356,11 +358,6 @@ class NpuFusedMoEKernel(BaseKernel):
             RuntimeError: If dependencies are not met.
         """
         model = kwargs.get("model", None)
-        if model is None:
-            raise ValueError(f"HFModel instance is required for {cls.__name__}.")
-
-        if not cls.check_deps():
-            raise RuntimeError("torch_npu is not available but NpuMoEFusedMoEKernel was called.")
 
         archs = getattr(model.config, "architectures", None) or []
         target_moe_mapping = None

src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_swiglu.py

@@ -25,10 +25,9 @@ import types
 
 import torch
 
-from ......accelerator.helper import DeviceType
+from ......accelerator.helper import DeviceType, get_current_accelerator
 from ......utils.types import HFModel
-from ...base import BaseKernel
+from ...base import BaseKernel, KernelPlugin
-from ...registry import register_kernel
 
 
 try:
@@ -86,7 +85,7 @@ def _npu_swiglu_gemma3ntext_forward(self, hidden_states):
     return down_proj
 
 
-@register_kernel
+@KernelPlugin("npu_fused_swiglu").register()
 class NpuSwiGluKernel(BaseKernel):
     """NPU Kernel for fused SwiGLU activation."""
 
@@ -121,11 +120,14 @@ class NpuSwiGluKernel(BaseKernel):
         }
     )
 
-    _kernel_id = "npu_fused_swiglu"
+    @staticmethod
-    _device = DeviceType.NPU
+    def check_device() -> None:
+        current = get_current_accelerator().type
+        if current != DeviceType.NPU:
+            raise RuntimeError(f"NpuSwiGluKernel requires NPU, current accelerator is {current}.")
 
-    @classmethod
+    @staticmethod
-    def apply(cls, **kwargs) -> "HFModel":
+    def _apply(**kwargs) -> "HFModel":
         """Applies the NPU fused SwiGLU kernel to the model.
 
         Args:
@@ -139,11 +141,6 @@ class NpuSwiGluKernel(BaseKernel):
             RuntimeError: If dependencies are not met.
         """
         model = kwargs.get("model", None)
-        if model is None:
-            raise ValueError(f"HFModel instance is required for {cls.__name__}.")
-
-        if not cls.check_deps():
-            raise RuntimeError("torch_npu is not available but NpuSwiGluKernel was called.")
 
         # Mapping of specific mlp modules to their corresponding kernel implementations
         kernel_mapping = {
@@ -158,7 +155,7 @@ class NpuSwiGluKernel(BaseKernel):
             # Match any module whose class name contains "MLP"
             if (
                 re.search(swiglu_pattern, module.__class__.__name__)
-                and module.__class__.__name__ in cls.expect_modules
+                and module.__class__.__name__ in NpuSwiGluKernel.expect_modules
             ):
                 # Bind function as an instance method to preserve `self` semantics
                 # and replace the original forward

src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/npu_rms_norm.py

@@ -26,10 +26,9 @@ import types
 import torch
 import torch.nn.functional as F
 
-from ......accelerator.helper import DeviceType
+from ......accelerator.helper import DeviceType, get_current_accelerator
 from ......utils.types import HFModel
-from ...base import BaseKernel
+from ...base import BaseKernel, KernelPlugin
-from ...registry import register_kernel
 
 
 try:
@@ -118,15 +117,18 @@ def npu_gated_rms_norm_forward(self, hidden_states, gate=None):
     return hidden_states.to(input_dtype)
 
 
-@register_kernel
+@KernelPlugin("npu_fused_rmsnorm").register()
 class NpuRMSNormKernel(BaseKernel):
     """NPU kernel wrapper for RMSNorm that applies the replacement within a model."""
 
-    _kernel_id = "npu_fused_rmsnorm"
+    @staticmethod
-    _device = DeviceType.NPU
+    def check_device() -> None:
+        current = get_current_accelerator().type
+        if current != DeviceType.NPU:
+            raise RuntimeError(f"NpuRMSNormKernel requires NPU, current accelerator is {current}.")
 
-    @classmethod
+    @staticmethod
-    def apply(cls, **kwargs) -> "HFModel":
+    def _apply(**kwargs) -> "HFModel":
         """Iterate the model and apply NPU-optimized forward to matched RMSNorm modules.
 
         Matches modules whose class name contains "RMSNorm" (case-insensitive) and binds
@@ -144,11 +146,6 @@ class NpuRMSNormKernel(BaseKernel):
             ValueError: If the model is not provided.
         """
         model = kwargs.get("model")
-        if model is None:
-            raise ValueError(f"HFModel instance is required for {cls.__name__}.")
-
-        if not cls.check_deps():
-            raise RuntimeError(f"torch_npu is not available but {cls.__name__} was called.")
 
         rms_norm_pattern = re.compile("RMSNorm", re.IGNORECASE)
 

src/llamafactory/v1/plugins/model_plugins/kernels/ops/rope/npu_rope.py

@@ -24,11 +24,10 @@ import sys
 
 import torch
 
-from ......accelerator.helper import DeviceType
+from ......accelerator.helper import DeviceType, get_current_accelerator
 from ......utils.logging import get_logger
 from ......utils.types import HFModel
-from ...base import BaseKernel
+from ...base import BaseKernel, KernelPlugin
-from ...registry import register_kernel
 
 
 logger = get_logger(__name__)
@@ -125,15 +124,18 @@ def _apply_multimodal_rotary_pos_emb_qwen25_vl(q, k, cos, sin, mrope_section, un
     return _apply_npu_rotary_emb(q, k, cos, sin)
 
 
-@register_kernel
+@KernelPlugin("npu_fused_rope").register()
 class NpuRoPEKernel(BaseKernel):
     """NPU Kernel for Rotary Position Embedding."""
 
-    _kernel_id = "npu_fused_rope"
+    @staticmethod
-    _device = DeviceType.NPU
+    def check_device() -> None:
+        current = get_current_accelerator().type
+        if current != DeviceType.NPU:
+            raise RuntimeError(f"NpuRoPEKernel requires NPU, current accelerator is {current}.")
 
-    @classmethod
+    @staticmethod
-    def apply(cls, **kwargs) -> "HFModel":
+    def _apply(**kwargs) -> "HFModel":
         """Apply RoPE acceleration by monkey-patching ``apply_rotary_pos_emb``.
 
         Iterates through the model's modules to find attention layers, identifies
@@ -151,12 +153,7 @@ class NpuRoPEKernel(BaseKernel):
             RuntimeError: If ``torch_npu`` is not available.
             ValueError: If the model is not provided.
         """
-        if not cls.check_deps():
-            raise RuntimeError(f"torch_npu is not available but {cls.__name__} was called.")
-
         model = kwargs.get("model", None)
-        if model is None:
-            raise ValueError(f"HFModel instance is required for {cls.__name__}.")
 
         _modules = set()
         for module in model.modules():

src/llamafactory/v1/plugins/model_plugins/kernels/registry.py

@@ -1,96 +0,0 @@
-# Copyright 2025 the LlamaFactory team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""The definition of kernel registry.
-
-Init Phase:
-1. Define kernel registry.
-2. Register kernels.
-
-"""
-
-from ....accelerator.helper import get_current_accelerator
-from .base import BaseKernel
-
-
-__all__ = ["Registry", "register_kernel"]
-
-
-class Registry:
-    """Registry for managing kernel implementations.
-
-    Storage structure: ``{ "kernel_id": Class }``
-    """
-
-    _kernels: dict[str, type[BaseKernel]] = {}
-
-    @classmethod
-    def register(cls, kernel_cls: type[BaseKernel]) -> type[BaseKernel] | None:
-        """Decorator to register a kernel class.
-
-        The class must inherit from :class:`BaseKernel` and specify ``_kernel_id`` and ``_device`` attributes.
-
-        Args:
-            kernel_cls (type[BaseKernel]): The kernel class to register.
-
-        Returns:
-            type[BaseKernel] | None: The registered kernel class if the device type matches the current accelerator
-
-        Raises:
-            TypeError: If the class does not inherit from :class:`BaseKernel`.
-            ValueError: If the kernel ID is missing or already registered.
-        """
-        if not issubclass(kernel_cls, BaseKernel):
-            raise TypeError(f"Class {kernel_cls} must inherit from BaseKernel")
-
-        kernel_id = kernel_cls.get_kernel_id()
-        device = kernel_cls.get_device()
-
-        # The device type of the current accelerator does not match the device type required by the kernel, skip registration
-        if get_current_accelerator().type not in device:
-            return
-
-        if not kernel_id:
-            raise ValueError(f"Kernel ID (_kernel_id) is needed for {kernel_cls} to register")
-
-        if kernel_id in cls._kernels:
-            raise ValueError(f"{kernel_id} already registered! The registered kernel is {cls._kernels[kernel_id]}")
-
-        cls._kernels[kernel_id] = kernel_cls
-        return kernel_cls
-
-    @classmethod
-    def get(cls, kernel_id: str) -> type[BaseKernel] | None:
-        """Retrieves a registered kernel implementation by its ID.
-
-        Args:
-            kernel_id (str): The ID of the kernel to retrieve.
-
-        Returns:
-            type[BaseKernel] | None: The kernel class if found, else ``None``.
-        """
-        return cls._kernels.get(kernel_id)
-
-    @classmethod
-    def get_registered_kernels(cls) -> dict[str, type[BaseKernel]]:
-        """Returns a dictionary of all registered kernels.
-
-        Returns:
-            dict[str, type[BaseKernel]]: Dictionary mapping kernel IDs to kernel classes.
-        """
-        return cls._kernels
-
-
-# export decorator alias
-register_kernel = Registry.register

src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py

@@ -37,8 +37,8 @@ logger = logging.get_logger(__name__)
 
 
 class SequenceParallelModelPlugin(BasePlugin):
-    def __call__(self, model, model_args):
+    def __call__(self, model, cp_size: int):
-        return super().__call__(model, model_args)
+        return super().__call__(model, cp_size)
 
 
 class SequenceParallelLossPlugin(BasePlugin):
@@ -82,10 +82,9 @@ def new_flash_attn_forward(
 
 
 @SequenceParallelModelPlugin("ulysses").register()
-def apply_sequence_parallel(model, model_args):
+def apply_sequence_parallel(model, cp_size: int):
     # Replace _flash_attention_forward with new_flash_attn_forward
     module = sys.modules[model.__module__]
-    cp_size = model_args.get("cp_size", 1)
 
     set_ulysses_sequence_parallel_group(DistributedInterface().get_group(Dim.CP))
 

src/llamafactory/v1/plugins/model_plugins/peft.py

@@ -13,7 +13,8 @@
 # limitations under the License.
 
 import re
-from typing import Literal, TypedDict, Union
+from dataclasses import dataclass, field
+from typing import Literal
 
 import torch
 from peft import LoraConfig, PeftModel, TaskType, get_peft_model
@@ -28,53 +29,59 @@ from ...utils.types import HFModel
 logger = logging.get_logger(__name__)
 
 
-class LoraConfigDict(TypedDict, total=False):
+@dataclass
-    name: Literal["lora"]
+class LoraParams:
+    """Typed configuration for the LoRA PEFT plugin."""
+
+    name: Literal["lora"] = "lora"
     """Plugin name."""
-    r: int
+    r: int = 8
-    """Lora rank."""
+    """LoRA rank."""
-    lora_alpha: int
+    lora_alpha: int = 16
-    """Lora alpha."""
+    """LoRA alpha."""
-    lora_dropout: float
+    lora_dropout: float = 0.05
-    """Lora dropout."""
+    """LoRA dropout."""
-    target_modules: Union[list[str], str]
+    target_modules: list[str] | str = "all"
     """Target modules."""
-    use_rslora: bool
+    use_rslora: bool = False
     """Use RS-LoRA."""
-    use_dora: bool
+    use_dora: bool = False
     """Use DoRA."""
-    modules_to_save: list[str]
+    modules_to_save: list[str] | None = None
     """Modules to save."""
-    adapter_name_or_path: Union[list[str], str]
+    adapter_name_or_path: list[str] | str | None = None
     """Path to the adapter(s)."""
-    export_dir: str
+    export_dir: str | None = None
     """Path to the export directory."""
-    export_size: int
+    export_size: int = 5
-    """Shard size for the export model."""
+    """Shard size for the exported model, in GB."""
-    export_hub_model_id: str
+    export_hub_model_id: str | None = None
-    """Hub model ID for the export model."""
+    """Hub model ID for the exported model."""
-    infer_dtype: Literal["auto", "float16", "float32", "bfloat16"]
+    infer_dtype: Literal["auto", "float16", "float32", "bfloat16"] = "auto"
-    """Inference data type for the export model."""
+    """Inference data type for the exported model."""
-    export_legacy_format: bool
+    export_legacy_format: bool = False
-    """Use legacy format for the export model."""
+    """Use legacy format for the exported model."""
 
 
-class FreezeConfigDict(TypedDict, total=False):
+@dataclass
-    name: Literal["freeze"]
+class FreezeParams:
+    """Typed configuration for the freeze PEFT plugin."""
+
+    name: Literal["freeze"] = "freeze"
     """Plugin name."""
-    freeze_trainable_layers: int
+    freeze_trainable_layers: int = 2
-    """Freeze trainable layers."""
+    """Number of trainable layers."""
-    freeze_trainable_modules: Union[list[str], str]
+    freeze_trainable_modules: list[str] | str = "all"
-    """Freeze trainable modules."""
+    """Trainable modules in the selected layers."""
-    freeze_extra_modules: list[str]
+    freeze_extra_modules: list[str] | str | None = field(default_factory=list)
-    """Freeze extra modules."""
+    """Extra non-hidden modules to train."""
-    cast_trainable_params_to_fp32: bool
+    cast_trainable_params_to_fp32: bool = True
-    """Cast trainable params to fp32."""
+    """Cast trainable parameters to float32."""
 
 
 class PeftPlugin(BasePlugin):
-    def __call__(self, model: HFModel, config: dict, is_train: bool) -> HFModel:
+    def __call__(self, model: HFModel, peft_config: dict, is_train: bool) -> HFModel:
-        return super().__call__(model, config, is_train)
+        return super().__call__(model, peft_config, is_train)
 
 
 def _find_all_linear_modules(model: HFModel) -> list[str]:
@@ -91,7 +98,7 @@ def _find_all_linear_modules(model: HFModel) -> list[str]:
     return list(module_names)
 
 
-def merge_adapters(model: HFModel, adapter_name_or_path: Union[list[str], str]) -> HFModel:
+def merge_adapters(model: HFModel, adapter_name_or_path: list[str] | str) -> HFModel:
     if not isinstance(adapter_name_or_path, list):
         adapter_name_or_path = [adapter_name_or_path]
 
@@ -103,7 +110,7 @@ def merge_adapters(model: HFModel, adapter_name_or_path: Union[list[str], str])
     return model
 
 
-def load_adapter(model: HFModel, adapter_name_or_path: Union[list[str], str], is_train: bool) -> HFModel:
+def load_adapter(model: HFModel, adapter_name_or_path: list[str] | str, is_train: bool) -> HFModel:
     r"""Loads adapter(s) into the model.
 
     Determine adapter usage based on mode:
@@ -149,15 +156,16 @@ def load_adapter(model: HFModel, adapter_name_or_path: Union[list[str], str], is
 
 
 @PeftPlugin("lora").register()
-def get_lora_model(model: HFModel, config: LoraConfigDict, is_train: bool = False) -> HFModel:
+def get_lora_model(model: HFModel, peft_config: dict | LoraParams, is_train: bool = False) -> HFModel:
-    adapter_name_or_path = config.get("adapter_name_or_path")
+    peft_config = PeftPlugin.parse_params(peft_config, LoraParams)
+    adapter_name_or_path = peft_config.adapter_name_or_path
 
     if adapter_name_or_path:
         return load_adapter(model, adapter_name_or_path, is_train)
 
     logger.info_rank0("Fine-tuning method: LoRA")
 
-    target_modules = config.get("target_modules", "all")
+    target_modules = peft_config.target_modules
 
     # Handle target modules
     if target_modules == "all":
@@ -175,19 +183,19 @@ def get_lora_model(model: HFModel, config: LoraConfigDict, is_train: bool = Fals
     else:
         task_type = TaskType.CAUSAL_LM
 
-    peft_config = LoraConfig(
+    lora_config = LoraConfig(
         task_type=task_type,
         inference_mode=not is_train,
-        r=config.get("r", 8),
+        r=peft_config.r,
-        lora_alpha=config.get("lora_alpha", 16),
+        lora_alpha=peft_config.lora_alpha,
-        lora_dropout=config.get("lora_dropout", 0.05),
+        lora_dropout=peft_config.lora_dropout,
-        use_rslora=config.get("use_rslora", False),
+        use_rslora=peft_config.use_rslora,
-        use_dora=config.get("use_dora", False),
+        use_dora=peft_config.use_dora,
         target_modules=target_modules,
-        modules_to_save=config.get("modules_to_save", None),
+        modules_to_save=peft_config.modules_to_save,
     )
 
-    model = get_peft_model(model, peft_config)
+    model = get_peft_model(model, lora_config)
 
     if is_train:
         model.print_trainable_parameters()
@@ -196,16 +204,17 @@ def get_lora_model(model: HFModel, config: LoraConfigDict, is_train: bool = Fals
 
 
 @PeftPlugin("freeze").register()
-def get_freeze_model(model: HFModel, config: FreezeConfigDict, is_train: bool = False) -> HFModel:
+def get_freeze_model(model: HFModel, peft_config: dict | FreezeParams, is_train: bool = False) -> HFModel:
+    peft_config = PeftPlugin.parse_params(peft_config, FreezeParams)
     logger.info_rank0("Fine-tuning method: Freeze")
 
     if not is_train:
         return model
 
-    freeze_trainable_layers = config.get("freeze_trainable_layers", 2)
+    freeze_trainable_layers = peft_config.freeze_trainable_layers
-    freeze_trainable_modules = config.get("freeze_trainable_modules", ["all"])
+    freeze_trainable_modules = peft_config.freeze_trainable_modules
-    freeze_extra_modules = config.get("freeze_extra_modules", [])
+    freeze_extra_modules = peft_config.freeze_extra_modules
-    cast_trainable_params_to_fp32 = config.get("cast_trainable_params_to_fp32", True)
+    cast_trainable_params_to_fp32 = peft_config.cast_trainable_params_to_fp32
 
     if isinstance(freeze_trainable_modules, str):
         freeze_trainable_modules = [module.strip() for module in freeze_trainable_modules.split(",")]
@@ -292,26 +301,16 @@ def get_freeze_model(model: HFModel, config: FreezeConfigDict, is_train: bool =
 def merge_and_export_model(args: InputArgument = None):
     model_args, _, _, _ = get_args(args)
 
-    export_config = model_args.peft_config
+    raw_config = model_args.peft_config
-    if export_config is None:
+    if raw_config is None:
         raise ValueError("Please specify peft_config to merge and export model.")
-
+    if raw_config.name != "lora":
-    export_dir = export_config.get("export_dir")
-    if export_dir is None:
-        raise ValueError("Please specify export_dir.")
-
-    export_size = export_config.get("export_size", 5)
-    export_hub_model_id = export_config.get("export_hub_model_id")
-    infer_dtype = export_config.get("infer_dtype", "auto")
-    export_legacy_format = export_config.get("export_legacy_format", False)
-
-    adapters = None
-    if export_config.get("name") == "lora":
-        adapters = export_config.get("adapter_name_or_path")
-    else:
         raise ValueError("Currently merge and export model function is only supported for lora.")
 
-    if adapters is None:
+    export_peft_config = PeftPlugin.parse_params(raw_config, LoraParams)
+    if export_peft_config.export_dir is None:
+        raise ValueError("Please specify export_dir.")
+    if export_peft_config.adapter_name_or_path is None:
         raise ValueError("Please set adapter_name_or_path to merge adapters into base model.")
 
     logger.info_rank0("Loading model for export...")
@@ -319,33 +318,33 @@ def merge_and_export_model(args: InputArgument = None):
     model = model_engine.model
     tokenizer = model_engine.processor
 
-    if infer_dtype == "auto":
+    if export_peft_config.infer_dtype == "auto":
         if model.config.torch_dtype == torch.float32 and torch.cuda.is_bf16_supported():
             model = model.to(torch.bfloat16)
             logger.info_rank0("Converted model to bfloat16.")
     else:
-        target_dtype = getattr(torch, infer_dtype)
+        target_dtype = getattr(torch, export_peft_config.infer_dtype)
         model = model.to(target_dtype)
-        logger.info_rank0(f"Converted model to {infer_dtype}.")
+        logger.info_rank0(f"Converted model to {export_peft_config.infer_dtype}.")
 
-    logger.info_rank0(f"Exporting model to {export_dir}...")
+    logger.info_rank0(f"Exporting model to {export_peft_config.export_dir}...")
     model.save_pretrained(
-        export_dir,
+        export_peft_config.export_dir,
-        max_shard_size=f"{export_size}GB",
+        max_shard_size=f"{export_peft_config.export_size}GB",
-        safe_serialization=not export_legacy_format,
+        safe_serialization=not export_peft_config.export_legacy_format,
     )
     if tokenizer is not None:
         try:
             if hasattr(tokenizer, "padding_side"):
                 tokenizer.padding_side = "left"
-            tokenizer.save_pretrained(export_dir)
+            tokenizer.save_pretrained(export_peft_config.export_dir)
         except Exception as e:
             logger.warning(f"Failed to save tokenizer: {e}")
 
-    if export_hub_model_id:
+    if export_peft_config.export_hub_model_id:
-        logger.info_rank0(f"Pushing to hub: {export_hub_model_id}...")
+        logger.info_rank0(f"Pushing to hub: {export_peft_config.export_hub_model_id}...")
-        model.push_to_hub(export_hub_model_id)
+        model.push_to_hub(export_peft_config.export_hub_model_id)
         if tokenizer is not None:
-            tokenizer.push_to_hub(export_hub_model_id)
+            tokenizer.push_to_hub(export_peft_config.export_hub_model_id)
 
     logger.info_rank0("Model exported successfully.")

src/llamafactory/v1/plugins/model_plugins/quantization.py

@@ -15,108 +15,104 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Any
+from dataclasses import dataclass
+from typing import Any, Literal
 
-import torch
-from transformers import BitsAndBytesConfig
-
-from ...accelerator.helper import get_current_device
-from ...config.model_args import ModelArguments
 from ...utils import logging
-from ...utils.packages import check_version
 from ...utils.plugin import BasePlugin
 
 
-if TYPE_CHECKING:
-    from transformers import PretrainedConfig, PreTrainedTokenizer
-
 logger = logging.get_logger(__name__)
 
 
 class QuantizationPlugin(BasePlugin):
-    r"""Plugin for model quantization."""
-
     def __call__(
         self,
-        init_kwargs: dict[str, Any] = None,
+        init_kwargs: dict[str, Any] | None = None,
-        config: "PretrainedConfig" = None,
+        quant_config=None,
-        tokenizer: "PreTrainedTokenizer" = None,
-        model_args: "ModelArguments" = None,
         is_trainable: bool = False,
     ) -> dict[str, Any]:
-        return super().__call__(
+        return super().__call__(init_kwargs, quant_config=quant_config, is_trainable=is_trainable)
-            init_kwargs, config=config, tokenizer=tokenizer, model_args=model_args, is_trainable=is_trainable
+
-        )
+
+@dataclass
+class BnbParams:
+    name: Literal["bnb", "auto"] = "bnb"
+    quantization_bit: int | None = None
+    compute_dtype: str | Any = "float16"
+    double_quantization: bool = True
+    quantization_type: str = "nf4"
+
+    def __post_init__(self) -> None:
+        import torch
+
+        if isinstance(self.compute_dtype, str):
+            dtype = getattr(torch, self.compute_dtype, None)
+            if not isinstance(dtype, torch.dtype):
+                raise ValueError(f"compute_dtype={self.compute_dtype!r} is not a torch dtype name.")
+            self.compute_dtype = dtype
+        elif not isinstance(self.compute_dtype, torch.dtype):
+            raise TypeError(f"compute_dtype must be str or torch.dtype, got {type(self.compute_dtype).__name__}.")
 
 
 @QuantizationPlugin("auto").register()
 def quantization_auto(
     init_kwargs: dict[str, Any],
-    **kwargs,
+    quant_config: dict | BnbParams,
+    is_trainable: bool = False,
 ) -> dict[str, Any]:
-    """Automatic quantization selection, only support bnb currently.
+    quant_config = QuantizationPlugin.parse_params(quant_config, BnbParams)
+    if quant_config.quantization_bit is None:
+        logger.warning_rank0("No quantization method applied.")
+        return init_kwargs
+    if quant_config.quantization_bit not in (4, 8):
+        raise ValueError(f"Unsupported quantization bit: {quant_config.quantization_bit} for auto quantization.")
 
-    Args:
+    logger.info_rank0(f"Loading {quant_config.quantization_bit}-bit quantized model.")
-        init_kwargs (dict[str, Any]): The kwargs for model initialization.
+    return QuantizationPlugin("bnb")(init_kwargs, quant_config=quant_config, is_trainable=is_trainable)
-        **kwargs: Keyword arguments containing the model.
-
-    Returns:
-        dict[str, Any]: The updated kwargs for model initialization.
-    """
-    model_args: ModelArguments = kwargs.get("model_args", None)
-    quant_config = model_args.quant_config
-
-    quantization_bit = quant_config.get("quantization_bit", None)
-    if quantization_bit is not None:
-        logger.info_rank0(f"Loading {quantization_bit}-bit quantized model.")
-        if quantization_bit in [8, 4]:
-            return quantization_with_bnb(init_kwargs, **kwargs)
-        else:
-            raise ValueError(f"Unsupported quantization bit: {quantization_bit} for auto quantization.")
-    logger.warning_rank0("No quantization method applied.")
-    return init_kwargs
 
 
 @QuantizationPlugin("bnb").register()
 def quantization_with_bnb(
     init_kwargs: dict[str, Any],
-    model_args: "ModelArguments" = None,
+    quant_config: dict | BnbParams,
-    **kwargs,
+    is_trainable: bool = False,
 ) -> dict[str, Any]:
-    r"""Quantization with BNB."""
+    from transformers import BitsAndBytesConfig
-    logger.info_rank0("Using Bitsandbytes quantization.")
+
-    quantization_bit = model_args.quant_config.get("quantization_bit", None)
+    from ...accelerator.helper import get_current_device
+    from ...utils.packages import check_version
+
+    quant_config = QuantizationPlugin.parse_params(quant_config, BnbParams)
+    quantization_bit = quant_config.quantization_bit
     if quantization_bit is None:
-        logger.warning_rank0("quantization_bit is not specified, default to 8-bit quantization.")
+        logger.warning_rank0("quantization_bit is not specified, default to 4-bit quantization.")
         quantization_bit = 4
-    assert quantization_bit in [8, 4], "Bitsandbytes only accepts 4-bit or 8-bit quantization."
+    if quantization_bit not in (4, 8):
+        raise ValueError("Bitsandbytes only accepts 4-bit or 8-bit quantization.")
+
+    logger.info_rank0("Using Bitsandbytes quantization.")
     if quantization_bit == 8:
         check_version("bitsandbytes>=0.37.0", mandatory=True)
         init_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_8bit=True)
-    elif quantization_bit == 4:
+    else:
         check_version("bitsandbytes>=0.39.0", mandatory=True)
         init_kwargs["quantization_config"] = BitsAndBytesConfig(
             load_in_4bit=True,
-            bnb_4bit_compute_dtype=model_args.quant_config.get("compute_dtype", torch.float16),
+            bnb_4bit_compute_dtype=quant_config.compute_dtype,
-            bnb_4bit_use_double_quant=model_args.quant_config.get("double_quantization", True),
+            bnb_4bit_use_double_quant=quant_config.double_quantization,
-            bnb_4bit_quant_type=model_args.quant_config.get("quantization_type", "nf4"),
+            bnb_4bit_quant_type=quant_config.quantization_type,
-            bnb_4bit_quant_storage=model_args.quant_config.get(
+            bnb_4bit_quant_storage=quant_config.compute_dtype,
-                "compute_dtype", torch.float16
-            ),  # crucial for fsdp+qlora
         )
-    else:
-        raise ValueError("Bitsandbytes only accepts 4-bit or 8-bit quantization.")
 
-    # TODO: improve deepspeed zero3 and fsdp detection.
+    if is_trainable:
-    if kwargs.get("is_trainable", False):
         logger.info_rank0("Detected inference mode, setting device_map for bitsandbytes quantization.")
-        init_kwargs["device_map"] = {"": get_current_device()}  # change auto device map for inference
+        init_kwargs["device_map"] = {"": get_current_device()}
     else:
         logger.info_rank0("Detected training mode, skip setting device_map for bitsandbytes quantization.")
-        if model_args.quant_config.get("quantization_bit") != 4:
+        if quantization_bit != 4:
             raise ValueError("Only 4-bit quantized model can use fsdp+qlora or auto device map.")
-
         check_version("bitsandbytes>=0.43.0", mandatory=True)
 
-    logger.info_rank0(f"Quantizing model to {model_args.quant_config.get('quantization_bit')} bit with bitsandbytes.")
+    logger.info_rank0(f"Quantizing model to {quantization_bit} bit with bitsandbytes.")
     return init_kwargs

src/llamafactory/v1/plugins/trainer_plugins/batching.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from abc import ABC, abstractmethod
 from collections.abc import Callable
 from math import ceil
 from typing import Any
@@ -22,34 +23,38 @@ from torch.utils.data import default_collate
 from ...utils.constants import IGNORE_INDEX
 from ...utils.helper import pad_and_truncate
 from ...utils.objects import StatefulBuffer
-from ...utils.plugin import BasePlugin
+from ...utils.plugin import BasePlugin, ensure_methods_implemented
 from ...utils.types import BatchInfo, BatchInput, DataLoader, ModelInput
 
 
 class BatchingPlugin(BasePlugin):
-    def get_data_provider_batch_size(self, batch_info: BatchInfo) -> int:
+    """Plugin family for batching strategy method groups."""
-        """Return the raw data provider batch size for this batching strategy."""
-        return self["get_data_provider_batch_size"](batch_info)
 
-    def compute_length(self, data_provider: DataLoader, batch_info: BatchInfo) -> int:
-        """Compute the length of the batch generator.
 
-        The approximate length is used to calculate the lr schedule.
+class BaseBatcher(ABC):
-        """
+    def __init_subclass__(cls, **kwargs) -> None:
-        return self["compute_length"](data_provider, batch_info)
+        super().__init_subclass__(**kwargs)
+        ensure_methods_implemented(cls)
 
+    @staticmethod
+    @abstractmethod
+    def get_data_provider_batch_size(batch_info: BatchInfo) -> int: ...
+
+    @staticmethod
+    @abstractmethod
+    def compute_length(data_provider: DataLoader, batch_info: BatchInfo) -> int: ...
+
+    @staticmethod
+    @abstractmethod
     def fill_buffer(
-        self,
         buffer: StatefulBuffer,
         batch_info: BatchInfo,
         next_samples: Callable[[bool], list[ModelInput] | None],
-    ) -> None:
+    ) -> None: ...
-        """Fill the buffer with data."""
-        return self["fill_buffer"](buffer, batch_info, next_samples)
 
-    def generate_batch(self, buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None:
+    @staticmethod
-        """Generate a batch from the buffer."""
+    @abstractmethod
-        return self["generate_batch"](buffer, batch_info)
+    def generate_batch(buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None: ...
 
 
 def _get_dynamic_micro_batch_sizes(samples: list[ModelInput], batch_info: BatchInfo) -> list[int]:
@@ -153,131 +158,131 @@ def _pack_padding_free_samples(samples: list[ModelInput], cutoff_len: int) -> Ba
     return {key: None if value is None else torch.tensor(value).unsqueeze(0) for key, value in packed.items()}
 
 
-@BatchingPlugin("padding_free").register("get_data_provider_batch_size")
+@BatchingPlugin("padding_free").register()
-def get_padding_free_data_provider_batch_size(batch_info: BatchInfo) -> int:
+class PaddingFreeBatcher(BaseBatcher):
-    return batch_info["micro_batch_size"] * batch_info["num_micro_batch"]
+    @staticmethod
+    def get_data_provider_batch_size(batch_info: BatchInfo) -> int:
+        return batch_info["micro_batch_size"] * batch_info["num_micro_batch"]
 
+    @staticmethod
+    def compute_length(data_provider: DataLoader, batch_info: BatchInfo) -> int:
+        return len(data_provider)
 
-@BatchingPlugin("padding_free").register("compute_length")
+    @staticmethod
-def compute_padding_free_length(data_provider: DataLoader, batch_info: BatchInfo) -> int:
+    def fill_buffer(
-    return len(data_provider)
+        buffer: StatefulBuffer,
+        batch_info: BatchInfo,
+        next_samples: Callable[[bool], list[ModelInput] | None],
+    ) -> None:
+        while len(buffer) < batch_info["micro_batch_size"] * batch_info["num_micro_batch"]:
+            samples = next_samples(False)
+            if samples is None:
+                break
 
+            buffer.put(samples)
 
-@BatchingPlugin("padding_free").register("fill_buffer")
+    @staticmethod
-def fill_padding_free_buffer(
+    def generate_batch(buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None:
-    buffer: StatefulBuffer,
+        micro_batch_size = batch_info["micro_batch_size"]
-    batch_info: BatchInfo,
+        num_micro_batch = batch_info["num_micro_batch"]
-    next_samples: Callable[[bool], list[ModelInput] | None],
+        cutoff_len = batch_info["cutoff_len"]
-) -> None:
+        batch_size = micro_batch_size * num_micro_batch
-    while len(buffer) < batch_info["micro_batch_size"] * batch_info["num_micro_batch"]:
+        if len(buffer) < batch_size:
-        samples = next_samples(False)
-        if samples is None:
-            break
-
-        buffer.put(samples)
-
-
-@BatchingPlugin("padding_free").register("generate_batch")
-def generate_padding_free_batch(buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None:
-    micro_batch_size = batch_info["micro_batch_size"]
-    num_micro_batch = batch_info["num_micro_batch"]
-    cutoff_len = batch_info["cutoff_len"]
-    batch_size = micro_batch_size * num_micro_batch
-    if len(buffer) < batch_size:
-        return None
-
-    samples = buffer.get(batch_size)
-    batch = []
-    for i in range(num_micro_batch):
-        micro_batch = samples[i * micro_batch_size : (i + 1) * micro_batch_size]
-        packed_micro_batch = _pack_padding_free_samples(micro_batch, cutoff_len)
-        if packed_micro_batch is None:
             return None
 
-        batch.append(packed_micro_batch)
+        samples = buffer.get(batch_size)
+        batch = []
+        for i in range(num_micro_batch):
+            micro_batch = samples[i * micro_batch_size : (i + 1) * micro_batch_size]
+            packed_micro_batch = _pack_padding_free_samples(micro_batch, cutoff_len)
+            if packed_micro_batch is None:
+                return None
 
-    return batch
+            batch.append(packed_micro_batch)
+
+        return batch
 
 
-@BatchingPlugin("dynamic_batching").register("get_data_provider_batch_size")
+@BatchingPlugin("dynamic_batching").register()
-def get_dynamic_batching_data_provider_batch_size(batch_info: BatchInfo) -> int:
+class DynamicBatcher(BaseBatcher):
-    return 1
+    @staticmethod
+    def get_data_provider_batch_size(batch_info: BatchInfo) -> int:
+        return 1
 
+    @staticmethod
+    def compute_length(data_provider: DataLoader, batch_info: BatchInfo) -> int:
+        batch_size = batch_info["micro_batch_size"] * batch_info["num_micro_batch"]
+        return ceil(len(data_provider) / batch_size)
 
-@BatchingPlugin("dynamic_batching").register("compute_length")
+    @staticmethod
-def compute_dynamic_batching_length(data_provider: DataLoader, batch_info: BatchInfo) -> int:
+    def fill_buffer(
-    batch_size = batch_info["micro_batch_size"] * batch_info["num_micro_batch"]
+        buffer: StatefulBuffer,
-    return ceil(len(data_provider) / batch_size)
+        batch_info: BatchInfo,
+        next_samples: Callable[[bool], list[ModelInput] | None],
+    ) -> None:
+        while len(_get_dynamic_micro_batch_sizes(buffer.samples, batch_info)) < batch_info["num_micro_batch"]:
+            samples = next_samples(True)
+            if samples is None:
+                break
 
+            buffer.put(samples)
 
-@BatchingPlugin("dynamic_batching").register("fill_buffer")
+    @staticmethod
-def fill_dynamic_batching_buffer(
+    def generate_batch(buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None:
-    buffer: StatefulBuffer,
+        micro_batch_sample_counts = _get_dynamic_micro_batch_sizes(buffer.samples, batch_info)
-    batch_info: BatchInfo,
+        if len(micro_batch_sample_counts) < batch_info["num_micro_batch"]:
-    next_samples: Callable[[bool], list[ModelInput] | None],
-) -> None:
-    while len(_get_dynamic_micro_batch_sizes(buffer.samples, batch_info)) < batch_info["num_micro_batch"]:
-        samples = next_samples(True)
-        if samples is None:
-            break
-
-        buffer.put(samples)
-
-
-@BatchingPlugin("dynamic_batching").register("generate_batch")
-def generate_dynamic_batching_batch(buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None:
-    micro_batch_sample_counts = _get_dynamic_micro_batch_sizes(buffer.samples, batch_info)
-    if len(micro_batch_sample_counts) < batch_info["num_micro_batch"]:
-        return None
-
-    batch = []
-    cutoff_len = batch_info["cutoff_len"]
-    for num_samples in micro_batch_sample_counts:
-        samples = buffer.get(num_samples)
-        batch.append(default_collate(pad_and_truncate(samples, cutoff_len)))
-
-    return batch
-
-
-@BatchingPlugin("dynamic_padding_free").register("get_data_provider_batch_size")
-def get_dynamic_padding_free_data_provider_batch_size(batch_info: BatchInfo) -> int:
-    return 1
-
-
-@BatchingPlugin("dynamic_padding_free").register("compute_length")
-def compute_dynamic_padding_free_length(data_provider: DataLoader, batch_info: BatchInfo) -> int:
-    batch_size = batch_info["micro_batch_size"] * batch_info["num_micro_batch"]
-    return ceil(len(data_provider) / batch_size)
-
-
-@BatchingPlugin("dynamic_padding_free").register("fill_buffer")
-def fill_dynamic_padding_free_buffer(
-    buffer: StatefulBuffer,
-    batch_info: BatchInfo,
-    next_samples: Callable[[bool], list[ModelInput] | None],
-) -> None:
-    while len(_get_dynamic_padding_free_micro_batch_sizes(buffer.samples, batch_info)) < batch_info["num_micro_batch"]:
-        samples = next_samples(True)
-        if samples is None:
-            break
-        buffer.put(samples)
-
-
-@BatchingPlugin("dynamic_padding_free").register("generate_batch")
-def generate_dynamic_padding_free_batch(buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None:
-    micro_batch_sample_counts = _get_dynamic_padding_free_micro_batch_sizes(buffer.samples, batch_info)
-    if len(micro_batch_sample_counts) < batch_info["num_micro_batch"]:
-        return None
-
-    batch = []
-    cutoff_len = batch_info["cutoff_len"]
-
-    for num_samples in micro_batch_sample_counts:
-        samples = buffer.get(num_samples)
-        packed_batch = _pack_padding_free_samples(samples, cutoff_len)
-        if packed_batch is None:
             return None
 
-        batch.append(packed_batch)
+        batch = []
+        cutoff_len = batch_info["cutoff_len"]
+        for num_samples in micro_batch_sample_counts:
+            samples = buffer.get(num_samples)
+            batch.append(default_collate(pad_and_truncate(samples, cutoff_len)))
 
-    return batch
+        return batch
+
+
+@BatchingPlugin("dynamic_padding_free").register()
+class DynamicPaddingFreeBatcher(BaseBatcher):
+    @staticmethod
+    def get_data_provider_batch_size(batch_info: BatchInfo) -> int:
+        return 1
+
+    @staticmethod
+    def compute_length(data_provider: DataLoader, batch_info: BatchInfo) -> int:
+        batch_size = batch_info["micro_batch_size"] * batch_info["num_micro_batch"]
+        return ceil(len(data_provider) / batch_size)
+
+    @staticmethod
+    def fill_buffer(
+        buffer: StatefulBuffer,
+        batch_info: BatchInfo,
+        next_samples: Callable[[bool], list[ModelInput] | None],
+    ) -> None:
+        while (
+            len(_get_dynamic_padding_free_micro_batch_sizes(buffer.samples, batch_info))
+            < batch_info["num_micro_batch"]
+        ):
+            samples = next_samples(True)
+            if samples is None:
+                break
+            buffer.put(samples)
+
+    @staticmethod
+    def generate_batch(buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None:
+        micro_batch_sample_counts = _get_dynamic_padding_free_micro_batch_sizes(buffer.samples, batch_info)
+        if len(micro_batch_sample_counts) < batch_info["num_micro_batch"]:
+            return None
+
+        batch = []
+        cutoff_len = batch_info["cutoff_len"]
+
+        for num_samples in micro_batch_sample_counts:
+            samples = buffer.get(num_samples)
+            packed_batch = _pack_padding_free_samples(samples, cutoff_len)
+            if packed_batch is None:
+                return None
+
+            batch.append(packed_batch)
+
+        return batch

src/llamafactory/v1/plugins/trainer_plugins/distributed/base.py

@@ -0,0 +1,50 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+from ....utils.plugin import ensure_methods_implemented
+
+
+if TYPE_CHECKING:
+    import torch
+
+    from ....utils.types import HFModel, Processor
+
+
+class BaseDistributed(ABC):
+    """Contract for distributed backend method groups."""
+
+    def __init_subclass__(cls, **kwargs) -> None:
+        super().__init_subclass__(**kwargs)
+        ensure_methods_implemented(cls)
+
+    @staticmethod
+    @abstractmethod
+    def shard_model(model: HFModel, dist_config: object, **kwargs) -> object: ...
+
+    @staticmethod
+    @abstractmethod
+    def save_model(model: HFModel, output_dir: str, processor: Processor) -> None: ...
+
+    @staticmethod
+    @abstractmethod
+    def save_checkpoint(model: HFModel, optimizer: torch.optim.Optimizer, ckpt_dir: str, **kwargs) -> None: ...
+
+    @staticmethod
+    @abstractmethod
+    def load_checkpoint(model: HFModel, optimizer: torch.optim.Optimizer, ckpt_dir: str, **kwargs) -> None: ...

src/llamafactory/v1/plugins/trainer_plugins/distributed/hub.py

@@ -1,94 +0,0 @@
-# Copyright 2025 the LlamaFactory team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-from typing import TYPE_CHECKING
-
-import torch
-
-from ....config.arg_utils import PluginConfig
-from ....utils.plugin import BasePlugin
-
-
-if TYPE_CHECKING:
-    from ....utils.types import HFModel, Processor
-
-
-class DistributedPlugin(BasePlugin):
-    def __call__(self, model: HFModel, dist_config: PluginConfig, **kwargs) -> HFModel:
-        return super().__call__(model, dist_config, **kwargs)
-
-
-@DistributedPlugin("fsdp2").register()
-def shard_model_fsdp2(model: HFModel, dist_config: PluginConfig, **kwargs) -> HFModel:
-    from .fsdp2 import FSDP2Engine
-
-    return FSDP2Engine(dist_config, bf16=bool(kwargs.get("bf16"))).shard_model(model)
-
-
-@DistributedPlugin("fsdp2").register("save_model")
-def save_model_fsdp2(model: HFModel, output_dir: str, processor: Processor) -> None:
-    from .fsdp2 import save_model
-
-    return save_model(model, output_dir, processor)
-
-
-@DistributedPlugin("fsdp2").register("save_checkpoint")
-def save_checkpoint_fsdp2(model: HFModel, optimizer: torch.optim.Optimizer, ckpt_dir: str, **kwargs) -> None:
-    from .fsdp2 import save_checkpoint
-
-    return save_checkpoint(model, optimizer, ckpt_dir, **kwargs)
-
-
-@DistributedPlugin("fsdp2").register("load_checkpoint")
-def load_checkpoint_fsdp2(model: HFModel, optimizer: torch.optim.Optimizer, ckpt_dir: str, **kwargs) -> None:
-    from .fsdp2 import load_checkpoint
-
-    return load_checkpoint(model, optimizer, ckpt_dir, **kwargs)
-
-
-@DistributedPlugin("deepspeed").register()
-def shard_model_deepspeed(model: HFModel, dist_config: PluginConfig, **kwargs) -> HFModel:
-    if dist_config.get("cp_size", 1) > 1:
-        raise ValueError("CP currently requires `dist_config.name: fsdp2`.")
-
-    from .deepspeed import DeepSpeedEngine
-
-    return DeepSpeedEngine(
-        dist_config,
-        num_micro_batch=kwargs.get("num_micro_batch"),
-        micro_batch_size=kwargs.get("micro_batch_size"),
-    ).shard_model(model)
-
-
-@DistributedPlugin("deepspeed").register("save_model")
-def save_model_deepspeed(model: HFModel, output_dir: str, processor: Processor) -> None:
-    from .deepspeed import save_model
-
-    return save_model(model, output_dir, processor)
-
-
-@DistributedPlugin("deepspeed").register("save_checkpoint")
-def save_checkpoint_deepspeed(model: HFModel, optimizer: torch.optim.Optimizer, ckpt_dir: str, **kwargs) -> None:
-    from .deepspeed import save_checkpoint
-
-    return save_checkpoint(model, optimizer, ckpt_dir, **kwargs)
-
-
-@DistributedPlugin("deepspeed").register("load_checkpoint")
-def load_checkpoint_deepspeed(model: HFModel, optimizer: torch.optim.Optimizer, ckpt_dir: str, **kwargs) -> None:
-    from .deepspeed import load_checkpoint
-
-    return load_checkpoint(model, optimizer, ckpt_dir, **kwargs)

src/llamafactory/v1/plugins/trainer_plugins/distributed/interface.py

@@ -0,0 +1,115 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Distributed backend plugin definitions.
+
+Backend-private params are parsed explicitly at ``shard_model``. ``DistributedInterface``
+reads mesh topology from ``TrainingArguments`` and never puts it in backend params.
+"""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass
+from typing import TYPE_CHECKING, Literal
+
+from ....utils.plugin import BasePlugin
+from .base import BaseDistributed
+
+
+if TYPE_CHECKING:
+    from ....config.arg_utils import PluginConfig
+    from ....utils.types import HFModel
+
+
+@dataclass
+class FSDP2Params:
+    name: Literal["fsdp2"] = "fsdp2"
+    reshard_after_forward: bool = True
+    offload_params: bool = False
+    pin_memory: bool = True
+    dcp_path: str | None = None
+
+
+@dataclass
+class DeepSpeedParams:
+    name: Literal["deepspeed"] = "deepspeed"
+    config_file: str = ""
+
+    def __post_init__(self) -> None:
+        if not self.config_file:
+            raise ValueError("DeepSpeed config_file is required.")
+
+
+class DistributedPlugin(BasePlugin):
+    """Plugin family for distributed training backends."""
+
+
+@DistributedPlugin("fsdp2").register()
+class FSDP2Distributed(BaseDistributed):
+    @staticmethod
+    def shard_model(model: HFModel, dist_config: PluginConfig | FSDP2Params, **kwargs) -> HFModel:
+        dist_config = DistributedPlugin.parse_params(dist_config, FSDP2Params)
+        from .fsdp2 import FSDP2Engine
+
+        return FSDP2Engine(asdict(dist_config), bf16=bool(kwargs.get("bf16"))).shard_model(model)
+
+    @staticmethod
+    def save_model(model, output_dir, processor) -> None:
+        from .fsdp2 import save_model
+
+        save_model(model, output_dir, processor)
+
+    @staticmethod
+    def save_checkpoint(model, optimizer, ckpt_dir, **kwargs) -> None:
+        from .fsdp2 import save_checkpoint
+
+        save_checkpoint(model, optimizer, ckpt_dir, **kwargs)
+
+    @staticmethod
+    def load_checkpoint(model, optimizer, ckpt_dir, **kwargs) -> None:
+        from .fsdp2 import load_checkpoint
+
+        load_checkpoint(model, optimizer, ckpt_dir, **kwargs)
+
+
+@DistributedPlugin("deepspeed").register()
+class DeepSpeedDistributed(BaseDistributed):
+    @staticmethod
+    def shard_model(model: HFModel, dist_config: PluginConfig | DeepSpeedParams, **kwargs) -> object:
+        dist_config = DistributedPlugin.parse_params(dist_config, DeepSpeedParams)
+        from .deepspeed import DeepSpeedEngine
+
+        return DeepSpeedEngine(
+            asdict(dist_config),
+            num_micro_batch=kwargs.get("num_micro_batch"),
+            micro_batch_size=kwargs.get("micro_batch_size"),
+        ).shard_model(model)
+
+    @staticmethod
+    def save_model(model, output_dir, processor) -> None:
+        from .deepspeed import save_model
+
+        save_model(model, output_dir, processor)
+
+    @staticmethod
+    def save_checkpoint(model, optimizer, ckpt_dir, **kwargs) -> None:
+        from .deepspeed import save_checkpoint
+
+        save_checkpoint(model, optimizer, ckpt_dir, **kwargs)
+
+    @staticmethod
+    def load_checkpoint(model, optimizer, ckpt_dir, **kwargs) -> None:
+        from .deepspeed import load_checkpoint
+
+        load_checkpoint(model, optimizer, ckpt_dir, **kwargs)

src/llamafactory/v1/trainers/dpo_trainer.py

@@ -62,10 +62,7 @@ def compute_sigmoid_dpo_loss(
     chosen_logratios = policy_chosen_logps - ref_chosen_logps
     rejected_logratios = policy_rejected_logps - ref_rejected_logps
     logits = chosen_logratios - rejected_logratios
-    return (
+    return -F.logsigmoid(beta * logits) * (1 - label_smoothing) - F.logsigmoid(-beta * logits) * label_smoothing
-        -F.logsigmoid(beta * logits) * (1 - label_smoothing)
-        - F.logsigmoid(-beta * logits) * label_smoothing
-    )
 
 
 def _validate_dpo_dataset_format(train_dataset: DataEngine, dataset_path: str) -> None:
@@ -97,8 +94,7 @@ class DPOTrainer(BaseTrainer):
         train_dataset,
         callbacks=None,
     ) -> None:
-        cp_size = args.dist_config.get("cp_size", 1) if args.dist_config is not None else 1
+        if args.cp_size > 1:
-        if cp_size > 1:
             raise NotImplementedError("DPO trainer currently only supports cp_size == 1.")
 
         self.pref_loss = args.pref_loss
@@ -378,7 +374,9 @@ class DPOTrainer(BaseTrainer):
 
         # Raw logits means (for logging)
         chosen_logits_mean = (shift_logits.mean(dim=-1) * chosen_logit_mask).sum() / (chosen_logit_mask.sum() + 1e-6)
-        rejected_logits_mean = (shift_logits.mean(dim=-1) * rejected_logit_mask).sum() / (rejected_logit_mask.sum() + 1e-6)
+        rejected_logits_mean = (shift_logits.mean(dim=-1) * rejected_logit_mask).sum() / (
+            rejected_logit_mask.sum() + 1e-6
+        )
 
         if self.pref_loss == "sigmoid":
             if not self._use_lora_ref and self.ref_model is None:
@@ -431,7 +429,7 @@ def run_dpo(args: InputArgument = None):
     model_args, data_args, training_args, _ = get_args(args)
     if getattr(training_args, "use_cpu", False):
         os.environ["FORCE_V1_CPU"] = "1"
-    DistributedInterface(training_args.dist_config)
+    DistributedInterface(training_args)
     train_dataset = DataEngine(data_args.train_dataset)
     _validate_dpo_dataset_format(train_dataset, data_args.train_dataset)
     model_engine = ModelEngine(model_args, is_train=True)

src/llamafactory/v1/trainers/rm_trainer.py

@@ -76,8 +76,7 @@ class RMTrainer(BaseTrainer):
         train_dataset,
         callbacks=None,
     ) -> None:
-        cp_size = args.dist_config.get("cp_size", 1) if args.dist_config is not None else 1
+        if args.cp_size > 1:
-        if cp_size > 1:
             raise NotImplementedError("RM trainer currently only supports cp_size == 1.")
 
         super().__init__(args, model, renderer, train_dataset, callbacks)
@@ -163,7 +162,7 @@ class RMTrainer(BaseTrainer):
 def run_rm(args: InputArgument = None):
     model_args, data_args, training_args, _ = get_args(args)
     model_args.model_class = ModelClass.CLS
-    DistributedInterface(training_args.dist_config)
+    DistributedInterface(training_args)
     train_dataset = DataEngine(data_args.train_dataset)
     _validate_rm_dataset_format(train_dataset, data_args.train_dataset)
     model_engine = ModelEngine(model_args, is_train=True)

src/llamafactory/v1/trainers/sft_trainer.py

@@ -31,7 +31,7 @@ class SFTTrainer(BaseTrainer):
 
 def run_sft(args: InputArgument = None):
     model_args, data_args, training_args, _ = get_args(args)
-    DistributedInterface(training_args.dist_config)
+    DistributedInterface(training_args)
     train_dataset = DataEngine(data_args.train_dataset)
     model_engine = ModelEngine(model_args, is_train=True)
     trainer = SFTTrainer(

src/llamafactory/v1/utils/plugin.py

@@ -12,96 +12,98 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Lightweight plugin routing and shared parameter parsing helpers."""
 
-from collections import defaultdict
+from __future__ import annotations
-from collections.abc import Callable
+
-from typing import Any
+from dataclasses import fields, is_dataclass
+from typing import Any, TypeVar
 
 from . import logging
 
 
 logger = logging.get_logger(__name__)
+ParamsT = TypeVar("ParamsT")
+
+
+def ensure_methods_implemented(cls: type) -> None:
+    """Raise when a static method-group implementation is incomplete."""
+    required: set[str] = set()
+    for base in cls.__mro__[1:]:
+        required |= getattr(base, "__abstractmethods__", frozenset())
+
+    missing = sorted(name for name in required if getattr(getattr(cls, name, None), "__isabstractmethod__", False))
+    if missing:
+        raise TypeError(f"{cls.__name__} does not implement all required methods: {missing}")
 
 
 class BasePlugin:
-    """Base class for plugins.
+    """Route a plugin name to one function or static method-group class.
 
-    A plugin is a callable object that can be registered and called by name.
+    Every plugin family subclass owns an isolated registry. Parameter schemas
-
+    deliberately do not live here; each plugin entrypoint parses its own config.
-    Example usage:
-    ```python
-    class PrintPlugin(BasePlugin):
-        def again(self):  # optional
-            self["again"]()
-
-
-    @PrintPlugin("hello").register()
-    def print_hello():
-        print("Hello world!")
-
-
-    @PrintPlugin("hello").register("again")
-    def print_hello_again():
-        print("Hello world! Again.")
-
-
-    PrintPlugin("hello")()
-    PrintPlugin("hello").again()
-    ```
     """
 
-    _registry: dict[str, dict[str, Callable]] = defaultdict(dict)
+    _registry: dict[str, Any] = {}
+
+    def __init_subclass__(cls, **kwargs) -> None:
+        super().__init_subclass__(**kwargs)
+        cls._registry = {}
 
     def __init__(self, name: str | None = None) -> None:
-        """Initialize the plugin with a name."""
         self.name = name
 
-    def register(self, method_name: str = "__call__") -> Callable:
+    def register(self):
-        """Decorator to register a function as a plugin."""
+        """Register one implementation object under this plugin name."""
         if self.name is None:
             raise ValueError("Plugin name should be specified.")
 
-        if method_name in self._registry[self.name]:
+        cls = type(self)
-            logger.warning_rank0_once(f"Method {method_name} of plugin {self.name} is already registered.")
+        if self.name in cls._registry:
+            logger.warning_rank0_once(f"Plugin {self.name!r} is already registered under {cls.__name__}.")
 
-        def decorator(func: Callable) -> Callable:
+        def decorator(obj: Any) -> Any:
-            self._registry[self.name][method_name] = func
+            cls._registry[self.name] = obj
-            return func
+            return obj
 
         return decorator
 
+    @classmethod
+    def parse_params(cls, config: Any, params_cls: type[ParamsT]) -> ParamsT:
+        """Strictly convert config to the params dataclass used by one plugin entrypoint."""
+        if not is_dataclass(params_cls):
+            raise TypeError(f"{cls.__name__} params must be a dataclass type, got {params_cls!r}.")
+        if isinstance(config, params_cls):
+            return config
+        if config is None:
+            values = {}
+        elif isinstance(config, dict):
+            values = dict(config)
+        else:
+            raise TypeError(
+                f"{cls.__name__} config must be a mapping or {params_cls.__name__}, got {type(config).__name__}."
+            )
+
+        known = {item.name for item in fields(params_cls)}
+        unknown = set(values) - known
+        if unknown:
+            raise ValueError(
+                f"Unknown params for {cls.__name__}.{params_cls.__name__}: {sorted(unknown)}. "
+                f"Expected: {sorted(known)}"
+            )
+
+        return params_cls(**values)
+
+    def _resolve(self) -> Any:
+        cls = type(self)
+        if self.name is None:
+            raise ValueError(f"{cls.__name__} must be constructed with a name.")
+        if self.name not in cls._registry:
+            raise ValueError(f"Plugin {self.name!r} is not registered under {cls.__name__}.")
+        return cls._registry[self.name]
+
     def __call__(self, *args, **kwargs) -> Any:
-        """Call the registered function with the given arguments."""
+        return self._resolve()(*args, **kwargs)
-        return self["__call__"](*args, **kwargs)
 
-    def __getattr__(self, method_name: str) -> Callable:
+    def __getattr__(self, attr: str) -> Any:
-        """Get the registered function with the given name."""
+        return getattr(self._resolve(), attr)
-        return self[method_name]
-
-    def __getitem__(self, method_name: str) -> Callable:
-        """Get the registered function with the given name."""
-        if method_name not in self._registry[self.name]:
-            raise ValueError(f"Method {method_name} of plugin {self.name} is not registered.")
-
-        return self._registry[self.name][method_name]
-
-
-if __name__ == "__main__":
-    """
-    python -m llamafactory.v1.utils.plugin
-    """
-
-    class PrintPlugin(BasePlugin):
-        def again(self):  # optional
-            self["again"]()
-
-    @PrintPlugin("hello").register()
-    def print_hello():
-        print("Hello world!")
-
-    @PrintPlugin("hello").register("again")
-    def print_hello_again():
-        print("Hello world! Again.")
-
-    PrintPlugin("hello")()
-    PrintPlugin("hello").again()

src/llamafactory/v1/utils/types.py

@@ -78,19 +78,6 @@ class DatasetInfo(TypedDict, total=False):
     """Is streaming dataset, default to False."""
 
 
-class DistributedConfig(TypedDict, total=False):
-    mp_replicate_size: NotRequired[int]
-    """Model parallel replicate size, default to 1."""
-    mp_shard_size: NotRequired[int]
-    """Model parallel shard size, default to world_size // mp_replicate_size."""
-    dp_size: NotRequired[int]
-    """Data parallel size, default to world_size // cp_size."""
-    cp_size: NotRequired[int]
-    """Context parallel size, default to 1."""
-    timeout: NotRequired[int]
-    """Timeout for distributed communication, default to 600."""
-
-
 class Content(TypedDict):
     type: Literal["text", "reasoning", "tool_call", "image_url", "video_url", "audio_url"]
     """Type of the content."""

tests_v1/config/test_args_parser.py

@@ -27,10 +27,9 @@ def test_get_args_from_yaml(tmp_path: Path):
         model_class: llm
         kernel_config:
             name: auto
-            include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null
         peft_config:
             name: lora
-            lora_rank: 0.8
+            r: 8
         quant_config: null
 
         ### data
@@ -60,9 +59,8 @@ def test_get_args_from_yaml(tmp_path: Path):
         assert data_args.train_dataset == "llamafactory/v1-sft-demo"
         assert model_args.model == "llamafactory/tiny-random-qwen3"
         assert model_args.kernel_config.name == "auto"
-        assert model_args.kernel_config.get("include_kernels") == "auto"
         assert model_args.peft_config.name == "lora"
-        assert model_args.peft_config.get("lora_rank") == 0.8
+        assert model_args.peft_config.get("r") == 8
         assert training_args.output_dir == "outputs/test_run"
         assert training_args.micro_batch_size == 1
         assert training_args.global_batch_size == 1

tests_v1/core/test_model_loader.py

@@ -30,9 +30,7 @@ def test_tiny_qwen():
 def test_tiny_qwen_with_kernel_plugin():
     from llamafactory.v1.plugins.model_plugins.kernels.ops.rms_norm.npu_rms_norm import npu_rms_norm_forward
 
-    model_args = ModelArguments(
+    model_args = ModelArguments(model="llamafactory/tiny-random-qwen3", kernel_config={"name": "auto"})
-        model="llamafactory/tiny-random-qwen3", kernel_config={"name": "auto", "include_kernels": "auto"}
-    )
     model_engine = ModelEngine(model_args)
     # test enable apply kernel plugin
     if hasattr(torch, "npu"):

tests_v1/plugins/model_plugins/test_kernel_plugin.py

@@ -30,13 +30,13 @@ def _apply_kernel(rank) -> None:
             if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"):
                 del sys.modules[k]
 
-        from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_default_kernels
+        from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels
 
         model = AutoModelForCausalLM.from_pretrained("llamafactory/tiny-random-qwen3")
         original_rmsnorm_forward = model.model.layers[0].input_layernorm.forward
         original_swiglu_forward = model.model.layers[0].mlp.forward
 
-        model = apply_default_kernels(model=model, include_kernels="npu_fused_rmsnorm")
+        model = apply_kernels(model=model, config={"name": "npu_fused_rmsnorm"})
 
         assert model.model.layers[0].input_layernorm.forward.__func__ is not original_rmsnorm_forward.__func__
         assert model.model.layers[0].mlp.forward.__func__ is original_swiglu_forward.__func__
@@ -53,13 +53,13 @@ def _apply_all_kernels(rank) -> None:
             if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"):
                 del sys.modules[k]
 
-        from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_default_kernels
+        from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels
 
         model = AutoModelForCausalLM.from_pretrained("llamafactory/tiny-random-qwen3")
         original_rmsnorm_forward = model.model.layers[0].input_layernorm.forward
         original_swiglu_forward = model.model.layers[0].mlp.forward
 
-        model = apply_default_kernels(model=model, include_kernels=True)
+        model = apply_kernels(model=model, config={"name": "auto"})
 
         assert model.model.layers[0].input_layernorm.forward.__func__ is not original_rmsnorm_forward.__func__
         assert model.model.layers[0].mlp.forward.__func__ is not original_swiglu_forward.__func__

tests_v1/plugins/model_plugins/test_ulysses_cp.py

@@ -18,6 +18,7 @@ import torch.multiprocessing as mp
 
 from llamafactory.v1.accelerator.interface import DistributedInterface
 from llamafactory.v1.config.model_args import ModelArguments
+from llamafactory.v1.config.training_args import TrainingArguments
 from llamafactory.v1.core.model_engine import ModelEngine
 from llamafactory.v1.plugins.model_plugins.parallelization.sequence_parallel import (
     SequenceParallelModelPlugin,
@@ -33,15 +34,14 @@ def _test_sequence_parallel_loss(
     with dist_env(local_rank, world_size, master_port):
         model_args = ModelArguments(model="llamafactory/tiny-random-qwen3")
 
-        # Initialize distributed interface with config
+        training_args = TrainingArguments(cp_mode="ulysses", cp_size=cp_size, dp_size=dp_size)
-        dist_config = {"cp_mode": "ulysses", "cp_size": cp_size, "dp_size": dp_size}
+        DistributedInterface(training_args)
-        DistributedInterface(dist_config)
 
         # Now create model engine
         model_engine = ModelEngine(model_args=model_args)
 
         # Apply sequence parallel plugin
-        SequenceParallelModelPlugin(dist_config.get("cp_mode", "ulysses"))(model_engine.model, dist_config)
+        SequenceParallelModelPlugin(training_args.cp_mode)(model_engine.model, training_args.cp_size)
 
         input_ids = torch.arange(1, batch_size * 5 + 1, dtype=torch.long).view(batch_size, 5)
         model_inputs = {

tests_v1/trainers/test_dpo_loss_precision.py

@@ -28,6 +28,7 @@ from llamafactory.v1.trainers.dpo_trainer import DPOTrainer, compute_sigmoid_dpo
 # Mock helpers
 # ==============================================================================
 
+
 def _make_mock_v1(
     pref_beta: float = 0.1,
     dpo_label_smoothing: float = 0.0,
@@ -67,6 +68,7 @@ R_REJECTED = torch.tensor([-3.2, -2.7, -4.2, -1.8])
 # Test 1 — Core loss correctness (pure function ↔ v1 instance ↔ v0/TRL)
 # ==============================================================================
 
+
 def test_sigmoid_dpo_loss_correctness():
     """Comprehensive correctness check for compute_sigmoid_dpo_loss and its wrapper."""
     # ---- 1a: pure function matches instance method ----
@@ -78,7 +80,12 @@ def test_sigmoid_dpo_loss_correctness():
     # ---- 1b: v1 matches v0 (TRL) on fixed inputs ----
     v0 = _make_mock_v0_dpo(beta=0.1)
     v0_losses, _, _ = CustomDPOTrainer.dpo_loss(
-        v0, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED, loss_type="sigmoid",
+        v0,
+        P_CHOSEN,
+        P_REJECTED,
+        R_CHOSEN,
+        R_REJECTED,
+        loss_type="sigmoid",
     )
     torch.testing.assert_close(actual, v0_losses, rtol=1e-6, atol=1e-6)
 
@@ -87,7 +94,12 @@ def test_sigmoid_dpo_loss_correctness():
         v0b = _make_mock_v0_dpo(beta=beta)
         v1b = _make_mock_v1(pref_beta=beta)
         vl, _, _ = CustomDPOTrainer.dpo_loss(
-            v0b, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED, loss_type="sigmoid",
+            v0b,
+            P_CHOSEN,
+            P_REJECTED,
+            R_CHOSEN,
+            R_REJECTED,
+            loss_type="sigmoid",
         )
         v1l = DPOTrainer._sigmoid_dpo_loss(v1b, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED)
         torch.testing.assert_close(v1l, vl, rtol=1e-6, atol=1e-6)
@@ -97,7 +109,12 @@ def test_sigmoid_dpo_loss_correctness():
         v0s = _make_mock_v0_dpo(beta=0.1, label_smoothing=ls)
         v1s = _make_mock_v1(pref_beta=0.1, dpo_label_smoothing=ls)
         vl, _, _ = CustomDPOTrainer.dpo_loss(
-            v0s, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED, loss_type="sigmoid",
+            v0s,
+            P_CHOSEN,
+            P_REJECTED,
+            R_CHOSEN,
+            R_REJECTED,
+            loss_type="sigmoid",
         )
         v1l = DPOTrainer._sigmoid_dpo_loss(v1s, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED)
         torch.testing.assert_close(v1l, vl, rtol=1e-6, atol=1e-6)
@@ -112,13 +129,17 @@ def test_sigmoid_dpo_loss_correctness():
     v1c = _make_mock_v1(pref_beta=0.1)
     loss_good = DPOTrainer._sigmoid_dpo_loss(
         v1c,
-        torch.tensor([-1.0]), torch.tensor([-10.0]),
+        torch.tensor([-1.0]),
-        torch.tensor([-3.0]), torch.tensor([-3.0]),
+        torch.tensor([-10.0]),
+        torch.tensor([-3.0]),
+        torch.tensor([-3.0]),
     )
     loss_bad = DPOTrainer._sigmoid_dpo_loss(
         v1c,
-        torch.tensor([-10.0]), torch.tensor([-1.0]),
+        torch.tensor([-10.0]),
-        torch.tensor([-3.0]), torch.tensor([-3.0]),
+        torch.tensor([-1.0]),
+        torch.tensor([-3.0]),
+        torch.tensor([-3.0]),
     )
     assert loss_good.item() < loss_bad.item()
 
@@ -147,6 +168,7 @@ def test_sigmoid_dpo_loss_correctness():
 # Test 2 — Random cross-validation & reward equivalence
 # ==============================================================================
 
+
 def test_cross_validate_and_rewards():
     """Randomised v0↔v1 cross-validation (50 seeds) + reward-margin check."""
     torch.manual_seed(42)
@@ -162,7 +184,12 @@ def test_cross_validate_and_rewards():
         v1 = _make_mock_v1(pref_beta=beta, dpo_label_smoothing=ls)
 
         v0_loss, _, _ = CustomDPOTrainer.dpo_loss(
-            v0, pc, pr, rc, rr, loss_type="sigmoid",
+            v0,
+            pc,
+            pr,
+            rc,
+            rr,
+            loss_type="sigmoid",
         )
         v1_loss = DPOTrainer._sigmoid_dpo_loss(v1, pc, pr, rc, rr)
         torch.testing.assert_close(v1_loss, v0_loss, rtol=1e-5, atol=1e-5)
@@ -184,6 +211,7 @@ def test_cross_validate_and_rewards():
 # Test 3 — End-to-end: log-prob extraction + synthetic batch + LD-DPO
 # ==============================================================================
 
+
 def _make_batch(num_pairs, seq_len, vocab_size, prompt_len=3, chosen_len=None, rejected_len=None):
     if chosen_len is None or rejected_len is None:
         rlen = (seq_len - prompt_len) // 2
@@ -198,8 +226,8 @@ def _make_batch(num_pairs, seq_len, vocab_size, prompt_len=3, chosen_len=None, r
     labels[:, :prompt_len] = IGNORE_INDEX
 
     token_type_ids = torch.zeros(num_pairs, actual, dtype=torch.long)
-    token_type_ids[:, prompt_len:prompt_len + chosen_len] = 1
+    token_type_ids[:, prompt_len : prompt_len + chosen_len] = 1
-    token_type_ids[:, prompt_len + chosen_len:] = 2
+    token_type_ids[:, prompt_len + chosen_len :] = 2
 
     torch.manual_seed(99)
     logits = torch.randn(num_pairs, actual, vocab_size)
@@ -226,7 +254,12 @@ def test_logp_extraction_and_e2e_loss():
 
     # --- unequal-length (LD-DPO) batch ---
     ids2, labels2, tt_ids2, logits2 = _make_batch(
-        1, 11, 64, prompt_len=2, chosen_len=6, rejected_len=3,
+        1,
+        11,
+        64,
+        prompt_len=2,
+        chosen_len=6,
+        rejected_len=3,
     )
     v1_ld = _make_mock_v1(pref_beta=0.1, ld_alpha=0.5)
 

tests_v1/trainers/test_fsdp2_dpo_trainer.py

@@ -33,7 +33,6 @@ template: qwen3_nothink
 
 kernel_config:
     name: auto
-    include_kernels: auto
 
 quant_config: null
 

tests_v1/trainers/test_fsdp2_sft_trainer.py

@@ -30,7 +30,6 @@ model_class: llm
 
 kernel_config:
     name: auto
-    include_kernels: auto
 
 quant_config: null