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 e14d52fd7..1e77d61a7 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,12 +1,8 @@ model: Qwen/Qwen3-0.6B model_class: llm -template: qwen3_nothink - - kernel_config: name: auto - include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null quant_config: null 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 ff0dc2cc2..18f79227d 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,11 +1,8 @@ model: Qwen/Qwen3-0.6B model_class: llm -template: qwen3_nothink - kernel_config: name: auto - include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null quant_config: null 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 9898a50d8..d13232f00 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,11 +1,8 @@ model: Qwen/Qwen3-0.6B model_class: llm -template: qwen3_nothink - kernel_config: name: auto - include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null quant_config: null 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 ec5babca0..c0c6c4163 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,11 +1,8 @@ model: Qwen/Qwen3-0.6B model_class: llm -template: qwen3_nothink - kernel_config: name: auto - include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null quant_config: null diff --git a/examples/v1/train_freeze/train_freeze_sft.yaml b/examples/v1/train_freeze/train_freeze_sft.yaml index a3e062ed9..924524dd9 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 @@ -12,7 +11,6 @@ peft_config: # Kernel Config kernel_config: name: auto - include_kernels: auto # FSDP Config dist_config: @@ -25,7 +23,6 @@ train_dataset: data/v1_sft_demo.yaml ### training output_dir: ./outputs/test_freeze micro_batch_size: 1 -global_batch_size: 4 cutoff_len: 2048 learning_rate: 2.0e-5 max_steps: 10 diff --git a/examples/v1/train_full/train_full_deepspeed.yaml b/examples/v1/train_full/train_full_deepspeed.yaml index 332c32717..fe495056f 100644 --- a/examples/v1/train_full/train_full_deepspeed.yaml +++ b/examples/v1/train_full/train_full_deepspeed.yaml @@ -1,10 +1,8 @@ model: Qwen/Qwen3-0.6B model_class: llm - kernel_config: name: auto - include_kernels: auto dist_config: name: deepspeed diff --git a/examples/v1/train_full/train_full_fsdp2.yaml b/examples/v1/train_full/train_full_fsdp2.yaml index c4796a7a2..27fefd55c 100644 --- a/examples/v1/train_full/train_full_fsdp2.yaml +++ b/examples/v1/train_full/train_full_fsdp2.yaml @@ -1,10 +1,8 @@ model: Qwen/Qwen3-0.6B model_class: llm - kernel_config: name: auto - include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null quant_config: null diff --git a/examples/v1/train_full/train_full_liger_kernel.yaml b/examples/v1/train_full/train_full_liger_kernel.yaml index b6face99c..f785c8a25 100644 --- a/examples/v1/train_full/train_full_liger_kernel.yaml +++ b/examples/v1/train_full/train_full_liger_kernel.yaml @@ -1,11 +1,8 @@ model: Qwen/Qwen3-0.6B model_class: llm -template: qwen3_nothink - kernel_config: name: liger_kernel - include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null quant_config: null diff --git a/examples/v1/train_full/train_full_ulysses_cp.yaml b/examples/v1/train_full/train_full_ulysses_cp.yaml index 6ed486a1e..1ee5b1439 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 32beec105..660bc29bc 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 @@ -14,7 +12,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 9596e6be5..0da1a1abd 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 @@ -13,7 +12,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 d9dc711a0..ff401844e 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 @@ -13,7 +12,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 b0dcc19d0..79a126c29 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 @@ -13,7 +12,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 a43885a55..530685551 100644 --- a/src/llamafactory/v1/accelerator/interface.py +++ b/src/llamafactory/v1/accelerator/interface.py @@ -29,16 +29,20 @@ And data parallelism types: from dataclasses import dataclass from datetime import timedelta from enum import StrEnum -from typing import Any, Optional +from typing import TYPE_CHECKING, Any, Optional from torch.distributed import barrier, destroy_process_group, init_process_group from torch.distributed.device_mesh import DeviceMesh, init_device_mesh from ..utils import logging -from ..utils.types import DistributedConfig, ProcessGroup, TensorLike +from ..utils.types import ProcessGroup, TensorLike from . import helper +if TYPE_CHECKING: + from ..config.training_args import TrainingArguments + + logger = logging.get_logger(__name__) @@ -128,12 +132,13 @@ class DistributedInterface: return cls._instance - def __init__(self, config: DistributedConfig | None = None) -> None: + def __init__( + self, + training_args: "TrainingArguments | None" = None, + ) -> None: if self._initialized: return - self.dist_config = config - helper.set_device_index() self._is_distributed = helper.is_distributed() self._rank = helper.get_rank() @@ -143,17 +148,17 @@ class DistributedInterface: self.current_device = helper.get_current_device() self.device_count = helper.get_device_count() - if config is None: + if training_args is None: self.strategy = DistributedStrategy() timeout = 18000 else: self.strategy = DistributedStrategy( - mp_replicate_size=config.get("mp_replicate_size", 1), - mp_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 8a9b1470d..f86f80a6d 100644 --- a/src/llamafactory/v1/config/training_args.py +++ b/src/llamafactory/v1/config/training_args.py @@ -76,7 +76,31 @@ class TrainingArguments: ) dist_config: PluginConfig | None = field( default=None, - metadata={"help": "Distribution configuration for training."}, + metadata={"help": "Distributed backend plugin configuration."}, + ) + dp_size: int | None = field( + default=None, + metadata={"help": "Data parallel size, default to world_size // cp_size."}, + ) + cp_size: int = field( + default=1, + metadata={"help": "Context parallel size."}, + ) + cp_mode: str = field( + default="ulysses", + metadata={"help": "Context parallel implementation."}, + ) + mp_replicate_size: int = field( + default=1, + metadata={"help": "Model parallel replicate size."}, + ) + mp_shard_size: int | None = field( + default=None, + metadata={"help": "Model parallel shard size, default to world_size // mp_replicate_size."}, + ) + dist_timeout: int = field( + default=18000, + metadata={"help": "Distributed process group initialization timeout in seconds."}, ) optim_config: PluginConfig | None = field( default=None, @@ -149,6 +173,12 @@ class TrainingArguments: self.dist_config = get_plugin_config(self.dist_config) self.optim_config = get_plugin_config(self.optim_config) self.lr_scheduler_config = get_plugin_config(self.lr_scheduler_config) + try: + from ..plugins.model_plugins.deepspeed_utils import register_deepspeed_dist_config + + register_deepspeed_dist_config(self.dist_config) + except ImportError: + pass # The optimizer learning rate has a single source of truth: ``learning_rate``. # Propagate it into ``optim_config["lr"]`` so optimizer plugins (e.g. Muon) pick it up diff --git a/src/llamafactory/v1/core/base_trainer.py b/src/llamafactory/v1/core/base_trainer.py index 77dad90c2..2dc998bc6 100644 --- a/src/llamafactory/v1/core/base_trainer.py +++ b/src/llamafactory/v1/core/base_trainer.py @@ -94,9 +94,12 @@ class BaseTrainer: dist_name = self.args.dist_config.name if self.args.dist_config is not None else None if dist_name == "deepspeed": - from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin + if self.args.cp_size > 1: + raise ValueError("Context parallelism currently requires `dist_config.name: fsdp2`.") - self._deepspeed_engine = DistributedPlugin("deepspeed")( + from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin + + self._deepspeed_engine = DistributedPlugin("deepspeed").shard_model( self.model, self.args.dist_config, num_micro_batch=self.train_batch_generator.num_micro_batch, @@ -139,7 +142,7 @@ class BaseTrainer: self.state.global_step = self.global_step self.state.epoch = self._resume_epoch - if self.args.dist_config is not None and self.args.dist_config.get("cp_size", 1) > 1: + if self.args.cp_size > 1: # qwen3.5 is not supported because of the different attention implementation, which will be supported in the future. if model.config.model_type == "qwen3_5": raise RuntimeError( @@ -152,7 +155,7 @@ class BaseTrainer: "Sequence parallelism requires flash attention. Please set `flash_attn: flash_attention_2`." ) - SequenceParallelModelPlugin(self.args.dist_config.get("cp_mode", "ulysses"))(model, self.args.dist_config) + SequenceParallelModelPlugin(self.args.cp_mode)(model, self.args.cp_size) def _create_batch_generator(self) -> None: if ( @@ -183,9 +186,9 @@ class BaseTrainer: device_ids = None if self.device.type == "cpu" else [self.device.index] self.model = DDP(self.model, device_ids=device_ids) else: - from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin + from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin - self.model = DistributedPlugin(self.args.dist_config.name)( + self.model = DistributedPlugin(self.args.dist_config.name).shard_model( self.model, self.args.dist_config, bf16=self.args.bf16, @@ -256,7 +259,7 @@ class BaseTrainer: step_valid_tokens = DistributedInterface().all_reduce(step_valid_tokens, op=ReduceOp.SUM) num_micro = len(micro_batches) for i, micro_batch in enumerate(micro_batches): - if self.args.dist_config and self.args.dist_config.get("cp_size", 1) > 1: + if self.args.cp_size > 1: from ..plugins.model_plugins.parallelization.sequence_parallel import ( SequenceParallelLossPlugin, ) @@ -352,7 +355,7 @@ class BaseTrainer: def save_model(self) -> None: """Save the model.""" if self.args.dist_config is not None and self.args.dist_config.name in ("deepspeed", "fsdp2"): - from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin + from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin DistributedPlugin(self.args.dist_config.name).save_model( self.model, self.args.output_dir, self.renderer.processor diff --git a/src/llamafactory/v1/core/data_engine.py b/src/llamafactory/v1/core/data_engine.py index ff8adaa45..810cb43ad 100644 --- a/src/llamafactory/v1/core/data_engine.py +++ b/src/llamafactory/v1/core/data_engine.py @@ -137,9 +137,7 @@ class DataEngine(Dataset): 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 4df8d060e..209cc494b 100644 --- a/src/llamafactory/v1/core/model_engine.py +++ b/src/llamafactory/v1/core/model_engine.py @@ -69,24 +69,25 @@ class ModelEngine: """Model configuration.""" self.renderer = Renderer(self.processor) """Renderer.""" - self._dist_config = DistributedInterface().dist_config - self._deepspeed_zero3_plugin = None self._deepspeed_zero3_enabled = False - if self.is_train and self._dist_config is not None and self._dist_config.get("name") == "deepspeed": + try: from ..plugins.model_plugins.deepspeed_utils import ( + is_deepspeed_zero3_enabled, setup_deepspeed_zero3_model_loading, teardown_deepspeed_zero3_model_loading, ) + self._deepspeed_zero3_enabled = self.is_train and is_deepspeed_zero3_enabled() + except ImportError: + pass + + if self._deepspeed_zero3_enabled: + plugin = setup_deepspeed_zero3_model_loading() try: - self._deepspeed_zero3_plugin = setup_deepspeed_zero3_model_loading(self.is_train, self._dist_config) - self._deepspeed_zero3_enabled = self._deepspeed_zero3_plugin is not None self.model = self._init_model() finally: - teardown_deepspeed_zero3_model_loading(self._deepspeed_zero3_plugin) - 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 @@ class ModelEngine: 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 @@ class ModelEngine: from ..plugins.model_plugins.peft import PeftPlugin - model = PeftPlugin(self.args.peft_config.name)(model, self.args.peft_config, self.is_train) + model = PeftPlugin(self.args.peft_config.name)( + model, + peft_config=self.args.peft_config, + is_train=self.is_train, + ) if self.args.kernel_config is not None: - from ..plugins.model_plugins.kernels.interface import KernelPlugin + from ..plugins.model_plugins.kernels.interface import apply_kernels - kernel_config = self.args.kernel_config - 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) + 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 293c3b1b3..048e61217 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 4dc5e2d68..87a62e87f 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 @@ class Renderer: 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 40ed4f01c..c0f18f44a 100644 --- a/src/llamafactory/v1/core/utils/checkpoint.py +++ b/src/llamafactory/v1/core/utils/checkpoint.py @@ -251,7 +251,7 @@ class TrainingCheckpointCoordinator: ) if self._dist_name in ("fsdp2", "deepspeed"): - from ...plugins.trainer_plugins.distributed.hub import DistributedPlugin + from ...plugins.trainer_plugins.distributed.interface import DistributedPlugin DistributedPlugin(self._dist_name).save_checkpoint( self._t.model, @@ -307,7 +307,7 @@ class TrainingCheckpointCoordinator: self._t._resume_epoch = metadata["epoch"] if self._dist_name in ("fsdp2", "deepspeed"): - from ...plugins.trainer_plugins.distributed.hub import DistributedPlugin + from ...plugins.trainer_plugins.distributed.interface import DistributedPlugin DistributedPlugin(self._dist_name).load_checkpoint( self._t.model, diff --git a/src/llamafactory/v1/plugins/model_plugins/deepspeed_utils.py b/src/llamafactory/v1/plugins/model_plugins/deepspeed_utils.py index de737ed6f..d31a2dad7 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 3916562b1..2174f0de4 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/base.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/base.py @@ -12,76 +12,40 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""The definition of base kernel class. - -Init Phase: -1. Define base kernel class. -2. Define abstract methods. - -""" - from abc import ABC, abstractmethod -from typing import Any -from ....accelerator.helper import DeviceType, get_current_accelerator +from ....utils.plugin import BasePlugin, ensure_methods_implemented from ....utils.types import HFModel +class KernelPlugin(BasePlugin): + """Plugin family for model kernel optimization classes.""" + + class BaseKernel(ABC): - r"""Base class for all kernel implementations. + """Template base for concrete kernel implementations.""" - Subclasses must implement the abstract methods and define the required class attributes. - """ + def __init_subclass__(cls, **kwargs) -> None: + super().__init_subclass__(**kwargs) + ensure_methods_implemented(cls) - _kernel_id: Any = "" # kernel ID, any hashable value to identify a kernel implementation - _device: list[DeviceType] = [DeviceType.CPU] # "cuda", "npu", "cpu", etc. - - @classmethod - def get_kernel_id(cls) -> str: - """Returns the unique identifier for the kernel.""" - return cls._kernel_id - - @classmethod - def get_device(cls) -> list[DeviceType]: - """Returns the device type list associated with the kernel (e.g., ["cuda", "npu", "cpu"]).""" - return cls._device - - @classmethod - def check_deps(cls) -> bool: - """Checks if the required dependencies for the kernel are available. - - Returns: - bool: ``True`` if dependencies are met, ``False`` otherwise. - - .. note:: - In explicit mode, if a user specifies an implementation but this check fails, - it should raise an error instead of silently switching. - Kernels can override this method to implement custom dependency checks. - """ - if get_current_accelerator().type not in cls._device: - return False - return True - - @classmethod + @staticmethod @abstractmethod + def check_device() -> None: ... + + @staticmethod + def check_deps() -> None: + pass + + @classmethod def apply(cls, **kwargs) -> HFModel: - """Applies the kernel optimization to the model. + cls.check_device() + cls.check_deps() + if kwargs.get("model") is None: + raise ValueError(f"HFModel instance is required for {cls.__name__}.") - Args: - **kwargs: Arbitrary keyword arguments, usually containing the model instance and the kernel configuration. + return cls._apply(**kwargs) - 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 + @staticmethod + @abstractmethod + def _apply(**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 1646448e9..beba83e6b 100644 --- a/src/llamafactory/v1/plugins/model_plugins/kernels/interface.py +++ b/src/llamafactory/v1/plugins/model_plugins/kernels/interface.py @@ -12,174 +12,63 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""The definition of kernel interface. +from typing import Any -Init Phase: -1. Scan all kernels. -2. Register default kernels. -3. Define kernel plugin. - -""" - -import importlib -from pathlib import Path - -from ....utils import logging -from ....utils.plugin import BasePlugin +from ....accelerator.helper import DeviceType, get_current_accelerator from ....utils.types import HFModel -from .registry import Registry +from .base import KernelPlugin + +# Import built-in implementations so their class decorators populate the registry. +from .liger_kernel_ops import LigerKernel # noqa: F401 +from .ops.mlp.cuda_fused_moe import CudaFusedMoEKernel # noqa: F401 +from .ops.mlp.npu_fused_moe import NpuFusedMoEKernel # noqa: F401 +from .ops.mlp.npu_swiglu import NpuSwiGluKernel # noqa: F401 +from .ops.rms_norm.npu_rms_norm import NpuRMSNormKernel # noqa: F401 +from .ops.rope.npu_rope import NpuRoPEKernel # noqa: F401 -logger = logging.get_logger(__name__) +_AUTO_KERNELS = { + DeviceType.NPU: ("npu_fused_moe", "npu_fused_rmsnorm", "npu_fused_rope", "npu_fused_swiglu"), +} -def scan_all_kernels(): - """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" - - if not ops_path.exists(): - return - - base_package = __package__ - - for file_path in ops_path.rglob("*.py"): - if file_path.name == "__init__.py": - continue - - # calculate the relative path: - # file_path = .../kernels_v2/ops/mlp/npu_swiglu.py - # rel_path = ops/mlp/npu_swiglu.py - rel_path = file_path.relative_to(Path(__file__).parent) - - # build module path: - module_name = ".".join(rel_path.parts)[:-3] - full_module_name = f"{base_package}.{module_name}" - - try: - importlib.import_module(full_module_name) - except Exception as e: - logger.warning(f"[Kernel Registry] Failed to import {full_module_name} when loading kernels: {e}") - - return Registry.get_registered_kernels() - - -default_kernels = scan_all_kernels() - - -def get_default_kernels(): - """Get a list of default registered kernel IDs. - - Returns: - list[str]: List of kernel IDs. - """ - return list(default_kernels.keys()) - - -def apply_kernel(kernel_id: str, **kwargs): - """Applies a specific kernel to the model. - - Args: - kernel_id (str): The ID of the kernel to apply. - **kwargs: Keyword arguments passed to the kernel application function. - Typically includes the model instance. - - Returns: - HFModel: The model with applied kernel. - """ - kernel = default_kernels.get(kernel_id) - if kernel is None: - raise ValueError(f"Kernel {kernel_id} not found") - - kernel.apply(**kwargs) - - -class KernelPlugin(BasePlugin): - """Plugin for managing kernel optimizations.""" - - pass - - -@KernelPlugin("auto").register() -def apply_default_kernels(model: HFModel, include_kernels: str = None) -> HFModel: - """Applies all default registered kernels to the model. - - Args: - model (HFModel): The model instance to apply kernels to. - include_kernels (str, optional): Comma-separated list of kernel IDs to apply. - If "auto" or True, applies all default kernels. - If None or False, no kernels are applied. - Defaults to None. - - Returns: - HFModel: The model with applied kernels. - """ - if not include_kernels: - return model - elif include_kernels == "auto" or include_kernels is True: - use_kernels = default_kernels.keys() - else: - use_kernels = include_kernels.split(",") # "kernel_id1,kernel_id2,kernel_id3" - - for kernel in use_kernels: - if kernel not in default_kernels: - raise ValueError(f"Kernel {kernel} not found") - - apply_kernel(kernel, model=model) +def _apply_auto_kernels(model: HFModel, **kwargs) -> HFModel: + device_type = get_current_accelerator().type + for kernel_name in _AUTO_KERNELS.get(device_type, ()): + model = KernelPlugin(kernel_name).apply(model=model, **kwargs) 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. +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.") - 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. + 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.") - Returns: - HFModel: The model with Liger kernel applied. - """ - if not include_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 + for name in names: + if name == "auto": + model = _apply_auto_kernels(model=model, config=config, require_logits=require_logits) + else: + model = KernelPlugin(name).apply(model=model, config=config, require_logits=require_logits) - 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 + + +def apply_v1_kernels(model: HFModel, use_v1_kernels: bool) -> HFModel: + """Apply v1 automatic kernels for the transitional v0 ``use_v1_kernels`` option.""" + if not use_v1_kernels: return model - return LigerKernel.apply(use_kernels=use_kernels, model=model, require_logits=require_logits) + return apply_kernels(model, {"name": "auto"}) + + +def apply_kernel(kernel_id: str, **kwargs) -> HFModel: + if kernel_id == "auto": + return _apply_auto_kernels(**kwargs) + + return KernelPlugin(kernel_id).apply(**kwargs) 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 1d6fd0dbd..744735415 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 @@ -25,7 +25,7 @@ import inspect from ....accelerator.helper import DeviceType, get_current_accelerator from ....utils.logging import get_logger from ....utils.types import HFModel -from .base import BaseKernel +from .base import BaseKernel, KernelPlugin logger = get_logger(__name__) @@ -41,26 +41,26 @@ _LIGER_FN_BY_MODEL_TYPE: dict[str, str] = { } +@KernelPlugin("liger_kernel").register() class LigerKernel(BaseKernel): """Liger Kernel for optimized model training.""" - _device = [DeviceType.CUDA, DeviceType.NPU] + @staticmethod + def check_device() -> None: + current = get_current_accelerator().type + if current not in (DeviceType.CUDA, DeviceType.NPU): + raise RuntimeError(f"LigerKernel requires CUDA or NPU, current accelerator is {current}.") - @classmethod - def check_deps(cls) -> bool: + @staticmethod + def check_deps() -> None: """Checks if the required dependencies for the kernel are available.""" try: import liger_kernel # noqa: F401 - - return super().check_deps() except ImportError: - logger.warning_rank0( - "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": + @staticmethod + def _apply(**kwargs) -> "HFModel": """Applies the Liger kernel to the model. Args: @@ -78,16 +78,12 @@ class LigerKernel(BaseKernel): RuntimeError: If dependencies are not met. """ model = kwargs.get("model") - use_kernels = kwargs.get("use_kernels", None) - 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()}" - ) + config = kwargs.get("config") + use_kernels = kwargs.get("use_kernels", "auto") require_logits = kwargs.get("require_logits", False) + if config is not None: + require_logits = config.get("require_logits", require_logits) model_type = getattr(model.config, "model_type", None) 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 1a2cc4b69..cfc75fd02 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 @@ -27,16 +27,22 @@ import types import torch import torch.nn.functional as F -from ......accelerator.helper import DeviceType +from ......accelerator.helper import DeviceType, get_current_accelerator from ......utils.types import HFModel -from ...base import BaseKernel -from ...registry import register_kernel -from .triton_grouped_gemm import ( - group_gemm_same_mn, - group_gemm_same_nk, - moe_gather, - moe_scatter, -) +from ...base import BaseKernel, KernelPlugin + + +try: + from .triton_grouped_gemm import ( + group_gemm_same_mn, + group_gemm_same_nk, + moe_gather, + moe_scatter, + ) +except ImportError as exc: + _TRITON_IMPORT_ERROR = exc +else: + _TRITON_IMPORT_ERROR = None logger = logging.getLogger(__name__) @@ -351,7 +357,7 @@ _TRITON_MOE_MAPPING: dict[str, dict[str, object]] = { # --------------------------------------------------------------------------- -@register_kernel +@KernelPlugin("cuda_fused_moe").register() class CudaFusedMoEKernel(BaseKernel): """Pure-Triton fused MoE kernel for NVIDIA CUDA GPUs. @@ -362,30 +368,20 @@ class CudaFusedMoEKernel(BaseKernel): Requires: CUDA GPU + Triton """ - _kernel_id = "cuda_fused_moe" - _device = DeviceType.CUDA + @staticmethod + def check_device() -> None: + current = get_current_accelerator().type + if current != DeviceType.CUDA: + raise RuntimeError(f"CudaFusedMoEKernel requires CUDA, current accelerator is {current}.") - @classmethod - def check_deps(cls) -> bool: - if not super().check_deps(): - return False - try: - import triton # noqa: F401 + @staticmethod + def check_deps() -> None: + if _TRITON_IMPORT_ERROR is not None: + raise RuntimeError("cuda_fused_moe requires Triton.") from _TRITON_IMPORT_ERROR - return True - except ImportError: - logger.info("cuda_fused_moe: Triton not available, kernel disabled.") - return False - - @classmethod - def apply(cls, **kwargs) -> HFModel: + @staticmethod + def _apply(**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 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 41cab6020..20c86a2dc 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 @@ -32,11 +32,10 @@ try: except ImportError: pass -from ......accelerator.helper import DeviceType +from ......accelerator.helper import DeviceType, get_current_accelerator from ......utils.packages import is_transformers_version_greater_than from ......utils.types import HFModel -from ...base import BaseKernel -from ...registry import register_kernel +from ...base import BaseKernel, KernelPlugin class GmmFunction(torch.autograd.Function): @@ -334,15 +333,18 @@ else: } -@register_kernel +@KernelPlugin("npu_fused_moe").register() class NpuFusedMoEKernel(BaseKernel): """NPU Fused MoE Kernel implementation.""" - _kernel_id = "npu_fused_moe" - _device = DeviceType.NPU + @staticmethod + def check_device() -> None: + current = get_current_accelerator().type + if current != DeviceType.NPU: + raise RuntimeError(f"NpuFusedMoEKernel requires NPU, current accelerator is {current}.") - @classmethod - def apply(cls, **kwargs) -> HFModel: + @staticmethod + def _apply(**kwargs) -> HFModel: """Applies the NPU fused MoE kernel to the model. Args: @@ -356,11 +358,6 @@ class NpuFusedMoEKernel(BaseKernel): RuntimeError: If dependencies are not met. """ model = kwargs.get("model", None) - if model is None: - raise ValueError(f"HFModel instance is required for {cls.__name__}.") - - if not cls.check_deps(): - raise RuntimeError("torch_npu is not available but NpuMoEFusedMoEKernel was called.") archs = getattr(model.config, "architectures", None) or [] target_moe_mapping = None 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 a45077bc0..e48de755b 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 @@ -25,10 +25,9 @@ import types import torch -from ......accelerator.helper import DeviceType +from ......accelerator.helper import DeviceType, get_current_accelerator from ......utils.types import HFModel -from ...base import BaseKernel -from ...registry import register_kernel +from ...base import BaseKernel, KernelPlugin try: @@ -86,7 +85,7 @@ def _npu_swiglu_gemma3ntext_forward(self, hidden_states): return down_proj -@register_kernel +@KernelPlugin("npu_fused_swiglu").register() class NpuSwiGluKernel(BaseKernel): """NPU Kernel for fused SwiGLU activation.""" @@ -121,11 +120,14 @@ class NpuSwiGluKernel(BaseKernel): } ) - _kernel_id = "npu_fused_swiglu" - _device = DeviceType.NPU + @staticmethod + def check_device() -> None: + current = get_current_accelerator().type + if current != DeviceType.NPU: + raise RuntimeError(f"NpuSwiGluKernel requires NPU, current accelerator is {current}.") - @classmethod - def apply(cls, **kwargs) -> "HFModel": + @staticmethod + def _apply(**kwargs) -> "HFModel": """Applies the NPU fused SwiGLU kernel to the model. Args: @@ -139,11 +141,6 @@ class NpuSwiGluKernel(BaseKernel): RuntimeError: If dependencies are not met. """ model = kwargs.get("model", None) - if model is None: - raise ValueError(f"HFModel instance is required for {cls.__name__}.") - - if not cls.check_deps(): - raise RuntimeError("torch_npu is not available but NpuSwiGluKernel was called.") # Mapping of specific mlp modules to their corresponding kernel implementations kernel_mapping = { @@ -158,7 +155,7 @@ class NpuSwiGluKernel(BaseKernel): # Match any module whose class name contains "MLP" if ( re.search(swiglu_pattern, module.__class__.__name__) - and module.__class__.__name__ in cls.expect_modules + and module.__class__.__name__ in NpuSwiGluKernel.expect_modules ): # Bind function as an instance method to preserve `self` semantics # and replace the original forward 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 04a4ad16f..3d06abc34 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 @@ -26,10 +26,9 @@ import types import torch import torch.nn.functional as F -from ......accelerator.helper import DeviceType +from ......accelerator.helper import DeviceType, get_current_accelerator from ......utils.types import HFModel -from ...base import BaseKernel -from ...registry import register_kernel +from ...base import BaseKernel, KernelPlugin try: @@ -118,15 +117,18 @@ def npu_gated_rms_norm_forward(self, hidden_states, gate=None): return hidden_states.to(input_dtype) -@register_kernel +@KernelPlugin("npu_fused_rmsnorm").register() class NpuRMSNormKernel(BaseKernel): """NPU kernel wrapper for RMSNorm that applies the replacement within a model.""" - _kernel_id = "npu_fused_rmsnorm" - _device = DeviceType.NPU + @staticmethod + def check_device() -> None: + current = get_current_accelerator().type + if current != DeviceType.NPU: + raise RuntimeError(f"NpuRMSNormKernel requires NPU, current accelerator is {current}.") - @classmethod - def apply(cls, **kwargs) -> "HFModel": + @staticmethod + def _apply(**kwargs) -> "HFModel": """Iterate the model and apply NPU-optimized forward to matched RMSNorm modules. Matches modules whose class name contains "RMSNorm" (case-insensitive) and binds @@ -144,11 +146,6 @@ class NpuRMSNormKernel(BaseKernel): ValueError: If the model is not provided. """ model = kwargs.get("model") - if model is None: - raise ValueError(f"HFModel instance is required for {cls.__name__}.") - - if not cls.check_deps(): - raise RuntimeError(f"torch_npu is not available but {cls.__name__} was called.") rms_norm_pattern = re.compile("RMSNorm", re.IGNORECASE) 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 0f9d4d7e2..3ee39e1d7 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 @@ -24,11 +24,10 @@ import sys import torch -from ......accelerator.helper import DeviceType +from ......accelerator.helper import DeviceType, get_current_accelerator from ......utils.logging import get_logger from ......utils.types import HFModel -from ...base import BaseKernel -from ...registry import register_kernel +from ...base import BaseKernel, KernelPlugin logger = get_logger(__name__) @@ -125,15 +124,18 @@ def _apply_multimodal_rotary_pos_emb_qwen25_vl(q, k, cos, sin, mrope_section, un return _apply_npu_rotary_emb(q, k, cos, sin) -@register_kernel +@KernelPlugin("npu_fused_rope").register() class NpuRoPEKernel(BaseKernel): """NPU Kernel for Rotary Position Embedding.""" - _kernel_id = "npu_fused_rope" - _device = DeviceType.NPU + @staticmethod + def check_device() -> None: + current = get_current_accelerator().type + if current != DeviceType.NPU: + raise RuntimeError(f"NpuRoPEKernel requires NPU, current accelerator is {current}.") - @classmethod - def apply(cls, **kwargs) -> "HFModel": + @staticmethod + def _apply(**kwargs) -> "HFModel": """Apply RoPE acceleration by monkey-patching ``apply_rotary_pos_emb``. Iterates through the model's modules to find attention layers, identifies @@ -151,12 +153,7 @@ class NpuRoPEKernel(BaseKernel): RuntimeError: If ``torch_npu`` is not available. ValueError: If the model is not provided. """ - if not cls.check_deps(): - raise RuntimeError(f"torch_npu is not available but {cls.__name__} was called.") - model = kwargs.get("model", None) - if model is None: - raise ValueError(f"HFModel instance is required for {cls.__name__}.") _modules = set() for module in model.modules(): 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 2921f0aed..000000000 --- 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 095379ae4..47a6a1d99 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 @@ logger = logging.get_logger(__name__) 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 9552de0b4..9030a6363 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 @@ from ...utils.types import HFModel 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.""" + 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.""" -class FreezeConfigDict(TypedDict, total=False): - name: Literal["freeze"] +@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 2ba74d47d..b889b5cf7 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. + quant_config = QuantizationPlugin.parse_params(quant_config, BnbParams) + if quant_config.quantization_bit is None: + logger.warning_rank0("No quantization method applied.") + return init_kwargs + if quant_config.quantization_bit not in (4, 8): + raise ValueError(f"Unsupported quantization bit: {quant_config.quantization_bit} for auto quantization.") - Args: - 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 + 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 715f4cc5c..4cbed315c 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 torch.utils.data import default_collate from ...utils.constants import IGNORE_INDEX from ...utils.helper import pad_and_truncate from ...utils.objects import StatefulBuffer -from ...utils.plugin import BasePlugin +from ...utils.plugin import BasePlugin, ensure_methods_implemented from ...utils.types import BatchInfo, BatchInput, DataLoader, ModelInput class BatchingPlugin(BasePlugin): - def get_data_provider_batch_size(self, batch_info: BatchInfo) -> int: - """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) + @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 + buffer.put(samples) -@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) - - -@BatchingPlugin("padding_free").register("generate_batch") -def generate_padding_free_batch(buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None: - micro_batch_size = batch_info["micro_batch_size"] - num_micro_batch = batch_info["num_micro_batch"] - cutoff_len = batch_info["cutoff_len"] - batch_size = micro_batch_size * num_micro_batch - if len(buffer) < batch_size: - return None - - samples = buffer.get(batch_size) - batch = [] - for i in range(num_micro_batch): - micro_batch = samples[i * micro_batch_size : (i + 1) * micro_batch_size] - packed_micro_batch = _pack_padding_free_samples(micro_batch, cutoff_len) - if packed_micro_batch is None: + @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) + samples = buffer.get(batch_size) + batch = [] + for i in range(num_micro_batch): + micro_batch = samples[i * micro_batch_size : (i + 1) * micro_batch_size] + packed_micro_batch = _pack_padding_free_samples(micro_batch, cutoff_len) + if packed_micro_batch is None: + return None - return batch + batch.append(packed_micro_batch) + + return batch -@BatchingPlugin("dynamic_batching").register("get_data_provider_batch_size") -def get_dynamic_batching_data_provider_batch_size(batch_info: BatchInfo) -> int: - return 1 +@BatchingPlugin("dynamic_batching").register() +class DynamicBatcher(BaseBatcher): + @staticmethod + def get_data_provider_batch_size(batch_info: BatchInfo) -> int: + return 1 + @staticmethod + def compute_length(data_provider: DataLoader, batch_info: BatchInfo) -> int: + batch_size = batch_info["micro_batch_size"] * batch_info["num_micro_batch"] + return ceil(len(data_provider) / batch_size) -@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) + @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 + buffer.put(samples) -@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 - - buffer.put(samples) - - -@BatchingPlugin("dynamic_batching").register("generate_batch") -def generate_dynamic_batching_batch(buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None: - micro_batch_sample_counts = _get_dynamic_micro_batch_sizes(buffer.samples, batch_info) - if len(micro_batch_sample_counts) < batch_info["num_micro_batch"]: - return None - - batch = [] - cutoff_len = batch_info["cutoff_len"] - for num_samples in micro_batch_sample_counts: - samples = buffer.get(num_samples) - batch.append(default_collate(pad_and_truncate(samples, cutoff_len))) - - return batch - - -@BatchingPlugin("dynamic_padding_free").register("get_data_provider_batch_size") -def get_dynamic_padding_free_data_provider_batch_size(batch_info: BatchInfo) -> int: - return 1 - - -@BatchingPlugin("dynamic_padding_free").register("compute_length") -def compute_dynamic_padding_free_length(data_provider: DataLoader, batch_info: BatchInfo) -> int: - batch_size = batch_info["micro_batch_size"] * batch_info["num_micro_batch"] - return ceil(len(data_provider) / batch_size) - - -@BatchingPlugin("dynamic_padding_free").register("fill_buffer") -def fill_dynamic_padding_free_buffer( - buffer: StatefulBuffer, - batch_info: BatchInfo, - next_samples: Callable[[bool], list[ModelInput] | None], -) -> None: - while len(_get_dynamic_padding_free_micro_batch_sizes(buffer.samples, batch_info)) < batch_info["num_micro_batch"]: - samples = next_samples(True) - if samples is None: - break - buffer.put(samples) - - -@BatchingPlugin("dynamic_padding_free").register("generate_batch") -def generate_dynamic_padding_free_batch(buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None: - micro_batch_sample_counts = _get_dynamic_padding_free_micro_batch_sizes(buffer.samples, batch_info) - if len(micro_batch_sample_counts) < batch_info["num_micro_batch"]: - return None - - batch = [] - cutoff_len = batch_info["cutoff_len"] - - for num_samples in micro_batch_sample_counts: - samples = buffer.get(num_samples) - packed_batch = _pack_padding_free_samples(samples, cutoff_len) - if packed_batch is None: + @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 - batch.append(packed_batch) + batch = [] + cutoff_len = batch_info["cutoff_len"] + for num_samples in micro_batch_sample_counts: + samples = buffer.get(num_samples) + batch.append(default_collate(pad_and_truncate(samples, cutoff_len))) - return batch + return batch + + +@BatchingPlugin("dynamic_padding_free").register() +class DynamicPaddingFreeBatcher(BaseBatcher): + @staticmethod + def get_data_provider_batch_size(batch_info: BatchInfo) -> int: + return 1 + + @staticmethod + def compute_length(data_provider: DataLoader, batch_info: BatchInfo) -> int: + batch_size = batch_info["micro_batch_size"] * batch_info["num_micro_batch"] + return ceil(len(data_provider) / batch_size) + + @staticmethod + def fill_buffer( + buffer: StatefulBuffer, + batch_info: BatchInfo, + next_samples: Callable[[bool], list[ModelInput] | None], + ) -> None: + while ( + len(_get_dynamic_padding_free_micro_batch_sizes(buffer.samples, batch_info)) + < batch_info["num_micro_batch"] + ): + samples = next_samples(True) + if samples is None: + break + buffer.put(samples) + + @staticmethod + def generate_batch(buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None: + micro_batch_sample_counts = _get_dynamic_padding_free_micro_batch_sizes(buffer.samples, batch_info) + if len(micro_batch_sample_counts) < batch_info["num_micro_batch"]: + return None + + batch = [] + cutoff_len = batch_info["cutoff_len"] + + for num_samples in micro_batch_sample_counts: + samples = buffer.get(num_samples) + packed_batch = _pack_padding_free_samples(samples, cutoff_len) + if packed_batch is None: + return None + + batch.append(packed_batch) + + return batch 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 000000000..2b05f76eb --- /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 b4d0babf0..000000000 --- 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 000000000..cc4ae06ee --- /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 d8f289b34..fb2dfb59d 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 @@ class DPOTrainer(BaseTrainer): train_dataset, callbacks=None, ) -> None: - cp_size = args.dist_config.get("cp_size", 1) if args.dist_config is not None else 1 - if 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 @@ class DPOTrainer(BaseTrainer): # Raw logits means (for logging) chosen_logits_mean = (shift_logits.mean(dim=-1) * chosen_logit_mask).sum() / (chosen_logit_mask.sum() + 1e-6) - rejected_logits_mean = (shift_logits.mean(dim=-1) * rejected_logit_mask).sum() / (rejected_logit_mask.sum() + 1e-6) + rejected_logits_mean = (shift_logits.mean(dim=-1) * rejected_logit_mask).sum() / ( + rejected_logit_mask.sum() + 1e-6 + ) if self.pref_loss == "sigmoid": if not self._use_lora_ref and self.ref_model is None: @@ -431,7 +429,7 @@ def run_dpo(args: InputArgument = None): model_args, data_args, training_args, _ = get_args(args) if getattr(training_args, "use_cpu", False): os.environ["FORCE_V1_CPU"] = "1" - DistributedInterface(training_args.dist_config) + DistributedInterface(training_args) train_dataset = DataEngine(data_args.train_dataset) _validate_dpo_dataset_format(train_dataset, data_args.train_dataset) model_engine = ModelEngine(model_args, is_train=True) diff --git a/src/llamafactory/v1/trainers/rm_trainer.py b/src/llamafactory/v1/trainers/rm_trainer.py index 6c78fa6fa..78694d47b 100644 --- a/src/llamafactory/v1/trainers/rm_trainer.py +++ b/src/llamafactory/v1/trainers/rm_trainer.py @@ -76,8 +76,7 @@ class RMTrainer(BaseTrainer): train_dataset, callbacks=None, ) -> None: - cp_size = args.dist_config.get("cp_size", 1) if args.dist_config is not None else 1 - if 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 @@ class RMTrainer(BaseTrainer): def run_rm(args: InputArgument = None): model_args, data_args, training_args, _ = get_args(args) model_args.model_class = ModelClass.CLS - DistributedInterface(training_args.dist_config) + DistributedInterface(training_args) train_dataset = DataEngine(data_args.train_dataset) _validate_rm_dataset_format(train_dataset, data_args.train_dataset) model_engine = ModelEngine(model_args, is_train=True) diff --git a/src/llamafactory/v1/trainers/sft_trainer.py b/src/llamafactory/v1/trainers/sft_trainer.py index 5b266c611..d84263283 100644 --- a/src/llamafactory/v1/trainers/sft_trainer.py +++ b/src/llamafactory/v1/trainers/sft_trainer.py @@ -31,7 +31,7 @@ class SFTTrainer(BaseTrainer): def run_sft(args: InputArgument = None): model_args, data_args, training_args, _ = get_args(args) - DistributedInterface(training_args.dist_config) + DistributedInterface(training_args) train_dataset = DataEngine(data_args.train_dataset) model_engine = ModelEngine(model_args, is_train=True) trainer = SFTTrainer( diff --git a/src/llamafactory/v1/utils/plugin.py b/src/llamafactory/v1/utils/plugin.py index 4c06c88d3..8debd5e7f 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") + + +def ensure_methods_implemented(cls: type) -> None: + """Raise when a static method-group implementation is incomplete.""" + required: set[str] = set() + for base in cls.__mro__[1:]: + required |= getattr(base, "__abstractmethods__", frozenset()) + + missing = sorted(name for name in required if getattr(getattr(cls, name, None), "__isabstractmethod__", False)) + if missing: + raise TypeError(f"{cls.__name__} does not implement all required methods: {missing}") class BasePlugin: - """Base class for plugins. + """Route a plugin name to one function or static method-group class. - A plugin is a callable object that can be registered and called by name. - - Example usage: - ```python - class PrintPlugin(BasePlugin): - def again(self): # optional - self["again"]() - - - @PrintPlugin("hello").register() - def print_hello(): - print("Hello world!") - - - @PrintPlugin("hello").register("again") - def print_hello_again(): - print("Hello world! Again.") - - - PrintPlugin("hello")() - PrintPlugin("hello").again() - ``` + 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 + @classmethod + def parse_params(cls, config: Any, params_cls: type[ParamsT]) -> ParamsT: + """Strictly convert config to the params dataclass used by one plugin entrypoint.""" + if not is_dataclass(params_cls): + raise TypeError(f"{cls.__name__} params must be a dataclass type, got {params_cls!r}.") + if isinstance(config, params_cls): + return config + if config is None: + values = {} + elif isinstance(config, dict): + values = dict(config) + else: + raise TypeError( + f"{cls.__name__} config must be a mapping or {params_cls.__name__}, got {type(config).__name__}." + ) + + known = {item.name for item in fields(params_cls)} + unknown = set(values) - known + if unknown: + raise ValueError( + f"Unknown params for {cls.__name__}.{params_cls.__name__}: {sorted(unknown)}. " + f"Expected: {sorted(known)}" + ) + + return params_cls(**values) + + def _resolve(self) -> Any: + cls = type(self) + if self.name is None: + raise ValueError(f"{cls.__name__} must be constructed with a name.") + if self.name not in cls._registry: + raise ValueError(f"Plugin {self.name!r} is not registered under {cls.__name__}.") + return cls._registry[self.name] + def __call__(self, *args, **kwargs) -> Any: - """Call the registered function with the given arguments.""" - return self["__call__"](*args, **kwargs) + return self._resolve()(*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!") - - @PrintPlugin("hello").register("again") - def print_hello_again(): - print("Hello world! Again.") - - 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 0cb57938b..2942a69b2 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 d60bbeb17..82d3a7bc2 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 96dbdd2c1..e354a28e8 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 a4207fed2..29822477f 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 746c2a0c2..2f0c5dbdf 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 @@ import torch.multiprocessing as mp from llamafactory.v1.accelerator.interface import DistributedInterface from llamafactory.v1.config.model_args import ModelArguments +from llamafactory.v1.config.training_args import TrainingArguments from llamafactory.v1.core.model_engine import ModelEngine from llamafactory.v1.plugins.model_plugins.parallelization.sequence_parallel import ( SequenceParallelModelPlugin, @@ -33,15 +34,14 @@ def _test_sequence_parallel_loss( with dist_env(local_rank, world_size, master_port): model_args = ModelArguments(model="llamafactory/tiny-random-qwen3") - # Initialize distributed interface with config - 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 3009406e8..96cbf100b 100644 --- a/tests_v1/trainers/test_dpo_loss_precision.py +++ b/tests_v1/trainers/test_dpo_loss_precision.py @@ -28,6 +28,7 @@ from llamafactory.v1.trainers.dpo_trainer import DPOTrainer, compute_sigmoid_dpo # Mock helpers # ============================================================================== + def _make_mock_v1( pref_beta: float = 0.1, dpo_label_smoothing: float = 0.0, @@ -67,6 +68,7 @@ R_REJECTED = torch.tensor([-3.2, -2.7, -4.2, -1.8]) # Test 1 — Core loss correctness (pure function ↔ v1 instance ↔ v0/TRL) # ============================================================================== + def test_sigmoid_dpo_loss_correctness(): """Comprehensive correctness check for compute_sigmoid_dpo_loss and its wrapper.""" # ---- 1a: pure function matches instance method ---- @@ -78,7 +80,12 @@ def test_sigmoid_dpo_loss_correctness(): # ---- 1b: v1 matches v0 (TRL) on fixed inputs ---- v0 = _make_mock_v0_dpo(beta=0.1) v0_losses, _, _ = CustomDPOTrainer.dpo_loss( - v0, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED, loss_type="sigmoid", + v0, + P_CHOSEN, + P_REJECTED, + R_CHOSEN, + R_REJECTED, + loss_type="sigmoid", ) torch.testing.assert_close(actual, v0_losses, rtol=1e-6, atol=1e-6) @@ -87,7 +94,12 @@ def test_sigmoid_dpo_loss_correctness(): v0b = _make_mock_v0_dpo(beta=beta) v1b = _make_mock_v1(pref_beta=beta) vl, _, _ = CustomDPOTrainer.dpo_loss( - v0b, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED, loss_type="sigmoid", + v0b, + P_CHOSEN, + P_REJECTED, + R_CHOSEN, + R_REJECTED, + loss_type="sigmoid", ) v1l = DPOTrainer._sigmoid_dpo_loss(v1b, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED) torch.testing.assert_close(v1l, vl, rtol=1e-6, atol=1e-6) @@ -97,7 +109,12 @@ def test_sigmoid_dpo_loss_correctness(): v0s = _make_mock_v0_dpo(beta=0.1, label_smoothing=ls) v1s = _make_mock_v1(pref_beta=0.1, dpo_label_smoothing=ls) vl, _, _ = CustomDPOTrainer.dpo_loss( - v0s, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED, loss_type="sigmoid", + v0s, + P_CHOSEN, + P_REJECTED, + R_CHOSEN, + R_REJECTED, + loss_type="sigmoid", ) v1l = DPOTrainer._sigmoid_dpo_loss(v1s, P_CHOSEN, P_REJECTED, R_CHOSEN, R_REJECTED) torch.testing.assert_close(v1l, vl, rtol=1e-6, atol=1e-6) @@ -112,13 +129,17 @@ def test_sigmoid_dpo_loss_correctness(): v1c = _make_mock_v1(pref_beta=0.1) loss_good = DPOTrainer._sigmoid_dpo_loss( v1c, - torch.tensor([-1.0]), torch.tensor([-10.0]), - torch.tensor([-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 62cefdb38..893614b0f 100644 --- a/tests_v1/trainers/test_fsdp2_dpo_trainer.py +++ b/tests_v1/trainers/test_fsdp2_dpo_trainer.py @@ -33,7 +33,6 @@ template: qwen3_nothink kernel_config: name: auto - include_kernels: auto quant_config: null diff --git a/tests_v1/trainers/test_fsdp2_sft_trainer.py b/tests_v1/trainers/test_fsdp2_sft_trainer.py index b5f1176ea..256b8098d 100644 --- a/tests_v1/trainers/test_fsdp2_sft_trainer.py +++ b/tests_v1/trainers/test_fsdp2_sft_trainer.py @@ -30,7 +30,6 @@ model_class: llm kernel_config: name: auto - include_kernels: auto quant_config: null