[v1] refactor registry plugin structure and params

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 @@
 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 @@ def __new__(cls, *args: Any, **kwargs: Any) -> "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 @@ def __init__(self, config: DistributedConfig | None = None) -> None:
         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_shard_size=config.get("mp_shard_size", None),
-                dp_size=config.get("dp_size", None),
-                cp_size=config.get("cp_size", 1),
+                mp_replicate_size=training_args.mp_replicate_size,
+                mp_shard_size=training_args.mp_shard_size,
+                dp_size=training_args.dp_size,
+                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

@@ -72,7 +72,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,
@@ -145,6 +169,12 @@ def __post_init__(self) -> None:
         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
 
         if str(self.batching_strategy) == str(BatchingStrategy.DYNAMIC_BATCHING):
             if self.max_steps is None or self.max_steps <= 0:

src/llamafactory/v1/core/base_trainer.py

@@ -93,9 +93,12 @@ def __init__(
         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,
@@ -138,7 +141,7 @@ def __init__(
         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(
@@ -151,7 +154,7 @@ def __init__(
                     "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 (
@@ -182,9 +185,9 @@ def _shard_model(self) -> None:
                 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,
@@ -255,7 +258,7 @@ def fit(self) -> None:
                 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,
                         )
@@ -281,7 +284,7 @@ def fit(self) -> None:
                 else:
                     grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.max_grad_norm).item()
 
-                    if self.args.dist_config and self.args.dist_config.get("cp_size", 1) > 1:
+                    if self.args.cp_size > 1:
                         grad_norm = grad_norm**2
                         grad_norm = DistributedInterface().all_reduce(grad_norm, op=ReduceOp.SUM, dim=Dim.CP)
                         grad_norm = grad_norm**0.5
@@ -342,7 +345,7 @@ def fit(self) -> None:
     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 @@ def _prefix_cuts(sample: Sample) -> list[int | None]:
         messages = sample.get("messages")
         if not messages:
             return [None]
-        cuts = [
-            i + 1 for i, m in enumerate(messages) if m["role"] == "assistant" and m.get("loss_weight", 1.0) > 1e-6
-        ]
+        cuts = [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: