423A35C7/LLaMA-Factory
1
0
Fork 0
mirror of https://github.com/hiyouga/LLaMA-Factory.git synced 2026-07-28 03:36:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Jiaqi
2026-07-24 15:23:21 +08:00
committed by GitHub
parent 19e9fe3ced
commit 3f77101580
50 changed files with 843 additions and 1007 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
examples/v1/train_batching_strategy/train_full_fsdp2_batching_normal.yaml
View File

@@ -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

3
examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_batching.yaml
View File

@@ -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

3
examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml
View File

@@ -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

3
examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml
View File

@@ -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

3
examples/v1/train_freeze/train_freeze_sft.yaml
View File

@@ -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

2
examples/v1/train_full/train_full_deepspeed.yaml
View File

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

2
examples/v1/train_full/train_full_fsdp2.yaml
View File

@@ -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

3
examples/v1/train_full/train_full_liger_kernel.yaml
View File

@@ -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

5
examples/v1/train_full/train_full_ulysses_cp.yaml
View File

@@ -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

3
examples/v1/train_lora/train_lora_dpo.yaml
View File

@@ -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:

2
examples/v1/train_lora/train_lora_sft.yaml
View File

@@ -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:

2
examples/v1/train_lora/train_lora_sft_rank0.yaml
View File

@@ -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:

2
examples/v1/train_qlora/quantization.yaml
View File

@@ -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:

27
src/llamafactory/v1/accelerator/interface.py
View File

@@ -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())

32
src/llamafactory/v1/config/training_args.py
View File

@@ -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

19
src/llamafactory/v1/core/base_trainer.py
View File

@@ -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

4
src/llamafactory/v1/core/data_engine.py
View File

@@ -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:

36
src/llamafactory/v1/core/model_engine.py
View File

@@ -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

9
src/llamafactory/v1/core/rendering/format.py
View File

@@ -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

10
src/llamafactory/v1/core/rendering/rendering.py
View File

@@ -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.

4
src/llamafactory/v1/core/utils/checkpoint.py
View File

@@ -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,

36
src/llamafactory/v1/plugins/model_plugins/deepspeed_utils.py
View File

@@ -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 "

86
src/llamafactory/v1/plugins/model_plugins/kernels/base.py
View File

@@ -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: ...

201
src/llamafactory/v1/plugins/model_plugins/kernels/interface.py
View File

@@ -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)

36
src/llamafactory/v1/plugins/model_plugins/kernels/liger_kernel_ops.py
View File

@@ -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)

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

@@ -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

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

@@ -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

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

@@ -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

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

@@ -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)

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

@@ -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():

96
src/llamafactory/v1/plugins/model_plugins/kernels/registry.py
View File

@@ -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

7
src/llamafactory/v1/plugins/model_plugins/parallelization/sequence_parallel.py
View File

@@ -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))

163
src/llamafactory/v1/plugins/model_plugins/peft.py
View File

@@ -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.")

120
src/llamafactory/v1/plugins/model_plugins/quantization.py
View File

@@ -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

261
src/llamafactory/v1/plugins/trainer_plugins/batching.py
View File

@@ -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

50
src/llamafactory/v1/plugins/trainer_plugins/distributed/base.py Normal file
View File

@@ -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: ...

94
src/llamafactory/v1/plugins/trainer_plugins/distributed/hub.py
View File

@@ -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)

115
src/llamafactory/v1/plugins/trainer_plugins/distributed/interface.py Normal file
View File

@@ -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)

14
src/llamafactory/v1/trainers/dpo_trainer.py
View File

@@ -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)

5
src/llamafactory/v1/trainers/rm_trainer.py
View File

@@ -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)

2
src/llamafactory/v1/trainers/sft_trainer.py
View File

@@ -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(

138
src/llamafactory/v1/utils/plugin.py
View File

@@ -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)

13
src/llamafactory/v1/utils/types.py
View File

@@ -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."""

6
tests_v1/config/test_args_parser.py
View File

@@ -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

4
tests_v1/core/test_model_loader.py
View File

@@ -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"):

8
tests_v1/plugins/model_plugins/test_kernel_plugin.py
View File

@@ -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__

8
tests_v1/plugins/model_plugins/test_ulysses_cp.py
View File

@@ -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 = {

55
tests_v1/trainers/test_dpo_loss_precision.py
View File

@@ -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)

1
tests_v1/trainers/test_fsdp2_dpo_trainer.py
View File

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

1
tests_v1/trainers/test_fsdp2_sft_trainer.py
View File

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