mirror of
https://github.com/hiyouga/LLaMA-Factory.git
synced 2026-08-17 13:35:44 +08:00
Compare commits
10 Commits
1b47415a2f
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f28afaf635 | ||
|
|
bc4b42cefc | ||
|
|
199b8873d7 | ||
|
|
0bbe481e6e | ||
|
|
63a89710c7 | ||
|
|
887b850813 | ||
|
|
84576b1408 | ||
|
|
713b5a3f95 | ||
|
|
62ae362455 | ||
|
|
3984675dd5 |
45
.github/workflows/docker_npu.yml
vendored
45
.github/workflows/docker_npu.yml
vendored
@@ -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
|
||||
|
||||
3
.github/workflows/docs.yml
vendored
3
.github/workflows/docs.yml
vendored
@@ -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 }}
|
||||
|
||||
9
.github/workflows/tests_npu.yml
vendored
9
.github/workflows/tests_npu.yml
vendored
@@ -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
106
README.md
@@ -1,9 +1,9 @@
|
||||

|
||||

|
||||
|
||||
[](https://github.com/hiyouga/LLaMA-Factory/stargazers)
|
||||
[](https://github.com/hiyouga/LLaMA-Factory/commits/main)
|
||||
[](https://github.com/hiyouga/LLaMA-Factory/graphs/contributors)
|
||||
[](https://github.com/hiyouga/LLaMA-Factory/actions/workflows/tests.yml)
|
||||
[](https://github.com/hiyouga/LlamaFactory/stargazers)
|
||||
[](https://github.com/hiyouga/LlamaFactory/commits/main)
|
||||
[](https://github.com/hiyouga/LlamaFactory/graphs/contributors)
|
||||
[](https://github.com/hiyouga/LlamaFactory/actions/workflows/tests.yml)
|
||||
[](https://pypi.org/project/llamafactory/)
|
||||
[](https://scholar.google.com/scholar?cites=12620864006390196564)
|
||||
[](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
|
||||
|
||||
|
||||
113
README_zh.md
113
README_zh.md
@@ -1,9 +1,9 @@
|
||||

|
||||

|
||||
|
||||
[](https://github.com/hiyouga/LLaMA-Factory/stargazers)
|
||||
[](https://github.com/hiyouga/LLaMA-Factory/commits/main)
|
||||
[](https://github.com/hiyouga/LLaMA-Factory/graphs/contributors)
|
||||
[](https://github.com/hiyouga/LLaMA-Factory/actions/workflows/tests.yml)
|
||||
[](https://github.com/hiyouga/LlamaFactory/stargazers)
|
||||
[](https://github.com/hiyouga/LlamaFactory/commits/main)
|
||||
[](https://github.com/hiyouga/LlamaFactory/graphs/contributors)
|
||||
[](https://github.com/hiyouga/LlamaFactory/actions/workflows/tests.yml)
|
||||
[](https://pypi.org/project/llamafactory/)
|
||||
[](https://scholar.google.com/scholar?cites=12620864006390196564)
|
||||
[](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.5(Qwen2 测试版)系列模型已在 LLaMA-Factory 中实现微调支持。详情请查阅该[博客页面](https://qwenlm.github.io/zh/blog/qwen1.5/)。
|
||||
[24/02/05] Qwen1.5(Qwen2 测试版)系列模型已在 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)
|
||||
|
||||
## 引用
|
||||
|
||||
|
||||
25
data/v1_multimodal_demo.jsonl
Normal file
25
data/v1_multimodal_demo.jsonl
Normal 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日进入太空的首位平民宇航员。他负责在轨操作空间科学实验有效载荷。"}]}]}
|
||||
|
||||
4
data/v1_multimodal_demo.yaml
Normal file
4
data/v1_multimodal_demo.yaml
Normal file
@@ -0,0 +1,4 @@
|
||||
multimodal_demo:
|
||||
path: data/v1_multimodal_demo.jsonl
|
||||
source: local
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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 的许可证不会替代或覆盖这些条款。
|
||||
|
||||
本镜像按“原样”提供,不附带任何明示或暗示的保证。用户需要自行验证软硬件兼容性、保障容器及运行配置的安全、遵守适用的许可证和法律,并在训练、评测或部署前审查模型与数据集的使用条款。
|
||||
|
||||
@@ -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:
|
||||
|
||||
173
docs/en/advanced/distributed/fsdpturbo-ep-efsdp.md
Normal file
173
docs/en/advanced/distributed/fsdpturbo-ep-efsdp.md
Normal 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.
|
||||
83
docs/en/advanced/ktransformers.md
Normal file
83
docs/en/advanced/ktransformers.md
Normal 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/`.
|
||||
@@ -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
|
||||
|
||||
|
||||
215
docs/zh/advanced/distributed/fsdpturbo-ep-efsdp.md
Normal file
215
docs/zh/advanced/distributed/fsdpturbo-ep-efsdp.md
Normal 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.post4,A5 环境为 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 | 关闭 | FLA(chunk size 16)/ eager | 100 | 1.3361 -> 0.0793 | 2.51 s/it | 通过并完成保存 |
|
||||
| Atlas 900 A3 SuperPoD | 1 | 16 | 1 | 关闭 | FLA(chunk size 16)/ fused | 100 | 1.3354 -> 0.1179 | 2.17 s/it | 通过并完成保存 |
|
||||
| Atlas 900 A3 SuperPoD | 2 | 4 | 2 | 关闭 | auto + FLA(chunk size 64)/ fused | 100 | 1.8114 -> 0.5260 | 7.65 s/it | 通过并完成保存 |
|
||||
| Atlas 900 A3 SuperPoD | 2 | 4 | 2 | 关闭 | auto + FLA(chunk 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.997,CP2/EP4/EFSDP2 eager/fused 为 0.977,说明两种 dispatcher 的
|
||||
优化轨迹一致。性能收益与切分有关:EP16 下 fused 比 eager 快约 13%,而加入 CP 和 EFSDP 后 fused
|
||||
比 eager 慢约 30%,因此不能把 fused 视为所有 mesh 的默认最优选择。
|
||||
|
||||
EP16 两组使用 global batch 16 和 cutoff length 256;CP2 两组使用 global batch 8 和 cutoff length
|
||||
128;A5 组使用 global batch 8 和 cutoff length 256。因此,首末 loss 用于验证各组自身的收敛趋势,
|
||||
不同切分组之间的绝对 loss 不应直接作为精度等价结论。
|
||||
119
docs/zh/advanced/ktransformers.md
Normal file
119
docs/zh/advanced/ktransformers.md
Normal file
@@ -0,0 +1,119 @@
|
||||
# KTransformers LoRA SFT
|
||||
|
||||
KTransformers(KT)将 MoE routed experts 放在 CPU 执行,LLaMA-Factory 继续负责数据、LoRA 参数和训练入口。
|
||||
当前生产范围是 routed-BF16 LoRA 与 routed-INT8 LoRA;Accelerate 配置只负责 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
|
||||
均包含非零更新。
|
||||
@@ -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
|
||||
|
||||
|
||||
6
examples/inference/mossvl.yaml
Normal file
6
examples/inference/mossvl.yaml
Normal 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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
14
examples/merge_lora/mossvl_lora_sft.yaml
Normal file
14
examples/merge_lora/mossvl_lora_sft.yaml
Normal 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
|
||||
57
examples/train_full/mossvl_full_sft.yaml
Normal file
57
examples/train_full/mossvl_full_sft.yaml
Normal 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
|
||||
54
examples/train_lora/mossvl_lora_sft.yaml
Normal file
54
examples/train_lora/mossvl_lora_sft.yaml
Normal 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
|
||||
@@ -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
|
||||
27
examples/v1/train_full/train_multimodal.yaml
Normal file
27
examples/v1/train_full/train_multimodal.yaml
Normal 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
|
||||
|
||||
2
requirements/fsdpturbo.txt
Normal file
2
requirements/fsdpturbo.txt
Normal 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
3
requirements/moss-vl.txt
Normal file
@@ -0,0 +1,3 @@
|
||||
transformers==4.57.1
|
||||
torchcodec==0.7.0
|
||||
joblib
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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": {
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"],
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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."
|
||||
)
|
||||
|
||||
@@ -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.")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
277
src/llamafactory/v1/core/utils/collation.py
Normal file
277
src/llamafactory/v1/core/utils/collation.py
Normal 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
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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.")
|
||||
@@ -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
|
||||
|
||||
@@ -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>")
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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():
|
||||
|
||||
631
tests/data/test_moss_vl_plugin.py
Normal file
631
tests/data/test_moss_vl_plugin.py
Normal 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]
|
||||
39
tests/data/test_moss_vl_training_configs.py
Normal file
39
tests/data/test_moss_vl_training_configs.py
Normal 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
|
||||
@@ -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)
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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}
|
||||
|
||||
|
||||
@@ -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"})
|
||||
|
||||
@@ -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
|
||||
):
|
||||
|
||||
@@ -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"
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user