[ci] pin ruff version and fix lint errors (#10681)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

.github/workflows/docker.yml

@@ -29,8 +29,6 @@ jobs:
       matrix:
         include:
           - device: "cuda"
-          - device: "npu-a2"
-          - device: "npu-a3"
 
     runs-on: ubuntu-latest
 
@@ -71,14 +69,6 @@ jobs:
           username: ${{ vars.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
 
-      - name: Login to Quay
-        if: ${{ github.event_name != 'pull_request' && startsWith(matrix.device, 'npu') }}
-        uses: docker/login-action@v3
-        with:
-          registry: quay.io
-          username: ${{ vars.QUAY_ASCEND_USERNAME }}
-          password: ${{ secrets.QUAY_ASCEND_TOKEN }}
-
       - name: Build and push Docker image (CUDA)
         if: ${{ matrix.device == 'cuda' }}
         uses: docker/build-push-action@v6
@@ -88,29 +78,3 @@ jobs:
           push: ${{ github.event_name != 'pull_request' }}
           tags: |
             docker.io/hiyouga/llamafactory:${{ steps.version.outputs.tag }}
-
-      - name: Build and push Docker image (NPU-A2)
-        if: ${{ matrix.device == 'npu-a2' }}
-        uses: docker/build-push-action@v6
-        with:
-          context: .
-          platforms: linux/amd64,linux/arm64
-          file: ./docker/docker-npu/Dockerfile
-          push: ${{ github.event_name != 'pull_request' }}
-          tags: |
-            docker.io/hiyouga/llamafactory:${{ steps.version.outputs.tag }}-npu-a2
-            quay.io/ascend/llamafactory:${{ steps.version.outputs.tag }}-npu-a2
-
-      - name: Build and push Docker image (NPU-A3)
-        if: ${{ matrix.device == 'npu-a3' }}
-        uses: docker/build-push-action@v6
-        with:
-          context: .
-          platforms: linux/amd64,linux/arm64
-          file: ./docker/docker-npu/Dockerfile
-          build-args: |
-            BASE_IMAGE=quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11
-          push: ${{ github.event_name != 'pull_request' }}
-          tags: |
-            docker.io/hiyouga/llamafactory:${{ steps.version.outputs.tag }}-npu-a3
-            quay.io/ascend/llamafactory:${{ steps.version.outputs.tag }}-npu-a3

.github/workflows/docker_npu.yml

@@ -0,0 +1,115 @@
+name: docker-npu
+
+on:
+  workflow_dispatch:
+  schedule:
+    - cron: "17 2 * * *"
+      timezone: "Asia/Shanghai"
+  release:
+    types:
+      - published
+
+jobs:
+  build:
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - device: "npu-a2"
+            os: "ubuntu"
+            base_image: "quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11"
+          - device: "npu-a3"
+            os: "ubuntu"
+            base_image: "quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11"
+          - device: "npu-a2"
+            os: "openeuler"
+            base_image: "quay.io/ascend/cann:9.0.0-910b-openeuler24.03-py3.11"
+          - device: "npu-a3"
+            os: "openeuler"
+            base_image: "quay.io/ascend/cann:9.0.0-a3-openeuler24.03-py3.11"
+
+    runs-on: ubuntu-latest
+
+    concurrency:
+      group: ${{ github.workflow }}-${{ github.ref }}-${{ matrix.device }}-${{ matrix.os }}
+      cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
+    environment:
+      name: docker
+      url: https://hub.docker.com/r/hiyouga/llamafactory
+
+    steps:
+      - name: Free up disk space
+        uses: jlumbroso/free-disk-space@v1.3.1
+        with:
+          tool-cache: true
+          docker-images: false
+
+      - name: Checkout
+        uses: actions/checkout@v6
+
+      - name: Get llamafactory version
+        id: version
+        run: |
+          if [ "${{ github.event_name }}" = "release" ]; then
+            echo "tag=$(grep -oP 'VERSION = "\K[^"]+' src/llamafactory/extras/env.py)" >> "$GITHUB_OUTPUT"
+          else
+            echo "tag=latest" >> "$GITHUB_OUTPUT"
+          fi
+
+      - name: Get NPU image tag
+        id: npu_tag
+        env:
+          BASE_IMAGE: ${{ matrix.base_image }}
+          DEVICE: ${{ matrix.device }}
+          MATRIX_OS: ${{ matrix.os }}
+          LLAMAFACTORY_VERSION: ${{ steps.version.outputs.tag }}
+        run: |
+          base_image_tag="${BASE_IMAGE##*:}"
+          cann_version="${base_image_tag%%-*}"
+          torch_npu_version="$(sed -nE 's/^torch[-_]npu==([0-9]+(\.[0-9]+)*).*/\1/p' requirements/npu.txt)"
+          accelerator="${DEVICE#npu-}"
+          accelerator="${accelerator^^}"
+          operating_system="$(grep -oE '(ubuntu|openeuler)' <<< "${base_image_tag}" | head -n 1)"
+          python_version="$(grep -oE 'py[0-9]+\.[0-9]+' <<< "${base_image_tag}" | head -n 1)"
+
+          if [[ -z "${cann_version}" || -z "${torch_npu_version}" || -z "${operating_system}" || -z "${python_version}" ]]; then
+            echo "Failed to derive the NPU image tag from ${BASE_IMAGE} and requirements/npu.txt" >&2
+            exit 1
+          fi
+
+          if [[ "${operating_system}" != "${MATRIX_OS}" ]]; then
+            echo "Operating system ${operating_system} derived from ${BASE_IMAGE} does not match matrix OS ${MATRIX_OS}" >&2
+            exit 1
+          fi
+
+          echo "tag=${LLAMAFACTORY_VERSION}-cann${cann_version}-torch_npu${torch_npu_version}-${accelerator}-${operating_system}-${python_version}" >> "$GITHUB_OUTPUT"
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
+      - name: Login to Docker Hub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ vars.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+
+      - name: Login to Quay
+        uses: docker/login-action@v3
+        with:
+          registry: quay.io
+          username: ${{ vars.QUAY_ASCEND_USERNAME }}
+          password: ${{ secrets.QUAY_ASCEND_TOKEN }}
+
+      - name: Build and push Docker image (${{ matrix.device }}-${{ matrix.os }})
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          platforms: linux/amd64,linux/arm64
+          file: ./docker/docker-npu/Dockerfile
+          build-args: |
+            BASE_IMAGE=${{ matrix.base_image }}
+          push: true
+          tags: |
+            docker.io/hiyouga/llamafactory:${{ steps.npu_tag.outputs.tag }}
+            quay.io/ascend/llamafactory:${{ steps.npu_tag.outputs.tag }}

.pre-commit-config.yaml

@@ -21,7 +21,7 @@ repos:
         args: [--py39-plus]
 
 -   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.13.2
+    rev: v0.15.5
     hooks:
     -   id: ruff
         args: [--fix]

Makefile

@@ -2,9 +2,12 @@
 
 check_dirs := scripts src tests tests_v1
 
+ruff_version := 0.15.5
+
 RUN := $(shell command -v uv >/dev/null 2>&1 && echo "uv run" || echo "")
 BUILD := $(shell command -v uv >/dev/null 2>&1 && echo "uv build" || echo "python -m build")
 TOOL := $(shell command -v uv >/dev/null 2>&1 && echo "uvx" || echo "")
+RUFF := $(shell command -v uv >/dev/null 2>&1 && echo "uvx ruff@$(ruff_version)" || echo "ruff")
 
 build:
 	$(BUILD)
@@ -17,12 +20,12 @@ license:
 	$(RUN) python3 tests/check_license.py $(check_dirs)
 
 quality:
-	$(TOOL) ruff check $(check_dirs)
-	$(TOOL) ruff format --check $(check_dirs)
+	$(RUFF) check $(check_dirs)
+	$(RUFF) format --check $(check_dirs)
 
 style:
-	$(TOOL) ruff check $(check_dirs) --fix
-	$(TOOL) ruff format $(check_dirs)
+	$(RUFF) check $(check_dirs) --fix
+	$(RUFF) format $(check_dirs)
 
 test:
 	WANDB_DISABLED=true $(RUN) pytest -vv --import-mode=importlib tests/ tests_v1/

README.md

@@ -604,19 +604,23 @@ To enable FlashAttention-2 on the Windows platform, please use the script from [
 
 <details><summary>For Ascend NPU users</summary>
 
-To install LLaMA Factory on Ascend NPU devices, please upgrade Python to version 3.10 or higher: `pip install -r requirements/npu.txt`. Additionally, you need to install the **Ascend CANN Toolkit and Kernels**. Please follow the [installation tutorial](https://llamafactory.readthedocs.io/en/latest/advanced/npu_installation.html).
+To install LLaMA Factory on Ascend NPU devices, please upgrade Python to version 3.10 or higher: `pip install -r requirements/npu.txt`. Additionally, you need to install the **Ascend CANN Toolkit and Kernels**. Please follow the [installation tutorial](https://llamafactory.readthedocs.io/en/latest/multibackend/npu/npu_installation.html).
 
 
 You can also download the pre-built Docker images:
 
 ```bash
 # Docker Hub
-docker pull hiyouga/llamafactory:latest-npu-a2
-docker pull hiyouga/llamafactory:latest-npu-a3
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-ubuntu-py3.11
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-openeuler-py3.11
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-openeuler-py3.11
 
 # quay.io
-docker pull quay.io/ascend/llamafactory:latest-npu-a2
-docker pull quay.io/ascend/llamafactory:latest-npu-a3
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-ubuntu-py3.11
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-openeuler-py3.11
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-openeuler-py3.11
 ```
 
 #### Install BitsAndBytes
@@ -697,12 +701,28 @@ docker compose up -d
 docker compose exec llamafactory bash
 ```
 
-For Ascend NPU users:
+For Ascend NPU users (A2 with Ubuntu by default):
 
 ```bash
 cd docker/docker-npu/
-docker compose up -d
-docker compose exec llamafactory bash
+docker compose up -d llamafactory-a2-ubuntu
+docker compose exec llamafactory-a2-ubuntu bash
+```
+
+Other NPU variants can be started with their corresponding profiles and services:
+
+```bash
+# A3 with Ubuntu
+docker compose --profile a3 up -d llamafactory-a3-ubuntu
+docker compose exec llamafactory-a3-ubuntu bash
+
+# A2 with openEuler
+docker compose --profile openeuler up -d llamafactory-a2-openeuler
+docker compose exec llamafactory-a2-openeuler bash
+
+# A3 with openEuler
+docker compose --profile a3-openeuler up -d llamafactory-a3-openeuler
+docker compose exec llamafactory-a3-openeuler bash
 ```
 
 For AMD ROCm users:

README_zh.md

@@ -605,18 +605,22 @@ pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/downl
 
 <details><summary>昇腾 NPU 用户指南</summary>
 
-在昇腾 NPU 设备上安装 LLaMA Factory 时，请升级 Python 到 3.10 及以上，并需要指定额外依赖项，使用 `pip install -r requirements/npu.txt` 命令安装。此外，还需要安装 **Ascend CANN Toolkit 与 Kernels**，安装方法请参考[安装教程](https://llamafactory.readthedocs.io/zh-cn/latest/advanced/npu_installation.html)。
+在昇腾 NPU 设备上安装 LLaMA Factory 时，请升级 Python 到 3.10 及以上，并需要指定额外依赖项，使用 `pip install -r requirements/npu.txt` 命令安装。此外，还需要安装 **Ascend CANN Toolkit 与 Kernels**，安装方法请参考[安装教程](https://llamafactory.readthedocs.io/zh-cn/latest/multibackend/npu/npu_installation.html)。
 
 您可以直接下载预安装的最新docker镜像：
 
 ```bash
 # Docker Hub
-docker pull hiyouga/llamafactory:latest-npu-a2
-docker pull hiyouga/llamafactory:latest-npu-a3
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-ubuntu-py3.11
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-openeuler-py3.11
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-openeuler-py3.11
 
 # quay.io
-docker pull quay.io/ascend/llamafactory:latest-npu-a2
-docker pull quay.io/ascend/llamafactory:latest-npu-a3
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-ubuntu-py3.11
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-openeuler-py3.11
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-openeuler-py3.11
 ```
 
 #### 安装 BitsAndBytes
@@ -697,12 +701,28 @@ docker compose up -d
 docker compose exec llamafactory bash
 ```
 
-昇腾 NPU 用户：
+昇腾 NPU 用户（默认使用 A2 和 Ubuntu）：
 
 ```bash
 cd docker/docker-npu/
-docker compose up -d
-docker compose exec llamafactory bash
+docker compose up -d llamafactory-a2-ubuntu
+docker compose exec llamafactory-a2-ubuntu bash
+```
+
+其他 NPU 组合可以通过对应的 profile 和服务启动：
+
+```bash
+# A3 + Ubuntu
+docker compose --profile a3 up -d llamafactory-a3-ubuntu
+docker compose exec llamafactory-a3-ubuntu bash
+
+# A2 + openEuler
+docker compose --profile openeuler up -d llamafactory-a2-openeuler
+docker compose exec llamafactory-a2-openeuler bash
+
+# A3 + openEuler
+docker compose --profile a3-openeuler up -d llamafactory-a3-openeuler
+docker compose exec llamafactory-a3-openeuler bash
 ```
 
 AMD ROCm 用户：

docker/docker-npu/OVERVIEW.md

@@ -0,0 +1,232 @@
+# LLaMA Factory for Ascend NPU
+
+LLaMA Factory Ascend NPU images provide a ready-to-use environment for fine-tuning, evaluating, and serving large language and multimodal models on Huawei Ascend Atlas NPUs. The images are based on Ascend CANN container images and include LLaMA Factory, Python, PyTorch, torch-npu, Triton Ascend, DeepSpeed, and the metric dependencies used by LLaMA Factory.
+
+For installation and troubleshooting details, see the [English NPU installation guide](https://llamafactory.readthedocs.io/en/latest/multibackend/npu/npu_installation.html).
+
+## Quick Reference
+
+- Image registries:
+  - `docker.io/hiyouga/llamafactory`
+  - `quay.io/ascend/llamafactory`
+- Dockerfile: `docker/docker-npu/Dockerfile`
+- Docker Compose file: `docker/docker-npu/docker-compose.yml`
+- Default base image: `quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11`
+- Supported accelerators: Ascend A2 and A3
+- Supported container operating systems: Ubuntu 22.04 and openEuler 24.03
+- Target CPU architectures: `linux/amd64` and `linux/arm64`
+- Exposed ports:
+  - `7860`: LLaMA Board Web UI
+  - `8000`: API service
+- Ascend environment script: `/usr/local/Ascend/ascend-toolkit/set_env.sh`
+
+The current image variants are:
+
+| Accelerator | Container OS | CANN base image |
+| --- | --- | --- |
+| A2 | Ubuntu 22.04 | `quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11` |
+| A3 | Ubuntu 22.04 | `quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11` |
+| A2 | openEuler 24.03 | `quay.io/ascend/cann:9.0.0-910b-openeuler24.03-py3.11` |
+| A3 | openEuler 24.03 | `quay.io/ascend/cann:9.0.0-a3-openeuler24.03-py3.11` |
+
+## Image Contents and Intended Use
+
+The image is intended for Ascend NPU training, fine-tuning, evaluation, Web UI, and API workflows supported by LLaMA Factory. It installs the following core components:
+
+| Component | Version or source |
+| --- | --- |
+| CANN | Inherited from the selected CANN 9.0.0 base image |
+| Python | Python 3.11, inherited from the base image |
+| PyTorch | `2.7.1` |
+| torch-npu | `2.7.1.post4` |
+| torchvision | `0.22.1` |
+| torchaudio | `2.7.1` |
+| Triton Ascend | `3.2.1` |
+| DeepSpeed | `>=0.10.0,<=0.18.4` |
+| LLaMA Factory | Installed from the repository build context |
+
+The image does not include model weights or datasets. Mount or download them separately and comply with their respective licenses and acceptable-use requirements.
+
+## Image Tags and Dockerfile Archive
+
+Images use the following tag format:
+
+```text
+<llamafactory-version>-cann<cann-version>-torch_npu<torch-npu-version>-<accelerator>-<os>-<python-version>
+```
+
+| Field | Example | Description |
+| --- | --- | --- |
+| `llamafactory-version` | `latest` or `0.9.6` | Non-release builds use `latest`; release builds use the LLaMA Factory version |
+| `cann-version` | `9.0.0` | Parsed from the CANN base image tag |
+| `torch-npu-version` | `2.7.1` | Parsed from `requirements/npu.txt`; a suffix such as `.post4` is not included in the image tag |
+| `accelerator` | `A2` or `A3` | Ascend hardware generation selected for the image |
+| `os` | `ubuntu` or `openeuler` | Container operating system family |
+| `python-version` | `py3.11` | Parsed from the CANN base image tag |
+
+Examples:
+
+```text
+latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+latest-cann9.0.0-torch_npu2.7.1-A3-openeuler-py3.11
+0.9.6-cann9.0.0-torch_npu2.7.1-A3-ubuntu-py3.11
+```
+
+The CPU architecture is not part of the tag. Published images are configured as multi-platform images, and Docker selects the `linux/amd64` or `linux/arm64` manifest for the host automatically.
+
+The Dockerfile and its distribution overview are archived together at:
+
+```text
+docker/docker-npu/
+├── Dockerfile
+├── OVERVIEW.md
+├── OVERVIEW.zh.md
+└── docker-compose.yml
+```
+
+## Quick Start
+
+### Prerequisites
+
+Before starting a container:
+
+1. Install an Ascend driver and firmware compatible with the CANN version in the image.
+2. Verify that `npu-smi info` works on the host.
+3. Install Docker with permission to access the required Ascend device nodes and driver files.
+
+Driver, firmware, CANN, torch-npu, and the target Ascend hardware must be mutually compatible.
+
+### Pull and Run
+
+The following example starts the latest A2 Ubuntu image with one NPU. Change the image tag and `/dev/davinci0` as needed.
+
+```bash
+export IMAGE=quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+
+docker pull "$IMAGE"
+
+docker run --rm -it \
+  --name llamafactory-npu \
+  --ipc=host \
+  --device=/dev/davinci0 \
+  --device=/dev/davinci_manager \
+  --device=/dev/devmm_svm \
+  --device=/dev/hisi_hdc \
+  -v /usr/local/dcmi:/usr/local/dcmi \
+  -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+  -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
+  -v /etc/ascend_install.info:/etc/ascend_install.info \
+  -v "$HOME/.cache/huggingface:/root/.cache/huggingface" \
+  -p 7860:7860 \
+  -p 8000:8000 \
+  "$IMAGE" \
+  bash
+```
+
+The host path for `npu-smi` may be `/usr/local/sbin/npu-smi` on some driver installations. Adjust the mount source when necessary. Add more `--device=/dev/davinci<N>` options to expose additional NPUs.
+
+Verify the runtime inside the container:
+
+```bash
+source /usr/local/Ascend/ascend-toolkit/set_env.sh
+npu-smi info
+python -c "import torch, torch_npu; print(torch.__version__, torch_npu.__version__, torch.npu.is_available())"
+llamafactory-cli help
+```
+
+Start LLaMA Board when needed:
+
+```bash
+llamafactory-cli webui
+```
+
+### Build Locally
+
+Run the build from the repository root. The following example builds the A3 openEuler variant:
+
+```bash
+docker build \
+  -f ./docker/docker-npu/Dockerfile \
+  --build-arg BASE_IMAGE=quay.io/ascend/cann:9.0.0-a3-openeuler24.03-py3.11 \
+  --build-arg PIP_INDEX=https://pypi.org/simple \
+  -t llamafactory:npu-a3-openeuler \
+  .
+```
+
+Available build arguments:
+
+| Argument | Default | Purpose |
+| --- | --- | --- |
+| `BASE_IMAGE` | A2 Ubuntu CANN 9.0.0 image | Selects the accelerator and container OS variant |
+| `PIP_INDEX` | `https://pypi.org/simple` | Selects the Python package index |
+| `PYTORCH_INDEX` | `https://download.pytorch.org/whl/cpu` | Selects the PyTorch wheel index used with torch-npu |
+| `HTTP_PROXY` | Empty | Provides an optional HTTP/HTTPS proxy during the build |
+
+Docker Compose can build and start each supported variant:
+
+```bash
+cd docker/docker-npu
+
+# A2 with Ubuntu
+docker compose up -d llamafactory-a2-ubuntu
+
+# A3 with Ubuntu
+docker compose --profile a3 up -d llamafactory-a3-ubuntu
+
+# A2 with openEuler
+docker compose --profile openeuler up -d llamafactory-a2-openeuler
+
+# A3 with openEuler
+docker compose --profile a3-openeuler up -d llamafactory-a3-openeuler
+```
+
+### Extend or Develop from the Image
+
+For interactive development, mount a local checkout and reinstall it in editable mode inside the container:
+
+```bash
+git clone https://github.com/hiyouga/LLaMA-Factory.git
+cd LLaMA-Factory
+
+# Add the same Ascend --device and driver mount options shown above.
+docker run --rm -it \
+  --ipc=host \
+  -v "$PWD:/workspace/LLaMA-Factory" \
+  -w /workspace/LLaMA-Factory \
+  "$IMAGE" \
+  bash
+
+pip install -e . --no-build-isolation
+```
+
+For a reproducible derived image, create a separate Dockerfile:
+
+```dockerfile
+FROM quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+
+COPY requirements-extension.txt /tmp/requirements-extension.txt
+RUN pip install --no-cache-dir -r /tmp/requirements-extension.txt
+
+COPY . /workspace/application
+WORKDIR /workspace/application
+```
+
+Pass Ascend devices and driver mounts when running the derived image; device access should not be embedded in the image itself.
+
+## Hardware Support and Compatibility Notes
+
+- A2 images use the `910b` CANN base image; A3 images use the `a3` CANN base image.
+- The image build targets both x86-64 (`linux/amd64`) and AArch64 (`linux/arm64`) hosts. This CPU architecture is independent of whether the accelerator is A2 or A3.
+- Ubuntu 22.04 and openEuler 24.03 refer to the operating system inside the container.
+- The current dependency baseline aligns PyTorch `2.7.1` with torch-npu `2.7.1.post4`. Upgrading either package independently may break compatibility.
+- Use a fixed release tag for reproducible production deployments. The `latest` tag can change after scheduled builds.
+- Legacy short tags such as `latest-npu-a2` do not encode the CANN, torch-npu, operating system, or Python versions. Prefer the full tag format documented above.
+- Validate the exact driver, firmware, CANN, and SoC combination before production deployment.
+
+## License and Disclaimer
+
+LLaMA Factory is distributed under the [Apache License 2.0](../../LICENSE).
+
+Ascend CANN, torch-npu, Triton Ascend, DeepSpeed, base operating-system packages, model weights, datasets, and other third-party components are governed by their respective licenses and terms. The LLaMA Factory license does not replace or override those terms.
+
+The image is provided on an "AS IS" basis, without warranties or conditions of any kind. Users are responsible for validating hardware and software compatibility, securing the container and its runtime configuration, complying with applicable licenses and laws, and reviewing model and dataset terms before training, evaluation, or deployment.

docker/docker-npu/OVERVIEW.zh.md

@@ -0,0 +1,232 @@
+# 面向昇腾 NPU 的 LLaMA Factory 镜像
+
+LLaMA Factory 昇腾 NPU 镜像面向华为昇腾 Atlas NPU，提供可直接用于大语言模型和多模态模型微调、评测与服务部署的运行环境。镜像基于昇腾 CANN 容器镜像构建，预装 LLaMA Factory、Python、PyTorch、torch-npu、Triton Ascend、DeepSpeed 和 LLaMA Factory 评测依赖。
+
+安装方法和问题排查请参考 [LLaMA Factory NPU 安装及配置文档](https://llamafactory.readthedocs.io/zh-cn/latest/multibackend/npu/npu_installation.html)。
+
+## 快速参考
+
+- 镜像仓库：
+  - `docker.io/hiyouga/llamafactory`
+  - `quay.io/ascend/llamafactory`
+- Dockerfile：`docker/docker-npu/Dockerfile`
+- Docker Compose 文件：`docker/docker-npu/docker-compose.yml`
+- 默认基础镜像：`quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11`
+- 支持的加速器：昇腾 A2、A3
+- 支持的容器操作系统：Ubuntu 22.04、openEuler 24.03
+- 目标 CPU 架构：`linux/amd64`、`linux/arm64`
+- 对外端口：
+  - `7860`：LLaMA Board Web UI
+  - `8000`：API 服务
+- 昇腾环境脚本：`/usr/local/Ascend/ascend-toolkit/set_env.sh`
+
+当前提供以下镜像组合：
+
+| 加速器 | 容器操作系统 | CANN 基础镜像 |
+| --- | --- | --- |
+| A2 | Ubuntu 22.04 | `quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11` |
+| A3 | Ubuntu 22.04 | `quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11` |
+| A2 | openEuler 24.03 | `quay.io/ascend/cann:9.0.0-910b-openeuler24.03-py3.11` |
+| A3 | openEuler 24.03 | `quay.io/ascend/cann:9.0.0-a3-openeuler24.03-py3.11` |
+
+## 镜像介绍
+
+该镜像用于运行 LLaMA Factory 支持的昇腾 NPU 训练、微调、评测、Web UI 和 API 服务，主要包含以下组件：
+
+| 组件 | 版本或来源 |
+| --- | --- |
+| CANN | 继承自所选 CANN 9.0.0 基础镜像 |
+| Python | Python 3.11，继承自基础镜像 |
+| PyTorch | `2.7.1` |
+| torch-npu | `2.7.1.post4` |
+| torchvision | `0.22.1` |
+| torchaudio | `2.7.1` |
+| Triton Ascend | `3.2.1` |
+| DeepSpeed | `>=0.10.0,<=0.18.4` |
+| LLaMA Factory | 从构建上下文中的仓库源码安装 |
+
+镜像不包含模型权重和数据集。请通过目录挂载或运行时下载的方式单独提供，并遵守对应的许可证和使用要求。
+
+## 镜像 Tag 说明与 Dockerfile 归档路径
+
+镜像使用以下 tag 格式：
+
+```text
+<llamafactory版本>-cann<CANN版本>-torch_npu<torch-npu版本>-<加速器>-<操作系统>-<Python版本>
+```
+
+| 字段 | 示例 | 说明 |
+| --- | --- | --- |
+| `llamafactory版本` | `latest` 或 `0.9.6` | 非 release 构建使用 `latest`，release 构建使用 LLaMA Factory 版本号 |
+| `CANN版本` | `9.0.0` | 从 CANN 基础镜像 tag 中提取 |
+| `torch-npu版本` | `2.7.1` | 从 `requirements/npu.txt` 中提取，镜像 tag 不包含 `.post4` 等后缀 |
+| `加速器` | `A2` 或 `A3` | 当前镜像所适配的昇腾硬件代际 |
+| `操作系统` | `ubuntu` 或 `openeuler` | 容器内操作系统类型 |
+| `Python版本` | `py3.11` | 从 CANN 基础镜像 tag 中提取 |
+
+示例：
+
+```text
+latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+latest-cann9.0.0-torch_npu2.7.1-A3-openeuler-py3.11
+0.9.6-cann9.0.0-torch_npu2.7.1-A3-ubuntu-py3.11
+```
+
+CPU 架构不写入 tag。发布镜像配置为多架构镜像，Docker 拉取时会根据宿主机自动选择 `linux/amd64` 或 `linux/arm64` 版本。
+
+Dockerfile 和用于镜像分发的概述文件在同一目录归档：
+
+```text
+docker/docker-npu/
+├── Dockerfile
+├── OVERVIEW.md
+├── OVERVIEW.zh.md
+└── docker-compose.yml
+```
+
+## 快速开始
+
+### 前置条件
+
+启动容器前需要：
+
+1. 在宿主机安装与镜像内 CANN 版本兼容的昇腾驱动和固件。
+2. 确认宿主机执行 `npu-smi info` 可以正常识别 NPU。
+3. 安装 Docker，并确保当前用户有权访问所需的昇腾设备节点和驱动文件。
+
+驱动、固件、CANN、torch-npu 与目标昇腾硬件需要保持兼容。
+
+### 拉取并运行镜像
+
+以下示例使用一张 NPU 启动最新的 A2 Ubuntu 镜像。请根据实际环境修改镜像 tag 和 `/dev/davinci0`。
+
+```bash
+export IMAGE=quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+
+docker pull "$IMAGE"
+
+docker run --rm -it \
+  --name llamafactory-npu \
+  --ipc=host \
+  --device=/dev/davinci0 \
+  --device=/dev/davinci_manager \
+  --device=/dev/devmm_svm \
+  --device=/dev/hisi_hdc \
+  -v /usr/local/dcmi:/usr/local/dcmi \
+  -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+  -v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
+  -v /etc/ascend_install.info:/etc/ascend_install.info \
+  -v "$HOME/.cache/huggingface:/root/.cache/huggingface" \
+  -p 7860:7860 \
+  -p 8000:8000 \
+  "$IMAGE" \
+  bash
+```
+
+部分驱动环境中的 `npu-smi` 位于 `/usr/local/sbin/npu-smi`，此时需要调整挂载源路径。使用多张 NPU 时，继续追加 `--device=/dev/davinci<N>` 参数。
+
+进入容器后验证运行环境：
+
+```bash
+source /usr/local/Ascend/ascend-toolkit/set_env.sh
+npu-smi info
+python -c "import torch, torch_npu; print(torch.__version__, torch_npu.__version__, torch.npu.is_available())"
+llamafactory-cli help
+```
+
+需要使用 LLaMA Board 时执行：
+
+```bash
+llamafactory-cli webui
+```
+
+### 本地构建
+
+在仓库根目录执行构建。以下示例构建 A3 openEuler 镜像：
+
+```bash
+docker build \
+  -f ./docker/docker-npu/Dockerfile \
+  --build-arg BASE_IMAGE=quay.io/ascend/cann:9.0.0-a3-openeuler24.03-py3.11 \
+  --build-arg PIP_INDEX=https://pypi.org/simple \
+  -t llamafactory:npu-a3-openeuler \
+  .
+```
+
+可用构建参数：
+
+| 参数 | 默认值 | 用途 |
+| --- | --- | --- |
+| `BASE_IMAGE` | A2 Ubuntu CANN 9.0.0 镜像 | 选择加速器和容器操作系统组合 |
+| `PIP_INDEX` | `https://pypi.org/simple` | 指定 Python 软件包索引 |
+| `PYTORCH_INDEX` | `https://download.pytorch.org/whl/cpu` | 指定配合 torch-npu 使用的 PyTorch wheel 索引 |
+| `HTTP_PROXY` | 空 | 构建期间可选的 HTTP/HTTPS 代理 |
+
+也可以通过 Docker Compose 构建并启动各个组合：
+
+```bash
+cd docker/docker-npu
+
+# A2 + Ubuntu
+docker compose up -d llamafactory-a2-ubuntu
+
+# A3 + Ubuntu
+docker compose --profile a3 up -d llamafactory-a3-ubuntu
+
+# A2 + openEuler
+docker compose --profile openeuler up -d llamafactory-a2-openeuler
+
+# A3 + openEuler
+docker compose --profile a3-openeuler up -d llamafactory-a3-openeuler
+```
+
+### 二次开发
+
+交互式开发时，可以将本地源码挂载到容器中，并在容器内以 editable 模式重新安装：
+
+```bash
+git clone https://github.com/hiyouga/LLaMA-Factory.git
+cd LLaMA-Factory
+
+# 同时添加前述昇腾 --device 和驱动目录挂载参数。
+docker run --rm -it \
+  --ipc=host \
+  -v "$PWD:/workspace/LLaMA-Factory" \
+  -w /workspace/LLaMA-Factory \
+  "$IMAGE" \
+  bash
+
+pip install -e . --no-build-isolation
+```
+
+需要可复现的派生镜像时，可以新建独立 Dockerfile：
+
+```dockerfile
+FROM quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+
+COPY requirements-extension.txt /tmp/requirements-extension.txt
+RUN pip install --no-cache-dir -r /tmp/requirements-extension.txt
+
+COPY . /workspace/application
+WORKDIR /workspace/application
+```
+
+运行派生镜像时仍需传入昇腾设备和驱动挂载参数，不应将设备访问配置固化到镜像中。
+
+## 硬件支持与兼容性说明
+
+- A2 镜像使用标记为 `910b` 的 CANN 基础镜像，A3 镜像使用标记为 `a3` 的 CANN 基础镜像。
+- 镜像构建目标同时包含 x86-64（`linux/amd64`）和 AArch64（`linux/arm64`）宿主机。CPU 架构与加速器属于 A2 还是 A3 无关。
+- Ubuntu 22.04 和 openEuler 24.03 指容器内部的操作系统。
+- 当前依赖基线将 PyTorch `2.7.1` 与 torch-npu `2.7.1.post4` 配套使用。单独升级其中一个软件包可能破坏兼容性。
+- 生产环境建议使用固定 release tag，以确保部署可复现；定时构建可能更新 `latest` tag。
+- `latest-npu-a2` 等旧式短 tag 没有体现 CANN、torch-npu、操作系统和 Python 版本，建议迁移到本文所述的完整 tag。
+- 正式部署前，请验证具体驱动、固件、CANN 和 SoC 组合的兼容性。
+
+## 许可证与免责声明
+
+LLaMA Factory 基于 [Apache License 2.0](../../LICENSE) 发布。
+
+昇腾 CANN、torch-npu、Triton Ascend、DeepSpeed、基础操作系统软件包、模型权重、数据集和其他第三方组件分别受其自身许可证与条款约束。LLaMA Factory 的许可证不会替代或覆盖这些条款。
+
+本镜像按“原样”提供，不附带任何明示或暗示的保证。用户需要自行验证软硬件兼容性、保障容器及运行配置的安全、遵守适用的许可证和法律，并在训练、评测或部署前审查模型与数据集的使用条款。

docker/docker-npu/docker-compose.yml

@@ -1,58 +1,80 @@
+x-build-args: &build-args
+  PIP_INDEX: https://pypi.org/simple
+
+x-build: &build
+  dockerfile: ./docker/docker-npu/Dockerfile
+  context: ../..
+
+x-npu-common: &npu-common
+  volumes:
+    - /usr/local/dcmi:/usr/local/dcmi
+    - /usr/local/bin/npu-smi:/usr/local/bin/npu-smi
+    - /usr/local/Ascend/driver:/usr/local/Ascend/driver
+    - /etc/ascend_install.info:/etc/ascend_install.info
+  ipc: host
+  tty: true
+  # shm_size: "16gb"  # ipc: host is set
+  stdin_open: true
+  command: bash
+  devices:
+    - /dev/davinci0
+    - /dev/davinci_manager
+    - /dev/devmm_svm
+    - /dev/hisi_hdc
+  restart: unless-stopped
+
 services:
-  llamafactory-a2:
+  llamafactory-a2-ubuntu:
+    <<: *npu-common
     build:
-      dockerfile: ./docker/docker-npu/Dockerfile
-      context: ../..
+      <<: *build
       args:
-        PIP_INDEX: https://pypi.org/simple
-    container_name: llamafactory-a2
-    image: llamafactory:npu-a2
-    volumes:
-      - /usr/local/dcmi:/usr/local/dcmi
-      - /usr/local/bin/npu-smi:/usr/local/bin/npu-smi
-      - /usr/local/Ascend/driver:/usr/local/Ascend/driver
-      - /etc/ascend_install.info:/etc/ascend_install.info
+        <<: *build-args
+        BASE_IMAGE: quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11
+    container_name: llamafactory-a2-ubuntu
+    image: llamafactory:npu-a2-ubuntu
     ports:
       - "7860:7860"
       - "8000:8000"
-    ipc: host
-    tty: true
-    # shm_size: "16gb"  # ipc: host is set
-    stdin_open: true
-    command: bash
-    devices:
-      - /dev/davinci0
-      - /dev/davinci_manager
-      - /dev/devmm_svm
-      - /dev/hisi_hdc
-    restart: unless-stopped
 
-  llamafactory-a3:
+  llamafactory-a3-ubuntu:
+    <<: *npu-common
     profiles: ["a3"]
     build:
-      dockerfile: ./docker/docker-npu/Dockerfile
-      context: ../..
+      <<: *build
       args:
+        <<: *build-args
         BASE_IMAGE: quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11
-        PIP_INDEX: https://pypi.org/simple
-    container_name: llamafactory-a3
-    image: llamafactory:npu-a3
-    volumes:
-      - /usr/local/dcmi:/usr/local/dcmi
-      - /usr/local/bin/npu-smi:/usr/local/bin/npu-smi
-      - /usr/local/Ascend/driver:/usr/local/Ascend/driver
-      - /etc/ascend_install.info:/etc/ascend_install.info
+    container_name: llamafactory-a3-ubuntu
+    image: llamafactory:npu-a3-ubuntu
     ports:
       - "7861:7860"
       - "8001:8000"
-    ipc: host
-    tty: true
-    # shm_size: "16gb"  # ipc: host is set
-    stdin_open: true
-    command: bash
-    devices:
-      - /dev/davinci0
-      - /dev/davinci_manager
-      - /dev/devmm_svm
-      - /dev/hisi_hdc
-    restart: unless-stopped
+
+  llamafactory-a2-openeuler:
+    <<: *npu-common
+    profiles: ["openeuler"]
+    build:
+      <<: *build
+      args:
+        <<: *build-args
+        BASE_IMAGE: quay.io/ascend/cann:9.0.0-910b-openeuler24.03-py3.11
+    container_name: llamafactory-a2-openeuler
+    image: llamafactory:npu-a2-openeuler
+    ports:
+      - "7862:7860"
+      - "8002:8000"
+
+  llamafactory-a3-openeuler:
+    <<: *npu-common
+    profiles: ["a3-openeuler"]
+    build:
+      <<: *build
+      args:
+        <<: *build-args
+        BASE_IMAGE: quay.io/ascend/cann:9.0.0-a3-openeuler24.03-py3.11
+    container_name: llamafactory-a3-openeuler
+    image: llamafactory:npu-a3-openeuler
+    ports:
+      - "7863:7860"
+      - "8003:8000"

examples/v1/train_batching_strategy/train_full_fsdp2_batching_normal.yaml

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

examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_batching.yaml

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

examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml

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

examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml

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

examples/v1/train_freeze/train_freeze_sft.yaml

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

examples/v1/train_full/train_full_deepspeed.yaml

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

examples/v1/train_full/train_full_fsdp2.yaml

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

examples/v1/train_full/train_full_liger_kernel.yaml

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

examples/v1/train_full/train_full_ulysses_cp.yaml

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

examples/v1/train_lora/train_lora_dpo.yaml

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

examples/v1/train_lora/train_lora_sft.yaml

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

examples/v1/train_lora/train_lora_sft_rank0.yaml

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

examples/v1/train_qlora/quantization.yaml

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

pyproject.toml

@@ -112,7 +112,7 @@ ignore = [
     "D105",  # no doc magic method
     "D107",  # no doc __init__
 ]
-extend-select = [
+select = [
     "C",      # complexity
     "E",      # error
     "F",      # pyflakes

src/llamafactory/model/adapter.py

@@ -198,12 +198,16 @@ def _setup_lora_tuning(
                 pass  # already loaded via load_unsloth_peft_model in loader.py
             else:
                 if model_args.use_unsloth:
-                    peft_model = load_unsloth_peft_model(config, model_args, finetuning_args, is_trainable=is_trainable)
+                    peft_model = load_unsloth_peft_model(
+                        config, model_args, finetuning_args, is_trainable=is_trainable
+                    )
                     if peft_model is not None:
                         model = peft_model
 
                 if not model_args.use_unsloth:  # unsloth was disabled or fell back
-                    model = PeftModel.from_pretrained(model, adapter_to_resume, is_trainable=is_trainable, **init_kwargs)
+                    model = PeftModel.from_pretrained(
+                        model, adapter_to_resume, is_trainable=is_trainable, **init_kwargs
+                    )
 
         logger.info_rank0("Loaded adapter(s): {}".format(",".join(model_args.adapter_name_or_path)))
 

src/llamafactory/model/model_utils/embedding.py

@@ -120,9 +120,7 @@ def _noisy_mean_initialization(
     avg_weight = _existing_embeddings(embed_weight, num_new_tokens, token_ids).mean(dim=0, keepdim=True)
 
     if token_ids:
-        noise_weight = torch.empty(
-            len(token_ids), embedding_dim, device=embed_weight.device, dtype=embed_weight.dtype
-        )
+        noise_weight = torch.empty(len(token_ids), embedding_dim, device=embed_weight.device, dtype=embed_weight.dtype)
         noise_weight.normal_(mean=0, std=(1.0 / math.sqrt(embedding_dim)))
         embed_weight[token_ids] = avg_weight + noise_weight
     else:
@@ -202,8 +200,7 @@ def _description_based_initialization(
             if len(valid_token_ids) == 0:
                 # Fallback: use mean of all existing embeddings
                 logger.warning_rank0(
-                    f"Description for token '{token_str}' contains no valid tokens. "
-                    "Using mean of existing embeddings."
+                    f"Description for token '{token_str}' contains no valid tokens. Using mean of existing embeddings."
                 )
                 base_embedding = fallback_embedding
             else:

src/llamafactory/train/hyper_parallel/trainer.py

@@ -278,7 +278,9 @@ class HyperParallelTrainer(CustomSeq2SeqTrainer):
         )
         logical_batches = len(batch_sampler) // self._cp_size
         dp_size = max(1, get_platform().get_world_size() // self._cp_size)
-        logical_length = logical_batches // dp_size if self.args.dataloader_drop_last else _ceil_div(logical_batches, dp_size)
+        logical_length = (
+            logical_batches // dp_size if self.args.dataloader_drop_last else _ceil_div(logical_batches, dp_size)
+        )
 
         dataloader_params = {
             "batch_sampler": batch_sampler,

src/llamafactory/train/tuner.py

@@ -90,9 +90,7 @@ def _training_function(config: dict[str, Any]) -> None:
 
     if finetuning_args.stage in ["pt", "sft"] and finetuning_args.use_hyper_parallel:
         if not is_hyper_parallel_available():
-            raise ImportError(
-                "hyper_parallel is not installed. Please install it with `pip install hyper_parallel`."
-            )
+            raise ImportError("hyper_parallel is not installed. Please install it with `pip install hyper_parallel`.")
         if finetuning_args.stage == "pt":
             from .hyper_parallel import run_pt as run_pt_hp
 

src/llamafactory/v1/accelerator/interface.py

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

src/llamafactory/v1/config/training_args.py

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

src/llamafactory/v1/core/base_trainer.py

@@ -94,9 +94,12 @@ class BaseTrainer:
         dist_name = self.args.dist_config.name if self.args.dist_config is not None else None
 
         if dist_name == "deepspeed":
-            from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin
+            if self.args.cp_size > 1:
+                raise ValueError("Context parallelism currently requires `dist_config.name: fsdp2`.")
 
-            self._deepspeed_engine = DistributedPlugin("deepspeed")(
+            from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
+
+            self._deepspeed_engine = DistributedPlugin("deepspeed").shard_model(
                 self.model,
                 self.args.dist_config,
                 num_micro_batch=self.train_batch_generator.num_micro_batch,
@@ -139,7 +142,7 @@ class BaseTrainer:
         self.state.global_step = self.global_step
         self.state.epoch = self._resume_epoch
 
-        if self.args.dist_config is not None and self.args.dist_config.get("cp_size", 1) > 1:
+        if self.args.cp_size > 1:
             # qwen3.5 is not supported because of the different attention implementation, which will be supported in the future.
             if model.config.model_type == "qwen3_5":
                 raise RuntimeError(
@@ -152,7 +155,7 @@ class BaseTrainer:
                     "Sequence parallelism requires flash attention. Please set `flash_attn: flash_attention_2`."
                 )
 
-            SequenceParallelModelPlugin(self.args.dist_config.get("cp_mode", "ulysses"))(model, self.args.dist_config)
+            SequenceParallelModelPlugin(self.args.cp_mode)(model, self.args.cp_size)
 
     def _create_batch_generator(self) -> None:
         if (
@@ -183,9 +186,9 @@ class BaseTrainer:
                 device_ids = None if self.device.type == "cpu" else [self.device.index]
                 self.model = DDP(self.model, device_ids=device_ids)
         else:
-            from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin
+            from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
 
-            self.model = DistributedPlugin(self.args.dist_config.name)(
+            self.model = DistributedPlugin(self.args.dist_config.name).shard_model(
                 self.model,
                 self.args.dist_config,
                 bf16=self.args.bf16,
@@ -256,7 +259,7 @@ class BaseTrainer:
                 step_valid_tokens = DistributedInterface().all_reduce(step_valid_tokens, op=ReduceOp.SUM)
                 num_micro = len(micro_batches)
                 for i, micro_batch in enumerate(micro_batches):
-                    if self.args.dist_config and self.args.dist_config.get("cp_size", 1) > 1:
+                    if self.args.cp_size > 1:
                         from ..plugins.model_plugins.parallelization.sequence_parallel import (
                             SequenceParallelLossPlugin,
                         )
@@ -291,9 +294,7 @@ class BaseTrainer:
                         # mp_shard=world); a separate CP reduce would over-count by sqrt(cp_size).
                         total_norm = total_norm.full_tensor()
                     # pass a Tensor: clip_grads_with_norm_ clamps max_norm / (total_norm + 1e-6).
-                    torch.nn.utils.clip_grads_with_norm_(
-                        self.model.parameters(), self.args.max_grad_norm, total_norm
-                    )
+                    torch.nn.utils.clip_grads_with_norm_(self.model.parameters(), self.args.max_grad_norm, total_norm)
                     grad_norm = total_norm.item()
 
                     if not torch.isfinite(torch.tensor(grad_norm)):  # type: ignore # pyright: ignore [reportUnknownReturnType]
@@ -352,7 +353,7 @@ class BaseTrainer:
     def save_model(self) -> None:
         """Save the model."""
         if self.args.dist_config is not None and self.args.dist_config.name in ("deepspeed", "fsdp2"):
-            from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin
+            from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
 
             DistributedPlugin(self.args.dist_config.name).save_model(
                 self.model, self.args.output_dir, self.renderer.processor

src/llamafactory/v1/core/data_engine.py

@@ -137,9 +137,7 @@ class DataEngine(Dataset):
         messages = sample.get("messages")
         if not messages:
             return [None]
-        cuts = [
-            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:

src/llamafactory/v1/core/model_engine.py

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

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

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

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

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

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

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

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

@@ -14,9 +14,32 @@
 
 import json
 from copy import deepcopy
+from functools import lru_cache
 from typing import Any
 
 
+_registered_dist_config: Any | None = None
+
+
+def register_deepspeed_dist_config(dist_config: Any | None) -> None:
+    """Register backend config before model loading without involving the accelerator."""
+    global _registered_dist_config
+    _registered_dist_config = dist_config
+    is_deepspeed_zero3_enabled.cache_clear()
+
+
+@lru_cache(maxsize=1)
+def is_deepspeed_zero3_enabled() -> bool:
+    dist_config = _registered_dist_config
+    if dist_config is None or getattr(dist_config, "name", None) != "deepspeed":
+        return False
+
+    config_file = dist_config.get("config_file")
+    if not config_file:
+        return False
+    return _load_deepspeed_config(config_file).get("zero_optimization", {}).get("stage") == 3
+
+
 def _normalize_precision_enabled(value: Any) -> bool | str:
     if isinstance(value, str):
         value_lower = value.lower()
@@ -69,18 +92,19 @@ def _load_deepspeed_config(config_file: str) -> dict[str, Any]:
         return json.load(f)
 
 
-def setup_deepspeed_zero3_model_loading(is_train: bool, dist_config: dict[str, Any] | None):
-    """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 "

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -37,8 +37,8 @@ logger = logging.get_logger(__name__)
 
 
 class SequenceParallelModelPlugin(BasePlugin):
-    def __call__(self, model, model_args):
-        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))
 

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

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

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

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

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

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

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

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

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

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

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

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

src/llamafactory/v1/plugins/trainer_plugins/optimizers/optimizer.py

@@ -19,6 +19,7 @@ from typing import TYPE_CHECKING
 from ....utils import logging
 from ....utils.plugin import BasePlugin
 
+
 if TYPE_CHECKING:
     from ....config.arg_utils import PluginConfig
     from ....utils.types import HFModel

src/llamafactory/v1/trainers/dpo_trainer.py

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

src/llamafactory/v1/trainers/rm_trainer.py

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

src/llamafactory/v1/trainers/sft_trainer.py

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

src/llamafactory/v1/utils/plugin.py

@@ -12,96 +12,98 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+"""Lightweight plugin routing and shared parameter parsing helpers."""
 
-from collections import defaultdict
-from 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)

src/llamafactory/v1/utils/types.py

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

src/llamafactory/webui/components/eval.py

@@ -60,7 +60,13 @@ def create_eval_tab(engine: "Engine") -> dict[str, "Component"]:
 
     input_elems.update({max_new_tokens, top_p, temperature, eval_seed, output_dir})
     elem_dict.update(
-        dict(max_new_tokens=max_new_tokens, top_p=top_p, temperature=temperature, eval_seed=eval_seed, output_dir=output_dir)
+        dict(
+            max_new_tokens=max_new_tokens,
+            top_p=top_p,
+            temperature=temperature,
+            eval_seed=eval_seed,
+            output_dir=output_dir,
+        )
     )
 
     with gr.Row():

tests_v1/config/test_args_parser.py

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

tests_v1/core/test_model_loader.py

@@ -30,9 +30,7 @@ def test_tiny_qwen():
 def test_tiny_qwen_with_kernel_plugin():
     from llamafactory.v1.plugins.model_plugins.kernels.ops.rms_norm.npu_rms_norm import npu_rms_norm_forward
 
-    model_args = ModelArguments(
-        model="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"):

tests_v1/plugins/model_plugins/test_kernel_plugin.py

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

tests_v1/plugins/model_plugins/test_ulysses_cp.py

@@ -18,6 +18,7 @@ import torch.multiprocessing as mp
 
 from llamafactory.v1.accelerator.interface import DistributedInterface
 from llamafactory.v1.config.model_args import ModelArguments
+from llamafactory.v1.config.training_args import TrainingArguments
 from llamafactory.v1.core.model_engine import ModelEngine
 from llamafactory.v1.plugins.model_plugins.parallelization.sequence_parallel import (
     SequenceParallelModelPlugin,
@@ -33,15 +34,14 @@ def _test_sequence_parallel_loss(
     with dist_env(local_rank, world_size, master_port):
         model_args = ModelArguments(model="llamafactory/tiny-random-qwen3")
 
-        # Initialize distributed interface with config
-        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 = {

tests_v1/trainers/test_dpo_loss_precision.py

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

tests_v1/trainers/test_fsdp2_dpo_trainer.py

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

tests_v1/trainers/test_fsdp2_sft_trainer.py

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