10 Commits

Author SHA1 Message Date
Hazeldxq
f28afaf635 [v1] add FSDPTurbo EP/EFSDP plugin for MoE training (#10676) 2026-08-13 20:45:55 +08:00
yyj
bc4b42cefc [train] Harden KTransformers MoE LoRA SFT integration (#10738) 2026-08-13 20:43:15 +08:00
浮梦
199b8873d7 [train] add fa3 supported (#10742) 2026-08-13 20:39:01 +08:00
xvxuopop
0bbe481e6e [docker] upgrade NPU images to CANN 9.1 and PyTorch 2.10 (#10729) 2026-08-10 11:20:41 +08:00
haqishen
63a89710c7 [data] pad position_ids on non-FA2 packing path (fixes rotary crash for Gemma-3/4) (#10737)
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-08-09 16:00:10 +08:00
richboyneedcash
887b850813 [assets] fix broken links in README (#10728)
Co-authored-by: richboyneedcash <273099414+richboyneedcash@users.noreply.github.com>
Co-authored-by: TRAE CLI <noreply@bytedance.com>
2026-08-06 15:09:01 +08:00
xvxuopop
84576b1408 [v1] refactor NPU kernel matching by model type (#10643) 2026-08-04 19:38:41 +08:00
SSSSuperC
713b5a3f95 [model] add MOSS-VL support (#10708) 2026-08-03 18:18:24 +08:00
浮梦
62ae362455 [v1] Support multimodal data training (#10656)
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-31 18:54:13 +08:00
Kyungmin Kim
3984675dd5 fix(ci): align workflow Python version with requires-python (#10707) 2026-07-31 16:56:00 +08:00
82 changed files with 5386 additions and 978 deletions

View File

@@ -17,16 +17,16 @@ jobs:
include:
- device: "npu-a2"
os: "ubuntu"
base_image: "quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11"
base_image: "quay.io/ascend/cann:9.1.0-910b-ubuntu22.04-py3.12"
- device: "npu-a3"
os: "ubuntu"
base_image: "quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11"
base_image: "quay.io/ascend/cann:9.1.0-a3-ubuntu22.04-py3.12"
- device: "npu-a2"
os: "openeuler"
base_image: "quay.io/ascend/cann:9.0.0-910b-openeuler24.03-py3.11"
base_image: "quay.io/ascend/cann:9.1.0-910b-openeuler24.03-py3.12"
- device: "npu-a3"
os: "openeuler"
base_image: "quay.io/ascend/cann:9.0.0-a3-openeuler24.03-py3.11"
base_image: "quay.io/ascend/cann:9.1.0-a3-openeuler24.03-py3.12"
runs-on: ubuntu-latest
@@ -48,7 +48,7 @@ jobs:
- name: Checkout
uses: actions/checkout@v6
- name: Get llamafactory version
- name: Get LlamaFactory version
id: version
run: |
if [ "${{ github.event_name }}" = "release" ]; then
@@ -61,29 +61,46 @@ jobs:
id: npu_tag
env:
BASE_IMAGE: ${{ matrix.base_image }}
DEVICE: ${{ matrix.device }}
MATRIX_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)"
torch_npu_version="$(sed -nE 's/^torch[-_]npu==([0-9]+(\.[0-9]+)*(\.post[0-9]+)?).*/\1/p' requirements/npu.txt)"
soc="$(grep -oE '(910b|a3)' <<< "${base_image_tag}" | head -n 1)"
operating_system="$(grep -oE '(ubuntu|openeuler)[0-9]+(\.[0-9]+)*' <<< "${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
if [[ -z "${cann_version}" || -z "${torch_npu_version}" || -z "${soc}" || -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
if [[ "${operating_system}" != "${MATRIX_OS}"* ]]; then
echo "Operating system ${operating_system} derived from ${BASE_IMAGE} does not match matrix OS family ${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"
case "${MATRIX_DEVICE}" in
npu-a2) expected_soc="910b" ;;
npu-a3) expected_soc="a3" ;;
*)
echo "Unsupported NPU device ${MATRIX_DEVICE}" >&2
exit 1
;;
esac
if [[ "${soc}" != "${expected_soc}" ]]; then
echo "SoC ${soc} derived from ${BASE_IMAGE} does not match matrix device ${MATRIX_DEVICE}" >&2
exit 1
fi
if [[ "${LLAMAFACTORY_VERSION}" == "latest" ]]; then
echo "tag=latest-${soc}-${MATRIX_OS}" >> "$GITHUB_OUTPUT"
else
echo "tag=${LLAMAFACTORY_VERSION}-cann${cann_version}-torch_npu${torch_npu_version}-${soc}-${operating_system}-${python_version}" >> "$GITHUB_OUTPUT"
fi
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3

View File

@@ -29,7 +29,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: '3.10'
python-version: '3.11'
- name: Install dependencies
run: |
@@ -66,6 +66,7 @@ jobs:
path: docs/_build/html
deploy:
if: github.event_name == 'push' && github.ref == 'refs/heads/main'
environment:
name: github-pages
url: ${{ steps.deployment.outputs.page_url }}

View File

@@ -8,6 +8,7 @@ on:
paths:
- "**/*.py"
- "pyproject.toml"
- "requirements/fsdpturbo.txt"
- "Makefile"
- ".github/workflows/*.yml"
pull_request:
@@ -16,6 +17,7 @@ on:
paths:
- "**/*.py"
- "pyproject.toml"
- "requirements/fsdpturbo.txt"
- "Makefile"
- ".github/workflows/*.yml"
@@ -25,11 +27,11 @@ jobs:
fail-fast: false
matrix:
python:
- "3.11"
- "3.12"
os:
- "linux-aarch64-a2-4"
pytorch_npu:
- "2.7.1"
- "2.10.0"
runs-on: ${{ matrix.os }}
@@ -38,7 +40,7 @@ jobs:
cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
container:
image: ascendai/cann:9.0.0-910b-ubuntu22.04-py3.11
image: ascendai/cann:9.1.0-910b-ubuntu22.04-py3.12
env:
HF_ENDPOINT: https://hf-mirror.com
HF_TOKEN: ${{ secrets.HF_TOKEN }}
@@ -68,6 +70,7 @@ jobs:
uv pip install -e .
uv pip install -r requirements/npu.txt
uv pip install -r requirements/dev.txt
uv pip install --no-deps -r requirements/fsdpturbo.txt
- name: Install node
run: |

106
README.md
View File

@@ -1,9 +1,9 @@
![# LLaMA Factory](assets/logo.png)
![# LlamaFactory](assets/logo.png)
[![GitHub Repo stars](https://img.shields.io/github/stars/hiyouga/LLaMA-Factory?style=social)](https://github.com/hiyouga/LLaMA-Factory/stargazers)
[![GitHub last commit](https://img.shields.io/github/last-commit/hiyouga/LLaMA-Factory)](https://github.com/hiyouga/LLaMA-Factory/commits/main)
[![GitHub contributors](https://img.shields.io/github/contributors/hiyouga/LLaMA-Factory?color=orange)](https://github.com/hiyouga/LLaMA-Factory/graphs/contributors)
[![GitHub workflow](https://github.com/hiyouga/LLaMA-Factory/actions/workflows/tests.yml/badge.svg)](https://github.com/hiyouga/LLaMA-Factory/actions/workflows/tests.yml)
[![GitHub Repo stars](https://img.shields.io/github/stars/hiyouga/LlamaFactory?style=social)](https://github.com/hiyouga/LlamaFactory/stargazers)
[![GitHub last commit](https://img.shields.io/github/last-commit/hiyouga/LlamaFactory)](https://github.com/hiyouga/LlamaFactory/commits/main)
[![GitHub contributors](https://img.shields.io/github/contributors/hiyouga/LlamaFactory?color=orange)](https://github.com/hiyouga/LlamaFactory/graphs/contributors)
[![GitHub workflow](https://github.com/hiyouga/LlamaFactory/actions/workflows/tests.yml/badge.svg)](https://github.com/hiyouga/LlamaFactory/actions/workflows/tests.yml)
[![PyPI](https://img.shields.io/pypi/v/llamafactory)](https://pypi.org/project/llamafactory/)
[![Citation](https://img.shields.io/badge/citation-1000+-green)](https://scholar.google.com/scholar?cites=12620864006390196564)
[![Docker Pulls](https://img.shields.io/docker/pulls/hiyouga/llamafactory)](https://hub.docker.com/r/hiyouga/llamafactory/tags)
@@ -95,7 +95,7 @@ Read technical notes:
- [Download from Modelers Hub](#download-from-modelers-hub)
- [Use W&B Logger](#use-wb-logger)
- [Use SwanLab Logger](#use-swanlab-logger)
- [Projects using LLaMA Factory](#projects-using-llama-factory)
- [Projects using LlamaFactory](#projects-using-llamafactory)
- [License](#license)
- [Citation](#citation)
- [Acknowledgement](#acknowledgement)
@@ -121,35 +121,35 @@ Read technical notes:
## Blogs
> [!TIP]
> Now we have a dedicated blog for LLaMA Factory!
> Now we have a dedicated blog for LlamaFactory!
>
> Website: https://blog.llamafactory.net/en/
- 💡 [KTransformers Fine-Tuning × LLaMA Factory: Fine-tuning 1000 Billion models with 2 4090-GPU + CPU](https://blog.llamafactory.net/en/posts/ktransformers/) (English)
- 💡 [Easy Dataset × LLaMA Factory: Enabling LLMs to Efficiently Learn Domain Knowledge](https://buaa-act.feishu.cn/wiki/GVzlwYcRFiR8OLkHbL6cQpYin7g) (English)
- 💡 [DataFlow × LLaMA Factory: Producing High-Quality Data for LLM Training with a Data Preparation Pipeline](https://wcny4qa9krto.feishu.cn/wiki/LWkkwTDBfiiRKqkDSvucG6yjnbW) (English) | [中文](https://wcny4qa9krto.feishu.cn/wiki/LlMxweUAJimrmykRD5qcGuswnHd)
- 💡 [DataFlex × LLaMA Factory: A Data-Centric Dynamic Training System Built on LLaMA-Factory](https://wcny4qa9krto.feishu.cn/wiki/OlREwPQWdi9K6ZkJNHIciLhtnkv) (English) | [中文](https://wcny4qa9krto.feishu.cn/wiki/H2A9wSsbCinzavkT2oyc2C5Vn0e)
- [A One-Stop Code-Free Model Reinforcement Learning and Deployment Platform based on LLaMA-Factory and EasyR1](https://aws.amazon.com/cn/blogs/china/building-llm-model-hub-based-on-llamafactory-and-easyr1/) (Chinese)
- [How Apoidea Group enhances visual information extraction from banking documents with multimodal models using LLaMA-Factory on Amazon SageMaker HyperPod](https://aws.amazon.com/cn/blogs/machine-learning/how-apoidea-group-enhances-visual-information-extraction-from-banking-documents-with-multimodal-models-using-llama-factory-on-amazon-sagemaker-hyperpod/) (English)
- 💡 [KTransformers Fine-Tuning × LlamaFactory: Fine-tuning 1000 Billion models with 2 4090-GPU + CPU](https://blog.llamafactory.net/en/posts/ktransformers/) (English)
- 💡 [Easy Dataset × LlamaFactory: Enabling LLMs to Efficiently Learn Domain Knowledge](https://buaa-act.feishu.cn/wiki/GVzlwYcRFiR8OLkHbL6cQpYin7g) (English)
- 💡 [DataFlow × LlamaFactory: Producing High-Quality Data for LLM Training with a Data Preparation Pipeline](https://wcny4qa9krto.feishu.cn/wiki/LWkkwTDBfiiRKqkDSvucG6yjnbW) (English) | [中文](https://wcny4qa9krto.feishu.cn/wiki/LlMxweUAJimrmykRD5qcGuswnHd)
- 💡 [DataFlex × LlamaFactory: A Data-Centric Dynamic Training System Built on LlamaFactory](https://wcny4qa9krto.feishu.cn/wiki/OlREwPQWdi9K6ZkJNHIciLhtnkv) (English) | [中文](https://wcny4qa9krto.feishu.cn/wiki/H2A9wSsbCinzavkT2oyc2C5Vn0e)
- [A One-Stop Code-Free Model Reinforcement Learning and Deployment Platform based on LlamaFactory and EasyR1](https://aws.amazon.com/cn/blogs/china/building-llm-model-hub-based-on-llamafactory-and-easyr1/) (Chinese)
- [How Apoidea Group enhances visual information extraction from banking documents with multimodal models using LlamaFactory on Amazon SageMaker HyperPod](https://aws.amazon.com/cn/blogs/machine-learning/how-apoidea-group-enhances-visual-information-extraction-from-banking-documents-with-multimodal-models-using-llama-factory-on-amazon-sagemaker-hyperpod/) (English)
<details><summary>All Blogs</summary>
- [LLaMA Factory: Fine-tuning the DeepSeek-R1-Distill-Qwen-7B Model for News Classifier](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory_deepseek_r1_distill_7b) (Chinese)
- [A One-Stop Code-Free Model Fine-Tuning \& Deployment Platform based on SageMaker and LLaMA-Factory](https://aws.amazon.com/cn/blogs/china/a-one-stop-code-free-model-fine-tuning-deployment-platform-based-on-sagemaker-and-llama-factory/) (Chinese)
- [LLaMA Factory Multi-Modal Fine-Tuning Practice: Fine-Tuning Qwen2-VL for Personal Tourist Guide](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory_qwen2vl) (Chinese)
- [LLaMA Factory: Fine-tuning Llama3 for Role-Playing](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory) (Chinese)
- [LlamaFactory: Fine-tuning the DeepSeek-R1-Distill-Qwen-7B Model for News Classifier](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory_deepseek_r1_distill_7b) (Chinese)
- [A One-Stop Code-Free Model Fine-Tuning \& Deployment Platform based on SageMaker and LlamaFactory](https://aws.amazon.com/cn/blogs/china/a-one-stop-code-free-model-fine-tuning-deployment-platform-based-on-sagemaker-and-llama-factory/) (Chinese)
- [LlamaFactory Multi-Modal Fine-Tuning Practice: Fine-Tuning Qwen2-VL for Personal Tourist Guide](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory_qwen2vl) (Chinese)
- [LlamaFactory: Fine-tuning Llama3 for Role-Playing](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory) (Chinese)
</details>
## Changelog
[25/10/26] We support Megatron-core training backend with [**mcore_adapter**](https://github.com/alibaba/ROLL/tree/main/mcore_adapter). See [PR #9237](https://github.com/hiyouga/LLaMA-Factory/pull/9237) to get started.
[25/10/26] We support Megatron-core training backend with [**mcore_adapter**](https://github.com/alibaba/ROLL/tree/main/mcore_adapter). See [PR #9237](https://github.com/hiyouga/LlamaFactory/pull/9237) to get started.
[25/08/22] We supported **[OFT](https://arxiv.org/abs/2306.07280)** and **[OFTv2](https://arxiv.org/abs/2506.19847)**. See [examples](examples/README.md) for usage.
[25/08/20] We supported fine-tuning the **[Intern-S1-mini](https://huggingface.co/internlm/Intern-S1-mini)** models. See [PR #8976](https://github.com/hiyouga/LLaMA-Factory/pull/8976) to get started.
[25/08/20] We supported fine-tuning the **[Intern-S1-mini](https://huggingface.co/internlm/Intern-S1-mini)** models. See [PR #8976](https://github.com/hiyouga/LlamaFactory/pull/8976) to get started.
[25/08/06] We supported fine-tuning the **[GPT-OSS](https://github.com/openai/gpt-oss)** models. See [PR #8826](https://github.com/hiyouga/LLaMA-Factory/pull/8826) to get started.
[25/08/06] We supported fine-tuning the **[GPT-OSS](https://github.com/openai/gpt-oss)** models. See [PR #8826](https://github.com/hiyouga/LlamaFactory/pull/8826) to get started.
<details><summary>Full Changelog</summary>
@@ -159,13 +159,13 @@ Read technical notes:
[25/04/21] We supported the **[Muon](https://github.com/KellerJordan/Muon)** optimizer. See [examples](examples/README.md) for usage. Thank [@tianshijing](https://github.com/tianshijing)'s PR.
[25/04/16] We supported fine-tuning the **[InternVL3](https://huggingface.co/OpenGVLab/InternVL3-8B)** model. See [PR #7258](https://github.com/hiyouga/LLaMA-Factory/pull/7258) to get started.
[25/04/16] We supported fine-tuning the **[InternVL3](https://huggingface.co/OpenGVLab/InternVL3-8B)** model. See [PR #7258](https://github.com/hiyouga/LlamaFactory/pull/7258) to get started.
[25/04/14] We supported fine-tuning the **[GLM-Z1](https://huggingface.co/THUDM/GLM-Z1-9B-0414)** and **[Kimi-VL](https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct)** models.
[25/04/06] We supported fine-tuning the **[Llama 4](https://ai.meta.com/blog/llama-4-multimodal-intelligence/)** model. See [PR #7611](https://github.com/hiyouga/LLaMA-Factory/pull/7611) to get started.
[25/04/06] We supported fine-tuning the **[Llama 4](https://ai.meta.com/blog/llama-4-multimodal-intelligence/)** model. See [PR #7611](https://github.com/hiyouga/LlamaFactory/pull/7611) to get started.
[25/03/31] We supported fine-tuning the **[Qwen2.5 Omni](https://qwenlm.github.io/blog/qwen2.5-omni/)** model. See [PR #7537](https://github.com/hiyouga/LLaMA-Factory/pull/7537) to get started.
[25/03/31] We supported fine-tuning the **[Qwen2.5 Omni](https://qwenlm.github.io/blog/qwen2.5-omni/)** model. See [PR #7537](https://github.com/hiyouga/LlamaFactory/pull/7537) to get started.
[25/03/15] We supported **[SGLang](https://github.com/sgl-project/sglang)** as inference backend. Try `infer_backend: sglang` to accelerate inference.
@@ -217,13 +217,13 @@ Read technical notes:
[24/04/26] We supported fine-tuning the **LLaVA-1.5** multimodal LLMs. See [examples](examples/README.md) for usage.
[24/04/22] We provided a **[Colab notebook](https://colab.research.google.com/drive/1eRTPn37ltBbYsISy9Aw2NuI2Aq5CQrD9?usp=sharing)** for fine-tuning the Llama-3 model on a free T4 GPU. Two Llama-3-derived models fine-tuned using LLaMA Factory are available at Hugging Face, check [Llama3-8B-Chinese-Chat](https://huggingface.co/shenzhi-wang/Llama3-8B-Chinese-Chat) and [Llama3-Chinese](https://huggingface.co/zhichen/Llama3-Chinese) for details.
[24/04/22] We provided a **[Colab notebook](https://colab.research.google.com/drive/1eRTPn37ltBbYsISy9Aw2NuI2Aq5CQrD9?usp=sharing)** for fine-tuning the Llama-3 model on a free T4 GPU. Two Llama-3-derived models fine-tuned using LlamaFactory are available at Hugging Face, check [Llama3-8B-Chinese-Chat](https://huggingface.co/shenzhi-wang/Llama3-8B-Chinese-Chat) and [Llama3-Chinese](https://huggingface.co/zhichen/Llama3-Chinese) for details.
[24/04/21] We supported **[Mixture-of-Depths](https://arxiv.org/abs/2404.02258)** according to [AstraMindAI's implementation](https://github.com/astramind-ai/Mixture-of-depths). See [examples](examples/README.md) for usage.
[24/04/16] We supported **[BAdam](https://arxiv.org/abs/2404.02827)** optimizer. See [examples](examples/README.md) for usage.
[24/04/16] We supported **[unsloth](https://github.com/unslothai/unsloth)**'s long-sequence training (Llama-2-7B-56k within 24GB). It achieves **117%** speed and **50%** memory compared with FlashAttention-2, more benchmarks can be found in [this page](https://github.com/hiyouga/LLaMA-Factory/wiki/Performance-comparison).
[24/04/16] We supported **[unsloth](https://github.com/unslothai/unsloth)**'s long-sequence training (Llama-2-7B-56k within 24GB). It achieves **117%** speed and **50%** memory compared with FlashAttention-2, more benchmarks can be found in [this page](https://github.com/hiyouga/LlamaFactory/wiki/Performance-comparison).
[24/03/31] We supported **[ORPO](https://arxiv.org/abs/2403.07691)**. See [examples](examples/README.md) for usage.
@@ -241,11 +241,11 @@ Read technical notes:
[24/02/15] We supported **block expansion** proposed by [LLaMA Pro](https://github.com/TencentARC/LLaMA-Pro). See [examples](examples/README.md) for usage.
[24/02/05] Qwen1.5 (Qwen2 beta version) series models are supported in LLaMA-Factory. Check this [blog post](https://qwenlm.github.io/blog/qwen1.5/) for details.
[24/02/05] Qwen1.5 (Qwen2 beta version) series models are supported in LlamaFactory. Check this [blog post](https://qwenlm.github.io/blog/qwen1.5/) for details.
[24/01/18] We supported **agent tuning** for most models, equipping model with tool using abilities by fine-tuning with `dataset: glaive_toolcall_en`.
[23/12/23] We supported **[unsloth](https://github.com/unslothai/unsloth)**'s implementation to boost LoRA tuning for the LLaMA, Mistral and Yi models. Try `use_unsloth: true` argument to activate unsloth patch. It achieves **170%** speed in our benchmark, check [this page](https://github.com/hiyouga/LLaMA-Factory/wiki/Performance-comparison) for details.
[23/12/23] We supported **[unsloth](https://github.com/unslothai/unsloth)**'s implementation to boost LoRA tuning for the LLaMA, Mistral and Yi models. Try `use_unsloth: true` argument to activate unsloth patch. It achieves **170%** speed in our benchmark, check [this page](https://github.com/hiyouga/LlamaFactory/wiki/Performance-comparison) for details.
[23/12/12] We supported fine-tuning the latest MoE model **[Mixtral 8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1)** in our framework. See hardware requirement [here](#hardware-requirement).
@@ -280,7 +280,7 @@ Read technical notes:
</details>
> [!TIP]
> If you cannot use the latest feature, please pull the latest code and install LLaMA-Factory again.
> If you cannot use the latest feature, please pull the latest code and install LlamaFactory again.
## Supported Models
@@ -604,23 +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/multibackend/npu/npu_installation.html).
To install LlamaFactory on Ascend NPU devices, please use Python 3.12 and install the extra dependencies with `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-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
docker pull hiyouga/llamafactory:latest-910b-ubuntu
docker pull hiyouga/llamafactory:latest-a3-ubuntu
docker pull hiyouga/llamafactory:latest-910b-openeuler
docker pull hiyouga/llamafactory:latest-a3-openeuler
# quay.io
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
docker pull quay.io/ascend/llamafactory:latest-910b-ubuntu
docker pull quay.io/ascend/llamafactory:latest-a3-ubuntu
docker pull quay.io/ascend/llamafactory:latest-910b-openeuler
docker pull quay.io/ascend/llamafactory:latest-a3-openeuler
```
#### Install BitsAndBytes
@@ -683,7 +683,7 @@ See [examples/README.md](examples/README.md) for advanced usage (including distr
> [!TIP]
> Use `llamafactory-cli help` to show help information.
>
> Read [FAQs](https://github.com/hiyouga/LLaMA-Factory/issues/4614) first if you encounter any problems.
> Read [FAQs](https://github.com/hiyouga/LlamaFactory/issues/4614) first if you encounter any problems.
### Fine-Tuning with LLaMA Board GUI (powered by [Gradio](https://github.com/gradio-app/gradio))
@@ -701,28 +701,26 @@ docker compose up -d
docker compose exec llamafactory bash
```
For Ascend NPU users (A2 with Ubuntu by default):
For Ascend NPU users:
```bash
cd docker/docker-npu/
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:
# A2 with Ubuntu
docker compose --profile a2-ubuntu up -d
docker compose --profile a2-ubuntu exec llamafactory-a2-ubuntu bash
```bash
# A3 with Ubuntu
docker compose --profile a3 up -d llamafactory-a3-ubuntu
docker compose exec llamafactory-a3-ubuntu bash
docker compose --profile a3-ubuntu up -d
docker compose --profile a3-ubuntu exec llamafactory-a3-ubuntu bash
# A2 with openEuler
docker compose --profile openeuler up -d llamafactory-a2-openeuler
docker compose exec llamafactory-a2-openeuler bash
docker compose --profile a2-openeuler up -d
docker compose --profile a2-openeuler 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
docker compose --profile a3-openeuler up -d
docker compose --profile a3-openeuler exec llamafactory-a3-openeuler bash
```
For AMD ROCm users:
@@ -864,7 +862,7 @@ When launching training tasks, you can log in to SwanLab in three ways:
2. Set the environment variable `SWANLAB_API_KEY` to your [API key](https://swanlab.cn/settings).
3. Use the `swanlab login` command to complete the login.
## Projects using LLaMA Factory
## Projects using LlamaFactory
If you have a project that should be incorporated, please contact via email or create a pull request.
@@ -944,7 +942,7 @@ If you have a project that should be incorporated, please contact via email or c
1. Sun et al. LAMBDA: A Large Model Based Data Agent. 2024. [[arxiv]](https://arxiv.org/abs/2407.17535)
1. Zhu et al. CollectiveSFT: Scaling Large Language Models for Chinese Medical Benchmark with Collective Instructions in Healthcare. 2024. [[arxiv]](https://arxiv.org/abs/2407.19705)
1. Yu et al. Correcting Negative Bias in Large Language Models through Negative Attention Score Alignment. 2024. [[arxiv]](https://arxiv.org/abs/2408.00137)
1. Xie et al. The Power of Personalized Datasets: Advancing Chinese Composition Writing for Elementary School through Targeted Model Fine-Tuning. IALP 2024. [[paper]](https://www.asianlp.sg/conferences/ialp2024/proceedings/papers/IALP2024_P055.pdf)
1. Xie et al. The Power of Personalized Datasets: Advancing Chinese Composition Writing for Elementary School through Targeted Model Fine-Tuning. IALP 2024. [[paper]](https://doi.org/10.1142/S2717554524500176)
1. Liu et al. Instruct-Code-Llama: Improving Capabilities of Language Model in Competition Level Code Generation by Online Judge Feedback. ICIC 2024. [[paper]](https://link.springer.com/chapter/10.1007/978-981-97-5669-8_11)
1. Wang et al. Cybernetic Sentinels: Unveiling the Impact of Safety Data Selection on Model Security in Supervised Fine-Tuning. ICIC 2024. [[paper]](https://link.springer.com/chapter/10.1007/978-981-97-5669-8_23)
1. Xia et al. Understanding the Performance and Estimating the Cost of LLM Fine-Tuning. 2024. [[arxiv]](https://arxiv.org/abs/2408.04693)
@@ -962,7 +960,7 @@ If you have a project that should be incorporated, please contact via email or c
1. **[Chinese-LLaVA-Med](https://github.com/BUAADreamer/Chinese-LLaVA-Med)**: A multimodal large language model specialized in Chinese medical domain, based on LLaVA-1.5-7B.
1. **[AutoRE](https://github.com/THUDM/AutoRE)**: A document-level relation extraction system based on large language models.
1. **[NVIDIA RTX AI Toolkit](https://github.com/NVIDIA/RTX-AI-Toolkit)**: SDKs for fine-tuning LLMs on Windows PC for NVIDIA RTX.
1. **[LazyLLM](https://github.com/LazyAGI/LazyLLM)**: An easy and lazy way for building multi-agent LLMs applications and supports model fine-tuning via LLaMA Factory.
1. **[LazyLLM](https://github.com/LazyAGI/LazyLLM)**: An easy and lazy way for building multi-agent LLMs applications and supports model fine-tuning via LlamaFactory.
1. **[RAG-Retrieval](https://github.com/NLPJCL/RAG-Retrieval)**: A full pipeline for RAG retrieval model fine-tuning, inference, and distillation. [[blog]](https://zhuanlan.zhihu.com/p/987727357)
1. **[360-LLaMA-Factory](https://github.com/Qihoo360/360-LLaMA-Factory)**: A modified library that supports long sequence SFT & DPO using ring attention.
1. **[Sky-T1](https://novasky-ai.github.io/posts/sky-t1/)**: An o1-like model fine-tuned by NovaSky AI with very small cost.
@@ -974,7 +972,7 @@ If you have a project that should be incorporated, please contact via email or c
This repository is licensed under the [Apache-2.0 License](LICENSE).
Please follow the model licenses to use the corresponding model weights: [BLOOM](https://huggingface.co/spaces/bigscience/license) / [DeepSeek](https://github.com/deepseek-ai/DeepSeek-LLM/blob/main/LICENSE-MODEL) / [Falcon](https://huggingface.co/tiiuae/falcon-180B/blob/main/LICENSE.txt) / [Gemma](https://ai.google.dev/gemma/terms) / [GLM-4](https://huggingface.co/THUDM/glm-4-9b/blob/main/LICENSE) / [GPT-2](https://github.com/openai/gpt-2/blob/master/LICENSE) / [Granite](LICENSE) / [InternLM](https://github.com/InternLM/InternLM#license) / [Llama](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) / [Llama 2](https://ai.meta.com/llama/license/) / [Llama 3](https://llama.meta.com/llama3/license/) / [Llama 4](https://github.com/meta-llama/llama-models/blob/main/models/llama4/LICENSE) / [MiniCPM](https://github.com/OpenBMB/MiniCPM/blob/main/MiniCPM%20Model%20License.md) / [Mistral/Mixtral/Pixtral](LICENSE) / [Phi-3/Phi-4](https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/blob/main/LICENSE) / [Qwen](https://github.com/QwenLM/Qwen/blob/main/Tongyi%20Qianwen%20LICENSE%20AGREEMENT) / [StarCoder 2](https://huggingface.co/spaces/bigcode/bigcode-model-license-agreement) / [TeleChat2](https://huggingface.co/Tele-AI/telechat-7B/blob/main/TeleChat%E6%A8%A1%E5%9E%8B%E7%A4%BE%E5%8C%BA%E8%AE%B8%E5%8F%AF%E5%8D%8F%E8%AE%AE.pdf) / [Yuan 2](https://github.com/IEIT-Yuan/Yuan-2.0/blob/main/LICENSE-Yuan)
Please follow the model licenses to use the corresponding model weights: [BLOOM](https://huggingface.co/spaces/bigscience/license) / [DeepSeek](https://github.com/deepseek-ai/DeepSeek-LLM/blob/main/LICENSE-MODEL) / [Falcon](https://huggingface.co/tiiuae/falcon-180B/blob/main/LICENSE.txt) / [Gemma](https://ai.google.dev/gemma/terms) / [GLM-4](https://huggingface.co/THUDM/glm-4-9b/blob/main/LICENSE) / [GPT-2](https://github.com/openai/gpt-2/blob/master/LICENSE) / [Granite](LICENSE) / [InternLM](https://github.com/InternLM/InternLM#license) / [Llama](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) / [Llama 2](https://ai.meta.com/llama/license/) / [Llama 3](https://llama.meta.com/llama3/license/) / [Llama 4](https://github.com/meta-llama/llama-models/blob/main/models/llama4/LICENSE) / [MiniCPM](https://github.com/OpenBMB/MiniCPM/blob/main/LICENSE) / [Mistral/Mixtral/Pixtral](LICENSE) / [Phi-3/Phi-4](https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/blob/main/LICENSE) / [Qwen](https://github.com/QwenLM/Qwen/blob/main/Tongyi%20Qianwen%20LICENSE%20AGREEMENT) / [StarCoder 2](https://huggingface.co/spaces/bigcode/bigcode-model-license-agreement) / [TeleChat2](https://huggingface.co/Tele-AI/telechat-7B/blob/main/TeleChat%E6%A8%A1%E5%9E%8B%E7%A4%BE%E5%8C%BA%E8%AE%B8%E5%8F%AF%E5%8D%8F%E8%AE%AE.pdf) / [Yuan 2](https://github.com/IEIT-Yuan/Yuan-2.0/blob/main/LICENSE-Yuan)
## Citation

View File

@@ -1,9 +1,9 @@
![# LLaMA Factory](assets/logo.png)
![# LlamaFactory](assets/logo.png)
[![GitHub Repo stars](https://img.shields.io/github/stars/hiyouga/LLaMA-Factory?style=social)](https://github.com/hiyouga/LLaMA-Factory/stargazers)
[![GitHub last commit](https://img.shields.io/github/last-commit/hiyouga/LLaMA-Factory)](https://github.com/hiyouga/LLaMA-Factory/commits/main)
[![GitHub contributors](https://img.shields.io/github/contributors/hiyouga/LLaMA-Factory?color=orange)](https://github.com/hiyouga/LLaMA-Factory/graphs/contributors)
[![GitHub workflow](https://github.com/hiyouga/LLaMA-Factory/actions/workflows/tests.yml/badge.svg)](https://github.com/hiyouga/LLaMA-Factory/actions/workflows/tests.yml)
[![GitHub Repo stars](https://img.shields.io/github/stars/hiyouga/LlamaFactory?style=social)](https://github.com/hiyouga/LlamaFactory/stargazers)
[![GitHub last commit](https://img.shields.io/github/last-commit/hiyouga/LlamaFactory)](https://github.com/hiyouga/LlamaFactory/commits/main)
[![GitHub contributors](https://img.shields.io/github/contributors/hiyouga/LlamaFactory?color=orange)](https://github.com/hiyouga/LlamaFactory/graphs/contributors)
[![GitHub workflow](https://github.com/hiyouga/LlamaFactory/actions/workflows/tests.yml/badge.svg)](https://github.com/hiyouga/LlamaFactory/actions/workflows/tests.yml)
[![PyPI](https://img.shields.io/pypi/v/llamafactory)](https://pypi.org/project/llamafactory/)
[![Citation](https://img.shields.io/badge/citation-1000+-green)](https://scholar.google.com/scholar?cites=12620864006390196564)
[![Docker Pulls](https://img.shields.io/docker/pulls/hiyouga/llamafactory)](https://hub.docker.com/r/hiyouga/llamafactory/tags)
@@ -86,7 +86,7 @@ https://github.com/user-attachments/assets/43b700c6-a178-41db-b1f8-8190a5d3fcfc
- [数据集](#数据集)
- [软硬件依赖](#软硬件依赖)
- [如何使用](#如何使用)
- [安装 LLaMA Factory](#安装-llama-factory)
- [安装 LlamaFactory](#安装-llamafactory)
- [数据准备](#数据准备)
- [快速开始](#快速开始)
- [LLaMA Board 可视化微调](#llama-board-可视化微调由-gradio-驱动)
@@ -96,7 +96,7 @@ https://github.com/user-attachments/assets/43b700c6-a178-41db-b1f8-8190a5d3fcfc
- [从魔乐社区下载](#从魔乐社区下载)
- [使用 W&B 面板](#使用-wb-面板)
- [使用 SwanLab 面板](#使用-swanlab-面板)
- [使用了 LLaMA Factory 的项目](#使用了-llama-factory-的项目)
- [使用了 LlamaFactory 的项目](#使用了-llamafactory-的项目)
- [协议](#协议)
- [引用](#引用)
- [致谢](#致谢)
@@ -122,35 +122,35 @@ https://github.com/user-attachments/assets/43b700c6-a178-41db-b1f8-8190a5d3fcfc
## 官方博客
> [!TIP]
> 我们现在拥有了 LLaMA Factory 的专属博客!
> 我们现在拥有了 LlamaFactory 的专属博客!
>
> 网站地址https://blog.llamafactory.net/
- 💡 [KTransformers Fine-Tuning × LLaMA Factory: 用2张4090级的GPU+CPU 微调 1000B规模的超大模型](https://swcil84qspu.feishu.cn/wiki/Z1sSwb2poijybxkyPEkcDG6enVc) (中文)
- 💡 [Easy Dataset × LLaMA Factory: 让大模型高效学习领域知识](https://buaa-act.feishu.cn/wiki/KY9xwTGs1iqHrRkjXBwcZP9WnL9)(中文)
- 💡 [DataFlow × LLaMA Factory: 利用数据准备流水线产出高质量数据训练 LLM](https://wcny4qa9krto.feishu.cn/wiki/LlMxweUAJimrmykRD5qcGuswnHd)(中文)| [English](https://wcny4qa9krto.feishu.cn/wiki/LWkkwTDBfiiRKqkDSvucG6yjnbW)
- 💡 [DataFlex × LLaMA Factory: 构建在 LLaMA-Factory 之上的以数据为中心的动态训练系统](https://wcny4qa9krto.feishu.cn/wiki/H2A9wSsbCinzavkT2oyc2C5Vn0e)(中文)| [English](https://wcny4qa9krto.feishu.cn/wiki/OlREwPQWdi9K6ZkJNHIciLhtnkv)
- [基于 LLaMA-Factory 和 EasyR1 打造一站式无代码大模型强化学习和部署平台 LLM Model Hub](https://aws.amazon.com/cn/blogs/china/building-llm-model-hub-based-on-llamafactory-and-easyr1/)(中文)
- [通过亚马逊 SageMaker HyperPod 上的 LLaMA-Factory 增强多模态模型银行文档的视觉信息提取](https://aws.amazon.com/cn/blogs/machine-learning/how-apoidea-group-enhances-visual-information-extraction-from-banking-documents-with-multimodal-models-using-llama-factory-on-amazon-sagemaker-hyperpod/)(英文)
- 💡 [KTransformers Fine-Tuning × LlamaFactory: 用2张4090级的GPU+CPU 微调 1000B规模的超大模型](https://swcil84qspu.feishu.cn/wiki/Z1sSwb2poijybxkyPEkcDG6enVc) (中文)
- 💡 [Easy Dataset × LlamaFactory: 让大模型高效学习领域知识](https://buaa-act.feishu.cn/wiki/KY9xwTGs1iqHrRkjXBwcZP9WnL9)(中文)
- 💡 [DataFlow × LlamaFactory: 利用数据准备流水线产出高质量数据训练 LLM](https://wcny4qa9krto.feishu.cn/wiki/LlMxweUAJimrmykRD5qcGuswnHd)(中文)| [English](https://wcny4qa9krto.feishu.cn/wiki/LWkkwTDBfiiRKqkDSvucG6yjnbW)
- 💡 [DataFlex × LlamaFactory: 构建在 LlamaFactory 之上的以数据为中心的动态训练系统](https://wcny4qa9krto.feishu.cn/wiki/H2A9wSsbCinzavkT2oyc2C5Vn0e)(中文)| [English](https://wcny4qa9krto.feishu.cn/wiki/OlREwPQWdi9K6ZkJNHIciLhtnkv)
- [基于 LlamaFactory 和 EasyR1 打造一站式无代码大模型强化学习和部署平台 LLM Model Hub](https://aws.amazon.com/cn/blogs/china/building-llm-model-hub-based-on-llamafactory-and-easyr1/)(中文)
- [通过亚马逊 SageMaker HyperPod 上的 LlamaFactory 增强多模态模型银行文档的视觉信息提取](https://aws.amazon.com/cn/blogs/machine-learning/how-apoidea-group-enhances-visual-information-extraction-from-banking-documents-with-multimodal-models-using-llama-factory-on-amazon-sagemaker-hyperpod/)(英文)
<details><summary>全部博客</summary>
- [LLaMA Factory微调 DeepSeek-R1-Distill-Qwen-7B 模型实现新闻标题分类器](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory_deepseek_r1_distill_7b)(中文)
- [基于 Amazon SageMaker 和 LLaMA-Factory 打造一站式无代码模型微调部署平台 Model Hub](https://aws.amazon.com/cn/blogs/china/a-one-stop-code-free-model-fine-tuning-deployment-platform-based-on-sagemaker-and-llama-factory/)(中文)
- [LLaMA Factory 多模态微调实践:微调 Qwen2-VL 构建文旅大模型](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory_qwen2vl)(中文)
- [LLaMA Factory微调 Llama3 模型实现角色扮演](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory)(中文)
- [LlamaFactory微调 DeepSeek-R1-Distill-Qwen-7B 模型实现新闻标题分类器](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory_deepseek_r1_distill_7b)(中文)
- [基于 Amazon SageMaker 和 LlamaFactory 打造一站式无代码模型微调部署平台 Model Hub](https://aws.amazon.com/cn/blogs/china/a-one-stop-code-free-model-fine-tuning-deployment-platform-based-on-sagemaker-and-llama-factory/)(中文)
- [LlamaFactory 多模态微调实践:微调 Qwen2-VL 构建文旅大模型](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory_qwen2vl)(中文)
- [LlamaFactory微调 Llama3 模型实现角色扮演](https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory)(中文)
</details>
## 更新日志
[25/10/26] 我们支持了Megatron-core作为训练后端和适配了[**mcore_adapter**](https://github.com/alibaba/ROLL/tree/main/mcore_adapter)。查看[PR #9237](https://github.com/hiyouga/LLaMA-Factory/pull/9237)以使用。
[25/10/26] 我们支持了Megatron-core作为训练后端和适配了[**mcore_adapter**](https://github.com/alibaba/ROLL/tree/main/mcore_adapter)。查看[PR #9237](https://github.com/hiyouga/LlamaFactory/pull/9237)以使用。
[25/08/22] 我们支持了 **[OFT](https://arxiv.org/abs/2306.07280)** 和 **[OFTv2](https://arxiv.org/abs/2506.19847)** 模型的微调。查看 [examples](examples/README.md) 以使用。
[25/08/20] 我们支持了 **[Intern-S1-mini](https://huggingface.co/internlm/Intern-S1-mini)** 模型的微调。查看 [PR #8976](https://github.com/hiyouga/LLaMA-Factory/pull/8976) 以使用。
[25/08/20] 我们支持了 **[Intern-S1-mini](https://huggingface.co/internlm/Intern-S1-mini)** 模型的微调。查看 [PR #8976](https://github.com/hiyouga/LlamaFactory/pull/8976) 以使用。
[25/08/06] 我们支持了 **[GPT-OSS](https://github.com/openai/gpt-oss)** 模型的微调。查看 [PR #8826](https://github.com/hiyouga/LLaMA-Factory/pull/8826) 以使用。
[25/08/06] 我们支持了 **[GPT-OSS](https://github.com/openai/gpt-oss)** 模型的微调。查看 [PR #8826](https://github.com/hiyouga/LlamaFactory/pull/8826) 以使用。
<details><summary>展开日志</summary>
@@ -160,13 +160,13 @@ https://github.com/user-attachments/assets/43b700c6-a178-41db-b1f8-8190a5d3fcfc
[25/04/21] 我们支持了 **[Muon](https://github.com/KellerJordan/Muon)** 优化器。详细用法请参照 [examples](examples/README_zh.md)。感谢 [@tianshijing](https://github.com/tianshijing) 的 PR。
[25/04/16] 我们支持了 **[InternVL3](https://huggingface.co/OpenGVLab/InternVL3-8B)** 模型的微调。查看 [PR #7258](https://github.com/hiyouga/LLaMA-Factory/pull/7258) 以使用。
[25/04/16] 我们支持了 **[InternVL3](https://huggingface.co/OpenGVLab/InternVL3-8B)** 模型的微调。查看 [PR #7258](https://github.com/hiyouga/LlamaFactory/pull/7258) 以使用。
[25/04/14] 我们支持了 **[GLM-Z1](https://huggingface.co/THUDM/GLM-Z1-9B-0414)** 和 **[Kimi-VL](https://huggingface.co/moonshotai/Kimi-VL-A3B-Instruct)** 模型的微调。
[25/04/06] 我们支持了 **[Llama 4](https://ai.meta.com/blog/llama-4-multimodal-intelligence/)** 模型的微调。查看 [PR #7611](https://github.com/hiyouga/LLaMA-Factory/pull/7611) 以使用。
[25/04/06] 我们支持了 **[Llama 4](https://ai.meta.com/blog/llama-4-multimodal-intelligence/)** 模型的微调。查看 [PR #7611](https://github.com/hiyouga/LlamaFactory/pull/7611) 以使用。
[25/03/31] 我们支持了 **[Qwen2.5 Omni](https://qwenlm.github.io/blog/qwen2.5-omni/)** 模型的微调。查看 [PR #7537](https://github.com/hiyouga/LLaMA-Factory/pull/7537) 以使用。
[25/03/31] 我们支持了 **[Qwen2.5 Omni](https://qwenlm.github.io/blog/qwen2.5-omni/)** 模型的微调。查看 [PR #7537](https://github.com/hiyouga/LlamaFactory/pull/7537) 以使用。
[25/03/15] 我们支持了 **[SGLang](https://github.com/sgl-project/sglang)** 推理后端,请使用 `infer_backend: sglang` 启用。
@@ -214,17 +214,17 @@ https://github.com/user-attachments/assets/43b700c6-a178-41db-b1f8-8190a5d3fcfc
[24/05/18] 我们支持了 **[KTO](https://arxiv.org/abs/2402.01306)** 偏好对齐算法。详细用法请参照 [examples](examples/README_zh.md)。
[24/05/14] 我们支持了昇腾 NPU 设备的训练和推理。详情请查阅[安装](#安装-llama-factory)部分。
[24/05/14] 我们支持了昇腾 NPU 设备的训练和推理。详情请查阅[安装](#安装-llamafactory)部分。
[24/04/26] 我们支持了多模态模型 **LLaVA-1.5** 的微调。详细用法请参照 [examples](examples/README_zh.md)。
[24/04/22] 我们提供了在免费 T4 GPU 上微调 Llama-3 模型的 **[Colab 笔记本](https://colab.research.google.com/drive/1d5KQtbemerlSDSxZIfAaWXhKr30QypiK?usp=sharing)**。Hugging Face 社区公开了两个利用 LLaMA Factory 微调的 Llama-3 模型,详情请见 [Llama3-8B-Chinese-Chat](https://huggingface.co/shenzhi-wang/Llama3-8B-Chinese-Chat) 和 [Llama3-Chinese](https://huggingface.co/zhichen/Llama3-Chinese)。
[24/04/22] 我们提供了在免费 T4 GPU 上微调 Llama-3 模型的 **[Colab 笔记本](https://colab.research.google.com/drive/1d5KQtbemerlSDSxZIfAaWXhKr30QypiK?usp=sharing)**。Hugging Face 社区公开了两个利用 LlamaFactory 微调的 Llama-3 模型,详情请见 [Llama3-8B-Chinese-Chat](https://huggingface.co/shenzhi-wang/Llama3-8B-Chinese-Chat) 和 [Llama3-Chinese](https://huggingface.co/zhichen/Llama3-Chinese)。
[24/04/21] 我们基于 [AstraMindAI 的仓库](https://github.com/astramind-ai/Mixture-of-depths)支持了 **[混合深度训练](https://arxiv.org/abs/2404.02258)**。详细用法请参照 [examples](examples/README_zh.md)。
[24/04/16] 我们支持了 **[BAdam](https://arxiv.org/abs/2404.02827)** 优化器。详细用法请参照 [examples](examples/README_zh.md)。
[24/04/16] 我们支持了 **[unsloth](https://github.com/unslothai/unsloth)** 的长序列训练24GB 可训练 Llama-2-7B-56k。该方法相比 FlashAttention-2 提供了 **117%** 的训练速度和 **50%** 的显存节约。更多数据请见[此页面](https://github.com/hiyouga/LLaMA-Factory/wiki/Performance-comparison)。
[24/04/16] 我们支持了 **[unsloth](https://github.com/unslothai/unsloth)** 的长序列训练24GB 可训练 Llama-2-7B-56k。该方法相比 FlashAttention-2 提供了 **117%** 的训练速度和 **50%** 的显存节约。更多数据请见[此页面](https://github.com/hiyouga/LlamaFactory/wiki/Performance-comparison)。
[24/03/31] 我们支持了 **[ORPO](https://arxiv.org/abs/2403.07691)**。详细用法请参照 [examples](examples/README_zh.md)。
@@ -242,11 +242,11 @@ https://github.com/user-attachments/assets/43b700c6-a178-41db-b1f8-8190a5d3fcfc
[24/02/15] 我们支持了 [LLaMA Pro](https://github.com/TencentARC/LLaMA-Pro) 提出的**块扩展**方法。详细用法请参照 [examples](examples/README_zh.md)。
[24/02/05] Qwen1.5Qwen2 测试版)系列模型已在 LLaMA-Factory 中实现微调支持。详情请查阅该[博客页面](https://qwenlm.github.io/zh/blog/qwen1.5/)。
[24/02/05] Qwen1.5Qwen2 测试版)系列模型已在 LlamaFactory 中实现微调支持。详情请查阅该[博客页面](https://qwenlm.github.io/zh/blog/qwen1.5/)。
[24/01/18] 我们针对绝大多数模型实现了 **Agent 微调**,微调时指定 `dataset: glaive_toolcall_zh` 即可使模型获得工具调用能力。
[23/12/23] 我们针对 LLaMA, Mistral 和 Yi 模型支持了 **[unsloth](https://github.com/unslothai/unsloth)** 的 LoRA 训练加速。请使用 `use_unsloth: true` 参数启用 unsloth 优化。该方法可提供 **170%** 的训练速度,详情请查阅[此页面](https://github.com/hiyouga/LLaMA-Factory/wiki/Performance-comparison)。
[23/12/23] 我们针对 LLaMA, Mistral 和 Yi 模型支持了 **[unsloth](https://github.com/unslothai/unsloth)** 的 LoRA 训练加速。请使用 `use_unsloth: true` 参数启用 unsloth 优化。该方法可提供 **170%** 的训练速度,详情请查阅[此页面](https://github.com/hiyouga/LlamaFactory/wiki/Performance-comparison)。
[23/12/12] 我们支持了微调最新的混合专家模型 **[Mixtral 8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1)**。硬件需求请查阅[此处](#硬件依赖)。
@@ -281,7 +281,7 @@ https://github.com/user-attachments/assets/43b700c6-a178-41db-b1f8-8190a5d3fcfc
</details>
> [!TIP]
> 如果您无法使用最新的功能,请尝试重新拉取代码并再次安装 LLaMA-Factory。
> 如果您无法使用最新的功能,请尝试重新拉取代码并再次安装 LlamaFactory。
## 模型
@@ -516,7 +516,7 @@ huggingface-cli login
## 如何使用
### 安装 LLaMA Factory
### 安装 LlamaFactory
> [!IMPORTANT]
> 此步骤为必需。
@@ -605,22 +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/multibackend/npu/npu_installation.html)。
在昇腾 NPU 设备上安装 LlamaFactory 时,请使用 Python 3.12使用 `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-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
docker pull hiyouga/llamafactory:latest-910b-ubuntu
docker pull hiyouga/llamafactory:latest-a3-ubuntu
docker pull hiyouga/llamafactory:latest-910b-openeuler
docker pull hiyouga/llamafactory:latest-a3-openeuler
# quay.io
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
docker pull quay.io/ascend/llamafactory:latest-910b-ubuntu
docker pull quay.io/ascend/llamafactory:latest-a3-ubuntu
docker pull quay.io/ascend/llamafactory:latest-910b-openeuler
docker pull quay.io/ascend/llamafactory:latest-a3-openeuler
```
#### 安装 BitsAndBytes
@@ -683,7 +683,7 @@ llamafactory-cli export examples/merge_lora/qwen3_lora_sft.yaml
> [!TIP]
> 使用 `llamafactory-cli help` 显示帮助信息。
>
> 遇到报错请先看[常见问题](https://github.com/hiyouga/LLaMA-Factory/issues/4614)。
> 遇到报错请先看[常见问题](https://github.com/hiyouga/LlamaFactory/issues/4614)。
### LLaMA Board 可视化微调(由 [Gradio](https://github.com/gradio-app/gradio) 驱动)
@@ -701,28 +701,26 @@ docker compose up -d
docker compose exec llamafactory bash
```
昇腾 NPU 用户(默认使用 A2 和 Ubuntu
昇腾 NPU 用户:
```bash
cd docker/docker-npu/
docker compose up -d llamafactory-a2-ubuntu
docker compose exec llamafactory-a2-ubuntu bash
```
其他 NPU 组合可以通过对应的 profile 和服务启动:
# A2 + Ubuntu
docker compose --profile a2-ubuntu up -d
docker compose --profile a2-ubuntu exec llamafactory-a2-ubuntu bash
```bash
# A3 + Ubuntu
docker compose --profile a3 up -d llamafactory-a3-ubuntu
docker compose exec llamafactory-a3-ubuntu bash
docker compose --profile a3-ubuntu up -d
docker compose --profile a3-ubuntu exec llamafactory-a3-ubuntu bash
# A2 + openEuler
docker compose --profile openeuler up -d llamafactory-a2-openeuler
docker compose exec llamafactory-a2-openeuler bash
docker compose --profile a2-openeuler up -d
docker compose --profile a2-openeuler exec llamafactory-a2-openeuler bash
# A3 + openEuler
docker compose --profile a3-openeuler up -d llamafactory-a3-openeuler
docker compose exec llamafactory-a3-openeuler bash
docker compose --profile a3-openeuler up -d
docker compose --profile a3-openeuler exec llamafactory-a3-openeuler bash
```
AMD ROCm 用户:
@@ -757,7 +755,6 @@ docker exec -it llamafactory bash
```bash
docker build -f ./docker/docker-npu/Dockerfile \
--build-arg PIP_INDEX=https://pypi.org/simple \
--build-arg EXTRAS=torch-npu,metrics \
-t llamafactory:latest .
docker run -dit --ipc=host \
@@ -867,7 +864,7 @@ swanlab_run_name: test_run # 可选
方式二:将环境变量 `SWANLAB_API_KEY` 设置为你的 [API 密钥](https://swanlab.cn/settings)。
方式三:启动前使用 `swanlab login` 命令完成登录。
## 使用了 LLaMA Factory 的项目
## 使用了 LlamaFactory 的项目
如果您有项目希望添加至下述列表,请通过邮件联系或者创建一个 PR。
@@ -947,7 +944,7 @@ swanlab_run_name: test_run # 可选
1. Sun et al. LAMBDA: A Large Model Based Data Agent. 2024. [[arxiv]](https://arxiv.org/abs/2407.17535)
1. Zhu et al. CollectiveSFT: Scaling Large Language Models for Chinese Medical Benchmark with Collective Instructions in Healthcare. 2024. [[arxiv]](https://arxiv.org/abs/2407.19705)
1. Yu et al. Correcting Negative Bias in Large Language Models through Negative Attention Score Alignment. 2024. [[arxiv]](https://arxiv.org/abs/2408.00137)
1. Xie et al. The Power of Personalized Datasets: Advancing Chinese Composition Writing for Elementary School through Targeted Model Fine-Tuning. IALP 2024. [[paper]](https://www.asianlp.sg/conferences/ialp2024/proceedings/papers/IALP2024_P055.pdf)
1. Xie et al. The Power of Personalized Datasets: Advancing Chinese Composition Writing for Elementary School through Targeted Model Fine-Tuning. IALP 2024. [[paper]](https://doi.org/10.1142/S2717554524500176)
1. Liu et al. Instruct-Code-Llama: Improving Capabilities of Language Model in Competition Level Code Generation by Online Judge Feedback. ICIC 2024. [[paper]](https://link.springer.com/chapter/10.1007/978-981-97-5669-8_11)
1. Wang et al. Cybernetic Sentinels: Unveiling the Impact of Safety Data Selection on Model Security in Supervised Fine-Tuning. ICIC 2024. [[paper]](https://link.springer.com/chapter/10.1007/978-981-97-5669-8_23)
1. Xia et al. Understanding the Performance and Estimating the Cost of LLM Fine-Tuning. 2024. [[arxiv]](https://arxiv.org/abs/2408.04693)
@@ -964,7 +961,7 @@ swanlab_run_name: test_run # 可选
1. **[Chinese-LLaVA-Med](https://github.com/BUAADreamer/Chinese-LLaVA-Med)**:中文多模态医学大模型,基于 LLaVA-1.5-7B 在中文多模态医疗数据上微调而得。
1. **[AutoRE](https://github.com/THUDM/AutoRE)**:基于大语言模型的文档级关系抽取系统。
1. **[NVIDIA RTX AI Toolkit](https://github.com/NVIDIA/RTX-AI-Toolkit)**:在 Windows 主机上利用英伟达 RTX 设备进行大型语言模型微调的开发包。
1. **[LazyLLM](https://github.com/LazyAGI/LazyLLM)**:一个低代码构建多 Agent 大模型应用的开发工具,支持基于 LLaMA Factory 的模型微调.
1. **[LazyLLM](https://github.com/LazyAGI/LazyLLM)**:一个低代码构建多 Agent 大模型应用的开发工具,支持基于 LlamaFactory 的模型微调.
1. **[RAG-Retrieval](https://github.com/NLPJCL/RAG-Retrieval)**:一个全链路 RAG 检索模型微调、推理和蒸馏代码库。[[blog]](https://zhuanlan.zhihu.com/p/987727357)
1. **[360-LLaMA-Factory](https://github.com/Qihoo360/360-LLaMA-Factory)**:一个魔改后的代码库,通过 Ring Attention 支持长序列的 SFT 和 DPO 训练。
1. **[Sky-T1](https://novasky-ai.github.io/posts/sky-t1/)**:由 NovaSky AI 微调的低成本类 o1 长推理模型。
@@ -976,7 +973,7 @@ swanlab_run_name: test_run # 可选
本仓库的代码依照 [Apache-2.0](LICENSE) 协议开源。
使用模型权重时,请遵循对应的模型协议:[BLOOM](https://huggingface.co/spaces/bigscience/license)/ [DeepSeek](https://github.com/deepseek-ai/DeepSeek-LLM/blob/main/LICENSE-MODEL) / [Falcon](https://huggingface.co/tiiuae/falcon-180B/blob/main/LICENSE.txt) / [Gemma](https://ai.google.dev/gemma/terms) / [GLM-4](https://huggingface.co/THUDM/glm-4-9b/blob/main/LICENSE) / [GPT-2](https://github.com/openai/gpt-2/blob/master/LICENSE) / [Granite](LICENSE) / [InternLM](https://github.com/InternLM/InternLM#license) / [Llama](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) / [Llama 2](https://ai.meta.com/llama/license/) / [Llama 3](https://llama.meta.com/llama3/license/) / [Llama 4](https://github.com/meta-llama/llama-models/blob/main/models/llama4/LICENSE) / [MiniCPM](https://github.com/OpenBMB/MiniCPM/blob/main/MiniCPM%20Model%20License.md) / [Mistral/Mixtral/Pixtral](LICENSE) / [Phi-3/Phi-4](https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/blob/main/LICENSE) / [Qwen](https://github.com/QwenLM/Qwen/blob/main/Tongyi%20Qianwen%20LICENSE%20AGREEMENT) / [StarCoder 2](https://huggingface.co/spaces/bigcode/bigcode-model-license-agreement) / [TeleChat2](https://huggingface.co/Tele-AI/telechat-7B/blob/main/TeleChat%E6%A8%A1%E5%9E%8B%E7%A4%BE%E5%8C%BA%E8%AE%B8%E5%8F%AF%E5%8D%8F%E8%AE%AE.pdf) / [Yuan 2](https://github.com/IEIT-Yuan/Yuan-2.0/blob/main/LICENSE-Yuan)
使用模型权重时,请遵循对应的模型协议:[BLOOM](https://huggingface.co/spaces/bigscience/license)/ [DeepSeek](https://github.com/deepseek-ai/DeepSeek-LLM/blob/main/LICENSE-MODEL) / [Falcon](https://huggingface.co/tiiuae/falcon-180B/blob/main/LICENSE.txt) / [Gemma](https://ai.google.dev/gemma/terms) / [GLM-4](https://huggingface.co/THUDM/glm-4-9b/blob/main/LICENSE) / [GPT-2](https://github.com/openai/gpt-2/blob/master/LICENSE) / [Granite](LICENSE) / [InternLM](https://github.com/InternLM/InternLM#license) / [Llama](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) / [Llama 2](https://ai.meta.com/llama/license/) / [Llama 3](https://llama.meta.com/llama3/license/) / [Llama 4](https://github.com/meta-llama/llama-models/blob/main/models/llama4/LICENSE) / [MiniCPM](https://github.com/OpenBMB/MiniCPM/blob/main/LICENSE) / [Mistral/Mixtral/Pixtral](LICENSE) / [Phi-3/Phi-4](https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/blob/main/LICENSE) / [Qwen](https://github.com/QwenLM/Qwen/blob/main/Tongyi%20Qianwen%20LICENSE%20AGREEMENT) / [StarCoder 2](https://huggingface.co/spaces/bigcode/bigcode-model-license-agreement) / [TeleChat2](https://huggingface.co/Tele-AI/telechat-7B/blob/main/TeleChat%E6%A8%A1%E5%9E%8B%E7%A4%BE%E5%8C%BA%E8%AE%B8%E5%8F%AF%E5%8D%8F%E8%AE%AE.pdf) / [Yuan 2](https://github.com/IEIT-Yuan/Yuan-2.0/blob/main/LICENSE-Yuan)
## 引用

View File

@@ -0,0 +1,25 @@
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}, {"type": "text", "value": "Who are they?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "They're Kane and Gretzka from Bayern Munich."}]}, {"role": "user", "content": [{"type": "text", "value": "What are they doing?"}, {"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}]}, {"role": "assistant", "content": [{"type": "text", "value": "They are celebrating on the soccer field."}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/2.jpg"}, {"type": "text", "value": "Who is he?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "He's Thomas Muller from Bayern Munich."}]}, {"role": "user", "content": [{"type": "text", "value": "Why is he on the ground?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "Because he's sliding on his knees to celebrate."}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/3.jpg"}, {"type": "text", "value": "Please describe this image"}]}, {"role": "assistant", "content": [{"type": "text", "value": "Chinese astronaut Gui Haichao is giving a speech."}]}, {"role": "user", "content": [{"type": "text", "value": "What has he accomplished?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "He was appointed to be a payload specialist on Shenzhou 16 mission in June 2022, thus becoming the first Chinese civilian of Group 3 in space on 30 May 2023. He is responsible for the on-orbit operation of space science experimental payloads."}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}, {"type": "text", "value": "他们是谁?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他们是拜仁慕尼黑的凯恩和格雷茨卡。"}]}, {"role": "user", "content": [{"type": "text", "value": "他们在做什么?"}, {"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他们在足球场上庆祝。"}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/2.jpg"}, {"type": "text", "value": "他是谁?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他是来自拜仁慕尼黑的托马斯·穆勒。"}]}, {"role": "user", "content": [{"type": "text", "value": "他为什么在地上?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "因为他正在双膝跪地滑行庆祝。"}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/3.jpg"}, {"type": "text", "value": "请描述这张图片"}]}, {"role": "assistant", "content": [{"type": "text", "value": "中国宇航员桂海潮正在讲话。"}]}, {"role": "user", "content": [{"type": "text", "value": "他取得过哪些成就?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他于2022年6月被任命为神舟十六号任务的有效载荷专家从而成为2023年5月30日进入太空的首位平民宇航员。他负责在轨操作空间科学实验有效载荷。"}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}, {"type": "text", "value": "Who are they?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "They're Kane and Gretzka from Bayern Munich."}]}, {"role": "user", "content": [{"type": "text", "value": "What are they doing?"}, {"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}]}, {"role": "assistant", "content": [{"type": "text", "value": "They are celebrating on the soccer field."}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/2.jpg"}, {"type": "text", "value": "Who is he?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "He's Thomas Muller from Bayern Munich."}]}, {"role": "user", "content": [{"type": "text", "value": "Why is he on the ground?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "Because he's sliding on his knees to celebrate."}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/3.jpg"}, {"type": "text", "value": "Please describe this image"}]}, {"role": "assistant", "content": [{"type": "text", "value": "Chinese astronaut Gui Haichao is giving a speech."}]}, {"role": "user", "content": [{"type": "text", "value": "What has he accomplished?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "He was appointed to be a payload specialist on Shenzhou 16 mission in June 2022, thus becoming the first Chinese civilian of Group 3 in space on 30 May 2023. He is responsible for the on-orbit operation of space science experimental payloads."}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}, {"type": "text", "value": "他们是谁?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他们是拜仁慕尼黑的凯恩和格雷茨卡。"}]}, {"role": "user", "content": [{"type": "text", "value": "他们在做什么?"}, {"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他们在足球场上庆祝。"}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/2.jpg"}, {"type": "text", "value": "他是谁?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他是来自拜仁慕尼黑的托马斯·穆勒。"}]}, {"role": "user", "content": [{"type": "text", "value": "他为什么在地上?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "因为他正在双膝跪地滑行庆祝。"}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/3.jpg"}, {"type": "text", "value": "请描述这张图片"}]}, {"role": "assistant", "content": [{"type": "text", "value": "中国宇航员桂海潮正在讲话。"}]}, {"role": "user", "content": [{"type": "text", "value": "他取得过哪些成就?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他于2022年6月被任命为神舟十六号任务的有效载荷专家从而成为2023年5月30日进入太空的首位平民宇航员。他负责在轨操作空间科学实验有效载荷。"}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}, {"type": "text", "value": "Who are they?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "They're Kane and Gretzka from Bayern Munich."}]}, {"role": "user", "content": [{"type": "text", "value": "What are they doing?"}, {"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}]}, {"role": "assistant", "content": [{"type": "text", "value": "They are celebrating on the soccer field."}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/2.jpg"}, {"type": "text", "value": "Who is he?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "He's Thomas Muller from Bayern Munich."}]}, {"role": "user", "content": [{"type": "text", "value": "Why is he on the ground?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "Because he's sliding on his knees to celebrate."}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/3.jpg"}, {"type": "text", "value": "Please describe this image"}]}, {"role": "assistant", "content": [{"type": "text", "value": "Chinese astronaut Gui Haichao is giving a speech."}]}, {"role": "user", "content": [{"type": "text", "value": "What has he accomplished?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "He was appointed to be a payload specialist on Shenzhou 16 mission in June 2022, thus becoming the first Chinese civilian of Group 3 in space on 30 May 2023. He is responsible for the on-orbit operation of space science experimental payloads."}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}, {"type": "text", "value": "他们是谁?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他们是拜仁慕尼黑的凯恩和格雷茨卡。"}]}, {"role": "user", "content": [{"type": "text", "value": "他们在做什么?"}, {"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他们在足球场上庆祝。"}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/2.jpg"}, {"type": "text", "value": "他是谁?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他是来自拜仁慕尼黑的托马斯·穆勒。"}]}, {"role": "user", "content": [{"type": "text", "value": "他为什么在地上?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "因为他正在双膝跪地滑行庆祝。"}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/3.jpg"}, {"type": "text", "value": "请描述这张图片"}]}, {"role": "assistant", "content": [{"type": "text", "value": "中国宇航员桂海潮正在讲话。"}]}, {"role": "user", "content": [{"type": "text", "value": "他取得过哪些成就?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他于2022年6月被任命为神舟十六号任务的有效载荷专家从而成为2023年5月30日进入太空的首位平民宇航员。他负责在轨操作空间科学实验有效载荷。"}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}, {"type": "text", "value": "Who are they?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "They're Kane and Gretzka from Bayern Munich."}]}, {"role": "user", "content": [{"type": "text", "value": "What are they doing?"}, {"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}]}, {"role": "assistant", "content": [{"type": "text", "value": "They are celebrating on the soccer field."}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/2.jpg"}, {"type": "text", "value": "Who is he?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "He's Thomas Muller from Bayern Munich."}]}, {"role": "user", "content": [{"type": "text", "value": "Why is he on the ground?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "Because he's sliding on his knees to celebrate."}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/3.jpg"}, {"type": "text", "value": "Please describe this image"}]}, {"role": "assistant", "content": [{"type": "text", "value": "Chinese astronaut Gui Haichao is giving a speech."}]}, {"role": "user", "content": [{"type": "text", "value": "What has he accomplished?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "He was appointed to be a payload specialist on Shenzhou 16 mission in June 2022, thus becoming the first Chinese civilian of Group 3 in space on 30 May 2023. He is responsible for the on-orbit operation of space science experimental payloads."}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}, {"type": "text", "value": "他们是谁?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他们是拜仁慕尼黑的凯恩和格雷茨卡。"}]}, {"role": "user", "content": [{"type": "text", "value": "他们在做什么?"}, {"type": "image_url", "value": "data/mllm_demo_data/1.jpg"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他们在足球场上庆祝。"}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/2.jpg"}, {"type": "text", "value": "他是谁?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他是来自拜仁慕尼黑的托马斯·穆勒。"}]}, {"role": "user", "content": [{"type": "text", "value": "他为什么在地上?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "因为他正在双膝跪地滑行庆祝。"}]}]}
{"messages": [{"role": "user", "content": [{"type": "image_url", "value": "data/mllm_demo_data/3.jpg"}, {"type": "text", "value": "请描述这张图片"}]}, {"role": "assistant", "content": [{"type": "text", "value": "中国宇航员桂海潮正在讲话。"}]}, {"role": "user", "content": [{"type": "text", "value": "他取得过哪些成就?"}]}, {"role": "assistant", "content": [{"type": "text", "value": "他于2022年6月被任命为神舟十六号任务的有效载荷专家从而成为2023年5月30日进入太空的首位平民宇航员。他负责在轨操作空间科学实验有效载荷。"}]}]}

View File

@@ -0,0 +1,4 @@
multimodal_demo:
path: data/v1_multimodal_demo.jsonl
source: local

View File

@@ -1,6 +1,6 @@
# https://hub.docker.com/r/ascendai/cann/tags
ARG BASE_IMAGE=quay.io/ascend/cann:9.0.0-910b-ubuntu22.04-py3.11
ARG BASE_IMAGE=quay.io/ascend/cann:9.1.0-910b-ubuntu22.04-py3.12
FROM ${BASE_IMAGE}
# Installation arguments

View File

@@ -1,6 +1,6 @@
# LLaMA Factory for Ascend NPU
# LlamaFactory Image 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.
LlamaFactory Ascend NPU images are designed for Huawei Ascend Atlas NPUs and provide a ready-to-use LlamaFactory environment. Built on Ascend CANN container images, they include Python, PyTorch, TorchNPU, DeepSpeed, LlamaFactory, and other components.
For installation and troubleshooting details, see the [English NPU installation guide](https://llamafactory.readthedocs.io/en/latest/multibackend/npu/npu_installation.html).
@@ -11,77 +11,68 @@ For installation and troubleshooting details, see the [English NPU installation
- `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:
The following `latest` NPU image tags are available:
| Accelerator | Container OS | CANN base image |
| Hardware series | Operating system | Tag |
| --- | --- | --- |
| 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` |
| A2 | Ubuntu 22.04 | `latest-910b-ubuntu` |
| A3 | Ubuntu 22.04 | `latest-a3-ubuntu` |
| A2 | openEuler 24.03 | `latest-910b-openeuler` |
| A3 | openEuler 24.03 | `latest-a3-openeuler` |
## Image Contents and Intended Use
## Image Overview
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:
The image includes the following core components:
| Component | Version or source |
| Component | Version |
| --- | --- |
| 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` |
| CANN | `9.1.0` |
| Python | `3.12` |
| PyTorch | `2.10.0` |
| TorchNPU | `2.10.0.post2` |
| torchvision / torchaudio | `0.25.0` / `2.10.0` |
| Transformers | Latest compatible version at build time |
| Triton Ascend | `3.2.1` |
| DeepSpeed | `>=0.10.0,<=0.18.4` |
| LLaMA Factory | Installed from the repository build context |
| DeepSpeed | Latest compatible version at build time |
| LlamaFactory | 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
## Image Tags
Images use the following tag format:
NPU `latest` and release tags use different formats; the following rules do not apply to CUDA images.
Non-release builds reuse the following short tags. Each scheduled build updates the image referenced by the corresponding tag:
```text
<llamafactory-version>-cann<cann-version>-torch_npu<torch-npu-version>-<accelerator>-<os>-<python-version>
latest-<chip>-<os>
```
| Field | Values | Description |
| --- | --- | --- |
| `chip` | `910b` or `a3` | Ascend chip model supported by the image |
| `os` | `ubuntu` or `openeuler` | Container operating system family |
Release builds use full tags:
```text
<LlamaFactory-version>-cann<CANN-version>-torch_npu<TorchNPU-version>-<chip>-<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 |
| `LlamaFactory-version` | `0.9.5` | LlamaFactory release version |
| `CANN-version` | `9.1.0` | Parsed from the CANN base image tag |
| `TorchNPU-version` | `2.10.0.post2` | Full TorchNPU version used by the image, including suffixes such as `.postN` |
| `chip` | `910b` or `a3` | Ascend chip model supported by the image |
| `os` | `ubuntu22.04` or `openeuler24.03` | Container operating system family and version |
| `Python-version` | `py3.12` | Parsed from the CANN base image tag |
Examples:
For example:
```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
0.9.5-cann9.1.0-torch_npu2.10.0.post2-a3-ubuntu22.04-py3.12
```
## Quick Start
@@ -94,33 +85,30 @@ Before starting a container:
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.
Driver, firmware, CANN, TorchNPU, 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.
The following example starts the latest A2 Ubuntu image with one NPU. Adjust `DOCKER_IMAGE` and the `--device` options for your environment.
```bash
export IMAGE=quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
docker pull "$IMAGE"
CONTAINER_NAME=llamafactory-npu
DOCKER_IMAGE=hiyouga/llamafactory:latest-910b-ubuntu
docker run --rm -it \
--name llamafactory-npu \
--ipc=host \
--net=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 /usr/local/dcmi:/usr/local/dcmi \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v "$HOME/.cache/huggingface:/root/.cache/huggingface" \
-p 7860:7860 \
-p 8000:8000 \
"$IMAGE" \
bash
-v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
-v /data:/data \
--name "$CONTAINER_NAME" \
"$DOCKER_IMAGE" \
/bin/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.
@@ -134,22 +122,16 @@ python -c "import torch, torch_npu; print(torch.__version__, torch_npu.__version
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:
Run the build from the repository root. The following example builds the A2 Ubuntu 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 BASE_IMAGE=quay.io/ascend/cann:9.1.0-910b-ubuntu22.04-py3.12 \
--build-arg PIP_INDEX=https://pypi.org/simple \
-t llamafactory:npu-a3-openeuler \
-t llamafactory:npu-910b-ubuntu \
.
```
@@ -157,76 +139,45 @@ Available build arguments:
| Argument | Default | Purpose |
| --- | --- | --- |
| `BASE_IMAGE` | A2 Ubuntu CANN 9.0.0 image | Selects the accelerator and container OS variant |
| `BASE_IMAGE` | `quay.io/ascend/cann:9.1.0-910b-ubuntu22.04-py3.12` | Selects the base image that matches the device model and container operating system |
| `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 |
| `PYTORCH_INDEX` | `https://download.pytorch.org/whl/cpu` | Selects the PyTorch wheel index used with TorchNPU |
| `HTTP_PROXY` | Empty | Provides an optional HTTP/HTTPS proxy during the build |
Docker Compose can build and start each supported variant:
### Start with Docker Compose
The preceding `docker build` command invokes the Dockerfile directly. It builds an image but does not start a container. Docker Compose does not use a separate build implementation: it reads the presets in `docker-compose.yml`, reuses the same Dockerfile, and selects a hardware-series and operating-system combination through a profile. Each `up -d` command below starts the selected container in the background. If the image is not available locally, Docker Compose builds it first:
```bash
cd docker/docker-npu
# A2 with Ubuntu
docker compose up -d llamafactory-a2-ubuntu
docker compose --profile a2-ubuntu up -d
# A3 with Ubuntu
docker compose --profile a3 up -d llamafactory-a3-ubuntu
docker compose --profile a3-ubuntu up -d
# A2 with openEuler
docker compose --profile openeuler up -d llamafactory-a2-openeuler
docker compose --profile a2-openeuler up -d
# A3 with openEuler
docker compose --profile a3-openeuler up -d llamafactory-a3-openeuler
docker compose --profile a3-openeuler up -d
```
### 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.
To build an image with Docker Compose without starting a container, use `docker compose --profile <profile> build`.
## 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.
- The image build targets both x86-64 (`linux/amd64`) and AArch64 (`linux/arm64`) hosts. The CPU architecture is independent of whether the hardware series 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.
- Legacy NPU tags are replaced by the `latest-<910b|a3>-<ubuntu|openeuler>` format.
- 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).
LlamaFactory 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.
Ascend CANN, TorchNPU, Triton Ascend, DeepSpeed, base operating-system packages, model weights, datasets, and other third-party components are governed by their respective licenses and terms. The LlamaFactory 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.

View File

@@ -1,8 +1,8 @@
# 面向昇腾 NPU 的 LLaMA Factory 镜像
# 面向昇腾 NPU 的 LlamaFactory 镜像
LLaMA Factory 昇腾 NPU 镜像面向华为昇腾 Atlas NPU提供可直接用于大语言模型和多模态模型微调、评测与服务部署的运行环境。镜像基于昇腾 CANN 容器镜像构建,预装 LLaMA Factory、Python、PyTorch、torch-npu、Triton Ascend、DeepSpeed 和 LLaMA Factory 评测依赖
LlamaFactory 昇腾 NPU 镜像面向华为昇腾 Atlas NPU提供可直接使用的 LlamaFactory 环境。镜像基于昇腾 CANN 容器镜像构建,预装 Python、PyTorch、TorchNPU、DeepSpeed、LlamaFactory 等组件
安装方法和问题排查请参考 [LLaMA Factory NPU 安装及配置文档](https://llamafactory.readthedocs.io/zh-cn/latest/multibackend/npu/npu_installation.html)。
安装方法和问题排查请参考 [LlamaFactory NPU 安装及配置文档](https://llamafactory.readthedocs.io/zh-cn/latest/multibackend/npu/npu_installation.html)。
## 快速参考
@@ -11,77 +11,68 @@ LLaMA Factory 昇腾 NPU 镜像面向华为昇腾 Atlas NPU提供可直接用
- `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`
当前提供以下镜像组合
当前提供以下 `latest` NPU 镜像 tag
| 加速器 | 容器操作系统 | CANN 基础镜像 |
| 硬件系列 | 操作系统 | Tag |
| --- | --- | --- |
| 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` |
| A2 | Ubuntu 22.04 | `latest-910b-ubuntu` |
| A3 | Ubuntu 22.04 | `latest-a3-ubuntu` |
| A2 | openEuler 24.03 | `latest-910b-openeuler` |
| A3 | openEuler 24.03 | `latest-a3-openeuler` |
## 镜像介绍
镜像用于运行 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` |
| CANN | `9.1.0` |
| Python | `3.12` |
| PyTorch | `2.10.0` |
| TorchNPU | `2.10.0.post2` |
| torchvision / torchaudio | `0.25.0` / `2.10.0` |
| Transformers | 构建时的最新兼容版本 |
| Triton Ascend | `3.2.1` |
| DeepSpeed | `>=0.10.0,<=0.18.4` |
| LLaMA Factory | 从构建上下文中的仓库源码安装 |
| DeepSpeed | 构建时的最新兼容版本 |
| LlamaFactory | 从构建上下文中的仓库源码安装 |
镜像不包含模型权重和数据集。请通过目录挂载或运行时下载的方式单独提供,并遵守对应的许可证和使用要求。
## 镜像 Tag 说明与 Dockerfile 归档路径
## 镜像 Tag 说明
镜像使用以下 tag 格式:
NPU 镜像的 `latest` 和 release tag 使用不同格式;以下规则不适用于 CUDA 镜像。
非 release 构建复用以下简短 tag每次定时构建会更新对应 tag 所指向的镜像:
```text
<llamafactory版本>-cann<CANN版本>-torch_npu<torch-npu版本>-<加速器>-<操作系统>-<Python版本>
latest-<芯片信息>-<操作系统>
```
| 字段 | 可选值 | 说明 |
| --- | --- | --- |
| `芯片信息` | `910b``a3` | 镜像所适配的昇腾芯片型号 |
| `操作系统` | `ubuntu``openeuler` | 容器操作系统类型 |
Release 构建使用完整 tag
```text
<LlamaFactory版本>-cann<CANN版本>-torch_npu<TorchNPU版本>-<芯片信息>-<操作系统>-<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 中提取 |
| `LlamaFactory版本` | `0.9.5` | LlamaFactory release 版本号 |
| `CANN版本` | `9.1.0` | 从 CANN 基础镜像 tag 中提取 |
| `TorchNPU版本` | `2.10.0.post2` | 镜像使用的 TorchNPU 完整版本,包含 `.postN` 等后缀 |
| `芯片信息` | `910b``a3` | 镜像所适配的昇腾芯片型号 |
| `操作系统` | `ubuntu22.04``openeuler24.03` | 容器操作系统类型和版本 |
| `Python版本` | `py3.12` | 从 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
0.9.5-cann9.1.0-torch_npu2.10.0.post2-a3-ubuntu22.04-py3.12
```
## 快速开始
@@ -94,33 +85,30 @@ docker/docker-npu/
2. 确认宿主机执行 `npu-smi info` 可以正常识别 NPU。
3. 安装 Docker并确保当前用户有权访问所需的昇腾设备节点和驱动文件。
驱动、固件、CANN、torch-npu 与目标昇腾硬件需要保持兼容。
驱动、固件、CANN、TorchNPU 与目标昇腾硬件需要保持兼容。
### 拉取并运行镜像
以下示例使用一张 NPU 启动最新的 A2 Ubuntu 镜像。请根据实际环境修改镜像 tag 和 `/dev/davinci0`
以下示例使用一张 NPU 启动最新的 A2 Ubuntu 镜像。请根据实际情况修改 ``DOCKER_IMAGE`` 和 ``device``
```bash
export IMAGE=quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
docker pull "$IMAGE"
CONTAINER_NAME=llamafactory-npu
DOCKER_IMAGE=hiyouga/llamafactory:latest-910b-ubuntu
docker run --rm -it \
--name llamafactory-npu \
--ipc=host \
--net=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 /usr/local/dcmi:/usr/local/dcmi \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v "$HOME/.cache/huggingface:/root/.cache/huggingface" \
-p 7860:7860 \
-p 8000:8000 \
"$IMAGE" \
bash
-v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
-v /data:/data \
--name "$CONTAINER_NAME" \
"$DOCKER_IMAGE" \
/bin/bash
```
部分驱动环境中的 `npu-smi` 位于 `/usr/local/sbin/npu-smi`,此时需要调整挂载源路径。使用多张 NPU 时,继续追加 `--device=/dev/davinci<N>` 参数。
@@ -134,22 +122,16 @@ python -c "import torch, torch_npu; print(torch.__version__, torch_npu.__version
llamafactory-cli help
```
需要使用 LLaMA Board 时执行:
### 本地构建镜像
```bash
llamafactory-cli webui
```
### 本地构建
在仓库根目录执行构建。以下示例构建 A3 openEuler 镜像:
在仓库根目录执行构建。以下示例构建 A2 Ubuntu 镜像:
```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 BASE_IMAGE=quay.io/ascend/cann:9.1.0-910b-ubuntu22.04-py3.12 \
--build-arg PIP_INDEX=https://pypi.org/simple \
-t llamafactory:npu-a3-openeuler \
-t llamafactory:npu-910b-ubuntu \
.
```
@@ -157,76 +139,45 @@ docker build \
| 参数 | 默认值 | 用途 |
| --- | --- | --- |
| `BASE_IMAGE` | A2 Ubuntu CANN 9.0.0 镜像 | 选择加速器和容器操作系统组合 |
| `BASE_IMAGE` | `quay.io/ascend/cann:9.1.0-910b-ubuntu22.04-py3.12` | 根据设备型号和容器操作系统选择对应的基础镜像 |
| `PIP_INDEX` | `https://pypi.org/simple` | 指定 Python 软件包索引 |
| `PYTORCH_INDEX` | `https://download.pytorch.org/whl/cpu` | 指定配合 torch-npu 使用的 PyTorch wheel 索引 |
| `PYTORCH_INDEX` | `https://download.pytorch.org/whl/cpu` | 指定配合 TorchNPU 使用的 PyTorch wheel 索引 |
| `HTTP_PROXY` | 空 | 构建期间可选的 HTTP/HTTPS 代理 |
也可以通过 Docker Compose 构建并启动各个组合:
### 通过 Docker Compose 启动
前面的 `docker build` 命令直接调用 Dockerfile只构建镜像不启动容器。Docker Compose 不使用另一套构建逻辑:它读取 `docker-compose.yml` 中的预设配置,复用同一个 Dockerfile并通过 profile 选择硬件系列和操作系统组合。下面的 `up -d` 会在后台启动容器若本地镜像不存在Docker Compose 会先构建镜像:
```bash
cd docker/docker-npu
# A2 + Ubuntu
docker compose up -d llamafactory-a2-ubuntu
docker compose --profile a2-ubuntu up -d
# A3 + Ubuntu
docker compose --profile a3 up -d llamafactory-a3-ubuntu
docker compose --profile a3-ubuntu up -d
# A2 + openEuler
docker compose --profile openeuler up -d llamafactory-a2-openeuler
docker compose --profile a2-openeuler up -d
# A3 + openEuler
docker compose --profile a3-openeuler up -d llamafactory-a3-openeuler
docker compose --profile a3-openeuler up -d
```
### 二次开发
交互式开发时,可以将本地源码挂载到容器中,并在容器内以 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
```
运行派生镜像时仍需传入昇腾设备和驱动挂载参数,不应将设备访问配置固化到镜像中。
如果只想通过 Docker Compose 构建镜像而不启动容器,请使用 `docker compose --profile <profile> build`。
## 硬件支持与兼容性说明
- A2 镜像使用标记为 `910b` 的 CANN 基础镜像A3 镜像使用标记为 `a3` 的 CANN 基础镜像。
- 镜像构建目标同时包含 x86-64`linux/amd64`)和 AArch64`linux/arm64`宿主机。CPU 架构与加速器属于 A2 还是 A3 无关。
- 镜像构建目标同时包含 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。
- 旧式 NPU tag 已由 `latest-<910b|a3>-<ubuntu|openeuler>` 格式取代。
- 正式部署前请验证具体驱动、固件、CANN 和 SoC 组合的兼容性。
## 许可证与免责声明
LLaMA Factory 基于 [Apache License 2.0](../../LICENSE) 发布。
LlamaFactory 基于 [Apache License 2.0](../../LICENSE) 发布。
昇腾 CANN、torch-npu、Triton Ascend、DeepSpeed、基础操作系统软件包、模型权重、数据集和其他第三方组件分别受其自身许可证与条款约束。LLaMA Factory 的许可证不会替代或覆盖这些条款。
昇腾 CANN、TorchNPU、Triton Ascend、DeepSpeed、基础操作系统软件包、模型权重、数据集和其他第三方组件分别受其自身许可证与条款约束。LlamaFactory 的许可证不会替代或覆盖这些条款。
本镜像按“原样”提供,不附带任何明示或暗示的保证。用户需要自行验证软硬件兼容性、保障容器及运行配置的安全、遵守适用的许可证和法律,并在训练、评测或部署前审查模型与数据集的使用条款。

View File

@@ -26,25 +26,26 @@ x-npu-common: &npu-common
services:
llamafactory-a2-ubuntu:
<<: *npu-common
profiles: ["a2-ubuntu"]
build:
<<: *build
args:
<<: *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
BASE_IMAGE: quay.io/ascend/cann:9.1.0-910b-ubuntu22.04-py3.12
container_name: llamafactory-910b-ubuntu
image: llamafactory:npu-910b-ubuntu
ports:
- "7860:7860"
- "8000:8000"
llamafactory-a3-ubuntu:
<<: *npu-common
profiles: ["a3"]
profiles: ["a3-ubuntu"]
build:
<<: *build
args:
<<: *build-args
BASE_IMAGE: quay.io/ascend/cann:9.0.0-a3-ubuntu22.04-py3.11
BASE_IMAGE: quay.io/ascend/cann:9.1.0-a3-ubuntu22.04-py3.12
container_name: llamafactory-a3-ubuntu
image: llamafactory:npu-a3-ubuntu
ports:
@@ -53,14 +54,14 @@ services:
llamafactory-a2-openeuler:
<<: *npu-common
profiles: ["openeuler"]
profiles: ["a2-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
BASE_IMAGE: quay.io/ascend/cann:9.1.0-910b-openeuler24.03-py3.12
container_name: llamafactory-910b-openeuler
image: llamafactory:npu-910b-openeuler
ports:
- "7862:7860"
- "8002:8000"
@@ -72,7 +73,7 @@ services:
<<: *build
args:
<<: *build-args
BASE_IMAGE: quay.io/ascend/cann:9.0.0-a3-openeuler24.03-py3.11
BASE_IMAGE: quay.io/ascend/cann:9.1.0-a3-openeuler24.03-py3.12
container_name: llamafactory-a3-openeuler
image: llamafactory:npu-a3-openeuler
ports:

View File

@@ -0,0 +1,173 @@
# FSDPTurbo EP/EFSDP and LlamaFactory FSDP2/CP Design
Chinese version: [FSDPTurbo EP/EFSDP 与 LlamaFactory FSDP2/CP 设计说明](../../../zh/advanced/distributed/fsdpturbo-ep-efsdp.md)
This document describes the current implementation of the `fsdpturbo` distributed plugin. Its core principle is a clear separation of responsibilities:
- FSDPTurbo owns expert parallelism (EP), expert parameter sharding (EFSDP), and device operator registration.
- LlamaFactory owns process initialization, the base DeviceMesh, outer FSDP2, CP, model initialization, and weight loading.
- The LlamaFactory integration layer combines the two parameter layouts and handles gradient norms across meshes.
## 1. Configuration Boundaries
Common parallel topology belongs to `TrainingArguments`, while FSDPTurbo-only settings remain in `dist_config`:
```yaml
cp_size: 1
dist_config:
name: fsdpturbo
ep_size: 16
ep_dispatcher: eager
```
The fields used by the minimal example have the following responsibilities:
- `ep_size`: expert-parallel group size.
- `ep_dispatcher`: FSDPTurbo EP dispatcher, which defaults to `eager`.
`dp_size`, `cp_size`, `cp_mode`, `mp_replicate_size`, `mp_shard_size`, and `dist_timeout` are common topology fields and therefore remain at the top level. `dist_config` is parsed strictly as `FSDPTurboParams`; putting a common topology field inside it is rejected instead of being silently ignored.
The top-level training option `bf16` controls FSDPTurbo parameter storage and compute dtype. The backend casts the model before FSDP materialization, so `ModelEngine` does not need to read distributed-backend configuration.
The following advanced fields are optional and are therefore omitted from the minimal YAML example above:
- `fsdp_ignored_modules`: additional modules excluded from the outer LlamaFactory FSDP2 path. Expert parameters selected by the model spec are automatically added to the ignored set by the integration layer, so normal configurations do not need to repeat them here.
- `hook_modules`: optional module patterns for FSDPTurbo EFSDP hooks. The default is an empty list.
- `fsdp_implementation`: the FSDPTurbo EFSDP implementation, either `native` or `custom`. The default is `native`.
The model spec determines the EFSDP targets. Non-expert parameters such as attention, embeddings, and the LM head do not enter the FSDPTurbo EFSDP plan. They remain managed by the outer LlamaFactory FSDP2 layer.
Model-specific module paths and preparation logic are managed exclusively by the `FSDPTurboEPModelSpec` registry. Built-in specs currently cover `qwen3_moe` and `qwen3_5_moe`; unregistered models fail with an explicit error. `ep_modules` and `ep_fsdp_modules` are not YAML options, and strict parameter parsing rejects them to prevent configuration from drifting away from the actual model structure.
## 2. Mesh Initialization
LlamaFactory's `DistributedInterface` initializes only its existing model and data meshes. It is unaware of EP and EFSDP and does not expose an extra mesh registration interface for distributed plugins. The FSDPTurbo expert topology is independently created and owned by `FSDPTurboParallelState` in the plugin module:
```text
run_sft / run_dpo / run_rm
-> DistributedInterface(training_args)
-> initialize LlamaFactory model/data meshes
-> DistributedPlugin("fsdpturbo").shard_model(...)
-> FSDPTurboFSDP2Engine.__init__()
-> FSDPTurboParallelState.initialize()
-> initialize and retain the expert parent mesh and submeshes
```
`FSDPTurboParallelState` creates a four-dimensional expert parent mesh:
```text
(edp, efsdp, ep, expert_cp)
```
Its current dimensions are calculated as follows:
```text
dp_size = world_size / cp_size
ep_fsdp_size = dp_size / ep_size
edp_size = dp_size / (ep_size * ep_fsdp_size)
mesh_shape = (edp_size, ep_fsdp_size, ep_size, cp_size)
```
The state object retains `edp_mesh`, `efsdp_mesh`, `ep_mesh`, and `expert_cp_mesh`. Model sharding and gradient norm logic inside the plugin read expert communication domains from this state, while other LlamaFactory backends do not need to implement or know about these interfaces. Initialization validates that `ep_size` is positive and divides `dp_size`; repeated initialization also rejects topology changes.
## 3. Model Sharding Order
The wrapping order must remain "expert side first, outer FSDP2 second":
```text
DistributedPlugin("fsdpturbo")
-> FSDPTurboFSDP2Engine.shard_model(model)
-> prepare_model_ep(model)
-> expert_parallelize_modules(model, ep_mesh, ep_plan)
-> expert_fully_shard_modules(model, efsdp_mesh, ep_plan, fsdp_plan)
-> collect expert parameters as ignored_params
-> FSDP2Engine.prepare_model(model, ignored_params=...)
-> apply outer fully_shard to the remaining Transformer Layers and root module
```
This prevents the same expert parameter from being managed by both EFSDP and outer FSDP2. Outer FSDP2 continues to reuse LlamaFactory's model initialization, checkpoint, and save flows.
The LlamaFactory integration layer accepts `eager`, `fused`, `mc2`, and `domino` and forwards the selected value unchanged to FSDPTurbo. Their implementation boundaries and current validation status differ:
| Dispatcher | Main path | Additional requirements | Validation in this PR |
| --- | --- | --- | --- |
| `eager` | Uses PyTorch implementations of permute, unpermute, and grouped matmul while tensors remain on the current accelerator, with standard AllToAll for token dispatch and combine | Minimal dependencies; serves as the reference implementation | End-to-end numerical and performance validation completed on Ascend A3 |
| `fused` | Keeps the same AllToAll topology while replacing permute, unpermute, and grouped matmul with device-fused operators | Requires matching device operators, dtypes, and layouts; local operators may fall back to eager when an expert receives no tokens | End-to-end numerical and performance validation completed on Ascend A3 |
| `mc2` | Uses dedicated operators that fuse AllToAllV with grouped matmul to reduce intermediate communication-computation overhead | Requires the MC2 NPU operators, an HCCL communicator, and their shape and dtype constraints | Implemented by FSDPTurbo but not validated end to end in this PR |
| `domino` | Splits the first dimension of the expert-module input into two slices and uses a separate communication stream and events to overlap AllToAll with expert computation | Requires asynchronous stream/event support and enough token work in both slices to amortize scheduling overhead | Implemented by FSDPTurbo but not validated end to end in this PR |
Only `eager` and `fused` are validated here because they cover the reference path and the commonly used A3 device-fused path, respectively, and therefore isolate and establish the correctness of the EP/EFSDP integration between LlamaFactory and FSDPTurbo. The current experiment matrix was not extended to `mc2` and `domino`: they add operator, communication-scheduling, and input-shape constraints that require separate numerical comparisons, long-run stability tests, and profiler analysis. They are accepted configuration choices, but the results in this PR should not be interpreted as evidence that they have reached the same stability, numerical, or performance level.
## 4. FSDPTurbo Dependency Entry Points
LlamaFactory imports each required object directly from the module that defines it:
```python
from fsdp_turbo.distributed.expert_parallel.expert_fully_shard_parallel import (
expert_fully_shard_modules,
)
from fsdp_turbo.distributed.expert_parallel.expert_parallel import expert_parallelize_modules
from fsdp_turbo.fsdp_turbo_config import EPPlanConfig, FSDPPlanConfig
from fsdp_turbo.utils.str_match import module_name_match
```
The imports occur inside `prepare_model_ep()`, so other distributed backends remain importable when FSDPTurbo is not installed. They intentionally bypass aggregate exports from `fsdp_turbo.distributed.__init__` to avoid extra dependencies and potential import cycles during package initialization.
## 5. Gradient Norms
Outer and expert parameters can belong to different DTensor meshes and therefore cannot be passed together to a single standard `clip_grad_norm_()` call. The `fsdpturbo` plugin groups parameters by their owning mesh and computes local p-power sums:
- Non-expert parameters are reduced over the DP and CP groups.
- Expert parameters are reduced over the EFSDP, EP, and expert-CP groups retained by `FSDPTurboParallelState`.
- After the global norm is assembled, the same clipping coefficient is applied to every local gradient.
A zero-gradient warmup runs during startup so that the required collectives are initialized before training begins. This is currently a backend-specific implementation for `fsdpturbo`; other backends retain their existing gradient norm paths until the upstream distributed plugin interface is decoupled.
## 6. Weight Loading
LlamaFactory retains the `init_on_meta` and safetensors loading flow. The parent `FSDP2Engine` loader dynamically invokes the FSDPTurbo engine override through `self._copy_weights(...)`, so the method is not dead code. It supports DTensors with multiple `Shard` placements by calculating the local slice for the current rank along each mesh dimension in sequence. Model save and checkpoint interfaces continue to reuse the LlamaFactory FSDP2 implementation.
## 7. Kernel Plugin
FLA operators do not belong in the distributed configuration. Operator selection is handled through an independent `kernel_config`:
```yaml
kernel_config:
name: auto, flash-linear-attention
include_kernels: chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
chunk_size: 32
```
The call path is:
```text
ModelEngine
-> apply_kernels("auto, flash-linear-attention")
-> accelerator-specific LlamaFactory auto kernels
-> KernelPlugin("flash-linear-attention").apply(...)
-> fsdp_turbo.ops.get_op()
-> FSDPTurbo device operator registry
-> fsdp_turbo.utils.patch.patch_model_members()
-> FLA backend implementation
```
`chunk_size` accepts `16`, `32`, and `64`, with a default of `64`. The kernel plugin and distributed plugin are independent. `name: flash-linear-attention` installs only the selected FLA operators. The comma-separated `name: auto, flash-linear-attention` form composes LlamaFactory's accelerator-specific automatic kernels with the FLA plugin before distributed sharding. LlamaFactory owns the operator-to-model-attribute mapping and `chunk_size` binding; FSDPTurbo owns device operator registration, selection, and generic callable patching. FLA stays explicit because it has optional external dependencies and is not part of the built-in `auto` set. FSDPTurbo subsequently replaces the target expert module's `forward`, so the final expert execution path is selected by `ep_dispatcher`; an MoE kernel applied during the auto stage is not retained as a separate second expert execution path.
## 8. CP Runtime Constraints and Validation Scope
When `init_on_meta` constructs the model, it must propagate `attn_implementation` in the same way as the `from_pretrained` path. Otherwise, the model falls back to a non-FlashAttention implementation and Ulysses CP cannot start. Before calling Hugging Face FlashAttention, Ulysses reconstructs the global attention mask. Only two-dimensional position IDs participate in packed-sequence detection. Multi-axis position IDs such as Qwen3.5 mRoPE have already been consumed by rotary embedding and must not be passed to the FlashAttention packed-sequence detection logic.
The current implementation has completed the following BF16 AdamW full SFT validations with Qwen3.5-35B-A3B on Atlas 900 A3 SuperPoD and Atlas 950 SuperPoD systems. This revalidation used FSDPTurbo `0e96fbc`. The A3 environment used CANN 9.0.0, PyTorch 2.7.1, and torch-npu 2.7.1.post4; the A5 environment used CANN 9.1.0-beta.3, PyTorch 2.10.0, and torch-npu 2.10.0.post2. Performance is calculated from the step 1 and step 100 log timestamps and excludes initialization and compilation before the first step as well as model saving after training:
| Machine | CP | EP | EFSDP | Checkpoint | Kernel / Dispatcher | Steps | Loss (first -> last) | Performance | Result |
| --- | ---: | ---: | ---: | --- | --- | ---: | --- | ---: | --- |
| Atlas 900 A3 SuperPoD | 1 | 16 | 1 | Off | FLA (chunk size 16) / eager | 100 | 1.3361 -> 0.0793 | 2.51 s/it | Passed and saved |
| Atlas 900 A3 SuperPoD | 1 | 16 | 1 | Off | FLA (chunk size 16) / fused | 100 | 1.3354 -> 0.1179 | 2.17 s/it | Passed and saved |
| Atlas 900 A3 SuperPoD | 2 | 4 | 2 | Off | auto + FLA (chunk size 64) / fused | 100 | 1.8114 -> 0.5260 | 7.65 s/it | Passed and saved |
| Atlas 900 A3 SuperPoD | 2 | 4 | 2 | Off | auto + FLA (chunk size 64) / eager | 100 | 1.8095 -> 0.5596 | 5.88 s/it | Passed and saved |
| Atlas 950 SuperPoD | 1 | 8 | 1 | Off | no kernel plugin configured / eager | 100 | 1.3575 -> 0.4439 | 2.68 s/it | Passed and saved |
Loss and gradient norm remained finite in all five runs, and every run completed 100 steps and model saving. With the same partition, the per-step loss correlation between eager and fused was 0.997 for EP16 and 0.977 for CP2/EP4/EFSDP2, which indicates consistent optimization trajectories. The performance effect depends on the partition: fused was about 13% faster than eager with EP16, but about 30% slower after adding CP and EFSDP. Fused therefore should not be treated as the default optimum for every mesh.
The EP16 runs used global batch 16 and cutoff length 256. The CP2 runs used global batch 8 and cutoff length 128. The A5 run used global batch 8 and cutoff length 256. The first-to-last loss validates convergence within each run; absolute loss values across different partition groups should not be used directly as a numerical-equivalence conclusion.

View File

@@ -0,0 +1,83 @@
# KTransformers LoRA SFT
KTransformers (KT) executes routed MoE experts on CPU while LLaMA-Factory remains responsible for data,
LoRA arguments, and the training entry point. The production scope is routed-BF16 and routed-INT8 LoRA.
KT has one user configuration source: the training YAML. Accelerate YAML contains FSDP2 settings only.
LLaMA-Factory derives LoRA rank, alpha, dropout, activation policy, and local runtime capacity.
```yaml
finetuning_type: lora
lora_rank: 8
lora_alpha: 16
lora_target: all
use_kt: true
disable_gradient_checkpointing: false
kt_cpu_activation: retain
kt_config:
kt_expert_weight_format: bf16
kt_backend: AMXBF16
kt_num_threads: 96
kt_tp_enabled: true
kt_threadpool_count: 2
kt_max_cache_depth: 2
```
Routed INT8 additionally requires matching expert and BF16 non-expert artifacts:
```yaml
kt_weight_path: /abs/path/to/routed-int8-experts
kt_non_expert_weight_path: /abs/path/to/bf16-non-expert-cache
kt_config:
kt_expert_weight_format: int8
kt_backend: auto
kt_weight_lifecycle: persistent
```
Launch the standard training entry point through Accelerate:
```bash
CUDA_VISIBLE_DEVICES=0,1 accelerate launch \
--config_file examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml \
src/train.py examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml
```
## Load a saved adapter
Use a local, complete KT adapter directory for chat or evaluation. Repeat the training LoRA shape (`finetuning_type`,
`lora_rank`, `lora_alpha`, and `lora_dropout`) and the KT base-weight settings. In particular, routed INT8 loading
must use the same `kt_weight_path` and `kt_non_expert_weight_path` as training.
```yaml
model_name_or_path: /abs/path/to/base-model
adapter_name_or_path: /abs/path/to/output/checkpoint-300
finetuning_type: lora
lora_rank: 8
lora_alpha: 16
lora_dropout: 0.0
use_kt: true
kt_cpu_activation: retain
kt_config:
kt_expert_weight_format: bf16
kt_backend: AMXBF16
kt_num_threads: 96
```
```bash
llamafactory-cli chat path/to/kt_adapter_infer.yaml
llamafactory-cli eval path/to/kt_adapter_eval.yaml
```
The directory must contain the standard PEFT adapter files and, when fused routed-expert LoRA is used,
`fused_expert_lora.safetensors` plus `kt_adapter_manifest.json`. LLaMA-Factory first loads the standard PEFT
adapter, then KT validates and restores the fused artifact. `adapter_folder` may select a local subdirectory;
paths outside the adapter root and Hub adapter IDs fail before model loading. Download a Hub bundle locally first.
For training resume, keep the original training YAML and use `resume_from_checkpoint`. The optimizer checkpoint
currently requires the same distributed world size. Missing, tampered, or mismatched artifacts fail closed instead
of falling back to the source checkpoint.
Do not combine KT with a second Transformers/FSDP checkpoint wrapper or Unsloth GC, and do not put `kt_config`
in the Accelerate YAML. See the BF16 and INT8 examples under `examples/ktransformers/train_lora/`.

View File

@@ -34,9 +34,11 @@ LlamaFactory Docs
advanced/lora-and-quantization/lora
advanced/lora-and-quantization/quantization
advanced/ktransformers
advanced/distributed/fsdp
advanced/distributed/deepspeed
advanced/distributed/parallel-dp-tp-ep-sp-cp
advanced/distributed/fsdpturbo-ep-efsdp
advanced/custom-kernels/triton
advanced/custom-kernels/fused-operators

View File

@@ -0,0 +1,215 @@
# FSDPTurbo EP/EFSDP 与 LlamaFactory FSDP2/CP 设计说明
English version: [FSDPTurbo EP/EFSDP and LlamaFactory FSDP2/CP Design](../../../en/advanced/distributed/fsdpturbo-ep-efsdp.md)
本文描述 `fsdpturbo` distributed plugin 的当前实现。核心原则是保持两侧职责清晰:
- FSDPTurbo 负责专家并行EP、专家参数分片EFSDP和设备算子注册。
- LlamaFactory 负责进程初始化、基础 DeviceMesh、外层 FSDP2、CP、模型初始化与权重加载。
- LlamaFactory 的集成层负责把两套参数布局组合起来,并处理跨 Mesh 的梯度范数。
## 1. 配置边界
公共并行拓扑放在 `TrainingArguments` 顶层FSDPTurbo 私有参数保留在 `dist_config`
```yaml
cp_size: 1
dist_config:
name: fsdpturbo
ep_size: 16
ep_dispatcher: eager
```
最小示例中的字段职责如下:
- `ep_size`:专家并行组大小。
- `ep_dispatcher`FSDPTurbo EP dispatcher默认为 `eager`
`dp_size``cp_size``cp_mode``mp_replicate_size``mp_shard_size``dist_timeout`
属于公共拓扑字段,继续放在顶层。`dist_config` 会被严格解析为 `FSDPTurboParams`;如果把公共拓扑
字段误放进去,会直接报错,而不是静默忽略。
顶层训练参数 `bf16` 同时控制 FSDPTurbo 的参数存储和计算 dtype。backend 会在 FSDP materialization
前完成模型 dtype 转换,因此 `ModelEngine` 不需要读取 distributed backend 配置。
以下高级字段为可选项,因此没有写入上面的最小 YAML 示例:
- `fsdp_ignored_modules`:额外排除在 LlamaFactory 外层 FSDP2 之外的模块。模型规格选中的专家参数
会被集成层自动加入忽略集合,普通配置无需重复填写。
- `hook_modules`FSDPTurbo EFSDP hook 的可选模块模式,默认为空列表。
- `fsdp_implementation`FSDPTurbo EFSDP 实现,可选 `native``custom`,默认为 `native`
EFSDP 的目标由模型规格决定。Attention、Embedding、LM Head 等非专家参数不进入 FSDPTurbo
EFSDP plan而是继续由 LlamaFactory 外层 FSDP2 管理。
模型相关的模块路径和准备逻辑统一由 `FSDPTurboEPModelSpec` 注册表管理。当前内置 `qwen3_moe`
`qwen3_5_moe`;未注册的模型会明确报错。`ep_modules``ep_fsdp_modules` 不属于 YAML
接口,严格参数解析会拒绝这两个字段,避免用户配置与模型实际结构失配。
## 2. Mesh 初始化
LlamaFactory 的 `DistributedInterface` 只初始化自身原有的 model/data mesh。它不感知 EP、EFSDP
也不为 distributed plugin 提供额外 mesh 注册接口。FSDPTurbo 的专家拓扑由插件文件内的
`FSDPTurboParallelState` 独立创建和持有:
```text
run_sft / run_dpo / run_rm
-> DistributedInterface(training_args)
-> 初始化 LlamaFactory model/data mesh
-> DistributedPlugin("fsdpturbo").shard_model(...)
-> FSDPTurboFSDP2Engine.__init__()
-> FSDPTurboParallelState.initialize()
-> 初始化并保存 expert parent mesh 及其子 mesh
```
`FSDPTurboParallelState` 创建专家侧四维父 Mesh
```text
(edp, efsdp, ep, expert_cp)
```
当前尺寸计算为:
```text
dp_size = world_size / cp_size
ep_fsdp_size = dp_size / ep_size
edp_size = dp_size / (ep_size * ep_fsdp_size)
mesh_shape = (edp_size, ep_fsdp_size, ep_size, cp_size)
```
状态对象保存 `edp_mesh``efsdp_mesh``ep_mesh``expert_cp_mesh`。插件内部模型切分和梯度范数
都从这个状态对象读取专家通信域LlamaFactory 其他 backend 不需要实现或感知这些接口。状态初始化
会校验 `ep_size` 为正数且能够整除 `dp_size`,重复初始化时也会拒绝拓扑发生变化。
## 3. 模型切分顺序
模型包装顺序必须保持为“专家侧优先,外层 FSDP2 随后”:
```text
DistributedPlugin("fsdpturbo")
-> FSDPTurboFSDP2Engine.shard_model(model)
-> prepare_model_ep(model)
-> expert_parallelize_modules(model, ep_mesh, ep_plan)
-> expert_fully_shard_modules(model, efsdp_mesh, ep_plan, fsdp_plan)
-> 收集专家参数作为 ignored_params
-> FSDP2Engine.prepare_model(model, ignored_params=...)
-> 对剩余 Transformer Layer 和根模块执行 outer fully_shard
```
这样可以避免同一专家参数同时被 EFSDP 和外层 FSDP2 管理。外层 FSDP2 仍复用 LlamaFactory
原有的初始化、checkpoint 和保存流程。
LlamaFactory 集成层接受 `eager``fused``mc2``domino`,并将选项原样传给
FSDPTurbo。这四种模式的实现边界和当前验证状态不同
| Dispatcher | 主要路径 | 额外要求 | 本 PR 验证状态 |
| --- | --- | --- | --- |
| `eager` | 使用 PyTorch 实现 permute、unpermute 和 grouped matmul张量仍在当前加速设备上通过标准 AllToAll 完成 token dispatch/combine | 依赖最少,用作参考实现 | 已在 A3 上完成精度和性能验证 |
| `fused` | 保持相同的 AllToAll 拓扑,将 permute、unpermute 和 grouped matmul 切换为设备融合算子 | 需要对应的设备算子、dtype 和 layout 支持;存在空专家时可回退到 eager 局部算子 | 已在 A3 上完成精度和性能验证 |
| `mc2` | 使用专用算子融合 AllToAllV 和 grouped matmul减少通信与计算之间的中间开销 | 依赖 MC2 NPU 算子、HCCL communicator 以及对应的 shape/dtype 约束 | FSDPTurbo 提供实现,本 PR 未做端到端验证 |
| `domino` | 将专家模块输入的第一维分成两片,使用独立通信流和 event 重叠 AllToAll 与专家计算 | 需要异步 stream/event 支持,且两个分片都要有足够的 token 工作量才能覆盖调度开销 | FSDPTurbo 提供实现,本 PR 未做端到端验证 |
当前只验证 `eager``fused`,是因为它们分别覆盖参考实现和 A3 常用设备融合路径,可用于隔离并验证
LlamaFactory 与 FSDPTurbo 之间的 EP/EFSDP 集成正确性。本次实验矩阵没有继续扩展到 `mc2`
`domino`:它们还引入了额外的算子、通信调度和输入形状约束,需要独立比较数值、长步稳定性和 profiler 结果。
因此,它们在配置接口上可选,但不应从本 PR 的实验结果推断为已达到相同的稳定性、精度或性能水平。
## 4. FSDPTurbo 依赖入口
LlamaFactory 从各功能的定义模块直接导入所需对象:
```python
from fsdp_turbo.distributed.expert_parallel.expert_fully_shard_parallel import (
expert_fully_shard_modules,
)
from fsdp_turbo.distributed.expert_parallel.expert_parallel import expert_parallelize_modules
from fsdp_turbo.fsdp_turbo_config import EPPlanConfig, FSDPPlanConfig
from fsdp_turbo.utils.str_match import module_name_match
```
导入发生在 `prepare_model_ep()` 内,因此没有安装 FSDPTurbo 时,其他 distributed backend 仍可正常导入。
这里不通过 `fsdp_turbo.distributed.__init__` 聚合导出,避免 package 初始化期间的额外依赖和潜在循环导入。
## 5. 梯度范数
外层参数和专家参数可能属于不同 DTensor Mesh不能直接放入一次标准 `clip_grad_norm_()`
`fsdpturbo` plugin 按参数所属 Mesh 分组计算局部 p 次方和:
- 非专家参数沿 DP 和 CP group 汇总。
- 专家参数沿 `FSDPTurboParallelState` 保存的 EFSDP、EP 和 expert-CP group 汇总。
- 汇总得到全局范数后,对所有本地梯度应用同一个 clipping coefficient。
启动阶段会执行一次零梯度 warmup使相关 collective 在正式训练前完成初始化。
当前这是 `fsdpturbo` backend 的专用实现;其他 backend 继续保留原有梯度范数路径,等待上游
distributed plugin 解耦后再统一公共接口。
## 6. 权重加载
LlamaFactory 保留 `init_on_meta` 和 safetensors 加载流程。父类 `FSDP2Engine` 的加载器通过
`self._copy_weights(...)` 动态调用 FSDPTurbo engine 的覆写实现,因此该方法不是未使用代码。
它支持包含多个 `Shard` placement 的 DTensor按各 Mesh 维度依次计算当前 rank 对应的本地切片。
模型保存和 checkpoint 接口继续复用 LlamaFactory FSDP2 实现。
## 7. Kernel plugin
FLA 算子不属于 distributed config。算子选择通过独立的 `kernel_config` 完成:
```yaml
kernel_config:
name: auto, flash-linear-attention
include_kernels: chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
chunk_size: 32
```
调用链如下:
```text
ModelEngine
-> apply_kernels("auto, flash-linear-attention")
-> LlamaFactory 当前加速器对应的 auto kernels
-> KernelPlugin("flash-linear-attention").apply(...)
-> fsdp_turbo.ops.get_op()
-> FSDPTurbo device operator registry
-> fsdp_turbo.utils.patch.patch_model_members()
-> FLA backend implementation
```
`chunk_size` 当前支持 `16``32``64`,默认值为 `64`。Kernel plugin 与 distributed plugin
彼此独立。`name: flash-linear-attention` 只安装所选 FLA 算子;逗号分隔的
`name: auto, flash-linear-attention` 会在分布式切分前组合 LlamaFactory 当前加速器的 auto kernels
与 FLA plugin。LlamaFactory 负责算子名到模型属性的映射和 `chunk_size` 参数绑定FSDPTurbo 负责设备
算子注册、选择和通用 callable patch。FLA 依赖可选的外部三方件,因此保持显式选择,不属于内置
`auto` 集合。FSDPTurbo
随后会替换目标专家模块的 `forward`,所以专家计算的最终路径由 `ep_dispatcher` 决定auto 阶段
应用的 MoE kernel 不会作为独立的第二条专家执行路径保留下来。
## 8. CP 运行约束与验证范围
`init_on_meta` 构造模型时必须与 `from_pretrained` 路径一样传递 `attn_implementation`,否则模型会退回
非 FlashAttention 实现Ulysses CP 无法启动。Ulysses 在调用 Hugging Face FlashAttention 前重建全局
attention mask只有二维 position IDs 才参与 packed-sequence 检测。Qwen3.5 mRoPE 等多轴 position IDs
已经在 rotary embedding 中消费,不应传入 FlashAttention 的 packed-sequence 检测逻辑。
当前实现已在 Atlas 900 A3 SuperPoD 和 Atlas 950 SuperPoD 上用 Qwen3.5-35B-A3B 完成以下
BF16、AdamW full SFT 验证。本次重验证使用 FSDPTurbo `0e96fbc`A3 环境为 CANN 9.0.0、
PyTorch 2.7.1 和 torch-npu 2.7.1.post4A5 环境为 CANN 9.1.0-beta.3、PyTorch 2.10.0 和
torch-npu 2.10.0.post2。表中性能按第 1 步至第 100 步的日志时间戳计算,不包含首步前的初始化、
编译和训练后的模型保存时间:
| 机器型号 | CP | EP | EFSDP | Checkpoint | Kernel / Dispatcher | 步数 | Loss首步 -> 末步) | 性能 | 结果 |
| --- | ---: | ---: | ---: | --- | --- | ---: | --- | ---: | --- |
| Atlas 900 A3 SuperPoD | 1 | 16 | 1 | 关闭 | FLAchunk size 16/ eager | 100 | 1.3361 -> 0.0793 | 2.51 s/it | 通过并完成保存 |
| Atlas 900 A3 SuperPoD | 1 | 16 | 1 | 关闭 | FLAchunk size 16/ fused | 100 | 1.3354 -> 0.1179 | 2.17 s/it | 通过并完成保存 |
| Atlas 900 A3 SuperPoD | 2 | 4 | 2 | 关闭 | auto + FLAchunk size 64/ fused | 100 | 1.8114 -> 0.5260 | 7.65 s/it | 通过并完成保存 |
| Atlas 900 A3 SuperPoD | 2 | 4 | 2 | 关闭 | auto + FLAchunk size 64/ eager | 100 | 1.8095 -> 0.5596 | 5.88 s/it | 通过并完成保存 |
| Atlas 950 SuperPoD | 1 | 8 | 1 | 关闭 | 未配置 kernel plugin / eager | 100 | 1.3575 -> 0.4439 | 2.68 s/it | 通过并完成保存 |
五组训练的 loss 和 grad norm 均保持有限,并完成 100 步及模型保存。同一切分下EP16 eager/fused
的逐步 loss 相关系数为 0.997CP2/EP4/EFSDP2 eager/fused 为 0.977,说明两种 dispatcher 的
优化轨迹一致。性能收益与切分有关EP16 下 fused 比 eager 快约 13%,而加入 CP 和 EFSDP 后 fused
比 eager 慢约 30%,因此不能把 fused 视为所有 mesh 的默认最优选择。
EP16 两组使用 global batch 16 和 cutoff length 256CP2 两组使用 global batch 8 和 cutoff length
128A5 组使用 global batch 8 和 cutoff length 256。因此首末 loss 用于验证各组自身的收敛趋势,
不同切分组之间的绝对 loss 不应直接作为精度等价结论。

View File

@@ -0,0 +1,119 @@
# KTransformers LoRA SFT
KTransformersKT将 MoE routed experts 放在 CPU 执行LLaMA-Factory 继续负责数据、LoRA 参数和训练入口。
当前生产范围是 routed-BF16 LoRA 与 routed-INT8 LoRAAccelerate 配置只负责 FSDP2不再保存 KT 参数。
## 安装检查
必须同时安装带 KT 公共接口的 `ktransformers``transformers-kt``accelerate-kt`。启动前可检查:
```bash
python - <<'PY'
from accelerate import Accelerator
from kt_kernel.sft import resolve_kt_pretrained_artifacts
from transformers import TrainingArguments
assert hasattr(TrainingArguments, "update_kt_config")
assert "adapter_only" in __import__("inspect").signature(Accelerator.get_state_dict).parameters
print(resolve_kt_pretrained_artifacts)
PY
```
## 配置
KT 只有一个用户配置源:训练 YAML。LoRA rank、alpha、dropout 和 runtime capacity 由 LLaMA-Factory
标准字段派生;不要在 `kt_config` 中重复填写。
BF16 示例:
```yaml
finetuning_type: lora
lora_rank: 8
lora_alpha: 16
lora_target: all
use_kt: true
disable_gradient_checkpointing: false
kt_cpu_activation: retain
kt_config:
kt_expert_weight_format: bf16
kt_backend: AMXBF16
kt_num_threads: 96
kt_tp_enabled: true
kt_threadpool_count: 2
kt_max_cache_depth: 2
```
INT8 还需要相互匹配的 routed expert 与 BF16 non-expert cache
```yaml
kt_weight_path: /abs/path/to/routed-int8-experts
kt_non_expert_weight_path: /abs/path/to/bf16-non-expert-cache
kt_config:
kt_expert_weight_format: int8
kt_backend: auto
kt_weight_lifecycle: persistent
```
完整配置见:
- `examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml`
- `examples/ktransformers/train_lora/deepseek_v3_int8_lora_sft_kt.yaml`
Activation 策略:
| `disable_gradient_checkpointing` | `kt_cpu_activation` | CPU / GPU |
| --- | --- | --- |
| `false` | `recompute` 或省略 | recompute / recompute |
| `false` | `retain` | retain / recompute |
| `true` | `retain` 或省略 | retain / retain |
| `true` | `recompute` | 不支持,启动前报错 |
## 启动与复用
```bash
CUDA_VISIBLE_DEVICES=0,1 accelerate launch \
--config_file examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml \
src/train.py examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml
```
输出 adapter 同时包含 standard PEFT 与 fused expert LoRA。
## 新进程加载
对话或评测必须使用本地的完整 KT adapter 目录,并重复训练时的 LoRA 形状配置:`finetuning_type`
`lora_rank``lora_alpha``lora_dropout`,以及相同的 KT base weight 配置。routed INT8 尤其要沿用训练时
`kt_weight_path``kt_non_expert_weight_path`
```yaml
model_name_or_path: /abs/path/to/base-model
adapter_name_or_path: /abs/path/to/output/checkpoint-300
finetuning_type: lora
lora_rank: 8
lora_alpha: 16
lora_dropout: 0.0
use_kt: true
kt_cpu_activation: retain
kt_config:
kt_expert_weight_format: bf16
kt_backend: AMXBF16
kt_num_threads: 96
```
```bash
llamafactory-cli chat path/to/kt_adapter_infer.yaml
llamafactory-cli eval path/to/kt_adapter_eval.yaml
```
目录必须包含 standard PEFT adapter 文件;使用 fused routed-expert LoRA 时,还必须包含
`fused_expert_lora.safetensors``kt_adapter_manifest.json`。LLaMA-Factory 先加载 standard PEFT随后由
KT 校验并恢复 fused artifact。`adapter_folder` 可以选择本地子目录;越出 adapter 根目录的路径和 Hub
adapter ID 会在加载模型前报错Hub bundle 需要先完整下载到本地。
续训应保留原训练 YAML并使用 `resume_from_checkpoint`。分布式 optimizer checkpoint 暂要求相同 world
size。artifact 缺失、hash 不匹配或来源模型不一致时会直接失败,不会退回源 checkpoint。
不要同时启用 Transformers/FSDP activation checkpointing、Unsloth GC也不要把 `kt_config` 放入
Accelerate YAML。每次训练都应确认 loss/grad finite、base model 未修改,并验证 standard/router/fused LoRA
均包含非零更新。

View File

@@ -34,9 +34,11 @@ LlamaFactory 文档
advanced/lora-and-quantization/lora
advanced/lora-and-quantization/quantization
advanced/ktransformers
advanced/distributed/fsdp
advanced/distributed/deepspeed
advanced/distributed/parallel-dp-tp-ep-sp-cp
advanced/distributed/fsdpturbo-ep-efsdp
advanced/custom-kernels/triton
advanced/custom-kernels/fused-operators

View File

@@ -0,0 +1,6 @@
### Install model-specific dependencies: `pip install -r requirements/moss-vl.txt`
model_name_or_path: OpenMOSS-Team/MOSS-VL-Instruct-0708
template: moss_vl
infer_backend: huggingface # choices: [huggingface, vllm, sglang, ktransformers]
trust_remote_code: true

View File

@@ -13,13 +13,3 @@ num_processes: 4 # Adjust based on your GPU count; 4 is suitable for 4 GPUs
rdzv_backend: static
same_network: true
use_cpu: false
kt_config:
enabled: true
kt_backend: AMXBF16 # Use with original BF16 expert weights.
kt_num_threads: 96
kt_tp_enabled: true
kt_threadpool_count: 2
kt_max_cache_depth: 2
kt_share_backward_bb: true
lora_rank: 8

View File

@@ -13,13 +13,3 @@ num_processes: 4 # Adjust based on your GPU count; 4 is suitable for 4 GPUs
rdzv_backend: static
same_network: true
use_cpu: false
kt_config:
enabled: true
kt_backend: AMXINT4 # Use with online-converted INT4 expert weights
kt_num_threads: 96
kt_tp_enabled: true
kt_threadpool_count: 2
kt_max_cache_depth: 2
kt_share_backward_bb: true
lora_rank: 8

View File

@@ -13,13 +13,3 @@ num_processes: 4 # Adjust based on your GPU count; 4 is suitable for 4 GPUs
rdzv_backend: static
same_network: true
use_cpu: false
kt_config:
enabled: true
kt_backend: AMXINT8 # Use with online-converted INT8 expert weights
kt_num_threads: 96
kt_tp_enabled: true
kt_threadpool_count: 2
kt_max_cache_depth: 2
kt_share_backward_bb: true
lora_rank: 8

View File

@@ -13,13 +13,3 @@ num_processes: 1 # Adjust based on your GPU count; 1 is suitable for 1 GPU
rdzv_backend: static
same_network: true
use_cpu: false
kt_config:
enabled: true
kt_backend: AMXINT8 # Use with online-converted INT8 expert weights
kt_num_threads: 96
kt_tp_enabled: true
kt_threadpool_count: 2
kt_max_cache_depth: 2
kt_share_backward_bb: true
lora_rank: 8

View File

@@ -13,13 +13,3 @@ num_processes: 8 # Adjust based on your GPU count; 8 is suitable for 8 GPUs
rdzv_backend: static
same_network: true
use_cpu: false
kt_config:
enabled: true
kt_backend: AMXINT8 # Use with online-converted INT8 expert weights
kt_num_threads: 96
kt_tp_enabled: true
kt_threadpool_count: 2
kt_max_cache_depth: 2
kt_share_backward_bb: true
lora_rank: 8

View File

@@ -0,0 +1,54 @@
### model
model_name_or_path: /path/to/DeepSeek-V3.1-source
trust_remote_code: true
### method
stage: sft
do_train: true
finetuning_type: lora
lora_rank: 8
lora_alpha: 16
lora_target: all
### dataset
dataset: identity, alpaca_en_demo
template: deepseek3
cutoff_len: 2048
max_samples: 100000
overwrite_cache: true
preprocessing_num_workers: 16
dataloader_num_workers: 4
### output
output_dir: saves/KT_FT_deepseekV3_int8
logging_steps: 10
save_steps: 500
plot_loss: true
overwrite_output_dir: true
save_only_model: false
report_to: none
### train
per_device_train_batch_size: 1
gradient_accumulation_steps: 1
learning_rate: 1.0e-4
num_train_epochs: 3.0
lr_scheduler_type: cosine
warmup_ratio: 0.1
bf16: true
ddp_timeout: 180000000
### ktransformers
use_kt: true
kt_cpu_activation: retain
kt_weight_path: /path/to/routed-int8-experts
kt_non_expert_weight_path: /path/to/bf16-non-expert-cache
kt_config:
kt_expert_weight_format: int8
kt_backend: auto
kt_weight_lifecycle: persistent
kt_num_threads: 96
kt_tp_enabled: true
kt_threadpool_count: 2
kt_max_cache_depth: 2
kt_share_backward_bb: true

View File

@@ -40,6 +40,13 @@ resume_from_checkpoint: null
### ktransformers
use_kt: true
# Pair with fsdp2_kt_bf16.yaml for original BF16 checkpoints.
# For pre-converted expert weights, uncomment kt_weight_path and use fsdp2_kt_int8.yaml or fsdp2_kt_int4.yaml.
# kt_weight_path: /path/to/DeepSeek-V3-AMXINT8
kt_cpu_activation: retain
kt_config:
kt_expert_weight_format: bf16
kt_backend: AMXBF16
kt_num_threads: 96
kt_tp_enabled: true
kt_threadpool_count: 2
kt_max_cache_depth: 2
kt_share_backward_bb: true
# The Accelerate YAML contains FSDP settings only. KT has a single configuration owner here.

View File

@@ -40,7 +40,14 @@ resume_from_checkpoint: null
### ktransformers
use_kt: true
# For original BF16 checkpoints, start with examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml.
# For pre-converted expert weights, uncomment kt_weight_path and use fsdp2_kt_int8.yaml or fsdp2_kt_int4.yaml.
# Pair the 397B path with fsdp2_kt_int8.yaml, tune cutoff_len to prepared weights and GPU memory.
# kt_weight_path: /path/to/Qwen3.5-MoE-AMXINT8
kt_cpu_activation: retain
kt_config:
kt_expert_weight_format: bf16
kt_backend: AMXBF16
kt_num_threads: 96
kt_tp_enabled: true
kt_threadpool_count: 2
kt_max_cache_depth: 2
kt_model_max_length: 2176 # Includes the text-only template's dummy-image tokens.
kt_share_backward_bb: true
# The Accelerate YAML contains FSDP settings only. KT has a single configuration owner here.

View File

@@ -0,0 +1,14 @@
### Install model-specific dependencies: `pip install -r requirements/moss-vl.txt`
### Note: DO NOT use quantized model or quantization_bit when merging lora adapters
### model
model_name_or_path: OpenMOSS-Team/MOSS-VL-Instruct-0708
adapter_name_or_path: saves/moss-vl-11b/lora/sft
template: moss_vl
trust_remote_code: true
### export
export_dir: saves/moss_vl_sft_merged
export_size: 5
export_device: cpu # choices: [cpu, auto]
export_legacy_format: false

View File

@@ -0,0 +1,57 @@
### Install model-specific dependencies: `pip install -r requirements/moss-vl.txt`
### model
model_name_or_path: OpenMOSS-Team/MOSS-VL-Instruct-0708
image_max_pixels: 262144
video_max_pixels: 16384
video_fps: 1.0
video_maxlen: 256
use_reentrant_gc: false
trust_remote_code: true
### method
stage: sft
do_train: true
finetuning_type: full
freeze_vision_tower: true
freeze_multi_modal_projector: true
freeze_language_model: false
deepspeed: examples/deepspeed/ds_z3_config.json
### dataset
dataset: mllm_demo,identity,alpaca_en_demo # video: mllm_video_demo
template: moss_vl
cutoff_len: 4096
max_samples: 1000
preprocessing_num_workers: 16
dataloader_num_workers: 4
packing: false
### output
output_dir: saves/moss-vl-11b/full/sft
logging_steps: 10
save_steps: 500
plot_loss: true
overwrite_output_dir: true
save_only_model: false
report_to: none # choices: [none, wandb, tensorboard, swanlab, mlflow]
### train
per_device_train_batch_size: 1
gradient_accumulation_steps: 1
gradient_checkpointing: true
gradient_checkpointing_kwargs:
use_reentrant: false
learning_rate: 1.0e-5
num_train_epochs: 3.0
lr_scheduler_type: cosine
warmup_ratio: 0.1
bf16: true
ddp_timeout: 180000000
resume_from_checkpoint: null
### eval
# val_size: 0.1
# per_device_eval_batch_size: 1
# eval_strategy: steps
# eval_steps: 500

View File

@@ -0,0 +1,54 @@
### Install model-specific dependencies: `pip install -r requirements/moss-vl.txt`
### model
model_name_or_path: OpenMOSS-Team/MOSS-VL-Instruct-0708
image_max_pixels: 262144
video_max_pixels: 16384
video_fps: 1.0
video_maxlen: 256
trust_remote_code: true
### method
stage: sft
do_train: true
finetuning_type: lora
lora_rank: 8
lora_target: all
freeze_vision_tower: true
freeze_multi_modal_projector: true
freeze_language_model: false
### dataset
dataset: mllm_demo,identity,alpaca_en_demo # video: mllm_video_demo
template: moss_vl
cutoff_len: 4096
max_samples: 1000
preprocessing_num_workers: 16
dataloader_num_workers: 4
packing: false
### output
output_dir: saves/moss-vl-11b/lora/sft
logging_steps: 10
save_steps: 500
plot_loss: true
overwrite_output_dir: true
save_only_model: false
report_to: none # choices: [none, wandb, tensorboard, swanlab, mlflow]
### train
per_device_train_batch_size: 2
gradient_accumulation_steps: 1
learning_rate: 1.0e-4
num_train_epochs: 3.0
lr_scheduler_type: cosine
warmup_ratio: 0.1
bf16: true
ddp_timeout: 180000000
resume_from_checkpoint: null
### eval
# val_size: 0.1
# per_device_eval_batch_size: 1
# eval_strategy: steps
# eval_steps: 500

View File

@@ -0,0 +1,32 @@
model: Qwen/Qwen3.5-35B-A3B
model_class: llm
kernel_config:
name: auto, flash-linear-attention
include_kernels: chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
chunk_size: 64
dist_config:
name: fsdpturbo
ep_size: 16
ep_dispatcher: eager
cp_size: 1
init_config:
name: init_on_meta
### data
train_dataset: data/v1_sft_demo.yaml
### training
output_dir: outputs/Qwen3.5-35B-A3B/full/sft
micro_batch_size: 1
cutoff_len: 256
learning_rate: 1.0e-4
bf16: true
max_steps: 5
### sample
sample_backend: hf
max_new_tokens: 128

View File

@@ -0,0 +1,27 @@
model: Qwen/Qwen3.5-0.8B
model_class: llm
kernel_config:
name: auto
quant_config: null
dist_config:
name: fsdp2
dcp_path: null
### data
train_dataset: data/v1_multimodal_demo.yaml
### training
output_dir: outputs/test_multimodal
micro_batch_size: 1
cutoff_len: 2048
learning_rate: 1.0e-4
max_steps: 5
### sample
sample_backend: hf
max_new_tokens: 128

View File

@@ -0,0 +1,2 @@
# Pin the FSDPTurbo API to a reproducible upstream main revision.
fsdp-turbo @ git+https://gitcode.com/Ascend/FSDPTurbo.git@d878dffdb1e0312dc098599f2b56810d6b592ee2

3
requirements/moss-vl.txt Normal file
View File

@@ -0,0 +1,3 @@
transformers==4.57.1
torchcodec==0.7.0
joblib

View File

@@ -1,5 +1,5 @@
torch==2.7.1
torch-npu==2.7.1.post4
torchvision==0.22.1
torchaudio==2.7.1
torch==2.10.0
torch-npu==2.10.0.post2
torchvision==0.25.0
torchaudio==2.10.0
decorator

View File

@@ -150,7 +150,9 @@ class MultiModalDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):
if isinstance(self.model, PeftModel):
self.model = self.model.base_model.model
if self.model is not None and hasattr(self.model, "get_rope_index"): # for qwen2vl mrope
if getattr(getattr(self.model, "config", None), "model_type", None) == "moss_vl":
self.get_rope_func = None # MOSS-VL computes its own XRoPE positions in model.forward.
elif self.model is not None and hasattr(self.model, "get_rope_index"): # for qwen2vl mrope
self.get_rope_func = self.model.get_rope_index # transformers < 4.52.0 or qwen2.5 omni
elif self.model is not None and hasattr(self.model, "model") and hasattr(self.model.model, "get_rope_index"):
self.get_rope_func = self.model.model.get_rope_index # transformers >= 4.52.0
@@ -322,6 +324,8 @@ class MultiModalDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):
)
def __call__(self, features: list[dict[str, Any]]) -> dict[str, "torch.Tensor"]:
model_type = getattr(getattr(self.model, "config", None), "model_type", None)
is_moss_vl = model_type == "moss_vl"
batch_images, batch_videos, batch_audios = [], [], []
batch_imglens, batch_vidlens, batch_audlens, batch_input_ids = [], [], [], []
packing_params_list: list[dict[str, Any] | None] = []
@@ -341,7 +345,10 @@ class MultiModalDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):
fake_input_ids = []
has_dummy_image = False
if (
self.template.mm_plugin.image_token is not None and sum(batch_imglens) == 0 and sum(batch_vidlens) == 0
self.template.mm_plugin.image_token is not None
and sum(batch_imglens) == 0
and sum(batch_vidlens) == 0
and not is_moss_vl # MOSS-VL builds one native zero-valued dummy per text-only sample in its plugin.
): # avoid process hanging in zero3/fsdp case
fake_messages = [{"role": "user", "content": IMAGE_PLACEHOLDER}]
fake_images = [Image.new("RGB", (64, 64), (255, 255, 255))]
@@ -416,7 +423,6 @@ class MultiModalDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):
features: dict[str, torch.Tensor] = super().__call__(features)
bsz, seq_len = features["input_ids"].shape[:2]
model_type = getattr(self.model.config, "model_type", None) if self.model is not None else None
is_omni = model_type in [
"qwen2_5_omni_thinker",
"qwen3_omni_moe_thinker",
@@ -461,12 +467,17 @@ class MultiModalDataCollatorForSeq2Seq(DataCollatorForSeq2Seq):
):
raise ValueError(f"{self.model.config.model_type} requires 3D position ids for mrope.")
if "cross_attention_mask" in mm_inputs: # for mllama inputs when pad_to_multiple_of is enabled
if (
"cross_attention_mask" in mm_inputs and mm_inputs["cross_attention_mask"].dtype != torch.bool
): # for mllama inputs when pad_to_multiple_of is enabled
cross_attention_mask = mm_inputs.pop("cross_attention_mask")
seq_len = features["input_ids"].size(1)
orig_len = cross_attention_mask.size(1)
mm_inputs["cross_attention_mask"] = F.pad(cross_attention_mask, (0, 0, 0, 0, 0, seq_len - orig_len))
if is_moss_vl:
mm_inputs = self.template.mm_plugin.post_process_mossvl_inputs(features, mm_inputs, self.processor)
features.update(mm_inputs)
if "image_bound" in features: # for minicpmv inputs
@@ -530,6 +541,17 @@ class SFTDataCollatorWith4DAttentionMask(MultiModalDataCollatorForSeq2Seq):
self._unpad_packed_features(features)
features["attention_mask"] = None # let transformers handle causal packed mask.
else:
# `DataCollatorForSeq2Seq(pad_to_multiple_of=...)` pads `input_ids`/`attention_mask`
# but leaves `position_ids` untouched (it is not in `model_input_names`). On the
# non-FA2 packing path we do not unpad, so `position_ids` stays shorter than
# `input_ids`, which makes cos/sin shorter than query and crashes
# `apply_rotary_pos_emb`. Right-pad `position_ids` to the padded length to match.
position_ids = features.get("position_ids")
if torch.is_tensor(position_ids):
pad_len = features["input_ids"].shape[-1] - position_ids.shape[-1]
if pad_len > 0:
features["position_ids"] = F.pad(position_ids, (0, pad_len), value=0)
for key, value in features.items(): # cast data dtype for paligemma
if torch.is_tensor(value) and torch.is_floating_point(value):

View File

@@ -472,6 +472,354 @@ class BasePlugin(MMPluginMixin):
return self._get_mm_inputs(images, videos, audios, processor)
@dataclass
class MossVLPlugin(BasePlugin):
vision_bos_token: str = "<|vision_start|>"
vision_eos_token: str = "<|vision_end|>"
time_bos_token: str = "<|time_start|>"
time_eos_token: str = "<|time_end|>"
@staticmethod
def _split_pixel_values(
pixel_values: "torch.Tensor",
grid_thw: "torch.Tensor",
) -> list["torch.Tensor"]:
patch_counts = [int(grid.prod().item()) for grid in grid_thw]
return list(torch.split(pixel_values, patch_counts))
@staticmethod
def _create_cross_attention_mask(
input_ids: Union[list[list[int]], "torch.Tensor"],
grid_thw: "torch.Tensor",
media_nums_per_sample: list[int],
image_token_id: int,
attention_mask: Optional["torch.Tensor"] = None,
padding_side: Literal["left", "right"] = "right",
) -> "torch.Tensor":
r"""Create the native MOSS-VL frame-level causal cross-attention mask."""
if isinstance(input_ids, list):
max_text_len = max(len(token_ids) for token_ids in input_ids)
input_ids_tensor = torch.full((len(input_ids), max_text_len), -1, dtype=torch.long)
attention_mask_tensor = torch.zeros_like(input_ids_tensor, dtype=torch.bool)
for batch_index, token_ids in enumerate(input_ids):
seq_len = len(token_ids)
start = max_text_len - seq_len if padding_side == "left" else 0
input_ids_tensor[batch_index, start : start + seq_len] = torch.tensor(token_ids, dtype=torch.long)
attention_mask_tensor[batch_index, start : start + seq_len] = True
else:
input_ids_tensor = input_ids
attention_mask_tensor = (
torch.ones_like(input_ids_tensor, dtype=torch.bool)
if attention_mask is None
else attention_mask.bool()
)
total_frames_per_sample = []
media_index = 0
for num_media in media_nums_per_sample:
sample_grid = grid_thw[media_index : media_index + num_media]
total_frames_per_sample.append(int(sample_grid[:, 0].sum().item()))
media_index += num_media
max_num_frames = max(total_frames_per_sample)
frame_indices = torch.arange(max_num_frames, device=input_ids_tensor.device).view(1, 1, -1)
visible_mask = (input_ids_tensor == image_token_id).cumsum(dim=1).unsqueeze(-1) > frame_indices
visible_mask &= attention_mask_tensor.unsqueeze(-1)
valid_frames = frame_indices < torch.tensor(
total_frames_per_sample,
device=input_ids_tensor.device,
).view(-1, 1, 1)
visible_mask &= valid_frames
return (~visible_mask).unsqueeze(1)
def _get_video_inputs(
self,
videos: list["VideoInput"],
processor: "MMProcessor",
return_metadata: bool,
) -> dict[str, Any]:
video_kwargs = {"return_tensors": "pt", "return_metadata": return_metadata}
if getattr(processor, "video_fps", None) is not None:
video_kwargs["video_fps"] = processor.video_fps
if getattr(processor, "video_maxlen", None) is not None:
video_kwargs["max_frames"] = processor.video_maxlen
video_min_pixels = getattr(processor, "video_min_pixels", None)
video_max_pixels = getattr(processor, "video_max_pixels", None)
if video_min_pixels is not None and video_max_pixels is not None:
video_kwargs["size"] = {
"shortest_edge": video_min_pixels,
"longest_edge": video_max_pixels,
}
return dict(processor.video_processor(videos=videos, **video_kwargs))
def _get_media_order_from_ids(
self,
input_ids: list[int],
processor: "MMProcessor",
num_images: int,
num_videos: int,
expected_video_frames: Optional[list[int]] = None,
) -> list[str]:
media_order = []
video_frame_counts = []
in_video = False
current_video_frames = 0
for token_id in input_ids:
if token_id == processor.vision_start_token_id:
if in_video:
raise ValueError(
"MOSS-VL encountered nested video token blocks after tokenization. "
"Please increase `cutoff_len` if a video placeholder was truncated."
)
media_order.append("video")
in_video = True
current_video_frames = 0
elif token_id == processor.vision_end_token_id:
if not in_video:
raise ValueError(
"MOSS-VL encountered a video end token without a matching start token after tokenization. "
"Please increase `cutoff_len` if a video placeholder was truncated."
)
video_frame_counts.append(current_video_frames)
in_video = False
elif token_id == processor.image_token_id:
if in_video:
current_video_frames += 1
else:
media_order.append("image")
if in_video:
raise ValueError(
"MOSS-VL encountered an incomplete video token block after tokenization. "
"Please increase `cutoff_len` or reduce `video_maxlen`."
)
if media_order.count("image") != num_images or media_order.count("video") != num_videos:
raise ValueError(
"MOSS-VL media tokens do not match the provided media after tokenization: "
f"order={media_order}, images={num_images}, videos={num_videos}. "
"Please increase `cutoff_len` if a visual placeholder was truncated."
)
if expected_video_frames is not None and video_frame_counts != expected_video_frames:
raise ValueError(
"MOSS-VL video frame tokens do not match the processed video after tokenization: "
f"tokens={video_frame_counts}, frames={expected_video_frames}. "
"Please increase `cutoff_len` or reduce `video_maxlen`."
)
return media_order
@override
def process_messages(
self,
messages: list[dict[str, str]],
images: list["ImageInput"],
videos: list["VideoInput"],
audios: list["AudioInput"],
processor: Optional["MMProcessor"],
) -> list[dict[str, str]]:
self._validate_input(processor, images, videos, audios)
self._validate_messages(messages, images, videos, audios)
messages = deepcopy(messages)
video_inputs = self._get_video_inputs(videos, processor, return_metadata=True) if videos else {}
video_grid_thw = video_inputs.get("video_grid_thw", [])
video_metadata = video_inputs.get("video_metadata", [])
video_index = 0
for message in messages:
content = message["content"]
content = content.replace(IMAGE_PLACEHOLDER, self.image_token)
while VIDEO_PLACEHOLDER in content:
metadata = video_metadata[video_index]
if metadata.fps is None:
metadata.fps = 24
timestamps = processor._calculate_timestamps(
metadata.frames_indices,
metadata.total_num_frames,
metadata.fps,
metadata.duration,
processor.video_processor.temporal_patch_size,
actual_timestamps=getattr(metadata, "actual_timestamps", None),
)
num_frames = int(video_grid_thw[video_index][0].item())
frame_tokens = [
f"{self.time_bos_token}{timestamps[frame_idx]:.1f} seconds{self.time_eos_token}{self.image_token}"
for frame_idx in range(num_frames)
]
video_tokens = f"{self.vision_bos_token}{''.join(frame_tokens)}{self.vision_eos_token}"
content = content.replace(VIDEO_PLACEHOLDER, video_tokens, 1)
video_index += 1
message["content"] = content
return messages
@override
def get_mm_inputs(
self,
images: list["ImageInput"],
videos: list["VideoInput"],
audios: list["AudioInput"],
imglens: list[int],
vidlens: list[int],
audlens: list[int],
batch_ids: list[list[int]],
processor: Optional["MMProcessor"],
) -> dict[str, Union[list[int], "torch.Tensor"]]:
self._validate_input(processor, images, videos, audios)
if audios:
raise ValueError("MOSS-VL does not support audio inputs.")
if not (len(imglens) == len(vidlens) == len(batch_ids)):
raise ValueError("MOSS-VL batch metadata must have one entry per sample.")
final_pixel_values = []
final_grid_thw = []
media_nums_per_sample = []
image_offset = 0
video_offset = 0
for imglen, vidlen, input_ids in zip(imglens, vidlens, batch_ids):
sample_images = images[image_offset : image_offset + imglen]
sample_videos = videos[video_offset : video_offset + vidlen]
image_offset += imglen
video_offset += vidlen
image_chunks, image_grids = [], []
if sample_images:
regularized_images = self._regularize_images(
sample_images,
image_max_pixels=2**63 - 1,
image_min_pixels=1,
)["images"]
image_kwargs = {"return_tensors": "pt"}
if getattr(processor, "image_min_pixels", None) is not None:
image_kwargs["min_pixels"] = processor.image_min_pixels
if getattr(processor, "image_max_pixels", None) is not None:
image_kwargs["max_pixels"] = processor.image_max_pixels
image_inputs = processor.image_processor(images=regularized_images, **image_kwargs)
image_grids = list(image_inputs["image_grid_thw"])
image_chunks = self._split_pixel_values(image_inputs["pixel_values"], image_inputs["image_grid_thw"])
video_chunks, video_grids = [], []
if sample_videos:
video_inputs = self._get_video_inputs(sample_videos, processor, return_metadata=False)
video_grids = list(video_inputs["video_grid_thw"])
video_chunks = self._split_pixel_values(
video_inputs["pixel_values_videos"],
video_inputs["video_grid_thw"],
)
media_order = self._get_media_order_from_ids(
input_ids,
processor,
imglen,
vidlen,
expected_video_frames=[int(grid[0].item()) for grid in video_grids],
)
if not media_order:
patch_size = getattr(processor.image_processor, "patch_size", None)
if patch_size is None: # lightweight/test processors without the native MOSS-VL contract
blank_image = Image.new("RGB", (128, 128), (255, 255, 255))
blank_inputs = processor.image_processor(images=[blank_image], return_tensors="pt")
final_pixel_values.append(blank_inputs["pixel_values"])
final_grid_thw.append(blank_inputs["image_grid_thw"][0])
else:
temporal_patch_size = getattr(processor.image_processor, "temporal_patch_size", None) or 1
merge_size = getattr(processor.image_processor, "merge_size", None) or 2
factor = patch_size * merge_size
side = math.ceil(128 / factor) * factor
grid_thw = torch.tensor([1, side // patch_size, side // patch_size])
feature_dim = 3 * temporal_patch_size * patch_size * patch_size
final_pixel_values.append(torch.zeros((int(grid_thw.prod()), feature_dim), dtype=torch.float32))
final_grid_thw.append(grid_thw)
media_nums_per_sample.append(1)
continue
image_index = 0
video_index = 0
for modality in media_order:
if modality == "image":
final_pixel_values.append(image_chunks[image_index])
final_grid_thw.append(image_grids[image_index])
image_index += 1
else:
final_pixel_values.append(video_chunks[video_index])
final_grid_thw.append(video_grids[video_index])
video_index += 1
media_nums_per_sample.append(len(media_order))
if image_offset != len(images) or video_offset != len(videos):
raise ValueError("MOSS-VL media lengths do not consume all provided inputs.")
mm_inputs = {
"pixel_values": torch.cat(final_pixel_values, dim=0),
"grid_thw": torch.stack(final_grid_thw),
"media_nums_per_sample": media_nums_per_sample,
}
mm_inputs["cross_attention_mask"] = self._create_cross_attention_mask(
batch_ids,
mm_inputs["grid_thw"],
media_nums_per_sample,
processor.image_token_id,
padding_side=processor.tokenizer.padding_side,
)
return mm_inputs
def post_process_mossvl_inputs(
self,
features: dict[str, "torch.Tensor"],
mm_inputs: dict[str, Any],
processor: "MMProcessor",
) -> dict[str, Any]:
r"""Create MOSS-VL batch-only inputs after the text batch has been padded."""
input_ids = features["input_ids"]
attention_mask = features["attention_mask"].bool()
mm_inputs["cross_attention_mask"] = self._create_cross_attention_mask(
input_ids,
mm_inputs["grid_thw"],
mm_inputs["media_nums_per_sample"],
processor.image_token_id,
attention_mask,
)
dummy_image_tokens = (input_ids == processor.image_token_id) & ~attention_mask
input_ids.masked_fill_(dummy_image_tokens, processor.tokenizer.pad_token_id)
labels = features.get("labels")
if labels is not None:
control_token_ids = {
processor.image_token_id,
processor.video_token_id,
processor.vision_start_token_id,
processor.vision_end_token_id,
processor.tokenizer.convert_tokens_to_ids(self.time_bos_token),
processor.tokenizer.convert_tokens_to_ids(self.time_eos_token),
}
for batch_index, token_ids in enumerate(input_ids):
in_vision = False
for token_index, token_id in enumerate(token_ids.tolist()):
if token_id == processor.vision_start_token_id:
in_vision = True
if in_vision or token_id in control_token_ids:
labels[batch_index, token_index] = IGNORE_INDEX
if token_id == processor.vision_end_token_id:
in_vision = False
# Native MOSS-VL labels_spans supervise through <|im_end|>, but not its trailing newline.
im_end_token_id = processor.tokenizer.convert_tokens_to_ids("<|im_end|>")
labels[:, 1:].masked_fill_(input_ids[:, :-1] == im_end_token_id, IGNORE_INDEX)
features.pop("position_ids", None)
return mm_inputs
@dataclass
class ErnieVLPlugin(BasePlugin):
@override
@@ -2911,6 +3259,7 @@ PLUGINS = {
"minicpm_v": MiniCPMVPlugin,
"minicpm_v_4_6": MiniCPMV4_6Plugin,
"mllama": MllamaPlugin,
"moss_vl": MossVLPlugin,
"paligemma": PaliGemmaPlugin,
"pixtral": PixtralPlugin,
"qwen2_audio": Qwen2AudioPlugin,

View File

@@ -333,6 +333,50 @@ class Template:
return modelfile
@dataclass
class MossVLTemplate(Template):
@override
def _encode(
self,
tokenizer: "PreTrainedTokenizer",
messages: list[dict[str, str]],
system: Optional[str],
tools: Optional[str],
) -> list[list[int]]:
system = system or self.default_system
encoded_messages = []
for i, message in enumerate(messages):
elements = []
if i == 0:
elements += self.format_prefix.apply()
if system or tools:
tool_text = self.format_tools.apply(content=tools)[0] if tools else ""
if tools and not system:
tool_text = tool_text.lstrip("\n")
elements += self.format_system.apply(content=(system + tool_text))
if message["role"] == Role.USER:
elements += self.format_user.apply(content=message["content"], idx=str(i // 2))
elif message["role"] == Role.ASSISTANT:
elements += self.format_assistant.apply(content=message["content"])
elif message["role"] == Role.OBSERVATION:
elements += self.format_observation.apply(content=message["content"])
elif message["role"] == Role.FUNCTION:
elements += self.format_function.apply(
content=message["content"],
thought_words=self.thought_words,
tool_call_words=self.tool_call_words,
)
else:
raise NotImplementedError("Unexpected role: {}".format(message["role"]))
encoded_messages.append(self._convert_elements_to_ids(tokenizer, elements))
return encoded_messages
@dataclass
class Llama2Template(Template):
r"""A template that fuse the system message to first user message."""
@@ -1526,6 +1570,32 @@ register_template(
)
# copied from qwen template
register_template(
name="moss_vl",
format_user=StringFormatter(slots=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]),
format_assistant=StringFormatter(slots=["{{content}}<|im_end|>\n"]),
format_system=StringFormatter(slots=["<|im_start|>system\n{{content}}<|im_end|>\n"]),
format_function=FunctionFormatter(slots=["{{content}}<|im_end|>\n"], tool_format="qwen"),
format_observation=StringFormatter(
slots=["<|im_start|>user\n<tool_response>\n{{content}}\n</tool_response><|im_end|>\n<|im_start|>assistant\n"]
),
format_tools=ToolFormatter(tool_format="qwen"),
stop_words=["<|im_end|>"],
replace_eos=True,
mm_plugin=get_mm_plugin(
name="moss_vl",
image_token="<|image_pad|>",
video_token="<|video_pad|>",
vision_bos_token="<|vision_start|>",
vision_eos_token="<|vision_end|>",
time_bos_token="<|time_start|>",
time_eos_token="<|time_end|>",
),
template_class=MossVLTemplate,
)
# copied from vicuna template
register_template(
name="llava",

View File

@@ -2198,6 +2198,17 @@ register_model_group(
)
register_model_group(
models={
"MOSS-VL-Instruct-0708": {
DownloadSource.DEFAULT: "OpenMOSS-Team/MOSS-VL-Instruct-0708",
},
},
template="moss_vl",
multimodal=True,
)
register_model_group(
models={
"OLMo-1B": {

View File

@@ -470,10 +470,23 @@ class KTransformersArguments:
default=False,
metadata={"help": "Whether to use KTransformers AMX MoE backend for SFT training."},
)
kt_cpu_activation: Literal["retain", "recompute"] | None = field(
default=None,
metadata={
"help": (
"Whether KTransformers retains CPU expert activations. Defaults to recompute while GPU "
"gradient checkpointing is enabled and retain otherwise."
)
},
)
kt_weight_path: str | None = field(
default=None,
metadata={"help": "Path to pre-quantized INT8 expert weights (.kt files)."},
)
kt_non_expert_weight_path: str | None = field(
default=None,
metadata={"help": "Path to the KT BF16 non-expert weight cache used with routed INT8 experts."},
)
kt_expert_checkpoint_path: str | None = field(
default=None,
metadata={"help": "Path to expert checkpoint (safetensors) for online conversion."},
@@ -490,52 +503,202 @@ class KTransformersArguments:
default=None,
metadata={"help": "Intermediate size for GPU-side LoRA Experts."},
)
_kt_inference_config: dict[str, Any] | None = field(default=None, init=False, repr=False)
_kt_config_handle: Any = field(default=None, init=False, repr=False)
_kt_adapter_artifact_path: str | None = field(default=None, init=False, repr=False)
def get_kt_config_dict(self, finetuning_args: Any, model_max_length: int | None) -> dict[str, Any]:
r"""Build KT config values from LLaMA-Factory model and LoRA arguments."""
kt_config = {
"kt_lora_rank": getattr(finetuning_args, "lora_rank", None),
"kt_lora_alpha": getattr(finetuning_args, "lora_alpha", None),
"kt_weight_path": self.kt_weight_path,
"kt_expert_checkpoint_path": self.kt_expert_checkpoint_path,
"kt_model_max_length": model_max_length,
"kt_use_lora_experts": self.kt_use_lora_experts,
"kt_lora_expert_num": self.kt_lora_expert_num,
"kt_lora_expert_intermediate_size": self.kt_lora_expert_intermediate_size,
_KT_DERIVED_KEYS = frozenset(
{
"enabled",
"kt_activation_policy",
"kt_expert_checkpoint_path",
"kt_full_weight_grad",
"kt_lora_alpha",
"kt_lora_dropout",
"kt_lora_expert_intermediate_size",
"kt_lora_expert_num",
"kt_lora_rank",
"kt_non_expert_weight_path",
"kt_skip_expert_loading",
"kt_train_mode",
"kt_use_lora_experts",
"kt_weight_path",
}
)
def __post_init__(self) -> None:
if self.kt_cpu_activation not in {None, "retain", "recompute"}:
raise ValueError("`kt_cpu_activation` must be `retain` or `recompute`.")
if not self.use_kt and self.kt_cpu_activation is not None:
raise ValueError("`kt_cpu_activation` is only valid when `use_kt: true`.")
def get_kt_activation_policy(self) -> dict[str, str]:
r"""Resolve LF's GPU checkpoint switch and KT's CPU activation setting."""
gpu_activation = "retain" if self.disable_gradient_checkpointing else "recompute"
cpu_activation = self.kt_cpu_activation or gpu_activation
if cpu_activation == "recompute" and gpu_activation == "retain":
raise ValueError(
"`kt_cpu_activation: recompute` requires GPU gradient checkpointing. "
"Set `disable_gradient_checkpointing: false` or use `kt_cpu_activation: retain`."
)
return {"cpu": cpu_activation, "gpu": gpu_activation}
@staticmethod
def _get_accelerator_kt_config(training_args: Any) -> Any:
accelerator_config = getattr(training_args, "accelerator_config", None)
if isinstance(accelerator_config, dict):
return accelerator_config.get("kt_config")
return getattr(accelerator_config, "kt_config", None)
def _normalize_advanced_kt_config(self, raw_config: Any) -> dict[str, Any]:
if raw_config is None:
return {}
if not isinstance(raw_config, dict):
raise TypeError("LLaMA-Factory `kt_config` must be a flat mapping.")
config = dict(raw_config)
conflicts = sorted(set(config) & self._KT_DERIVED_KEYS)
if conflicts:
raise ValueError(f"These `kt_config` values are derived from LLaMA-Factory arguments: {conflicts}.")
return config
def _get_advanced_kt_config(self, training_args: Any) -> dict[str, Any]:
raw_config = getattr(training_args, "kt_config", None)
accelerator_config = self._get_accelerator_kt_config(training_args)
if raw_config is None:
if accelerator_config is not None:
raise ValueError(
"Put KTransformers settings in the LLaMA-Factory training YAML `kt_config`; "
"remove `kt_config` from the Accelerate config."
)
return {}
if accelerator_config is not None and accelerator_config != raw_config:
raise ValueError("LLaMA-Factory YAML and Accelerate config cannot define different KT settings.")
return self._normalize_advanced_kt_config(raw_config)
def configure_kt_checkpointing(self, training_args: Any) -> None:
r"""Keep LLaMA-Factory as the single gradient-checkpointing entry point."""
if self.use_unsloth or self.use_unsloth_gc:
raise ValueError("KTransformers cannot be combined with Unsloth checkpoint wrapping.")
if getattr(training_args, "gradient_checkpointing", False):
raise ValueError(
"KTransformers uses LLaMA-Factory's `disable_gradient_checkpointing`; "
"remove `gradient_checkpointing: true`."
)
if getattr(training_args, "gradient_checkpointing_kwargs", None) is not None:
raise ValueError("KTransformers supplies its checkpoint context; remove `gradient_checkpointing_kwargs`.")
fsdp_config = getattr(training_args, "fsdp_config", None)
if isinstance(fsdp_config, dict) and fsdp_config.get("activation_checkpointing"):
raise ValueError("Disable FSDP activation checkpointing when using KTransformers.")
if os.environ.get("FSDP_ACTIVATION_CHECKPOINTING", "false").lower() in {"1", "true", "yes"}:
raise ValueError("Disable FSDP activation checkpointing when using KTransformers.")
self.get_kt_activation_policy()
if not self.disable_gradient_checkpointing:
self.use_reentrant_gc = False
training_args.gradient_checkpointing = False
training_args.gradient_checkpointing_kwargs = None
def get_kt_config_dict(
self,
finetuning_args: Any,
model_max_length: int | None,
advanced_config: dict[str, Any] | None = None,
) -> dict[str, Any]:
r"""Map LLaMA-Factory-owned training values to the public KT configuration."""
if getattr(finetuning_args, "finetuning_type", None) != "lora":
raise ValueError("KTransformers thin integration currently supports LoRA finetuning only.")
kt_config = dict(advanced_config or {})
configured_capacity = kt_config.pop("kt_model_max_length", None)
if configured_capacity is not None:
try:
configured_capacity = int(configured_capacity)
except (TypeError, ValueError) as exc:
raise ValueError("`kt_model_max_length` must be a positive integer.") from exc
if configured_capacity <= 0:
raise ValueError("`kt_model_max_length` must be a positive integer.")
kt_config.update(
{
"kt_lora_rank": getattr(finetuning_args, "lora_rank", None),
"kt_lora_alpha": getattr(finetuning_args, "lora_alpha", None),
"kt_lora_dropout": getattr(finetuning_args, "lora_dropout", None),
"kt_weight_path": self.kt_weight_path,
"kt_non_expert_weight_path": self.kt_non_expert_weight_path,
"kt_expert_checkpoint_path": self.kt_expert_checkpoint_path,
"kt_model_max_length": max(model_max_length or 0, configured_capacity or 0) or None,
"kt_use_lora_experts": self.kt_use_lora_experts,
"kt_lora_expert_num": self.kt_lora_expert_num,
"kt_lora_expert_intermediate_size": self.kt_lora_expert_intermediate_size,
"kt_activation_policy": self.get_kt_activation_policy(),
"kt_train_mode": "lora",
"kt_full_weight_grad": False,
}
)
return {key: value for key, value in kt_config.items() if value is not None}
def _resolve_kt_adapter_artifact_dir(self, operation: str) -> str | None:
if not self.adapter_name_or_path:
return None
if len(self.adapter_name_or_path) != 1:
raise ValueError("KTransformers accepts a single `adapter_name_or_path`.")
adapter_root = os.path.realpath(os.path.expanduser(self.adapter_name_or_path[0]))
adapter_dir = adapter_root
if self.adapter_folder:
adapter_dir = os.path.realpath(os.path.join(adapter_root, self.adapter_folder))
if os.path.commonpath((adapter_root, adapter_dir)) != adapter_root:
raise ValueError("`adapter_folder` must stay inside the KT adapter directory.")
if not os.path.isdir(adapter_dir):
raise ValueError(f"KTransformers {operation} requires a local adapter directory.")
return adapter_dir
def apply_kt_config(self, finetuning_args: Any, training_args: Any, model_max_length: int | None) -> None:
r"""Apply LLaMA-Factory KT args to transformers/accelerate KT integration points."""
if not self.use_kt:
return
kt_config = self.get_kt_config_dict(finetuning_args, model_max_length)
env_mapping = {
"kt_weight_path": "ACCELERATE_KT_WEIGHT_PATH",
"kt_expert_checkpoint_path": "ACCELERATE_KT_EXPERT_CHECKPOINT_PATH",
"kt_model_max_length": "ACCELERATE_KT_MODEL_MAX_LENGTH",
"kt_lora_rank": "ACCELERATE_KT_LORA_RANK",
"kt_lora_alpha": "ACCELERATE_KT_LORA_ALPHA",
"kt_use_lora_experts": "ACCELERATE_KT_USE_LORA_EXPERTS",
"kt_lora_expert_num": "ACCELERATE_KT_LORA_EXPERT_NUM",
"kt_lora_expert_intermediate_size": "ACCELERATE_KT_LORA_EXPERT_INTERMEDIATE_SIZE",
}
for key, env_key in env_mapping.items():
value = kt_config.get(key)
if value is not None:
os.environ[env_key] = str(value)
hf_kt = getattr(training_args, "hf_kt_config", None)
if hf_kt is None or not hasattr(hf_kt, "_kt_config") or not isinstance(hf_kt._kt_config, dict):
return
hf_kt._kt_config.update(kt_config)
gc_enabled = getattr(training_args, "gradient_checkpointing", False) or not getattr(
self, "disable_gradient_checkpointing", True
self.configure_kt_checkpointing(training_args)
kt_config = self.get_kt_config_dict(
finetuning_args,
model_max_length,
self._get_advanced_kt_config(training_args),
)
if gc_enabled:
hf_kt._kt_config.setdefault("kt_share_cache_pool", True)
update_kt_config = getattr(training_args, "update_kt_config", None)
if not callable(update_kt_config):
raise RuntimeError(
"The installed Transformers-KT does not provide `TrainingArguments.update_kt_config()`."
)
adapter_dir = self._resolve_kt_adapter_artifact_dir("training")
update_kt_config(kt_config, adapter_name_or_path=adapter_dir)
def configure_kt_loading(self, finetuning_args: Any, model_max_length: int | None) -> None:
r"""Configure KT model loading for inference and evaluation."""
if not self.use_kt:
if self._kt_inference_config is not None:
raise ValueError("`kt_config` requires `use_kt: true`.")
return
if self.infer_backend != EngineName.HF:
raise ValueError("KTransformers inference requires `infer_backend: huggingface`.")
adapter_dir = self._resolve_kt_adapter_artifact_dir("inference")
try:
from transformers.integrations.kt import configure_kt
except (ImportError, ModuleNotFoundError) as exc:
raise RuntimeError("The installed Transformers-KT does not provide `configure_kt()`.") from exc
kt_config = self.get_kt_config_dict(
finetuning_args,
model_max_length,
self._normalize_advanced_kt_config(self._kt_inference_config),
)
self._kt_adapter_artifact_path = adapter_dir
self._kt_config_handle = configure_kt(kt_config)
@dataclass
@@ -580,6 +743,7 @@ class ModelArguments(
ExportArguments.__post_init__(self)
VllmArguments.__post_init__(self)
SGLangArguments.__post_init__(self)
KTransformersArguments.__post_init__(self)
@classmethod
def copyfrom(cls, source: "Self", **kwargs) -> "Self":

View File

@@ -18,6 +18,7 @@
import json
import os
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional
@@ -48,6 +49,14 @@ logger = logging.get_logger(__name__)
check_dependencies()
@dataclass
class _KTransformersRuntimeArguments:
kt_config: dict[str, Any] | None = field(
default=None,
metadata={"help": "Advanced KTransformers settings used during inference or evaluation."},
)
_TRAIN_ARGS = [
ModelArguments,
DataArguments,
@@ -56,9 +65,9 @@ _TRAIN_ARGS = [
GeneratingArguments,
]
_TRAIN_CLS = tuple[ModelArguments, DataArguments, TrainingArguments, FinetuningArguments, GeneratingArguments]
_INFER_ARGS = [ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]
_INFER_ARGS = [ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments, _KTransformersRuntimeArguments]
_INFER_CLS = tuple[ModelArguments, DataArguments, FinetuningArguments, GeneratingArguments]
_EVAL_ARGS = [ModelArguments, DataArguments, EvaluationArguments, FinetuningArguments]
_EVAL_ARGS = [ModelArguments, DataArguments, EvaluationArguments, FinetuningArguments, _KTransformersRuntimeArguments]
_EVAL_CLS = tuple[ModelArguments, DataArguments, EvaluationArguments, FinetuningArguments]
if is_mcore_adapter_available() and is_env_enabled("USE_MCA"):
@@ -117,6 +126,26 @@ def read_args(args: dict[str, Any] | list[str] | None = None) -> dict[str, Any]
return sys.argv[1:]
def _get_kt_runtime_capacity(
data_args: "DataArguments",
training_args: "TrainingArguments",
finetuning_args: "FinetuningArguments",
) -> int:
r"""Return the largest local token batch submitted to a KT expert."""
tokens_per_sample = data_args.cutoff_len
if finetuning_args.stage == "sft" and data_args.packing:
tokens_per_sample += 1
if finetuning_args.stage == "sft" and training_args.do_train:
tokens_per_sample = ((tokens_per_sample + 7) // 8) * 8
local_batch_sizes = [1]
if training_args.do_train:
local_batch_sizes.append(training_args.per_device_train_batch_size)
if training_args.do_eval or training_args.do_predict:
local_batch_sizes.append(training_args.per_device_eval_batch_size)
return tokens_per_sample * max(local_batch_sizes)
def _parse_args(
parser: "HfArgumentParser", args: dict[str, Any] | list[str] | None = None, allow_extra_keys: bool = False
) -> tuple[Any]:
@@ -340,13 +369,21 @@ def _configure_mbridge_training_args(training_args, data_args, finetuning_args)
def _parse_infer_args(args: dict[str, Any] | list[str] | None = None) -> _INFER_CLS:
parser = HfArgumentParser(_INFER_ARGS)
allow_extra_keys = is_env_enabled("ALLOW_EXTRA_ARGS")
return _parse_args(parser, args, allow_extra_keys=allow_extra_keys)
model_args, data_args, finetuning_args, generating_args, kt_args = _parse_args(
parser, args, allow_extra_keys=allow_extra_keys
)
model_args._kt_inference_config = kt_args.kt_config
return model_args, data_args, finetuning_args, generating_args
def _parse_eval_args(args: dict[str, Any] | list[str] | None = None) -> _EVAL_CLS:
parser = HfArgumentParser(_EVAL_ARGS)
allow_extra_keys = is_env_enabled("ALLOW_EXTRA_ARGS")
return _parse_args(parser, args, allow_extra_keys=allow_extra_keys)
model_args, data_args, eval_args, finetuning_args, kt_args = _parse_args(
parser, args, allow_extra_keys=allow_extra_keys
)
model_args._kt_inference_config = kt_args.kt_config
return model_args, data_args, eval_args, finetuning_args
def get_ray_args(args: dict[str, Any] | list[str] | None = None) -> RayArguments:
@@ -605,10 +642,10 @@ def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS
elif training_args.fp16:
model_args.compute_dtype = torch.float16
data_args.packing = data_args.packing if data_args.packing is not None else finetuning_args.stage == "pt"
model_args.device_map = {"": get_current_device()}
model_args.model_max_length = data_args.cutoff_len
model_args.block_diag_attn = data_args.neat_packing
data_args.packing = data_args.packing if data_args.packing is not None else finetuning_args.stage == "pt"
# Log on each process the small summary
logger.info(
@@ -620,7 +657,11 @@ def get_train_args(args: dict[str, Any] | list[str] | None = None) -> _TRAIN_CLS
transformers.set_seed(training_args.seed)
if model_args.use_kt:
model_args.apply_kt_config(finetuning_args, training_args, model_args.model_max_length)
model_args.apply_kt_config(
finetuning_args,
training_args,
_get_kt_runtime_capacity(data_args, training_args, finetuning_args),
)
return model_args, data_args, training_args, finetuning_args, generating_args
@@ -657,6 +698,8 @@ def get_infer_args(args: dict[str, Any] | list[str] | None = None) -> _INFER_CLS
else:
model_args.device_map = "auto"
model_args.configure_kt_loading(finetuning_args, data_args.cutoff_len)
return model_args, data_args, finetuning_args, generating_args
@@ -675,6 +718,7 @@ def get_eval_args(args: dict[str, Any] | list[str] | None = None) -> _EVAL_CLS:
_check_extra_dependencies(model_args, finetuning_args)
model_args.device_map = "auto"
model_args.configure_kt_loading(finetuning_args, data_args.cutoff_len)
transformers.set_seed(eval_args.seed)

View File

@@ -138,6 +138,12 @@ def _setup_freeze_tuning(
logger.info_rank0("Set trainable layers: {}".format(",".join(trainable_layers)))
def _load_kt_inference_adapter_artifacts(model: "PreTrainedModel", adapter_path: str) -> None:
from kt_kernel.sft import load_kt_adapter_artifacts
load_kt_adapter_artifacts(model, adapter_path)
def _setup_lora_tuning(
config: "PretrainedConfig",
model: "PreTrainedModel",
@@ -185,6 +191,8 @@ def _setup_lora_tuning(
"revision": model_args.model_revision,
"token": model_args.hf_hub_token,
}
if model_args.use_kt:
init_kwargs["autocast_adapter_dtype"] = cast_trainable_params_to_fp32
for adapter in adapter_to_merge:
model: LoraModel = PeftModel.from_pretrained(model, adapter, **init_kwargs)
@@ -209,6 +217,12 @@ def _setup_lora_tuning(
model, adapter_to_resume, is_trainable=is_trainable, **init_kwargs
)
if model_args.use_kt and not is_trainable:
adapter_path = model_args._kt_adapter_artifact_path
if adapter_path is None:
raise RuntimeError("KT adapter artifacts were not resolved before model loading.")
_load_kt_inference_adapter_artifacts(model, adapter_path)
logger.info_rank0("Loaded adapter(s): {}".format(",".join(model_args.adapter_name_or_path)))
if is_trainable and adapter_to_resume is None: # create new lora weights while training
@@ -264,7 +278,7 @@ def _setup_lora_tuning(
raise ValueError("KTransformers only supports LoRA finetuning.")
peft_config = LoraConfig(task_type=TaskType.CAUSAL_LM, inference_mode=False, **peft_kwargs)
model = get_peft_model(model, peft_config)
model = get_peft_model(model, peft_config, autocast_adapter_dtype=cast_trainable_params_to_fp32)
elif model_args.use_unsloth:
if finetuning_args.finetuning_type == "oft":
raise ValueError("Unsloth is currently not supported for OFT.")

View File

@@ -217,9 +217,9 @@ def load_model(
"You are try to using future feature about kernels, please note that this feature "
"is not supported for all models. If get any error, please disable this feature, or report the issue."
)
from ..v1.plugins.model_plugins.kernels.interface import apply_default_kernels
from ..v1.plugins.model_plugins.kernels.interface import apply_v1_kernels
model = apply_default_kernels(model, include_kernels=model_args.use_v1_kernels)
model = apply_v1_kernels(model, use_v1_kernels=model_args.use_v1_kernels)
trainable_params, all_param = count_parameters(model)
if is_trainable:

View File

@@ -82,6 +82,14 @@ def configure_attn_implementation(config: "PretrainedConfig", model_args: "Model
return
requested_attn_implementation = "flash_attention_2"
elif model_args.flash_attn == AttentionFunction.FA3:
from transformers.utils import is_flash_attn_3_available
if not is_flash_attn_3_available():
logger.warning_rank0("FlashAttention-3 is not installed.")
return
requested_attn_implementation = "flash_attention_3"
else:
raise NotImplementedError(f"Unknown attention type: {model_args.flash_attn}")
@@ -109,6 +117,8 @@ def print_attn_implementation(config: "PretrainedConfig") -> None:
if attn_implementation == "flash_attention_2":
logger.info_rank0("Using FlashAttention-2 for faster training and inference.")
elif attn_implementation == "flash_attention_3":
logger.info_rank0("Using FlashAttention-3 for faster training and inference.")
elif attn_implementation == "sdpa":
logger.info_rank0("Using torch SDPA for faster training and inference.")
else:

View File

@@ -40,6 +40,23 @@ if TYPE_CHECKING:
logger = logging.get_logger(__name__)
def _get_gradient_checkpointing_kwargs(model_args: "ModelArguments") -> dict[str, Any]:
r"""Build checkpoint kwargs through KT's public activation-context provider."""
if not model_args.use_kt:
return {"use_reentrant": model_args.use_reentrant_gc}
policy = model_args.get_kt_activation_policy()
if policy["gpu"] != "recompute":
return {"use_reentrant": False}
try:
from kt_kernel.sft import get_activation_checkpoint_context_fn
except (ImportError, ModuleNotFoundError) as exc:
raise RuntimeError("The installed kt-kernel does not provide the activation checkpoint context API.") from exc
return {"use_reentrant": False, "context_fn": get_activation_checkpoint_context_fn()}
def get_unsloth_gradient_checkpointing_func() -> Callable:
class UnslothGradientCheckpointing(torch.autograd.Function):
r"""Saves VRAM by smartly offloading to RAM."""
@@ -172,7 +189,7 @@ def prepare_model_for_training(model: "PreTrainedModel", model_args: "ModelArgum
)
model.gradient_checkpointing_enable = MethodType(gradient_checkpointing_enable, model)
model.gradient_checkpointing_enable(
gradient_checkpointing_kwargs={"use_reentrant": model_args.use_reentrant_gc}
gradient_checkpointing_kwargs=_get_gradient_checkpointing_kwargs(model_args)
)
setattr(model.config, "use_cache", False) # turn off when gradient checkpointing is enabled
logger.info_rank0("Gradient checkpointing enabled.")

View File

@@ -40,6 +40,10 @@ if TYPE_CHECKING:
logger = logging.get_logger(__name__)
def _uses_kt_non_expert_cache(model_args: "ModelArguments") -> bool:
return model_args.use_kt and bool(model_args.kt_non_expert_weight_path)
def _get_quantization_dataset(tokenizer: "PreTrainedTokenizer", model_args: "ModelArguments") -> list[dict[str, Any]]:
r"""Prepare the tokenized dataset to perform AutoGPTQ. Do not use tensor output for JSON serialization."""
if os.path.isfile(model_args.export_quantization_dataset):
@@ -108,6 +112,13 @@ def configure_quantization(
init_kwargs["ignore_mismatched_sizes"] = True
if quant_method == QuantizationMethod.FP8:
if _uses_kt_non_expert_cache(model_args):
if model_args.quantization_bit is not None:
raise ValueError("`quantization_bit` cannot be combined with KT weight caches.")
logger.info_rank0("Skipping source FP8 dequantization because KT weight caches are configured.")
return
from transformers import FineGrainedFP8Config
quant_config = FineGrainedFP8Config(dequantize=True)

View File

@@ -56,7 +56,7 @@ class CompositeModel:
)
break
if project_module is not None:
if isinstance(project_module, torch.nn.Module):
mm_projectors.append(project_module)
return mm_projectors
@@ -344,6 +344,15 @@ _register_composite_model(
)
_register_composite_model(
model_type="moss_vl",
projector_keys=["model.visual.merger", "model.separator_token"],
vision_model_keys=["model.visual.pos_embed", "model.visual.patch_embed", "model.visual.blocks"],
language_model_keys=["model.language_model", "lm_head"],
lora_conflict_keys=["patch_embed"],
)
_register_composite_model(
model_type="mllama",
vision_model_keys=["vision_model"],

View File

@@ -23,7 +23,7 @@ from transformers.modeling_utils import is_fsdp_enabled
from transformers.utils import is_torch_cuda_available, is_torch_npu_available
from ..extras import logging
from ..extras.misc import infer_optim_dtype
from ..extras.misc import check_version, infer_optim_dtype
from ..extras.packages import is_transformers_version_greater_than
from .model_utils.attention import configure_attn_implementation, print_attn_implementation
from .model_utils.checkpointing import prepare_model_for_training
@@ -418,6 +418,10 @@ def patch_config(
"pip install git+https://github.com/huggingface/transformers.git@3c2517727ce28a30f5044e01663ee204deb1cdbe"
)
if getattr(config, "model_type", None) == "moss_vl":
check_version("transformers==4.57.1", mandatory=True)
check_version("torchcodec==0.7.0", mandatory=True)
if getattr(config, "model_type", None) == "qwen3_omni_moe":
patch_qwen3_omni_moe_thinker_text_sparse_moe_block()

View File

@@ -43,7 +43,7 @@ from ..utils.callbacks import (
TrainerCallback,
TrainerState,
)
from ..utils.helper import compute_valid_tokens
from ..utils.helper import compute_valid_tokens, is_tokenizer, model_uses_mrope
from ..utils.types import BatchInput, HFModel, ModelOutput, Tensor, TorchDataset
from .rendering import Renderer
from .utils.batching import BatchGenerator
@@ -75,6 +75,7 @@ class BaseTrainer:
self.dp_size = DistributedInterface().get_world_size(Dim.DP)
self.cp_size = DistributedInterface().get_world_size(Dim.CP)
self.model_input_names = self.renderer.processor.model_input_names
self._uses_mrope = model_uses_mrope(self.model.config)
self._create_batch_generator()
# Calculate num_training_steps: max_steps takes priority if set
@@ -89,6 +90,9 @@ class BaseTrainer:
if self.args.enable_activation_checkpointing:
self.model.gradient_checkpointing_enable({"use_reentrant": False})
# Note: under FSDP2 bf16, encoder-tower nn.LayerNorms are made dtype-safe for the
# checkpoint recompute inside the FSDP2 engine (see fsdp2.py prepare_model), so the
# tower keeps activation checkpointing too.
self._deepspeed_engine = None
dist_name = self.args.dist_config.name if self.args.dist_config is not None else None
@@ -184,7 +188,11 @@ class BaseTrainer:
"dist_config is None but distributed training is enabled; falling back to DistributedDataParallel."
)
device_ids = None if self.device.type == "cpu" else [self.device.index]
self.model = DDP(self.model, device_ids=device_ids)
# Multimodal models invoke the vision tower only when a step carries media; a
# globally media-less step leaves vision params unused, which trips DDP's default
# all-params-used assertion. (FSDP tolerates a uniform skip; DDP does not.)
find_unused = not is_tokenizer(self.renderer.processor)
self.model = DDP(self.model, device_ids=device_ids, find_unused_parameters=find_unused)
else:
from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
@@ -224,6 +232,9 @@ class BaseTrainer:
model_inputs = {
k: v.to(self.device, non_blocking=True) for k, v in batch.items() if isinstance(v, torch.Tensor)
}
# Let mRoPE models build their own multimodal 3D position ids (see _uses_mrope in __init__).
if self._uses_mrope:
model_inputs.pop("position_ids", None)
labels = batch["labels"].to(self.device, non_blocking=True)
outputs: ModelOutput = model(**model_inputs)
logits = outputs.logits.float()
@@ -283,19 +294,24 @@ class BaseTrainer:
# deepspeed: engine.step() already ran inside backward at the sync boundary
grad_norm = self._deepspeed_engine.get_grad_norm()
else:
# FSDP2 shards params/grads across the fsdp mesh, so clip_grad_norm_ returns a
# per-rank local shard norm (global / sqrt(shard_size)): reported grad_norm then
# scales as 1/sqrt(dp_size) and the clip coefficient is applied per-shard. Reduce
# to the true global norm first, then clip with it.
grads = [p.grad for p in self.model.parameters() if p.grad is not None]
total_norm = torch.nn.utils.get_total_norm(grads)
if isinstance(total_norm, DTensor):
# full_tensor all-reduces across the fsdp mesh (spans CP under default
# 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)
grad_norm = total_norm.item()
dist_name = self.args.dist_config.name if self.args.dist_config else None
if dist_name == "fsdpturbo":
from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
grad_norm = DistributedPlugin(dist_name).clip_grad_norm(self.model, self.args.max_grad_norm)
else:
# FSDP2 shards params/grads across the fsdp mesh, so clip_grad_norm_ returns a
# per-rank local shard norm. Materialize the true global norm before clipping.
grads = [p.grad for p in self.model.parameters() if p.grad is not None]
total_norm = torch.nn.utils.get_total_norm(grads)
if isinstance(total_norm, DTensor):
# full_tensor all-reduces across the fsdp mesh (spans CP under default
# mp_shard=world); a separate CP reduce would over-count by sqrt(cp_size).
total_norm = total_norm.full_tensor()
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]
logger.warning_rank0(f"Gradient norm is not finite: {grad_norm}")
@@ -352,7 +368,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"):
if self.args.dist_config is not None and self.args.dist_config.name in ("deepspeed", "fsdp2", "fsdpturbo"):
from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
DistributedPlugin(self.args.dist_config.name).save_model(

View File

@@ -112,8 +112,7 @@ class ModelEngine:
if self.args.custom_chat_template:
if not is_tokenizer(self.processor):
self.processor.chat_template = self.args.custom_chat_template
else:
tokenizer.chat_template = self.args.custom_chat_template
tokenizer.chat_template = self.args.custom_chat_template
def _init_model_config(self) -> HFConfig:
"""Init model config."""
@@ -150,8 +149,19 @@ class ModelEngine:
if self.args.model_class == ModelClass.LLM:
from transformers import AutoModelForCausalLM, AutoModelForImageTextToText
if type(self.model_config) in AutoModelForImageTextToText._model_mapping.keys():
# AutoModelForMultimodalLM (audio / other multimodal LMs, e.g. Qwen2-Audio) was added in
# a newer transformers; fall back gracefully when it is absent (e.g. 4.57.1).
try:
from transformers import AutoModelForMultimodalLM
except ImportError:
AutoModelForMultimodalLM = None
cfg_type = type(self.model_config)
if cfg_type in AutoModelForImageTextToText._model_mapping.keys():
AutoClass = AutoModelForImageTextToText
elif AutoModelForMultimodalLM is not None and cfg_type in AutoModelForMultimodalLM._model_mapping.keys():
# Audio / other multimodal LMs (e.g. Qwen2-Audio) live here, not in CausalLM.
AutoClass = AutoModelForMultimodalLM
else:
AutoClass = AutoModelForCausalLM
@@ -173,7 +183,7 @@ class ModelEngine:
if init_device.type == DeviceType.META:
assert self.args.quant_config is None, "Quantization is not supported with meta device."
with init_empty_weights():
model = AutoClass.from_config(self.model_config)
model = AutoClass.from_config(self.model_config, attn_implementation=self.args.flash_attn)
else:
model = AutoClass.from_pretrained(
self.args.model,
@@ -187,6 +197,10 @@ class ModelEngine:
init_mode = self.args.init_config.name if self.args.init_config is not None else "init_on_default"
model._init_mode = init_mode
if hasattr(model, "thinker"):
model = model.thinker
model._init_mode = init_mode
if self.args.peft_config is None:
if self.is_train:
logger.info_rank0("Fine-tuning mode: full tuning")

View File

@@ -1,4 +1,4 @@
# Copyright 2025 the LlamaFactory team.
# Copyright 2026 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.
@@ -14,13 +14,15 @@
"""Message <-> HF-template plumbing for rendering.
Pure, stateless helpers: convert v1 ``Message`` to HF chat-template format. No tokenization policy
decisions live here -- only mechanical conversion used by ``rendering.py``.
Pure, stateless helpers: convert v1 ``Message`` to HF chat-template format, extract/count media, and
guard media placeholder counts. No tokenization policy decisions live here -- only mechanical
conversion used by ``rendering.py``.
"""
import json
from ...utils.types import Message
from ...utils.helper import get_tokenizer
from ...utils.types import Message, Processor
_FALLBACK_CHATML_JINJA = (
@@ -33,28 +35,59 @@ _FALLBACK_CHATML_JINJA = (
)
def _to_hf_messages(messages: list[Message]) -> list[dict]:
def _to_hf_messages(messages: list[Message], is_multimodal: bool = False) -> list[dict]:
"""Convert v1 Message format to HF format for apply_chat_template."""
hf_messages = []
for message in messages:
tool_calls: list[dict] = []
reasoning_content = ""
text = ""
for content in message["content"]:
if content["type"] == "text":
text += content["value"]
elif content["type"] == "reasoning":
reasoning_content += content["value"]
elif content["type"] == "tool_call":
try:
tc = json.loads(content["value"])
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"]}})
hf_msg = {"role": message["role"], "content": text}
if is_multimodal:
hf_content = []
for content in message["content"]:
if content["type"] == "text":
hf_content.append({"type": "text", "text": content["value"]})
elif content["type"] == "reasoning":
reasoning_content += content["value"]
elif content["type"] == "tool_call":
try:
tc = json.loads(content["value"])
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"]}}
)
elif content["type"] == "image_url":
hf_content.append({"type": "image", "image": content["value"]})
elif content["type"] == "video_url":
hf_content.append({"type": "video", "video": content["value"]})
elif content["type"] == "audio_url":
hf_content.append({"type": "audio", "audio": content["value"]})
hf_msg = {"role": message["role"], "content": hf_content}
else:
text = ""
for content in message["content"]:
if content["type"] == "text":
text += content["value"]
elif content["type"] == "reasoning":
reasoning_content += content["value"]
elif content["type"] == "tool_call":
try:
tc = json.loads(content["value"])
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"]}}
)
hf_msg = {"role": message["role"], "content": text}
if tool_calls:
hf_msg["tool_calls"] = tool_calls
@@ -63,3 +96,75 @@ def _to_hf_messages(messages: list[Message]) -> list[dict]:
hf_messages.append(hf_msg)
return hf_messages
def _extract_media_from_messages(messages: list[Message]) -> tuple[list, list, list]:
"""Extract image, video and audio paths/values from messages in order."""
images, videos, audios = [], [], []
for message in messages:
for content in message["content"]:
if content["type"] == "image_url":
images.append(content["value"])
elif content["type"] == "video_url":
videos.append(content["value"])
elif content["type"] == "audio_url":
audios.append(content["value"])
return images, videos, audios
def _count_media_in_messages(messages: list[Message]) -> tuple[int, int, int]:
"""Count total images, videos and audios in messages."""
n_images, n_videos, n_audios = 0, 0, 0
for message in messages:
for content in message["content"]:
if content["type"] == "image_url":
n_images += 1
elif content["type"] == "video_url":
n_videos += 1
elif content["type"] == "audio_url":
n_audios += 1
return n_images, n_videos, n_audios
def _load_audios(values: list, sampling_rate: int) -> list:
"""Load audio inputs into mono waveforms resampled to ``sampling_rate``."""
import numpy as np
import torchaudio
results = []
for value in values:
if isinstance(value, np.ndarray):
results.append(value)
continue
waveform, sr = torchaudio.load(value)
if waveform.shape[0] > 1: # downmix to mono
waveform = waveform.mean(dim=0, keepdim=True)
if sr != sampling_rate:
waveform = torchaudio.functional.resample(waveform, sr, sampling_rate)
results.append(waveform.squeeze(0).numpy())
return results
def _check_placeholder_counts(
processor: "Processor", full_text: str, n_images: int, n_videos: int, n_audios: int = 0
) -> None:
"""Guard: every media placeholder in the rendered text must originate from a media block."""
tokenizer = get_tokenizer(processor)
for attr, count, kind in (
("image_token_id", n_images, "image"),
("video_token_id", n_videos, "video"),
("audio_token_id", n_audios, "audio"),
):
tid = getattr(processor, attr, None)
if tid is None:
tid = getattr(tokenizer, attr, None)
if tid is None:
continue
placeholder = tokenizer.convert_ids_to_tokens(tid)
seen = full_text.count(placeholder)
if seen != count:
raise ValueError(
f"{kind} placeholder count ({seen}) != number of {kind} blocks ({count}); "
"media must be provided via image_url/video_url content blocks."
)

View File

@@ -1,4 +1,4 @@
# Copyright 2025 the LlamaFactory team.
# Copyright 2026 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.
@@ -19,24 +19,29 @@ sibling modules:
- ``format`` -- v1<->HF message conversion
- ``escape`` -- special-token escaping (prompt-injection hardening)
Assistant supervision is located WITHOUT a per-model marker table: a training sample is rendered
so that its last message is the supervised assistant turn, and that turn's token span is recovered
by a single prompt/full difference -- encode the prompt (everything up to and including the
assistant role header, via ``add_generation_prompt=True``) and the full sequence, then the tail of
the full sequence that the prompt does not cover is exactly this turn. Multi-turn conversations are
split into one sample per supervised turn (see ``process_samples``) so the supervised turn is always
the last one; this keeps the diff on the only boundary that is prefix-stable across chat templates
(appending the final assistant turn never restripts earlier turns), so models with reasoning-history
stripping (e.g. Qwen3 ``<think>``) are handled correctly without hard-coding role markers.
Note: ``position_ids`` are assigned by ``process_samples`` (1-based); multimodal (mrope) position
ids are expected to be recomputed by the model/trainer.
"""
import json
import numpy as np
import torch
from ...utils.constants import IGNORE_INDEX
from ...utils.helper import get_tokenizer
from ...utils.helper import get_tokenizer, is_tokenizer
from ...utils.types import Message, ModelInput, Processor, Sample
from ..utils.collation import _MULTIMODAL_PASSTHROUGH_KEYS
from .escape import _escape_special, _escape_special_in_messages, _special_token_strings
from .format import _FALLBACK_CHATML_JINJA, _to_hf_messages
from .format import (
_FALLBACK_CHATML_JINJA,
_check_placeholder_counts,
_count_media_in_messages,
_extract_media_from_messages,
_load_audios,
_to_hf_messages,
)
def _render_messages(
@@ -46,20 +51,23 @@ def _render_messages(
is_generate: bool = False,
**kwargs,
) -> ModelInput:
r"""Render messages using the model's own chat template.
r"""Render messages using the model's own chat template, locating supervision by a prompt/full diff.
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
is_multimodal = not is_tokenizer(processor)
template_caller = processor if is_multimodal else tokenizer
if not getattr(template_caller, "chat_template", None):
template_caller.chat_template = _FALLBACK_CHATML_JINJA
# 0. Neutralize special-token strings in user-controlled text (no-op for normal data).
specials = _special_token_strings(tokenizer)
special_ids = {tid for tid, t in tokenizer.added_tokens_decoder.items() if getattr(t, "special", False)}
messages = _escape_special_in_messages(messages, specials, special_ids, tokenizer)
hf_messages = _to_hf_messages(messages)
hf_messages = _to_hf_messages(messages, is_multimodal=is_multimodal)
tools_parsed = None
if tools:
@@ -70,36 +78,76 @@ def _render_messages(
raise ValueError(f"tools is not valid JSON: {tools!r}") from e
if not isinstance(tools_parsed, list):
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
if not is_generate and hf_messages and hf_messages[-1]["role"] == "assistant":
kwargs["enable_thinking"] = bool(hf_messages[-1].get("reasoning_content"))
def _encode(hf_msgs: list[dict], src_msgs: list[Message], add_generation_prompt: bool):
"""Render + tokenize, expanding media via the processor. Returns (input_ids, mm_outputs)."""
text = template_caller.apply_chat_template(
hf_msgs, tokenize=False, add_generation_prompt=add_generation_prompt, tools=tools_parsed, **kwargs
)
return tokenizer(text, add_special_tokens=False)["input_ids"]
if is_multimodal and _count_media_in_messages(src_msgs) != (0, 0, 0):
images, videos, audios = _extract_media_from_messages(src_msgs)
# Every placeholder must come from a media block (escaping broke any literal ones).
_check_placeholder_counts(processor, text, len(images), len(videos), len(audios))
proc_kwargs = {"return_tensors": "pt"}
if images:
proc_kwargs["images"] = images
if videos:
proc_kwargs["videos"] = videos
if audios:
# Audio processors want decoded waveforms at the model's sampling rate, not paths.
proc_kwargs["audio"] = _load_audios(audios, processor.feature_extractor.sampling_rate)
mm_outputs = processor(text=text, **proc_kwargs)
return mm_outputs["input_ids"][0].tolist(), mm_outputs
return tokenizer(text, add_special_tokens=False)["input_ids"], None
# 1. Full sequence, used verbatim.
input_ids = _encode(hf_messages, add_generation_prompt=is_generate)
# 1. Full sequence (used verbatim), plus its multimodal feature outputs.
input_ids, outputs = _encode(hf_messages, messages, add_generation_prompt=is_generate)
n = len(input_ids)
def _attach_multimodal(result: ModelInput) -> None:
if outputs is None:
return
for key in _MULTIMODAL_PASSTHROUGH_KEYS:
if key in outputs:
result[key] = outputs[key]
mm_type_ids = outputs["mm_token_type_ids"][0].tolist() if "mm_token_type_ids" in outputs else None
for attr, marker in (("image_token_id", 1), ("video_token_id", 2), ("audio_token_id", 3)):
token_id = getattr(processor, attr, None)
if token_id is None:
token_id = getattr(tokenizer, attr, None)
if token_id is None or token_id not in input_ids:
continue
if mm_type_ids is not None and marker in mm_type_ids:
continue
if mm_type_ids is None:
mm_type_ids = [0] * len(input_ids)
mm_type_ids = [marker if tid == token_id else t for t, tid in zip(mm_type_ids, input_ids)]
if mm_type_ids is not None:
result["mm_token_type_ids"] = mm_type_ids
if is_generate:
# Generation prompt only -- nothing is supervised.
return ModelInput(
result = ModelInput(
input_ids=input_ids,
attention_mask=[1] * n,
labels=[IGNORE_INDEX] * n,
loss_weights=[0.0] * n,
)
_attach_multimodal(result)
return result
# 2. Locate the supervised (last) assistant turn by a prompt/full diff (no marker table).
if not messages or messages[-1]["role"] != "assistant":
raise ValueError(
"training render expects the last message to be the supervised assistant turn; "
"multi-turn conversations are split per turn in process_samples."
)
prompt_ids = _encode(hf_messages[:-1], add_generation_prompt=True)
prompt_ids, _ = _encode(hf_messages[:-1], messages[:-1], add_generation_prompt=True)
if input_ids[: len(prompt_ids)] != prompt_ids:
# The prompt must be a token-prefix of the full sequence for the diff to be valid. If a
# template re-renders earlier turns when the final turn is appended, fail loud rather than
@@ -117,16 +165,21 @@ def _render_messages(
labels.append(tid if supervised else IGNORE_INDEX)
loss_weights.append(weight)
return ModelInput(
result = ModelInput(
input_ids=input_ids,
attention_mask=[1] * n,
labels=labels,
loss_weights=loss_weights,
)
_attach_multimodal(result)
return result
class Renderer:
def __init__(self, processor: Processor) -> None:
def __init__(self, processor: Processor, config=None):
# ``config`` is accepted for call-site compatibility (ModelEngine passes the model config)
# but is no longer needed: supervision is located by a prompt/full diff, not a per-model
# marker table, so the renderer is model-agnostic.
self.processor = processor
def render_messages(
@@ -152,6 +205,61 @@ class Renderer:
"""
return _render_messages(self.processor, messages, tools, is_generate, **kwargs)
def get_dummy_media_fragment(self, modality: str) -> dict:
"""Build (and cache) a minimal valid media fragment for ``modality`` ("image"|"video"|"audio")."""
if modality not in ("image", "video", "audio"):
raise ValueError(f"Unsupported dummy media modality: {modality!r} (expected image/video/audio).")
if is_tokenizer(self.processor):
raise RuntimeError("Cannot build a dummy media fragment for a text-only processor.")
if not hasattr(self, "_dummy_fragments"):
self._dummy_fragments: dict[str, dict] = {}
if modality in self._dummy_fragments:
return self._dummy_fragments[modality]
from PIL import Image as _PILImage
if modality == "image":
media_block = {"type": "image_url", "value": _PILImage.new("RGB", (64, 64))}
target, presence_key = 1, "pixel_values"
elif modality == "video":
# A minimal clip: the temporal patch size is typically 2, so provide two frames.
media_block = {"type": "video_url", "value": np.zeros((2, 64, 64, 3), dtype=np.uint8)}
target, presence_key = 2, "pixel_values_videos"
else:
# A short synthetic waveform at the model's sampling rate; the feature extractor pads it.
sr = self.processor.feature_extractor.sampling_rate
media_block = {"type": "audio_url", "value": np.zeros(sr // 10, dtype=np.float32)}
target, presence_key = 3, "input_features"
messages: list[Message] = [
{"role": "user", "content": [media_block]},
{"role": "assistant", "content": [{"type": "text", "value": "ok"}]},
]
rendered = self.render_messages(messages)
mm_type_ids = rendered.get("mm_token_type_ids")
if not mm_type_ids or target not in mm_type_ids or presence_key not in rendered:
raise RuntimeError(f"Processor did not emit {modality} placeholder tokens for the dummy sample.")
positions = [i for i, t in enumerate(mm_type_ids) if t == target]
# Include the surrounding start/end delimiters (vision_start/end or audio_bos/eos) so the
# fragment matches exactly what the template emits around real media.
lo = max(positions[0] - 1, 0)
hi = min(positions[-1] + 2, len(rendered["input_ids"]))
fragment: dict = {
"input_ids": list(rendered["input_ids"][lo:hi]),
"mm_token_type_ids": list(mm_type_ids[lo:hi]),
}
for key in _MULTIMODAL_PASSTHROUGH_KEYS:
if key in rendered:
fragment[key] = rendered[key]
self._dummy_fragments[modality] = fragment
return fragment
def process_samples(self, samples: list[Sample]) -> list[ModelInput]:
"""Process samples to model input.
@@ -189,6 +297,17 @@ class Renderer:
model_input["position_ids"] = list(range(1, len(chosen_input["input_ids"]) + 1)) + list(
range(1, len(rejected_input["input_ids"]) + 1)
)
for key in _MULTIMODAL_PASSTHROUGH_KEYS:
tensors = [inp[key] for inp in (chosen_input, rejected_input) if key in inp]
if tensors:
model_input[key] = torch.cat(tensors, dim=0)
if "mm_token_type_ids" in chosen_input or "mm_token_type_ids" in rejected_input:
chosen_mm = chosen_input.get("mm_token_type_ids", [0] * len(chosen_input["input_ids"]))
rejected_mm = rejected_input.get("mm_token_type_ids", [0] * len(rejected_input["input_ids"]))
model_input["mm_token_type_ids"] = chosen_mm + rejected_mm
rendered.append(model_input)
else:
raise ValueError("No valid messages or chosen_messages/rejected_messages found in sample.")

View File

@@ -1,4 +1,4 @@
# Copyright 2025 the LlamaFactory team.
# Copyright 2026 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.
@@ -31,13 +31,16 @@ from torch.utils.data import default_collate
from torchdata.stateful_dataloader import StatefulDataLoader
from torchdata.stateful_dataloader.sampler import StatefulDistributedSampler
from ...accelerator.helper import ReduceOp
from ...accelerator.interface import Dim, DistributedInterface
from ...config import BatchingStrategy
from ...utils import logging
from ...utils.helper import pad_and_truncate
from ...utils.constants import IGNORE_INDEX
from ...utils.helper import is_tokenizer
from ...utils.objects import StatefulBuffer
from ...utils.types import BatchInfo, BatchInput, ModelInput, TorchDataset
from ...utils.types import BatchInfo, BatchInput, ModelInput, Tensor, TorchDataset
from ..rendering import Renderer
from .collation import _MULTIMODAL_PASSTHROUGH_KEYS, pad_and_truncate
logger = logging.get_logger(__name__)
@@ -45,7 +48,82 @@ logger = logging.get_logger(__name__)
__all__ = ["BatchGenerator"]
def default_collate_fn(buffer: StatefulBuffer, batch_info: BatchInfo) -> list[BatchInput] | None:
# (modality, presence/feature key, grid key, mm_token_type_ids marker) for encoder-tower alignment.
# The presence key is what survives collation when the modality is present; the grid key is unused
# here (kept for parity with the collation specs). Audio carries no grid -- feature_attention_mask
# rides along as a passthrough feature.
_ALIGN_MODALITIES = (
("image", "pixel_values", "image_grid_thw", 1),
("video", "pixel_values_videos", "video_grid_thw", 2),
("audio", "input_features", "feature_attention_mask", 3),
)
def _collate_micro_batch(micro_batch: list[ModelInput], cutoff_len: int) -> BatchInput:
"""Pad/truncate then collate one micro batch (text fields stacked, MM features dim-0 concat)."""
padded = pad_and_truncate(micro_batch, cutoff_len)
standard_samples = [{k: v for k, v in s.items() if k not in _MULTIMODAL_PASSTHROUGH_KEYS} for s in padded]
collated = default_collate(standard_samples)
for key in _MULTIMODAL_PASSTHROUGH_KEYS:
tensors = [s[key] for s in padded if key in s]
if tensors:
collated[key] = torch.cat(tensors, dim=0)
return collated
def _inject_dummy_into_collated(collated: BatchInput, fragment: dict, marker: int) -> None:
"""Append a zero-loss dummy media fragment to an already-collated micro batch, in place.
Operates *after* pad_and_truncate so it reflects post-truncation presence: an image whose
placeholder tokens were partially cut is deleted by ``_align_multimodal_on_truncation``,
turning that sample text-only -- which must be detected here (not before truncation) or the
vision-tower call count still desyncs across ranks.
The dummy tokens are appended (extra columns) into row 0 only; other rows get padding there.
Causal attention keeps every real token's logits unchanged; the dummy carries IGNORE_INDEX
labels and zero loss weight, so it contributes nothing to the loss while forcing the (FSDP-
sharded) vision tower to run.
"""
bsz, seqlen = collated["input_ids"].shape
frag_ids = torch.tensor(fragment["input_ids"], dtype=collated["input_ids"].dtype)
frag_len = frag_ids.numel()
frag_mm = torch.tensor(fragment["mm_token_type_ids"], dtype=torch.long)
new_len = seqlen + frag_len
def _grow(tensor: Tensor, pad_value, row0_tail=None) -> Tensor:
out = torch.full((bsz, new_len), pad_value, dtype=tensor.dtype)
out[:, :seqlen] = tensor
if row0_tail is not None:
out[0, seqlen:] = row0_tail.to(tensor.dtype)
return out
collated["input_ids"] = _grow(collated["input_ids"], 0, frag_ids)
collated["attention_mask"] = _grow(collated["attention_mask"], 0)
collated["attention_mask"][0, seqlen:] = 1
collated["labels"] = _grow(collated["labels"], IGNORE_INDEX) # dummy region stays ignored
collated["loss_weights"] = _grow(collated["loss_weights"], 0.0)
if "position_ids" in collated:
pos = _grow(collated["position_ids"], 0)
pos[0, seqlen:] = torch.arange(seqlen + 1, new_len + 1, dtype=pos.dtype)
collated["position_ids"] = pos
mm = collated.get("mm_token_type_ids")
if mm is not None:
collated["mm_token_type_ids"] = _grow(mm, 0, frag_mm)
else:
mm = torch.zeros((bsz, new_len), dtype=torch.long)
mm[0, seqlen:] = frag_mm
collated["mm_token_type_ids"] = mm
for key, value in fragment.items():
if key in ("input_ids", "mm_token_type_ids"):
continue
collated[key] = torch.cat([collated[key], value], dim=0) if key in collated else value
def default_collate_fn(
buffer: StatefulBuffer, batch_info: BatchInfo, renderer: Renderer | None = None
) -> 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"]
@@ -54,10 +132,24 @@ def default_collate_fn(buffer: StatefulBuffer, batch_info: BatchInfo) -> list[Ba
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]
batch.append(default_collate(pad_and_truncate(micro_batch, cutoff_len)))
micro_batches = [samples[i * micro_batch_size : (i + 1) * micro_batch_size] for i in range(num_micro_batch)]
# Collate first; presence is judged on the *post-truncation* result, since truncation can
# delete a partially-cut image and turn a sample text-only (see _inject_dummy_into_collated).
batch = [_collate_micro_batch(mb, cutoff_len) for mb in micro_batches]
if renderer is not None and not is_tokenizer(renderer.processor):
present = torch.zeros((num_micro_batch, len(_ALIGN_MODALITIES)), dtype=torch.int64)
for i, collated in enumerate(batch):
for m, (_, pixel_key, _, _) in enumerate(_ALIGN_MODALITIES):
present[i, m] = int(pixel_key in collated)
present = DistributedInterface().all_reduce(present, op=ReduceOp.MAX, dim=Dim.DP)
for i, collated in enumerate(batch):
for m, (modality, pixel_key, _, marker) in enumerate(_ALIGN_MODALITIES):
if present[i, m] and pixel_key not in collated:
_inject_dummy_into_collated(collated, renderer.get_dummy_media_fragment(modality), marker)
return batch
@@ -227,8 +319,17 @@ class BatchGenerator(Iterator):
def _generate_batch(self) -> list[BatchInput] | None:
if self.batching_strategy == BatchingStrategy.NORMAL:
return default_collate_fn(self._buffer, self._batch_info)
return default_collate_fn(self._buffer, self._batch_info, self.renderer)
else:
# Non-NORMAL strategies (dynamic / padding_free) collate ragged pixel tensors with a
# bare default_collate and have no vision-tower alignment, so multimodal data would
# crash or hang. Fail loud instead of silently mishandling it.
if any(k in s for s in self._buffer.samples for k in _MULTIMODAL_PASSTHROUGH_KEYS):
raise NotImplementedError(
f"batching_strategy={self.batching_strategy.value!r} does not support multimodal data; "
"use the NORMAL strategy for image/video training."
)
from ...plugins.trainer_plugins.batching import BatchingPlugin
return BatchingPlugin(self.batching_strategy).generate_batch(self._buffer, self._batch_info)

View File

@@ -250,7 +250,7 @@ class TrainingCheckpointCoordinator:
num_training_steps=self._t.num_training_steps,
)
if self._dist_name in ("fsdp2", "deepspeed"):
if self._dist_name in ("fsdp2", "fsdpturbo", "deepspeed"):
from ...plugins.trainer_plugins.distributed.interface import DistributedPlugin
DistributedPlugin(self._dist_name).save_checkpoint(
@@ -306,7 +306,7 @@ class TrainingCheckpointCoordinator:
self._t.global_step = metadata["global_step"]
self._t._resume_epoch = metadata["epoch"]
if self._dist_name in ("fsdp2", "deepspeed"):
if self._dist_name in ("fsdp2", "fsdpturbo", "deepspeed"):
from ...plugins.trainer_plugins.distributed.interface import DistributedPlugin
DistributedPlugin(self._dist_name).load_checkpoint(

View File

@@ -0,0 +1,277 @@
# Copyright 2026 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.
"""Batch collation utils: padding/truncation and multimodal-feature alignment.
These operate on already-rendered ``ModelInput`` dicts (token lists + pixel tensors) and produce
padded ``BatchInput`` tensors. They are pure batching concerns -- independent of how a sample was
rendered -- and are consumed by the batch generators in ``core/utils/batching.py`` and
``plugins/trainer_plugins/batching.py``. Kept out of ``rendering.py`` so that file is only about
turning messages into a single tokenized sample.
"""
import torch
from ...utils.constants import IGNORE_INDEX
from ...utils.types import BatchInput, ModelInput, Tensor
# Multimodal feature keys the processor emits per sample. They are NOT padded/stacked like text
# fields: pixel/audio-feature tensors are ragged (variable patch / frame counts), so the collators
# concatenate them along dim 0 instead. Shared by rendering (which copies them verbatim from the
# processor) and the collators (which merge them across a micro batch).
_MULTIMODAL_PASSTHROUGH_KEYS = frozenset(
{
"pixel_values",
"image_grid_thw",
"pixel_values_videos",
"video_grid_thw",
"second_per_grid_ts", # Qwen2.5-VL name for the video temporal grid spacing
"video_second_per_grid", # Qwen2.5-Omni name for the same (fed to get_rope_index)
"input_features",
"feature_attention_mask",
}
)
def _pad_and_truncate(tensor: Tensor, max_seqlen: int, pad_value: int = 0) -> Tensor:
if tensor.shape[-1] >= max_seqlen:
return tensor[..., :max_seqlen]
pad_shape = list(tensor.shape)
pad_shape[-1] = max_seqlen - tensor.shape[-1]
pad_tensor = torch.full(pad_shape, pad_value, dtype=tensor.dtype, device=tensor.device)
return torch.cat([tensor, pad_tensor], dim=-1)
def _align_grid_media(
sample: ModelInput,
mm_type_ids: list[int],
max_length: int,
*,
target: int,
grid_key: str,
pixel_key: str,
) -> list[int]:
"""Trim and zero one modality's orphaned tokens for a single sample.
Layout-agnostic: a media item's placeholder tokens may be a single contiguous run or split
into per-frame sub-runs; completeness is decided per token *position*, so both are handled
identically.
Returns the (possibly updated) ``mm_token_type_ids`` so chained calls see earlier zeroing.
"""
if grid_key not in sample or pixel_key not in sample:
return mm_type_ids
grid = sample[grid_key]
n_items = len(grid)
if n_items == 0:
return mm_type_ids
positions = [i for i, t in enumerate(mm_type_ids) if t == target]
patches_per_item = [int(grid[i].prod()) for i in range(n_items)]
total_patches = sum(patches_per_item)
total_tokens = len(positions)
# merge_size**2 = pixel patches per placeholder token, derived from the data. Bail out
# untouched if the sample is inconsistent.
if total_tokens == 0 or total_patches % total_tokens != 0:
return mm_type_ids
merge_sq = total_patches // total_tokens
tokens_per_item = [p // merge_sq for p in patches_per_item]
if sum(tokens_per_item) != total_tokens:
return mm_type_ids
# Each item owns a contiguous slice of `positions`; it is complete iff its last
# placeholder token lands inside the kept window [0, max_length).
n_complete = 0
cum = 0
for n_i in tokens_per_item:
if positions[cum + n_i - 1] < max_length:
n_complete += 1
cum += n_i
else:
break
if n_complete >= n_items:
return mm_type_ids
# Trim pixel features and grid to the complete prefix.
keep_patches = sum(patches_per_item[:n_complete])
sample[pixel_key] = sample[pixel_key][:keep_patches]
sample[grid_key] = grid[:n_complete]
# Zero out orphaned placeholder tokens that fall inside the kept window; tokens
# beyond max_length are removed by truncation anyway (positions are sorted).
input_ids = list(sample["input_ids"])
mm_type_ids = list(mm_type_ids)
labels = list(sample["labels"]) if "labels" in sample else None
loss_weights = list(sample["loss_weights"]) if "loss_weights" in sample else None
for pos in positions[cum:]:
if pos >= max_length:
break
input_ids[pos] = 0
mm_type_ids[pos] = 0
if labels is not None:
labels[pos] = IGNORE_INDEX
if loss_weights is not None:
loss_weights[pos] = 0.0
sample["input_ids"] = input_ids
sample["mm_token_type_ids"] = mm_type_ids
if labels is not None:
sample["labels"] = labels
if loss_weights is not None:
sample["loss_weights"] = loss_weights
return mm_type_ids
def _align_audio(sample: ModelInput, mm_type_ids: list[int], max_length: int, *, target: int = 3) -> list[int]:
"""Trim and zero orphaned audio tokens for a single sample on truncation.
Returns the (possibly updated) ``mm_token_type_ids``.
"""
if "input_features" not in sample or "feature_attention_mask" not in sample:
return mm_type_ids
n_items = sample["input_features"].shape[0]
if n_items == 0:
return mm_type_ids
positions = [i for i, t in enumerate(mm_type_ids) if t == target]
if not positions:
return mm_type_ids
# Group the marked positions into maximal contiguous runs; each run is one audio's token span.
runs: list[tuple[int, int]] = []
run_start = prev = positions[0]
for pos in positions[1:]:
if pos != prev + 1:
runs.append((run_start, prev))
run_start = pos
prev = pos
runs.append((run_start, prev))
# Layout must match the feature rows one-to-one, else bail rather than corrupt the mapping.
if len(runs) != n_items:
return mm_type_ids
# An audio is complete iff its last placeholder token lands inside the kept window.
n_complete = 0
for _start, end in runs:
if end < max_length:
n_complete += 1
else:
break
if n_complete >= n_items:
return mm_type_ids
# Trim feature rows to the complete prefix.
sample["input_features"] = sample["input_features"][:n_complete]
sample["feature_attention_mask"] = sample["feature_attention_mask"][:n_complete]
# Zero out orphaned placeholder tokens that fall inside the kept window; tokens beyond
# max_length are removed by truncation anyway.
input_ids = list(sample["input_ids"])
mm_type_ids = list(mm_type_ids)
labels = list(sample["labels"]) if "labels" in sample else None
loss_weights = list(sample["loss_weights"]) if "loss_weights" in sample else None
for start, end in runs[n_complete:]:
for pos in range(start, end + 1):
if pos >= max_length:
break
input_ids[pos] = 0
mm_type_ids[pos] = 0
if labels is not None:
labels[pos] = IGNORE_INDEX
if loss_weights is not None:
loss_weights[pos] = 0.0
sample["input_ids"] = input_ids
sample["mm_token_type_ids"] = mm_type_ids
if labels is not None:
sample["labels"] = labels
if loss_weights is not None:
sample["loss_weights"] = loss_weights
return mm_type_ids
def _align_multimodal_on_truncation(sample: ModelInput, max_length: int) -> ModelInput:
"""Remove orphaned multimodal data when the sequence will be truncated.
When cutoff_len truncates input_ids, media whose placeholder tokens are partially cut lose
their token<->feature correspondence. Trims pixel_values/grid_thw (vision) and
input_features/feature_attention_mask (audio) to the complete items and zeros out orphaned
placeholder tokens so the model ignores them.
"""
mm_type_ids = sample.get("mm_token_type_ids")
if mm_type_ids is None:
return sample
sample = dict(sample)
mm_type_ids = _align_grid_media(
sample, mm_type_ids, max_length, target=1, grid_key="image_grid_thw", pixel_key="pixel_values"
)
mm_type_ids = _align_grid_media(
sample, mm_type_ids, max_length, target=2, grid_key="video_grid_thw", pixel_key="pixel_values_videos"
)
mm_type_ids = _align_audio(sample, mm_type_ids, max_length, target=3)
# Remove empty multimodal fields entirely
if "image_grid_thw" in sample and len(sample["image_grid_thw"]) == 0:
del sample["pixel_values"]
del sample["image_grid_thw"]
if "video_grid_thw" in sample and len(sample["video_grid_thw"]) == 0:
del sample["pixel_values_videos"]
del sample["video_grid_thw"]
if "input_features" in sample and sample["input_features"].shape[0] == 0:
del sample["input_features"]
del sample["feature_attention_mask"]
return sample
def pad_and_truncate(samples: list[ModelInput], max_seqlen: int) -> list[BatchInput]:
max_length = min(max(len(sample["input_ids"]) for sample in samples), max_seqlen)
padded_samples = []
for sample in samples:
# Align multimodal fields before truncation: remove images/videos whose
# placeholder tokens would be partially cut, preventing pixel<->token mismatch.
if len(sample["input_ids"]) > max_length and any(k in sample for k in _MULTIMODAL_PASSTHROUGH_KEYS):
sample = _align_multimodal_on_truncation(sample, max_length)
padded_sample = {}
for key, value in sample.items():
if key in _MULTIMODAL_PASSTHROUGH_KEYS:
padded_sample[key] = value
continue
if "label" in key:
pad_value = IGNORE_INDEX
else:
pad_value = 0
if not isinstance(value, str):
padded_sample[key] = _pad_and_truncate(torch.tensor(value), max_length, pad_value)
else:
padded_sample[key] = value
padded_samples.append(padded_sample)
return padded_samples

View File

@@ -14,11 +14,13 @@
import json
import re
from typing import Any, Literal, NotRequired, TypedDict
from ...utils import logging
from ...utils.constants import AUDIO_PLACEHOLDER, IMAGE_PLACEHOLDER, VIDEO_PLACEHOLDER
from ...utils.plugin import BasePlugin
from ...utils.types import DPOSample, Sample, SFTSample, ToolCall
from ...utils.types import Content, DPOSample, Sample, SFTSample, ToolCall
logger = logging.get_logger(__name__)
@@ -29,6 +31,9 @@ class AlpacaSample(TypedDict, total=False):
instruction: str
input: NotRequired[str]
output: str
images: NotRequired[list[str] | str]
videos: NotRequired[list[str] | str]
audios: NotRequired[list[str] | str]
SharegptMessage = TypedDict(
@@ -40,6 +45,9 @@ SharegptMessage = TypedDict(
class SharegptSample(TypedDict, total=False):
conversations: list[SharegptMessage]
tools: NotRequired[str]
images: NotRequired[list[str] | str]
videos: NotRequired[list[str] | str]
audios: NotRequired[list[str] | str]
class OpenaiMessage(TypedDict, total=False):
@@ -54,6 +62,65 @@ class OpenaiSample(TypedDict, total=False):
class PairSample(TypedDict, total=False):
chosen: list[OpenaiMessage]
rejected: list[OpenaiMessage]
images: NotRequired[list[str] | str]
videos: NotRequired[list[str] | str]
audios: NotRequired[list[str] | str]
# Inline media tag -> v1 content block type, and the raw-sample column holding the paths.
_MEDIA_SPECS: tuple[tuple[str, str, str], ...] = (
(IMAGE_PLACEHOLDER, "image_url", "images"),
(VIDEO_PLACEHOLDER, "video_url", "videos"),
(AUDIO_PLACEHOLDER, "audio_url", "audios"),
)
_TAG_TO_BLOCK = {tag: block_type for tag, block_type, _col in _MEDIA_SPECS}
_TAG_PATTERN = re.compile("(" + "|".join(re.escape(tag) for tag, _b, _c in _MEDIA_SPECS) + ")")
def _as_media_list(value: Any) -> list:
"""Normalize a media column value into a list of paths/URLs (None -> [], scalar -> [scalar])."""
if value is None:
return []
if isinstance(value, (list, tuple)):
return list(value)
return [value]
def _build_media_iters(raw_sample: dict[str, Any]) -> dict[str, Any]:
"""Build per-modality path iterators from a raw sample's media columns."""
return {tag: iter(_as_media_list(raw_sample.get(col))) for tag, _block_type, col in _MEDIA_SPECS}
def _to_content_blocks(text: str, media_iters: dict[str, Any]) -> list[Content]:
"""Split ``text`` on inline media placeholders, interleaving media-url content blocks.
Each placeholder consumes the next path from its modality iterator (in document order). Plain
text with no placeholders yields a single text block (byte-identical to the legacy behavior).
Raises on an unmatched placeholder (more tags than media files).
"""
if not _TAG_PATTERN.search(text):
return [{"type": "text", "value": text}]
blocks: list[Content] = []
for segment in _TAG_PATTERN.split(text):
block_type = _TAG_TO_BLOCK.get(segment)
if block_type is not None:
try:
path = next(media_iters[segment])
except StopIteration:
raise ValueError(f"More {segment} tags than provided media files.") from None
blocks.append({"type": block_type, "value": path})
elif segment:
blocks.append({"type": "text", "value": segment})
return blocks
def _assert_media_consumed(media_iters: dict[str, Any]) -> None:
"""Ensure every media file was referenced by a tag (fewer tags than media -> error)."""
for tag, media_iter in media_iters.items():
unused = len(list(media_iter))
if unused:
raise ValueError(f"Fewer {tag} tags than provided media files ({unused} unused).")
class DataConverterPlugin(BasePlugin):
@@ -76,6 +143,7 @@ def alpaca_converter(raw_sample: AlpacaSample) -> SFTSample:
SFTSample: SFT sample.
"""
messages = []
media_iters = _build_media_iters(raw_sample)
if "system" in raw_sample:
messages.append(
{"role": "system", "content": [{"type": "text", "value": raw_sample["system"]}], "loss_weight": 0.0}
@@ -85,9 +153,9 @@ def alpaca_converter(raw_sample: AlpacaSample) -> SFTSample:
messages.append(
{
"role": "user",
"content": [
{"type": "text", "value": raw_sample.get("instruction", "") + raw_sample.get("input", "")}
],
"content": _to_content_blocks(
raw_sample.get("instruction", "") + raw_sample.get("input", ""), media_iters
),
"loss_weight": 0.0,
}
)
@@ -97,6 +165,7 @@ def alpaca_converter(raw_sample: AlpacaSample) -> SFTSample:
{"role": "assistant", "content": [{"type": "text", "value": raw_sample["output"]}], "loss_weight": 1.0}
)
_assert_media_consumed(media_iters)
return {"messages": messages}
@@ -121,6 +190,7 @@ def sharegpt_converter(raw_sample: SharegptSample) -> SFTSample:
}
sample = {}
messages = []
media_iters = _build_media_iters(raw_sample)
for message in raw_sample.get("conversations", []):
tag = message["from"]
if tag not in tag_mapping:
@@ -146,11 +216,12 @@ def sharegpt_converter(raw_sample: SharegptSample) -> SFTSample:
messages.append(
{
"role": tag_mapping[tag],
"content": [{"type": "text", "value": message["value"]}],
"content": _to_content_blocks(message["value"], media_iters),
"loss_weight": 1.0 if tag == "gpt" else 0.0,
}
)
_assert_media_consumed(media_iters)
sample["messages"] = messages
tools = raw_sample.get("tools")
@@ -178,6 +249,8 @@ def pair_converter(raw_sample: PairSample) -> DPOSample:
"""
def process_message(raw_messages: list[OpenaiMessage]):
# chosen and rejected share the sample's media; each side consumes its own iterators.
media_iters = _build_media_iters(raw_sample)
messages = []
for message in raw_messages:
if message["role"] == "tool":
@@ -201,11 +274,12 @@ def pair_converter(raw_sample: PairSample) -> DPOSample:
messages.append(
{
"role": message["role"],
"content": [{"type": "text", "value": message["content"]}],
"content": _to_content_blocks(message["content"], media_iters),
"loss_weight": 1.0 if message["role"] == "assistant" else 0.0,
}
)
_assert_media_consumed(media_iters)
return messages
sample = {}
@@ -221,3 +295,4 @@ def pair_converter(raw_sample: PairSample) -> DPOSample:
logger.warning_rank0(f"Invalid tools format: {str(tools)}")
return sample

View File

@@ -20,6 +20,7 @@ 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.linear_attention.fla import FlashLinearAttentionKernel # 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

View File

@@ -0,0 +1,103 @@
# 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.
"""Flash Linear Attention kernel plugin backed by FSDPTurbo's operator registry."""
from functools import partial
from ......accelerator.helper import DeviceType, get_current_accelerator
from ......utils import logging
from ......utils.types import HFModel
from ...base import BaseKernel, KernelPlugin
logger = logging.get_logger(__name__)
CHUNK_GATED_DELTA_RULE = "chunk_gated_delta_rule"
FUSED_RECURRENT_GATED_DELTA_RULE = "fused_recurrent_gated_delta_rule"
FLASH_LINEAR_ATTENTION_KERNELS = (
CHUNK_GATED_DELTA_RULE,
FUSED_RECURRENT_GATED_DELTA_RULE,
)
FLA_MODULE_ATTRIBUTES = {
CHUNK_GATED_DELTA_RULE: "chunk_gated_delta_rule",
FUSED_RECURRENT_GATED_DELTA_RULE: "recurrent_gated_delta_rule",
}
SUPPORTED_CHUNK_SIZES = (16, 32, 64)
@KernelPlugin("flash-linear-attention").register()
class FlashLinearAttentionKernel(BaseKernel):
"""Install selected FLA callables through FSDPTurbo's device operator registry."""
@staticmethod
def check_device() -> None:
current = get_current_accelerator().type
if current not in (DeviceType.CUDA, DeviceType.NPU):
raise RuntimeError(f"FlashLinearAttentionKernel requires CUDA or NPU, current accelerator is {current}.")
@staticmethod
def check_deps() -> None:
try:
import fla.ops.gated_delta_rule # noqa: F401
import fsdp_turbo.ops.fla # noqa: F401
from fsdp_turbo.ops.registry import get_op # noqa: F401
from fsdp_turbo.utils.patch import patch_model_members # noqa: F401
except ImportError as exc:
raise RuntimeError("Flash Linear Attention and FSDPTurbo are required for this kernel.") from exc
@staticmethod
def _apply(**kwargs) -> HFModel:
model = kwargs["model"]
config = kwargs.get("config") or {}
include_kernels = config.get("include_kernels", "auto")
chunk_size = config.get("chunk_size", 64)
if include_kernels == "auto" or include_kernels is True:
selected = list(FLASH_LINEAR_ATTENTION_KERNELS)
elif isinstance(include_kernels, str):
selected = [name.strip() for name in include_kernels.split(",") if name.strip()]
else:
raise TypeError("kernel_config.include_kernels must be 'auto' or a comma-separated string.")
if not selected:
raise ValueError("kernel_config.include_kernels must select at least one FLA kernel.")
unsupported = set(selected).difference(FLASH_LINEAR_ATTENTION_KERNELS)
if unsupported:
raise ValueError(f"Unsupported Flash Linear Attention kernels: {sorted(unsupported)}")
if isinstance(chunk_size, bool) or not isinstance(chunk_size, int) or chunk_size not in SUPPORTED_CHUNK_SIZES:
raise ValueError(f"chunk_size must be one of {SUPPORTED_CHUNK_SIZES}, got {chunk_size!r}.")
from fsdp_turbo.ops.registry import get_op
from fsdp_turbo.utils.patch import patch_model_members
patched = 0
named_modules = tuple(model.named_modules())
for op_name in selected:
module_attribute = FLA_MODULE_ATTRIBUTES[op_name]
op = get_op(op_name)
configured_op = partial(op, chunk_size=chunk_size) if op_name == CHUNK_GATED_DELTA_RULE else op
targets = {
f"{type(module).__module__}.{type(module).__name__}.{module_attribute}"
for _, module in named_modules
if callable(getattr(module, module_attribute, None))
}
matched = patch_model_members(model, sorted(targets), configured_op) if targets else 0
if matched == 0:
raise RuntimeError(f"FLA operator `{op_name}` did not match any model module attributes.")
patched += matched
logger.info_rank0(f"Flash Linear Attention kernels updated {patched} module callables: {selected}.")
return model

View File

@@ -29,15 +29,21 @@ import torch.nn.functional as F
try:
import torch_npu
except ImportError:
pass
except ImportError as exc:
_TORCH_NPU_IMPORT_ERROR = exc
else:
_TORCH_NPU_IMPORT_ERROR = None
from ......accelerator.helper import DeviceType, get_current_accelerator
from ......utils.logging import get_logger
from ......utils.packages import is_transformers_version_greater_than
from ......utils.types import HFModel
from ...base import BaseKernel, KernelPlugin
logger = get_logger(__name__)
class GmmFunction(torch.autograd.Function):
"""Custom autograd function for NPU Grouped Matrix Multiplication (GMM)."""
@@ -49,7 +55,7 @@ class GmmFunction(torch.autograd.Function):
ctx: Context object to save tensors for backward pass.
x (Tensor): Input tensor.
weight (Tensor): Weight tensor.
group_list (list): List of group sizes.
group_list (Tensor): Number of tokens assigned to each expert.
Returns:
Tensor: The result of the grouped matrix multiplication.
@@ -174,14 +180,14 @@ class HybridGmmFunction(torch.autograd.Function):
return (None, *grad_x_list, *grad_w_list)
class NpuMoeFused:
"""Container for NPU fused MoE forward functions."""
class NpuMoeFusedV4:
"""Container for Transformers v4 NPU fused MoE forward functions."""
@staticmethod
def npu_moe_experts_forward(
def stacked_experts_forward(
self, hidden_states: torch.Tensor, routing_weights: torch.Tensor, router_indices: torch.Tensor
) -> torch.Tensor:
"""Forward pass for MoE experts using NPU fused operations.
"""Forward pass for Transformers v4 MoE experts using NPU fused operations.
Args:
self: The MoE layer instance.
@@ -197,7 +203,9 @@ class NpuMoeFused:
permuted_hidden_states, row_ids_map = torch_npu.npu_moe_token_permute(
hidden_states, router_indices.to(torch.int32)
)
tokens_per_expert = torch.histc(router_indices, bins=self.num_experts, min=0, max=self.num_experts)
tokens_per_expert = torch.histc(
router_indices.float(), bins=self.num_experts, min=0, max=self.num_experts
).long()
intermediate_hidden_states = GmmFunction.apply(permuted_hidden_states, self.gate_up_proj, tokens_per_expert)
intermediate_activations = torch_npu.npu_swiglu(intermediate_hidden_states, dim=-1)
output = GmmFunction.apply(intermediate_activations, self.down_proj, tokens_per_expert)
@@ -206,61 +214,33 @@ class NpuMoeFused:
return next_states
@staticmethod
def npu_moe_sparse_block_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
r"""Forward pass for sparse MoE block using NPU optimization.
def stacked_sparse_block_forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
r"""Forward pass for Transformers v4 sparse MoE block using NPU optimization.
Args:
self: The MoE sparse block instance.
hidden_states (Tensor): Input hidden states.
Returns:
Tensor: The routed output.
tuple: A tuple containing the routed output and router logits.
"""
batch_size = hidden_states.shape[0]
hidden_states = hidden_states.reshape(-1, self.hidden_size)
router_logits = self.gate(hidden_states)
routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float)
routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
routing_weights, router_indices = torch.topk(routing_weights, self.top_k, dim=-1)
routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
routing_weights = routing_weights.to(hidden_states.dtype)
hidden_states = hidden_states.reshape(batch_size, -1, self.hidden_size)
routed_out = self.experts(hidden_states, routing_weights, router_indices)
return routed_out
return routed_out, router_logits
@staticmethod
def npu_moe_experts_v5_forward(
self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor
) -> torch.Tensor:
"""Forward pass for Transformers v5+ MoE experts using NPU fused operations.
Transformers v5 stores expert weights in F.linear layout:
gate_up_proj: [num_experts, 2 * intermediate_dim, hidden_dim]
down_proj: [num_experts, hidden_dim, intermediate_dim]
The NPU grouped matmul path expects matmul layout, so both weights are transposed.
"""
hidden_states = hidden_states.reshape(-1, self.hidden_dim)
permuted_hidden_states, row_ids_map = torch_npu.npu_moe_token_permute(
hidden_states, top_k_index.to(torch.int32)
)
tokens_per_expert = torch.histc(top_k_index.float(), bins=self.num_experts, min=0, max=self.num_experts).long()
gate_up_proj = self.gate_up_proj.transpose(1, 2)
down_proj = self.down_proj.transpose(1, 2)
intermediate_hidden_states = GmmFunction.apply(permuted_hidden_states, gate_up_proj, tokens_per_expert)
intermediate_activations = torch_npu.npu_swiglu(intermediate_hidden_states, dim=-1)
output = GmmFunction.apply(intermediate_activations, down_proj, tokens_per_expert)
return torch_npu.npu_moe_token_unpermute(output, row_ids_map, probs=top_k_weights)
class Qwen3NpuMoeFused:
"""Container for Qwen3 NPU fused MoE forward functions."""
@staticmethod
def qwen3moe_sparse_moe_block_forward(self, hidden_states: torch.Tensor):
"""Forward pass for Qwen3 sparse MoE block using NPU fused operations.
def sparse_block_forward(self, hidden_states: torch.Tensor):
"""Forward pass for a Transformers v4 list-backed sparse MoE block using NPU fused operations.
Args:
self: The Qwen3 MoE block instance.
self: The sparse MoE block instance.
hidden_states (Tensor): Input hidden states.
Returns:
@@ -304,33 +284,90 @@ class Qwen3NpuMoeFused:
next_states = next_states.view(batch_size, sequence_length, -1)
return next_states, router_logits
@staticmethod
def shared_sparse_block_forward(self, hidden_states: torch.Tensor):
"""Forward pass for a Transformers v4 sparse MoE block with a shared expert."""
next_states, router_logits = NpuMoeFusedV4.sparse_block_forward(self, hidden_states)
# moe patch config mapping
if is_transformers_version_greater_than("5.0.0"):
kernel_moe_mapping = {
"Qwen3MoeForCausalLM": {
"Qwen3MoeExperts": NpuMoeFused.npu_moe_experts_v5_forward,
},
"Qwen3VLMoeForConditionalGeneration": {
"Qwen3VLMoeTextExperts": NpuMoeFused.npu_moe_experts_v5_forward,
},
"Qwen3_5MoeForCausalLM": {
"Qwen3_5MoeExperts": NpuMoeFused.npu_moe_experts_v5_forward,
},
"Qwen3_5MoeForConditionalGeneration": {
"Qwen3_5MoeExperts": NpuMoeFused.npu_moe_experts_v5_forward,
},
}
else:
kernel_moe_mapping = {
"Qwen3MoeForCausalLM": {
"Qwen3MoeSparseMoeBlock": Qwen3NpuMoeFused.qwen3moe_sparse_moe_block_forward,
},
"Qwen3VLMoeForConditionalGeneration": {
"Qwen3VLMoeTextExperts": NpuMoeFused.npu_moe_experts_forward,
"Qwen3VLMoeTextSparseMoeBlock": NpuMoeFused.npu_moe_sparse_block_forward,
},
}
shared_expert_output = self.shared_expert(hidden_states)
shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output
next_states = next_states + shared_expert_output
return next_states, router_logits
class NpuMoeFusedV5:
"""Container for Transformers v5 NPU fused MoE forward functions."""
@staticmethod
def experts_forward(
self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor
) -> torch.Tensor:
"""Forward pass for Transformers v5+ MoE experts using NPU fused operations.
Transformers v5 stores expert weights in F.linear layout:
gate_up_proj: [num_experts, 2 * intermediate_dim, hidden_dim]
down_proj: [num_experts, hidden_dim, intermediate_dim]
The NPU grouped matmul path expects matmul layout, so both weights are transposed.
"""
hidden_states = hidden_states.reshape(-1, self.hidden_dim)
permuted_hidden_states, row_ids_map = torch_npu.npu_moe_token_permute(
hidden_states, top_k_index.to(torch.int32)
)
tokens_per_expert = torch.histc(top_k_index.float(), bins=self.num_experts, min=0, max=self.num_experts).long()
gate_up_proj = self.gate_up_proj.transpose(1, 2)
down_proj = self.down_proj.transpose(1, 2)
intermediate_hidden_states = GmmFunction.apply(permuted_hidden_states, gate_up_proj, tokens_per_expert)
intermediate_activations = torch_npu.npu_swiglu(intermediate_hidden_states, dim=-1)
output = GmmFunction.apply(intermediate_activations, down_proj, tokens_per_expert)
return torch_npu.npu_moe_token_unpermute(output, row_ids_map, probs=top_k_weights)
_V4_MODEL_TYPE_TO_PATCHES = {
"qwen3_moe": {
"Qwen3MoeSparseMoeBlock": NpuMoeFusedV4.sparse_block_forward,
},
"qwen3_next": {
"Qwen3NextSparseMoeBlock": NpuMoeFusedV4.shared_sparse_block_forward,
},
"qwen3_omni_moe": {
"Qwen3OmniMoeThinkerTextSparseMoeBlock": NpuMoeFusedV4.sparse_block_forward,
"Qwen3OmniMoeTalkerTextSparseMoeBlock": NpuMoeFusedV4.shared_sparse_block_forward,
},
"qwen3_omni_moe_thinker": {
"Qwen3OmniMoeThinkerTextSparseMoeBlock": NpuMoeFusedV4.sparse_block_forward,
},
"qwen3_vl_moe": {
"Qwen3VLMoeTextExperts": NpuMoeFusedV4.stacked_experts_forward,
"Qwen3VLMoeTextSparseMoeBlock": NpuMoeFusedV4.stacked_sparse_block_forward,
},
}
_V5_MODEL_TYPE_TO_PATCHES = {
"qwen3_moe": {
"Qwen3MoeExperts": NpuMoeFusedV5.experts_forward,
},
"qwen3_next": {
"Qwen3NextExperts": NpuMoeFusedV5.experts_forward,
},
"qwen3_omni_moe": {
"Qwen3OmniMoeThinkerTextExperts": NpuMoeFusedV5.experts_forward,
"Qwen3OmniMoeTalkerTextExperts": NpuMoeFusedV5.experts_forward,
},
"qwen3_omni_moe_thinker": {
"Qwen3OmniMoeThinkerTextExperts": NpuMoeFusedV5.experts_forward,
},
"qwen3_vl_moe": {
"Qwen3VLMoeTextExperts": NpuMoeFusedV5.experts_forward,
},
"qwen3_5_moe": {
"Qwen3_5MoeExperts": NpuMoeFusedV5.experts_forward,
},
}
_MODEL_TYPE_TO_PATCHES = (
_V5_MODEL_TYPE_TO_PATCHES if is_transformers_version_greater_than("5.0.0") else _V4_MODEL_TYPE_TO_PATCHES
)
@KernelPlugin("npu_fused_moe").register()
@@ -343,6 +380,17 @@ class NpuFusedMoEKernel(BaseKernel):
if current != DeviceType.NPU:
raise RuntimeError(f"NpuFusedMoEKernel requires NPU, current accelerator is {current}.")
@staticmethod
def check_deps() -> None:
if _TORCH_NPU_IMPORT_ERROR is not None:
raise RuntimeError("NpuFusedMoEKernel requires torch_npu.") from _TORCH_NPU_IMPORT_ERROR
@staticmethod
def _get_patch_forward(model_type: str, module: torch.nn.Module):
"""Return the version-specific NPU forward function for a matched MoE module."""
model_patches = _MODEL_TYPE_TO_PATCHES.get(model_type, {})
return model_patches.get(module.__class__.__name__)
@staticmethod
def _apply(**kwargs) -> HFModel:
"""Applies the NPU fused MoE kernel to the model.
@@ -352,27 +400,21 @@ class NpuFusedMoEKernel(BaseKernel):
Returns:
HFModel: The model with patched MoE forward functions.
Raises:
ValueError: If the model is not provided.
RuntimeError: If dependencies are not met.
"""
model = kwargs.get("model", None)
model = kwargs["model"]
archs = getattr(model.config, "architectures", None) or []
target_moe_mapping = None
for arch in archs:
if arch in kernel_moe_mapping:
target_moe_mapping = kernel_moe_mapping[arch]
break
if target_moe_mapping is None:
model_type = getattr(model.config, "model_type", None)
if model_type not in _MODEL_TYPE_TO_PATCHES:
return model
patched_count = 0
for module in model.modules():
class_name = module.__class__.__name__
if class_name in target_moe_mapping:
new_forward_func = target_moe_mapping[class_name]
module.forward = types.MethodType(new_forward_func, module)
patch_forward = NpuFusedMoEKernel._get_patch_forward(model_type, module)
if patch_forward is not None:
module.forward = types.MethodType(patch_forward, module)
patched_count += 1
if patched_count:
logger.info_rank0(f"Applied NPU fused MoE kernel to {patched_count} modules for model type: {model_type}.")
return model

View File

@@ -20,20 +20,24 @@ Init Phase:
"""
import re
import types
import torch
from ......accelerator.helper import DeviceType, get_current_accelerator
from ......utils.logging import get_logger
from ......utils.types import HFModel
from ...base import BaseKernel, KernelPlugin
logger = get_logger(__name__)
try:
import torch_npu
except ImportError:
pass
except ImportError as exc:
_TORCH_NPU_IMPORT_ERROR = exc
else:
_TORCH_NPU_IMPORT_ERROR = None
def npu_swiglu_forward(self, hidden_state):
@@ -51,81 +55,69 @@ def npu_swiglu_forward(self, hidden_state):
)
def _npu_swiglu_glm4_forward(self, hidden_states):
"""SwiGLU forward pass for GLM4 on NPU.
Args:
self: The GLM4 MLP layer instance.
hidden_states (Tensor): Input hidden states.
Returns:
Tensor: Output of SwiGLU.
"""
up_states = self.gate_up_proj(hidden_states)
gate, up_states = up_states.chunk(2, dim=-1)
return self.down_proj(torch_npu.npu_swiglu(torch.cat((gate, up_states), dim=-1), dim=-1))
def _npu_swiglu_gemma3ntext_forward(self, hidden_states):
"""SwiGLU forward pass for Gemma3nText on NPU.
Args:
self: The Gemma3nText MLP layer instance.
hidden_states (Tensor): Input hidden states.
Returns:
Tensor: Output of SwiGLU.
"""
gate_proj = self.gate_proj(hidden_states)
if self.activation_sparsity > 0.0:
gate_proj = self._gaussian_topk(gate_proj)
down_proj = self.down_proj(
torch_npu.npu_swiglu(torch.cat((gate_proj, self.up_proj(hidden_states)), dim=-1), dim=-1)
)
return down_proj
_MODEL_TYPE_TO_PATCHES = {
"qwen3": {
"Qwen3MLP": npu_swiglu_forward,
},
"qwen3_moe": {
"Qwen3MoeMLP": npu_swiglu_forward,
},
"qwen3_next": {
"Qwen3NextMLP": npu_swiglu_forward,
},
"qwen3_omni_moe": {
"Qwen3OmniMoeThinkerTextMLP": npu_swiglu_forward,
"Qwen3OmniMoeMLP": npu_swiglu_forward,
"Qwen3OmniMoeTalkerTextMLP": npu_swiglu_forward,
"Qwen3OmniMoeCode2WavMlp": npu_swiglu_forward,
},
"qwen3_omni_moe_thinker": {
"Qwen3OmniMoeThinkerTextMLP": npu_swiglu_forward,
},
"qwen3_vl": {
"Qwen3VLTextMLP": npu_swiglu_forward,
},
"qwen3_vl_moe": {
"Qwen3VLMoeTextMLP": npu_swiglu_forward,
},
"qwen3_5": {
"Qwen3_5MLP": npu_swiglu_forward,
},
"qwen3_5_moe": {
"Qwen3_5MoeMLP": npu_swiglu_forward,
},
}
@KernelPlugin("npu_fused_swiglu").register()
class NpuSwiGluKernel(BaseKernel):
"""NPU Kernel for fused SwiGLU activation."""
# just support apply to the following module layers
expect_modules = frozenset(
{
"Qwen3VLMoeTextMLP",
"Qwen3VLTextMLP",
"Qwen3OmniMoeThinkerTextMLP",
"Qwen3OmniMoeMLP",
"Qwen3OmniMoeTalkerTextMLP",
"Qwen3OmniMoeCode2WavMlp",
"Qwen3NextMLP",
"Qwen3MoeMLP",
"Qwen3MLP",
"Qwen2MLP",
"Qwen2MoeMLP",
"Qwen2_5_VLMLP",
"Qwen2_5OmniMLP",
"Llama4TextMLP",
"LlamaMLP",
"Glm4MLP",
"Glm4MoeMLP",
"Glm4vMoeTextMLP",
"Gemma3MLP",
"Gemma2MLP",
"Gemma3nTextMLP",
"Phi3MLP",
"DeepseekV2MLP",
"DeepseekV3MLP",
"SeedOssMLP",
}
)
@staticmethod
def check_device() -> None:
current = get_current_accelerator().type
if current != DeviceType.NPU:
raise RuntimeError(f"NpuSwiGluKernel requires NPU, current accelerator is {current}.")
@staticmethod
def check_deps() -> None:
if _TORCH_NPU_IMPORT_ERROR is not None:
raise RuntimeError("NpuSwiGluKernel requires torch_npu.") from _TORCH_NPU_IMPORT_ERROR
@staticmethod
def _get_patch_forward(model_type: str, module: torch.nn.Module):
"""Return the NPU forward function for a matched SwiGLU MLP module."""
model_patches = _MODEL_TYPE_TO_PATCHES.get(model_type, {})
patch_forward = model_patches.get(module.__class__.__name__)
if patch_forward is None:
return None
config = getattr(module, "config", None)
if getattr(config, "hidden_act", None) != "silu":
return None
return patch_forward
@staticmethod
def _apply(**kwargs) -> "HFModel":
"""Applies the NPU fused SwiGLU kernel to the model.
@@ -135,31 +127,21 @@ class NpuSwiGluKernel(BaseKernel):
Returns:
HFModel: The model with patched SwiGLU forward functions.
Raises:
ValueError: If the model is not provided.
RuntimeError: If dependencies are not met.
"""
model = kwargs.get("model", None)
model = kwargs["model"]
# Mapping of specific mlp modules to their corresponding kernel implementations
kernel_mapping = {
"Glm4MLP": _npu_swiglu_glm4_forward,
"Glm4vTextMLP": _npu_swiglu_glm4_forward,
"Phi3MLP": _npu_swiglu_glm4_forward,
"Gemma3nTextMLP": _npu_swiglu_gemma3ntext_forward,
}
model_type = getattr(model.config, "model_type", None)
if model_type not in _MODEL_TYPE_TO_PATCHES:
return model
swiglu_pattern = re.compile("MLP", re.IGNORECASE)
for name, module in model.named_modules():
# Match any module whose class name contains "MLP"
if (
re.search(swiglu_pattern, module.__class__.__name__)
and module.__class__.__name__ in NpuSwiGluKernel.expect_modules
):
# Bind function as an instance method to preserve `self` semantics
# and replace the original forward
kernel_func = kernel_mapping.get(module.__class__.__name__, npu_swiglu_forward)
module.forward = types.MethodType(kernel_func, module)
patched_count = 0
for module in model.modules():
patch_forward = NpuSwiGluKernel._get_patch_forward(model_type, module)
if patch_forward is not None:
module.forward = types.MethodType(patch_forward, module)
patched_count += 1
if patched_count:
logger.info_rank0(f"Applied NPU SwiGLU kernel to {patched_count} modules for model type: {model_type}.")
return model

View File

@@ -20,54 +20,32 @@ Init Phase:
"""
import re
import types
import torch
import torch.nn.functional as F
from ......accelerator.helper import DeviceType, get_current_accelerator
from ......utils.logging import get_logger
from ......utils.types import HFModel
from ...base import BaseKernel, KernelPlugin
logger = get_logger(__name__)
try:
import torch_npu
except ImportError:
pass
def _should_use_residual_rmsnorm(module):
"""Detect whether the module uses residual RMSNorm parameterization.
Residual RMSNorm uses ``scale = 1.0 + weight`` where weight is initialized to 0,
while standard RMSNorm uses ``scale = weight`` where weight is initialized to 1.
Args:
module (nn.Module): The RMSNorm module to check.
Returns:
bool: ``True`` if the module uses residual parameterization, ``False`` otherwise.
.. note::
This must follow the module's forward semantics. Do not infer it from trained
weight values because standard RMSNorm weights can also be close to zero.
"""
residual_rmsnorm_classes = {
"Qwen3_5RMSNorm",
"Qwen3_5MoeRMSNorm",
"Qwen3NextRMSNorm",
}
class_name = module.__class__.__name__
return class_name in residual_rmsnorm_classes
except ImportError as exc:
_TORCH_NPU_IMPORT_ERROR = exc
else:
_TORCH_NPU_IMPORT_ERROR = None
def npu_rms_norm_forward(self, hidden_states):
"""NPU forward implementation for standard RMSNorm.
Args:
self (nn.Module): The RMSNorm module instance with ``weight`` and ``variance_epsilon``.
self (nn.Module): The RMSNorm module instance with ``weight`` and either ``variance_epsilon`` or ``eps``.
hidden_states (Tensor): Input hidden states tensor.
Returns:
@@ -75,48 +53,108 @@ def npu_rms_norm_forward(self, hidden_states):
"""
_eps = getattr(self, "variance_epsilon", None) or getattr(self, "eps", 1e-6)
if hasattr(self, "weight") and self.weight is not None:
if getattr(self, "_npu_use_residual_rmsnorm", False):
effective_weight = 1.0 + self.weight.float()
else:
effective_weight = self.weight.float()
else:
effective_weight = None
weight = getattr(self, "weight", None)
if weight is None:
raise RuntimeError(f"{self.__class__.__name__} has no RMSNorm weight for NPU RMSNorm kernel.")
if effective_weight is not None:
return torch_npu.npu_rms_norm(hidden_states, effective_weight.to(hidden_states.dtype), epsilon=_eps)[0]
else:
return torch_npu.npu_rms_norm(hidden_states, self.weight, epsilon=_eps)[0]
effective_weight = weight.float()
return torch_npu.npu_rms_norm(hidden_states, effective_weight.to(hidden_states.dtype), epsilon=_eps)[0]
def npu_residual_rms_norm_forward(self, hidden_states):
"""NPU forward implementation for residual RMSNorm.
Residual RMSNorm uses ``scale = 1.0 + weight`` where ``weight`` is initialized
to 0 in the original transformers implementation.
Args:
self (nn.Module): The residual RMSNorm module with ``weight`` and either ``variance_epsilon`` or ``eps``.
hidden_states (Tensor): Input hidden states tensor.
Returns:
Tensor: Normalized tensor consistent with residual RMSNorm behavior.
"""
_eps = getattr(self, "variance_epsilon", None) or getattr(self, "eps", 1e-6)
weight = getattr(self, "weight", None)
if weight is None:
raise RuntimeError(f"{self.__class__.__name__} has no RMSNorm weight for NPU RMSNorm kernel.")
effective_weight = 1.0 + weight.float()
return torch_npu.npu_rms_norm(hidden_states, effective_weight.to(hidden_states.dtype), epsilon=_eps)[0]
def npu_gated_rms_norm_forward(self, hidden_states, gate=None):
"""NPU forward implementation for Gated RMSNorm with high-precision FP32 computation.
This function performs RMSNorm and gated SiLU multiplication in FP32 for numerical
stability. Unlike standard RMSNorm, Gated RMSNorm in Qwen3.5 uses standard
parameterization (``scale = weight`` where weight is initialized to 1), so the
residual weight adjustment (``1.0 + weight``) is not applied here.
stability. The supported gated RMSNorm modules use ``scale = weight`` with weight
initialized to 1, unlike the residual RMSNorm variants that use ``1.0 + weight``.
Args:
self (nn.Module): The Gated RMSNorm module instance.
hidden_states (Tensor): Input hidden states tensor.
gate (Tensor, optional): Gate tensor for SiLU activation. Defaults to ``None``.
gate (Tensor): Gate tensor for SiLU activation.
Returns:
Tensor: Output tensor cast back to the original input dtype.
Raises:
ValueError: If the gate tensor is not provided.
"""
if gate is None:
raise ValueError(f"{self.__class__.__name__} requires a gate tensor for NPU Gated RMSNorm.")
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
_eps = getattr(self, "variance_epsilon", None) or getattr(self, "eps", 1e-6)
hidden_states = torch_npu.npu_rms_norm(hidden_states, self.weight.float(), epsilon=_eps)[0]
if gate is not None:
hidden_states = hidden_states * F.silu(gate.to(torch.float32))
hidden_states = hidden_states * F.silu(gate.to(torch.float32))
return hidden_states.to(input_dtype)
_MODEL_TYPE_TO_PATCHES = {
"qwen3": {
"Qwen3RMSNorm": npu_rms_norm_forward,
},
"qwen3_moe": {
"Qwen3MoeRMSNorm": npu_rms_norm_forward,
},
"qwen3_next": {
"Qwen3NextRMSNorm": npu_residual_rms_norm_forward,
"Qwen3NextRMSNormGated": npu_gated_rms_norm_forward,
},
"qwen3_omni_moe": {
"Qwen3OmniMoeThinkerTextRMSNorm": npu_rms_norm_forward,
"Qwen3OmniMoeTextRMSNorm": npu_rms_norm_forward,
"Qwen3OmniMoeRMSNorm": npu_rms_norm_forward,
"Qwen3OmniMoeCode2WavRMSNorm": npu_rms_norm_forward,
},
"qwen3_omni_moe_thinker": {
"Qwen3OmniMoeThinkerTextRMSNorm": npu_rms_norm_forward,
"Qwen3OmniMoeTextRMSNorm": npu_rms_norm_forward,
},
"qwen3_vl": {
"Qwen3VLTextRMSNorm": npu_rms_norm_forward,
},
"qwen3_vl_moe": {
"Qwen3VLMoeTextRMSNorm": npu_rms_norm_forward,
},
"qwen3_5": {
"Qwen3_5RMSNorm": npu_residual_rms_norm_forward,
"Qwen3_5RMSNormGated": npu_gated_rms_norm_forward,
},
"qwen3_5_moe": {
"Qwen3_5MoeRMSNorm": npu_residual_rms_norm_forward,
"Qwen3_5MoeRMSNormGated": npu_gated_rms_norm_forward,
},
}
@KernelPlugin("npu_fused_rmsnorm").register()
class NpuRMSNormKernel(BaseKernel):
"""NPU kernel wrapper for RMSNorm that applies the replacement within a model."""
@@ -127,34 +165,45 @@ class NpuRMSNormKernel(BaseKernel):
if current != DeviceType.NPU:
raise RuntimeError(f"NpuRMSNormKernel requires NPU, current accelerator is {current}.")
@staticmethod
def check_deps() -> None:
if _TORCH_NPU_IMPORT_ERROR is not None:
raise RuntimeError("NpuRMSNormKernel requires torch_npu.") from _TORCH_NPU_IMPORT_ERROR
@staticmethod
def _get_patch_forward(model_type: str, module: torch.nn.Module):
"""Return the NPU forward function for a matched RMSNorm module."""
model_patches = _MODEL_TYPE_TO_PATCHES.get(model_type, {})
return model_patches.get(module.__class__.__name__)
@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
the appropriate NPU-optimized forward function as an instance method via
``types.MethodType`` to replace the original ``forward``.
Matches modules configured for the current model type, then binds the corresponding
NPU-optimized forward function as an instance method via ``types.MethodType`` to
replace the original ``forward``.
Args:
**kwargs: Keyword arguments containing the model.
Returns:
HFModel: The model with NPU fused RMSNorm.
Raises:
RuntimeError: If ``torch_npu`` is not available.
ValueError: If the model is not provided.
"""
model = kwargs.get("model")
model = kwargs["model"]
rms_norm_pattern = re.compile("RMSNorm", re.IGNORECASE)
model_type = getattr(model.config, "model_type", None)
if model_type not in _MODEL_TYPE_TO_PATCHES:
return model
for _, module in model.named_modules():
if re.search(rms_norm_pattern, module.__class__.__name__):
if "Gated" in module.__class__.__name__:
module.forward = types.MethodType(npu_gated_rms_norm_forward, module)
else:
module._npu_use_residual_rmsnorm = _should_use_residual_rmsnorm(module)
module.forward = types.MethodType(npu_rms_norm_forward, module)
patched_count = 0
for module in model.modules():
patch_forward = NpuRMSNormKernel._get_patch_forward(model_type, module)
if patch_forward is not None:
module.forward = types.MethodType(patch_forward, module)
patched_count += 1
if patched_count:
logger.info_rank0(f"Applied NPU RMSNorm kernel to {patched_count} modules for model type: {model_type}.")
return model

View File

@@ -20,7 +20,7 @@ Init Phase:
"""
import sys
import importlib
import torch
@@ -34,16 +34,18 @@ logger = get_logger(__name__)
try:
import torch_npu
except ImportError:
pass
except ImportError as exc:
_TORCH_NPU_IMPORT_ERROR = exc
else:
_TORCH_NPU_IMPORT_ERROR = None
def _apply_npu_rotary_emb(q, k, cos, sin):
"""Apply NPU-accelerated rotary embedding with automatic Partial RoPE detection.
This function automatically detects whether to use Partial RoPE or Full RoPE
based on the dimension ratio between ``cos/sin`` and ``q/k`` tensors, ensuring
compatibility with future model versions without hardcoding.
Partial RoPE is detected when the ``cos/sin`` width is smaller than the ``q/k``
head dimension. The leading rotary dimensions are transformed and any trailing
dimensions are passed through unchanged.
Args:
q (Tensor): Query tensor.
@@ -61,14 +63,14 @@ def _apply_npu_rotary_emb(q, k, cos, sin):
q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
q_embed = torch_npu.npu_rotary_mul(q_rot, cos, sin).to(q.dtype)
k_embed = torch_npu.npu_rotary_mul(k_rot, cos, sin).to(k.dtype)
q_embed = torch_npu.npu_rotary_mul(q_rot, cos, sin, "half").to(q.dtype)
k_embed = torch_npu.npu_rotary_mul(k_rot, cos, sin, "half").to(k.dtype)
q_embed = torch.cat([q_embed, q_pass], dim=-1)
k_embed = torch.cat([k_embed, k_pass], dim=-1)
else:
q_embed = torch_npu.npu_rotary_mul(q, cos, sin).to(q.dtype)
k_embed = torch_npu.npu_rotary_mul(k, cos, sin).to(k.dtype)
q_embed = torch_npu.npu_rotary_mul(q, cos, sin, "half").to(q.dtype)
k_embed = torch_npu.npu_rotary_mul(k, cos, sin, "half").to(k.dtype)
return q_embed, k_embed
@@ -84,44 +86,43 @@ def _apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
k (Tensor): Key tensor.
cos (Tensor): Cosine part of embedding.
sin (Tensor): Sine part of embedding.
position_ids (Tensor, optional): Position IDs. Defaults to ``None``.
position_ids (Tensor | int, optional): Ignored Transformers v4 position IDs, or the Transformers v5
``unsqueeze_dim`` when supplied as the fifth positional argument.
unsqueeze_dim (int): Dimension to unsqueeze cos and sin. Defaults to 1.
Returns:
tuple[Tensor, Tensor]: The embedded query and key tensors ``(q_embed, k_embed)``.
"""
# In transformers v5, the fifth positional argument is ``unsqueeze_dim``.
if isinstance(position_ids, int):
unsqueeze_dim = position_ids
cos = cos.unsqueeze(unsqueeze_dim)
sin = sin.unsqueeze(unsqueeze_dim)
return _apply_npu_rotary_emb(q, k, cos, sin)
def _apply_multimodal_rotary_pos_emb_qwen25_vl(q, k, cos, sin, mrope_section, unsqueeze_dim=1):
"""Apply Rotary Position Embedding with multimodal sections (Qwen2-VL) on NPU.
This function supports Partial RoPE for multimodal inputs with automatic dimension
detection, ensuring compatibility with future model versions.
Args:
q (Tensor): Query tensor.
k (Tensor): Key tensor.
cos (Tensor): Cosine part of embedding.
sin (Tensor): Sine part of embedding.
mrope_section (list[int]): Multimodal RoPE section sizes.
unsqueeze_dim (int): Dimension to unsqueeze cos and sin. Defaults to 1.
Returns:
tuple[Tensor, Tensor]: The embedded query and key tensors ``(q_embed, k_embed)``.
"""
mrope_section = mrope_section * 2
cos = torch.cat([m[i % 3] for i, m in enumerate(cos.split(mrope_section, dim=-1))], dim=-1).unsqueeze(
unsqueeze_dim
)
sin = torch.cat([m[i % 3] for i, m in enumerate(sin.split(mrope_section, dim=-1))], dim=-1).unsqueeze(
unsqueeze_dim
def _default_rope_patch(module_type: str):
return (
(
f"transformers.models.{module_type}.modeling_{module_type}",
(("apply_rotary_pos_emb", _apply_rotary_pos_emb),),
),
)
return _apply_npu_rotary_emb(q, k, cos, sin)
_MODEL_TYPE_TO_PATCHES = {
"qwen3": _default_rope_patch("qwen3"),
"qwen3_moe": _default_rope_patch("qwen3_moe"),
"qwen3_next": _default_rope_patch("qwen3_next"),
"qwen3_omni_moe": _default_rope_patch("qwen3_omni_moe"),
"qwen3_omni_moe_thinker": _default_rope_patch("qwen3_omni_moe"),
"qwen3_vl": _default_rope_patch("qwen3_vl"),
"qwen3_vl_moe": _default_rope_patch("qwen3_vl_moe"),
"qwen3_5": _default_rope_patch("qwen3_5"),
"qwen3_5_moe": _default_rope_patch("qwen3_5_moe"),
}
@KernelPlugin("npu_fused_rope").register()
@@ -135,50 +136,58 @@ class NpuRoPEKernel(BaseKernel):
raise RuntimeError(f"NpuRoPEKernel requires NPU, current accelerator is {current}.")
@staticmethod
def _apply(**kwargs) -> "HFModel":
"""Apply RoPE acceleration by monkey-patching ``apply_rotary_pos_emb``.
def check_deps() -> None:
if _TORCH_NPU_IMPORT_ERROR is not None:
raise RuntimeError("NpuRoPEKernel requires torch_npu.") from _TORCH_NPU_IMPORT_ERROR
Iterates through the model's modules to find attention layers, identifies
the module where they are defined, and replaces the original
``apply_rotary_pos_emb`` function in that module's namespace with the
NPU-accelerated version.
@staticmethod
def _apply_model_patches(model_type: str) -> int:
patches = _MODEL_TYPE_TO_PATCHES.get(model_type)
if patches is None:
return 0
patched_count = 0
for module_name, replacements in patches:
try:
target_module = importlib.import_module(module_name)
except Exception as e:
logger.warning_rank0_once(f"Failed to import {module_name} for NPU RoPE kernel: {e}")
continue
for target_function_name, replacement in replacements:
if not hasattr(target_module, target_function_name):
logger.warning_rank0_once(f"{module_name} has no {target_function_name}, skip NPU RoPE patch.")
continue
if getattr(target_module, target_function_name) is replacement:
continue
setattr(target_module, target_function_name, replacement)
patched_count += 1
return patched_count
@staticmethod
def _apply(**kwargs) -> "HFModel":
"""Apply RoPE acceleration by monkey-patching rotary embedding functions.
Selects the target transformers modeling module from ``model.config.model_type``
and replaces its rotary embedding helper with the NPU-accelerated version.
Args:
**kwargs: Keyword arguments containing the model.
Returns:
HFModel: The model with patched RoPE functions.
Raises:
RuntimeError: If ``torch_npu`` is not available.
ValueError: If the model is not provided.
"""
model = kwargs.get("model", None)
model = kwargs["model"]
_modules = set()
for module in model.modules():
if "Attention" in module.__class__.__name__:
module_name = module.__class__.__module__
if module_name in _modules:
continue
try:
target_module = sys.modules[module_name]
if hasattr(target_module, "apply_rotary_pos_emb"):
if getattr(target_module, "apply_rotary_pos_emb") is not _apply_rotary_pos_emb:
setattr(target_module, "apply_rotary_pos_emb", _apply_rotary_pos_emb)
_modules.add(module_name)
if hasattr(target_module, "apply_multimodal_rotary_pos_emb"):
if (
getattr(target_module, "apply_multimodal_rotary_pos_emb")
is not _apply_multimodal_rotary_pos_emb_qwen25_vl
):
setattr(
target_module,
"apply_multimodal_rotary_pos_emb",
_apply_multimodal_rotary_pos_emb_qwen25_vl,
)
_modules.add(module_name)
except Exception as e:
logger.warning_rank0_once(f"Failed to apply RoPE kernel to module {module_name}: {e}")
model_type = getattr(model.config, "model_type", None)
if model_type not in _MODEL_TYPE_TO_PATCHES:
return model
patched_count = NpuRoPEKernel._apply_model_patches(model_type)
if patched_count:
logger.info_rank0(f"Applied NPU RoPE kernel to {patched_count} functions for model type: {model_type}.")
return model

View File

@@ -52,6 +52,14 @@ def get_ulysses_sequence_parallel_rank(group: ProcessGroup = None) -> int:
return dist.get_rank(group) if group else 0
def _get_text_position_ids(position_ids: Optional[Tensor]) -> Optional[Tensor]:
# Transformers < 5.4 broadcasts Qwen3.5 text positions over the mRoPE axes.
if position_ids is not None and position_ids.ndim == 3 and position_ids.stride(0) == 0:
position_ids = position_ids[0]
return position_ids.contiguous() if position_ids is not None and position_ids.ndim == 2 else None
class UlyssesAttention(torch.nn.Module):
"""Initialization.
@@ -123,8 +131,8 @@ class UlyssesAttention(torch.nn.Module):
softmax_scale = q.shape[-1] ** -0.5
sp_world_size = get_ulysses_sequence_parallel_world_size(self.spg)
local_position_ids = position_ids
# HF FlashAttention only uses 2-D position IDs to detect packed sequences.
position_ids = _get_text_position_ids(position_ids)
if position_ids is not None:
global_position_ids = [torch.empty_like(position_ids) for _ in range(sp_world_size)]
dist.all_gather(global_position_ids, position_ids, group=self.spg)
@@ -144,13 +152,11 @@ class UlyssesAttention(torch.nn.Module):
# contribute an all-ones shard.
if torch.any(torch.stack(global_has_attention_mask)):
if attention_mask is None:
if local_position_ids is not None:
attention_mask = torch.ones_like(local_position_ids, dtype=torch.int64)
else:
attention_mask = torch.ones(query.shape[0], query.shape[1], dtype=torch.int64, device=query.device)
attention_mask = torch.ones(query.shape[0], query.shape[1], dtype=torch.int64, device=query.device)
else:
attention_mask = attention_mask.to(torch.int64)
attention_mask = attention_mask.contiguous()
global_attention_mask = [torch.empty_like(attention_mask) for _ in range(sp_world_size)]
dist.all_gather(global_attention_mask, attention_mask, group=self.spg)
attention_mask = torch.cat(global_attention_mask, dim=1).contiguous()

View File

@@ -1,4 +1,4 @@
# Copyright 2025 the LlamaFactory team.
# Copyright 2026 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.
@@ -20,8 +20,8 @@ from typing import Any
import torch
from torch.utils.data import default_collate
from ...core.utils.collation import pad_and_truncate
from ...utils.constants import IGNORE_INDEX
from ...utils.helper import pad_and_truncate
from ...utils.objects import StatefulBuffer
from ...utils.plugin import BasePlugin, ensure_methods_implemented
from ...utils.types import BatchInfo, BatchInput, DataLoader, ModelInput

View File

@@ -34,6 +34,14 @@ from ...model_plugins.deepspeed_utils import infer_deepspeed_mixed_precision
logger = get_logger(__name__)
# ZeRO-3 bucket sizes that accelerate derives from the model's hidden size
_ZERO3_BUCKET_FORMULAS = {
"reduce_bucket_size": lambda hidden: hidden * hidden,
"stage3_prefetch_bucket_size": lambda hidden: int(0.9 * hidden * hidden),
"stage3_param_persistence_threshold": lambda hidden: 10 * hidden,
}
class DeepSpeedEngine:
"""DeepSpeed integration using accelerate's built-in capabilities.
@@ -84,6 +92,7 @@ class DeepSpeedEngine:
Internally calls deepspeed.initialize() and wraps the returned objects.
"""
self._fill_zero3_bucket_sizes(model)
if lr_scheduler is not None:
model, optimizer, lr_scheduler = self.accelerator.prepare(model, optimizer, lr_scheduler)
else:
@@ -94,6 +103,23 @@ class DeepSpeedEngine:
logger.info_rank0("Model, optimizer, and lr_scheduler prepared via accelerate")
return model, optimizer, lr_scheduler
def _fill_zero3_bucket_sizes(self, model: HFModel) -> None:
"""Fill ZeRO-3 ``auto`` bucket sizes that accelerate cannot infer for multimodal models."""
zero_config = self.accelerator.state.deepspeed_plugin.deepspeed_config.get("zero_optimization", {})
auto_keys = [key for key in _ZERO3_BUCKET_FORMULAS if zero_config.get(key) == "auto"]
if not auto_keys:
return
config = model.config
text_config = config.get_text_config() if hasattr(config, "get_text_config") else config
hidden_size = getattr(text_config, "hidden_size", None)
if hidden_size is None:
return
for key in auto_keys:
zero_config[key] = _ZERO3_BUCKET_FORMULAS[key](hidden_size)
logger.info_rank0(f"Resolved ZeRO-3 {auto_keys} from text-config hidden_size={hidden_size}.")
def backward(self, loss: torch.Tensor) -> None:
"""Backward pass using accelerate.
@@ -108,7 +134,7 @@ class DeepSpeedEngine:
"""Get the global gradient norm from the DeepSpeed engine."""
engine_wrapper = getattr(self.accelerator, "deepspeed_engine_wrapped", None)
if engine_wrapper is not None:
return engine_wrapper.engine.get_global_grad_norm() or 0.0
return float(engine_wrapper.engine.get_global_grad_norm() or 0.0)
return 0.0

View File

@@ -1,4 +1,4 @@
# Copyright 2025 the LlamaFactory team.
# Copyright 2026 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.
@@ -73,20 +73,50 @@ def _make_safetensor_loader(checkpoint_file: str, tensor_key: str):
return _load_tensor
def get_transformer_layer_cls(model: HFModel) -> type[nn.Module] | None:
def _cast_norm_input_to_weight_dtype(module: nn.Module, args: tuple):
"""forward-pre-hook: cast a norm layer's input to its weight dtype."""
if not args:
return None
x = args[0]
weight = getattr(module, "weight", None)
if isinstance(x, torch.Tensor) and weight is not None and x.dtype != weight.dtype:
return (x.to(weight.dtype), *args[1:])
return None
def _make_norms_dtype_safe(model: HFModel) -> int:
"""Register the dtype-safe hook on every dtype-strict ``nn.LayerNorm`` in the model."""
n = 0
for module in model.modules():
if isinstance(module, nn.LayerNorm):
module.register_forward_pre_hook(_cast_norm_input_to_weight_dtype)
n += 1
return n
def get_transformer_layer_cls(model: HFModel) -> set[type[nn.Module]]:
classes: set[type[nn.Module]] = set()
for module in model.modules():
for attr in ("layers", "blocks"):
seq = getattr(module, attr, None)
if isinstance(seq, nn.ModuleList) and len(seq) > 0:
classes.add(type(seq[0]))
if classes:
return classes
no_split_modules = getattr(model, "_no_split_modules", None)
if no_split_modules:
if isinstance(no_split_modules, (list, tuple)):
for name, module in model.named_modules():
for cls_name in no_split_modules:
if module.__class__.__name__ == cls_name:
return module.__class__
if hasattr(model, "model") and hasattr(model.model, "layers"):
return type(model.model.layers[0])
if hasattr(model, "layers"):
return type(model.layers[0])
found: dict[str, type[nn.Module]] = {}
for _, module in model.named_modules():
cls_name = module.__class__.__name__
if cls_name in no_split_modules and cls_name not in found:
found[cls_name] = module.__class__
if len(found) == len(no_split_modules):
break
if found:
return set(found.values())
return None
return set()
def save_model(model: HFModel, output_dir: str, processor: Processor) -> None:
@@ -190,22 +220,26 @@ class FSDP2Engine:
def is_lora_module_wrap(self, model) -> bool:
return any(isinstance(module, LoraLayer) for module in model.modules())
def prepare_model(self, model: HFModel) -> HFModel:
def prepare_model(self, model: HFModel, ignored_params: set[nn.Parameter] | None = None) -> HFModel:
if self.fsdp_mesh is None:
logger.warning("No FSDP Mesh available, skipping FSDP wrapping.")
return model
mp_policy = self.get_mp_policy()
layer_cls = get_transformer_layer_cls(model)
transformer_layer_cls_to_wrap = get_transformer_layer_cls(model)
if layer_cls is None:
if not transformer_layer_cls_to_wrap:
logger.warning(
"Could not identify Transformer Layer class, applying FSDP to the whole model structure only."
)
transformer_layer_cls_to_wrap = set()
else:
logger.info(f"Applying per-layer FSDP to {layer_cls.__name__}")
transformer_layer_cls_to_wrap = {layer_cls}
names = ", ".join(cls.__name__ for cls in transformer_layer_cls_to_wrap)
logger.info(f"Applying per-layer FSDP to: {names}")
def _ignored_params_for(module: nn.Module) -> set[nn.Parameter] | None:
if not ignored_params:
return None
return ignored_params.intersection(module.parameters()) or None
if self.is_lora_module_wrap(model):
lora_modules = []
@@ -222,6 +256,7 @@ class FSDP2Engine:
reshard_after_forward=self.reshard_after_forward,
mp_policy=mp_policy,
offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None,
ignored_params=_ignored_params_for(module),
)
logger.info("Applying FSDP wrap for LoRA layer separately.")
@@ -242,6 +277,7 @@ class FSDP2Engine:
reshard_after_forward=self.reshard_after_forward,
mp_policy=mp_policy,
offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None,
ignored_params=_ignored_params_for(module),
)
# BaseTrainer is the single source of truth for gradient checkpointing.
@@ -259,12 +295,18 @@ class FSDP2Engine:
model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
if self.mixed_precision == "bf16":
n_patched = _make_norms_dtype_safe(model)
if self.rank == 0 and n_patched:
logger.info(f"Made {n_patched} nn.LayerNorm(s) dtype-safe for bf16 checkpointing.")
fully_shard(
model,
mesh=self.fsdp_mesh,
reshard_after_forward=self.reshard_after_forward,
mp_policy=mp_policy,
offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None,
ignored_params=_ignored_params_for(model),
)
return model

View File

@@ -0,0 +1,444 @@
# 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 collections.abc import Callable
import torch
from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
from ....accelerator.interface import Dim, DistributedInterface
from ....utils.logging import get_logger
from ....utils.types import HFModel
from .fsdp2 import FSDP2Engine
logger = get_logger(__name__)
class FSDPTurboParallelState:
"""Own FSDPTurbo's expert topology without extending LlamaFactory's global interface."""
EDP = "edp"
EFSDP = "efsdp"
EP = "ep"
EXPERT_CP = "expert_cp"
def __init__(self) -> None:
self._initialized = False
self.dp_size = 1
self.cp_size = 1
self.ep_size = 1
self.efsdp_size = 1
self.edp_size = 1
self.expert_mesh: DeviceMesh | None = None
self.edp_mesh: DeviceMesh | None = None
self.efsdp_mesh: DeviceMesh | None = None
self.ep_mesh: DeviceMesh | None = None
self.expert_cp_mesh: DeviceMesh | None = None
@property
def initialized(self) -> bool:
return self._initialized
def initialize(self, dist_interface: DistributedInterface, dist_config: dict) -> None:
dp_size = dist_interface.get_world_size(Dim.DP)
cp_size = dist_interface.strategy.cp_size
ep_size = int(dist_config.get("ep_size", 1))
if ep_size < 1:
raise ValueError(f"ep_size must be positive, got {ep_size}.")
if dp_size % ep_size != 0:
raise ValueError(f"dp_size must be divisible by ep_size, got {dp_size} % {ep_size} != 0.")
topology = (dp_size, cp_size, ep_size)
if self._initialized:
current_topology = (self.dp_size, self.cp_size, self.ep_size)
if topology != current_topology:
raise RuntimeError(
f"FSDPTurbo parallel state is already initialized with {current_topology}, got {topology}."
)
return
self.dp_size = dp_size
self.cp_size = cp_size
self.ep_size = ep_size
if ep_size > 1:
self.efsdp_size = dp_size // ep_size
self.edp_size = dp_size // (ep_size * self.efsdp_size)
if dist_interface.get_device_mesh(Dim.DP) is None:
raise RuntimeError("FSDPTurbo expert parallelism requires an initialized distributed device mesh.")
self.expert_mesh = init_device_mesh(
device_type=dist_interface.current_device.type,
mesh_shape=(self.edp_size, self.efsdp_size, self.ep_size, self.cp_size),
mesh_dim_names=(self.EDP, self.EFSDP, self.EP, self.EXPERT_CP),
)
self.edp_mesh = self.expert_mesh[self.EDP]
self.efsdp_mesh = self.expert_mesh[self.EFSDP]
self.ep_mesh = self.expert_mesh[self.EP]
self.expert_cp_mesh = self.expert_mesh[self.EXPERT_CP]
self._initialized = True
_FSDPTURBO_PARALLEL_STATE = FSDPTurboParallelState()
def get_fsdpturbo_parallel_state() -> FSDPTurboParallelState:
return _FSDPTURBO_PARALLEL_STATE
def _grad_to_local_fp32(grad: torch.Tensor) -> torch.Tensor:
from torch.distributed._tensor import DTensor
local_grad = grad.to_local() if isinstance(grad, DTensor) else grad
return local_grad.detach().to(torch.float32)
def _local_pth_sum(parameters: list[torch.nn.Parameter], norm_type: float, device: torch.device) -> torch.Tensor:
total = torch.zeros((), device=device, dtype=torch.float32)
for param in parameters:
grad = getattr(param, "grad", None)
if grad is None:
continue
total = total + torch.norm(_grad_to_local_fp32(grad), p=norm_type).pow(norm_type)
return total
def _allreduce_sum_(value: torch.Tensor, groups: list[object]) -> torch.Tensor:
import torch.distributed as dist
for group in groups:
if group is not None:
dist.all_reduce(value, op=dist.ReduceOp.SUM, group=group)
return value
def clip_grad_norm_(model: HFModel, max_norm: float, **kwargs) -> float:
"""CP-aware grad norm clipping for FSDPTurbo EP + EFSDP + outer FSDP2.
Avoids torch.nn.utils.get_total_norm() since mixed DTensor meshes
(`dp` vs `efsdp`/`ep`) may hit DTensor stack propagation failures.
"""
from torch.distributed._tensor import DTensor
norm_type = float(kwargs.get("norm_type", 2.0))
dist_interface = DistributedInterface()
parallel_state = get_fsdpturbo_parallel_state()
if not parallel_state.initialized:
raise RuntimeError("FSDPTurbo parallel state must be initialized before clipping gradients.")
device = dist_interface.current_device
dp_group = dist_interface.get_group(Dim.DP)
cp_group = dist_interface.get_group(Dim.CP) if dist_interface.strategy.cp_size > 1 else None
ep_group = parallel_state.ep_mesh.get_group() if parallel_state.ep_mesh is not None else None
efsdp_group = parallel_state.efsdp_mesh.get_group() if parallel_state.efsdp_mesh is not None else None
expert_cp_group = (
parallel_state.expert_cp_mesh.get_group()
if parallel_state.expert_cp_mesh is not None and parallel_state.cp_size > 1
else None
)
ep_params: list[torch.nn.Parameter] = []
non_ep_params: list[torch.nn.Parameter] = []
for param in model.parameters():
grad = getattr(param, "grad", None)
if grad is None:
continue
mesh_names = set(getattr(getattr(grad, "device_mesh", None), "mesh_dim_names", ()) or ())
is_ep_side = isinstance(grad, DTensor) and bool(mesh_names & {parallel_state.EP, parallel_state.EFSDP})
if is_ep_side:
ep_params.append(param)
else:
non_ep_params.append(param)
if not ep_params and not non_ep_params:
return 0.0
total_pth = torch.zeros((), device=device, dtype=torch.float32)
if non_ep_params:
non_ep_pth = _local_pth_sum(non_ep_params, norm_type, device)
total_pth = total_pth + _allreduce_sum_(non_ep_pth, [dp_group, cp_group])
if ep_params:
ep_pth = _local_pth_sum(ep_params, norm_type, device)
total_pth = total_pth + _allreduce_sum_(ep_pth, [efsdp_group, ep_group, expert_cp_group])
total_norm = total_pth.pow(1.0 / norm_type)
clip_coef = min(max_norm / (float(total_norm.item()) + 1e-6), 1.0)
if clip_coef < 1.0:
for param in ep_params + non_ep_params:
grad = getattr(param, "grad", None)
if grad is not None:
grad.detach().mul_(clip_coef)
return float(total_norm.item())
def _get_model_type(model: HFModel) -> str | None:
return getattr(getattr(model, "config", None), "model_type", None)
class FSDPTurboEPModelSpec:
_registry: dict[str, "FSDPTurboEPModelSpec"] = {}
def __init__(
self,
ep_modules: list[str],
ep_fsdp_modules: list[str] | None = None,
prepare_fn: Callable[[HFModel], HFModel] | None = None,
) -> None:
self.ep_modules = ep_modules
self.ep_fsdp_modules = ep_fsdp_modules
self.prepare_fn = prepare_fn
@classmethod
def register(
cls,
model_type: str,
ep_modules: list[str],
ep_fsdp_modules: list[str] | None = None,
):
def decorator(fn):
cls._registry[model_type] = cls(
ep_modules=ep_modules,
ep_fsdp_modules=ep_fsdp_modules,
prepare_fn=fn,
)
return fn
return decorator
@classmethod
def get(cls, model: HFModel) -> "FSDPTurboEPModelSpec | None":
model_type = _get_model_type(model)
if model_type is None:
return None
return cls._registry.get(model_type)
def prepare(self, model: HFModel) -> HFModel:
if self.prepare_fn is None:
return model
return self.prepare_fn(model)
@FSDPTurboEPModelSpec.register(
"qwen3_moe",
ep_modules=["model.layers.{*}.mlp.experts"],
ep_fsdp_modules=["model.layers.{*}.mlp"],
)
def _prepare_qwen3_moe_for_ep(model: HFModel) -> HFModel:
prepared = 0
for module in model.modules():
if not all(hasattr(module, attr) for attr in ("gate_up_proj", "down_proj", "hidden_dim", "num_experts")):
continue
# FSDPTurbo's eager EP dispatcher expects sparse expert blocks to expose `hidden_size`.
if not hasattr(module, "hidden_size"):
module.hidden_size = module.hidden_dim
prepared += 1
if prepared:
logger.info_rank0(f"FSDPTurbo EP adapter: prepared {prepared} sparse expert modules for Transformers 5.x.")
else:
logger.info_rank0("FSDPTurbo EP adapter did not find a sparse expert module requiring preparation.")
return model
@FSDPTurboEPModelSpec.register(
"qwen3_5_moe",
ep_modules=["model.language_model.layers.{*}.mlp.experts"],
ep_fsdp_modules=["model.language_model.layers.{*}.mlp"],
)
def _prepare_qwen3_5_moe_for_ep(model: HFModel) -> HFModel:
return model
class FSDPTurboFSDP2Engine(FSDP2Engine):
"""FSDPTurbo EP adapter that reuses LlamaFactory's init/load flow.
Design:
- FSDPTurbo owns EP / EFSDP only.
- LlamaFactory owns FSDP / CP / init-load lifecycle.
"""
def __init__(self, dist_config: dict, bf16: bool = False):
self.dist_config = dist_config
super().__init__(dist_config, bf16=bf16)
self.parallel_state = get_fsdpturbo_parallel_state()
self.parallel_state.initialize(self.dist_interface, self.dist_config)
self.ep_size = self.parallel_state.ep_size
self.ep_fsdp_size = self.parallel_state.efsdp_size
dp_mesh = self.dist_interface.get_device_mesh(Dim.DP)
if dp_mesh is not None:
self.fsdp_mesh = dp_mesh
logger.info(f"Using DP-orthogonal FSDP mesh: {self.fsdp_mesh}")
@staticmethod
def _get_ep_fsdp_modules(spec: FSDPTurboEPModelSpec) -> list[str]:
if spec.ep_fsdp_modules is not None:
return spec.ep_fsdp_modules
ep_fsdp_modules = []
for module in spec.ep_modules:
if module.endswith(".experts"):
ep_fsdp_modules.append(module.removesuffix(".experts"))
else:
ep_fsdp_modules.append(module)
return ep_fsdp_modules
def shard_model(self, model: HFModel) -> HFModel:
"""Set storage dtype before FSDP materialization without leaking backend config into ModelEngine."""
param_dtype = torch.bfloat16 if self.mixed_precision == "bf16" else torch.float32
model = model.to(param_dtype)
logger.info_rank0(f"Using {param_dtype} for FSDPTurbo full tuning.")
return super().shard_model(model)
def _copy_weights(self, param, loaded_tensor):
"""Copy full checkpoint tensors into mixed-mesh DTensors from the inherited loader."""
from torch.distributed._tensor import DTensor, Shard
if loaded_tensor.dtype != param.dtype:
loaded_tensor = loaded_tensor.to(param.dtype)
if isinstance(param, DTensor):
local_tensor = param.to_local()
shard_placements = [
(i, placement) for i, placement in enumerate(param.placements) if isinstance(placement, Shard)
]
if not shard_placements:
local_tensor.copy_(loaded_tensor)
return
mesh = param.device_mesh
my_coordinate = mesh.get_coordinate()
if my_coordinate is None:
return
sliced_tensor = loaded_tensor
for mesh_dim, shard_placement in shard_placements:
dim = shard_placement.dim
rank_in_dim = my_coordinate[mesh_dim]
world_size_in_dim = mesh.size(mesh_dim)
full_size = sliced_tensor.shape[dim]
chunk_size = (full_size + world_size_in_dim - 1) // world_size_in_dim
start = rank_in_dim * chunk_size
end = min(start + chunk_size, full_size)
if start >= full_size:
return
sliced_tensor = sliced_tensor.narrow(dim, start, end - start)
slices = [slice(None)] * local_tensor.ndim
for _, shard_placement in shard_placements:
dim = shard_placement.dim
slices[dim] = slice(0, sliced_tensor.shape[dim])
local_tensor[tuple(slices)].copy_(sliced_tensor)
return
param.data.copy_(loaded_tensor)
def prepare_model_ep(self, model: HFModel) -> tuple[HFModel, set]:
"""Apply FSDPTurbo EP/EFSDP and return parameters excluded from outer FSDP."""
from fsdp_turbo.distributed.expert_parallel.expert_fully_shard_parallel import (
expert_fully_shard_modules,
)
from fsdp_turbo.distributed.expert_parallel.expert_parallel import expert_parallelize_modules
from fsdp_turbo.fsdp_turbo_config import EPPlanConfig, FSDPPlanConfig
from fsdp_turbo.utils.str_match import module_name_match
spec = FSDPTurboEPModelSpec.get(model)
if spec is None:
raise ValueError(f"No FSDPTurbo EP spec is registered for model_type={_get_model_type(model)}.")
ep_modules = spec.ep_modules
model = spec.prepare(model)
if self.ep_size > 1:
ep_plan = EPPlanConfig(
apply_modules=ep_modules,
dispatcher=self.dist_config.get("ep_dispatcher", "eager"),
apply_efsdp_modules=self._get_ep_fsdp_modules(spec),
)
ep_plan.gradient_divide_factor = float(self.ep_size * self.parallel_state.efsdp_size)
fsdp_plan = FSDPPlanConfig(
# FSDPTurbo uses this plan only to place EFSDP hooks and select its
# implementation. EFSDP targets come from ep_plan.apply_efsdp_modules.
apply_modules={},
hook_modules=self.dist_config.get("hook_modules", []),
fsdp_implementation=self.dist_config.get("fsdp_implementation", "native"),
)
ep_mesh = self.parallel_state.ep_mesh
efsdp_mesh = self.parallel_state.efsdp_mesh
if ep_mesh is None:
raise RuntimeError("FSDPTurbo EP mesh is not initialized.")
if self.ep_fsdp_size > 1 and efsdp_mesh is None:
raise RuntimeError("FSDPTurbo EFSDP mesh is not initialized.")
if self.rank == 0:
logger.info("Applying FSDPTurbo EP backend.")
logger.info(f"FSDPTurbo EP apply patterns: {ep_modules}")
logger.info(f"FSDPTurbo EP device mesh: {ep_mesh}")
logger.info(f"FSDPTurbo EP gradient divide factor: {ep_plan.gradient_divide_factor}")
model = expert_parallelize_modules(model, ep_mesh, ep_plan)
if self.ep_fsdp_size > 1:
if self.rank == 0:
logger.info(f"FSDPTurbo EFSDP apply patterns: {ep_plan.apply_efsdp_modules}")
logger.info(f"FSDPTurbo EFSDP device mesh: {efsdp_mesh}")
model = expert_fully_shard_modules(model, efsdp_mesh, ep_plan, fsdp_plan)
# Collect ignored params for the outer FSDP wrap
fsdp_ignored_modules = list(self.dist_config.get("fsdp_ignored_modules", []))
if self.ep_size > 1:
fsdp_ignored_modules.extend(ep_modules)
ignored_params = set()
if fsdp_ignored_modules:
for name, module in model.named_modules():
for pattern in fsdp_ignored_modules:
if module_name_match(pattern, name):
ignored_params.update(list(module.parameters(recurse=True)))
if ignored_params and self.rank == 0:
logger.info(f"FSDPTurbo FSDP2: Ignoring {len(ignored_params)} EP parameters in outer FSDP.")
return model, ignored_params
def prepare_model(self, model: HFModel) -> HFModel:
# Apply FSDPTurbo EP first, then shard the remaining parameters with LlamaFactory FSDP2.
model, ignored_params = self.prepare_model_ep(model)
return super().prepare_model(model, ignored_params=ignored_params)
def _warmup_grad_norm(self, model: HFModel) -> None:
"""Warm up collectives without stacking gradients from different DTensor meshes."""
if self.fsdp_mesh is None:
return
logger.info_rank0("Warming up FSDPTurbo mixed-mesh grad norm computation...")
for param in model.parameters():
if param.requires_grad:
param.grad = torch.zeros_like(param)
with torch.no_grad():
clip_grad_norm_(model, 1.0)
for param in model.parameters():
if param.requires_grad:
param.grad = None
logger.info_rank0("FSDPTurbo mixed-mesh grad norm warmup completed.")

View File

@@ -20,7 +20,7 @@ reads mesh topology from ``TrainingArguments`` and never puts it in backend para
from __future__ import annotations
from dataclasses import asdict, dataclass
from dataclasses import asdict, dataclass, field
from typing import TYPE_CHECKING, Literal
from ....utils.plugin import BasePlugin
@@ -41,6 +41,24 @@ class FSDP2Params:
dcp_path: str | None = None
@dataclass
class FSDPTurboParams:
name: Literal["fsdpturbo"] = "fsdpturbo"
reshard_after_forward: bool = True
offload_params: bool = False
pin_memory: bool = True
dcp_path: str | None = None
ep_size: int = 1
ep_dispatcher: str = "eager"
fsdp_ignored_modules: list[str] = field(default_factory=list)
hook_modules: list[str] = field(default_factory=list)
fsdp_implementation: str = "native"
def __post_init__(self) -> None:
if self.ep_size < 1:
raise ValueError(f"ep_size must be positive, got {self.ep_size}.")
@dataclass
class DeepSpeedParams:
name: Literal["deepspeed"] = "deepspeed"
@@ -83,6 +101,40 @@ class FSDP2Distributed(BaseDistributed):
load_checkpoint(model, optimizer, ckpt_dir, **kwargs)
@DistributedPlugin("fsdpturbo").register()
class FSDPTurboDistributed(BaseDistributed):
@staticmethod
def shard_model(model: HFModel, dist_config: PluginConfig | FSDPTurboParams, **kwargs) -> HFModel:
dist_config = DistributedPlugin.parse_params(dist_config, FSDPTurboParams)
from .fsdpturbo import FSDPTurboFSDP2Engine
return FSDPTurboFSDP2Engine(asdict(dist_config), bf16=bool(kwargs.get("bf16"))).shard_model(model)
@staticmethod
def clip_grad_norm(model: HFModel, max_norm: float, **kwargs) -> float:
from .fsdpturbo import clip_grad_norm_
return clip_grad_norm_(model, max_norm, **kwargs)
@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

View File

@@ -1,4 +1,4 @@
# Copyright 2025 the LlamaFactory team.
# Copyright 2026 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.
@@ -12,4 +12,10 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import os
IGNORE_INDEX = -100
IMAGE_PLACEHOLDER = os.getenv("IMAGE_PLACEHOLDER", "<image>")
VIDEO_PLACEHOLDER = os.getenv("VIDEO_PLACEHOLDER", "<video>")
AUDIO_PLACEHOLDER = os.getenv("AUDIO_PLACEHOLDER", "<audio>")

View File

@@ -1,4 +1,4 @@
# Copyright 2025 the LlamaFactory team.
# Copyright 2026 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.
@@ -23,7 +23,7 @@ from transformers import set_seed as hf_set_seed
from ..accelerator.helper import is_torch_npu_available
from ..accelerator.interface import DistributedInterface
from .constants import IGNORE_INDEX
from .types import BatchInput, ModelInput, Processor, Tensor
from .types import BatchInput, Processor
def enable_full_determinism(seed: int) -> None:
@@ -79,37 +79,6 @@ def get_tokenizer(processor: Processor) -> PreTrainedTokenizer:
return processor.tokenizer if hasattr(processor, "tokenizer") else processor
def _pad_and_truncate(tensor: Tensor, max_seqlen: int, pad_value: int = 0) -> Tensor:
if tensor.shape[-1] >= max_seqlen:
return tensor[..., :max_seqlen]
pad_shape = list(tensor.shape)
pad_shape[-1] = max_seqlen - tensor.shape[-1]
pad_tensor = torch.full(pad_shape, pad_value, dtype=tensor.dtype, device=tensor.device)
return torch.cat([tensor, pad_tensor], dim=-1)
def pad_and_truncate(samples: list[ModelInput], max_seqlen: int) -> list[BatchInput]:
max_length = min(max(len(sample["input_ids"]) for sample in samples), max_seqlen)
padded_samples = []
for sample in samples:
padded_sample = {}
for key, value in sample.items():
if "label" in key:
pad_value = IGNORE_INDEX
else:
pad_value = 0
if not isinstance(value, str):
padded_sample[key] = _pad_and_truncate(torch.tensor(value), max_length, pad_value)
else:
padded_sample[key] = value
padded_samples.append(padded_sample)
return padded_samples
def compute_valid_tokens(batches: list[BatchInput]) -> int:
"""Compute valid tokens in batches.
@@ -125,3 +94,15 @@ def compute_valid_tokens(batches: list[BatchInput]) -> int:
for batch in batches
if "labels" in batch
)
def model_uses_mrope(config) -> bool:
"""Whether the model uses multimodal RoPE (3D position ids built from grid_thw).
Detected from the (text) config's rope settings carrying an ``mrope_section`` (Qwen2.5-VL /
Qwen3-VL / Qwen3.5 family). Such models compute their own multimodal position ids inside
``forward`` when ``position_ids`` is not provided.
"""
text_config = getattr(config, "text_config", config)
rope = getattr(text_config, "rope_scaling", None) or getattr(text_config, "rope_parameters", None)
return isinstance(rope, dict) and "mrope_section" in rope

View File

@@ -1,4 +1,4 @@
# Copyright 2025 the LlamaFactory team.
# Copyright 2026 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.
@@ -141,6 +141,20 @@ class ModelInput(TypedDict, total=False):
"""Position ids for the model (optional)."""
token_type_ids: NotRequired[list[int]]
"""Token type ids used in DPO, 1 represents the chosen messages, 2 represents the rejected messages."""
pixel_values: NotRequired[Any]
"""Pixel values for vision models."""
image_grid_thw: NotRequired[Any]
"""Image grid (temporal, height, width) for vision models."""
pixel_values_videos: NotRequired[Any]
"""Pixel values for video inputs."""
video_grid_thw: NotRequired[Any]
"""Video grid (temporal, height, width) for video models."""
input_features: NotRequired[Any]
"""Audio input features (e.g. mel spectrogram) for audio models."""
feature_attention_mask: NotRequired[Any]
"""Attention mask over the audio input features."""
mm_token_type_ids: NotRequired[list[int]]
"""Multimodal token type ids: 0=text, 1=image, 2=video, 3=audio."""
class BatchInput(TypedDict, total=False):
@@ -156,6 +170,20 @@ class BatchInput(TypedDict, total=False):
"""Position ids for the model (optional)."""
token_type_ids: NotRequired[Tensor]
"""Token type ids used in DPO, 1 represents the chosen messages, 2 represents the rejected messages."""
pixel_values: NotRequired[Tensor]
"""Pixel values for vision models."""
image_grid_thw: NotRequired[Tensor]
"""Image grid (temporal, height, width) for vision models."""
pixel_values_videos: NotRequired[Tensor]
"""Pixel values for video inputs."""
video_grid_thw: NotRequired[Tensor]
"""Video grid (temporal, height, width) for video models."""
input_features: NotRequired[Tensor]
"""Audio input features (e.g. mel spectrogram) for audio models."""
feature_attention_mask: NotRequired[Tensor]
"""Attention mask over the audio input features."""
mm_token_type_ids: NotRequired[Tensor]
"""Multimodal token type ids: 0=text, 1=image, 2=video, 3=audio."""
class BatchInfo(TypedDict):

View File

@@ -13,6 +13,7 @@
# limitations under the License.
import os
from types import SimpleNamespace
from typing import TYPE_CHECKING, Any
import numpy as np
@@ -417,6 +418,24 @@ def test_qwen2_vl_plugin():
_check_plugin(**check_inputs)
def test_moss_vl_plugin():
messages = [
{"role": "user", "content": "First <image>, finally <image>."},
{"role": "assistant", "content": "Done."},
]
expected_messages = [
{"role": "user", "content": "First <|image_pad|>, finally <|image_pad|>."},
{"role": "assistant", "content": "Done."},
]
processor = SimpleNamespace(image_processor=object(), video_processor=object())
plugin = get_mm_plugin(name="moss_vl", image_token="<|image_pad|>", video_token="<|video_pad|>")
processed_messages = plugin.process_messages(messages, [object(), object()], [], [], processor)
assert processed_messages == expected_messages
assert messages[0]["content"] == "First <image>, finally <image>."
@pytest.mark.runs_on(["cpu", "mps"])
@pytest.mark.skipif(not is_transformers_version_greater_than("4.57.0"), reason="Requires transformers>=4.57.0")
def test_qwen3_vl_plugin():

View File

@@ -0,0 +1,631 @@
# 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 types import SimpleNamespace
import pytest
import torch
from PIL import Image
from llamafactory.data.collator import MultiModalDataCollatorForSeq2Seq
from llamafactory.data.mm_plugin import get_mm_plugin
from llamafactory.data.processor.supervised import SupervisedDatasetProcessor
from llamafactory.extras.constants import IGNORE_INDEX
IMAGE_TOKEN_ID = 101
VIDEO_TOKEN_ID = 102
VISION_START_TOKEN_ID = 103
VISION_END_TOKEN_ID = 104
TIME_START_TOKEN_ID = 105
TIME_END_TOKEN_ID = 106
IM_END_TOKEN_ID = 107
class _ImageProcessor:
def __init__(self):
self.calls = []
def __call__(self, images, return_tensors, **kwargs):
self.calls.append({"return_tensors": return_tensors, **kwargs})
values = []
for image in images:
marker = image.getpixel((0, 0))[0] + 1
values.append(torch.full((1, 3), marker, dtype=torch.float32))
return {
"pixel_values": torch.cat(values),
"image_grid_thw": torch.tensor([[1, 1, 1]] * len(images)),
}
class _VideoProcessor:
temporal_patch_size = 1
def __init__(self):
self.calls = []
def __call__(self, videos, return_tensors, return_metadata, **kwargs):
self.calls.append(
{
"return_tensors": return_tensors,
"return_metadata": return_metadata,
**kwargs,
}
)
result = {
"pixel_values_videos": torch.cat(
[torch.full((2, 3), 9 + index, dtype=torch.float32) for index in range(len(videos))]
),
"video_grid_thw": torch.tensor([[2, 1, 1]] * len(videos)),
}
if return_metadata:
result["video_metadata"] = [
SimpleNamespace(frames_indices=[0, 2], total_num_frames=2, fps=2.0, duration=2.0) for _ in videos
]
return result
class _Tokenizer:
pad_token_id = 0
padding_side = "right"
_token_ids = {
"<|time_start|>": TIME_START_TOKEN_ID,
"<|time_end|>": TIME_END_TOKEN_ID,
"<|im_end|>": IM_END_TOKEN_ID,
}
def convert_tokens_to_ids(self, token):
return self._token_ids[token]
def pad(self, features, padding, max_length, pad_to_multiple_of, return_tensors):
del padding, max_length, return_tensors
sequence_length = max(len(feature["input_ids"]) for feature in features)
if pad_to_multiple_of is not None:
sequence_length = ((sequence_length + pad_to_multiple_of - 1) // pad_to_multiple_of) * pad_to_multiple_of
padded = {"input_ids": [], "attention_mask": []}
for feature in features:
pad_length = sequence_length - len(feature["input_ids"])
if self.padding_side == "right":
padded["input_ids"].append(feature["input_ids"] + [self.pad_token_id] * pad_length)
padded["attention_mask"].append(feature["attention_mask"] + [0] * pad_length)
else:
padded["input_ids"].append([self.pad_token_id] * pad_length + feature["input_ids"])
padded["attention_mask"].append([0] * pad_length + feature["attention_mask"])
return {key: torch.tensor(value) for key, value in padded.items()}
class _Processor:
image_token_id = IMAGE_TOKEN_ID
video_token_id = VIDEO_TOKEN_ID
vision_start_token_id = VISION_START_TOKEN_ID
vision_end_token_id = VISION_END_TOKEN_ID
def __init__(self):
self.image_processor = _ImageProcessor()
self.video_processor = _VideoProcessor()
self.tokenizer = _Tokenizer()
@staticmethod
def _calculate_timestamps(*args, **kwargs):
del args, kwargs
return [0.0, 1.0]
def _get_plugin():
return get_mm_plugin(
name="moss_vl",
image_token="<|image_pad|>",
video_token="<|video_pad|>",
vision_bos_token="<|vision_start|>",
vision_eos_token="<|vision_end|>",
time_bos_token="<|time_start|>",
time_eos_token="<|time_end|>",
)
def _video_ids(seed):
return [
VISION_START_TOKEN_ID,
TIME_START_TOKEN_ID,
seed,
TIME_END_TOKEN_ID,
IMAGE_TOKEN_ID,
TIME_START_TOKEN_ID,
seed + 1,
TIME_END_TOKEN_ID,
IMAGE_TOKEN_ID,
VISION_END_TOKEN_ID,
]
def _left_pad(sequences, pad_value):
max_len = max(map(len, sequences))
return torch.tensor([[pad_value] * (max_len - len(sequence)) + sequence for sequence in sequences])
def test_moss_vl_process_messages_expands_video_frames():
plugin = _get_plugin()
processor = _Processor()
messages = [
{"role": "user", "content": "First <image>, then <video>, finally <image>."},
{"role": "assistant", "content": "Done."},
]
images = [Image.new("RGB", (2, 2)), Image.new("RGB", (2, 2), (2, 0, 0))]
processed = plugin.process_messages(messages, images, ["video.mp4"], [], processor)
video_tokens = (
"<|vision_start|>"
"<|time_start|>0.0 seconds<|time_end|><|image_pad|>"
"<|time_start|>1.0 seconds<|time_end|><|image_pad|>"
"<|vision_end|>"
)
assert processed[0]["content"] == (f"First <|image_pad|>, then {video_tokens}, finally <|image_pad|>.")
assert messages[0]["content"] == "First <image>, then <video>, finally <image>."
@pytest.mark.parametrize(
("content", "images", "videos", "error"),
[
("Missing media: <image>.", [], [], "number of images does not match"),
("Missing media: <video>.", [], [], "number of videos does not match"),
],
)
def test_moss_vl_rejects_placeholder_count_mismatch(content, images, videos, error):
plugin = _get_plugin()
with pytest.raises(ValueError, match=error):
plugin.process_messages(
[{"role": "user", "content": content}],
images,
videos,
[],
_Processor(),
)
def test_moss_vl_process_messages_expands_multiple_videos_in_order():
plugin = _get_plugin()
messages = [{"role": "user", "content": "Compare <video> with <video>."}]
processed = plugin.process_messages(messages, [], ["first.mp4", "second.mp4"], [], _Processor())
frame_tokens = (
"<|vision_start|>"
"<|time_start|>0.0 seconds<|time_end|><|image_pad|>"
"<|time_start|>1.0 seconds<|time_end|><|image_pad|>"
"<|vision_end|>"
)
assert processed[0]["content"] == f"Compare {frame_tokens} with {frame_tokens}."
def test_moss_vl_forwards_spatial_pixel_limits_to_native_processors():
plugin = _get_plugin()
processor = _Processor()
processor.image_min_pixels = 1024
processor.image_max_pixels = 262144
processor.video_min_pixels = 256
processor.video_max_pixels = 16384
processor.video_fps = 1.0
processor.video_maxlen = 8
image = Image.new("RGB", (1024, 1024))
plugin.process_messages(
[{"role": "user", "content": "Compare <image> and <video>."}],
[image],
["video.mp4"],
[],
processor,
)
plugin.get_mm_inputs(
[image],
["video.mp4"],
[],
[1],
[1],
[0],
[[IMAGE_TOKEN_ID, *_video_ids(201)]],
processor,
)
assert processor.image_processor.calls == [
{
"return_tensors": "pt",
"min_pixels": 1024,
"max_pixels": 262144,
}
]
assert processor.video_processor.calls == [
{
"return_tensors": "pt",
"return_metadata": True,
"video_fps": 1.0,
"max_frames": 8,
"size": {"shortest_edge": 256, "longest_edge": 16384},
},
{
"return_tensors": "pt",
"return_metadata": False,
"video_fps": 1.0,
"max_frames": 8,
"size": {"shortest_edge": 256, "longest_edge": 16384},
},
]
def test_moss_vl_rejects_invalid_batch_metadata():
plugin = _get_plugin()
processor = _Processor()
image = Image.new("RGB", (2, 2))
with pytest.raises(ValueError, match="batch metadata must have one entry per sample"):
plugin.get_mm_inputs([image], [], [], [1], [], [0], [[IMAGE_TOKEN_ID]], processor)
with pytest.raises(ValueError, match="media lengths do not consume all provided inputs"):
plugin.get_mm_inputs([image], [], [], [0], [0], [0], [[201]], processor)
def test_moss_vl_rejects_truncated_media_tokens():
plugin = _get_plugin()
with pytest.raises(ValueError, match="increase `cutoff_len`"):
plugin.get_mm_inputs(
[Image.new("RGB", (2, 2))],
[],
[],
[1],
[0],
[0],
[[201, 202]],
_Processor(),
)
def test_moss_vl_rejects_incomplete_video_token_block():
plugin = _get_plugin()
truncated_video_ids = _video_ids(201)[:-1]
with pytest.raises(ValueError, match="incomplete video token block"):
plugin.get_mm_inputs(
[],
["video.mp4"],
[],
[0],
[1],
[0],
[truncated_video_ids],
_Processor(),
)
def test_moss_vl_rejects_video_frame_token_count_mismatch():
plugin = _get_plugin()
incomplete_frame_ids = [VISION_START_TOKEN_ID, IMAGE_TOKEN_ID, VISION_END_TOKEN_ID]
with pytest.raises(ValueError, match="video frame tokens do not match"):
plugin.get_mm_inputs(
[],
["video.mp4"],
[],
[0],
[1],
[0],
[incomplete_frame_ids],
_Processor(),
)
def test_moss_vl_media_order_batch_mask_and_labels():
plugin = _get_plugin()
processor = _Processor()
images = [Image.new("RGB", (2, 2)), Image.new("RGB", (2, 2), (2, 0, 0))]
first_ids = [
IMAGE_TOKEN_ID,
201,
VISION_START_TOKEN_ID,
TIME_START_TOKEN_ID,
202,
TIME_END_TOKEN_ID,
IMAGE_TOKEN_ID,
TIME_START_TOKEN_ID,
203,
TIME_END_TOKEN_ID,
IMAGE_TOKEN_ID,
VISION_END_TOKEN_ID,
IMAGE_TOKEN_ID,
204,
]
second_ids = [301, 302]
assert plugin._get_media_order_from_ids(first_ids, processor, 2, 1) == ["image", "video", "image"]
mm_inputs = plugin.get_mm_inputs(
images=images,
videos=["video.mp4"],
audios=[],
imglens=[2, 0],
vidlens=[1, 0],
audlens=[0, 0],
batch_ids=[first_ids, second_ids],
processor=processor,
)
assert mm_inputs["grid_thw"].tolist() == [[1, 1, 1], [2, 1, 1], [1, 1, 1], [1, 1, 1]]
assert mm_inputs["media_nums_per_sample"] == [3, 1]
assert mm_inputs["pixel_values"][:, 0].tolist() == [1.0, 9.0, 9.0, 3.0, 256.0]
pre_padding_mask = mm_inputs["cross_attention_mask"]
assert pre_padding_mask.shape == (2, 1, len(first_ids), 4)
assert pre_padding_mask[0, 0, 0].tolist() == [False, True, True, True]
assert pre_padding_mask[0, 0, 10].tolist() == [False, False, False, True]
assert pre_padding_mask[0, 0, 12].tolist() == [False, False, False, False]
assert pre_padding_mask[1].all()
seq_len = len(first_ids)
input_ids = torch.tensor([first_ids, [0] * (seq_len - 2) + second_ids])
attention_mask = torch.tensor([[1] * seq_len, [0] * (seq_len - 2) + [1, 1]])
labels = input_ids.clone()
labels[attention_mask == 0] = IGNORE_INDEX
features = {
"input_ids": input_ids,
"attention_mask": attention_mask,
"labels": labels,
"position_ids": torch.arange(seq_len).repeat(2, 1),
}
mm_inputs = plugin.post_process_mossvl_inputs(features, mm_inputs, processor)
mask = mm_inputs["cross_attention_mask"]
assert mask.shape == (2, 1, seq_len, 4)
assert mask[0, 0, 0].tolist() == [False, True, True, True]
assert mask[0, 0, 10].tolist() == [False, False, False, True]
assert mask[0, 0, 12].tolist() == [False, False, False, False]
assert mask[1].all()
assert "position_ids" not in features
assert not torch.any((features["input_ids"] == IMAGE_TOKEN_ID) & ~features["attention_mask"].bool())
assert features["labels"][0, 1].item() == 201
assert features["labels"][0, 13].item() == 204
assert features["labels"][0, 0].item() == IGNORE_INDEX
assert features["labels"][0, 12].item() == IGNORE_INDEX
assert torch.all(features["labels"][0, 2:12] == IGNORE_INDEX)
def test_moss_vl_complex_batch_keeps_media_and_masks_sample_local():
plugin = _get_plugin()
processor = _Processor()
batch_ids = [
[IMAGE_TOKEN_ID, 211, IMAGE_TOKEN_ID, 212],
[221, *_video_ids(222), 223, *_video_ids(224), 225],
[IMAGE_TOKEN_ID, 231, *_video_ids(232), 233, IMAGE_TOKEN_ID, 234],
[241, 242, 243],
]
images = [Image.new("RGB", (2, 2), (marker, 0, 0)) for marker in range(4)]
mm_inputs = plugin.get_mm_inputs(
images=images,
videos=["first.mp4", "second.mp4", "third.mp4"],
audios=[],
imglens=[2, 0, 2, 0],
vidlens=[0, 2, 1, 0],
audlens=[0, 0, 0, 0],
batch_ids=batch_ids,
processor=processor,
)
assert mm_inputs["grid_thw"].tolist() == [
[1, 1, 1],
[1, 1, 1],
[2, 1, 1],
[2, 1, 1],
[1, 1, 1],
[2, 1, 1],
[1, 1, 1],
[1, 1, 1],
]
assert mm_inputs["media_nums_per_sample"] == [2, 2, 3, 1]
assert mm_inputs["pixel_values"][:, 0].tolist() == [
1.0,
2.0,
9.0,
9.0,
10.0,
10.0,
3.0,
9.0,
9.0,
4.0,
256.0,
]
input_ids = _left_pad(batch_ids, 0)
attention_mask = _left_pad([[1] * len(ids) for ids in batch_ids], 0)
labels = input_ids.clone()
labels[attention_mask == 0] = IGNORE_INDEX
features = {
"input_ids": input_ids,
"attention_mask": attention_mask,
"labels": labels,
"position_ids": torch.arange(input_ids.shape[1]).repeat(len(batch_ids), 1),
}
plugin.post_process_mossvl_inputs(features, mm_inputs, processor)
cross_mask = mm_inputs["cross_attention_mask"]
assert cross_mask.shape == (4, 1, input_ids.shape[1], 4)
assert (~cross_mask[0]).sum().item() > 0
assert (~cross_mask[1]).sum().item() > 0
assert (~cross_mask[2]).sum().item() > 0
assert cross_mask[3].all()
assert cross_mask[0, ..., 2:].all()
assert not cross_mask[1, ..., :4].all()
assert not cross_mask[2, ..., :4].all()
assert features["labels"][3, -3:].tolist() == [241, 242, 243]
assert torch.all(features["labels"][features["attention_mask"] == 0] == IGNORE_INDEX)
assert "position_ids" not in features
def test_moss_vl_supervised_processor_to_collator_mixed_batch(monkeypatch):
plugin = _get_plugin()
processor = _Processor()
tokenizer = processor.tokenizer
template = SimpleNamespace(mm_plugin=plugin)
dataset_processor = SupervisedDatasetProcessor(
template=template,
tokenizer=tokenizer,
processor=processor,
data_args=SimpleNamespace(),
)
first_ids = [IMAGE_TOKEN_ID, 211, *_video_ids(212), IMAGE_TOKEN_ID, 214]
second_ids = [221, 222]
def encode_example(prompt, **kwargs):
del kwargs
input_ids = first_ids if "<image>" in prompt[0]["content"] else second_ids
return input_ids, input_ids.copy()
monkeypatch.setattr(dataset_processor, "_encode_data_example", encode_example)
examples = {
"_prompt": [
[{"role": "user", "content": "Compare <image>, <video>, and <image>."}],
[{"role": "user", "content": "Text-only question."}],
],
"_response": [
[{"role": "assistant", "content": "Mixed answer."}],
[{"role": "assistant", "content": "Text answer."}],
],
"_system": ["", ""],
"_tools": ["", ""],
"_images": [
[Image.new("RGB", (2, 2)), Image.new("RGB", (2, 2), (2, 0, 0))],
None,
],
"_videos": [["video.mp4"], None],
"_audios": [None, None],
}
model_inputs = dataset_processor.preprocess_dataset(examples)
assert "media_order" not in model_inputs
collator = MultiModalDataCollatorForSeq2Seq(
tokenizer=tokenizer,
model=SimpleNamespace(config=SimpleNamespace(model_type="moss_vl")),
template=template,
processor=processor,
label_pad_token_id=IGNORE_INDEX,
)
features = [
{key: values[index] for key, values in model_inputs.items()} for index in range(len(model_inputs["input_ids"]))
]
batch = collator(features)
assert batch["grid_thw"].tolist() == [[1, 1, 1], [2, 1, 1], [1, 1, 1], [1, 1, 1]]
assert batch["media_nums_per_sample"] == [3, 1]
assert batch["pixel_values"][:, 0].tolist() == [1.0, 9.0, 9.0, 3.0, 256.0]
assert batch["cross_attention_mask"].shape == (2, 1, len(first_ids), 4)
assert batch["cross_attention_mask"][1].all()
assert torch.all(batch["labels"][1, len(second_ids) :] == IGNORE_INDEX)
assert "position_ids" not in batch
def test_moss_vl_generate_collator_keeps_left_padded_cross_attention_mask():
plugin = _get_plugin()
processor = _Processor()
processor.tokenizer.padding_side = "left"
template = SimpleNamespace(mm_plugin=plugin)
batch_ids = [
[IMAGE_TOKEN_ID, 211],
[301, IMAGE_TOKEN_ID, 302, 303],
]
features = [
{
"input_ids": input_ids,
"attention_mask": [1] * len(input_ids),
"labels": input_ids.copy(),
"images": [Image.new("RGB", (2, 2))],
}
for input_ids in batch_ids
]
collator = MultiModalDataCollatorForSeq2Seq(
tokenizer=processor.tokenizer,
model=SimpleNamespace(config=SimpleNamespace(model_type="moss_vl")),
template=template,
processor=processor,
label_pad_token_id=IGNORE_INDEX,
pad_to_multiple_of=8,
)
batch = collator(features)
assert batch["cross_attention_mask"].shape == (2, 1, 8, 1)
assert batch["cross_attention_mask"][0, 0, :, 0].tolist() == [True] * 6 + [False, False]
assert batch["cross_attention_mask"][1, 0, :, 0].tolist() == [True] * 5 + [False, False, False]
def test_moss_vl_predict_collator_uses_precomputed_cross_attention_mask_without_model():
plugin = _get_plugin()
processor = _Processor()
template = SimpleNamespace(mm_plugin=plugin)
batch_ids = [
[IMAGE_TOKEN_ID, 211],
[301, IMAGE_TOKEN_ID, 302, 303],
]
features = [
{
"input_ids": input_ids,
"attention_mask": [1] * len(input_ids),
"labels": input_ids.copy(),
"images": [Image.new("RGB", (2, 2))],
}
for input_ids in batch_ids
]
collator = MultiModalDataCollatorForSeq2Seq(
tokenizer=processor.tokenizer,
model=None,
template=template,
processor=processor,
label_pad_token_id=IGNORE_INDEX,
)
batch = collator(features)
assert batch["cross_attention_mask"].shape == (2, 1, 4, 1)
assert batch["cross_attention_mask"][0, 0, :, 0].tolist() == [False, False, True, True]
assert batch["cross_attention_mask"][1, 0, :, 0].tolist() == [True, False, False, False]
def test_moss_vl_masks_only_the_token_after_im_end():
plugin = _get_plugin()
processor = _Processor()
input_ids = torch.tensor([[301, IM_END_TOKEN_ID, 302, 303]])
features = {
"input_ids": input_ids,
"attention_mask": torch.ones_like(input_ids),
"labels": input_ids.clone(),
}
mm_inputs = plugin.get_mm_inputs([], [], [], [0], [0], [0], [input_ids[0].tolist()], processor)
plugin.post_process_mossvl_inputs(features, mm_inputs, processor)
assert features["labels"].tolist() == [[301, IM_END_TOKEN_ID, IGNORE_INDEX, 303]]
def test_moss_vl_native_text_dummy_shape_and_values():
plugin = _get_plugin()
processor = _Processor()
processor.image_processor = SimpleNamespace(patch_size=16, temporal_patch_size=1, merge_size=2)
mm_inputs = plugin.get_mm_inputs([], [], [], [0], [0], [0], [[301]], processor)
assert mm_inputs["grid_thw"].tolist() == [[1, 8, 8]]
assert mm_inputs["pixel_values"].shape == (64, 768)
assert torch.count_nonzero(mm_inputs["pixel_values"]).item() == 0
assert mm_inputs["media_nums_per_sample"] == [1]

View File

@@ -0,0 +1,39 @@
# 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 pathlib import Path
import yaml
ROOT = Path(__file__).resolve().parents[2]
def test_moss_vl_training_configs_are_unpacked_and_additive():
lora = yaml.safe_load((ROOT / "examples/train_lora/mossvl_lora_sft.yaml").read_text())
full = yaml.safe_load((ROOT / "examples/train_full/mossvl_full_sft.yaml").read_text())
assert lora["template"] == full["template"] == "moss_vl"
assert lora["packing"] is full["packing"] is False
assert lora["per_device_train_batch_size"] == 2
assert full["per_device_train_batch_size"] == 1
assert full["use_reentrant_gc"] is False
assert full["gradient_checkpointing"] is True
assert full["gradient_checkpointing_kwargs"] == {"use_reentrant": False}
for config in (lora, full):
assert config["model_name_or_path"] == "OpenMOSS-Team/MOSS-VL-Instruct-0708"
assert not any("/inspire/" in str(value) or "/tmp/" in str(value) for value in config.values())
assert config["freeze_vision_tower"] is True
assert config["freeze_multi_modal_projector"] is True
assert config["freeze_language_model"] is False

View File

@@ -19,7 +19,13 @@ import pytest
from transformers import AutoTokenizer
from llamafactory.data import get_template_and_fix_tokenizer
from llamafactory.data.template import parse_template
from llamafactory.data.template import TEMPLATES, parse_template
from llamafactory.extras.constants import (
DEFAULT_TEMPLATE,
MULTIMODAL_SUPPORTED_MODELS,
SUPPORTED_MODELS,
DownloadSource,
)
from llamafactory.extras.packages import is_transformers_version_greater_than
from llamafactory.hparams import DataArguments
@@ -91,6 +97,22 @@ def _check_template(
_check_tokenization(tokenizer, (prompt_ids, answer_ids), (prompt_str, answer_str))
def test_moss_vl_registration():
model_name = "MOSS-VL-Instruct-0708"
assert model_name in SUPPORTED_MODELS
assert SUPPORTED_MODELS[model_name][DownloadSource.DEFAULT] == "OpenMOSS-Team/MOSS-VL-Instruct-0708"
assert DEFAULT_TEMPLATE[model_name] == "moss_vl"
assert model_name in MULTIMODAL_SUPPORTED_MODELS
assert TEMPLATES["moss_vl"].mm_plugin.__class__.__name__ == "MossVLPlugin"
assert TEMPLATES["moss_vl"].mm_plugin.image_token == "<|image_pad|>"
assert TEMPLATES["moss_vl"].mm_plugin.video_token == "<|video_pad|>"
assert TEMPLATES["moss_vl"].mm_plugin.vision_bos_token == "<|vision_start|>"
assert TEMPLATES["moss_vl"].mm_plugin.vision_eos_token == "<|vision_end|>"
assert TEMPLATES["moss_vl"].mm_plugin.time_bos_token == "<|time_start|>"
assert TEMPLATES["moss_vl"].mm_plugin.time_eos_token == "<|time_end|>"
@pytest.mark.runs_on(["cpu", "mps"])
def test_encode_oneturn():
tokenizer = AutoTokenizer.from_pretrained(TINY_LLAMA3)

View File

@@ -13,6 +13,7 @@
# limitations under the License.
import os
from types import SimpleNamespace
import pytest
import torch
@@ -21,7 +22,131 @@ from transformers import AutoConfig, AutoModelForImageTextToText
from llamafactory.extras.packages import is_transformers_version_greater_than
from llamafactory.hparams import FinetuningArguments, ModelArguments
from llamafactory.model.adapter import init_adapter
from llamafactory.model.adapter import _setup_freeze_tuning, _setup_full_tuning, init_adapter
from llamafactory.model.model_utils.misc import find_all_linear_modules
from llamafactory.model.model_utils.visual import COMPOSITE_MODELS, autocast_projector_dtype, patch_target_modules
class _MossVLFixture(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.config = SimpleNamespace(
model_type="moss_vl",
text_config=SimpleNamespace(num_hidden_layers=2),
)
self.model = torch.nn.Module()
self.model.separator_token = torch.nn.Parameter(torch.empty(4))
self.model.visual = torch.nn.Module()
self.model.visual.pos_embed = torch.nn.Embedding(4, 4)
self.model.visual.patch_embed = torch.nn.Module()
self.model.visual.patch_embed.proj = torch.nn.Linear(4, 4)
self.model.visual.blocks = torch.nn.ModuleList([self._make_block(), self._make_block()])
self.model.visual.merger = torch.nn.Module()
self.model.visual.merger.linear_fc1 = torch.nn.Linear(4, 4)
self.model.language_model = torch.nn.Module()
self.model.language_model.layers = torch.nn.ModuleList([self._make_layer(), self._make_layer()])
self.lm_head = torch.nn.Linear(4, 4)
@staticmethod
def _make_block() -> torch.nn.Module:
block = torch.nn.Module()
block.attn = torch.nn.Module()
block.attn.qkv = torch.nn.Linear(4, 4)
return block
@staticmethod
def _make_layer() -> torch.nn.Module:
layer = torch.nn.Module()
layer.self_attn = torch.nn.Module()
layer.self_attn.q_proj = torch.nn.Linear(4, 4)
return layer
@pytest.mark.parametrize("freeze_vision_tower", (False, True))
@pytest.mark.parametrize("freeze_multi_modal_projector", (False, True))
@pytest.mark.parametrize("freeze_language_model", (False, True))
def test_moss_vl_full(
freeze_vision_tower: bool,
freeze_multi_modal_projector: bool,
freeze_language_model: bool,
):
model = _MossVLFixture()
finetuning_args = FinetuningArguments(
finetuning_type="full",
freeze_vision_tower=freeze_vision_tower,
freeze_multi_modal_projector=freeze_multi_modal_projector,
freeze_language_model=freeze_language_model,
)
_setup_full_tuning(model, finetuning_args, is_trainable=True, cast_trainable_params_to_fp32=False)
for name, param in model.named_parameters():
if name.startswith("model.visual.merger") or name == "model.separator_token":
assert param.requires_grad != freeze_multi_modal_projector
elif name.startswith("model.visual"):
assert param.requires_grad != freeze_vision_tower
else:
assert param.requires_grad != freeze_language_model
@pytest.mark.parametrize("freeze_multi_modal_projector", (False, True))
def test_moss_vl_freeze(freeze_multi_modal_projector: bool):
model = _MossVLFixture()
finetuning_args = FinetuningArguments(
finetuning_type="freeze",
freeze_trainable_layers=1,
freeze_vision_tower=True,
freeze_multi_modal_projector=freeze_multi_modal_projector,
freeze_language_model=False,
)
_setup_freeze_tuning(model, finetuning_args, is_trainable=True, cast_trainable_params_to_fp32=False)
assert model.model.separator_token.requires_grad != freeze_multi_modal_projector
assert model.model.visual.merger.linear_fc1.weight.requires_grad != freeze_multi_modal_projector
assert model.model.visual.patch_embed.proj.weight.requires_grad is False
assert model.model.language_model.layers[0].self_attn.q_proj.weight.requires_grad is False
assert model.model.language_model.layers[1].self_attn.q_proj.weight.requires_grad is True
@pytest.mark.parametrize("freeze_vision_tower", (False, True))
def test_moss_vl_lora_target_all(freeze_vision_tower: bool):
model = _MossVLFixture()
finetuning_args = FinetuningArguments(
finetuning_type="lora",
lora_target="all",
freeze_vision_tower=freeze_vision_tower,
freeze_multi_modal_projector=True,
freeze_language_model=False,
)
target_modules = find_all_linear_modules(model, freeze_vision_tower)
target_modules = patch_target_modules(model, finetuning_args, target_modules)
assert any(name.startswith("model.language_model") and name.endswith("q_proj") for name in target_modules)
assert any(name.startswith("model.visual.blocks") and name.endswith("qkv") for name in target_modules) != (
freeze_vision_tower
)
assert all("patch_embed" not in name for name in target_modules)
assert all("merger" not in name for name in target_modules)
assert all("lm_head" not in name for name in target_modules)
def test_moss_vl_projector_modules():
model = _MossVLFixture()
composite_model = COMPOSITE_MODELS["moss_vl"]
assert composite_model.projector_keys == ["model.visual.merger", "model.separator_token"]
assert composite_model.get_projectors(model) == [model.model.visual.merger]
def test_moss_vl_quantized_projector_hook_skips_parameter():
model = _MossVLFixture()
model.quantization_method = "bitsandbytes"
autocast_projector_dtype(model, SimpleNamespace(compute_dtype=torch.float16))
assert len(model.model.visual.merger._forward_hooks) == 1
@pytest.mark.parametrize("freeze_vision_tower", (False, True))

View File

@@ -26,7 +26,9 @@ def test_get_args_from_yaml(tmp_path: Path):
trust_remote_code: true
model_class: llm
kernel_config:
name: auto
name: auto, flash-linear-attention
include_kernels: chunk_gated_delta_rule, fused_recurrent_gated_delta_rule
chunk_size: 32
peft_config:
name: lora
r: 8
@@ -58,7 +60,11 @@ def test_get_args_from_yaml(tmp_path: Path):
model_args, data_args, training_args, sample_args = get_args()
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.name == "auto, flash-linear-attention"
assert model_args.kernel_config.get("include_kernels") == (
"chunk_gated_delta_rule, fused_recurrent_gated_delta_rule"
)
assert model_args.kernel_config.get("chunk_size") == 32
assert model_args.peft_config.name == "lora"
assert model_args.peft_config.get("r") == 8
assert training_args.output_dir == "outputs/test_run"
@@ -68,3 +74,16 @@ def test_get_args_from_yaml(tmp_path: Path):
assert training_args.bf16 is False
assert training_args.dist_config is None
assert sample_args.sample_backend == "hf"
def test_qwen35_fsdpturbo_example_uses_v1_arguments():
config_file = (
Path(__file__).parents[2] / "examples" / "v1" / "train_full" / "train_full_qwen3_moe_fsdpturbo_ep_fsdp.yaml"
)
with patch.object(sys, "argv", ["test_args_parser.py", str(config_file)]):
model_args, _, training_args, _ = get_args()
assert model_args.model == "Qwen/Qwen3.5-35B-A3B"
assert model_args.custom_chat_template is None
assert training_args.dist_config.name == "fsdpturbo"

View File

@@ -370,3 +370,208 @@ def test_dynamic_padding_free_fill_buffer_restarts_until_micro_batch_is_complete
assert len(batch) == 1
assert batch[0]["input_ids"].shape == (1, 18)
assert len(batch_generator._buffer) == 1
def _image_fragment(n_pad: int = 4, merge_sq: int = 4):
"""Hand-crafted image fragment: vision_start + n_pad image_pad + vision_end."""
import torch
pad, vstart, vend = 9, 8, 7
return {
"input_ids": [vstart] + [pad] * n_pad + [vend],
"mm_token_type_ids": [0] + [1] * n_pad + [0],
"pixel_values": torch.zeros((n_pad * merge_sq, 16), dtype=torch.float32),
"image_grid_thw": torch.tensor([[1, 2, n_pad * 2]], dtype=torch.long),
}
def _text_sample(n: int, base: int = 100):
s = _make_model_input(n, start=base)
s["position_ids"] = list(range(1, n + 1))
return s
def test_inject_appends_zero_loss_dummy_into_collated_text_batch():
import torch
from llamafactory.v1.core.utils.batching import _collate_micro_batch, _inject_dummy_into_collated
collated = _collate_micro_batch([_text_sample(20), _text_sample(8)], cutoff_len=4096)
assert "pixel_values" not in collated
bsz, seqlen = collated["input_ids"].shape
frag = _image_fragment(n_pad=4)
fl = len(frag["input_ids"])
_inject_dummy_into_collated(collated, frag, marker=1)
new_len = seqlen + fl
# every sequence field grew by the fragment length, batch size unchanged
for key in ("input_ids", "attention_mask", "labels", "loss_weights", "position_ids", "mm_token_type_ids"):
assert collated[key].shape == (bsz, new_len)
# dummy lives only in row 0's tail; other rows are padding (attention 0) there
assert collated["input_ids"][0, seqlen:].tolist() == frag["input_ids"]
assert collated["attention_mask"][0, seqlen:].tolist() == [1] * fl
assert collated["attention_mask"][1, seqlen:].tolist() == [0] * fl
# zero loss contribution
assert collated["labels"][0, seqlen:].tolist() == [IGNORE_INDEX] * fl
assert torch.all(collated["loss_weights"][:, seqlen:] == 0.0)
assert collated["mm_token_type_ids"][0, seqlen:].tolist() == frag["mm_token_type_ids"]
# pixel features carried verbatim
assert torch.equal(collated["pixel_values"], frag["pixel_values"])
assert torch.equal(collated["image_grid_thw"], frag["image_grid_thw"])
def test_inject_video_concatenates_alongside_existing_image():
"""Injecting a missing modality leaves the other modality's features intact (dim-0 cat)."""
import torch
from llamafactory.v1.core.utils.batching import _collate_micro_batch, _inject_dummy_into_collated
img = _text_sample(10)
img["pixel_values"] = torch.ones((8, 16), dtype=torch.float32)
img["image_grid_thw"] = torch.tensor([[1, 2, 4]], dtype=torch.long)
img["mm_token_type_ids"] = [0] * 10
collated = _collate_micro_batch([img], cutoff_len=4096)
video_frag = {
"input_ids": [8, 6, 6, 7],
"mm_token_type_ids": [0, 2, 2, 0],
"pixel_values_videos": torch.zeros((8, 16), dtype=torch.float32),
"video_grid_thw": torch.tensor([[1, 2, 4]], dtype=torch.long),
}
_inject_dummy_into_collated(collated, video_frag, marker=2)
# image features untouched, video features added
assert torch.equal(collated["pixel_values"], torch.ones((8, 16)))
assert collated["pixel_values_videos"].shape[0] == 8
assert collated["video_grid_thw"].shape[0] == 1
assert collated["mm_token_type_ids"][0, -4:].tolist() == [0, 2, 2, 0]
def test_collate_creates_mm_token_type_ids_for_pure_text_then_inject():
"""A pure-text micro batch has no mm_token_type_ids; injection must create it."""
from llamafactory.v1.core.utils.batching import _collate_micro_batch, _inject_dummy_into_collated
collated = _collate_micro_batch([_text_sample(12)], cutoff_len=4096)
assert "mm_token_type_ids" not in collated
seqlen = collated["input_ids"].shape[1]
frag = _image_fragment(n_pad=3)
_inject_dummy_into_collated(collated, frag, marker=1)
assert "mm_token_type_ids" in collated
assert collated["mm_token_type_ids"].shape == collated["input_ids"].shape
# original region all zero (text), dummy region carries the markers
assert collated["mm_token_type_ids"][0, :seqlen].tolist() == [0] * seqlen
assert collated["mm_token_type_ids"][0, seqlen:].tolist() == frag["mm_token_type_ids"]
def _audio_fragment(n_tok: int = 2, n_frames: int = 3000):
"""Hand-crafted audio fragment: audio_bos + n_tok AUDIO + audio_eos, with feature rows."""
import torch
aud, bos, eos = 50, 51, 52
return {
"input_ids": [bos] + [aud] * n_tok + [eos],
"mm_token_type_ids": [0] + [3] * n_tok + [0],
"input_features": torch.zeros((1, 128, n_frames), dtype=torch.float32),
"feature_attention_mask": torch.ones((1, n_frames), dtype=torch.long),
}
def test_inject_audio_dummy_into_text_batch():
"""A pure-text micro batch gets an audio dummy appended so the audio tower fires on every rank."""
import torch
from llamafactory.v1.core.utils.batching import _collate_micro_batch, _inject_dummy_into_collated
collated = _collate_micro_batch([_text_sample(12)], cutoff_len=4096)
assert "input_features" not in collated
seqlen = collated["input_ids"].shape[1]
frag = _audio_fragment(n_tok=2)
fl = len(frag["input_ids"])
_inject_dummy_into_collated(collated, frag, marker=3)
# audio feature tensors carried verbatim; placeholder tokens marked 3 in the dummy tail
assert torch.equal(collated["input_features"], frag["input_features"])
assert torch.equal(collated["feature_attention_mask"], frag["feature_attention_mask"])
assert collated["mm_token_type_ids"][0, seqlen:].tolist() == frag["mm_token_type_ids"]
# zero loss contribution from the dummy
assert collated["labels"][0, seqlen:].tolist() == [IGNORE_INDEX] * fl
assert torch.all(collated["loss_weights"][:, seqlen:] == 0.0)
def test_audio_truncation_drops_orphaned_item_and_zeros_tokens():
"""Truncating mid-audio trims the orphaned feature row and zeros its in-window tokens."""
import torch
from llamafactory.v1.core.utils.collation import _align_multimodal_on_truncation
aud = 50
# text(2) + [audio#0: 4 tok] + text(1) + [audio#1: 4 tok] + text(1)
input_ids = [1, 2] + [aud] * 4 + [3] + [aud] * 4 + [4]
mm = [0, 0] + [3] * 4 + [0] + [3] * 4 + [0]
sample = {
"input_ids": input_ids,
"labels": input_ids.copy(),
"loss_weights": [1.0] * len(input_ids),
"mm_token_type_ids": mm,
"input_features": torch.zeros((2, 128, 10), dtype=torch.float32),
"feature_attention_mask": torch.ones((2, 10), dtype=torch.long),
}
# audio#1 occupies positions 7..10; cut at 9 so its last token (10) is orphaned, audio#0 intact
out = _align_multimodal_on_truncation(dict(sample), max_length=9)
assert out["input_features"].shape[0] == 1 # only the complete audio#0 survives
assert out["feature_attention_mask"].shape[0] == 1
# audio#0 tokens (positions 2..5) untouched
assert all(out["input_ids"][i] == aud and out["mm_token_type_ids"][i] == 3 for i in range(2, 6))
# audio#1's in-window tokens (positions 7,8) zeroed + delabeled (positions >= 9 cut by truncation)
for i in (7, 8):
assert out["input_ids"][i] == 0
assert out["mm_token_type_ids"][i] == 0
assert out["labels"][i] == IGNORE_INDEX
assert out["loss_weights"][i] == 0.0
def test_audio_truncation_keeps_all_when_complete():
"""No trimming when the cut falls after every audio's last token."""
import torch
from llamafactory.v1.core.utils.collation import _align_multimodal_on_truncation
aud = 50
input_ids = [1] + [aud] * 4 + [2]
sample = {
"input_ids": input_ids,
"labels": input_ids.copy(),
"loss_weights": [1.0] * len(input_ids),
"mm_token_type_ids": [0] + [3] * 4 + [0],
"input_features": torch.zeros((1, 128, 10), dtype=torch.float32),
"feature_attention_mask": torch.ones((1, 10), dtype=torch.long),
}
out = _align_multimodal_on_truncation(dict(sample), max_length=6)
assert out["input_features"].shape[0] == 1
assert out["input_ids"] == input_ids
def test_drop_unsupervised_samples():
"""Samples whose supervised tokens fall entirely beyond cutoff_len are dropped (warn once)."""
from types import SimpleNamespace
def _s(weights): # a sample's input_ids length matches its loss_weights length
return {"input_ids": list(range(len(weights))), "loss_weights": weights}
gen = SimpleNamespace(cutoff_len=4, _warned_truncation=False)
samples = [
_s([0.0, 0.0, 1.0, 1.0]), # fits cutoff (len 4), supervised -> kept
_s([0.0, 0.0, 0.0, 0.0, 1.0, 1.0]), # len 6 > 4, supervision only beyond cutoff -> dropped
_s([1.0, 1.0]), # short, fully supervised -> kept
_s([0.0, 0.0, 1.0, 1.0, 1.0, 1.0]), # len 6 > 4 but supervision within cutoff -> kept
]
kept = BatchGenerator._drop_unsupervised(gen, samples)
assert kept == [samples[0], samples[2], samples[3]]
assert gen._warned_truncation is True

View File

@@ -71,6 +71,148 @@ def test_sharegpt_converter():
assert DataConverterPlugin("sharegpt")(example) == expected_data
def test_sharegpt_converter_multimodal():
example = {
"conversations": [
{"from": "human", "value": "What is <image> and what happens in <video>?"},
{"from": "gpt", "value": "An image and a video."},
],
"images": ["/p/a.jpg"],
"videos": ["/p/v.mp4"],
}
expected_data = {
"messages": [
{
"role": "user",
"content": [
{"type": "text", "value": "What is "},
{"type": "image_url", "value": "/p/a.jpg"},
{"type": "text", "value": " and what happens in "},
{"type": "video_url", "value": "/p/v.mp4"},
{"type": "text", "value": "?"},
],
"loss_weight": 0.0,
},
{"role": "assistant", "content": [{"type": "text", "value": "An image and a video."}], "loss_weight": 1.0},
]
}
assert DataConverterPlugin("sharegpt")(example) == expected_data
def test_sharegpt_converter_multiple_images_in_order():
# images are a sample-level list consumed by <image> tags in document order across turns
example = {
"conversations": [
{"from": "human", "value": "<image><image>Compare these."},
{"from": "gpt", "value": "Done."},
],
"images": ["/p/a.jpg", "/p/b.jpg"],
}
user = DataConverterPlugin("sharegpt")(example)["messages"][0]
assert user["content"] == [
{"type": "image_url", "value": "/p/a.jpg"},
{"type": "image_url", "value": "/p/b.jpg"},
{"type": "text", "value": "Compare these."},
]
def test_sharegpt_converter_no_media_unchanged():
# backward compatibility: a scalar (non-list) image column and no tags is normalized; with no
# media columns at all the output is byte-identical to the text-only path.
example = {"conversations": [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "yo"}]}
assert DataConverterPlugin("sharegpt")(example) == {
"messages": [
{"role": "user", "content": [{"type": "text", "value": "hi"}], "loss_weight": 0.0},
{"role": "assistant", "content": [{"type": "text", "value": "yo"}], "loss_weight": 1.0},
]
}
def test_alpaca_converter_multimodal():
example = {"instruction": "Describe <image>", "input": "", "output": "ok", "images": ["/p/a.jpg"]}
user = DataConverterPlugin("alpaca")(example)["messages"][0]
assert user["content"] == [
{"type": "text", "value": "Describe "},
{"type": "image_url", "value": "/p/a.jpg"},
]
def test_pair_converter_multimodal_shared_media():
# chosen and rejected each reference the same sample-level image
example = {
"chosen": [
{"role": "user", "content": "Look at <image>"},
{"role": "assistant", "content": "good"},
],
"rejected": [
{"role": "user", "content": "Look at <image>"},
{"role": "assistant", "content": "bad"},
],
"images": ["/p/a.jpg"],
}
out = DataConverterPlugin("pair")(example)
for side in ("chosen_messages", "rejected_messages"):
assert out[side][0]["content"] == [
{"type": "text", "value": "Look at "},
{"type": "image_url", "value": "/p/a.jpg"},
]
def test_converter_media_count_mismatch():
# more tags than media files
with pytest.raises(ValueError, match="More <image> tags"):
DataConverterPlugin("sharegpt")(
{
"conversations": [{"from": "human", "value": "<image><image>"}, {"from": "gpt", "value": "x"}],
"images": ["/p/a.jpg"],
}
)
# fewer tags than media files
with pytest.raises(ValueError, match="Fewer <image> tags"):
DataConverterPlugin("sharegpt")(
{
"conversations": [{"from": "human", "value": "<image>"}, {"from": "gpt", "value": "x"}],
"images": ["/p/a.jpg", "/p/b.jpg"],
}
)
def test_converter_audio_column_and_tag():
# an <audio> tag consumes the next path from the audios column, lifted into an audio_url block
example = {
"conversations": [
{"from": "human", "value": "hear <audio>What is this?"},
{"from": "gpt", "value": "A bell."},
],
"audios": ["/p/a.wav"],
}
user = DataConverterPlugin("sharegpt")(example)["messages"][0]
assert user["content"] == [
{"type": "text", "value": "hear "},
{"type": "audio_url", "value": "/p/a.wav"},
{"type": "text", "value": "What is this?"},
]
def test_converter_audio_count_mismatch():
# more audio tags than files
with pytest.raises(ValueError, match="More <audio> tags"):
DataConverterPlugin("sharegpt")(
{
"conversations": [{"from": "human", "value": "<audio><audio>"}, {"from": "gpt", "value": "x"}],
"audios": ["/p/a.wav"],
}
)
# fewer audio tags than files
with pytest.raises(ValueError, match="Fewer <audio> tags"):
DataConverterPlugin("sharegpt")(
{
"conversations": [{"from": "human", "value": "<audio>"}, {"from": "gpt", "value": "x"}],
"audios": ["/p/a.wav", "/p/b.wav"],
}
)
@pytest.mark.parametrize("num_samples", [16])
def test_pair_converter(num_samples: int):
data_args = DataArguments(train_dataset="llamafactory/v1-dataset-info/orca-dpo-pairs.yaml")
@@ -117,3 +259,4 @@ def test_pair_converter(num_samples: int):
],
}
assert data_engine[index] == {"_dataset_name": "tiny_dataset", **expected_data}

View File

@@ -13,30 +13,51 @@
# limitations under the License.
import sys
from functools import partial
from unittest.mock import MagicMock, patch
import pytest
import torch.multiprocessing as mp
from torch import nn
from transformers import AutoModelForCausalLM
def _original_fla_op(*args, **kwargs):
return args, kwargs
class _LinearAttention(nn.Module):
def __init__(self):
super().__init__()
self.chunk_gated_delta_rule = _original_fla_op
self.recurrent_gated_delta_rule = _original_fla_op
class _FLAModel(nn.Module):
def __init__(self):
super().__init__()
self.linear_attn = _LinearAttention()
def _apply_kernel(rank) -> None:
with patch("torch.accelerator.current_accelerator") as mock_get_accelerator:
mock_device = MagicMock()
setattr(mock_device, "type", "npu")
mock_get_accelerator.return_value = mock_device
# reload kernel modules to respect mocked accelerator
for k in list(sys.modules.keys()):
if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"):
del sys.modules[k]
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_kernels(model=model, config={"name": "npu_fused_rmsnorm"})
with patch.dict(sys.modules, {"torch_npu": MagicMock()}):
# Reload kernel modules so dependency checks use the mocked NPU environment.
for k in list(sys.modules.keys()):
if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"):
del sys.modules[k]
from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels
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__
@@ -48,18 +69,19 @@ def _apply_all_kernels(rank) -> None:
setattr(mock_device, "type", "npu")
mock_get_accelerator.return_value = mock_device
# reload kernel modules to respect mocked accelerator
for k in list(sys.modules.keys()):
if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"):
del sys.modules[k]
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_kernels(model=model, config={"name": "auto"})
with patch.dict(sys.modules, {"torch_npu": MagicMock()}):
# Reload kernel modules so dependency checks use the mocked NPU environment.
for k in list(sys.modules.keys()):
if k.startswith("llamafactory.v1.plugins.model_plugins.kernels"):
del sys.modules[k]
from llamafactory.v1.plugins.model_plugins.kernels.interface import apply_kernels
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__
@@ -71,3 +93,62 @@ def test_apply_kernel():
def test_apply_all_kernels():
mp.spawn(_apply_all_kernels)
@pytest.mark.runs_on(["npu"])
def test_flash_linear_attention_kernels_compose_with_auto(monkeypatch):
import fsdp_turbo.ops.fla # noqa: F401
from fsdp_turbo.ops import get_op
from llamafactory.v1.plugins.model_plugins.kernels import interface
from llamafactory.v1.plugins.model_plugins.kernels.ops.linear_attention.fla import (
FlashLinearAttentionKernel,
)
model = _FLAModel()
auto_calls = []
monkeypatch.setattr(
interface,
"_apply_auto_kernels",
lambda model, **kwargs: auto_calls.append((model, kwargs)) or model,
)
# FLA execution is outside this bridge test; its external runtime is not required.
monkeypatch.setattr(FlashLinearAttentionKernel, "check_deps", staticmethod(lambda: None))
config = {
"name": "auto, flash-linear-attention",
"include_kernels": "fused_recurrent_gated_delta_rule, chunk_gated_delta_rule",
"chunk_size": 32,
}
assert interface.apply_kernels(model, config) is model
assert auto_calls == [(model, {"config": config, "require_logits": False})]
assert get_op("chunk_gated_delta_rule").__module__ == "fsdp_turbo.ops.fla"
chunk_op = model.linear_attn.chunk_gated_delta_rule
assert isinstance(chunk_op, partial)
assert chunk_op.func.__module__ == "fsdp_turbo.ops.fla"
assert chunk_op.keywords == {"chunk_size": 32}
assert model.linear_attn.recurrent_gated_delta_rule.__module__ == "fsdp_turbo.ops.fla"
with pytest.raises(RuntimeError, match="did not match any model module attributes"):
FlashLinearAttentionKernel.apply(
model=nn.Linear(2, 2),
config={"include_kernels": "chunk_gated_delta_rule", "chunk_size": 32},
)
def test_flash_linear_attention_kernel_validates_config(monkeypatch):
from llamafactory.v1.plugins.model_plugins.kernels.ops.linear_attention.fla import (
FlashLinearAttentionKernel,
)
model = nn.Sequential(nn.Linear(2, 2))
monkeypatch.setattr(FlashLinearAttentionKernel, "check_device", staticmethod(lambda: None))
monkeypatch.setattr(FlashLinearAttentionKernel, "check_deps", staticmethod(lambda: None))
with pytest.raises(ValueError, match="chunk_size"):
FlashLinearAttentionKernel.apply(model=model, config={"include_kernels": "auto", "chunk_size": 48})
with pytest.raises(ValueError, match="Unsupported Flash Linear Attention kernels"):
FlashLinearAttentionKernel.apply(model=model, config={"include_kernels": "not_a_kernel"})

View File

@@ -20,6 +20,7 @@ from llamafactory.v1.accelerator.interface import DistributedInterface
from llamafactory.v1.config.model_args import ModelArguments
from llamafactory.v1.config.training_args import TrainingArguments
from llamafactory.v1.core.model_engine import ModelEngine
from llamafactory.v1.plugins.model_plugins.parallelization import ulysses
from llamafactory.v1.plugins.model_plugins.parallelization.sequence_parallel import (
SequenceParallelModelPlugin,
sequence_parallel_loss,
@@ -28,6 +29,39 @@ from llamafactory.v1.utils.env import find_available_port
from llamafactory.v1.utils.pytest import dist_env
def test_qwen3_5_broadcast_position_ids_keep_packed_boundaries(monkeypatch: pytest.MonkeyPatch):
local_position_ids = torch.tensor([[0, 1, 0]])
remote_position_ids = torch.tensor([[1, 2, 3]])
mrope_position_ids = local_position_ids.unsqueeze(0).expand(3, -1, -1)
captured = {}
monkeypatch.setattr(ulysses.SeqAllToAll4D, "apply", lambda _, tensor, *__: tensor)
monkeypatch.setattr(ulysses, "get_ulysses_sequence_parallel_world_size", lambda _: 2)
def fake_all_gather(outputs, tensor, **_):
outputs[0].copy_(tensor)
outputs[1].copy_(remote_position_ids if tensor.shape == local_position_ids.shape else tensor)
def fake_attention(query, _key, _value, _attention_mask, **kwargs):
captured["position_ids"] = kwargs["position_ids"]
return query
monkeypatch.setattr(ulysses.dist, "all_gather", fake_all_gather)
attention = ulysses.UlyssesAttention(sequence_process_group=object(), attn_fn=fake_attention)
hidden_states = torch.zeros(1, 3, 2, 4)
attention(hidden_states, hidden_states, hidden_states, None, 6, position_ids=mrope_position_ids)
assert captured["position_ids"].tolist() == [[0, 1, 0, 1, 2, 3]]
assert captured["position_ids"].is_contiguous()
def test_true_mrope_position_ids_are_not_used_as_packed_boundaries():
mrope_position_ids = torch.tensor([[[0, 1, 2]], [[0, 1, 1]], [[0, 1, 0]]])
assert ulysses._get_text_position_ids(mrope_position_ids) is None
def _test_sequence_parallel_loss(
local_rank: int, world_size: int, master_port: int, cp_size: int, dp_size: int, batch_size: int
):

View File

@@ -0,0 +1,134 @@
# 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 types import SimpleNamespace
import pytest
import torch
from llamafactory.v1.plugins.trainer_plugins.distributed import fsdpturbo as fsdpturbo_module
from llamafactory.v1.plugins.trainer_plugins.distributed.fsdpturbo import (
FSDPTurboEPModelSpec,
FSDPTurboFSDP2Engine,
FSDPTurboParallelState,
)
from llamafactory.v1.plugins.trainer_plugins.distributed.interface import (
DistributedPlugin,
FSDPTurboParams,
)
class _Model(torch.nn.Module):
def __init__(self, model_type: str):
super().__init__()
self.config = SimpleNamespace(model_type=model_type)
def test_qwen35_ep_model_spec():
spec = FSDPTurboEPModelSpec.get(_Model("qwen3_5_moe"))
assert spec is not None
assert spec.ep_modules == ["model.language_model.layers.{*}.mlp.experts"]
assert spec.ep_fsdp_modules == ["model.language_model.layers.{*}.mlp"]
def test_fsdpturbo_uses_class_plugin_and_strict_backend_params():
plugin = DistributedPlugin("fsdpturbo")
params = plugin.parse_params({"name": "fsdpturbo", "ep_size": 4}, FSDPTurboParams)
assert params.ep_size == 4
assert callable(plugin.shard_model)
assert callable(plugin.clip_grad_norm)
with pytest.raises(ValueError, match="Unknown params"):
plugin.parse_params({"name": "fsdpturbo", "cp_size": 2}, FSDPTurboParams)
for key in ("ep_modules", "ep_fsdp_modules"):
with pytest.raises(ValueError, match="Unknown params"):
plugin.parse_params({"name": "fsdpturbo", key: ["model.layers.*.mlp"]}, FSDPTurboParams)
def test_fsdpturbo_sets_storage_dtype_inside_backend(monkeypatch):
from llamafactory.v1.plugins.trainer_plugins.distributed.fsdp2 import FSDP2Engine
monkeypatch.setattr(FSDP2Engine, "shard_model", lambda self, model: model)
engine = object.__new__(FSDPTurboFSDP2Engine)
engine.mixed_precision = "bf16"
model = torch.nn.Linear(2, 2, dtype=torch.float32)
assert engine.shard_model(model).weight.dtype == torch.bfloat16
def test_fsdpturbo_sets_public_efsdp_gradient_divide_factor(monkeypatch):
expert_parallel_module = pytest.importorskip("fsdp_turbo.distributed.expert_parallel.expert_parallel")
expert_fully_shard_module = pytest.importorskip(
"fsdp_turbo.distributed.expert_parallel.expert_fully_shard_parallel"
)
captured = {}
monkeypatch.setattr(expert_parallel_module, "expert_parallelize_modules", lambda model, mesh, plan: model)
def _expert_fully_shard_modules(model, mesh, ep_plan, fsdp_plan):
captured["gradient_divide_factor"] = ep_plan.gradient_divide_factor
return model
monkeypatch.setattr(expert_fully_shard_module, "expert_fully_shard_modules", _expert_fully_shard_modules)
engine = object.__new__(FSDPTurboFSDP2Engine)
engine.dist_config = {"ep_dispatcher": "eager"}
engine.ep_size = 4
engine.ep_fsdp_size = 2
engine.parallel_state = SimpleNamespace(efsdp_size=2, ep_mesh=object(), efsdp_mesh=object())
engine.rank = 0
engine.prepare_model_ep(_Model("qwen3_5_moe"))
assert captured["gradient_divide_factor"] == 8.0
def test_fsdpturbo_owns_expert_mesh_topology(monkeypatch):
calls = []
class _Mesh:
def __init__(self, name="expert"):
self.name = name
def __getitem__(self, name):
return _Mesh(name)
def _init_device_mesh(**kwargs):
calls.append(kwargs)
return _Mesh()
class _DistributedInterface:
current_device = torch.device("cpu")
strategy = SimpleNamespace(cp_size=1)
def get_world_size(self, dim):
return 16
def get_device_mesh(self, dim):
return _Mesh("dp")
monkeypatch.setattr(fsdpturbo_module, "init_device_mesh", _init_device_mesh)
state = FSDPTurboParallelState()
state.initialize(_DistributedInterface(), {"ep_size": 8})
assert calls == [
{
"device_type": "cpu",
"mesh_shape": (1, 2, 8, 1),
"mesh_dim_names": ("edp", "efsdp", "ep", "expert_cp"),
}
]
assert state.ep_mesh.name == "ep"
assert state.efsdp_mesh.name == "efsdp"
assert state.expert_cp_mesh.name == "expert_cp"

View File

@@ -29,8 +29,6 @@ model: Qwen/Qwen3-0.6B
trust_remote_code: true
model_class: llm
template: qwen3_nothink
kernel_config:
name: auto
@@ -41,7 +39,7 @@ dist_config:
dcp_path: null
init_config:
name: init_on_meta
name: init_on_default
# PEFT Configuration
peft_config: