From 26d949ae87879d0e402c95a910db098daaf4d238 Mon Sep 17 00:00:00 2001 From: jiaqiw09 Date: Sun, 12 Jul 2026 14:22:00 +0800 Subject: [PATCH 1/2] refactor register and params --- .../train_full_fsdp2_batching_normal.yaml | 1 - .../train_full_fsdp2_dynamic_batching.yaml | 1 - ...train_full_fsdp2_dynamic_padding_free.yaml | 1 - .../train_full_fsdp2_padding_free.yaml | 1 - .../v1/train_freeze/train_freeze_sft.yaml | 1 - .../v1/train_full/train_full_deepspeed.yaml | 1 - examples/v1/train_full/train_full_fsdp2.yaml | 1 - .../train_full/train_full_liger_kernel.yaml | 1 - .../v1/train_full/train_full_ulysses_cp.yaml | 5 +- examples/v1/train_lora/train_lora_dpo.yaml | 1 - examples/v1/train_lora/train_lora_sft.yaml | 1 - .../v1/train_lora/train_lora_sft_rank0.yaml | 1 - examples/v1/train_qlora/quantization.yaml | 1 - src/llamafactory/v1/accelerator/interface.py | 27 +- src/llamafactory/v1/config/training_args.py | 32 ++- src/llamafactory/v1/core/base_trainer.py | 21 +- src/llamafactory/v1/core/data_engine.py | 4 +- src/llamafactory/v1/core/model_engine.py | 38 ++- src/llamafactory/v1/core/rendering/format.py | 9 +- .../v1/core/rendering/rendering.py | 10 +- src/llamafactory/v1/core/utils/checkpoint.py | 4 +- .../plugins/model_plugins/deepspeed_utils.py | 36 ++- .../v1/plugins/model_plugins/kernels/base.py | 72 ++---- .../model_plugins/kernels/interface.py | 190 +++++--------- .../model_plugins/kernels/liger_kernel_ops.py | 19 +- .../kernels/ops/mlp/cuda_fused_moe.py | 18 +- .../kernels/ops/mlp/npu_fused_moe.py | 7 +- .../kernels/ops/mlp/npu_swiglu.py | 7 +- .../kernels/ops/rms_norm/npu_rms_norm.py | 7 +- .../kernels/ops/rope/npu_rope.py | 7 +- .../plugins/model_plugins/kernels/registry.py | 96 ------- .../parallelization/sequence_parallel.py | 7 +- .../v1/plugins/model_plugins/peft.py | 167 ++++++------ .../v1/plugins/model_plugins/quantization.py | 122 +++++---- .../v1/plugins/trainer_plugins/batching.py | 243 +++++++++--------- .../trainer_plugins/distributed/base.py | 50 ++++ .../trainer_plugins/distributed/hub.py | 94 ------- .../trainer_plugins/distributed/interface.py | 115 +++++++++ src/llamafactory/v1/trainers/dpo_trainer.py | 14 +- src/llamafactory/v1/trainers/rm_trainer.py | 5 +- src/llamafactory/v1/trainers/sft_trainer.py | 2 +- src/llamafactory/v1/utils/plugin.py | 134 +++++----- src/llamafactory/v1/utils/types.py | 13 - tests_v1/config/test_args_parser.py | 6 +- tests_v1/core/test_model_loader.py | 4 +- .../model_plugins/test_kernel_plugin.py | 8 +- .../plugins/model_plugins/test_ulysses_cp.py | 8 +- tests_v1/trainers/test_dpo_loss_precision.py | 55 +++- tests_v1/trainers/test_fsdp2_dpo_trainer.py | 1 - tests_v1/trainers/test_fsdp2_sft_trainer.py | 1 - 50 files changed, 771 insertions(+), 899 deletions(-) delete mode 100644 src/llamafactory/v1/plugins/model_plugins/kernels/registry.py create mode 100644 src/llamafactory/v1/plugins/trainer_plugins/distributed/base.py delete mode 100644 src/llamafactory/v1/plugins/trainer_plugins/distributed/hub.py create mode 100644 src/llamafactory/v1/plugins/trainer_plugins/distributed/interface.py diff --git a/examples/v1/train_batching_strategy/train_full_fsdp2_batching_normal.yaml b/examples/v1/train_batching_strategy/train_full_fsdp2_batching_normal.yaml index e14d52fd77..854c9a6639 100644 --- a/examples/v1/train_batching_strategy/train_full_fsdp2_batching_normal.yaml +++ b/examples/v1/train_batching_strategy/train_full_fsdp2_batching_normal.yaml @@ -6,7 +6,6 @@ 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 diff --git a/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_batching.yaml b/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_batching.yaml index ff0dc2cc2c..98e46c9e92 100644 --- a/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_batching.yaml +++ b/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_batching.yaml @@ -5,7 +5,6 @@ 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 diff --git a/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml b/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml index 9898a50d89..8bd674459b 100644 --- a/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml +++ b/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml @@ -5,7 +5,6 @@ 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 diff --git a/examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml b/examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml index ec5babca00..0d5577ae5c 100644 --- a/examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml +++ b/examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml @@ -5,7 +5,6 @@ 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 diff --git a/examples/v1/train_freeze/train_freeze_sft.yaml b/examples/v1/train_freeze/train_freeze_sft.yaml index a3e062ed9f..27bdc20838 100644 --- a/examples/v1/train_freeze/train_freeze_sft.yaml +++ b/examples/v1/train_freeze/train_freeze_sft.yaml @@ -12,7 +12,6 @@ peft_config: # Kernel Config kernel_config: name: auto - include_kernels: auto # FSDP Config dist_config: diff --git a/examples/v1/train_full/train_full_deepspeed.yaml b/examples/v1/train_full/train_full_deepspeed.yaml index 332c327177..482053513d 100644 --- a/examples/v1/train_full/train_full_deepspeed.yaml +++ b/examples/v1/train_full/train_full_deepspeed.yaml @@ -4,7 +4,6 @@ model_class: llm kernel_config: name: auto - include_kernels: auto dist_config: name: deepspeed diff --git a/examples/v1/train_full/train_full_fsdp2.yaml b/examples/v1/train_full/train_full_fsdp2.yaml index c4796a7a25..98e0c29772 100644 --- a/examples/v1/train_full/train_full_fsdp2.yaml +++ b/examples/v1/train_full/train_full_fsdp2.yaml @@ -4,7 +4,6 @@ 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 diff --git a/examples/v1/train_full/train_full_liger_kernel.yaml b/examples/v1/train_full/train_full_liger_kernel.yaml index b6face99c7..0725c6ab8b 100644 --- a/examples/v1/train_full/train_full_liger_kernel.yaml +++ b/examples/v1/train_full/train_full_liger_kernel.yaml @@ -5,7 +5,6 @@ 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 diff --git a/examples/v1/train_full/train_full_ulysses_cp.yaml b/examples/v1/train_full/train_full_ulysses_cp.yaml index 6ed486a1e9..1ee5b1439b 100644 --- a/examples/v1/train_full/train_full_ulysses_cp.yaml +++ b/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 diff --git a/examples/v1/train_lora/train_lora_dpo.yaml b/examples/v1/train_lora/train_lora_dpo.yaml index 32beec105a..af892063e8 100644 --- a/examples/v1/train_lora/train_lora_dpo.yaml +++ b/examples/v1/train_lora/train_lora_dpo.yaml @@ -14,7 +14,6 @@ peft_config: # Kernel Config kernel_config: name: auto - include_kernels: auto # FSDP Config dist_config: diff --git a/examples/v1/train_lora/train_lora_sft.yaml b/examples/v1/train_lora/train_lora_sft.yaml index 9596e6be5c..9193125aea 100644 --- a/examples/v1/train_lora/train_lora_sft.yaml +++ b/examples/v1/train_lora/train_lora_sft.yaml @@ -13,7 +13,6 @@ peft_config: # Kernel Config kernel_config: name: auto - include_kernels: auto # FSDP Config dist_config: diff --git a/examples/v1/train_lora/train_lora_sft_rank0.yaml b/examples/v1/train_lora/train_lora_sft_rank0.yaml index d9dc711a07..fd60141b6b 100644 --- a/examples/v1/train_lora/train_lora_sft_rank0.yaml +++ b/examples/v1/train_lora/train_lora_sft_rank0.yaml @@ -13,7 +13,6 @@ peft_config: # Kernel Config kernel_config: name: auto - include_kernels: auto # FSDP Config dist_config: diff --git a/examples/v1/train_qlora/quantization.yaml b/examples/v1/train_qlora/quantization.yaml index b0dcc19d0f..2ff0f581bf 100644 --- a/examples/v1/train_qlora/quantization.yaml +++ b/examples/v1/train_qlora/quantization.yaml @@ -13,7 +13,6 @@ peft_config: # Kernel Config kernel_config: name: auto - include_kernels: auto # FSDP Config dist_config: diff --git a/src/llamafactory/v1/accelerator/interface.py b/src/llamafactory/v1/accelerator/interface.py index a43885a552..530685551f 100644 --- a/src/llamafactory/v1/accelerator/interface.py +++ b/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()) diff --git a/src/llamafactory/v1/config/training_args.py b/src/llamafactory/v1/config/training_args.py index a69060f69e..df850deae6 100644 --- a/src/llamafactory/v1/config/training_args.py +++ b/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: diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index eda95c7693..0c6bf22bcc 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/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 diff --git a/src/llamafactory/v1/core/data_engine.py b/src/llamafactory/v1/core/data_engine.py index ff8adaa450..810cb43adf 100644 --- a/src/llamafactory/v1/core/data_engine.py +++ b/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: diff --git a/src/llamafactory/v1/core/model_engine.py b/src/llamafactory/v1/core/model_engine.py index 4df8d060e5..209cc494b4 100644 --- a/src/llamafactory/v1/core/model_engine.py +++ b/src/llamafactory/v1/core/model_engine.py @@ -69,24 +69,25 @@ def __init__( """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) - self._deepspeed_zero3_plugin = None - self._deepspeed_zero3_enabled = False + teardown_deepspeed_zero3_model_loading(plugin) else: self.model = self._init_model() @@ -142,9 +143,7 @@ def _init_model(self) -> HFModel: init_kwargs = QuantizationPlugin(self.args.quant_config.name)( init_kwargs=init_kwargs, - config=self.model_config, - tokenizer=self.processor, - model_args=self.args, + quant_config=self.args.quant_config, is_trainable=self.is_train, ) @@ -200,17 +199,16 @@ def _init_model(self) -> HFModel: 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 - - kernel_config = self.args.kernel_config - 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) + from ..plugins.model_plugins.kernels.interface import apply_kernels + + model = apply_kernels(model, self.args.kernel_config, require_logits=self.is_train) return model diff --git a/src/llamafactory/v1/core/rendering/format.py b/src/llamafactory/v1/core/rendering/format.py index 293c3b1b33..048e612179 100644 --- a/src/llamafactory/v1/core/rendering/format.py +++ b/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( - 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"]}} - ) + raise ValueError(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"]}}) 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 - diff --git a/src/llamafactory/v1/core/rendering/rendering.py b/src/llamafactory/v1/core/rendering/rendering.py index 4dc5e2d689..87a62e87f2 100644 --- a/src/llamafactory/v1/core/rendering/rendering.py +++ b/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 @@ def render_messages( Returns: ModelInput with input_ids, attention_mask, labels, and loss_weights. """ - return _render_messages( - self.processor, - messages, - tools, - is_generate, - **kwargs - ) + return _render_messages(self.processor, messages, tools, is_generate, **kwargs) def process_samples(self, samples: list[Sample]) -> list[ModelInput]: """Process samples to model input. diff --git a/src/llamafactory/v1/core/utils/checkpoint.py b/src/llamafactory/v1/core/utils/checkpoint.py index 40ed4f01c6..c0f18f44a6 100644 --- a/src/llamafactory/v1/core/utils/checkpoint.py +++ b/src/llamafactory/v1/core/utils/checkpoint.py @@ -251,7 +251,7 @@ def save(self, epoch: int) -> None: ) 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 @@ def resume(self, ckpt_path: str) -> None: 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, diff --git a/src/llamafactory/v1/plugins/model_plugins/deepspeed_utils.py b/src/llamafactory/v1/plugins/model_plugins/deepspeed_utils.py index de737ed6f7..d31a2dad71 100644 --- a/src/llamafactory/v1/plugins/model_plugins/deepspeed_utils.py +++ b/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): - """Enable transformers' ZeRO-3-aware model loading for the current thread.""" - config_file = dist_config.get("config_file") +def setup_deepspeed_zero3_model_loading(): + """Enable ZeRO-3-aware model loading for the registered backend config.""" + 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 " diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/base.py b/src/llamafactory/v1/plugins/model_plugins/kernels/base.py index 3916562b19..345eb6a1af 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/base.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/base.py @@ -12,76 +12,34 @@ # 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 ensure_methods_implemented from ....utils.types import HFModel 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. - """ + _device: DeviceType | list[DeviceType] - _kernel_id: Any = "" # kernel ID, any hashable value to identify a kernel implementation - _device: list[DeviceType] = [DeviceType.CPU] # "cuda", "npu", "cpu", etc. + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + ensure_methods_implemented(cls) @classmethod - def get_kernel_id(cls) -> str: - """Returns the unique identifier for the kernel.""" - return cls._kernel_id + def check_deps(cls) -> None: + supported = cls._device if isinstance(cls._device, list) else [cls._device] + current = get_current_accelerator().type + if current not in supported: + raise RuntimeError(f"Kernel {cls.__name__!r} requires {supported}, current accelerator is {current}.") @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 + def apply(cls, **kwargs) -> HFModel: + cls.check_deps() + return cls._apply(**kwargs) @classmethod @abstractmethod - def apply(cls, **kwargs) -> HFModel: - """Applies the kernel optimization to the model. - - Args: - **kwargs: Arbitrary keyword arguments, usually containing the model instance and the kernel configuration. - - Returns: - HFModel: The model with the kernel applied. - - 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 + def _apply(cls, **kwargs) -> HFModel: ... diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/interface.py b/src/llamafactory/v1/plugins/model_plugins/kernels/interface.py index 1646448e94..4653f4ee23 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/interface.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/interface.py @@ -12,174 +12,104 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""The definition of kernel interface. - -Init Phase: -1. Scan all kernels. -2. Register default kernels. -3. Define kernel plugin. - -""" - -import importlib -from pathlib import Path +from typing import Any from ....utils import logging from ....utils.plugin import BasePlugin from ....utils.types import HFModel -from .registry import Registry logger = logging.get_logger(__name__) -def scan_all_kernels(): - """Scan all kernels in the ``ops`` directory. - - Scans the ``ops`` directory for all ``.py`` files and attempts to import them. - 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" +class KernelPlugin(BasePlugin): + """Plugin family for model kernel optimizations.""" - if not ops_path.exists(): - return - base_package = __package__ +@KernelPlugin("npu_fused_rmsnorm").register() +def apply_npu_fused_rmsnorm(model: HFModel, **kwargs) -> HFModel: + from .ops.rms_norm.npu_rms_norm import NpuRMSNormKernel - for file_path in ops_path.rglob("*.py"): - if file_path.name == "__init__.py": - continue + return NpuRMSNormKernel.apply(model=model, **kwargs) - # 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}" +@KernelPlugin("npu_fused_rope").register() +def apply_npu_fused_rope(model: HFModel, **kwargs) -> HFModel: + from .ops.rope.npu_rope import NpuRoPEKernel - 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 NpuRoPEKernel.apply(model=model, **kwargs) - return Registry.get_registered_kernels() +@KernelPlugin("npu_fused_swiglu").register() +def apply_npu_fused_swiglu(model: HFModel, **kwargs) -> HFModel: + from .ops.mlp.npu_swiglu import NpuSwiGluKernel -default_kernels = scan_all_kernels() + return NpuSwiGluKernel.apply(model=model, **kwargs) -def get_default_kernels(): - """Get a list of default registered kernel IDs. +@KernelPlugin("npu_fused_moe").register() +def apply_npu_fused_moe(model: HFModel, **kwargs) -> HFModel: + from .ops.mlp.npu_fused_moe import NpuFusedMoEKernel - Returns: - list[str]: List of kernel IDs. - """ - return list(default_kernels.keys()) + return NpuFusedMoEKernel.apply(model=model, **kwargs) -def apply_kernel(kernel_id: str, **kwargs): - """Applies a specific kernel to the model. +@KernelPlugin("cuda_fused_moe").register() +def apply_cuda_fused_moe(model: HFModel, **kwargs) -> HFModel: + from .ops.mlp.cuda_fused_moe import CudaFusedMoEKernel - 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. + return CudaFusedMoEKernel.apply(model=model, **kwargs) - 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) +@KernelPlugin("liger_kernel").register() +def apply_liger_kernels(model: HFModel, config: dict[str, Any] | None = None, **kwargs) -> HFModel: + try: + from .liger_kernel_ops import LigerKernel + except ImportError as exc: + logger.warning_rank0(f"[Kernel] Failed to import liger_kernel ops, skip. Error: {exc}") + return model + require_logits = kwargs.get("require_logits", False) + if config is not None: + require_logits = config.get("require_logits", require_logits) + return LigerKernel.apply(use_kernels="auto", model=model, require_logits=require_logits) -class KernelPlugin(BasePlugin): - """Plugin for managing kernel optimizations.""" - pass +_AUTO_KERNELS = ("npu_fused_moe", "npu_fused_rmsnorm", "npu_fused_rope", "npu_fused_swiglu") @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" +def apply_auto_kernels(model: HFModel, **kwargs) -> HFModel: + for kernel_name in _AUTO_KERNELS: + try: + model = KernelPlugin(kernel_name)(model=model, **kwargs) + except RuntimeError: + continue + return model - for kernel in use_kernels: - if kernel not in default_kernels: - raise ValueError(f"Kernel {kernel} not found") - apply_kernel(kernel, model=model) +def apply_kernels(model: HFModel, config: dict[str, Any], require_logits: bool = False) -> HFModel: + """Apply the comma-separated kernel names selected by ``kernel_config.name``.""" + kernel_names = config.get("name") + if not isinstance(kernel_names, str): + raise TypeError("kernel_config.name must be a string.") + names = [name.strip() for name in kernel_names.split(",") if name.strip()] + if not names: + raise ValueError("kernel_config.name must contain at least one kernel name.") + + for name in names: + model = KernelPlugin(name)(model=model, config=config, require_logits=require_logits) return model -@KernelPlugin("liger_kernel").register() -def apply_liger_kernels( - model: HFModel, - include_kernels: str = None, - require_logits: bool = False, -) -> HFModel: - """Applies Liger kernel to the model. - - Args: - model (HFModel): The model instance to apply kernels to. - include_kernels (str, optional): If ``"auto"`` or ``True``, apply Liger with - 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: - HFModel: The model with Liger kernel applied. - """ - if not include_kernels: +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 - 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: - from .liger_kernel_ops import LigerKernel - except ImportError as e: - logger.warning_rank0(f"[Kernel] Failed to import liger_kernel ops, skip. Error: {e}") - return model + return apply_kernels(model, {"name": "auto"}) + - return LigerKernel.apply(use_kernels=use_kernels, model=model, require_logits=require_logits) +def apply_kernel(kernel_id: str, **kwargs) -> HFModel: + return KernelPlugin(kernel_id)(**kwargs) diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/liger_kernel_ops.py b/src/llamafactory/v1/plugins/model_plugins/kernels/liger_kernel_ops.py index 1d6fd0dbd3..0068bb55c7 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/liger_kernel_ops.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/liger_kernel_ops.py @@ -22,7 +22,7 @@ import inspect -from ....accelerator.helper import DeviceType, get_current_accelerator +from ....accelerator.helper import DeviceType from ....utils.logging import get_logger from ....utils.types import HFModel from .base import BaseKernel @@ -47,20 +47,16 @@ class LigerKernel(BaseKernel): _device = [DeviceType.CUDA, DeviceType.NPU] @classmethod - def check_deps(cls) -> bool: + def check_deps(cls) -> None: """Checks if the required dependencies for the kernel are available.""" + super().check_deps() try: import liger_kernel # noqa: F401 - - return super().check_deps() except ImportError: - logger.warning_rank0( - "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 + raise RuntimeError("Liger kernel is not installed.") from None @classmethod - def apply(cls, **kwargs) -> "HFModel": + def _apply(cls, **kwargs) -> "HFModel": """Applies the Liger kernel to the model. Args: @@ -82,11 +78,6 @@ def apply(cls, **kwargs) -> "HFModel": if model is None: 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) model_type = getattr(model.config, "model_type", None) diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/cuda_fused_moe.py b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/cuda_fused_moe.py index 1a2cc4b69f..c78c4d3a01 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/cuda_fused_moe.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/cuda_fused_moe.py @@ -30,7 +30,6 @@ from ......accelerator.helper import DeviceType from ......utils.types import HFModel from ...base import BaseKernel -from ...registry import register_kernel from .triton_grouped_gemm import ( group_gemm_same_mn, group_gemm_same_nk, @@ -351,7 +350,6 @@ def _triton_moe_sparse_block_forward(self, hidden_states: torch.Tensor): # --------------------------------------------------------------------------- -@register_kernel class CudaFusedMoEKernel(BaseKernel): """Pure-Triton fused MoE kernel for NVIDIA CUDA GPUs. @@ -366,27 +364,19 @@ class CudaFusedMoEKernel(BaseKernel): _device = DeviceType.CUDA @classmethod - def check_deps(cls) -> bool: - if not super().check_deps(): - return False + def check_deps(cls) -> None: + super().check_deps() try: import triton # noqa: F401 - - return True except ImportError: - logger.info("cuda_fused_moe: Triton not available, kernel disabled.") - return False + raise RuntimeError("cuda_fused_moe requires Triton.") from None @classmethod - def apply(cls, **kwargs) -> HFModel: + 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 for arch in archs: diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_fused_moe.py b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_fused_moe.py index 41cab6020a..0145f5a7cb 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_fused_moe.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_fused_moe.py @@ -36,7 +36,6 @@ from ......utils.packages import is_transformers_version_greater_than from ......utils.types import HFModel from ...base import BaseKernel -from ...registry import register_kernel class GmmFunction(torch.autograd.Function): @@ -334,7 +333,6 @@ def qwen3moe_sparse_moe_block_forward(self, hidden_states: torch.Tensor): } -@register_kernel class NpuFusedMoEKernel(BaseKernel): """NPU Fused MoE Kernel implementation.""" @@ -342,7 +340,7 @@ class NpuFusedMoEKernel(BaseKernel): _device = DeviceType.NPU @classmethod - def apply(cls, **kwargs) -> HFModel: + def _apply(cls, **kwargs) -> HFModel: """Applies the NPU fused MoE kernel to the model. Args: @@ -359,9 +357,6 @@ def apply(cls, **kwargs) -> HFModel: 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 for arch in archs: diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_swiglu.py b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_swiglu.py index a45077bc00..6b75768b49 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_swiglu.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/mlp/npu_swiglu.py @@ -28,7 +28,6 @@ from ......accelerator.helper import DeviceType from ......utils.types import HFModel from ...base import BaseKernel -from ...registry import register_kernel try: @@ -86,7 +85,6 @@ def _npu_swiglu_gemma3ntext_forward(self, hidden_states): return down_proj -@register_kernel class NpuSwiGluKernel(BaseKernel): """NPU Kernel for fused SwiGLU activation.""" @@ -125,7 +123,7 @@ class NpuSwiGluKernel(BaseKernel): _device = DeviceType.NPU @classmethod - def apply(cls, **kwargs) -> "HFModel": + def _apply(cls, **kwargs) -> "HFModel": """Applies the NPU fused SwiGLU kernel to the model. Args: @@ -142,9 +140,6 @@ def apply(cls, **kwargs) -> "HFModel": 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 = { "Glm4MLP": _npu_swiglu_glm4_forward, diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/npu_rms_norm.py b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/npu_rms_norm.py index 04a4ad16f9..8ed045187e 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/npu_rms_norm.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rms_norm/npu_rms_norm.py @@ -29,7 +29,6 @@ from ......accelerator.helper import DeviceType from ......utils.types import HFModel from ...base import BaseKernel -from ...registry import register_kernel try: @@ -118,7 +117,6 @@ def npu_gated_rms_norm_forward(self, hidden_states, gate=None): return hidden_states.to(input_dtype) -@register_kernel class NpuRMSNormKernel(BaseKernel): """NPU kernel wrapper for RMSNorm that applies the replacement within a model.""" @@ -126,7 +124,7 @@ class NpuRMSNormKernel(BaseKernel): _device = DeviceType.NPU @classmethod - def apply(cls, **kwargs) -> "HFModel": + def _apply(cls, **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 @@ -147,9 +145,6 @@ def apply(cls, **kwargs) -> "HFModel": 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) for _, module in model.named_modules(): diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rope/npu_rope.py b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rope/npu_rope.py index 0f9d4d7e22..019ae9fdf5 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rope/npu_rope.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/ops/rope/npu_rope.py @@ -28,7 +28,6 @@ from ......utils.logging import get_logger from ......utils.types import HFModel from ...base import BaseKernel -from ...registry import register_kernel logger = get_logger(__name__) @@ -125,7 +124,6 @@ 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 class NpuRoPEKernel(BaseKernel): """NPU Kernel for Rotary Position Embedding.""" @@ -133,7 +131,7 @@ class NpuRoPEKernel(BaseKernel): _device = DeviceType.NPU @classmethod - def apply(cls, **kwargs) -> "HFModel": + def _apply(cls, **kwargs) -> "HFModel": """Apply RoPE acceleration by monkey-patching ``apply_rotary_pos_emb``. Iterates through the model's modules to find attention layers, identifies @@ -151,9 +149,6 @@ def apply(cls, **kwargs) -> "HFModel": 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__}.") diff --git a/src/llamafactory/v1/plugins/model_plugins/kernels/registry.py b/src/llamafactory/v1/plugins/model_plugins/kernels/registry.py deleted file mode 100644 index 2921f0aed4..0000000000 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/registry.py +++ /dev/null @@ -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 diff --git a/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py b/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py index 8d5073bb48..a2afe5ba23 100644 --- a/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py +++ b/src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py @@ -37,8 +37,8 @@ class SequenceParallelModelPlugin(BasePlugin): - def __call__(self, model, model_args): - return super().__call__(model, model_args) + def __call__(self, model, cp_size: int): + 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)) diff --git a/src/llamafactory/v1/plugins/model_plugins/peft.py b/src/llamafactory/v1/plugins/model_plugins/peft.py index 9552de0b48..9030a6363d 100644 --- a/src/llamafactory/v1/plugins/model_plugins/peft.py +++ b/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 @@ logger = logging.get_logger(__name__) -class LoraConfigDict(TypedDict, total=False): - name: Literal["lora"] +@dataclass +class LoraParams: + """Typed configuration for the LoRA PEFT plugin.""" + + name: Literal["lora"] = "lora" """Plugin name.""" - r: int - """Lora rank.""" - lora_alpha: int - """Lora alpha.""" - lora_dropout: float - """Lora dropout.""" - target_modules: Union[list[str], str] + r: int = 8 + """LoRA rank.""" + lora_alpha: int = 16 + """LoRA alpha.""" + lora_dropout: float = 0.05 + """LoRA dropout.""" + 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 - """Shard size for the export model.""" - export_hub_model_id: str - """Hub model ID for the export model.""" - infer_dtype: Literal["auto", "float16", "float32", "bfloat16"] - """Inference data type for the export model.""" - export_legacy_format: bool - """Use legacy format for the export model.""" - - -class FreezeConfigDict(TypedDict, total=False): - name: Literal["freeze"] + export_size: int = 5 + """Shard size for the exported model, in GB.""" + export_hub_model_id: str | None = None + """Hub model ID for the exported model.""" + infer_dtype: Literal["auto", "float16", "float32", "bfloat16"] = "auto" + """Inference data type for the exported model.""" + export_legacy_format: bool = False + """Use legacy format for the exported model.""" + + +@dataclass +class FreezeParams: + """Typed configuration for the freeze PEFT plugin.""" + + name: Literal["freeze"] = "freeze" """Plugin name.""" - freeze_trainable_layers: int - """Freeze trainable layers.""" - freeze_trainable_modules: Union[list[str], str] - """Freeze trainable modules.""" - freeze_extra_modules: list[str] - """Freeze extra modules.""" - cast_trainable_params_to_fp32: bool - """Cast trainable params to fp32.""" + freeze_trainable_layers: int = 2 + """Number of trainable layers.""" + freeze_trainable_modules: list[str] | str = "all" + """Trainable modules in the selected layers.""" + freeze_extra_modules: list[str] | str | None = field(default_factory=list) + """Extra non-hidden modules to train.""" + cast_trainable_params_to_fp32: bool = True + """Cast trainable parameters to float32.""" class PeftPlugin(BasePlugin): - def __call__(self, model: HFModel, config: dict, is_train: bool) -> HFModel: - return super().__call__(model, config, is_train) + def __call__(self, model: HFModel, peft_config: dict, is_train: bool) -> HFModel: + 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: - adapter_name_or_path = config.get("adapter_name_or_path") +def get_lora_model(model: HFModel, peft_config: dict | LoraParams, is_train: bool = False) -> HFModel: + 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), - lora_alpha=config.get("lora_alpha", 16), - lora_dropout=config.get("lora_dropout", 0.05), - use_rslora=config.get("use_rslora", False), - use_dora=config.get("use_dora", False), + r=peft_config.r, + lora_alpha=peft_config.lora_alpha, + lora_dropout=peft_config.lora_dropout, + use_rslora=peft_config.use_rslora, + 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_modules = config.get("freeze_trainable_modules", ["all"]) - freeze_extra_modules = config.get("freeze_extra_modules", []) - cast_trainable_params_to_fp32 = config.get("cast_trainable_params_to_fp32", True) + freeze_trainable_layers = peft_config.freeze_trainable_layers + freeze_trainable_modules = peft_config.freeze_trainable_modules + freeze_extra_modules = peft_config.freeze_extra_modules + 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 - if export_config is None: + raw_config = model_args.peft_config + if raw_config is None: raise ValueError("Please specify peft_config to merge and export model.") - - 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: + if raw_config.name != "lora": 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, - max_shard_size=f"{export_size}GB", - safe_serialization=not export_legacy_format, + export_peft_config.export_dir, + max_shard_size=f"{export_peft_config.export_size}GB", + 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: - logger.info_rank0(f"Pushing to hub: {export_hub_model_id}...") - model.push_to_hub(export_hub_model_id) + if export_peft_config.export_hub_model_id: + logger.info_rank0(f"Pushing to hub: {export_peft_config.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.") diff --git a/src/llamafactory/v1/plugins/model_plugins/quantization.py b/src/llamafactory/v1/plugins/model_plugins/quantization.py index 2ba74d47d2..b889b5cf72 100644 --- a/src/llamafactory/v1/plugins/model_plugins/quantization.py +++ b/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, - config: "PretrainedConfig" = None, - tokenizer: "PreTrainedTokenizer" = None, - model_args: "ModelArguments" = None, + init_kwargs: dict[str, Any] | None = None, + quant_config=None, is_trainable: bool = False, ) -> dict[str, Any]: - return super().__call__( - init_kwargs, config=config, tokenizer=tokenizer, model_args=model_args, is_trainable=is_trainable - ) + return super().__call__(init_kwargs, quant_config=quant_config, 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. - - Args: - init_kwargs (dict[str, Any]): The kwargs for model initialization. - **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 + 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.") + + logger.info_rank0(f"Loading {quant_config.quantization_bit}-bit quantized model.") + return QuantizationPlugin("bnb")(init_kwargs, quant_config=quant_config, is_trainable=is_trainable) @QuantizationPlugin("bnb").register() def quantization_with_bnb( init_kwargs: dict[str, Any], - model_args: "ModelArguments" = None, - **kwargs, + quant_config: dict | BnbParams, + is_trainable: bool = False, ) -> dict[str, Any]: - r"""Quantization with BNB.""" - logger.info_rank0("Using Bitsandbytes quantization.") - quantization_bit = model_args.quant_config.get("quantization_bit", None) + from transformers import BitsAndBytesConfig + + 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_use_double_quant=model_args.quant_config.get("double_quantization", True), - bnb_4bit_quant_type=model_args.quant_config.get("quantization_type", "nf4"), - bnb_4bit_quant_storage=model_args.quant_config.get( - "compute_dtype", torch.float16 - ), # crucial for fsdp+qlora + bnb_4bit_compute_dtype=quant_config.compute_dtype, + bnb_4bit_use_double_quant=quant_config.double_quantization, + bnb_4bit_quant_type=quant_config.quantization_type, + bnb_4bit_quant_storage=quant_config.compute_dtype, ) - else: - raise ValueError("Bitsandbytes only accepts 4-bit or 8-bit quantization.") - # TODO: improve deepspeed zero3 and fsdp detection. - if kwargs.get("is_trainable", False): + if is_trainable: 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 diff --git a/src/llamafactory/v1/plugins/trainer_plugins/batching.py b/src/llamafactory/v1/plugins/trainer_plugins/batching.py index 715f4cc5c9..4cbed315cb 100644 --- a/src/llamafactory/v1/plugins/trainer_plugins/batching.py +++ b/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 ...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: - """Return the raw data provider batch size for this batching strategy.""" - return self["get_data_provider_batch_size"](batch_info) + """Plugin family for batching strategy method groups.""" - 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. - """ - return self["compute_length"](data_provider, batch_info) +class BaseBatcher(ABC): + def __init_subclass__(cls, **kwargs) -> None: + 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: - """Fill the buffer with data.""" - return self["fill_buffer"](buffer, batch_info, next_samples) + ) -> None: ... - def generate_batch(self, buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None: - """Generate a batch from the buffer.""" - return self["generate_batch"](buffer, batch_info) + @staticmethod + @abstractmethod + 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") -def get_padding_free_data_provider_batch_size(batch_info: BatchInfo) -> int: - return batch_info["micro_batch_size"] * batch_info["num_micro_batch"] +@BatchingPlugin("padding_free").register() +class PaddingFreeBatcher(BaseBatcher): + @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") -def compute_padding_free_length(data_provider: DataLoader, batch_info: BatchInfo) -> int: - return len(data_provider) - - -@BatchingPlugin("padding_free").register("fill_buffer") -def fill_padding_free_buffer( - 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) - + @staticmethod + def fill_buffer( + 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 -@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 + buffer.put(samples) - 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: + @staticmethod + def generate_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 - batch.append(packed_micro_batch) - - return batch - - -@BatchingPlugin("dynamic_batching").register("get_data_provider_batch_size") -def get_dynamic_batching_data_provider_batch_size(batch_info: BatchInfo) -> int: - return 1 - - -@BatchingPlugin("dynamic_batching").register("compute_length") -def compute_dynamic_batching_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) + 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) -@BatchingPlugin("dynamic_batching").register("fill_buffer") -def fill_dynamic_batching_buffer( - buffer: StatefulBuffer, - 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 + return batch - buffer.put(samples) +@BatchingPlugin("dynamic_batching").register() +class DynamicBatcher(BaseBatcher): + @staticmethod + def get_data_provider_batch_size(batch_info: BatchInfo) -> int: + return 1 -@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 + @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) - 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))) + @staticmethod + def fill_buffer( + buffer: StatefulBuffer, + 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 - return batch + buffer.put(samples) + @staticmethod + def generate_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 -@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 + 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("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() +class DynamicPaddingFreeBatcher(BaseBatcher): + @staticmethod + def get_data_provider_batch_size(batch_info: BatchInfo) -> int: + return 1 -@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) + @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) -@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 + @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"] + 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 + 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.append(packed_batch) - return batch + return batch diff --git a/src/llamafactory/v1/plugins/trainer_plugins/distributed/base.py b/src/llamafactory/v1/plugins/trainer_plugins/distributed/base.py new file mode 100644 index 0000000000..2b05f76eb2 --- /dev/null +++ b/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: ... diff --git a/src/llamafactory/v1/plugins/trainer_plugins/distributed/hub.py b/src/llamafactory/v1/plugins/trainer_plugins/distributed/hub.py deleted file mode 100644 index b4d0babf0d..0000000000 --- a/src/llamafactory/v1/plugins/trainer_plugins/distributed/hub.py +++ /dev/null @@ -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) diff --git a/src/llamafactory/v1/plugins/trainer_plugins/distributed/interface.py b/src/llamafactory/v1/plugins/trainer_plugins/distributed/interface.py new file mode 100644 index 0000000000..cc4ae06ee8 --- /dev/null +++ b/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) diff --git a/src/llamafactory/v1/trainers/dpo_trainer.py b/src/llamafactory/v1/trainers/dpo_trainer.py index d8f289b345..fb2dfb59da 100644 --- a/src/llamafactory/v1/trainers/dpo_trainer.py +++ b/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 ( - -F.logsigmoid(beta * logits) * (1 - label_smoothing) - - F.logsigmoid(-beta * logits) * label_smoothing - ) + return -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 @@ def __init__( train_dataset, callbacks=None, ) -> None: - cp_size = args.dist_config.get("cp_size", 1) if args.dist_config is not None else 1 - if cp_size > 1: + if args.cp_size > 1: raise NotImplementedError("DPO trainer currently only supports cp_size == 1.") self.pref_loss = args.pref_loss @@ -378,7 +374,9 @@ def compute_loss(self, batch: BatchInput) -> Tensor: # 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) diff --git a/src/llamafactory/v1/trainers/rm_trainer.py b/src/llamafactory/v1/trainers/rm_trainer.py index 6c78fa6fa4..78694d47bc 100644 --- a/src/llamafactory/v1/trainers/rm_trainer.py +++ b/src/llamafactory/v1/trainers/rm_trainer.py @@ -76,8 +76,7 @@ def __init__( train_dataset, callbacks=None, ) -> None: - cp_size = args.dist_config.get("cp_size", 1) if args.dist_config is not None else 1 - if cp_size > 1: + if args.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 @@ def compute_loss(self, batch: BatchInput) -> Tensor: 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) diff --git a/src/llamafactory/v1/trainers/sft_trainer.py b/src/llamafactory/v1/trainers/sft_trainer.py index 5b266c611b..d842632838 100644 --- a/src/llamafactory/v1/trainers/sft_trainer.py +++ b/src/llamafactory/v1/trainers/sft_trainer.py @@ -31,7 +31,7 @@ def compute_loss(self, batch: BatchInput) -> Tensor: 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( diff --git a/src/llamafactory/v1/utils/plugin.py b/src/llamafactory/v1/utils/plugin.py index 4c06c88d33..8debd5e7ff 100644 --- a/src/llamafactory/v1/utils/plugin.py +++ b/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 collections.abc import Callable -from typing import Any +from __future__ import annotations + +from dataclasses import fields, is_dataclass +from typing import Any, TypeVar from . import logging logger = logging.get_logger(__name__) +ParamsT = TypeVar("ParamsT") -class BasePlugin: - """Base class for plugins. - - A plugin is a callable object that can be registered and called by name. - - Example usage: - ```python - class PrintPlugin(BasePlugin): - def again(self): # optional - self["again"]() - +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()) - @PrintPlugin("hello").register() - def print_hello(): - print("Hello world!") + 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}") - @PrintPlugin("hello").register("again") - def print_hello_again(): - print("Hello world! Again.") - +class BasePlugin: + """Route a plugin name to one function or static method-group class. - PrintPlugin("hello")() - PrintPlugin("hello").again() - ``` + Every plugin family subclass owns an isolated registry. Parameter schemas + deliberately do not live here; each plugin entrypoint parses its own config. """ - _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: - """Decorator to register a function as a plugin.""" + def register(self): + """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]: - logger.warning_rank0_once(f"Method {method_name} of plugin {self.name} is already registered.") + cls = type(self) + 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: - self._registry[self.name][method_name] = func - return func + def decorator(obj: Any) -> Any: + cls._registry[self.name] = obj + return obj return decorator - def __call__(self, *args, **kwargs) -> Any: - """Call the registered function with the given arguments.""" - return self["__call__"](*args, **kwargs) - - def __getattr__(self, method_name: str) -> Callable: - """Get the registered function with the given name.""" - 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!") + @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] - @PrintPlugin("hello").register("again") - def print_hello_again(): - print("Hello world! Again.") + def __call__(self, *args, **kwargs) -> Any: + return self._resolve()(*args, **kwargs) - PrintPlugin("hello")() - PrintPlugin("hello").again() + def __getattr__(self, attr: str) -> Any: + return getattr(self._resolve(), attr) diff --git a/src/llamafactory/v1/utils/types.py b/src/llamafactory/v1/utils/types.py index 0cb57938b3..2942a69b2e 100644 --- a/src/llamafactory/v1/utils/types.py +++ b/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.""" diff --git a/tests_v1/config/test_args_parser.py b/tests_v1/config/test_args_parser.py index d60bbeb17a..82d3a7bc21 100644 --- a/tests_v1/config/test_args_parser.py +++ b/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 diff --git a/tests_v1/core/test_model_loader.py b/tests_v1/core/test_model_loader.py index 96dbdd2c13..e354a28e86 100644 --- a/tests_v1/core/test_model_loader.py +++ b/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="llamafactory/tiny-random-qwen3", kernel_config={"name": "auto", "include_kernels": "auto"} - ) + model_args = ModelArguments(model="llamafactory/tiny-random-qwen3", kernel_config={"name": "auto"}) model_engine = ModelEngine(model_args) # test enable apply kernel plugin if hasattr(torch, "npu"): diff --git a/tests_v1/plugins/model_plugins/test_kernel_plugin.py b/tests_v1/plugins/model_plugins/test_kernel_plugin.py index a4207fed2a..29822477f3 100644 --- a/tests_v1/plugins/model_plugins/test_kernel_plugin.py +++ b/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__ diff --git a/tests_v1/plugins/model_plugins/test_ulysses_cp.py b/tests_v1/plugins/model_plugins/test_ulysses_cp.py index 746c2a0c28..2f0c5dbdf1 100644 --- a/tests_v1/plugins/model_plugins/test_ulysses_cp.py +++ b/tests_v1/plugins/model_plugins/test_ulysses_cp.py @@ -18,6 +18,7 @@ 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 - dist_config = {"cp_mode": "ulysses", "cp_size": cp_size, "dp_size": dp_size} - DistributedInterface(dist_config) + training_args = TrainingArguments(cp_mode="ulysses", cp_size=cp_size, dp_size=dp_size) + DistributedInterface(training_args) # 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 = { diff --git a/tests_v1/trainers/test_dpo_loss_precision.py b/tests_v1/trainers/test_dpo_loss_precision.py index 3009406e81..96cbf100b5 100644 --- a/tests_v1/trainers/test_dpo_loss_precision.py +++ b/tests_v1/trainers/test_dpo_loss_precision.py @@ -28,6 +28,7 @@ # Mock helpers # ============================================================================== + def _make_mock_v1( pref_beta: float = 0.1, dpo_label_smoothing: float = 0.0, @@ -67,6 +68,7 @@ def _make_mock_v0_dpo(beta: float = 0.1, label_smoothing: float = 0.0) -> Simple # 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([-3.0]), torch.tensor([-3.0]), + torch.tensor([-1.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([-3.0]), torch.tensor([-3.0]), + torch.tensor([-10.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 + chosen_len:] = 2 + token_type_ids[:, prompt_len : prompt_len + chosen_len] = 1 + 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) diff --git a/tests_v1/trainers/test_fsdp2_dpo_trainer.py b/tests_v1/trainers/test_fsdp2_dpo_trainer.py index 62cefdb38b..893614b0fd 100644 --- a/tests_v1/trainers/test_fsdp2_dpo_trainer.py +++ b/tests_v1/trainers/test_fsdp2_dpo_trainer.py @@ -33,7 +33,6 @@ def test_fsdp2_dpo_trainer(tmp_path: Path): kernel_config: name: auto - include_kernels: auto quant_config: null diff --git a/tests_v1/trainers/test_fsdp2_sft_trainer.py b/tests_v1/trainers/test_fsdp2_sft_trainer.py index b5f1176ea9..256b8098d1 100644 --- a/tests_v1/trainers/test_fsdp2_sft_trainer.py +++ b/tests_v1/trainers/test_fsdp2_sft_trainer.py @@ -30,7 +30,6 @@ def test_fsdp2_sft_trainer(tmp_path: Path): kernel_config: name: auto - include_kernels: auto quant_config: null From 5735f9e64e6edc2601d8700d6d27fe6984f7d1aa Mon Sep 17 00:00:00 2001 From: jiaqiw09 Date: Sun, 12 Jul 2026 15:51:43 +0800 Subject: [PATCH 2/2] update examples --- .../train_full_fsdp2_batching_normal.yaml | 3 --- .../train_full_fsdp2_dynamic_batching.yaml | 2 -- .../train_full_fsdp2_dynamic_padding_free.yaml | 2 -- .../train_batching_strategy/train_full_fsdp2_padding_free.yaml | 2 -- examples/v1/train_freeze/train_freeze_sft.yaml | 2 -- examples/v1/train_full/train_full_deepspeed.yaml | 1 - examples/v1/train_full/train_full_fsdp2.yaml | 1 - examples/v1/train_full/train_full_liger_kernel.yaml | 2 -- examples/v1/train_lora/train_lora_dpo.yaml | 2 -- examples/v1/train_lora/train_lora_sft.yaml | 1 - examples/v1/train_lora/train_lora_sft_rank0.yaml | 1 - examples/v1/train_qlora/quantization.yaml | 1 - 12 files changed, 20 deletions(-) diff --git a/examples/v1/train_batching_strategy/train_full_fsdp2_batching_normal.yaml b/examples/v1/train_batching_strategy/train_full_fsdp2_batching_normal.yaml index 854c9a6639..1e77d61a71 100644 --- a/examples/v1/train_batching_strategy/train_full_fsdp2_batching_normal.yaml +++ b/examples/v1/train_batching_strategy/train_full_fsdp2_batching_normal.yaml @@ -1,9 +1,6 @@ model: Qwen/Qwen3-0.6B model_class: llm -template: qwen3_nothink - - kernel_config: name: auto diff --git a/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_batching.yaml b/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_batching.yaml index 98e46c9e92..18f79227d9 100644 --- a/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_batching.yaml +++ b/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_batching.yaml @@ -1,8 +1,6 @@ model: Qwen/Qwen3-0.6B model_class: llm -template: qwen3_nothink - kernel_config: name: auto diff --git a/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml b/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml index 8bd674459b..d13232f008 100644 --- a/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml +++ b/examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml @@ -1,8 +1,6 @@ model: Qwen/Qwen3-0.6B model_class: llm -template: qwen3_nothink - kernel_config: name: auto diff --git a/examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml b/examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml index 0d5577ae5c..c0c6c41635 100644 --- a/examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml +++ b/examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml @@ -1,8 +1,6 @@ model: Qwen/Qwen3-0.6B model_class: llm -template: qwen3_nothink - kernel_config: name: auto diff --git a/examples/v1/train_freeze/train_freeze_sft.yaml b/examples/v1/train_freeze/train_freeze_sft.yaml index 27bdc20838..924524dd95 100644 --- a/examples/v1/train_freeze/train_freeze_sft.yaml +++ b/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 @@ -24,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 diff --git a/examples/v1/train_full/train_full_deepspeed.yaml b/examples/v1/train_full/train_full_deepspeed.yaml index 482053513d..fe495056f4 100644 --- a/examples/v1/train_full/train_full_deepspeed.yaml +++ b/examples/v1/train_full/train_full_deepspeed.yaml @@ -1,7 +1,6 @@ model: Qwen/Qwen3-0.6B model_class: llm - kernel_config: name: auto diff --git a/examples/v1/train_full/train_full_fsdp2.yaml b/examples/v1/train_full/train_full_fsdp2.yaml index 98e0c29772..27fefd55c0 100644 --- a/examples/v1/train_full/train_full_fsdp2.yaml +++ b/examples/v1/train_full/train_full_fsdp2.yaml @@ -1,7 +1,6 @@ model: Qwen/Qwen3-0.6B model_class: llm - kernel_config: name: auto diff --git a/examples/v1/train_full/train_full_liger_kernel.yaml b/examples/v1/train_full/train_full_liger_kernel.yaml index 0725c6ab8b..f785c8a252 100644 --- a/examples/v1/train_full/train_full_liger_kernel.yaml +++ b/examples/v1/train_full/train_full_liger_kernel.yaml @@ -1,8 +1,6 @@ model: Qwen/Qwen3-0.6B model_class: llm -template: qwen3_nothink - kernel_config: name: liger_kernel diff --git a/examples/v1/train_lora/train_lora_dpo.yaml b/examples/v1/train_lora/train_lora_dpo.yaml index af892063e8..660bc29bc4 100644 --- a/examples/v1/train_lora/train_lora_dpo.yaml +++ b/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 diff --git a/examples/v1/train_lora/train_lora_sft.yaml b/examples/v1/train_lora/train_lora_sft.yaml index 9193125aea..0da1a1abd1 100644 --- a/examples/v1/train_lora/train_lora_sft.yaml +++ b/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 diff --git a/examples/v1/train_lora/train_lora_sft_rank0.yaml b/examples/v1/train_lora/train_lora_sft_rank0.yaml index fd60141b6b..ff401844e9 100644 --- a/examples/v1/train_lora/train_lora_sft_rank0.yaml +++ b/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 diff --git a/examples/v1/train_qlora/quantization.yaml b/examples/v1/train_qlora/quantization.yaml index 2ff0f581bf..79a126c292 100644 --- a/examples/v1/train_qlora/quantization.yaml +++ b/examples/v1/train_qlora/quantization.yaml @@ -1,7 +1,6 @@ model: Qwen/Qwen3-0.6B model_class: llm - # PEFT Configuration peft_config: name: lora