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

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

.claude/skills/llamafactory-sft/SKILL.md

@@ -0,0 +1,366 @@
+---
+name: llamafactory-sft
+description: One-stop guided LlamaFactory SFT workflow — data prep, model prep, fine-tuning config/method selection, background training, loss visualization, effect validation and model export. TRIGGER when the user wants to fine-tune / SFT a model with LlamaFactory, run a full LlamaFactory training pipeline, or asks to "微调" / "用 llamafactory 做 sft" / "train a LoRA". SKIP for general questions about the codebase or non-training tasks.
+---
+
+# LlamaFactory One-Stop SFT Workflow
+
+Guide the user through a complete LlamaFactory SFT run: **data prep → model prep → fine-tuning config → background training → loss visualization → (optional) effect validation → (optional) model export**.
+
+The core of this skill is **interactive guidance**: at every key decision point use `AskUserQuestion` and never make irreversible assumptions on the user's behalf.
+
+The working directory defaults to the current directory (`./`, the repo root). Run all commands via `llamafactory-cli` (or `lmf`).
+
+> **Artifacts directory (keep generated files out of model / checkpoint folders).** All skill-generated **yaml configs** (train / inference / export) and **run logs** (download / train / export) MUST be written to a single dedicated directory, NOT inside `./models/` (downloaded weights) or the training `output_dir` (saved checkpoints). Use `./llamafactory_runs/<run_id>/` where `<run_id>` is `<model>_<method>_YYYYMMDD_HHMM` (e.g. `./llamafactory_runs/qwen3-4b_lora_20260615_0807/`). Create it once at the start of the CLI route and place every config + log file there. The training `output_dir` (the actual model checkpoints) stays separate under `saves/...`, and downloaded base models stay under `./models/...`. This keeps configs/logs, checkpoints, and weights cleanly separated.
+
+> Language: communicate with the user in whatever language they request, or otherwise follow the active language setting/preference. Do not hardcode a fixed conversation language.
+
+---
+
+## Progress board (show only when progress changes)
+
+This workflow has many stages, so the user must always be able to see **where we are**: which steps are done, which is in progress, and which remain.
+
+**Rule:** Render the **progress board** only when the board state changes, or when the user explicitly asks for status/progress. A board state changes when any step marker changes (`[ ]` / `[~]` / `[x]` / `[-]`) or when Stage 0 changes the branch shape (e.g. skipped validation/export for "SFT only"). Do **not** repeat the board on routine status updates, repeated confirmations, log polling, or consecutive messages within the same stage if the markers are unchanged.
+
+Use these markers:
+- `[x]` done
+- `[~]` in progress (the step you are working on right now)
+- `[ ]` not started
+- `[-]` skipped (e.g. validation/export when the user chose "SFT only", or download when the model is already local)
+
+Render it as a compact checklist inside a fenced code block. **Always wrap the entire progress board in triple backticks** (use `text` as the info string, or no info string) so Markdown preserves every line break and marker exactly. **The board's language must follow the active conversation language / the user's request — it is NOT required to be Chinese.** The example below is in Chinese only for illustration; render the title and step labels in whatever language the user is using (e.g. English: a "Progress" board with "0. Confirm overall flow", "1. Data preparation", ...):
+
+```
+进度看板
+[x] 0. 确认整体流程（范围 / 执行方式）
+[x] 1. 数据准备
+[~] 2. 模型准备
+[ ] 3. 微调配置
+[ ] 4. SFT 训练
+[ ] 5. 效果验证
+[ ] 6. 模型导出
+```
+
+Guidelines:
+- The step list above is the canonical set. Loss visualization is part of **SFT 训练** (it happens automatically once training finishes) — you know to handle it, but do NOT show it as a separate line on the board.
+- Mark steps `[-]` instead of dropping them when they don't apply to the chosen flow (e.g. "SFT only" skips 5 & 6; WebUI route collapses 3–6 into the UI).
+- After Stage 0, adapt the board to the chosen branch (mark skipped steps `[-]`) and keep that shape for the rest of the run.
+- Do **not** render the board as plain paragraph text or as Markdown task-list checkboxes (`- [x] ...`); both can change the intended layout or marker semantics.
+- Keep it terse — one line per step. When the board is shown, put it at the **top** of your message, then continue with the actual content / question below it.
+- Update markers as soon as a step's status changes; show the updated board once at that transition, and never show a stale board.
+
+---
+
+## Stage 0: Confirm the overall flow
+
+Before doing anything, use a **single `AskUserQuestion` call that asks all three of the most important branches together** (one question object each, in the same call) so the user can settle the whole shape of the run in one step:
+
+1. **Flow scope**:
+   - SFT only (data prep + model prep + fine-tuning)
+   - Full flow (also includes effect validation + model export)
+
+2. **Execution mode**:
+   - Run everything via CLI commands
+   - Launch a **WebUI** for the SFT / validation / export parts and let the user operate it
+
+3. **Fine-tuning type** — which SFT fine-tuning type to use, so you can match the corresponding official example yaml directory:
+   - **LoRA** → base config from `examples/train_lora/` (recommended default)
+   - **Full** (full-parameter) → base config from `examples/train_full/`
+   - **QLoRA** (quantized LoRA) → base config from `examples/train_qlora/`
+
+Notes on the fine-tuning-type answer:
+- The type only matters for the **all-CLI** route. If the user picked **WebUI** in question 2, **ignore their fine-tuning-type answer** (the type is selected inside the UI) — it does no harm to have asked.
+- For the **all-CLI** route, record the chosen type and use it (together with the model family resolved in Stage 2) to pick the base example yaml in Stage 3.
+
+Record the user's choices and branch the later stages accordingly.
+
+---
+
+## Stage 1: Data preparation
+
+Use `AskUserQuestion` to ask about the data source:
+
+- **Use LlamaFactory built-in data**: list selectable datasets from `data/dataset_info.json` (e.g. `identity`, `alpaca_zh_demo`, `alpaca_en_demo`, ...). Allow multi-select and join into `dataset: a,b,c`.
+- **Use custom data**: ask for the data file path(s) (multiple allowed).
+
+### Custom data validation & onboarding
+
+For each custom data file:
+
+1. **Read and validate the format.** LlamaFactory supports `alpaca` and `sharegpt` formats; file types may be json/jsonl/csv/parquet/arrow.
+   - **Alpaca format** key fields: `instruction` (required), `input` (optional), `output` (required), `system`/`history` (optional).
+   - **ShareGPT format** key fields: a `conversations` list whose elements contain `from` (human/gpt) and `value`; multimodal data adds `images`/`videos`/`audios`.
+   - Full field definitions are in `data/README.md` / `data/README_zh.md`; read them when needed.
+2. **If it does not meet the requirements, fix it**: with the user's consent, convert the data into a compliant format (write a NEW file, never destroy the user's original data in place).
+3. **Copy it into the `data/` directory.**
+4. **Update `data/dataset_info.json`**: add a dataset description entry. Minimal form:
+   ```json
+   "my_dataset": { "file_name": "my_dataset.json" }
+   ```
+   For sharegpt or custom column names, also add `formatting` / `columns` / `tags`.
+   When editing this JSON, keep it valid (use Edit for precise insertion; verify braces/commas).
+
+### identity.json and other templated datasets
+
+If the user picks `identity.json` (contains template variables like `{{name}}`, `{{author}}`), use `AskUserQuestion` to ask whether to do a global replacement. If yes:
+- **Use `AskUserQuestion` to collect the concrete replacement values for the `{{name}}` and `{{author}}` variables** (one question per variable, e.g. "模型名 {{name}}" and "作者/机构 {{author}}"), plus any other template variables the file contains. Do not assume or hardcode these values — always ask the user.
+- **Recommend copying first, then replacing** (e.g. `data/identity_custom.json`) to avoid polluting the repo's built-in file, and register the new name in `dataset_info.json`.
+- Use Edit with `replace_all` to perform the replacement.
+
+---
+
+## Stage 2: Model preparation
+
+Use `AskUserQuestion` to ask for the model choice, and **offer a default recommendation** (e.g. `Qwen/Qwen3-4B-Instruct-2507`, matching `examples/train_lora/qwen3_lora_sft.yaml`).
+
+Once the model is decided:
+1. **Verify the model type against LlamaFactory's own model list, and use it to pick the matching default yaml.** Before downloading or configuring anything, look up the chosen model in LlamaFactory's built-in registry rather than guessing:
+   - `SUPPORTED_MODELS` and `DEFAULT_TEMPLATE` live in `src/llamafactory/extras/constants.py` (registered via `register_model_group(...)`, each group sharing one `template=`). This is the authoritative list of which models LlamaFactory supports and which template each uses.
+   - Match the user's model (by HF/ModelScope id or model name) to an entry there to learn its **model family / type and default `template`** (e.g. a Qwen3 model → `qwen3` / `qwen3_nothink`). Cross-check `src/llamafactory/data/template.py` (the `TEMPLATES` dict) if needed.
+   - **Use the resolved family/template — together with the fine-tuning type chosen in Stage 0 (LoRA → `examples/train_lora/`, Full → `examples/train_full/`, QLoRA → `examples/train_qlora/`) — to select the closest official example yaml** (and the matching `examples/inference/`, `examples/merge_lora/`) as the base config for Stage 3 — e.g. a Qwen3 model with LoRA → `examples/train_lora/qwen3_lora_sft.yaml`. If there is no exact example for that family/type, pick the nearest supported one and note that you adapted it.
+   - **If the model is NOT in the supported list** (no entry / no matching template), clearly tell the user it is unsupported, so they can switch to a supported model or define a custom template — do not silently proceed.
+2. **Check whether it already exists locally** (HF cache `~/.cache/huggingface/hub`, or a user-specified local path).
+   - **Verify completeness, not just existence — beware empty shells.** A cached directory can exist while being only a few KB (an empty shell where the real download never finished). After finding a local copy, verify it: total size is in the expected GB range, all `*.safetensors` shards referenced by `model.safetensors.index.json` are present, and there are no `*.incomplete` files. Only treat the model as available if it passes; otherwise treat it as NOT downloaded and proceed to download.
+3. **If not present locally**: first use `AskUserQuestion` to confirm the **download source**:
+   - **Hugging Face Hub**
+   - **ModelScope** (often faster in mainland China)
+   - **Let the agent decide** (pick automatically based on network reachability / region)
+
+   Then **start a download task concurrently** via `Agent` / Bash with `run_in_background`:
+   ```bash
+   # Hugging Face (modern CLI; `huggingface-cli` is deprecated — use `hf download`)
+   # Enable Xet high-performance transfer for much faster downloads:
+   HF_XET_HIGH_PERFORMANCE=1 hf download <model_id> --local-dir <path>
+   # ModelScope (set USE_MODELSCOPE_HUB=1 to also let LlamaFactory auto-download at train time)
+   modelscope download --model <model_id> --local_dir <path>
+   ```
+   Notes:
+   - `huggingface-cli download` is **deprecated** in recent `huggingface_hub` versions and may just print help text — always use `hf download`.
+   - The old `HF_HUB_ENABLE_HF_TRANSFER=1` is **deprecated** too; use `HF_XET_HIGH_PERFORMANCE=1` instead.
+   - If a download source is slow (e.g. ModelScope < ~1 MB/s), consider switching to the other source rather than waiting hours. In practice HF + Xet is often dramatically faster.
+
+   You may instead rely on auto-download at train time (set `USE_MODELSCOPE_HUB=1` for ModelScope).
+   Keep advancing the config work while it downloads, and report download progress to the user periodically (interval ≤ 100 seconds).
+
+---
+
+## Stage 2.5: GPU selection (detect free GPU(s) once before running)
+
+**Before the first command that uses the GPU** (SFT training, and later validation / export), **detect GPU status once, pick the device(s) to use, and reuse that exact choice** for SFT, validation, and export. Never assume a GPU is free — other jobs may already be using the machine. **Detect only once here** — do NOT re-detect before every run; the index/indices chosen now are reused for the whole workflow.
+
+1. **Detect GPUs and their load.** Try AMD/ROCm first, then NVIDIA. Wrap the call in `timeout` so a live-refreshing monitor can never hang the agent (plain `amd-smi` without a subcommand only prints help and lacks the `VRAM_USAGE` / `GFX%` snapshot, so keep the `monitor` subcommand for the one-shot table):
+   ```bash
+   timeout 10 amd-smi monitor 2>/dev/null || rocm-smi 2>/dev/null || nvidia-smi
+   ```
+   Read each GPU's **VRAM usage** and **utilization**. Treat a GPU as **free** when its VRAM usage is near-empty (e.g. ≲ 1 GB) and utilization is low (e.g. ≲ a few %). Also list any running training/inference processes if helpful (e.g. `pgrep -af "llamafactory-cli"`).
+2. **Decide which device(s) to use:**
+   - **No free GPU:** do NOT silently queue onto a busy card — tell the user which GPUs are busy (and roughly by how much VRAM / what is running) and use `AskUserQuestion` to let them choose: wait for one to free up, share a partially-used GPU anyway, or specify a particular index.
+   - **Exactly one free GPU:** use it (single-card).
+   - **Multiple free GPUs:** use `AskUserQuestion` to ask the user whether to run **single-card** or **multi-card** (list the free indices). If the user picks multi-card, use the selected free indices together.
+3. **Record the chosen index/indices and reuse them for every GPU command** (training, validation, export) by exporting the platform's visibility env var:
+   - **AMD/ROCm:** `HIP_VISIBLE_DEVICES=<idx>` (multi-card: comma-separated, e.g. `HIP_VISIBLE_DEVICES=0,1`)
+   - **NVIDIA/CUDA:** `CUDA_VISIBLE_DEVICES=<idx>` (multi-card: comma-separated)
+
+   Caveat: the visibility-env index follows the smi enumeration, and the selected device becomes index 0 *inside* the process. The mapping between `HIP_VISIBLE_DEVICES` and the `amd-smi` physical order may not be 1:1 — after launching, confirm the **intended** card's VRAM actually rises in `amd-smi` (and not some other card) before trusting it.
+
+Tell the user **which GPU(s) you selected and why** (free vs busy, single vs multi), and surface it in the relevant status updates.
+
+Notes:
+- **Progress board:** GPU selection is part of **模型准备** prep work — handle it before the first run, but do NOT add a separate board line for it.
+- **WebUI route:** you don't drive the runs yourself, but still detect and **tell the user which GPU is free** so they can set it in the UI / launch env.
+
+---
+
+## Stage 3A: WebUI route (user chose WebUI)
+
+If the user chose WebUI in Stage 0:
+
+```bash
+llamafactory-cli webui
+```
+
+- Launch in the background, redirecting logs to a log file.
+- Read the actual listening **port** from the log (default 7860) and tell the user.
+- If on a remote SSH environment, explain port forwarding:
+  ```bash
+  ssh -L 7860:localhost:7860 user@remote_host
+  ```
+  then open `http://localhost:7860` in the local browser.
+- Tell the user that subsequent SFT / validation / export are done in the UI.
+- **After the UI is up, print a "建议在 WebUI 中填写的配置 / Suggested WebUI settings" table** that maps the choices already made in earlier stages onto the fields the user will fill in the WebUI, so they don't have to remember them. Derive the values from what was resolved earlier — do NOT re-ask. Include at least:
+
+  | WebUI 字段 | 建议值 | 来源 |
+  |-----------|--------|------|
+  | Model path (模型路径) | `<resolved local model path or HF/MS id>` | Stage 2 模型准备 |
+  | Template (对话模板) | `<resolved template>` | Stage 2 注册表解析（与训练一致） |
+  | Dataset (数据集) | `<dataset a,b,c>` | Stage 1 数据准备 |
+  | Finetuning method (微调方法) | `<lora / full / qlora>` | Stage 0 微调类型 |
+
+  Add any other already-known values that map to UI fields (e.g. dataset dir if custom data was onboarded). Make clear these are **suggestions to enter in the UI** (LlamaFactory's WebUI does not auto-load them), and that the user can still adjust everything in the interface. If something was not resolved yet (e.g. the model path because the user deferred to train-time auto-download), say so instead of inventing a value.
+
+---
+
+## Stage 3B: CLI route (user chose all-CLI)
+
+If the user chose all-CLI:
+
+1. Using the official example that matches **both the chosen model family AND the fine-tuning type selected in Stage 0** (LoRA → `examples/train_lora/`, Full → `examples/train_full/`, QLoRA → `examples/train_qlora/`; e.g. a Qwen3 + LoRA run → `examples/train_lora/qwen3_lora_sft.yaml`) as a template, **match the default yaml to the chosen model**. **Prefer the official example's default parameters** — treat the example yaml as the source of truth and change as little as possible. Any field may still be changed when the user's needs or data selection call for it (including `template`), but **every change must be tracked and justified** (see step 3).
+   - If the repo has no ready-made config for that model, check whether it is supported (template list). If unsupported, report back to the user.
+2. **Print the key parameters to the user for confirmation** as a table: `model_name_or_path`, `stage: sft`, `finetuning_type` (lora/full/qlora), `lora_rank`/`lora_target`, `dataset`, `template`, `cutoff_len`, `output_dir`, `per_device_train_batch_size`, `gradient_accumulation_steps`, `learning_rate`, `num_train_epochs`, `bf16`, etc.
+3. **If you changed ANY value away from the official example's default**, then *after* the full parameter table, also show a **separate "差异 / Diff vs. default" table** listing only the changed fields, in the form `| parameter | default value | new value | reason |`. The reason must state *why* it changed — e.g. user-specified, required adaptation to the chosen model/dataset, or another concrete cause. This makes every deviation explicit and reviewable. If nothing was changed from the defaults, say so explicitly and omit the diff table.
+4. Use `AskUserQuestion` to ask whether the fine-tuning **method/parameters** need adjusting (finetuning_type, lora_rank, learning rate, number of epochs, etc.), and modify as needed. When you do adjust, update the diff table accordingly.
+
+> **Small-dataset hyperparameter reminder.** For very small datasets (e.g. `identity` with ~91 rows), the example defaults can produce too few steps: total_steps ≈ ceil(num_rows / (batch_size × grad_accum)) × epochs. If total_steps is tiny (single digits) or `logging_steps` ≥ total_steps, you will get too few loss points to plot a curve, and the model may underfit. Conversely, cranking epochs very high on a single tiny dataset causes **overfitting / catastrophic forgetting** (the model memorizes identity but its general language ability degrades — garbled or wrong-language output). Recommended mitigation: mix in a general dataset (e.g. `identity_custom,alpaca_zh_demo,alpaca_en_demo`), keep epochs moderate, and set `logging_steps` small enough to capture several points. Surface this trade-off to the user when relevant.
+
+---
+
+## Stage 4: Generate yaml and run training
+
+1. **Generate the final yaml** into the **artifacts directory** (`./llamafactory_runs/<run_id>/`, see the top-of-file convention), NOT into the checkpoint or model folders. Use a **distinguishing suffix** so multiple runs do not collide — combine the **model name + fine-tuning method + date/time stamp**, e.g. `llamafactory_runs/<run_id>/sft.yaml` where `<run_id>` = `<model>_<method>_YYYYMMDD_HHMM` (build the stamp with `date +%Y%m%d_%H%M`). The yaml's `output_dir` (the actual checkpoints) is a separate location under `saves/<model>/<method>/sft_YYYYMMDD_HHMM`. Show the yaml to the user for **final confirmation**. Make sure `plot_loss: true` so loss can be plotted later.
+2. After confirmation, **run SFT in the background on the GPU(s) chosen in Stage 2.5**, writing the log into the same artifacts directory:
+   ```bash
+   TS=$(date +%Y%m%d_%H%M)
+   RUN_ID="<model>_<method>_${TS}"
+   RUN_DIR="llamafactory_runs/${RUN_ID}"
+   mkdir -p "${RUN_DIR}"
+   # Prefix with the selected GPU's visibility env var (HIP_VISIBLE_DEVICES for AMD, CUDA_VISIBLE_DEVICES for NVIDIA):
+   HIP_VISIBLE_DEVICES=<idx> nohup llamafactory-cli train "${RUN_DIR}/sft.yaml" > "${RUN_DIR}/train.log" 2>&1 &
+   ```
+   Use Bash with `run_in_background`, or `nohup ... &`.
+   - **Beware the "fake completion" notification.** When you launch training with `nohup ... &` (or a backgrounded wrapper), the wrapper shell exits immediately and the harness may emit a `<task-notification> ... completed (exit code 0)` event — but the **real training process is still running** (it was detached). Do NOT treat that notification as "training finished". Training is only truly done when you confirm it via the actual process/log: `pgrep -f "llamafactory-cli train"` shows no process AND the log contains `Training completed` / a `train_runtime` line / the final `train metrics`. Until then, keep polling.
+3. **Periodically check task status and report to the user** (interval ≤ 100 seconds): tail the log file, check the process is alive (and optionally `amd-smi` / `nvidia-smi`). **Each status report must include BOTH the current run state AND the current loss.**
+   - Progress (`current step / total steps`) comes from the tqdm progress-bar lines, which use `\r`; convert `\r`→`\n` first, e.g. `tr '\r' '\n' < "${RUN_DIR}/train.log" | grep -oE "[0-9]+/[0-9]+ \[[^]]*\]" | tail -1` (reference the log path directly — a literal `<log>` placeholder would be parsed by bash as a redirection operator and error out).
+   - Loss is logged as dict lines like `{'loss': '5.227', 'grad_norm': '4.634', 'learning_rate': '9.924e-05', 'epoch': '5'}` (emitted every `logging_steps`). Extract the latest with e.g. `grep -aoE "\{'loss':[^}]*\}" "${RUN_DIR}/train.log" | tail -1`, and report the latest `loss` value (optionally with `epoch`) to the user alongside the step progress.
+
+---
+
+## Stage 5: Loss visualization
+
+After training finishes:
+- `output_dir` will contain `training_loss.png` (because `plot_loss: true`); **just tell the user the image path** — do NOT render the loss curve on the command line.
+
+---
+
+## Stage 6: Effect validation (optional, if user chose full flow)
+
+Load the fine-tuned checkpoint for inference validation.
+
+**Up front — before asking the user for test questions — tell them about interactive self-testing.** Explain that interactive `chat` can't be driven in this agent environment, so validation here uses a non-interactive script, but if they want to test interactively themselves they can run the following in their own terminal (mention this once at the very start of this stage, before step 2's question, and again in the closing summary after validation finishes):
+```bash
+llamafactory-cli chat llamafactory_runs/<run_id>/infer.yaml
+```
+
+1. Generate an inference config into the **artifacts directory** (`./llamafactory_runs/<run_id>/infer.yaml`). **The config differs by the `finetuning_type` chosen in Stage 0 — branch accordingly:**
+
+   - **LoRA / QLoRA** (the training `output_dir` is a LoRA *adapter*, not a full model) — base it on `examples/inference/qwen3_lora_sft.yaml` and use `adapter_name_or_path`:
+     ```yaml
+     model_name_or_path: <base_model>
+     adapter_name_or_path: <output_dir>   # LoRA adapter path (the training output_dir)
+     template: <template>                 # MUST match the template used at training time
+     infer_backend: huggingface
+     trust_remote_code: true
+     ```
+   - **Full** (the training `output_dir` is already a complete set of model weights — there is NO adapter) — base it on `examples/inference/qwen3_full_sft.yaml` and point `model_name_or_path` directly at the `output_dir`, with **NO `adapter_name_or_path` field at all**:
+     ```yaml
+     model_name_or_path: <output_dir>     # the full fine-tuned model (training output_dir)
+     template: <template>                 # MUST match the template used at training time
+     infer_backend: huggingface
+     trust_remote_code: true
+     ```
+     Adding an `adapter_name_or_path` for a Full run is wrong (there is no adapter to load) and will fail or silently load nothing.
+
+   - **Keep `template` identical to the training config.** A mismatch (e.g. trained with `qwen3_nothink` but inferred with `qwen3`) activates think / tool-call special tokens the model never saw in training and produces garbled output.
+   - **Show the inference config as a table and get user confirmation before running.** Just like the training yaml, after generating `infer.yaml` print its key parameters as a table (for LoRA/QLoRA: `model_name_or_path`, `adapter_name_or_path`, `template`, `infer_backend`, `trust_remote_code`; for Full: `model_name_or_path`, `template`, `infer_backend`, `trust_remote_code`) and use `AskUserQuestion` to let the user confirm (or request changes) before you run any inference. Do not start validation until the user confirms.
+2. Ask the user to provide **test text** (and test image path(s) for multimodal models).
+3. **Use a non-interactive batch inference script** (preferred). `llamafactory-cli chat` is an **interactive** REPL and cannot be driven in this agent / background environment, so validation is done with a short script (write it into the artifacts dir) that loads `ChatModel` with the **same args as `infer.yaml`** (so it must follow the same LoRA/QLoRA-vs-Full branching — include `adapter_name_or_path` only for LoRA/QLoRA, never for Full) and feeds the test prompts, e.g. for a LoRA/QLoRA run:
+   ```python
+   from llamafactory.chat import ChatModel
+   chat = ChatModel({
+       "model_name_or_path": "<base_model>",
+       "adapter_name_or_path": "<output_dir>",  # LoRA/QLoRA only — OMIT this key for Full
+       "template": "<template>",
+       "infer_backend": "huggingface",
+       "trust_remote_code": True,
+   })
+   for q in ["你是谁？", "Who are you?"]:
+       print(q, "->", chat.chat([{"role": "user", "content": q}])[0].response_text)
+   ```
+   For a **Full** run, drop the `adapter_name_or_path` key and set `"model_name_or_path": "<output_dir>"`.
+   Running this in the foreground is fine (only training really needs background); **run it on the GPU(s) chosen in Stage 2.5** by prefixing the command with the visibility env var, e.g. `HIP_VISIBLE_DEVICES=<idx> python <script>` (or `CUDA_VISIBLE_DEVICES=<idx>` on NVIDIA); just **report the outputs to the user** when it finishes.
+   - **After validation finishes, state the interactive self-test command again** (the same `llamafactory-cli chat llamafactory_runs/<run_id>/infer.yaml` shown at the start of this stage), so the user has it both before and after validation.
+4. Multimodal: pass images together with the prompt and show the model's answer.
+
+---
+
+## Stage 7: Model export (optional)
+
+If the user wants to export, **the meaning of "export" depends on the `finetuning_type` chosen in Stage 0 — branch accordingly:**
+
+- **LoRA / QLoRA** — export *merges* the LoRA adapter into the base model and writes a standalone full model. This is the classic `merge_lora` flow.
+- **Full** — there is **no adapter and nothing to merge**; the training `output_dir` is already a complete model. "Export" here just re-saves / tidies those full weights (plus tokenizer, generation config, etc.) into a clean user-specified directory. Do NOT describe this as "merging".
+
+1. Ask for the **export name/directory** (user-specified). This is where the exported model weights go (e.g. under `./models/...`) — it is the model output, separate from the artifacts directory.
+   - For **LoRA / QLoRA**, **prefer a `merged` suffix** in the recommended name (e.g. `./models/<base_model>-merged`, or `<base_model>-<identity>-merged`) so the merged output is clearly distinguishable from the base weights.
+   - For **Full**, `merged` is misleading (nothing was merged) — recommend a plain descriptive suffix instead (e.g. `./models/<base_model>-sft`, or `<base_model>-<identity>`).
+   - The user can always override.
+2. Generate an export config into the **artifacts directory** (`./llamafactory_runs/<run_id>/export.yaml`), based on `examples/merge_lora/*.yaml`. **The config differs by finetuning_type:**
+
+   - **LoRA / QLoRA** — point `model_name_or_path` at the base model and `adapter_name_or_path` at the adapter (`output_dir`):
+     ```yaml
+     model_name_or_path: <base_model>
+     adapter_name_or_path: <output_dir>
+     template: <template>                 # MUST match the training template
+     trust_remote_code: true
+     export_dir: <user-specified name>
+     export_size: 5
+     export_device: auto                  # auto = use the GPU chosen in Stage 2.5; cpu is much slower
+     export_legacy_format: false
+     ```
+     Note: when merging LoRA/QLoRA, do **not** load a quantized model or set `quantization_bit` — merging into a quantized base produces a broken model. Merge into the full-precision base, then quantize separately if needed.
+   - **Full** — point `model_name_or_path` directly at the `output_dir` and use **NO `adapter_name_or_path`**:
+     ```yaml
+     model_name_or_path: <output_dir>     # the full fine-tuned model (training output_dir)
+     template: <template>                 # MUST match the training template
+     trust_remote_code: true
+     export_dir: <user-specified name>
+     export_size: 5
+     export_device: auto                  # auto = use the GPU chosen in Stage 2.5; cpu is much slower
+     export_legacy_format: false
+     ```
+   - **`export_device` controls whether the merge/save runs on GPU or CPU.** `export_device: cpu` does the whole merge on CPU — it does **not** use the GPU even if you set a visibility env var, and for some models it can be very slow or appear to stall. To actually use the GPU chosen in Stage 2.5, set **`export_device: auto`** (it will place the model on the visible GPU). Prefer `auto` when a GPU is free; only fall back to `cpu` when no GPU is available (and warn the user it will be slow).
+   - **Show the export config as a table and get user confirmation before running.** Just like the training yaml, after generating `export.yaml` print its key parameters as a table (for LoRA/QLoRA include `adapter_name_or_path`; for Full omit it — list `model_name_or_path`, `template`, `export_dir`, `export_size`, `export_device`, `export_legacy_format`, etc.) and use `AskUserQuestion` to let the user confirm (or request changes) before you run the export. Do not start the export until the user confirms.
+3. Run (export is usually quick — foreground is fine; only training really needs background). **Run it on the GPU(s) chosen in Stage 2.5** by prefixing with the visibility env var, and tee the output into the artifacts dir for the record:
+   ```bash
+   # HIP_VISIBLE_DEVICES for AMD, CUDA_VISIBLE_DEVICES for NVIDIA; pair with export_device: auto to use the GPU.
+   HIP_VISIBLE_DEVICES=<idx> llamafactory-cli export llamafactory_runs/<run_id>/export.yaml 2>&1 | tee llamafactory_runs/<run_id>/export.log
+   ```
+4. When export completes, tell the user the final model path.
+
+---
+
+## Final summary (end of the whole workflow)
+
+When all chosen stages are finished, give a concise closing summary and then **stop**. The summary may include ONLY:
+
+- the produced artifacts (config/log dir, LoRA adapter, loss curve image path, merged-model export dir);
+- the training result (steps / epochs / final loss / runtime);
+- the validation result (if validation ran);
+- the export directory (if export ran);
+- the ready-to-run command for interactive self-testing (e.g. `llamafactory-cli chat ...`).
+
+**Do NOT propose or offer any follow-up work** — no suggestions about quantization, GGUF conversion, deployment, serving as an API, further training, or anything similar. Do not end with questions like "需要我帮你……吗？". Simply describe what was produced and stop.
+
+---
+
+## General requirements
+
+- Ask first with `AskUserQuestion` at every irreversible or ambiguous decision; do not assume.
+- **Confirm before every yaml write — no exceptions.** Before writing ANY yaml config to disk (train / inference / export), you MUST first print its key parameters as a table AND call `AskUserQuestion` to get the user's explicit confirmation. Only write the file after the user confirms. This applies to every config in every stage, not just training — do not "save time" by writing infer.yaml or export.yaml directly. If the user requests changes, update the table and re-confirm before writing.
+- Run **long-running** commands (model download, training) **in the background + report periodically** (interval ≤ 100 seconds). Inference validation and export are usually quick and do NOT need background — run them in the foreground.
+- **Detect GPU status once before running and pin the free device(s)** (see Stage 2.5). Detect a single time at the start, pick the device(s) — if multiple are free, ask single-vs-multi — then reuse the same index/indices for SFT / validation / export via the platform's visibility env var (`HIP_VISIBLE_DEVICES` for AMD, `CUDA_VISIBLE_DEVICES` for NVIDIA). Do NOT re-detect before each run. Never silently launch onto a busy card — if none are free, ask the user. For export to actually use the GPU, also set `export_device: auto` (not `cpu`).
+- **Keep generated yaml configs and logs in the artifacts directory** (`./llamafactory_runs/<run_id>/`); do not write them into `./models/` (downloaded weights) or the training `output_dir` (saved checkpoints).
+- **Prefer official-example defaults; track every deviation.** When generating the training yaml, start from the official example and change as little as possible. List any changed field in a "diff vs. default" table with a concrete reason (user-specified, model/dataset adaptation, etc.).
+- **Keep `template` consistent across train / inference / export.** Use the template from the official example unless there's a tracked reason to change it; a train/infer mismatch produces garbled output.
+- Never destroy the user's original data files in place; do not pollute the repo's built-in `data/*.json` — prefer copies.
+- Keep `dataset_info.json` valid JSON after edits.

.github/workflows/docker.yml

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

.github/workflows/docker_npu.yml

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

.pre-commit-config.yaml

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

Makefile

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

README.md

@@ -19,7 +19,21 @@
 [![Open in Studios](https://img.shields.io/badge/ModelScope-Open%20in%20Studios-blue)](https://modelscope.cn/studios/hiyouga/LLaMA-Board)
 [![Open in Novita](https://img.shields.io/badge/Novita-Deploy%20Template-blue)](https://novita.ai/templates-library/105981?sharer=88115474-394e-4bda-968e-b88e123d0c47)
 
-### Used by [Amazon](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/), [NVIDIA](https://developer.nvidia.com/rtx/ai-toolkit), [Aliyun](https://help.aliyun.com/zh/pai/use-cases/fine-tune-a-llama-3-model-with-llama-factory), etc.
+### Used by [Amazon](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/), [NVIDIA](https://build.nvidia.com/spark/llama-factory), [Aliyun](https://help.aliyun.com/zh/pai/use-cases/fine-tune-a-llama-3-model-with-llama-factory), etc.
+
+----
+
+<div align="center" markdown="1">
+
+### Check our new open-source project —<br>🐧 [PenguinHarness](https://github.com/Prism-Shadow/penguin-harness): Your desktop agent that automatically builds agents for just $0.02 of tokens!
+
+Follow our project: https://github.com/Prism-Shadow/penguin-harness
+
+</div>
+
+https://github.com/user-attachments/assets/9b7033e8-f08a-4c3f-bd33-547896664e6e
+
+----
 
 <div align="center" markdown="1">
 
@@ -32,7 +46,7 @@
 
 ### Easily fine-tune 100+ large language models with zero-code [CLI](#quickstart) and [Web UI](#fine-tuning-with-llama-board-gui-powered-by-gradio)
 
-![GitHub Trend](https://trendshift.io/api/badge/repositories/4535)
+![GitHub Trend](https://trendshift.io/api/badge/repositories/17371)
 
 </div>
 
@@ -50,10 +64,12 @@ Start local training:
 Start cloud training:
 - **Colab (free)**: https://colab.research.google.com/drive/1eRTPn37ltBbYsISy9Aw2NuI2Aq5CQrD9?usp=sharing
 - **PAI-DSW (free trial)**: https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory
+- **AMD GPU Cloud (free credits)**: https://github.com/AMD-AIM/AMD_Developers_Notebooks/blob/main/en/AMD_developer_LLaMAFactory_note_en.md
 
 Read technical notes:
 - **Documentation (WIP)**: https://llamafactory.readthedocs.io/en/latest/
 - **Documentation (AMD GPU)**: https://rocm.docs.amd.com/projects/ai-developer-hub/en/latest/notebooks/fine_tune/llama_factory_llama3.html
+- **Documentation (ASCEND NPU)**: https://llamafactory.readthedocs.io/en/latest/multibackend/npu/index.html
 - **Official Blog**: https://blog.llamafactory.net/en/
 
 > [!NOTE]
@@ -111,6 +127,8 @@ Read technical notes:
 
 - 💡 [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)
 
@@ -297,8 +315,9 @@ Read technical notes:
 | [LLaVA-NeXT](https://huggingface.co/llava-hf)                     | 7B/8B/13B/34B/72B/110B           | llava_next           |
 | [LLaVA-NeXT-Video](https://huggingface.co/llava-hf)               | 7B/34B                           | llava_next_video     |
 | [MiMo](https://huggingface.co/XiaomiMiMo)                         | 7B/309B                          | mimo/mimo_v2         |
-| [MiniCPM 4](https://huggingface.co/openbmb)                       | 0.5B/8B                          | cpm4                 |
+| [MiniCPM 4/5](https://huggingface.co/openbmb)                     | 0.5B/1B/8B                       | cpm4/empty           |
 | [MiniCPM-o/MiniCPM-V 4.5](https://huggingface.co/openbmb)         | 8B/9B                            | minicpm_o/minicpm_v  |
+| [MiniCPM-V 4.6](https://huggingface.co/openbmb)                   | 3B/8B                            | minicpm_v_4_6        |
 | [MiniMax-M1/MiniMax-M2](https://huggingface.co/MiniMaxAI/models)  | 229B/456B                        | minimax1/minimax2    |
 | [Ministral 3](https://huggingface.co/mistralai)                   | 3B/8B/14B                        | ministral3           |
 | [Mistral/Mixtral](https://huggingface.co/mistralai)               | 7B/8x7B/8x22B                    | mistral              |
@@ -554,7 +573,24 @@ Try `dataloader_num_workers: 0` if you encounter `Can't pickle local object` err
 
 #### Install BitsAndBytes
 
-If you want to enable the quantized LoRA (QLoRA) on the Windows platform, you need to install a pre-built version of `bitsandbytes` library, which supports CUDA 11.1 to 12.2, please select the appropriate [release version](https://github.com/jllllll/bitsandbytes-windows-webui/releases/tag/wheels) based on your CUDA version.
+To enable Quantized LoRA (QLoRA) on Windows, you need to install bitsandbytes.
+
+For most users, it is recommended to install the latest official release:
+
+```bash
+pip install bitsandbytes
+```
+
+If you are using uv to manage your virtual environment, it is recommended to install bitsandbytes after installing the GPU-enabled version of PyTorch:
+
+```bash
+uv pip install bitsandbytes --no-deps
+```
+
+[!IMPORTANT]
+Pay attention to the CUDA Toolkit version when installing bitsandbytes. Official bitsandbytes releases are built for specific CUDA Toolkit versions. On Windows x86-64, separate builds are currently provided for CUDA 11.8–12.6 and CUDA 12.8–12.9. Support for NVIDIA RTX 50 Series GPUs (e.g., RTX 5060 Ti, sm_120) requires the CUDA 12.8–12.9 builds.
+
+If your environment uses an older CUDA version, or you need compatibility with older Windows / PyTorch combinations, you can install the third-party precompiled Windows wheel:
 
 ```bash
 pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.41.2.post2-py3-none-win_amd64.whl
@@ -568,19 +604,23 @@ To enable FlashAttention-2 on the Windows platform, please use the script from [
 
 <details><summary>For Ascend NPU users</summary>
 
-To install LLaMA Factory on Ascend NPU devices, please upgrade Python to version 3.10 or higher: `pip install -r requirements/npu.txt`. Additionally, you need to install the **Ascend CANN Toolkit and Kernels**. Please follow the [installation tutorial](https://llamafactory.readthedocs.io/en/latest/advanced/npu_installation.html).
+To install LLaMA Factory on Ascend NPU devices, please upgrade Python to version 3.10 or higher: `pip install -r requirements/npu.txt`. Additionally, you need to install the **Ascend CANN Toolkit and Kernels**. Please follow the [installation tutorial](https://llamafactory.readthedocs.io/en/latest/multibackend/npu/npu_installation.html).
 
 
 You can also download the pre-built Docker images:
 
 ```bash
 # Docker Hub
-docker pull hiyouga/llamafactory:latest-npu-a2
-docker pull hiyouga/llamafactory:latest-npu-a3
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-ubuntu-py3.11
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-openeuler-py3.11
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-openeuler-py3.11
 
 # quay.io
-docker pull quay.io/ascend/llamafactory:latest-npu-a2
-docker pull quay.io/ascend/llamafactory:latest-npu-a3
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-ubuntu-py3.11
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-openeuler-py3.11
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-openeuler-py3.11
 ```
 
 #### Install BitsAndBytes
@@ -661,12 +701,28 @@ docker compose up -d
 docker compose exec llamafactory bash
 ```
 
-For Ascend NPU users:
+For Ascend NPU users (A2 with Ubuntu by default):
 
 ```bash
 cd docker/docker-npu/
-docker compose up -d
-docker compose exec llamafactory bash
+docker compose up -d llamafactory-a2-ubuntu
+docker compose exec llamafactory-a2-ubuntu bash
+```
+
+Other NPU variants can be started with their corresponding profiles and services:
+
+```bash
+# A3 with Ubuntu
+docker compose --profile a3 up -d llamafactory-a3-ubuntu
+docker compose exec llamafactory-a3-ubuntu bash
+
+# A2 with openEuler
+docker compose --profile openeuler up -d llamafactory-a2-openeuler
+docker compose exec llamafactory-a2-openeuler bash
+
+# A3 with openEuler
+docker compose --profile a3-openeuler up -d llamafactory-a3-openeuler
+docker compose exec llamafactory-a3-openeuler bash
 ```
 
 For AMD ROCm users:
@@ -824,7 +880,7 @@ If you have a project that should be incorporated, please contact via email or c
 1. Choi et al. FACT-GPT: Fact-Checking Augmentation via Claim Matching with LLMs. 2024. [[arxiv]](https://arxiv.org/abs/2402.05904)
 1. Zhang et al. AutoMathText: Autonomous Data Selection with Language Models for Mathematical Texts. 2024. [[arxiv]](https://arxiv.org/abs/2402.07625)
 1. Lyu et al. KnowTuning: Knowledge-aware Fine-tuning for Large Language Models. 2024. [[arxiv]](https://arxiv.org/abs/2402.11176)
-1. Yang et al. LaCo: Large Language Model Pruning via Layer Collaps. 2024. [[arxiv]](https://arxiv.org/abs/2402.11187)
+1. Yang et al. LaCo: Large Language Model Pruning via Layer Collapse. 2024. [[arxiv]](https://arxiv.org/abs/2402.11187)
 1. Bhardwaj et al. Language Models are Homer Simpson! Safety Re-Alignment of Fine-tuned Language Models through Task Arithmetic. 2024. [[arxiv]](https://arxiv.org/abs/2402.11746)
 1. Yang et al. Enhancing Empathetic Response Generation by Augmenting LLMs with Small-scale Empathetic Models. 2024. [[arxiv]](https://arxiv.org/abs/2402.11801)
 1. Yi et al. Generation Meets Verification: Accelerating Large Language Model Inference with Smart Parallel Auto-Correct Decoding. ACL 2024 Findings. [[arxiv]](https://arxiv.org/abs/2402.11809)
@@ -939,7 +995,3 @@ If this work is helpful, please kindly cite as:
 ## Acknowledgement
 
 This repo benefits from [PEFT](https://github.com/huggingface/peft), [TRL](https://github.com/huggingface/trl), [QLoRA](https://github.com/artidoro/qlora) and [FastChat](https://github.com/lm-sys/FastChat). Thanks for their wonderful works.
-
-## Star History
-
-![Star History Chart](https://api.star-history.com/svg?repos=hiyouga/LLaMA-Factory&type=Date)

README_zh.md

@@ -19,7 +19,21 @@
 [![Open in Studios](https://img.shields.io/badge/ModelScope-Open%20in%20Studios-blue)](https://modelscope.cn/studios/hiyouga/LLaMA-Board)
 [![Open in Novita](https://img.shields.io/badge/Novita-Deploy%20Template-blue)](https://novita.ai/templates-library/105981?sharer=88115474-394e-4bda-968e-b88e123d0c47)
 
-### 获得[亚马逊](https://aws.amazon.com/cn/blogs/china/a-one-stop-code-free-model-fine-tuning-deployment-platform-based-on-sagemaker-and-llama-factory/)、[英伟达](https://developer.nvidia.cn/rtx/ai-toolkit)、[阿里云](https://help.aliyun.com/zh/pai/use-cases/fine-tune-a-llama-3-model-with-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/)、[英伟达](https://build.nvidia.com/spark/llama-factory)、[阿里云](https://help.aliyun.com/zh/pai/use-cases/fine-tune-a-llama-3-model-with-llama-factory)等的应用。
+
+----
+
+<div align="center" markdown="1">
+
+### 欢迎关注我们全新的开源项目——<br>🐧 [PenguinHarness](https://github.com/Prism-Shadow/penguin-harness)：只需 0.2 元的 Token，即可自动构建 Agent 的桌面级 Agent！
+
+点击关注项目：https://github.com/Prism-Shadow/penguin-harness
+
+</div>
+
+https://github.com/user-attachments/assets/604eb626-0a5d-4a62-87e3-14ebade1cd5f
+
+----
 
 <div align="center" markdown="1">
 
@@ -32,7 +46,7 @@
 
 ### 使用零代码[命令行](#快速开始)与 [Web UI](#llama-board-可视化微调由-gradio-驱动) 轻松微调百余种大模型
 
-![GitHub Trend](https://trendshift.io/api/badge/repositories/4535)
+![GitHub Trend](https://trendshift.io/api/badge/repositories/17371)
 
 </div>
 
@@ -50,12 +64,13 @@ https://github.com/user-attachments/assets/43b700c6-a178-41db-b1f8-8190a5d3fcfc
 开始云端训练：
 - **Colab（免费）**：https://colab.research.google.com/drive/1d5KQtbemerlSDSxZIfAaWXhKr30QypiK?usp=sharing
 - **PAI-DSW（免费试用）**：https://gallery.pai-ml.com/#/preview/deepLearning/nlp/llama_factory
+- **AMD GPU Cloud (免费试用)**: https://github.com/AMD-AIM/AMD_Developers_Notebooks/blob/main/zh/AMD_developer_LLaMAFactory_note_zh.md
 
 阅读技术文档：
 - **入门教程**：https://zhuanlan.zhihu.com/p/695287607
 - **微调视频教程**：https://www.bilibili.com/video/BV1djgRzxEts/
 - **框架文档**：https://llamafactory.readthedocs.io/zh-cn/latest/
-- **框架文档（昇腾 NPU）**：https://ascend.github.io/docs/sources/llamafactory/
+- **框架文档（昇腾 NPU）**：https://llamafactory.readthedocs.io/zh-cn/latest/multibackend/npu/index.html
 - **官方博客**：https://blog.llamafactory.net/
 
 > [!NOTE]
@@ -113,6 +128,8 @@ https://github.com/user-attachments/assets/43b700c6-a178-41db-b1f8-8190a5d3fcfc
 
 - 💡 [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/)（英文）
 
@@ -299,8 +316,9 @@ https://github.com/user-attachments/assets/43b700c6-a178-41db-b1f8-8190a5d3fcfc
 | [LLaVA-NeXT](https://huggingface.co/llava-hf)                     | 7B/8B/13B/34B/72B/110B           | llava_next           |
 | [LLaVA-NeXT-Video](https://huggingface.co/llava-hf)               | 7B/34B                           | llava_next_video     |
 | [MiMo](https://huggingface.co/XiaomiMiMo)                         | 7B/309B                          | mimo/mimo_v2         |
-| [MiniCPM 4](https://huggingface.co/openbmb)                       | 0.5B/8B                          | cpm4                 |
+| [MiniCPM 4/5](https://huggingface.co/openbmb)                     | 0.5B/1B/8B                       | cpm4/empty           |
 | [MiniCPM-o/MiniCPM-V 4.5](https://huggingface.co/openbmb)         | 8B/9B                            | minicpm_o/minicpm_v  |
+| [MiniCPM-V 4.6](https://huggingface.co/openbmb)                   | 3B/8B                            | minicpm_v_4_6        |
 | [MiniMax-M1/MiniMax-M2](https://huggingface.co/MiniMaxAI/models)  | 229B/456B                        | minimax1/minimax2    |
 | [Ministral 3](https://huggingface.co/mistralai)                   | 3B/8B/14B                        | ministral3           |
 | [Mistral/Mixtral](https://huggingface.co/mistralai)               | 7B/8x7B/8x22B                    | mistral              |
@@ -556,7 +574,24 @@ python -c "import torch; print(torch.cuda.is_available())"
 
 #### 安装 BitsAndBytes
 
-如果要在 Windows 平台上开启量化 LoRA（QLoRA），需要安装预编译的 `bitsandbytes` 库, 支持 CUDA 11.1 到 12.2, 请根据您的 CUDA 版本情况选择适合的[发布版本](https://github.com/jllllll/bitsandbytes-windows-webui/releases/tag/wheels)。
+如果要在 Windows 平台上开启量化 LoRA（QLoRA），需要安装 bitsandbytes。
+
+对于大多数用户，建议优先使用官方发布的最新版本：
+
+```bash
+pip install bitsandbytes
+```
+
+如果使用 uv 管理虚拟环境，建议在安装好 GPU 版本 PyTorch 之后，再执行：
+
+```bash
+uv pip install bitsandbytes --no-deps
+```
+
+> [!IMPORTANT]
+> 安装 bitsandbytes 时，请注意 CUDA Toolkit 版本。bitsandbytes 的官方发布包是按不同 CUDA 版本分别构建的，Windows x86-64 目前提供了面向 CUDA 11.8–12.6 和 CUDA 12.8–12.9 的不同构建；其中支持 RTX 50 系列（如 RTX 5060 Ti，sm_120）的构建对应 CUDA 12.8–12.9。
+
+若当前环境的 CUDA 版本较旧，或者需要兼容较老的 Windows / PyTorch 组合，可以使用第三方预编译版本：
 
 ```bash
 pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/download/wheels/bitsandbytes-0.41.2.post2-py3-none-win_amd64.whl
@@ -570,18 +605,22 @@ pip install https://github.com/jllllll/bitsandbytes-windows-webui/releases/downl
 
 <details><summary>昇腾 NPU 用户指南</summary>
 
-在昇腾 NPU 设备上安装 LLaMA Factory 时，请升级 Python 到 3.10 及以上，并需要指定额外依赖项，使用 `pip install -r requirements/npu.txt` 命令安装。此外，还需要安装 **Ascend CANN Toolkit 与 Kernels**，安装方法请参考[安装教程](https://llamafactory.readthedocs.io/zh-cn/latest/advanced/npu_installation.html)。
+在昇腾 NPU 设备上安装 LLaMA Factory 时，请升级 Python 到 3.10 及以上，并需要指定额外依赖项，使用 `pip install -r requirements/npu.txt` 命令安装。此外，还需要安装 **Ascend CANN Toolkit 与 Kernels**，安装方法请参考[安装教程](https://llamafactory.readthedocs.io/zh-cn/latest/multibackend/npu/npu_installation.html)。
 
 您可以直接下载预安装的最新docker镜像：
 
 ```bash
 # Docker Hub
-docker pull hiyouga/llamafactory:latest-npu-a2
-docker pull hiyouga/llamafactory:latest-npu-a3
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-ubuntu-py3.11
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-openeuler-py3.11
+docker pull hiyouga/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-openeuler-py3.11
 
 # quay.io
-docker pull quay.io/ascend/llamafactory:latest-npu-a2
-docker pull quay.io/ascend/llamafactory:latest-npu-a3
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-ubuntu-py3.11
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-ubuntu-py3.11
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A2-openeuler-py3.11
+docker pull quay.io/ascend/llamafactory:latest-cann9.0.0-torch_npu2.7.1-A3-openeuler-py3.11
 ```
 
 #### 安装 BitsAndBytes
@@ -662,12 +701,28 @@ docker compose up -d
 docker compose exec llamafactory bash
 ```
 
-昇腾 NPU 用户：
+昇腾 NPU 用户（默认使用 A2 和 Ubuntu）：
 
 ```bash
 cd docker/docker-npu/
-docker compose up -d
-docker compose exec llamafactory bash
+docker compose up -d llamafactory-a2-ubuntu
+docker compose exec llamafactory-a2-ubuntu bash
+```
+
+其他 NPU 组合可以通过对应的 profile 和服务启动：
+
+```bash
+# A3 + Ubuntu
+docker compose --profile a3 up -d llamafactory-a3-ubuntu
+docker compose exec llamafactory-a3-ubuntu bash
+
+# A2 + openEuler
+docker compose --profile openeuler up -d llamafactory-a2-openeuler
+docker compose exec llamafactory-a2-openeuler bash
+
+# A3 + openEuler
+docker compose --profile a3-openeuler up -d llamafactory-a3-openeuler
+docker compose exec llamafactory-a3-openeuler bash
 ```
 
 AMD ROCm 用户：
@@ -828,7 +883,7 @@ swanlab_run_name: test_run # 可选
 1. Choi et al. FACT-GPT: Fact-Checking Augmentation via Claim Matching with LLMs. 2024. [[arxiv]](https://arxiv.org/abs/2402.05904)
 1. Zhang et al. AutoMathText: Autonomous Data Selection with Language Models for Mathematical Texts. 2024. [[arxiv]](https://arxiv.org/abs/2402.07625)
 1. Lyu et al. KnowTuning: Knowledge-aware Fine-tuning for Large Language Models. 2024. [[arxiv]](https://arxiv.org/abs/2402.11176)
-1. Yang et al. LaCo: Large Language Model Pruning via Layer Collaps. 2024. [[arxiv]](https://arxiv.org/abs/2402.11187)
+1. Yang et al. LaCo: Large Language Model Pruning via Layer Collapse. 2024. [[arxiv]](https://arxiv.org/abs/2402.11187)
 1. Bhardwaj et al. Language Models are Homer Simpson! Safety Re-Alignment of Fine-tuned Language Models through Task Arithmetic. 2024. [[arxiv]](https://arxiv.org/abs/2402.11746)
 1. Yang et al. Enhancing Empathetic Response Generation by Augmenting LLMs with Small-scale Empathetic Models. 2024. [[arxiv]](https://arxiv.org/abs/2402.11801)
 1. Yi et al. Generation Meets Verification: Accelerating Large Language Model Inference with Smart Parallel Auto-Correct Decoding. ACL 2024 Findings. [[arxiv]](https://arxiv.org/abs/2402.11809)
@@ -942,7 +997,3 @@ swanlab_run_name: test_run # 可选
 ## 致谢
 
 本项目受益于 [PEFT](https://github.com/huggingface/peft)、[TRL](https://github.com/huggingface/trl)、[QLoRA](https://github.com/artidoro/qlora) 和 [FastChat](https://github.com/lm-sys/FastChat)，感谢以上诸位作者的付出。
-
-## Star History
-
-![Star History Chart](https://api.star-history.com/svg?repos=hiyouga/LLaMA-Factory&type=Date)

docker/docker-npu/Dockerfile

@@ -36,6 +36,7 @@ COPY . /app
 RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh
 RUN pip uninstall -y torch torchvision torchaudio
 RUN pip install --no-cache-dir -r requirements/npu.txt --index-url "${PYTORCH_INDEX}"
+RUN pip install --no-cache-dir -r requirements/triton_ascend.txt
 RUN pip install --no-cache-dir -r requirements/deepspeed.txt
 RUN pip install --no-cache-dir -e . --no-build-isolation && \
     pip install --no-cache-dir -r requirements/metrics.txt --no-build-isolation

docker/docker-npu/OVERVIEW.md

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

docker/docker-npu/OVERVIEW.zh.md

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

docker/docker-npu/docker-compose.yml

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

docs/zh/getting-started.md

@@ -1,16 +1,15 @@
-# Getting Started
+# 快速开始
 
+LLaMA Factory 是一个高效、灵活的大模型微调框架，支持 100+ 种主流大语言模型的微调训练。本文档将帮助您快速上手使用 LLaMA Factory。
 
-## 训练方法
+## 支持的训练方法
 
 |          方法          |     全参数训练      |    部分参数训练     |       LoRA         |       QLoRA        |
 |:---------------------:| ------------------ | ------------------ | ------------------ | ------------------ |
-|      指令监督微调       | :white_check_mark: |  |  | |
-|      奖励模型训练       |  |  |  | |
-|        DPO 训练        |  |  |  | |
-
-
+|   指令监督微调 (SFT)   | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
+|      DPO 训练          | :white_check_mark: | :white_check_mark: | :white_check_mark: | :white_check_mark: |
 
+> **提示**: v1 版本目前支持 SFT 和 DPO 两种训练方法，均支持多种加速特性，包括 DeepSpeed、FSDP、FlashAttention-2 等。
 
 ## 软件依赖
 
@@ -32,40 +31,131 @@
 |   flash-attn(NVIDIA GPU)   | 2.5.6  | 2.7.2  |
 
 
-## 如何使用
-
-### 安装 LLaMA Factory
+## 安装 LLaMA Factory
 
 > [!IMPORTANT]
-> 此步骤为必需。
+> 此步骤为必需。请确保您的环境满足上述软件依赖要求。
 
-#### 从源码安装
+### 从源码安装（推荐）
 
 ```bash
-git clone --depth 1 https://github.com/hiyouga/LlamaFactory.git
-cd LlamaFactory
+git clone --depth 1 https://github.com/hiyouga/LLaMA-Factory.git
+cd LLaMA-Factory
 pip install -e .
 ```
 
+### 使用 pip 安装
 
-### 数据准备
+```bash
+pip install llamafactory
+```
 
-关于数据集文件的格式，请参考 [data-preparation/README.md](data-preparation/README.md) 的内容。你可以使用 HuggingFace / ModelScope 上的数据集或加载本地数据集。
+### 可选依赖
+
+如果您需要使用特定的加速特性，可以安装相应的依赖：
+
+```bash
+# 安装 FlashAttention-2 支持
+pip install flash-attn --no-build-isolation
+
+# 安装 DeepSpeed 支持
+pip install deepspeed
+
+# 安装 Unsloth 支持（用于加速 LoRA 训练）
+pip install unsloth
+```
+
+## 数据准备
+
+LLaMA Factory 支持多种数据格式，包括 JSON、JSONL、CSV 等。关于数据集文件的详细格式说明，请参考 [数据准备指南](../../data/README_zh.md)。
+
+### 使用内置数据集
+
+LLaMA Factory 提供了多个内置数据集用于快速测试，您可以在 `data/dataset_info.json` 中查看所有可用的数据集。
+
+### 使用自定义数据集
+
+您可以使用 HuggingFace / ModelScope 上的数据集或加载本地数据集。
 
 > [!NOTE]
-> 使用自定义数据集或自定义数据集格式时，请参照 [data-preparation/README.md](data-preparation/README.md) 进行配置，如有必要，请重新实现自定义数据集的数据处理逻辑，包括对应的`converter`。
+> 使用自定义数据集或自定义数据集格式时，请参照 [数据准备指南](../../data/README_zh.md) 进行配置。如有必要，请重新实现自定义数据集的数据处理逻辑，包括对应的 `converter`。
 
-您也可以使用 **[Easy Dataset](https://github.com/ConardLi/easy-dataset)**、**[DataFlow](https://github.com/OpenDCAI/DataFlow)** 和 **[GraphGen](https://github.com/open-sciencelab/GraphGen)** 构建用于微调的合成数据。
+### 数据构建工具
 
-### 快速开始
+您也可以使用以下工具构建用于微调的合成数据：
+- **[Easy Dataset](https://github.com/ConardLi/easy-dataset)** - 易于使用的数据集构建工具
+- **[DataFlow](https://github.com/OpenDCAI/DataFlow)** - 高质量数据准备管道
+- **[GraphGen](https://github.com/open-sciencelab/GraphGen)** - 基于图的数据生成工具
 
-下面的命令展示了对 Qwen3-0.6B 模型使用 FSDP2 进行 全参**微调**，两行命令等价。
+## 快速开始
+
+### 命令行训练
+
+下面的命令展示了对 Qwen3-0.6B 模型使用 FSDP2 进行全参数微调：
 
 ```bash
 export USE_V1=1
 llamafactory-cli sft examples/v1/train_full/train_full_fsdp2.yaml
-llamafactory-cli train examples/v1/train_full/train_full_fsdp2.yaml
-
 ```
 
-高级用法请参考 [advanced](./advanced/README.md)（包括多卡多机微调、分布式、Lora、量化、以及各种加速特性等）。
+> **提示**: `llamafactory-cli sft` 和 `llamafactory-cli train` 命令等价。
+
+### Web UI 训练
+
+LLaMA Factory 提供了直观的 Web 界面（LLaMA Board），您可以通过图形界面进行训练：
+
+```bash
+llamafactory-cli webui
+```
+
+在浏览器中打开 http://localhost:7860 即可开始使用。
+
+### 推理部署
+
+训练完成后，您可以使用以下命令部署模型：
+
+```bash
+# 使用 vLLM 后端进行高性能推理
+llamafactory-cli chat --model_name_or_path path/to/your/model --template qwen --infer_backend vllm
+
+# 使用 HuggingFace 后端进行推理
+llamafactory-cli chat --model_name_or_path path/to/your/model --template qwen
+```
+
+## 进阶用法
+
+高级用法请参考 [进阶指南](./advanced/README.md)，包括：
+- 多卡多机分布式训练
+- LoRA/QLoRA 微调
+- 模型量化（AWQ/GPTQ/LLM.int8 等）
+- 各种加速特性（DeepSpeed、FSDP、FlashAttention 等）
+- 多模态模型微调
+
+## 常见问题
+
+### 1. 内存不足怎么办？
+
+- 使用 LoRA 或 QLoRA 代替全参数训练
+- 减小 `batch_size` 和 `cutoff_len`
+- 启用 `gradient_checkpointing`
+- 使用 DeepSpeed ZeRO-2 或 ZeRO-3
+
+### 2. 如何选择合适的训练方法？
+
+- **SFT（指令微调）**: 最常用的方法，适用于大多数场景，通过监督数据训练模型
+- **DPO（直接偏好优化）**: 用于对齐人类偏好，提升模型输出质量，无需训练奖励模型
+
+### 3. 训练完成后如何评估模型？
+
+```bash
+llamafactory-cli eval --model_name_or_path path/to/your/model --template qwen --dataset mmlu
+```
+
+更多评估方法请参考 [评估指南](./advanced/evaluation.md)。
+
+## 获取帮助
+
+如果您在使用过程中遇到问题：
+- 查看 [GitHub Issues](https://github.com/hiyouga/LLaMA-Factory/issues)
+- 加入 [Discord 社区](https://discord.gg/rKfvV9r9FK)
+- 加入微信群（扫描二维码）

examples/accelerate/fsdp2_config_qwen35_moe.yaml

@@ -0,0 +1,20 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: FSDP
+downcast_bf16: 'no'
+fsdp_config:
+  fsdp_version: 2
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_transformer_layer_cls_to_wrap: Qwen3_5MoeDecoderLayer,Qwen3_5MoeVisionBlock
+  fsdp_cpu_ram_efficient_loading: true
+  fsdp_offload_params: false
+  fsdp_reshard_after_forward: true
+  fsdp_state_dict_type: FULL_STATE_DICT
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8  # Change to match your NPU count (e.g., 8 for A2, 16 for A3)
+rdzv_backend: static
+same_network: true
+use_cpu: false

examples/ascend/qwen3_5moe_lora_sft_fsdp2.yaml

@@ -0,0 +1,51 @@
+# Start FSDP2 full fine-tuning on Ascend NPU
+# Usage:
+#   accelerate launch \
+#       --config_file examples/accelerate/fsdp2_config_qwen35_moe.yaml \
+#       src/train.py examples/ascend/qwen3_5moe_lora_sft_fsdp2.yaml
+#
+# Note: Change `num_processes` in fsdp2_config_qwen35_moe.yaml to match your NPU count
+
+### model
+model_name_or_path: Qwen/Qwen3.5-35B-A3B
+trust_remote_code: true
+use_v1_kernels: false
+flash_attn: fa2
+
+### method
+stage: sft
+do_train: true
+finetuning_type: lora
+lora_rank: 8
+lora_target: all
+
+### dataset
+dataset: alpaca_en_demo
+template: qwen3_5_nothink
+cutoff_len: 2048
+max_samples: 1000
+overwrite_cache: true
+preprocessing_num_workers: 16
+dataloader_num_workers: 4
+packing: false
+
+### output
+output_dir: saves/Qwen3.5-35B/lora/sft
+logging_steps: 1
+save_steps: 2000
+max_steps: 2000
+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
+learning_rate: 1.0e-5
+lr_scheduler_type: cosine
+warmup_ratio: 0.1
+bf16: true
+ddp_timeout: 1800
+resume_from_checkpoint: null
+disable_gradient_checkpointing: true

examples/v1/train_batching_strategy/train_full_fsdp2_batching_normal.yaml

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

examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_batching.yaml

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

examples/v1/train_batching_strategy/train_full_fsdp2_dynamic_padding_free.yaml

@@ -1,11 +1,8 @@
 model: Qwen/Qwen3-0.6B
 model_class: llm
 
-template: qwen3_nothink
-
 kernel_config:
   name: auto
-  include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null
 
 quant_config: null
 
@@ -20,7 +17,7 @@ train_dataset: data/v1_sft_demo.yaml
 output_dir: outputs/test_fsdp2
 micro_batch_size: 4
 batching_strategy: dynamic_padding_free
-flash_attn: flash_attention2
+flash_attn: flash_attention_2
 cutoff_len: 2048
 learning_rate: 1.0e-4
 max_steps: 10

examples/v1/train_batching_strategy/train_full_fsdp2_padding_free.yaml

@@ -1,11 +1,8 @@
 model: Qwen/Qwen3-0.6B
 model_class: llm
 
-template: qwen3_nothink
-
 kernel_config:
   name: auto
-  include_kernels: auto # choice: null/true/false/auto/kernel_id1,kernel_id2,kernel_id3, default is null
 
 quant_config: null
 
@@ -20,7 +17,7 @@ train_dataset: data/v1_sft_demo.yaml
 output_dir: outputs/test_fsdp2
 micro_batch_size: 4
 batching_strategy: padding_free
-flash_attn: flash_attention2
+flash_attn: flash_attention_2
 cutoff_len: 2048
 learning_rate: 1.0e-4
 max_steps: 10

examples/v1/train_freeze/train_freeze_sft.yaml

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

examples/v1/train_full/train_full_deepspeed.yaml

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

examples/v1/train_full/train_full_fsdp2.yaml

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

examples/v1/train_full/train_full_liger_kernel.yaml

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

examples/v1/train_full/train_full_muon.yaml

@@ -0,0 +1,29 @@
+model: Qwen/Qwen3-0.6B
+model_class: llm
+
+
+dist_config:
+  name: fsdp2
+
+optim_config:
+  name: muon
+  wd: 0.1
+  momentum: 0.95
+  nesterov: true
+  ns_steps: 5
+  adamw_betas: [0.9, 0.95]
+  adamw_eps: 1.0e-8
+
+### data
+train_dataset: data/v1_sft_demo.yaml
+
+### training
+output_dir: outputs/test_muon
+micro_batch_size: 1
+cutoff_len: 2048
+learning_rate: 1.0e-5
+max_steps: 10
+
+### sample
+sample_backend: hf
+max_new_tokens: 128

examples/v1/train_full/train_full_ulysses_cp.yaml

@@ -2,14 +2,15 @@ model: Qwen/Qwen3-0.6B
 trust_remote_code: true
 model_class: llm
 
-template: qwen3_nothink
+flash_attn: flash_attention_2
 
 # FSDP Config
 dist_config:
   name: fsdp2
   dcp_path: null
-  cp_mode: ulysses
-  cp_size: 2
+
+cp_mode: ulysses
+cp_size: 2
 
 ### data
 train_dataset: data/v1_sft_demo.yaml

examples/v1/train_lora/train_lora_dpo.yaml

@@ -0,0 +1,33 @@
+model: Qwen/Qwen3-4B
+model_class: llm
+
+# PEFT Configuration
+peft_config:
+  name: lora
+  r: 16
+  lora_alpha: 32
+  lora_dropout: 0.05
+  target_modules: all
+
+# Kernel Config
+kernel_config:
+  name: auto
+
+# FSDP Config
+dist_config:
+  name: fsdp2
+  dcp_path: null
+
+### data
+train_dataset: data/v1_dpo_demo.yaml
+
+### training
+output_dir: ./outputs/test_lora
+micro_batch_size: 1
+cutoff_len: 2048
+learning_rate: 1.0e-5
+max_steps: 10
+
+### sample
+sample_backend: hf
+max_new_tokens: 128

examples/v1/train_lora/train_lora_sft.yaml

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

examples/v1/train_lora/train_lora_sft_rank0.yaml

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

examples/v1/train_qlora/quantization.yaml

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

pyproject.toml

@@ -40,7 +40,7 @@ dependencies = [
     "torch>=2.4.0",
     "torchvision>=0.19.0",
     "torchaudio>=2.4.0",
-    "transformers>=4.55.0,<=5.6.0,!=4.52.0,!=4.57.0",
+    "transformers>=4.55.0,<=5.8.0,!=4.57.0,!=5.6.0",
     "datasets>=2.16.0,<=4.0.0",
     "accelerate>=1.3.0,<=1.11.0",
     "peft>=0.18.0,<=0.18.1",
@@ -112,7 +112,7 @@ ignore = [
     "D105",  # no doc magic method
     "D107",  # no doc __init__
 ]
-extend-select = [
+select = [
     "C",      # complexity
     "E",      # error
     "F",      # pyflakes

requirements/triton_ascend.txt

@@ -0,0 +1,2 @@
+--extra-index-url https://triton-ascend.osinfra.cn/pypi/simple
+triton-ascend==3.2.1

src/llamafactory/chat/hf_engine.py

@@ -19,7 +19,7 @@ from threading import Thread
 from typing import TYPE_CHECKING, Any, Optional, Union
 
 import torch
-from transformers import GenerationConfig, TextIteratorStreamer
+from transformers import GenerationConfig, TextIteratorStreamer, set_seed
 from typing_extensions import override
 
 from ..data import get_template_and_fix_tokenizer
@@ -128,6 +128,7 @@ class HuggingfaceEngine(BaseEngine):
         skip_special_tokens: Optional[bool] = input_kwargs.pop("skip_special_tokens", None)
         max_length: Optional[int] = input_kwargs.pop("max_length", None)
         max_new_tokens: Optional[int] = input_kwargs.pop("max_new_tokens", None)
+        seed: Optional[int] = input_kwargs.pop("seed", None)
         stop: Optional[Union[str, list[str]]] = input_kwargs.pop("stop", None)
 
         if stop is not None:
@@ -177,6 +178,8 @@ class HuggingfaceEngine(BaseEngine):
             attention_mask=attention_mask,
             generation_config=GenerationConfig(**generating_args),
         )
+        if seed is not None:
+            gen_kwargs["_seed"] = seed
 
         mm_inputs = template.mm_plugin.get_mm_inputs(**mm_input_dict, batch_ids=[prompt_ids], processor=processor)
         for key, value in mm_inputs.items():
@@ -237,6 +240,10 @@ class HuggingfaceEngine(BaseEngine):
             audios,
             input_kwargs,
         )
+        seed = gen_kwargs.pop("_seed", None)
+        if seed is not None:
+            set_seed(seed)
+
         generate_output = model.generate(**gen_kwargs)
         if isinstance(generate_output, tuple):
             generate_output = generate_output[1][0]  # post-process the minicpm_o output
@@ -292,6 +299,10 @@ class HuggingfaceEngine(BaseEngine):
             audios,
             input_kwargs,
         )
+        seed = gen_kwargs.pop("_seed", None)
+        if seed is not None:
+            set_seed(seed)
+
         streamer = TextIteratorStreamer(
             tokenizer,
             skip_prompt=True,

src/llamafactory/chat/sglang_engine.py

@@ -171,6 +171,7 @@ class SGLangEngine(BaseEngine):
         skip_special_tokens: Optional[bool] = input_kwargs.pop("skip_special_tokens", None)
         max_length: Optional[int] = input_kwargs.pop("max_length", None)
         max_new_tokens: Optional[int] = input_kwargs.pop("max_new_tokens", None)
+        seed: Optional[int] = input_kwargs.pop("seed", None)
         stop: Optional[Union[str, list[str]]] = input_kwargs.pop("stop", None)
 
         if num_return_sequences != 1:
@@ -205,6 +206,8 @@ class SGLangEngine(BaseEngine):
             if skip_special_tokens is not None
             else self.generating_args["skip_special_tokens"],
         }
+        if seed is not None:
+            sampling_params["seed"] = seed
 
         def stream_request():
             json_data = {

src/llamafactory/chat/vllm_engine.py

@@ -144,6 +144,7 @@ class VllmEngine(BaseEngine):
         skip_special_tokens: Optional[bool] = input_kwargs.pop("skip_special_tokens", None)
         max_length: Optional[int] = input_kwargs.pop("max_length", None)
         max_new_tokens: Optional[int] = input_kwargs.pop("max_new_tokens", None)
+        seed: Optional[int] = input_kwargs.pop("seed", None)
         stop: Optional[Union[str, list[str]]] = input_kwargs.pop("stop", None)
 
         if length_penalty is not None:
@@ -163,7 +164,7 @@ class VllmEngine(BaseEngine):
         if max_new_tokens:
             max_tokens = max_new_tokens
 
-        sampling_params = SamplingParams(
+        sampling_kwargs = dict(
             n=num_return_sequences,
             repetition_penalty=(
                 repetition_penalty if repetition_penalty is not None else self.generating_args["repetition_penalty"]
@@ -179,6 +180,10 @@ class VllmEngine(BaseEngine):
             if skip_special_tokens is not None
             else self.generating_args["skip_special_tokens"],
         )
+        if seed is not None:
+            sampling_kwargs["seed"] = seed
+
+        sampling_params = SamplingParams(**sampling_kwargs)
 
         multi_modal_data = {}
         if images is not None:  # add image features

src/llamafactory/data/mm_plugin.py

@@ -929,7 +929,7 @@ class InternVLPlugin(BasePlugin):
         messages = deepcopy(messages)
         mm_inputs = self._get_mm_inputs(images, videos, audios, processor)
 
-        image_pixel_patch_list = mm_inputs.get("image_num_patches")  # pathes of images
+        image_pixel_patch_list = mm_inputs.get("image_num_patches")  # patches of images
         video_num_patches = mm_inputs.get("video_num_patches")  # all patches for frames of videos
         video_patch_indices = mm_inputs.get("video_patch_indices")  # num frames of per video
 
@@ -2379,16 +2379,13 @@ class Qwen3VLPlugin(Qwen2VLPlugin):
 
             image_grid_thw = mm_inputs.get("image_grid_thw", [])
             video_grid_thw = mm_inputs.get("video_grid_thw", [])
-            num_frames = video_grid_thw[0][0] if len(video_grid_thw) > 0 else 0  # hard code for now
             video_metadata = mm_inputs.get("video_metadata", [])
 
         else:
             image_grid_thw = [None] * len(images)
             video_grid_thw = [None] * len(videos)
-            num_frames = 0
-            timestamps = [0]
 
-        for idx, message in enumerate(messages):
+        for message in messages:
             content = message["content"]
             while IMAGE_PLACEHOLDER in content:
                 image_seqlen = (
@@ -2403,19 +2400,17 @@ class Qwen3VLPlugin(Qwen2VLPlugin):
 
             while VIDEO_PLACEHOLDER in content:
                 if self.expand_mm_tokens:
-                    metadata = video_metadata[idx]
+                    video_grid = video_grid_thw[num_video_tokens]
+                    num_frames = int(video_grid[0].item())
+                    metadata = video_metadata[num_video_tokens]
                     timestamps = processor._calculate_timestamps(
                         metadata.frames_indices,
                         metadata.fps,
-                        video_processor.merge_size,
+                        getattr(video_processor, "temporal_patch_size", 2),
                     )
                     video_structure = ""
+                    video_seqlen = int((video_grid[1:].prod() // video_merge_length).item())
                     for frame_index in range(num_frames):
-                        video_seqlen = (
-                            video_grid_thw[num_video_tokens][1:].prod() // video_merge_length
-                            if self.expand_mm_tokens
-                            else 1
-                        )
                         timestamp_sec = timestamps[frame_index]
                         frame_structure = (
                             f"<{timestamp_sec:.1f} seconds>"
@@ -2672,7 +2667,7 @@ class Qwen2OmniPlugin(Qwen2VLPlugin):
 
             if (
                 use_audio_in_video and len(audios) and len(videos)
-            ):  # if use the audio of video # deal video token and audio token togather
+            ):  # if use the audio of video # deal video token and audio token together
                 if len(videos) != len(audios):
                     raise ValueError(
                         f"Number of videos ({len(videos)}) must match number of audios ({len(audios)}) when using audio in video."

src/llamafactory/data/template.py

@@ -655,8 +655,10 @@ def get_template_and_fix_tokenizer(tokenizer: "PreTrainedTokenizer", data_args:
 
     if isinstance(template, ReasoningTemplate):
         logger.warning_rank0(
-            "You are using reasoning template, "
-            "please add `_nothink` suffix if the model is not a reasoning model. "
+            "You are using reasoning template. "
+            "If the base model is NOT a reasoning model (i.e., it has a separate Instruct variant), "
+            "please add `_nothink` suffix to disable thinking. "
+            "For reasoning-only model families (e.g., Qwen3.6), the suffix is not needed. "
             "e.g., qwen3_vl_nothink"
         )
         template.enable_thinking = data_args.enable_thinking
@@ -1033,7 +1035,7 @@ register_template(
     format_assistant=StringFormatter(slots=["{{content}}<turn|>\n"]),
     format_system=StringFormatter(
         slots=["<|turn>system\n<|think|>{{content}}<turn|>\n"]
-    ),  #  default thought singal contained
+    ),  #  default thought signal contained
     format_observation=StringFormatter(
         slots=["<|turn>tool\n{{content}}<turn|>\n<|turn>model\n"]
     ),  # seem not consistent with the chattemplate
@@ -1059,7 +1061,7 @@ register_template(
     format_assistant=StringFormatter(slots=["{{content}}<turn|>\n"]),
     format_system=StringFormatter(
         slots=["<|turn>system\n<|think|>{{content}}<turn|>\n"]
-    ),  #  default thought singal contained
+    ),  #  default thought signal contained
     format_observation=StringFormatter(slots=["<|turn>tool\n{{content}}<turn|>\n<|turn>model\n"]),
     format_tools=ToolFormatter(tool_format="gemma4"),
     format_function=FunctionFormatter(slots=["<|tool>{{content}}<tool|>"], tool_format="gemma4"),
@@ -1274,6 +1276,30 @@ register_template(
 )
 
 
+# The following two templates are copied from the official Hy-MT2 chat templates:
+# https://github.com/Tencent-Hunyuan/Hy-MT2/blob/main/train/llama_factory_support/hy_dense_template.py
+register_template(
+    name="hy_dense_1_8b",
+    format_user=StringFormatter(slots=["<｜hy_User｜>{{content}}"]),
+    format_assistant=StringFormatter(slots=["<｜hy_Assistant｜>{{content}}"]),
+    format_system=StringFormatter(slots=["{{content}}<｜hy_place▁holder▁no▁3｜>"]),
+    format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
+    stop_words=["<｜hy_place▁holder▁no▁2｜>"],
+    efficient_eos=True,
+)
+
+
+register_template(
+    name="hy_dense_7b",
+    format_user=StringFormatter(slots=["{{content}}<|extra_0|>"]),
+    format_assistant=StringFormatter(slots=["{{content}}"]),
+    format_system=StringFormatter(slots=["{{content}}<|extra_4|>"]),
+    format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
+    stop_words=["<|eos|>"],
+    efficient_eos=True,
+)
+
+
 register_template(
     name="intern2",
     format_user=StringFormatter(slots=["<|im_start|>user\n{{content}}<|im_end|>\n<|im_start|>assistant\n"]),

src/llamafactory/extras/constants.py

@@ -886,6 +886,9 @@ register_model_group(
         "Gemma-4-E4B-Thinking": {
             DownloadSource.DEFAULT: "google/gemma-4-E4B-it",
         },
+        "Gemma-4-12B-Thinking": {
+            DownloadSource.DEFAULT: "google/gemma-4-12B-it",
+        },
     },
     template="gemma4n",
     multimodal=True,
@@ -1245,6 +1248,28 @@ register_model_group(
 )
 
 
+register_model_group(
+    models={
+        "Hy-MT2-1.8B-Instruct": {
+            DownloadSource.DEFAULT: "tencent/Hy-MT2-1.8B",
+            DownloadSource.MODELSCOPE: "Tencent-Hunyuan/Hy-MT2-1.8B",
+        },
+    },
+    template="hy_dense_1_8b",
+)
+
+
+register_model_group(
+    models={
+        "Hy-MT2-7B-Instruct": {
+            DownloadSource.DEFAULT: "tencent/Hy-MT2-7B",
+            DownloadSource.MODELSCOPE: "Tencent-Hunyuan/Hy-MT2-7B",
+        },
+    },
+    template="hy_dense_7b",
+)
+
+
 register_model_group(
     models={
         "HY-MT1.5-1.8B-Instruct": {
@@ -1912,6 +1937,17 @@ register_model_group(
 )
 
 
+register_model_group(
+    models={
+        "MiniCPM5-1B-Chat": {
+            DownloadSource.DEFAULT: "openbmb/MiniCPM5-1B",
+            DownloadSource.MODELSCOPE: "OpenBMB/MiniCPM5-1B",
+        },
+    },
+    template="empty",
+)
+
+
 register_model_group(
     models={
         "MiniCPM-o-2.6": {
@@ -3336,6 +3372,19 @@ register_model_group(
 )
 
 
+# Qwen-AgentWorld: language world model based on Qwen3.5-35B-A3B (MoE).
+# Reference: https://github.com/QwenLM/Qwen-AgentWorld
+register_model_group(
+    models={
+        "Qwen-AgentWorld-35B-A3B-Thinking": {
+            DownloadSource.DEFAULT: "Qwen/Qwen-AgentWorld-35B-A3B",
+            DownloadSource.MODELSCOPE: "Qwen/Qwen-AgentWorld-35B-A3B",
+        },
+    },
+    template="qwen3",
+)
+
+
 register_model_group(
     models={
         "Vicuna-v1.5-7B-Chat": {

src/llamafactory/extras/env.py

@@ -19,7 +19,7 @@
 from collections import OrderedDict
 
 
-VERSION = "0.9.5"
+VERSION = "0.9.6.dev0"
 
 
 def print_env() -> None:

src/llamafactory/extras/misc.py

@@ -94,7 +94,7 @@ def check_version(requirement: str, mandatory: bool = False) -> None:
 
 def check_dependencies() -> None:
     r"""Check the version of the required packages."""
-    check_version("transformers>=4.55.0,<=5.6.0")
+    check_version("transformers>=4.55.0,<=5.8.0,!=4.57.0,!=5.6.0")
     check_version("datasets>=2.16.0,<=4.0.0")
     check_version("accelerate>=1.3.0,<=1.15.0")
     check_version("peft>=0.18.0,<=0.20.0")

src/llamafactory/hparams/evaluation_args.py

@@ -44,7 +44,7 @@ class EvaluationArguments:
     )
     n_shot: int = field(
         default=5,
-        metadata={"help": "Number of examplars for few-shot learning."},
+        metadata={"help": "Number of exemplars for few-shot learning."},
     )
     save_dir: str | None = field(
         default=None,

src/llamafactory/hparams/finetuning_args.py

@@ -385,7 +385,7 @@ class BAdamArgument:
             "help": (
                 "The mode of the mask for BAdam optimizer. "
                 "`adjacent` means that the trainable parameters are adjacent to each other, "
-                "`scatter` means that trainable parameters are randomly choosed from the weight."
+                "`scatter` means that trainable parameters are randomly chosen from the weight."
             )
         },
     )
@@ -487,7 +487,7 @@ class FinetuningArguments(
         metadata={
             "help": (
                 "Whether or not to use HyperParallel distributed training backend (FSDP/TP). "
-                "Only supported for the 'sft' stage with full fine-tuning."
+                "Only supported for the 'pt' and 'sft' stages with full fine-tuning."
             )
         },
     )
@@ -500,6 +500,10 @@ class FinetuningArguments(
             )
         },
     )
+    hyper_parallel_cp_size: int = field(
+        default=1,
+        metadata={"help": "Context parallel size used when `use_hyper_parallel=True`."},
+    )
     use_muon: bool = field(
         default=False,
         metadata={"help": "Whether or not to use the Muon optimizer."},
@@ -526,7 +530,7 @@ class FinetuningArguments(
     )
     freeze_vision_tower: bool = field(
         default=True,
-        metadata={"help": "Whether ot not to freeze the vision tower in MLLM training."},
+        metadata={"help": "Whether or not to freeze the vision tower in MLLM training."},
     )
     freeze_multi_modal_projector: bool = field(
         default=True,
@@ -576,6 +580,7 @@ class FinetuningArguments(
         assert self.finetuning_type in ["lora", "oft", "freeze", "full"], "Invalid fine-tuning method."
         assert self.ref_model_quantization_bit in [None, 8, 4], "We only accept 4-bit or 8-bit quantization."
         assert self.reward_model_quantization_bit in [None, 8, 4], "We only accept 4-bit or 8-bit quantization."
+        assert self.hyper_parallel_cp_size > 0, "`hyper_parallel_cp_size` must be greater than 0."
 
         if self.stage == "ppo" and self.reward_model is None:
             raise ValueError("`reward_model` is necessary for PPO training.")

src/llamafactory/model/adapter.py

@@ -194,10 +194,20 @@ def _setup_lora_tuning(
             logger.info_rank0(f"Merged {len(adapter_to_merge)} adapter(s).")
 
         if adapter_to_resume is not None:  # resume lora training
-            if model_args.use_unsloth:
-                model = load_unsloth_peft_model(config, model_args, finetuning_args, is_trainable=is_trainable)
+            if isinstance(model, PeftModel):
+                pass  # already loaded via load_unsloth_peft_model in loader.py
             else:
-                model = PeftModel.from_pretrained(model, adapter_to_resume, is_trainable=is_trainable, **init_kwargs)
+                if model_args.use_unsloth:
+                    peft_model = load_unsloth_peft_model(
+                        config, model_args, finetuning_args, is_trainable=is_trainable
+                    )
+                    if peft_model is not None:
+                        model = peft_model
+
+                if not model_args.use_unsloth:  # unsloth was disabled or fell back
+                    model = PeftModel.from_pretrained(
+                        model, adapter_to_resume, is_trainable=is_trainable, **init_kwargs
+                    )
 
         logger.info_rank0("Loaded adapter(s): {}".format(",".join(model_args.adapter_name_or_path)))
 

src/llamafactory/model/loader.py

@@ -34,7 +34,7 @@ from .adapter import init_adapter
 from .model_utils.liger_kernel import apply_liger_kernel
 from .model_utils.misc import register_autoclass
 from .model_utils.mod import convert_pretrained_model_to_mod, load_mod_pretrained_model
-from .model_utils.unsloth import load_unsloth_pretrained_model
+from .model_utils.unsloth import load_unsloth_peft_model, load_unsloth_pretrained_model
 from .model_utils.valuehead import load_valuehead_params
 from .patcher import patch_config, patch_model, patch_processor, patch_tokenizer, patch_valuehead_model
 
@@ -142,14 +142,13 @@ def load_model(
     apply_liger_kernel(config, model_args, is_trainable, require_logits=(finetuning_args.stage not in ["pt", "sft"]))
 
     model = None
-    lazy_load = False
     if model_args.use_unsloth:
         if model_args.adapter_name_or_path is not None:
-            lazy_load = True
+            model = load_unsloth_peft_model(config, model_args, finetuning_args, is_trainable=is_trainable)
         elif is_trainable:
             model = load_unsloth_pretrained_model(config, model_args, finetuning_args)
 
-    if model is None and not lazy_load:
+    if model is None:
         init_kwargs["config"] = config
         init_kwargs["pretrained_model_name_or_path"] = model_args.model_name_or_path
         init_kwargs["torch_dtype"] = "auto"
@@ -176,9 +175,8 @@ def load_model(
         if model_args.mixture_of_depths == "convert":
             model = convert_pretrained_model_to_mod(model, config, model_args)
 
-    if not lazy_load:
-        patch_model(model, tokenizer, model_args, is_trainable, add_valuehead)
-        register_autoclass(config, model, tokenizer)
+    patch_model(model, tokenizer, model_args, is_trainable, add_valuehead)
+    register_autoclass(config, model, tokenizer)
 
     model = init_adapter(config, model, model_args, finetuning_args, is_trainable)
 

src/llamafactory/model/model_utils/embedding.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import math
+from collections.abc import Iterable
 from contextlib import nullcontext
 from typing import TYPE_CHECKING, Optional
 
@@ -29,7 +30,81 @@ if TYPE_CHECKING:
 logger = logging.get_logger(__name__)
 
 
-def _noisy_mean_initialization(embed_weight: "torch.Tensor", num_new_tokens: int) -> None:
+def get_embedding_vocab_size(model: "PreTrainedModel") -> int:
+    r"""Get the vocab size from the input embedding layer.
+
+    Handles DeepSpeed ZeRO-3 parameter sharding by gathering the embedding weight
+    before reading its size.
+    """
+    embedding = model.get_input_embeddings()
+    if is_deepspeed_zero3_enabled():
+        import deepspeed  # type: ignore
+
+        with deepspeed.zero.GatheredParameters([embedding.weight]):
+            return embedding.weight.size(0)
+
+    return embedding.weight.size(0)
+
+
+def _resolve_new_token_ids(
+    new_tokens: Optional[Iterable[str]],
+    tokenizer: "PreTrainedTokenizer",
+    embed_size: int,
+) -> Optional[list[int]]:
+    r"""Resolve the explicit embedding-row IDs of the newly added tokens.
+
+    Relying on ``embed_weight[-num_new_tokens:]`` to locate new tokens is unsafe when
+    the model embedding was already padded beyond the tokenizer vocab (e.g. Qwen2.5-VL
+    has vocab 151665 but embedding 151936). In that case the appended tokens land
+    inside the original padding zone and the tail slice points at the wrong rows.
+
+    Args:
+        new_tokens: Iterable of the newly added token strings.
+        tokenizer: The tokenizer instance.
+        embed_size: Current embedding size (upper bound for valid token IDs).
+
+    Returns:
+        A sorted list of unique, in-range token IDs, or ``None`` when no tokens are
+        given so that callers can fall back to the tail-slice behaviour.
+    """
+    if not new_tokens:
+        return None
+
+    unk_token_id = getattr(tokenizer, "unk_token_id", None)
+    token_ids: set[int] = set()
+    for token_str in new_tokens:
+        token_id = tokenizer.convert_tokens_to_ids(token_str)
+        if token_id is None or token_id == unk_token_id or not (0 <= token_id < embed_size):
+            logger.warning_rank0(f"Token '{token_str}' not found or out of range, skipping during init.")
+            continue
+
+        token_ids.add(token_id)
+
+    return sorted(token_ids) or None
+
+
+def _existing_embeddings(
+    embed_weight: "torch.Tensor", num_new_tokens: int, new_token_ids: Optional[list[int]]
+) -> "torch.Tensor":
+    """Return the rows treated as 'existing' embeddings used as the init baseline.
+
+    Prefers excluding the explicit new-token rows (robust to padding). Falls back to
+    dropping the last ``num_new_tokens`` rows when no explicit IDs are available.
+    """
+    if new_token_ids:
+        mask = torch.ones(embed_weight.size(0), dtype=torch.bool, device=embed_weight.device)
+        mask[torch.as_tensor(new_token_ids, device=embed_weight.device, dtype=torch.long)] = False
+        return embed_weight[mask]
+
+    if num_new_tokens > 0:
+        return embed_weight[:-num_new_tokens]
+
+    return embed_weight
+
+
+def _noisy_mean_initialization(
+    embed_weight: "torch.Tensor", num_new_tokens: int, token_ids: Optional[list[int]] = None
+) -> None:
     """Initialize new token embeddings with mean + Gaussian noise.
 
     This is the default initialization method used by LlamaFactory.
@@ -37,12 +112,21 @@ def _noisy_mean_initialization(embed_weight: "torch.Tensor", num_new_tokens: int
     Args:
         embed_weight: The embedding weight matrix to initialize (shape: [vocab_size, embedding_dim])
         num_new_tokens: Number of new tokens added at the end of the embedding matrix
+        token_ids: Explicit token IDs to initialize. When provided, these exact rows are
+            written (robust to padding). When ``None``, falls back to the last
+            ``num_new_tokens`` rows.
     """
     embedding_dim = embed_weight.size(1)
-    avg_weight = embed_weight[:-num_new_tokens].mean(dim=0, keepdim=True)
-    noise_weight = torch.empty_like(embed_weight[-num_new_tokens:])
-    noise_weight.normal_(mean=0, std=(1.0 / math.sqrt(embedding_dim)))
-    embed_weight[-num_new_tokens:] = avg_weight + noise_weight
+    avg_weight = _existing_embeddings(embed_weight, num_new_tokens, token_ids).mean(dim=0, keepdim=True)
+
+    if token_ids:
+        noise_weight = torch.empty(len(token_ids), embedding_dim, device=embed_weight.device, dtype=embed_weight.dtype)
+        noise_weight.normal_(mean=0, std=(1.0 / math.sqrt(embedding_dim)))
+        embed_weight[token_ids] = avg_weight + noise_weight
+    else:
+        noise_weight = torch.empty_like(embed_weight[-num_new_tokens:])
+        noise_weight.normal_(mean=0, std=(1.0 / math.sqrt(embedding_dim)))
+        embed_weight[-num_new_tokens:] = avg_weight + noise_weight
 
 
 def _description_based_initialization(
@@ -51,6 +135,7 @@ def _description_based_initialization(
     descriptions: dict[str, str],
     tokenizer: "PreTrainedTokenizer",
     model: "PreTrainedModel",
+    new_token_ids: Optional[list[int]] = None,
     add_noise: bool = False,
 ) -> None:
     """Initialize new token embeddings based on textual descriptions.
@@ -61,6 +146,9 @@ def _description_based_initialization(
     3. Averages them to initialize the new token's embedding
     4. Optionally adds Gaussian noise
 
+    New tokens are placed by their resolved token ID rather than by tail slicing,
+    so the initialization is correct even when the embedding matrix was padded.
+
     Args:
         embed_weight: The embedding weight matrix to initialize (shape: [vocab_size, embedding_dim])
         num_new_tokens: Number of new tokens added
@@ -68,6 +156,8 @@ def _description_based_initialization(
                       e.g., {"<think>": "A token representing reasoning process"}
         tokenizer: The tokenizer instance
         model: The model instance (used to get input embeddings)
+        new_token_ids: IDs of all newly added tokens. Used to exclude not-yet-initialized
+            rows when averaging description-token embeddings (robust to embedding padding).
         add_noise: Whether to add Gaussian noise to the initialization
 
     Example:
@@ -77,38 +167,54 @@ def _description_based_initialization(
         }
     """
     embedding_dim = embed_weight.size(1)
+    vocab_size = embed_weight.size(0)
+    unk_token_id = getattr(tokenizer, "unk_token_id", None)
+    device = embed_weight.device
+
+    # The set of rows that are NOT yet initialized (the newly added tokens). Description
+    # tokens that fall into this set must be excluded, otherwise we would average garbage.
+    # `num_new_tokens` (the padded resize delta) is NOT a reliable boundary, so rely on
+    # the explicit IDs, falling back to resolving them from the description keys.
+    if new_token_ids is None:
+        new_token_ids = _resolve_new_token_ids(descriptions.keys(), tokenizer, vocab_size)
+
+    new_id_set = set(new_token_ids or [])
+    fallback_embedding = _existing_embeddings(embed_weight, num_new_tokens, new_token_ids).mean(dim=0)
+
+    for token_str, desc in descriptions.items():
+        # Resolve token ID for correct placement (robust to embedding padding)
+        token_id = tokenizer.convert_tokens_to_ids(token_str)
+        if token_id is None or token_id == unk_token_id or not (0 <= token_id < vocab_size):
+            logger.warning_rank0(f"desc_init: token '{token_str}' not found or out of range, skipping.")
+            continue
 
-    for i, desc in enumerate(descriptions.values()):
         # Tokenize description text
         tokens = tokenizer(desc, return_tensors="pt", add_special_tokens=False)
 
         with torch.no_grad():
-            token_ids = tokens["input_ids"][0]
-            # Move to the same device as embed_weight
-            device = embed_weight.device
-            token_ids = token_ids.to(device)
+            token_ids = tokens["input_ids"][0].tolist()
 
-            # Filter out new tokens (they don't have valid embeddings yet)
-            valid_token_ids = token_ids[token_ids < (len(tokenizer) - num_new_tokens)]
+            # Keep only description tokens that already have a meaningful embedding.
+            valid_token_ids = [tid for tid in token_ids if tid not in new_id_set and 0 <= tid < vocab_size]
 
             if len(valid_token_ids) == 0:
                 # Fallback: use mean of all existing embeddings
                 logger.warning_rank0(
-                    f"Description for token {i + 1}/{num_new_tokens} contains no valid tokens. "
-                    "Using mean of existing embeddings."
+                    f"Description for token '{token_str}' contains no valid tokens. Using mean of existing embeddings."
                 )
-                base_embedding = embed_weight[:-num_new_tokens].mean(dim=0)
+                base_embedding = fallback_embedding
             else:
                 # Get embeddings of description tokens and average them
-                token_embeds = model.get_input_embeddings()(valid_token_ids)
+                valid_ids_tensor = torch.as_tensor(valid_token_ids, device=device, dtype=torch.long)
+                token_embeds = model.get_input_embeddings()(valid_ids_tensor)
                 base_embedding = token_embeds.mean(dim=0)
 
             # Add noise if requested (ensure correct device and dtype)
             if add_noise:
                 noise = torch.randn_like(base_embedding) * (1.0 / math.sqrt(embedding_dim))
-                embed_weight[-num_new_tokens + i] = base_embedding + noise
+                embed_weight[token_id] = base_embedding + noise
             else:
-                embed_weight[-num_new_tokens + i] = base_embedding
+                embed_weight[token_id] = base_embedding
 
 
 def _initialize_embeddings(
@@ -118,6 +224,7 @@ def _initialize_embeddings(
     new_special_tokens_config: Optional[dict],
     tokenizer: "PreTrainedTokenizer",
     model: "PreTrainedModel",
+    new_token_ids: Optional[list[int]] = None,
 ) -> None:
     """Single source of truth for embedding initialization.
 
@@ -130,16 +237,18 @@ def _initialize_embeddings(
         new_special_tokens_config: Config dict with token descriptions (required for desc_init methods)
         tokenizer: The tokenizer instance
         model: The model instance
+        new_token_ids: Explicit IDs of the newly added tokens (robust to embedding padding).
+            When ``None``, the init helpers fall back to the last ``num_new_tokens`` rows.
     """
     if init_method == "desc_init" and new_special_tokens_config:
         logger.info_rank0("Using semantic initialization (desc_init) for new special tokens")
         _description_based_initialization(
-            embed_weight, num_new_tokens, new_special_tokens_config, tokenizer, model, add_noise=False
+            embed_weight, num_new_tokens, new_special_tokens_config, tokenizer, model, new_token_ids, add_noise=False
         )
     elif init_method == "desc_init_w_noise" and new_special_tokens_config:
         logger.info_rank0("Using semantic initialization with noise (desc_init_w_noise) for new special tokens")
         _description_based_initialization(
-            embed_weight, num_new_tokens, new_special_tokens_config, tokenizer, model, add_noise=True
+            embed_weight, num_new_tokens, new_special_tokens_config, tokenizer, model, new_token_ids, add_noise=True
         )
     else:
         if init_method != "noise_init":
@@ -147,20 +256,28 @@ def _initialize_embeddings(
                 f"init_method='{init_method}' requires descriptions config, falling back to 'noise_init'"
             )
         logger.info_rank0("Using noisy mean initialization (noise_init) for new special tokens")
-        _noisy_mean_initialization(embed_weight, num_new_tokens)
+        _noisy_mean_initialization(embed_weight, num_new_tokens, token_ids=new_token_ids)
 
 
 def resize_embedding_layer(
     model: "PreTrainedModel",
     tokenizer: "PreTrainedTokenizer",
+    new_tokens: Optional[Iterable[str]] = None,
     new_special_tokens_config: Optional[dict] = None,
     init_special_tokens: str = "noise_init",
 ) -> None:
-    r"""Resize token embeddings and initialize new tokens.
+    r"""Resize token embeddings (when needed) and initialize the newly added tokens.
+
+    Resizing and initialization are decoupled: even when the tokenizer vocab fits inside
+    the model's existing (padded) embedding matrix and no resize is triggered, the newly
+    added tokens still occupy uninitialized rows and must be initialized. We therefore
+    resolve the explicit row IDs of ``new_tokens`` and always initialize those rows.
 
     Args:
         model: The model to resize
         tokenizer: The tokenizer (used to get target vocab size)
+        new_tokens: Iterable of the newly added token strings. Used to locate the exact
+            embedding rows to initialize, which is robust to pre-existing embedding padding.
         new_special_tokens_config: Optional dict with token descriptions for semantic initialization
         init_special_tokens: Initialization method ('noise_init', 'desc_init', 'desc_init_w_noise')
     """
@@ -175,44 +292,70 @@ def resize_embedding_layer(
     else:
         context_maybe_zero3 = nullcontext()
 
-    with context_maybe_zero3:
-        current_embedding_size = model.get_input_embeddings().weight.size(0)
+    current_embedding_size = get_embedding_vocab_size(model)
+    needs_resize = len(tokenizer) > current_embedding_size
 
-    if len(tokenizer) > current_embedding_size:
+    if needs_resize:
         if getattr(model, "quantization_method", None):
             raise ValueError("Cannot resize embedding layers of a quantized model.")
 
         if not isinstance(model.get_output_embeddings(), torch.nn.Linear):
             raise ValueError("Current model does not support resizing embedding layers.")
 
-        model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=64)
-        with context_maybe_zero3:
-            new_embedding_size = model.get_input_embeddings().weight.size(0)
-            num_new_tokens = new_embedding_size - current_embedding_size
+        # mean_resizing=False preserves the original embedding distribution exactly.
+        # HuggingFace's default mean_resizing=True re-samples new rows from the mean/covariance
+        # of existing embeddings, which conflicts with our explicit initialization below.
+        model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=64, mean_resizing=False)
+
+    with context_maybe_zero3:
+        new_embedding_size = model.get_input_embeddings().weight.size(0)
+        num_new_tokens = new_embedding_size - current_embedding_size
+
+        # Resolve the exact rows of the new tokens. This works whether or not a resize was
+        # triggered (e.g. tokens added into a model's pre-existing padding zone).
+        new_token_ids = _resolve_new_token_ids(new_tokens, tokenizer, new_embedding_size)
+
+        if num_new_tokens <= 0 and not new_token_ids:
+            return
+
+        if needs_resize:
             logger.info_rank0(
                 f"Resizing embeddings: {current_embedding_size} -> {new_embedding_size} (+{num_new_tokens} tokens)"
             )
+        else:
+            logger.info_rank0(
+                f"No resize needed (vocab fits in padded embedding {new_embedding_size}); "
+                f"initializing {len(new_token_ids or [])} new token(s) in place."
+            )
 
-            # Initialize input embeddings
+        # Initialize input embeddings
+        _initialize_embeddings(
+            model.get_input_embeddings().weight.data,
+            num_new_tokens,
+            init_special_tokens,
+            new_special_tokens_config,
+            tokenizer,
+            model,
+            new_token_ids=new_token_ids,
+        )
+
+        # Initialize output embeddings if not tied
+        if model.get_output_embeddings() is not None and not model.config.tie_word_embeddings:
             _initialize_embeddings(
-                model.get_input_embeddings().weight.data,
+                model.get_output_embeddings().weight.data,
                 num_new_tokens,
                 init_special_tokens,
                 new_special_tokens_config,
                 tokenizer,
                 model,
+                new_token_ids=new_token_ids,
             )
 
-            # Initialize output embeddings if not tied
-            if model.get_output_embeddings() is not None and not model.config.tie_word_embeddings:
-                _initialize_embeddings(
-                    model.get_output_embeddings().weight.data,
-                    num_new_tokens,
-                    init_special_tokens,
-                    new_special_tokens_config,
-                    tokenizer,
-                    model,
-                )
-
+    if needs_resize:
         model.config.vocab_size = new_embedding_size
+        # Also update the nested text_config for VL models (e.g., Qwen2.5-VL, LLaVA),
+        # otherwise config.vocab_size and config.text_config.vocab_size become inconsistent.
+        if hasattr(model.config, "text_config") and hasattr(model.config.text_config, "vocab_size"):
+            model.config.text_config.vocab_size = new_embedding_size
+
         logger.info_rank0(f"Resized token embeddings from {current_embedding_size} to {new_embedding_size}.")

src/llamafactory/model/model_utils/liger_kernel.py

@@ -16,6 +16,7 @@ import inspect
 from typing import TYPE_CHECKING
 
 from ...extras import logging
+from ...extras.misc import get_device_name
 
 
 if TYPE_CHECKING:
@@ -81,6 +82,8 @@ def apply_liger_kernel(
         from liger_kernel.transformers import apply_liger_kernel_to_qwen3_next as apply_liger_kernel
     elif model_type == "qwen3_5":
         from liger_kernel.transformers import apply_liger_kernel_to_qwen3_5 as apply_liger_kernel
+    elif model_type == "qwen3_5_moe":
+        from liger_kernel.transformers import apply_liger_kernel_to_qwen3_5_moe as apply_liger_kernel
     elif model_type == "gpt_oss":
         try:
             from liger_kernel.transformers import apply_liger_kernel_to_gpt_oss as apply_liger_kernel
@@ -97,5 +100,12 @@ def apply_liger_kernel(
     else:
         kwargs = {}
 
+    if get_device_name() == "npu":
+        import torch
+
+        if "Ascend910" not in torch.npu.get_device_name(0):
+            kwargs["swiglu"] = False
+            kwargs["fused_linear_cross_entropy"] = False
+
     apply_liger_kernel(**kwargs)
     logger.info_rank0("Liger kernel has been applied to the model.")

src/llamafactory/model/model_utils/unsloth.py

@@ -84,8 +84,12 @@ def load_unsloth_peft_model(
     model_args: "ModelArguments",
     finetuning_args: "FinetuningArguments",
     is_trainable: bool,
-) -> "PreTrainedModel":
-    r"""Load peft model with unsloth. Used in both training and inference."""
+) -> Optional["PreTrainedModel"]:
+    r"""Load peft model with unsloth. Used in both training and inference.
+
+    Returns None if unsloth does not support the model type, and sets
+    model_args.use_unsloth = False so callers can fall back to standard loading.
+    """
     from unsloth import FastLanguageModel  # type: ignore
 
     unsloth_kwargs = _get_unsloth_kwargs(config, model_args.adapter_name_or_path[0], model_args, finetuning_args)
@@ -95,7 +99,9 @@ def load_unsloth_peft_model(
 
         model, _ = FastLanguageModel.from_pretrained(**unsloth_kwargs)
     except NotImplementedError:
-        raise ValueError("Unsloth does not support model type {}.".format(getattr(config, "model_type", None)))
+        logger.warning_rank0("Unsloth does not support model type {}.".format(getattr(config, "model_type", None)))
+        model_args.use_unsloth = False
+        return None
 
     if not is_trainable:
         FastLanguageModel.for_inference(model)

src/llamafactory/model/patcher.py

@@ -20,6 +20,7 @@ from peft import PeftModel
 from transformers import GenerationMixin, PreTrainedModel, PreTrainedTokenizerBase
 from transformers.integrations import is_deepspeed_zero3_enabled
 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
@@ -84,7 +85,60 @@ def _check_fla_dependencies() -> None:
         ) from exc
 
 
-def patch_qwen3_5_forward(model: "PreTrainedModel") -> None:
+def patch_qwen3_5_forward_npu(model: "PreTrainedModel") -> None:
+    """Patch for Qwen3.5 models on NPU by importing torch_npu to enable torch.cuda compatibility.
+
+    On NPU, torch.cuda operations will fail unless torch_npu is imported.
+    torch_npu provides compatibility layer that maps torch.cuda calls to NPU operations.
+
+    Also replaces chunk_gated_delta_rule with NPU-compatible implementation.
+    """
+    import importlib.metadata
+
+    if "Ascend910" not in torch.npu.get_device_name(0):
+        logger.warning_rank0("Currently only 910B series NPUs are supported for the NPU GDN patch.")
+        return
+
+    try:
+        importlib.metadata.version("triton_ascend")
+    except importlib.metadata.PackageNotFoundError:
+        logger.warning_rank0(
+            "triton_ascend not installed, skipping NPU GDN patch. "
+            "To enable it on NPU, reinstall Triton with the Ascend build: "
+            "`pip uninstall -y triton && pip install -r requirements/triton_ascend.txt`. "
+            "Note: triton and triton_ascend cannot coexist — triton must be uninstalled first."
+        )
+        return
+
+    logger.info_rank0("triton_ascend detected for NPU compatibility.")
+
+    from ..third_party.triton.chunk_gated_delta_rule import chunk_gated_delta_rule as npu_chunk_gated_delta_rule
+
+    if model.config.architectures[0] == "Qwen3_5MoeForConditionalGeneration":
+        try:
+            # Qwen3.5-MoE structure: model.model.language_model.layers
+            for layer in model.model.language_model.layers:
+                if hasattr(layer, "linear_attn"):
+                    layer.linear_attn.chunk_gated_delta_rule = npu_chunk_gated_delta_rule
+
+            logger.info_rank0(
+                "Replaced chunk_gated_delta_rule with NPU-compatible implementation for Qwen3.5-MoE model."
+            )
+        except Exception as e:
+            logger.warning_rank0(f"Failed to replace chunk_gated_delta_rule for NPU: {e}")
+    elif model.config.architectures[0] == "Qwen3_5ForConditionalGeneration":
+        try:
+            # Qwen3.5 structure: model.model.layers
+            for layer in model.model.layers:
+                if hasattr(layer, "linear_attn"):
+                    layer.linear_attn.chunk_gated_delta_rule = npu_chunk_gated_delta_rule
+
+            logger.info_rank0("Replaced chunk_gated_delta_rule with NPU-compatible implementation for Qwen3.5 model.")
+        except Exception as e:
+            logger.warning_rank0(f"Failed to replace chunk_gated_delta_rule for NPU: {e}")
+
+
+def patch_qwen3_5_forward_gpu(model: "PreTrainedModel") -> None:
     """Patch the forward method of Qwen3_5ForConditionalGeneration to support cu_seqlens input only patch when do training.
 
     Refer to: https://github.com/axolotl-ai-cloud/axolotl/blob/main/src/axolotl/monkeypatch/models/qwen3_5/modeling.py.
@@ -403,9 +457,14 @@ def patch_model(
         prepare_valuehead_model(model)
 
     if model_args.resize_vocab:
+        # Pass the explicit list of newly added tokens so their exact embedding rows can be
+        # located and initialized, even when they land in a model's pre-existing padding zone.
+        new_tokens = (model_args.add_tokens or []) + (model_args.add_special_tokens or [])
+
         resize_embedding_layer(
             model,
             tokenizer,
+            new_tokens=new_tokens or None,
             new_special_tokens_config=getattr(model_args, "_special_token_descriptions", None),
             init_special_tokens=model_args.init_special_tokens,
         )
@@ -421,8 +480,12 @@ def patch_model(
         autocast_projector_dtype(model, model_args)
         add_z3_leaf_module(model)
 
-        if getattr(model.config, "model_type", None) in ["qwen3_5", "qwen3_5_moe"] and model_args.flash_attn == "fa2":
-            patch_qwen3_5_forward(model)
+        if getattr(model.config, "model_type", None) in ["qwen3_5", "qwen3_5_moe"]:
+            if is_torch_npu_available():
+                patch_qwen3_5_forward_npu(model)
+            elif is_torch_cuda_available() and model_args.flash_attn == "fa2":
+                # this is the patch for packing/neat_packing for GPU GDN. And when setting packing, flash_attn must be fa2.
+                patch_qwen3_5_forward_gpu(model)
 
     if not model_args.use_unsloth:
         print_attn_implementation(model.config)

src/llamafactory/third_party/triton/chunk_delta_h.py

@@ -0,0 +1,594 @@
+# Copyright 2025 the LlamaFactory team.
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+#
+# 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 typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from .utils import get_autotune_config, get_npu_properties, prepare_chunk_indices, prepare_chunk_offsets
+
+
+CUBE_CORE_NUM = get_npu_properties()["num_aicore"]
+
+
+@triton.heuristics(
+    {
+        "USE_G": lambda args: args["g"] is not None,
+        "USE_GK": lambda args: args["gk"] is not None,
+        "USE_INITIAL_STATE": lambda args: args["h0"] is not None,
+        "STORE_FINAL_STATE": lambda args: args["ht"] is not None,
+        "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.autotune(
+    configs=get_autotune_config(multibuffer_list=(False,)),
+    key=["H", "K", "V", "BT"],
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
+    k,
+    v,
+    w,
+    v_new,
+    g,
+    gk,
+    h,
+    h0,
+    ht,
+    cu_seqlens,
+    chunk_offsets,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BV: tl.constexpr,
+    NT: tl.constexpr,
+    USE_G: tl.constexpr,
+    USE_GK: tl.constexpr,
+    USE_INITIAL_STATE: tl.constexpr,
+    STORE_FINAL_STATE: tl.constexpr,
+    SAVE_NEW_VALUE: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    T_all = T
+    NT_all = NT
+    i_v, i_nh = tl.program_id(0), tl.program_id(1)
+    i_n, i_h = i_nh // H, i_nh % H
+    if IS_VARLEN:
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        T = eos - bos
+        NT = tl.cdiv(T, BT)
+        boh = tl.load(chunk_offsets + i_n).to(tl.int32)
+    else:
+        bos, eos = i_n * T, i_n * T + T
+        NT = tl.cdiv(T, BT)
+        boh = i_n * NT
+
+    # Initialize hidden states
+    b_h1 = tl.zeros([64, BV], dtype=tl.float32)
+    if K > 64:
+        b_h2 = tl.zeros([64, BV], dtype=tl.float32)
+    if K > 128:
+        b_h3 = tl.zeros([64, BV], dtype=tl.float32)
+    if K > 192:
+        b_h4 = tl.zeros([64, BV], dtype=tl.float32)
+
+    if IS_VARLEN:
+        v = v + (i_h * T_all + bos) * V
+        k = k + (i_h * T_all + bos) * K
+        w = w + (i_h * T_all + bos) * K
+        g = g + i_h * T_all + bos
+        h = h + (i_h * NT_all + boh) * K * V
+        if SAVE_NEW_VALUE:
+            v_new_base = v_new + (i_h * T_all + bos) * V
+    else:
+        v = v + (i_n * H + i_h) * T * V
+        k = k + (i_n * H + i_h) * T * K
+        w = w + (i_n * H + i_h) * T * K
+        g = g + (i_n * H + i_h) * T
+        h = h + (i_n * H + i_h) * NT * K * V
+        if SAVE_NEW_VALUE:
+            v_new_base = v_new + (i_n * H + i_h) * T * V
+
+    if USE_INITIAL_STATE:
+        h0_ptr = h0 + i_nh * K * V
+    if STORE_FINAL_STATE:
+        ht_ptr = ht + i_nh * K * V
+
+    # Load initial state
+    if USE_INITIAL_STATE:
+        p_h0_1 = tl.make_block_ptr(h0_ptr, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))
+        b_h1 += tl.load(p_h0_1, boundary_check=(0, 1)).to(tl.float32)
+        if K > 64:
+            p_h0_2 = tl.make_block_ptr(h0_ptr, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0))
+            b_h2 += tl.load(p_h0_2, boundary_check=(0, 1)).to(tl.float32)
+        if K > 128:
+            p_h0_3 = tl.make_block_ptr(h0_ptr, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0))
+            b_h3 += tl.load(p_h0_3, boundary_check=(0, 1)).to(tl.float32)
+        if K > 192:
+            p_h0_4 = tl.make_block_ptr(h0_ptr, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0))
+            b_h4 += tl.load(p_h0_4, boundary_check=(0, 1)).to(tl.float32)
+
+    # Main recurrence over chunks
+    for i_t in range(NT):
+        # Store current hidden state h_t
+        p_h1 = tl.make_block_ptr(h + i_t * K * V, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))
+        tl.store(p_h1, b_h1.to(p_h1.dtype.element_ty), boundary_check=(0, 1))
+        if K > 64:
+            p_h2 = tl.make_block_ptr(h + i_t * K * V, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0))
+            tl.store(p_h2, b_h2.to(p_h2.dtype.element_ty), boundary_check=(0, 1))
+        if K > 128:
+            p_h3 = tl.make_block_ptr(h + i_t * K * V, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0))
+            tl.store(p_h3, b_h3.to(p_h3.dtype.element_ty), boundary_check=(0, 1))
+        if K > 192:
+            p_h4 = tl.make_block_ptr(h + i_t * K * V, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0))
+            tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1))
+
+        # Compute v_residual = v - w @ h
+        p_w = tl.make_block_ptr(w, (T, K), (K, 1), (i_t * BT, 0), (BT, 64), (1, 0))
+        b_w = tl.load(p_w, boundary_check=(0, 1))
+        b_v = tl.dot(b_w, b_h1.to(b_w.dtype))
+        if K > 64:
+            p_w = tl.make_block_ptr(w, (T, K), (K, 1), (i_t * BT, 64), (BT, 64), (1, 0))
+            b_w = tl.load(p_w, boundary_check=(0, 1))
+            b_v += tl.dot(b_w, b_h2.to(b_w.dtype))
+        if K > 128:
+            p_w = tl.make_block_ptr(w, (T, K), (K, 1), (i_t * BT, 128), (BT, 64), (1, 0))
+            b_w = tl.load(p_w, boundary_check=(0, 1))
+            b_v += tl.dot(b_w, b_h3.to(b_w.dtype))
+        if K > 192:
+            p_w = tl.make_block_ptr(w, (T, K), (K, 1), (i_t * BT, 192), (BT, 64), (1, 0))
+            b_w = tl.load(p_w, boundary_check=(0, 1))
+            b_v += tl.dot(b_w, b_h4.to(b_w.dtype))
+
+        p_v = tl.make_block_ptr(v, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v
+
+        if SAVE_NEW_VALUE:
+            p_v_new = tl.make_block_ptr(v_new_base, (T, V), (V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+            tl.store(p_v_new, b_v.to(p_v_new.dtype.element_ty), boundary_check=(0, 1))
+
+        last_idx = min((i_t + 1) * BT, T) - 1
+
+        # Apply output gate g
+        if USE_G:
+            m_t = (i_t * BT + tl.arange(0, BT)).to(tl.float32) < T
+            b_g_last = tl.load(g + last_idx)
+            p_g = tl.make_block_ptr(g, (T,), (1,), (i_t * BT,), (BT,), (0,))
+            b_g = tl.load(p_g, boundary_check=(0,))
+            b_v *= (m_t * tl.exp(b_g_last - b_g))[:, None]
+            b_g_last_exp = tl.exp(b_g_last)
+            b_h1 *= b_g_last_exp
+            if K > 64:
+                b_h2 *= b_g_last_exp
+            if K > 128:
+                b_h3 *= b_g_last_exp
+            if K > 192:
+                b_h4 *= b_g_last_exp
+
+        # Apply key gate gk
+        if USE_GK:
+            o_k1 = tl.arange(0, 64).to(tl.float32)
+            gk_base_ptr = gk + (i_n * H + i_h) * T * K
+            b_gk_last1 = tl.load(gk_base_ptr + last_idx * K + o_k1, mask=(o_k1 < K), other=0.0)
+            b_h1 *= tl.exp(b_gk_last1)[:, None]
+            if K > 64:
+                o_k2 = 64 + o_k1
+                b_gk_last2 = tl.load(gk_base_ptr + last_idx * K + o_k2, mask=(o_k2 < K), other=0.0)
+                b_h2 *= tl.exp(b_gk_last2)[:, None]
+            if K > 128:
+                o_k3 = 128 + o_k1
+                b_gk_last3 = tl.load(gk_base_ptr + last_idx * K + o_k3, mask=(o_k3 < K), other=0.0)
+                b_h3 *= tl.exp(b_gk_last3)[:, None]
+            if K > 192:
+                o_k4 = 192 + o_k1
+                b_gk_last4 = tl.load(gk_base_ptr + last_idx * K + o_k4, mask=(o_k4 < K), other=0.0)
+                b_h4 *= tl.exp(b_gk_last4)[:, None]
+
+        b_v = b_v.to(k.dtype.element_ty)
+
+        # Update hidden state: h += k @ v
+        p_k = tl.make_block_ptr(k, (K, T), (1, K), (0, i_t * BT), (64, BT), (0, 1))
+        b_k = tl.load(p_k, boundary_check=(0, 1))
+        if USE_GK:
+            p_gk = tl.make_block_ptr(gk_base_ptr, (K, T), (1, K), (0, i_t * BT), (64, BT), (0, 1))
+            b_k = (b_k * tl.exp(b_gk_last1[:, None] - tl.load(p_gk, boundary_check=(0, 1)))).to(b_k.dtype)
+        b_h1 += tl.dot(b_k, b_v)
+
+        if K > 64:
+            p_k = tl.make_block_ptr(k, (K, T), (1, K), (64, i_t * BT), (64, BT), (0, 1))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            if USE_GK:
+                p_gk = tl.make_block_ptr(gk_base_ptr, (K, T), (1, K), (64, i_t * BT), (64, BT), (0, 1))
+                b_k = (b_k * tl.exp(b_gk_last2[:, None] - tl.load(p_gk, boundary_check=(0, 1)))).to(b_k.dtype)
+            b_h2 += tl.dot(b_k, b_v)
+
+        if K > 128:
+            p_k = tl.make_block_ptr(k, (K, T), (1, K), (128, i_t * BT), (64, BT), (0, 1))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            if USE_GK:
+                p_gk = tl.make_block_ptr(gk_base_ptr, (K, T), (1, K), (128, i_t * BT), (64, BT), (0, 1))
+                b_k = (b_k * tl.exp(b_gk_last3[:, None] - tl.load(p_gk, boundary_check=(0, 1)))).to(b_k.dtype)
+            b_h3 += tl.dot(b_k, b_v)
+
+        if K > 192:
+            p_k = tl.make_block_ptr(k, (K, T), (1, K), (192, i_t * BT), (64, BT), (0, 1))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            if USE_GK:
+                p_gk = tl.make_block_ptr(gk_base_ptr, (K, T), (1, K), (192, i_t * BT), (64, BT), (0, 1))
+                b_k = (b_k * tl.exp(b_gk_last4[:, None] - tl.load(p_gk, boundary_check=(0, 1)))).to(b_k.dtype)
+            b_h4 += tl.dot(b_k, b_v)
+
+    # Store final state
+    if STORE_FINAL_STATE:
+        p_ht = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))
+        tl.store(p_ht, b_h1.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
+        if K > 64:
+            p_ht = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0))
+            tl.store(p_ht, b_h2.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
+        if K > 128:
+            p_ht = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0))
+            tl.store(p_ht, b_h3.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
+        if K > 192:
+            p_ht = tl.make_block_ptr(ht_ptr, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0))
+            tl.store(p_ht, b_h4.to(p_ht.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_gated_delta_rule_fwd_h(
+    k: torch.Tensor,
+    w: torch.Tensor,
+    u: torch.Tensor,
+    g: Optional[torch.Tensor] = None,
+    gk: Optional[torch.Tensor] = None,
+    initial_state: Optional[torch.Tensor] = None,
+    output_final_state: bool = False,
+    chunk_size: int = 64,  # default:64
+    save_new_value: bool = True,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, T, H, K, V = *k.shape, u.shape[-1]
+    BT = chunk_size
+
+    chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None
+    # N: the actual number of sequences in the batch with either equal or variable lengths
+    if cu_seqlens is None:
+        N, NT, chunk_offsets = B, triton.cdiv(T, BT), None
+    else:
+        N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT)
+    assert K <= 256, "current kernel does not support head dimension larger than 256."
+
+    h = k.new_empty(B, NT, H, K, V).permute(0, 2, 1, 3, 4).contiguous()
+    final_state = k.new_empty(N, H, K, V, dtype=torch.float32) if output_final_state else None
+
+    BV = 128
+
+    v_new = torch.empty_like(u).permute(0, 2, 1, 3).contiguous() if save_new_value else None
+    k = k.permute(0, 2, 1, 3).contiguous()
+    w = w.permute(0, 2, 1, 3).contiguous()
+    u = u.permute(0, 2, 1, 3).contiguous()
+    g = g.permute(0, 2, 1).contiguous()
+    chunk_gated_delta_rule_fwd_kernel_h_blockdim64[(triton.cdiv(V, BV), N * H)](
+        k=k,
+        v=u,
+        w=w,
+        v_new=v_new,
+        g=g,
+        gk=gk,
+        h=h,
+        h0=initial_state,
+        ht=final_state,
+        cu_seqlens=cu_seqlens,
+        chunk_offsets=chunk_offsets,
+        T=T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BV=BV,
+        NT=NT,
+    )
+    h = h.permute(0, 2, 1, 3, 4).contiguous()
+    v_new = v_new.permute(0, 2, 1, 3).contiguous()
+    return h, v_new, final_state
+
+
+@triton.heuristics(
+    {
+        "USE_G": lambda args: args["g"] is not None,
+        "USE_GK": lambda args: args["gk"] is not None,
+        "USE_INITIAL_STATE": lambda args: args["dh0"] is not None,
+        "USE_FINAL_STATE_GRADIENT": lambda args: args["dht"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.autotune(
+    configs=get_autotune_config(multibuffer_list=(True, False)),
+    key=["H", "K", "V", "BT", "BV", "USE_G", "IS_VARLEN"],
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64(
+    q,
+    k,
+    w,
+    g,
+    gk,
+    dht,
+    dh0,
+    do,
+    dh,
+    dv,
+    dv2,
+    cu_seqlens,
+    chunk_offsets,
+    scale,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BV: tl.constexpr,
+    USE_G: tl.constexpr,
+    USE_GK: tl.constexpr,
+    USE_INITIAL_STATE: tl.constexpr,
+    USE_FINAL_STATE_GRADIENT: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    T_all = T
+    i_v, i_nh = tl.program_id(0), tl.program_id(1)
+    i_n, i_h = i_nh // H, i_nh % H
+    if IS_VARLEN:
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        T = eos - bos
+        NT = tl.cdiv(T, BT)
+        boh = tl.load(chunk_offsets + i_n).to(tl.int32)
+    else:
+        bos, eos = i_n * T, i_n * T + T
+        NT = tl.cdiv(T, BT)
+        boh = i_n * NT
+
+    b_dh1 = tl.zeros([64, BV], dtype=tl.float32)
+    if K > 64:
+        b_dh2 = tl.zeros([64, BV], dtype=tl.float32)
+    if K > 128:
+        b_dh3 = tl.zeros([64, BV], dtype=tl.float32)
+    if K > 192:
+        b_dh4 = tl.zeros([64, BV], dtype=tl.float32)
+
+    q += (bos * H + i_h) * K
+    k += (bos * H + i_h) * K
+    w += (bos * H + i_h) * K
+    do += (bos * H + i_h) * V
+    dv += (bos * H + i_h) * V
+    dv2 += (bos * H + i_h) * V
+    dh += (boh * H + i_h) * K * V
+    if USE_GK:
+        gk += (bos * H + i_h) * K
+
+    if USE_INITIAL_STATE:
+        dh0 += i_nh * K * V
+    if USE_FINAL_STATE_GRADIENT:
+        dht += i_nh * K * V
+
+    stride_v = H * V
+    stride_h = H * K * V
+    stride_k = H * K
+
+    if USE_FINAL_STATE_GRADIENT:
+        p_dht1 = tl.make_block_ptr(dht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))
+        b_dh1 += tl.load(p_dht1, boundary_check=(0, 1))
+        if K > 64:
+            p_dht2 = tl.make_block_ptr(dht, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0))
+            b_dh2 += tl.load(p_dht2, boundary_check=(0, 1))
+        if K > 128:
+            p_dht3 = tl.make_block_ptr(dht, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0))
+            b_dh3 += tl.load(p_dht3, boundary_check=(0, 1))
+        if K > 192:
+            p_dht4 = tl.make_block_ptr(dht, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0))
+            b_dh4 += tl.load(p_dht4, boundary_check=(0, 1))
+
+    for i_t in range(NT - 1, -1, -1):
+        p_dh1 = tl.make_block_ptr(dh + i_t * stride_h, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))
+        tl.store(p_dh1, b_dh1.to(p_dh1.dtype.element_ty), boundary_check=(0, 1))
+        if K > 64:
+            p_dh2 = tl.make_block_ptr(dh + i_t * stride_h, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0))
+            tl.store(p_dh2, b_dh2.to(p_dh2.dtype.element_ty), boundary_check=(0, 1))
+        if K > 128:
+            p_dh3 = tl.make_block_ptr(dh + i_t * stride_h, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0))
+            tl.store(p_dh3, b_dh3.to(p_dh3.dtype.element_ty), boundary_check=(0, 1))
+        if K > 192:
+            p_dh4 = tl.make_block_ptr(dh + i_t * stride_h, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0))
+            tl.store(p_dh4, b_dh4.to(p_dh4.dtype.element_ty), boundary_check=(0, 1))
+
+        last_idx = min((i_t + 1) * BT, T) - 1
+        if USE_G:
+            if IS_VARLEN:
+                bos_g = i_h * T_all + bos
+            else:
+                bos_g = (i_n * H + i_h) * T_all
+            bg_last = tl.load(g + bos_g + last_idx)
+            bg_last_exp = tl.exp(bg_last)
+            p_g = tl.make_block_ptr(
+                base=g + bos_g, shape=(T,), strides=(1,), offsets=(i_t * BT,), block_shape=(BT,), order=(0,)
+            )
+            b_g = tl.load(p_g, boundary_check=(0,))
+            b_g_exp = tl.exp(b_g)
+
+        p_dv = tl.make_block_ptr(dv, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        p_dv2 = tl.make_block_ptr(dv2, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+        p_do = tl.make_block_ptr(do, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+
+        b_do = tl.load(p_do, boundary_check=(0, 1))
+
+        # Update dv
+        p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0))
+        b_k = tl.load(p_k, boundary_check=(0, 1))
+        if USE_GK:
+            o_k1 = tl.arange(0, 64)
+            b_gk_last1 = tl.load(gk + last_idx * H * K + o_k1, mask=(o_k1 < K), other=0.0)
+        b_dv = tl.dot(b_k, b_dh1.to(b_k.dtype))
+
+        if K > 64:
+            p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            if USE_GK:
+                o_k2 = 64 + o_k1
+                b_gk_last2 = tl.load(gk + last_idx * H * K + o_k2, mask=(o_k2 < K), other=0.0)
+            b_dv += tl.dot(b_k, b_dh2.to(b_k.dtype))
+
+        if K > 128:
+            p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 128), (BT, 64), (1, 0))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            if USE_GK:
+                o_k3 = 128 + o_k1
+                b_gk_last3 = tl.load(gk + last_idx * H * K + o_k3, mask=(o_k3 < K), other=0.0)
+            b_dv += tl.dot(b_k, b_dh3.to(b_k.dtype))
+
+        if K > 192:
+            p_k = tl.make_block_ptr(k, (T, K), (stride_k, 1), (i_t * BT, 192), (BT, 64), (1, 0))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            if USE_GK:
+                o_k4 = 192 + o_k1
+                b_gk_last4 = tl.load(gk + last_idx * H * K + o_k4, mask=(o_k4 < K), other=0.0)
+            b_dv += tl.dot(b_k, b_dh4.to(b_k.dtype))
+
+        if USE_G:
+            m_t = (i_t * BT + tl.arange(0, BT)).to(tl.float32) < T
+            b_dv *= (m_t * tl.exp(bg_last - b_g))[:, None]
+        b_dv += tl.load(p_dv, boundary_check=(0, 1))
+
+        tl.store(p_dv2, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))
+        # Update dh
+        p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1))
+        p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1))
+        b_w = tl.load(p_w, boundary_check=(0, 1))
+        b_q = tl.load(p_q, boundary_check=(0, 1))
+        if USE_G:
+            b_dh1 *= bg_last_exp
+            b_q = b_q * b_g_exp[None, :]
+        if USE_GK:
+            b_dh1 *= tl.exp(b_gk_last1[:, None])
+        b_dh1 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))
+        if K > 64:
+            p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1))
+            p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1))
+            b_q = tl.load(p_q, boundary_check=(0, 1))
+            b_w = tl.load(p_w, boundary_check=(0, 1))
+            if USE_G:
+                b_dh2 *= bg_last_exp
+                b_q = b_q * b_g_exp[None, :]
+            if USE_GK:
+                b_dh2 *= tl.exp(b_gk_last2[:, None])
+            b_dh2 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))
+        if K > 128:
+            p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1))
+            p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1))
+            b_q = tl.load(p_q, boundary_check=(0, 1))
+            b_w = tl.load(p_w, boundary_check=(0, 1))
+            if USE_G:
+                b_dh3 *= bg_last_exp
+                b_q = b_q * b_g_exp[None, :]
+            if USE_GK:
+                b_dh3 *= tl.exp(b_gk_last3[:, None])
+            b_dh3 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))
+        if K > 192:
+            p_q = tl.make_block_ptr(q, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1))
+            p_w = tl.make_block_ptr(w, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1))
+            b_q = tl.load(p_q, boundary_check=(0, 1))
+            b_w = tl.load(p_w, boundary_check=(0, 1))
+            if USE_G:
+                b_dh4 *= bg_last_exp
+                b_q = b_q * b_g_exp[None, :]
+            if USE_GK:
+                b_dh4 *= tl.exp(b_gk_last4[:, None])
+            b_dh4 += tl.dot(b_q.to(b_q.dtype), b_do.to(b_q.dtype)) * scale - tl.dot(b_w, b_dv.to(b_w.dtype))
+
+    if USE_INITIAL_STATE:
+        p_dh0 = tl.make_block_ptr(dh0, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))
+        tl.store(p_dh0, b_dh1.to(p_dh0.dtype.element_ty), boundary_check=(0, 1))
+        if K > 64:
+            p_dh1 = tl.make_block_ptr(dh0, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0))
+            tl.store(p_dh1, b_dh2.to(p_dh1.dtype.element_ty), boundary_check=(0, 1))
+        if K > 128:
+            p_dh2 = tl.make_block_ptr(dh0, (K, V), (V, 1), (128, i_v * BV), (64, BV), (1, 0))
+            tl.store(p_dh2, b_dh3.to(p_dh2.dtype.element_ty), boundary_check=(0, 1))
+        if K > 192:
+            p_dh3 = tl.make_block_ptr(dh0, (K, V), (V, 1), (192, i_v * BV), (64, BV), (1, 0))
+            tl.store(p_dh3, b_dh4.to(p_dh3.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_gated_delta_rule_bwd_dhu(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    w: torch.Tensor,
+    do: torch.Tensor,
+    dv: torch.Tensor,
+    g: torch.Tensor | None = None,
+    gk: torch.Tensor | None = None,
+    h0: torch.Tensor | None = None,
+    dht: torch.Tensor | None = None,
+    scale: float | None = None,
+    cu_seqlens: torch.LongTensor | None = None,
+    chunk_size: int = 64,  # SY: remove this argument and force chunk size 64?
+    chunk_indices: torch.LongTensor | None = None,
+    use_exp2: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    B, T, H, K, V = *q.shape, do.shape[-1]
+    # N: the actual number of sequences in the batch with either equal or variable lengths
+    BT = 64
+    assert K <= 256, "current kernel does not support head dimension being larger than 256."
+
+    if chunk_indices is None and cu_seqlens is not None:
+        chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size)
+    if cu_seqlens is None:
+        N, NT, chunk_offsets = B, triton.cdiv(T, BT), None
+    else:
+        N, NT, chunk_offsets = len(cu_seqlens) - 1, len(chunk_indices), prepare_chunk_offsets(cu_seqlens, BT)
+
+    dh = q.new_empty(B, NT, H, K, V)
+    dh0 = torch.empty_like(h0, dtype=torch.float32) if h0 is not None else None
+    dv2 = torch.empty_like(dv)
+
+    BV = 128
+
+    g = g.permute(0, 2, 1).contiguous()
+
+    chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64[(triton.cdiv(V, BV), N * H)](
+        q=q,
+        k=k,
+        w=w,
+        g=g,
+        gk=gk,
+        dht=dht,
+        dh0=dh0,
+        do=do,
+        dh=dh,
+        dv=dv,
+        dv2=dv2,
+        cu_seqlens=cu_seqlens,
+        chunk_offsets=chunk_offsets,
+        scale=scale,
+        T=T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BV=BV,
+    )
+    return dh, dh0, dv2

src/llamafactory/third_party/triton/chunk_gated_delta_rule.py

@@ -0,0 +1,347 @@
+# Copyright 2025 the LlamaFactory team.
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+#
+# 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.
+
+import warnings
+from typing import Optional
+
+import torch
+
+from .chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h
+from .chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o
+from .chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
+from .cumsum import chunk_local_cumsum
+from .solve_tril import solve_tril
+from .utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
+from .wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd
+
+
+def chunk_gated_delta_rule_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor,
+    output_final_state: bool,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    chunk_size: int = 64,
+):
+    g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens, head_first=False)
+    # obtain WY representation. u is actually the new v.
+    A = chunk_scaled_dot_kkt_fwd(
+        k=k, g=g, beta=beta, cu_seqlens=cu_seqlens, chunk_size=chunk_size, output_dtype=torch.float32
+    )
+    A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)
+    w, u = recompute_w_u_fwd(
+        k=k,
+        v=v,
+        beta=beta,
+        A=A,
+        g=g,
+        cu_seqlens=cu_seqlens,
+    )
+    h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
+        k=k,
+        w=w,
+        u=u,
+        g=g,
+        initial_state=initial_state,
+        output_final_state=output_final_state,
+        chunk_size=chunk_size,
+        cu_seqlens=cu_seqlens,
+    )
+    o = chunk_fwd_o(
+        q=q,
+        k=k,
+        v=v_new,
+        h=h,
+        g=g,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_size=chunk_size,
+    )
+    return g, o, A, final_state
+
+
+def chunk_gated_delta_rule_bwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    A: torch.Tensor,
+    scale: float,
+    initial_state: torch.Tensor,
+    do: torch.Tensor,
+    dht: torch.Tensor,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    chunk_size: int = 64,
+):
+    w, u = recompute_w_u_fwd(
+        k=k,
+        v=v,
+        beta=beta,
+        A=A,
+        g=g,
+        cu_seqlens=cu_seqlens,
+    )
+    h, v_new, _ = chunk_gated_delta_rule_fwd_h(
+        k=k,
+        w=w,
+        u=u,
+        g=g,
+        initial_state=initial_state,
+        output_final_state=False,
+        cu_seqlens=cu_seqlens,
+        chunk_size=chunk_size,
+    )
+    dv = chunk_bwd_dv_local(
+        q=q,
+        k=k,
+        g=g,
+        do=do,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_size=chunk_size,
+    )
+    dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu(
+        q=q,
+        k=k,
+        w=w,
+        g=g,
+        h0=initial_state,
+        dht=dht,
+        do=do,
+        dv=dv,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_size=chunk_size,
+    )
+    dq, dk, dw, dg = chunk_bwd_dqkwg(
+        q=q,
+        k=k,
+        v=v_new,
+        w=w,
+        g=g,
+        h=h,
+        dv=dv,
+        do=do,
+        dh=dh,
+        chunk_size=chunk_size,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+    )
+    dk2, dv, db, dg2 = prepare_wy_repr_bwd(
+        k=k, v=v, beta=beta, g=g, A=A, dw=dw, du=dv, cu_seqlens=cu_seqlens, chunk_size=chunk_size
+    )
+    dk.add_(dk2)
+    dg.add_(dg2)
+    if dg.dtype != torch.float32:
+        raise ValueError(f"dg current type is {dg.dtype} , should be float32")
+    dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, head_first=False)
+    return dq, dk, dv, db, dg, dh0
+
+
+class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
+    @staticmethod
+    @input_guard
+    @autocast_custom_fwd
+    def forward(
+        ctx,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        g: torch.Tensor,
+        beta: torch.Tensor,
+        scale: float,
+        initial_state: torch.Tensor,
+        output_final_state: bool,
+        cu_seqlens: Optional[torch.LongTensor] = None,
+        use_qk_l2norm_in_kernel: bool = False,
+        chunk_size: int = 64,
+    ):
+        q_rstd, k_rstd = None, None
+        g, o, A, final_state = chunk_gated_delta_rule_fwd(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            scale=scale,
+            initial_state=initial_state,
+            output_final_state=output_final_state,
+            cu_seqlens=cu_seqlens,
+            chunk_size=chunk_size,
+        )
+        ctx.save_for_backward(q, q_rstd, k, k_rstd, v, g, beta, A, initial_state, cu_seqlens)
+        ctx.scale = scale
+        ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel
+        ctx.chunk_size = chunk_size
+        return o.to(q.dtype), final_state
+
+    @staticmethod
+    @input_guard
+    @autocast_custom_bwd
+    def backward(ctx, do: torch.Tensor, dht: torch.Tensor):
+        q, q_rstd, k, k_rstd, v, g, beta, A, initial_state, cu_seqlens = ctx.saved_tensors
+        dq, dk, dv, db, dg, dh0 = chunk_gated_delta_rule_bwd(
+            q=q,
+            k=k,
+            v=v,
+            g=g,
+            beta=beta,
+            A=A,
+            scale=ctx.scale,
+            initial_state=initial_state,
+            do=do,
+            dht=dht,
+            cu_seqlens=cu_seqlens,
+            chunk_size=ctx.chunk_size,
+        )
+        return dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta), None, dh0, None, None, None, None
+
+
+@torch.compiler.disable
+def chunk_gated_delta_rule(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    scale: float = None,
+    initial_state: torch.Tensor = None,
+    output_final_state: bool = False,
+    use_qk_l2norm_in_kernel: bool = False,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    chunk_size: int = 64,
+    head_first: bool = False,
+):
+    r"""Args:
+        q (torch.Tensor):
+            queries of shape `[B, T, H, K]`.
+        k (torch.Tensor):
+            keys of shape `[B, T, H, K]`.
+        v (torch.Tensor):
+            values of shape `[B, T, H, V]`.
+        g (torch.Tensor):
+            (forget) gating tensor (in log space!) of shape `[B, T, H]`.
+        beta (torch.Tensor):
+            betas of shape `[B, T, H]`.
+        scale (Optional[float]):
+            Scale factor for the RetNet attention scores.
+            If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
+        initial_state (Optional[torch.Tensor]):
+            Initial state of shape `[N, H, K, V]` for `N` input sequences.
+            For equal-length input sequences, `N` equals the batch size `B`.
+            Default: `None`.
+        output_final_state (Optional[bool]):
+            Whether to output the final state of shape `[N, H, K, V]`. Default: `False`.
+        use_qk_l2norm_in_kernel (bool):
+            Whether to apply L2norm to the q/k tensor internally. Default: `False`.
+        cu_seqlens (torch.LongTensor):
+            Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
+            consistent with the FlashAttention API.
+        head_first (Optional[bool]):
+            Whether the inputs are in the head-first format. Default: `False`.
+            This argument has been deprecated.
+
+    Returns:
+        o (torch.Tensor):
+            Outputs of shape `[B, T, H, V]`.
+        final_state (torch.Tensor):
+            Final state of shape `[N, H, K, V]` if `output_final_state=True` else `None`.
+
+    Examples::
+        >>> import torch
+        >>> import torch.nn.functional as F
+        >>> from einops import rearrange
+        >>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule
+        # inputs with equal lengths
+        >>> B, T, H, K, V = 4, 2048, 4, 512, 512
+        >>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')
+        >>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1)
+        >>> v = torch.randn(B, T, H, V, dtype=torch.bfloat16, device='cuda')
+        >>> beta = torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda').sigmoid()
+        >>> g = F.logsigmoid(torch.rand(B, T, H, dtype=torch.bfloat16, device='cuda'))
+        >>> h0 = torch.randn(B, H, K, V, dtype=torch.bfloat16, device='cuda')
+        >>> o, ht = chunk_gated_delta_rule(
+            q, k, v, g, beta,
+            initial_state=h0,
+            output_final_state=True
+        )
+        # for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
+        >>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))
+        # for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
+        >>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)
+        >>> o, ht = chunk_gated_delta_rule(
+            q, k, v, g, beta,
+            initial_state=h0,
+            output_final_state=True,
+            cu_seqlens=cu_seqlens
+        )
+    """  # noqa: D205
+    if q.dtype != k.dtype or k.dtype != v.dtype:
+        raise ValueError(
+            f"q current type is {q.dtype} , k current type is {k.dtype} ,v current type is {v.dtype} , they should are equal"
+        )
+    if q.dtype == torch.float32:
+        raise ValueError("ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16.")
+    if len(beta.shape) != 3:
+        raise ValueError(
+            f"beta current shape len is {len(beta.shape)}, beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise."
+        )
+
+    if head_first:
+        warnings.warn(
+            "head_first is deprecated and will be removed in a future version. "
+            "Please use head_first=False for now instead."
+        )
+    if not head_first and q.shape[1] < q.shape[2]:
+        warnings.warn(
+            f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). "
+            "This may indicate the inputs were passed in head-first format [B, H, T, ...] "
+            "when head_first=False was specified. "
+            "Please verify your input tensor format matches the expected shape [B, T, H, ...]."
+        )
+    if cu_seqlens is not None:
+        if q.shape[0] != 1:
+            raise ValueError(
+                f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
+                f"Please flatten variable-length inputs before processing."
+            )
+        if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1:
+            raise ValueError(
+                f"The number of initial states is expected to be equal to the number of input sequences, "
+                f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}."
+            )
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+
+    def l2norm(x: torch.FloatTensor, dim: int = -1, eps: float = 1e-6):
+        """This function is intended to align with the l2norm implementation in the FLA library."""
+        original_dtype = x.dtype
+        inv_norm = torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps)
+        # Counteract verl's autocast promotion (bf16 -> fp32) by restoring original dtype
+        return (x * inv_norm).to(original_dtype)
+
+    if use_qk_l2norm_in_kernel:
+        q = l2norm(q, dim=-1, eps=1e-6)
+        k = l2norm(k, dim=-1, eps=1e-6)
+
+    o, final_state = ChunkGatedDeltaRuleFunction.apply(
+        q, k, v, g, beta, scale, initial_state, output_final_state, cu_seqlens, False, chunk_size
+    )
+    return o, final_state

src/llamafactory/third_party/triton/chunk_o.py

@@ -0,0 +1,617 @@
+# Copyright 2025 the LlamaFactory team.
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+#
+# 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 typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from .utils import exp, prepare_chunk_indices, prepare_chunk_offsets
+
+
+@triton.heuristics(
+    {
+        "USE_G": lambda args: args["g"] is not None,
+        "USE_G_GAMMA": lambda args: args["g_gamma"] is not None,
+        "USE_DW": lambda args: args["dw"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_bwd_kernel_dqkwg(
+    q,
+    k,
+    v,
+    h,
+    g,
+    g_gamma,
+    do,
+    dh,
+    dq,
+    dk,
+    dg,
+    w,
+    dv,
+    dw,
+    cu_seqlens,
+    chunk_indices,
+    scale,
+    B: tl.constexpr,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_G: tl.constexpr,
+    USE_G_GAMMA: tl.constexpr,
+    USE_DW: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    gdiff,
+):
+    i_t, i_b = tl.program_id(0), tl.program_id(1)
+    T_max = T
+    if IS_VARLEN:
+        i_tg = i_t
+        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        total = B * T_max
+        T = eos - bos
+    else:
+        NT = tl.cdiv(T, BT)
+        i_tg = i_b * NT + i_t
+        bos, eos = i_b * T, i_b * T + T
+        total = B * T_max
+
+    NK = tl.cdiv(K, BK)
+    for i_k in range(NK):
+        if USE_G:
+            dg_k = dg + i_k * total * H
+
+        for i_h in range(H):
+            v_h = v + (bos * H + i_h) * V
+            do_h = do + (bos * H + i_h) * V
+            h_h = h + (i_tg * H + i_h).to(tl.int64) * K * V
+            dh_h = dh + (i_tg * H + i_h).to(tl.int64) * K * V
+            q_h = q + (bos * H + i_h) * K
+            k_h = k + (bos * H + i_h) * K
+            dq_h = dq + (bos * H + i_h) * K
+            dk_h = dk + (bos * H + i_h) * K
+
+            if USE_DW:
+                w_h = w + (bos * H + i_h) * K  # noqa: F841
+                dw_h = dw + (bos * H + i_h) * K
+                dv_h = dv + (bos * H + i_h) * V
+
+            if USE_G:
+                if IS_VARLEN:
+                    dg_h = dg_k + i_h * T_max + bos
+                    g_h = g + i_h * T_max + bos
+                else:
+                    dg_h = dg_k + (i_b * H + i_h) * T_max
+                    g_h = g + (i_b * H + i_h) * T_max
+                b_dg_last = tl.zeros(
+                    [
+                        1,
+                    ],
+                    dtype=tl.float32,
+                )
+
+            if USE_G_GAMMA:
+                b_gamma = tl.load(g_gamma + i_h)
+                b_g = b_gamma * (tl.arange(0, BT) + 1)
+                b_g_last = b_gamma * min(BT, T - i_t * BT)
+
+            b_dq = tl.zeros([BT, BK], dtype=tl.float32)
+            b_dk = tl.zeros([BT, BK], dtype=tl.float32)
+            b_ds = tl.zeros([BT, BT], dtype=tl.float32)
+            b_dw = tl.zeros([BT, BK], dtype=tl.float32) if USE_DW else None
+
+            for i_v in range(tl.cdiv(V, BV)):
+                p_v = tl.make_block_ptr(v_h, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+                p_do = tl.make_block_ptr(do_h, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+                p_h = tl.make_block_ptr(h_h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1))
+                p_dh = tl.make_block_ptr(dh_h, (V, K), (1, V), (i_v * BV, i_k * BK), (BV, BK), (0, 1))
+
+                b_v = tl.load(p_v, boundary_check=(0, 1))
+                b_do = tl.load(p_do, boundary_check=(0, 1))
+                b_h = tl.load(p_h, boundary_check=(0, 1))
+                b_dh = tl.load(p_dh, boundary_check=(0, 1))
+
+                if USE_G:
+                    b_dg_last += tl.sum(b_h * b_dh)
+
+                b_ds += tl.dot(b_do, tl.trans(b_v))
+                b_dq += tl.dot(b_do, b_h.to(b_do.dtype))
+                b_dk += tl.dot(b_v, b_dh.to(b_v.dtype))
+
+                if USE_DW:
+                    p_dv = tl.make_block_ptr(dv_h, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+                    b_dv = tl.load(p_dv, boundary_check=(0, 1))
+                    b_dw += tl.dot(b_dv.to(b_v.dtype), b_h.to(b_v.dtype))
+
+            if USE_DW:
+                p_dw = tl.make_block_ptr(dw_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+                tl.store(p_dw, -b_dw.to(p_dw.dtype.element_ty), boundary_check=(0, 1))
+
+            tl.debug_barrier()
+
+            p_q = tl.make_block_ptr(q_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+            p_k = tl.make_block_ptr(k_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+            b_q = tl.load(p_q, boundary_check=(0, 1))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+
+            p_dq = tl.make_block_ptr(dq_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+            p_dk = tl.make_block_ptr(dk_h, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+
+            o_t = i_t * BT + tl.arange(0, BT)
+            m_t = o_t < T
+            m_A = (o_t[:, None] >= o_t[None, :]) & (m_t[:, None] & m_t)
+
+            if USE_G:
+                b_dg = tl.zeros(
+                    [
+                        BT,
+                    ],
+                    dtype=tl.float32,
+                )
+                p_g = tl.make_block_ptr(g_h, (T,), (1,), (i_t * BT,), (BT,), (0,))
+                b_g = tl.load(p_g, boundary_check=(0,))
+                b_g_last = tl.load(g_h + (min(i_t * BT + BT, T) - 1) * 1)
+                b_dg_last *= tl.exp(b_g_last)
+
+                b_dq = b_dq * tl.exp(b_g)[:, None] * scale
+                b_dg += tl.sum(b_dq * b_q, axis=1)
+
+                b_dk = b_dk * tl.where(m_t, tl.exp(-b_g + b_g_last), 0)[:, None]
+                b_dg -= tl.sum(b_k * b_dk, axis=1)
+                b_dg_last += tl.sum(b_dk * b_k)
+
+                if IS_VARLEN:
+                    b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale
+                else:
+                    p_gdiff = tl.make_block_ptr(
+                        gdiff + i_b * H * NT * BT * BT + i_h * NT * BT * BT + i_t * BT * BT,
+                        (BT, BT),
+                        (BT, 1),
+                        (0, 0),
+                        (BT, BT),
+                        (1, 0),
+                    )
+                    gdiff_ = tl.load(p_gdiff)
+                    b_ds = b_ds * gdiff_ * scale
+
+                b_ds2 = b_ds * tl.dot(b_q, tl.trans(b_k))
+                b_dg += tl.sum(b_ds2, axis=1)
+                b_dg -= tl.sum(b_ds2, axis=0)
+
+                b_ds = b_ds.to(b_k.dtype)
+                b_dq += tl.dot(b_ds, b_k)
+                b_dk += tl.dot(tl.trans(b_ds), b_q)
+                p_dg = tl.make_block_ptr(dg_h, (T,), (1,), (i_t * BT,), (BT,), (0,))
+
+                last_index_local = min(BT, T - i_t * BT) - 1
+                if last_index_local >= 0:
+                    is_last_mask = tl.arange(0, BT) == last_index_local
+                    b_dg = tl.where(is_last_mask, b_dg + b_dg_last, b_dg)
+                else:
+                    pass
+
+                tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
+                tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
+                tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,))
+
+            elif USE_G_GAMMA:
+                b_dq = b_dq * exp(b_g)[:, None] * scale
+                b_dk = b_dk * tl.where(m_t, exp(-b_g + b_g_last), 0)[:, None]
+                b_ds = tl.where(m_A, b_ds * exp(b_g[:, None] - b_g[None, :]), 0) * scale
+                b_ds = b_ds.to(b_k.dtype)
+                b_dq += tl.dot(b_ds, b_k)
+                b_dk += tl.dot(tl.trans(b_ds), b_q)
+                tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
+                tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
+
+            else:
+                b_ds = tl.where(m_A, b_ds, 0)
+                b_ds = b_ds.to(b_k.dtype)
+                b_dq += tl.dot(b_ds, b_k)
+                b_dk += tl.dot(tl.trans(b_ds), b_q) * scale
+                b_dq *= scale
+                tl.store(p_dq, b_dq.to(p_dq.dtype.element_ty), boundary_check=(0, 1))
+                tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.heuristics(
+    {
+        "USE_G": lambda args: args["g"] is not None,
+        "USE_G_GAMMA": lambda args: args["g_gamma"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_bwd_kernel_dv_local(
+    q,
+    k,
+    g,
+    g_gamma,
+    do,
+    dv,
+    cu_seqlens,
+    chunk_indices,
+    scale,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_G: tl.constexpr,
+    USE_G_GAMMA: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_t, i_b = tl.program_id(0), tl.program_id(1)
+    T_max = T
+
+    if IS_VARLEN:
+        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    for i_h in range(H):
+        offset_kh = (bos * H + i_h) * K
+        offset_vh = (bos * H + i_h) * V
+
+        b_A = tl.zeros([BT, BT], dtype=tl.float32)
+        for i_k in range(tl.cdiv(K, BK)):
+            p_k = tl.make_block_ptr(k + offset_kh, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
+            p_q = tl.make_block_ptr(q + offset_kh, (K, T), (1, H * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1))
+            b_q = tl.load(p_q, boundary_check=(0, 1))
+            b_k = tl.load(p_k, boundary_check=(0, 1))
+            b_A += tl.dot(b_k, b_q)
+
+        if USE_G:
+            if IS_VARLEN:
+                offset_g = i_h * T_max + bos
+            else:
+                offset_g = i_b * H * T_max + i_h * T_max
+
+            p_g = tl.make_block_ptr(g + offset_g, (T,), (1,), (i_t * BT,), (BT,), (0,))
+            b_g = tl.load(p_g, boundary_check=(0,))
+
+        if USE_G_GAMMA:
+            b_gamma = tl.load(g_gamma + i_h)
+            b_g = b_gamma * (tl.arange(0, BT) + 1)
+
+        o_t = i_t * BT + tl.arange(0, BT)
+        m_t = o_t < T
+        m_A = (o_t[:, None] <= o_t[None, :]) & (m_t[:, None] & m_t)
+
+        if USE_G:
+            b_A = tl.where(m_A, b_A * tl.exp(b_g[None, :] - b_g[:, None]) * scale, 0).to(do.dtype.element_ty)
+        else:
+            b_A = tl.where(m_A, b_A * scale, 0).to(do.dtype.element_ty)
+
+        for i_v in range(tl.cdiv(V, BV)):
+            p_do = tl.make_block_ptr(do + offset_vh, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+            p_dv = tl.make_block_ptr(dv + offset_vh, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+            b_do = tl.load(p_do, boundary_check=(0, 1))
+            b_dv = tl.dot(b_A.to(b_do.dtype), b_do)
+            tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.heuristics(
+    {
+        "USE_G": lambda args: args["g"] is not None,
+        "USE_G_GAMMA": lambda args: args["g_gamma"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_fwd_kernel_o(
+    q,
+    k,
+    v,
+    h,
+    g,
+    g_gamma,
+    o,
+    cu_seqlens,
+    chunk_offsets,
+    scale,
+    T,
+    H: tl.constexpr,
+    N: tl.constexpr,
+    Hg: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_G: tl.constexpr,
+    USE_G_GAMMA: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    T_max = T
+    for i_v in range(tl.cdiv(V, BV)):
+        for i_n in range(N):
+            if IS_VARLEN:
+                bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+                T = eos - bos
+                NT = tl.cdiv(T, BT)
+                boh = tl.load(chunk_offsets + i_n).to(tl.int64)
+            else:
+                bos, eos = i_n * T, i_n * T + T
+                NT = tl.cdiv(T, BT)
+                boh = i_n * NT
+
+            core_id = tl.program_id(0)
+            total_cores = tl.num_programs(0)
+            base_chunks_per_pid = NT // total_cores
+            remainder = NT % total_cores
+
+            if core_id < remainder:
+                chunks_this_pid = base_chunks_per_pid + 1
+                start_idx = core_id * chunks_this_pid
+            else:
+                chunks_this_pid = base_chunks_per_pid
+                start_idx = core_id * base_chunks_per_pid + remainder
+
+            # offset calculation
+            for i_h in range(0, H):
+                q_offset = (bos * Hg + i_h // (H // Hg)) * K
+                k_offset = (bos * Hg + i_h // (H // Hg)) * K
+                v_offset = (bos * H + i_h) * V
+                o_offset = (bos * H + i_h) * V
+
+                for i_t in range(start_idx, start_idx + chunks_this_pid):
+                    i_tg = boh + i_t
+                    h_base = h + (i_tg * H + i_h).to(tl.int64) * K * V
+                    b_o = tl.zeros([BT, BV], dtype=tl.float32)
+                    b_A = tl.zeros([BT, BT], dtype=tl.float32)
+                    for i_k in range(tl.cdiv(K, BK)):
+                        p_q = tl.make_block_ptr(
+                            q + q_offset, (T, K), (Hg * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)
+                        )
+                        p_k = tl.make_block_ptr(
+                            k + k_offset, (K, T), (1, Hg * K), (i_k * BK, i_t * BT), (BK, BT), (0, 1)
+                        )
+                        p_h = tl.make_block_ptr(h_base, (K, V), (V, 1), (i_k * BK, i_v * BV), (BK, BV), (1, 0))
+                        b_q = tl.load(p_q, boundary_check=(0, 1))
+                        b_k = tl.load(p_k, boundary_check=(0, 1))
+                        b_h = tl.load(p_h, boundary_check=(0, 1))
+
+                        # [BT, BK] @ [BK, BV] -> [BT, BV]
+                        b_o += tl.dot(b_q, b_h)
+                        # [BT, BK] @ [BK, BT] -> [BT, BT]
+                        b_A += tl.dot(b_q, b_k)
+
+                    if USE_G:
+                        if IS_VARLEN:
+                            p_g = tl.make_block_ptr(g + bos + i_h * T_max, (T,), (1,), (i_t * BT,), (BT,), (0,))
+                        else:
+                            p_g = tl.make_block_ptr(g + bos * H + i_h * T_max, (T,), (1,), (i_t * BT,), (BT,), (0,))
+                        b_g = tl.load(p_g, boundary_check=(0,))
+                        b_o = b_o * exp(b_g)[:, None]
+                        b_A = b_A * exp(b_g[:, None] - b_g[None, :])
+                    if USE_G_GAMMA:
+                        b_gamma = tl.load(g_gamma + i_h)
+                        b_g = b_gamma * (tl.arange(0, BT) + 1)
+
+                    o_i = tl.arange(0, BT)
+                    m_A = o_i[:, None] >= o_i[None, :]
+                    b_A = tl.where(m_A, b_A, 0)
+
+                    p_v = tl.make_block_ptr(v + v_offset, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+                    p_o = tl.make_block_ptr(o + o_offset, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
+                    b_v = tl.load(p_v, boundary_check=(0, 1))
+
+                    # to fix mma -> mma layout conversion
+                    # already solved by triton v3.2 or higher
+                    b_o = b_o * scale + tl.dot(b_A.to(b_v.dtype), b_v) * scale
+                    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))
+
+
+def chunk_bwd_dqkwg(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    do: torch.Tensor,
+    h: torch.Tensor,
+    dh: torch.Tensor,
+    g: Optional[torch.Tensor] = None,
+    g_gamma: Optional[torch.Tensor] = None,
+    dv: Optional[torch.Tensor] = None,
+    w: Optional[torch.Tensor] = None,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    chunk_size: int = 64,
+    scale: float = 1.0,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+    chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    BK = 128 if cu_seqlens is None else 64
+    BV = 64
+    NK = triton.cdiv(K, BK)
+    dq = torch.empty_like(q)
+    dk = torch.empty_like(k)
+    g = g.transpose(1, 2).contiguous()
+    dg = torch.empty(NK, *g.shape, dtype=torch.float32, device=g.device) if g is not None else None
+    dw = torch.empty_like(w) if w is not None else None
+    grid = (NT, B)
+
+    if cu_seqlens is None:
+        if NT * BT == T:
+            g_ = g.reshape(B, H, NT, BT)
+            g_diff = g_[:, :, :, :, None] - g_[:, :, :, None, :]
+            g_diff = g_diff.clamp(-60, 60).exp()
+            g_diff[:, :, :] *= torch.tril(torch.ones(BT, BT), diagonal=0).to(g.device)
+        else:
+            diff = NT * BT - T
+            g_ = torch.cat((g, torch.zeros(B, H, diff).to(g.device)), dim=-1).reshape(B, H, NT, BT)
+            g_diff = g_[:, :, :, :, None] - g_[:, :, :, None, :]
+            g_diff = g_diff.clamp(-60, 60).exp()
+            g_diff[:, :, :] *= torch.tril(torch.ones(BT, BT), diagonal=0).to(g.device)
+            bias = torch.arange(0, BT).to(g.device)
+            o_t = (NT - 1) * BT + bias
+            m_t = o_t < T
+            m_A = m_t[:, None] & m_t
+            g_diff[:, :, -1] *= m_A
+    else:
+        g_diff = None
+
+    chunk_bwd_kernel_dqkwg[grid](
+        q=q,
+        k=k,
+        v=v,
+        h=h,
+        g=g,
+        g_gamma=g_gamma,
+        do=do,
+        dh=dh,
+        dv=dv,
+        w=w,
+        dw=dw,
+        dq=dq,
+        dk=dk,
+        dg=dg,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        scale=scale,
+        B=B,
+        T=T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+        gdiff=g_diff,
+    )
+
+    if dg is not None:
+        dg = dg.sum(0)
+        dg = dg.transpose(1, 2).contiguous()
+    return dq, dk, dw, dg
+
+
+def chunk_bwd_dv_local(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    do: torch.Tensor,
+    g: Optional[torch.Tensor] = None,
+    g_gamma: Optional[torch.Tensor] = None,
+    scale: float = None,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    chunk_size: int = 64,
+) -> torch.Tensor:
+    B, T, H, K, V = *k.shape, do.shape[-1]
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+    chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+
+    BK = 128
+    BV = 128
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    g = g.transpose(1, 2).contiguous()
+    dv = torch.empty_like(do)
+    grid = (NT, B)
+    chunk_bwd_kernel_dv_local[grid](
+        q=q,
+        k=k,
+        g=g,
+        g_gamma=g_gamma,
+        do=do,
+        dv=dv,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        scale=scale,
+        T=T,
+        H=H,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+    )
+    return dv
+
+
+def chunk_fwd_o(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    h: torch.Tensor,
+    g: Optional[torch.Tensor] = None,
+    g_gamma: Optional[torch.Tensor] = None,
+    scale: Optional[float] = None,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    chunk_size: int = 64,
+) -> torch.Tensor:
+    B, T, Hg, K, V = *q.shape, v.shape[-1]
+    H = v.shape[-2]
+    BT = min(chunk_size, max(16, triton.next_power_of_2(T)))
+    chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)  # noqa: F841
+    if scale is None:
+        scale = k.shape[-1] ** -0.5
+
+    o = torch.empty_like(v)
+    if cu_seqlens is None:
+        N, chunk_offsets = B, None
+    else:
+        N, chunk_offsets = (
+            len(cu_seqlens) - 1,
+            prepare_chunk_offsets(cu_seqlens, BT),
+        )
+
+    def grid(meta):
+        return (triton.cdiv(V, meta["BV"]), N * H)
+
+    g = g.transpose(1, 2).contiguous()
+    h = h.contiguous()
+    CV_kernel_num = 24
+    chunk_fwd_kernel_o[(CV_kernel_num,)](
+        q,
+        k,
+        v,
+        h,
+        g,
+        g_gamma,
+        o,
+        cu_seqlens,
+        chunk_offsets,
+        scale,
+        T=T,
+        H=H,
+        N=N,
+        Hg=Hg,
+        K=K,
+        V=V,
+        BT=BT,
+        BK=128,
+        BV=128,
+    )
+    return o
+
+
+bwd_chunk_dqkwg = chunk_bwd_dqkwg
+bwd_chunk_dv_local = chunk_bwd_dv_local

src/llamafactory/third_party/triton/chunk_scaled_dot_kkt.py

@@ -0,0 +1,359 @@
+# Copyright 2025 the LlamaFactory team.
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+#
+# 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 typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from .utils import prepare_chunk_indices
+
+
+@triton.heuristics(
+    {
+        "USE_G": lambda args: args["g"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_scaled_dot_kkt_fwd_kernel(
+    k,
+    g,
+    beta,
+    A,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    USE_G: tl.constexpr,
+    NT,
+    B,
+    TOTAL_TASKS,
+):
+    core_id = tl.program_id(0)
+    num_blocks = tl.num_programs(0)
+    T_max = T
+
+    base_tasks_per_block = TOTAL_TASKS // num_blocks
+    remainder_tasks = TOTAL_TASKS % num_blocks
+
+    if core_id < remainder_tasks:
+        tasks_this_core = base_tasks_per_block + 1
+        start_idx = core_id * tasks_this_core
+    else:
+        tasks_this_core = base_tasks_per_block
+        start_idx = core_id * base_tasks_per_block + remainder_tasks
+
+    for idx in range(start_idx, start_idx + tasks_this_core):
+        i_b = idx // NT
+        local_idx = idx % NT
+
+        if IS_VARLEN:
+            i_n = tl.load(chunk_indices + local_idx * 2).to(tl.int32)
+            i_t = tl.load(chunk_indices + local_idx * 2 + 1).to(tl.int32)
+            bos = tl.load(cu_seqlens + i_n).to(tl.int32)
+            eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+            T_local = eos - bos
+        else:
+            bos, eos = 0, T
+            i_t = local_idx
+            T_local = T
+
+        for i_h in range(H):
+            k_batch_off = i_b * T_max * H * K
+            beta_batch_off = i_b * H * T_max
+            g_batch_off = i_b * H * T_max
+            A_batch_off = i_b * T_max * H * BT
+
+            p_beta = tl.make_block_ptr(
+                beta + beta_batch_off + bos + i_h * T_max, (T_local,), (1,), (i_t * BT,), (BT,), (0,)
+            )
+            b_beta = tl.load(p_beta, boundary_check=(0,))
+
+            b_A = tl.zeros([BT, BT], dtype=tl.float32)
+            for i_k in range(tl.cdiv(K, BK)):
+                p_k = tl.make_block_ptr(
+                    k + k_batch_off + (bos * H + i_h) * K,
+                    (T_local, K),
+                    (H * K, 1),
+                    (i_t * BT, i_k * BK),
+                    (BT, BK),
+                    (1, 0),
+                )
+                b_k = tl.load(p_k, boundary_check=(0, 1))
+                dot_product = tl.dot(b_k, tl.trans(b_k))
+
+                o_t = i_t * BT + tl.arange(0, BT)
+                o_t = o_t.to(tl.float32)
+                T_mask = (o_t < T_local).to(tl.float32)
+
+                row_indices = tl.arange(0, BT)[:, None]
+                col_indices = tl.arange(0, BT)[None, :]
+                tril_mask = (row_indices > col_indices).to(tl.float32)
+                tril_mask = tril_mask * T_mask[:, None]
+                masked_dot = dot_product * tril_mask
+                b_A += masked_dot
+
+            if USE_G:
+                p_g = tl.make_block_ptr(
+                    g + g_batch_off + bos + i_h * T_max, (T_local,), (1,), (i_t * BT,), (BT,), (0,)
+                )
+                b_g = tl.load(p_g, boundary_check=(0,))
+                b_g_diff = b_g[:, None] - b_g[None, :]
+                b_g_diff = tl.minimum(tl.maximum(b_g_diff, -50.0), 50.0)
+                b_A *= tl.exp(b_g_diff)
+            b_A *= b_beta[:, None]
+
+            p_A = tl.make_block_ptr(
+                A + A_batch_off + (bos * H + i_h) * BT, (T_local, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0)
+            )
+            tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
+@triton.autotune(configs=[triton.Config({"BK": BK}) for BK in [32, 64]], key=["BC"])
+@triton.jit(do_not_specialize=["T"])
+def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter(
+    k,
+    g,
+    beta,
+    A,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    BC: tl.constexpr,
+    BK: tl.constexpr,
+    NC: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_t, i_c, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_i, i_j = i_c // NC, i_c % NC
+
+    for i_h in range(H):
+        if IS_VARLEN:
+            i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
+            bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+            T_val = eos - bos
+        else:
+            bos, eos = i_b * T, i_b * T + T
+            T_val = T
+
+        should_compute = (i_t * BT + i_i * BC < T_val) and (i_i > i_j)
+
+        if should_compute:
+            k_ptr = k + (bos * H + i_h) * K
+            g_ptr = g + (bos * H + i_h) * K
+            A_ptr = A + (bos * H + i_h) * BT
+
+            p_beta = tl.make_block_ptr(beta + bos * H + i_h, (T_val,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,))
+            b_beta = tl.load(p_beta, boundary_check=(0,))
+
+            b_A = tl.zeros([BC, BC], dtype=tl.float32)
+            for i_k in range(tl.cdiv(K, BK)):
+                p_k = tl.make_block_ptr(
+                    k_ptr, (T_val, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)
+                )
+                p_g = tl.make_block_ptr(
+                    g_ptr, (T_val, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)
+                )
+                b_kt = tl.make_block_ptr(
+                    k_ptr, (K, T_val), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)
+                )
+                p_gk = tl.make_block_ptr(
+                    g_ptr, (K, T_val), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)
+                )
+
+                o_k = i_k * BK + tl.arange(0, BK)
+                m_k = o_k < K
+                b_gn = tl.load(g_ptr + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0)
+                b_g = tl.load(p_g, boundary_check=(0, 1))
+                b_k = tl.load(p_k, boundary_check=(0, 1)) * tl.exp(b_g - b_gn[None, :])
+                b_gk = tl.load(p_gk, boundary_check=(0, 1))
+                b_kt = tl.load(b_kt, boundary_check=(0, 1)) * tl.exp(b_gn[:, None] - b_gk)
+                b_A += tl.dot(b_k, b_kt)
+            b_A *= b_beta[:, None]
+
+            p_A = tl.make_block_ptr(A_ptr, (T_val, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0))
+            tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1))
+
+
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
+@triton.jit(do_not_specialize=["T"])
+def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra(
+    k,
+    g,
+    beta,
+    A,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    BC: tl.constexpr,
+    BK: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    i_t, i_i, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+
+    for i_h in range(H):
+        if IS_VARLEN:
+            i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
+            bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+            T_val = eos - bos
+        else:
+            bos, eos = i_b * T, i_b * T + T
+            T_val = T
+
+        should_compute = i_t * BT + i_i * BC < T_val
+
+        if should_compute:
+            o_i = tl.arange(0, BC)
+            o_k = tl.arange(0, BK)
+            m_k = o_k < K
+            m_A = (i_t * BT + i_i * BC + o_i) < T_val
+            o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC
+
+            p_k = tl.make_block_ptr(
+                k + (bos * H + i_h) * K, (T_val, K), (H * K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)
+            )
+            p_g = tl.make_block_ptr(
+                g + (bos * H + i_h) * K, (T_val, K), (H * K, 1), (i_t * BT + i_i * BC, 0), (BC, BK), (1, 0)
+            )
+            p_beta = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h
+
+            b_k = tl.load(p_k, boundary_check=(0, 1)) * tl.load(p_beta, mask=m_A, other=0)[:, None]
+            b_g = tl.load(p_g, boundary_check=(0, 1))
+
+            p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k
+            p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k
+
+            for j in range(0, min(BC, T_val - i_t * BT - i_i * BC)):
+                b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32)
+                b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32)
+                b_A = tl.sum(b_k * b_kt[None, :] * tl.exp(b_g - b_gk[None, :]), 1)
+                # 转化成f32
+                o_i_tmp = o_i.to(tl.float32)
+                b_A = tl.where(o_i_tmp > j, b_A, 0.0)
+
+                tl.store(A + o_A + j, b_A, mask=m_A)
+                p_kt += H * K
+                p_gk += H * K
+
+
+def chunk_scaled_dot_kkt_fwd(
+    k: torch.Tensor,
+    g: Optional[torch.Tensor] = None,
+    gk: Optional[torch.Tensor] = None,
+    beta: Optional[torch.Tensor] = None,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    chunk_size: int = 64,
+    output_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    r"""Compute beta * K * K^T.
+
+    Args:
+        k (torch.Tensor):
+            The key tensor of shape `[B, T, H, K]`.
+        beta (torch.Tensor):
+            The beta tensor of shape `[B, T, H]`.
+        g (torch.Tensor):
+            The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`.
+        gk (torch.Tensor):
+            The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`.
+        cu_seqlens (torch.LongTensor):
+            The cumulative sequence lengths of the input tensor.
+            Default: None
+        chunk_size (int):
+            The chunk size. Default: 64.
+        output_dtype (torch.dtype):
+            The dtype of the output tensor. Default: `torch.float32`
+
+    Returns:
+        beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size.
+    """
+    B, T, H, K = k.shape
+    BT = chunk_size
+    chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+    beta = beta.transpose(1, 2).contiguous()
+    g = g.transpose(1, 2).contiguous()
+    BK = 128
+    kernel_num = 24
+
+    if gk is None:
+        A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype)
+        chunk_scaled_dot_kkt_fwd_kernel[(kernel_num,)](
+            k=k,
+            g=g,
+            beta=beta,
+            A=A,
+            cu_seqlens=cu_seqlens,
+            chunk_indices=chunk_indices,
+            T=T,
+            H=H,
+            K=K,
+            BT=BT,
+            BK=BK,
+            NT=NT,
+            B=B,
+            TOTAL_TASKS=B * NT,
+        )
+        return A
+
+    BC = min(16, BT)
+    NC = triton.cdiv(BT, BC)
+    BK = max(triton.next_power_of_2(K), 16)
+    A = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype)
+    grid = (NT, NC * NC, B)
+    chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter[grid](
+        k=k,
+        g=gk,
+        beta=beta,
+        A=A,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        K=K,
+        BT=BT,
+        BC=BC,
+        NC=NC,
+    )
+
+    grid = (NT, NC, B)
+    chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra[grid](
+        k=k,
+        g=gk,
+        beta=beta,
+        A=A,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        K=K,
+        BT=BT,
+        BC=BC,
+        BK=BK,
+    )
+    return A

src/llamafactory/third_party/triton/cumsum.py

@@ -0,0 +1,147 @@
+# Copyright 2025 the LlamaFactory team.
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+#
+# 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 typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from .utils import prepare_chunk_indices
+
+
+@triton.heuristics(
+    {"HAS_SCALE": lambda args: args["scale"] is not None, "IS_VARLEN": lambda args: args["cu_seqlens"] is not None}
+)
+@triton.jit(do_not_specialize=["T"])
+def chunk_local_cumsum_scalar_kernel(
+    s,
+    o,
+    scale,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    B: tl.constexpr,
+    H: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    REVERSE: tl.constexpr,
+    HAS_SCALE: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    HEAD_FIRST: tl.constexpr,
+    CHUNK_SIZE: tl.constexpr = 64,
+):
+    i_block, i_b = tl.program_id(0), tl.program_id(1)
+    N_CHUNKS: tl.constexpr = BLOCK_T // CHUNK_SIZE
+
+    if IS_VARLEN:
+        i_s, i_block = (
+            tl.load(chunk_indices + i_block * 2).to(tl.int32),
+            tl.load(chunk_indices + i_block * 2 + 1).to(tl.int32),
+        )
+
+        bos, eos = tl.load(cu_seqlens + i_s).to(tl.int32), tl.load(cu_seqlens + i_s + 1).to(tl.int32)
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    ptr_s = tl.make_block_ptr(s + bos * H, (T, H), (H, 1), (i_block * BLOCK_T, 0), (BLOCK_T, H), (1, 0))
+    ptr_o = tl.make_block_ptr(o + bos * H, (T, H), (H, 1), (i_block * BLOCK_T, 0), (BLOCK_T, H), (1, 0))
+    b_s = tl.load(ptr_s, boundary_check=(0,)).to(tl.float32)
+    b_s = tl.reshape(b_s, (N_CHUNKS, CHUNK_SIZE, H))
+    b_s = tl.trans(b_s, (1, 0, 2))
+    b_o = tl.cumsum(b_s, axis=0)
+    if REVERSE:
+        b_z = tl.sum(b_s, axis=0)
+        b_o = -b_o + b_z[None] + b_s
+    if HAS_SCALE:
+        b_o *= scale
+    b_o = tl.trans(b_o, (1, 0, 2))
+    b_o = tl.reshape(b_o, (BLOCK_T, H))
+
+    tl.store(ptr_o, b_o.to(ptr_o.dtype.element_ty), boundary_check=(0,))
+    return
+
+
+def chunk_local_cumsum_scalar(
+    g: torch.Tensor,
+    chunk_size: int,
+    reverse: bool = False,
+    scale: float = None,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    head_first: bool = False,
+    output_dtype: Optional[torch.dtype] = torch.float,
+) -> torch.Tensor:
+    B, T, H = g.shape
+    if chunk_size != 2 ** (chunk_size.bit_length() - 1):
+        raise ValueError(f"chunk_size must be a power of 2, chunk_size is{chunk_size}")
+    # We adjust the tiling strategy to prevent overflow in in backward passes and context parallel scenarios
+    #  while maximizing UB utilization where possible.
+    # The tiling strategy is as follows:
+    # 1. BT must be greater than or equal to chunk_size.
+    # 2. UB estimation varies directly with H.
+    # 3. BT in reverse mode is smaller than in forward mode.
+    BT = max(chunk_size, triton.next_power_of_2((1 << 11 if reverse else 1 << 12) // H))
+    chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+    g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
+    grid = (NT, B)
+    chunk_local_cumsum_scalar_kernel[grid](
+        s=g_org,
+        o=g,
+        scale=scale,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        B=B,
+        H=H,
+        BLOCK_T=BT,
+        HEAD_FIRST=head_first,
+        REVERSE=reverse,
+        CHUNK_SIZE=chunk_size,
+    )
+    return g
+
+
+def chunk_local_cumsum(
+    g: torch.Tensor,
+    chunk_size: int,
+    reverse: bool = False,
+    scale: float = None,
+    cu_seqlens: Optional[torch.Tensor] = None,
+    head_first: bool = False,
+    output_dtype: Optional[torch.dtype] = torch.float,
+    **kwargs,
+) -> torch.Tensor:
+    if cu_seqlens is not None:
+        if g.shape[0] != 1:
+            raise ValueError(
+                f"Only batch size 1 is supported when cu_seqlens are provided, current size is{g.shape[0]}"
+            )
+    if len(g.shape) == 3:
+        return chunk_local_cumsum_scalar(
+            g=g,
+            chunk_size=chunk_size,
+            reverse=reverse,
+            scale=scale,
+            cu_seqlens=cu_seqlens,
+            head_first=head_first,
+            output_dtype=output_dtype,
+        )
+    else:
+        raise ValueError(
+            f"Unsupported input shape {g.shape}, "
+            f"which should be (B, T, H, D) if `head_first=False` "
+            f"or (B, H, T, D) otherwise"
+        )

src/llamafactory/third_party/triton/solve_tril.py

@@ -0,0 +1,272 @@
+# Copyright 2025 the LlamaFactory team.
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
+#
+# 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.
+
+
+import os
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from .utils import input_guard, make_tensor_descriptor, prepare_chunk_indices
+
+
+FLA_TRIL_PRECISION = os.environ.get("FLA_TRIL_PRECISION", "ieee")
+
+
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
+@triton.jit(do_not_specialize=["T", "TPP"])
+def solve_tril_16x16_kernel(
+    A,
+    Ai,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    TPP: tl.constexpr,
+    USE_TMA: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    DOT_PRECISION: tl.constexpr,
+):
+    pid_t, pid_bh = tl.program_id(0), tl.program_id(1)
+    i_b, i_h = pid_bh // H, pid_bh % H
+
+    base_t = pid_t * TPP
+
+    if IS_VARLEN:
+        i_n = tl.load(chunk_indices + base_t * 2).to(tl.int32)
+        bos = tl.load(cu_seqlens + i_n).to(tl.int32)
+        eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        T_eff = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+        T_eff = T
+
+    o_i = tl.arange(0, 16)  # noqa: F841
+    o_i_fp32 = tl.arange(0, 16).to(tl.float32)
+    m_A = o_i_fp32[:, None] > o_i_fp32[None, :]
+    m_I = o_i_fp32[:, None] == o_i_fp32[None, :]
+
+    A = A + (bos * H + i_h) * BT
+    Ai = Ai + (bos * H + i_h) * BT
+
+    for tpp in tl.static_range(0, TPP):
+        tile_t = base_t + tpp
+        tile_row = tile_t * 16
+
+        offset = (tile_t * 16) % BT
+
+        if not USE_TMA:
+            p_A = tl.make_block_ptr(A, (T_eff, BT), (H * BT, 1), (tile_row, offset), (16, 16), (1, 0))
+            b_A_raw = tl.load(p_A, boundary_check=(0, 1)).to(tl.float32)
+        else:
+            desc = make_tensor_descriptor(A, [T_eff, BT], [H * BT, 1], [16, 16])
+            desc_o = make_tensor_descriptor(Ai, [T_eff, 16], [H * 16, 1], [16, 16])
+            b_A_raw = desc.load([tile_row, offset]).to(tl.float32)
+
+        b_A_neg = -b_A_raw
+        b_A = b_A_neg * m_A
+        for i in range(2, min(16, T_eff - tile_row)):
+            slice_res = tl.extract_slice(b_A_neg, [i, 0], [1, 16], [1, 1])
+            b_a_val = tl.reshape(slice_res, (16,), can_reorder=True)
+            dot_prod = tl.sum(b_a_val[:, None] * b_A, 0)
+            b_a_update = b_a_val + dot_prod
+            b_A = tl.where((o_i_fp32 == i)[:, None], b_a_update, b_A)
+        b_A += m_I
+
+        if not USE_TMA:
+            p_Ai = tl.make_block_ptr(Ai, (T_eff, 16), (H * 16, 1), (tile_row, 0), (16, 16), (1, 0))
+            tl.store(p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+        else:
+            desc_o.store([tile_row, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne"))
+
+
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
+@triton.jit(do_not_specialize=["T", "TPP"])
+def merge_16x16_to_32x32_inverse_kernel(
+    A,
+    Ai,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    TPP: tl.constexpr,
+    USE_TMA: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    DOT_PRECISION: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    i_b, i_h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+
+    o_i = tl.arange(0, 16)
+    m_A = o_i[:, None] > o_i[None, :]
+    m_I = o_i[:, None] == o_i[None, :]
+    A += (bos * H + i_h) * BT
+    Ai += (bos * H + i_h) * BT
+
+    if not USE_TMA:
+        p_A_11 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0))
+        p_A_22 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0))
+        b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32)
+        b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32)
+    else:
+        desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16])
+        desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16])
+        b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32)
+        b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32)
+
+    b_Ai_11 = -tl.where(m_A, b_Ai_11, 0)
+    b_Ai_22 = -tl.where(m_A, b_Ai_22, 0)
+
+    for i in range(2, min(16, T - i_t * BT)):
+        b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i)
+        b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0)
+        b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11)
+    for i in range(16 + 2, min(32, T - i_t * BT)):
+        b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16)
+        b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0)
+        b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22)
+
+    b_Ai_11 += m_I
+    b_Ai_22 += m_I
+
+    if not USE_TMA:
+        p_A_21 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0))
+        b_A_21 = tl.load(p_A_21, boundary_check=(0, 1)).to(tl.float32)
+    else:
+        b_A_21 = desc.load([i_t * BT + 16, 0]).to(tl.float32)
+
+    b_Ai_21 = -tl.dot(tl.dot(b_Ai_22, b_A_21, input_precision=DOT_PRECISION), b_Ai_11, input_precision=DOT_PRECISION)
+
+    if not USE_TMA:
+        p_Ai_11 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0))
+        p_Ai_21 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0))
+        p_Ai_22 = tl.make_block_ptr(Ai, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0))
+        tl.store(p_Ai_11, b_Ai_11.to(p_Ai_11.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+        tl.store(p_Ai_22, b_Ai_22.to(p_Ai_22.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+        tl.store(p_Ai_21, b_Ai_21.to(p_Ai_21.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    else:
+        desc_o.store([i_t * BT + 0, 0], b_Ai_11.to(desc_o.dtype, fp_downcast_rounding="rtne"))
+        desc_o.store([i_t * BT + 16, 0], b_Ai_21.to(desc_o.dtype, fp_downcast_rounding="rtne"))
+        desc_o.store([i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne"))
+
+
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
+@triton.jit(do_not_specialize=["T"])
+def solve_tril_64x64_kernel(
+    A,
+    Ai,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    BT: tl.constexpr,
+    USE_TMA: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+    DOT_PRECISION: tl.constexpr,
+):
+    i_t, i_bh = tl.program_id(0), tl.program_id(1)
+    i_b, i_h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        T = eos - bos
+    else:
+        bos, eos = i_b * T, i_b * T + T
+    o_i = tl.arange(0, 64)
+    m_I = o_i[:, None] == o_i[None, :]
+
+    A = A + (bos * H + i_h) * BT
+    Ai = Ai + (bos * H + i_h) * 64
+
+    offset = (i_t * 64) % BT
+    if not USE_TMA:
+        p_A = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * 64, offset), (64, 64), (1, 0))
+        b_A = -tl.load(p_A, boundary_check=(0, 1)).to(tl.float32)
+    else:
+        desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [64, 64])
+        desc_o = make_tensor_descriptor(Ai, [T, 64], [H * 64, 1], [64, 64])
+        b_A = -desc.load([i_t * 64, offset]).to(tl.float32)
+
+    for i in range(2, min(64, T - i_t * 64)):
+        b_a = -tl.load(A + (i_t * 64 + i) * H * BT + o_i + offset)
+        b_a = b_a + tl.sum(b_a[:, None] * b_A, 0)
+        b_A = tl.where((o_i == i)[:, None], b_a, b_A)
+    b_A += m_I
+    if not USE_TMA:
+        p_Ai = tl.make_block_ptr(Ai, (T, 64), (H * 64, 1), (i_t * 64, 0), (64, 64), (1, 0))
+        tl.store(p_Ai, b_A.to(p_Ai.dtype.element_ty, fp_downcast_rounding="rtne"), boundary_check=(0, 1))
+    else:
+        desc_o.store([i_t * 64, 0], b_A.to(desc_o.dtype, fp_downcast_rounding="rtne"))
+
+
+@input_guard
+def solve_tril(
+    A: torch.Tensor, cu_seqlens: Optional[torch.Tensor] = None, output_dtype: torch.dtype = torch.float
+) -> torch.Tensor:
+    """Compute the inverse of the matrix I + A
+    A should be strictly lower triangular, i.e., A.triu() == 0.
+
+    Args:
+        A (torch.Tensor):
+            [B, T, H, BT], where BT should only be 16, 32, or 64.
+        cu_seqlens (torch.Tensor):
+            The cumulative sequence lengths of the input tensor. Default: `None`.
+        output_dtype (torch.dtype):
+            The dtype of the output tensor. Default: `torch.float`.
+            If `None`, the output dtype will be the same as the input dtype.
+
+    Returns:
+        (I + A)^-1 with the same shape as A
+    """  # noqa: D205
+    if A.shape[-1] not in [16, 32, 64]:
+        raise ValueError(f"A shape BT should in [16,32, 64], but current is {A.shape[-1]}")
+    output_dtype = A.dtype if output_dtype is None else output_dtype
+
+    B, T, H, BT = A.shape
+    chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
+
+    Ai = torch.zeros_like(A, dtype=output_dtype)
+
+    if BT == 16:
+        merge_fn = solve_tril_16x16_kernel
+    elif BT == 32:
+        merge_fn = merge_16x16_to_32x32_inverse_kernel
+    elif BT == 64:
+        merge_fn = solve_tril_64x64_kernel
+
+    merge_fn[NT, B * H](
+        A=A,
+        Ai=Ai,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        H=H,
+        BT=BT,
+        USE_TMA=False,
+        DOT_PRECISION=FLA_TRIL_PRECISION,
+    )
+    return Ai

src/llamafactory/third_party/triton/utils.py

@@ -0,0 +1,359 @@
+# Copyright 2025 the LlamaFactory team.
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+#
+# 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.
+
+import contextlib
+import functools
+import itertools
+import logging
+import os
+import warnings
+from collections.abc import Callable
+from enum import Enum
+from typing import Any, Optional
+
+import torch
+import triton
+import triton.language as tl
+import triton.language.extra.libdevice as tldevice
+import triton.runtime.driver as driver
+from packaging import version
+
+
+logger = logging.getLogger(__name__)
+
+FLA_CI_ENV = os.getenv("FLA_CI_ENV") == "1"
+
+
+def tensor_cache(fn: Optional[Callable[..., torch.Tensor]] = None, *, maxsize: int = 1) -> Any:
+    """A decorator that caches the most recent results of a function with tensor inputs.
+
+    This decorator will store the outputs of the decorated function for the most recent
+    set of input tensors, up to `maxsize` entries. If the function is called again with
+    the same input tensors, it will return the cached result.
+
+    When maxsize=1 (default), the behavior is identical to caching only the most recent result.
+    Can be used as @tensor_cache or @tensor_cache(maxsize=n).
+
+    Args:
+        fn (Callable[..., torch.Tensor], optional):
+            The function to be decorated when used without parentheses.
+        maxsize (int):
+            Maximum number of input combinations to cache. Default is 1.
+
+    Returns:
+        Callable[..., torch.Tensor]:
+        A wrapped version of the input function with caching.
+    """
+    if maxsize < 1:
+        raise ValueError("maxsize must be at least 1")
+
+    def _is_match(a: Any, b: Any) -> bool:
+        if isinstance(a, torch.Tensor) and isinstance(b, torch.Tensor):
+            return a is b
+        try:
+            return a == b
+        except Exception:
+            return a is b
+
+    def _make_wrapper(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]:
+        cache: list = []
+
+        @functools.wraps(fn)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            for i, (cached_args, cached_kwargs, cached_result) in enumerate(cache):
+                if len(args) == len(cached_args) and len(kwargs) == len(cached_kwargs):
+                    if all(_is_match(a, b) for a, b in zip(args, cached_args)) and all(
+                        k in cached_kwargs and _is_match(v, cached_kwargs[k]) for k, v in kwargs.items()
+                    ):
+                        if i != 0:
+                            cache.insert(0, cache.pop(i))
+                        return cached_result
+
+            result = fn(*args, **kwargs)
+            cache.insert(0, (args, kwargs, result))
+            if len(cache) > maxsize:
+                cache.pop()
+            return result
+
+        return wrapper
+
+    if fn is not None:
+        return _make_wrapper(fn)
+    return _make_wrapper
+
+
+@tensor_cache
+def prepare_lens(cu_seqlens: torch.LongTensor) -> torch.LongTensor:
+    return cu_seqlens[1:] - cu_seqlens[:-1]
+
+
+@tensor_cache(maxsize=3)
+def prepare_chunk_indices(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor:
+    indices = torch.cat([torch.arange(n) for n in triton.cdiv(prepare_lens(cu_seqlens), chunk_size).tolist()])
+    return torch.stack([indices.eq(0).cumsum(0) - 1, indices], 1).to(cu_seqlens)
+
+
+def get_abs_err(x, y):
+    return (x.detach() - y.detach()).flatten().abs().max().item()
+
+
+def get_err_ratio(x, y):
+    err = (x.detach() - y.detach()).flatten().square().mean().sqrt().item()
+    base = (x.detach()).flatten().square().mean().sqrt().item()
+    return err / (base + 1e-8)
+
+
+def assert_close(prefix, ref, tri, ratio, warning=False, err_atol=1e-6):
+    abs_atol = get_abs_err(ref, tri)
+    msg = f"{prefix:>16} diff: {abs_atol:.6f} ratio: {get_err_ratio(ref, tri):.6f}"
+    logger.info(msg)
+    error_rate = get_err_ratio(ref, tri)
+    if abs_atol <= err_atol:
+        return
+    if warning or (FLA_CI_ENV and (error_rate < 0.01 or abs_atol <= 0.3)):
+        if error_rate > ratio:
+            warnings.warn(msg)
+    else:
+        assert error_rate < ratio, msg
+
+
+if hasattr(triton.language, "_experimental_make_tensor_descriptor"):
+    # For Triton 3.3.x
+    make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor
+elif hasattr(triton.language, "make_tensor_descriptor"):
+    # For Triton 3.4.x and later
+    make_tensor_descriptor = triton.language.make_tensor_descriptor
+else:
+    """
+    Fallback implementation when TMA is not supported.
+    Returns None to indicate TMA descriptors are unavailable.
+    Just make triton compiler happy.
+    """
+
+    @triton.jit
+    def make_tensor_descriptor(
+        base,
+        shape,
+        strides,
+        block_shape,
+        _builder=None,
+    ):
+        return None
+
+
+@functools.cache
+def get_available_device() -> str:
+    try:
+        return triton.runtime.driver.active.get_current_target().backend
+    except BaseException:
+        _cpu_device_warning()
+        return "cpu"
+
+
+def map_triton_backend_to_torch_device() -> str:
+    backend = get_available_device()  # 'cuda' | 'hip' | 'xpu' | 'cpu' | ...
+    return {"cuda": "cuda", "hip": "cuda", "xpu": "xpu"}.get(backend, backend)
+
+
+device = get_available_device() if get_available_device() != "hip" else "cuda"
+device_torch_lib = getattr(torch, device)
+device_platform = get_available_device()
+is_amd = device_platform == "hip"
+is_nvidia = device_platform == "cuda"
+is_nvidia_hopper = is_nvidia and (
+    "NVIDIA H" in torch.cuda.get_device_name(0) or torch.cuda.get_device_capability()[0] >= 9
+)
+
+is_tf32_supported = is_nvidia and torch.cuda.get_device_capability(0)[0] >= 8
+is_tma_supported = (
+    (is_nvidia and torch.cuda.get_device_capability(0)[0] >= 9)
+    and os.environ.get("FLA_NO_USE_TMA", "0") != "1"
+    and (
+        hasattr(triton.language, "_experimental_make_tensor_descriptor")
+        or hasattr(triton.language, "make_tensor_descriptor")
+    )
+)
+
+if is_nvidia and not is_tf32_supported:
+    # Make old card happy, since triton will use tf32 by default.
+    # This is a workaround for old nvidia card.
+    os.environ["TRITON_F32_DEFAULT"] = "ieee"
+
+
+@functools.cache
+def check_pytorch_version(version_s: str = "2.4") -> bool:
+    return version.parse(torch.__version__) >= version.parse(version_s)
+
+
+if check_pytorch_version("2.4"):
+    device = "cuda" if device == "cpu" else device
+    autocast_custom_fwd = functools.partial(torch.amp.custom_fwd, device_type=device)
+    autocast_custom_bwd = functools.partial(torch.amp.custom_bwd, device_type=device)
+
+    def custom_device_ctx(index: int):
+        return device_torch_lib.device(index)
+else:
+    assert device == "cuda", "Only cuda device is supported for PyTorch version < 2.4.0."
+    autocast_custom_fwd = device_torch_lib.amp.custom_fwd
+    autocast_custom_bwd = device_torch_lib.amp.custom_bwd
+
+    def custom_device_ctx(index: int):
+        return torch.cuda.device(index)
+
+
+def input_guard(fn: Callable[..., torch.Tensor]) -> Callable[..., torch.Tensor]:
+    """A decorator to make sure all input tensors are contiguous and set the device based on input tensors."""
+
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        contiguous_args = (i if not isinstance(i, torch.Tensor) else i.contiguous() for i in args)
+        contiguous_kwargs = {k: (v if not isinstance(v, torch.Tensor) else v.contiguous()) for k, v in kwargs.items()}
+
+        tensor = None
+        for arg in args:
+            if isinstance(arg, torch.Tensor):
+                tensor = arg
+                break
+        if tensor is None:
+            for value in kwargs.values():
+                if isinstance(value, torch.Tensor):
+                    tensor = value
+                    break
+
+        if tensor is not None:
+            ctx = custom_device_ctx(tensor.device.index)
+        else:
+            ctx = contextlib.nullcontext()
+
+        with ctx:
+            return fn(*contiguous_args, **contiguous_kwargs)
+
+    return wrapper
+
+
+def _cpu_device_warning():
+    warnings.warn(("Triton is not supported on current platform, roll back to CPU."), stacklevel=1)
+
+
+@tensor_cache
+def prepare_chunk_offsets(cu_seqlens: torch.LongTensor, chunk_size: int) -> torch.LongTensor:
+    return torch.cat([cu_seqlens.new_tensor([0]), triton.cdiv(prepare_lens(cu_seqlens), chunk_size)]).cumsum(-1)
+
+
+if os.environ.get("FLA_USE_FAST_OPS", "0") == "1":
+    exp = tldevice.fast_expf
+    exp2 = tldevice.exp2
+    log = tldevice.fast_logf
+    log2 = tldevice.fast_log2f
+else:
+    exp = tl.exp
+    exp2 = tl.math.exp2
+    log = tl.log
+    log2 = tl.log2
+
+
+def get_all_max_shared_mem():
+    try:
+        return [
+            triton.runtime.driver.active.utils.get_device_properties(i)["max_shared_mem"]
+            for i in range(device_torch_lib.device_count())
+        ]
+    except BaseException:
+        _cpu_device_warning()
+        return [-1]
+
+
+class Backend(Enum):
+    ADA = 101376  # RTX 4090
+    AMPERE = 166912  # A100
+    HOPPER = 232448  # H100
+    DEFAULT = 102400  # Default
+
+    @classmethod
+    def get_shared_memory(cls, arch: str) -> int:
+        try:
+            return cls[arch.upper()].value
+        except KeyError:
+            return cls.DEFAULT.value
+
+
+@functools.cache
+def check_shared_mem(arch: str = "none", tensor_idx: int = 0) -> bool:
+    try:
+        device_shared_mem_list = get_all_max_shared_mem()
+        max_shared_memory = device_shared_mem_list[tensor_idx]
+        return max_shared_memory >= Backend.get_shared_memory(arch)
+    except Exception:
+        return False
+
+
+def get_autotune_config(
+    multibuffer_list: tuple = (False,),
+    unit_flag_list: tuple = (False,),
+    limit_auto_multi_buffer_only_for_local_buffer_list: tuple = (False,),
+    limit_auto_multi_buffer_of_local_buffer_list: tuple = ("no-l0c",),
+    set_workspace_multibuffer_list: tuple = (2, 4),
+    enable_hivm_auto_cv_balance_list: tuple = (True,),
+    tile_mix_vector_loop_num_list: tuple = (2, 4),
+    tile_mix_cube_loop_num_list: tuple = (2, 4),
+):
+    configs = []
+    for (
+        multibuffer,
+        unit_flag,
+        limit_auto_multi_buffer_only_for_local_buffer,
+        limit_auto_multi_buffer_of_local_buffer,
+    ) in itertools.product(
+        list(multibuffer_list),
+        list(unit_flag_list),
+        list(limit_auto_multi_buffer_only_for_local_buffer_list),
+        list(limit_auto_multi_buffer_of_local_buffer_list),
+    ):
+        base_config_dict = {
+            "multibuffer": multibuffer,
+            "unit_flag": unit_flag,
+            "limit_auto_multi_buffer_only_for_local_buffer": limit_auto_multi_buffer_only_for_local_buffer,
+            "limit_auto_multi_buffer_of_local_buffer": limit_auto_multi_buffer_of_local_buffer,
+        }
+
+        if limit_auto_multi_buffer_only_for_local_buffer:
+            configs.append(triton.Config(base_config_dict))
+        else:
+            for (
+                set_workspace_multibuffer,
+                enable_hivm_auto_cv_balance,
+                tile_mix_vector_loop,
+                tile_mix_cube_loop,
+            ) in itertools.product(
+                list(set_workspace_multibuffer_list),
+                list(enable_hivm_auto_cv_balance_list),
+                list(tile_mix_vector_loop_num_list),
+                list(tile_mix_cube_loop_num_list),
+            ):
+                full_config_dict = base_config_dict.copy()
+                full_config_dict.update(
+                    {
+                        "set_workspace_multibuffer": set_workspace_multibuffer,
+                        "enable_hivm_auto_cv_balance": enable_hivm_auto_cv_balance,
+                        "tile_mix_vector_loop": tile_mix_vector_loop,
+                        "tile_mix_cube_loop": tile_mix_cube_loop,
+                    }
+                )
+                configs.append(triton.Config(full_config_dict))
+    return configs
+
+
+def get_npu_properties():
+    return driver.active.utils.get_device_properties(torch.npu.current_device())

src/llamafactory/third_party/triton/wy_fast.py

@@ -0,0 +1,387 @@
+# Copyright 2025 the LlamaFactory team.
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+# Copyright (c) 2026, Huawei Technologies Co., Ltd.  All rights reserved.
+#
+# 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 typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from .utils import exp, prepare_chunk_indices
+
+
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
+@triton.jit(do_not_specialize=["T"])
+def prepare_wy_repr_bwd_kernel(
+    k,
+    v,
+    beta,
+    g,
+    A,
+    dw,
+    du,
+    dk,
+    dv,
+    dbeta,
+    dg,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    B,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    NT: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    core_id = tl.program_id(0)
+    total_cores = tl.num_programs(0)
+    T_max = T
+
+    base_chunks_per_pid = NT // total_cores
+    remainder_chunks = NT % total_cores
+
+    if core_id < remainder_chunks:
+        chunks_this_pid = base_chunks_per_pid + 1
+        start_idx = core_id * chunks_this_pid
+    else:
+        chunks_this_pid = base_chunks_per_pid
+        start_idx = core_id * chunks_this_pid + remainder_chunks
+
+    for idx in range(start_idx, start_idx + chunks_this_pid):
+        for i_b in range(B):
+            if IS_VARLEN:
+                i_n, i_t = (
+                    tl.load(chunk_indices + idx * 2).to(tl.int32),
+                    tl.load(chunk_indices + idx * 2 + 1).to(tl.int32),
+                )
+                bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+                T = eos - bos
+            else:
+                i_t = idx
+                bos, eos = i_b * T, i_b * T + T
+
+            o_t = i_t * BT + tl.arange(0, BT)
+            m_t = o_t < T
+            m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t)
+            for i_h in range(0, H):
+                if IS_VARLEN:
+                    offset = bos + i_h * T_max
+                else:
+                    offset = bos * H + i_h * T_max
+
+                p_beta = tl.make_block_ptr(beta + offset, (T,), (1,), (i_t * BT,), (BT,), (0,))
+                p_g = tl.make_block_ptr(g + offset, (T,), (1,), (i_t * BT,), (BT,), (0,))
+                p_A = tl.make_block_ptr(
+                    A + (bos * H + i_h) * BT, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1)
+                )
+
+                b_A = tl.load(p_A, boundary_check=(0, 1))
+                b_beta = tl.load(p_beta, boundary_check=(0,))
+                b_g = tl.load(p_g, boundary_check=(0,))
+                b_g_exp = tl.exp(b_g)
+
+                b_dbeta = tl.zeros([BT], dtype=tl.float32)
+                b_dA = tl.zeros([BT, BT], dtype=tl.float32)
+                b_dg = tl.zeros([BT], dtype=tl.float32)
+
+                for i_k in range(tl.cdiv(K, BK)):
+                    p_k = tl.make_block_ptr(
+                        k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)
+                    )
+                    p_dk = tl.make_block_ptr(
+                        dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)
+                    )
+                    p_dw = tl.make_block_ptr(
+                        dw + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)
+                    )
+                    b_k = tl.load(p_k, boundary_check=(0, 1))
+                    b_k_beta_g = (b_k * b_beta[:, None] * b_g_exp[:, None]).to(b_k.dtype)
+                    b_dw = tl.load(p_dw, boundary_check=(0, 1))
+                    b_dA += tl.dot(b_dw, tl.trans(b_k_beta_g))
+                    b_dk_beta_g = tl.dot(b_A, b_dw)
+                    b_dk = b_dk_beta_g * b_beta[:, None] * b_g_exp[:, None]
+                    b_dbeta += tl.sum(b_dk_beta_g * b_k * b_g_exp[:, None], 1)
+                    b_dg += tl.sum(b_dk_beta_g * b_k * b_g_exp[:, None] * b_beta[:, None], 1)
+                    tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
+
+                for i_v in range(tl.cdiv(V, BV)):
+                    p_v = tl.make_block_ptr(
+                        v + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)
+                    )
+                    p_dv = tl.make_block_ptr(
+                        dv + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)
+                    )
+                    p_du = tl.make_block_ptr(
+                        du + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)
+                    )
+                    b_v = tl.load(p_v, boundary_check=(0, 1))
+                    b_v_beta = (b_v * b_beta[:, None]).to(b_v.dtype)
+                    b_du = tl.load(p_du, boundary_check=(0, 1))
+                    b_dA += tl.dot(b_du, tl.trans(b_v_beta))
+                    b_dv_beta = tl.dot(b_A, b_du)
+                    b_dv = b_dv_beta * b_beta[:, None]
+                    b_dbeta += tl.sum(b_dv_beta * b_v, 1)
+                    tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))
+
+                b_dA = tl.where(m_A, b_dA, 0)
+                b_dA = tl.dot(b_dA.to(b_A.dtype), b_A)
+                b_dA = tl.dot(b_A, b_dA.to(b_A.dtype))
+                b_dA = tl.where(m_A, -b_dA * exp(b_g[:, None] - b_g[None, :]), 0)
+                b_dA = b_dA.to(k.dtype.element_ty)
+                b_A = tl.zeros([BT, BT], dtype=tl.float32)
+
+                for i_k in range(tl.cdiv(K, BK)):
+                    p_k = tl.make_block_ptr(
+                        k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)
+                    )
+                    p_dk = tl.make_block_ptr(
+                        dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)
+                    )
+                    b_k = tl.load(p_k, boundary_check=(0, 1))
+                    b_dk = tl.load(p_dk, boundary_check=(0, 1))
+                    b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype)
+                    b_A += tl.dot(b_k_beta, tl.trans(b_k))
+                    b_dk_beta = tl.dot(b_dA, b_k)
+                    b_dbeta += tl.sum(b_dk_beta * b_k, 1)
+                    b_dk += tl.dot(tl.trans(b_dA), b_k_beta)
+                    b_dk += b_dk_beta * b_beta[:, None]
+                    tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
+
+                b_dA_A = b_dA * b_A
+                b_dg += tl.sum(b_dA_A, axis=1) - tl.sum(b_dA_A, axis=0)
+                p_dg = tl.make_block_ptr(dg + offset, (T,), (1,), (i_t * BT,), (BT,), (0,))
+                p_dbeta = tl.make_block_ptr(dbeta + offset, (T,), (1,), (i_t * BT,), (BT,), (0,))
+                tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,))
+                tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,))
+
+
+@triton.heuristics(
+    {
+        "USE_G": lambda args: args["g"] is not None,
+        "USE_GK": lambda args: args["gk"] is not None,
+        "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
+    }
+)
+@triton.jit(do_not_specialize=["T"])
+def recompute_w_u_fwd_kernel(
+    k,
+    v,
+    beta,
+    w,
+    u,
+    A,
+    g,
+    gk,
+    cu_seqlens,
+    chunk_indices,
+    T_tmp,
+    B,
+    H: tl.constexpr,
+    K: tl.constexpr,
+    V: tl.constexpr,
+    NT: tl.constexpr,
+    BT: tl.constexpr,
+    BK: tl.constexpr,
+    BV: tl.constexpr,
+    USE_G: tl.constexpr,
+    USE_GK: tl.constexpr,
+    IS_VARLEN: tl.constexpr,
+):
+    core_id = tl.program_id(0)
+    total_cores = tl.num_programs(0)
+    T_max = T_tmp
+
+    base_chunks_per_pid = NT // total_cores
+    remainder_chunks = NT % total_cores
+
+    if core_id < remainder_chunks:
+        chunks_this_pid = base_chunks_per_pid + 1
+        start_idx = core_id * chunks_this_pid
+    else:
+        chunks_this_pid = base_chunks_per_pid
+        start_idx = core_id * chunks_this_pid + remainder_chunks
+
+    for idx in range(start_idx, start_idx + chunks_this_pid):
+        for i_b in range(B):
+            for i_h in range(0, H):
+                if IS_VARLEN:
+                    i_n, i_t = (
+                        tl.load(chunk_indices + idx * 2).to(tl.int32),
+                        tl.load(chunk_indices + idx * 2 + 1).to(tl.int32),
+                    )
+                    bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+                    offset = bos + i_h * T_max
+                    T = eos - bos
+                else:
+                    T = T_tmp
+                    i_t = idx
+                    bos, eos = i_b * T, i_b * T + T
+                    offset = bos * H + i_h * T_max
+
+                p_beta = tl.make_block_ptr(beta + offset, (T,), (1,), (i_t * BT,), (BT,), (0,))
+                b_beta = tl.load(p_beta, boundary_check=(0,))
+
+                p_A = tl.make_block_ptr(
+                    A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)
+                )
+                b_A = tl.load(p_A, boundary_check=(0, 1))
+
+                for i_v in range(tl.cdiv(V, BV)):
+                    p_v = tl.make_block_ptr(
+                        v + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)
+                    )
+                    p_u = tl.make_block_ptr(
+                        u + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)
+                    )
+                    b_v = tl.load(p_v, boundary_check=(0, 1))
+                    b_vb = (b_v * b_beta[:, None]).to(b_v.dtype)
+                    b_u = tl.dot(b_A, b_vb, allow_tf32=False)
+                    tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1))
+
+                if USE_G:
+                    p_g = tl.make_block_ptr(g + offset, (T,), (1,), (i_t * BT,), (BT,), (0,))
+                    b_g = tl.exp(tl.load(p_g, boundary_check=(0,)))
+
+                for i_k in range(tl.cdiv(K, BK)):
+                    p_k = tl.make_block_ptr(
+                        k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)
+                    )
+                    p_w = tl.make_block_ptr(
+                        w + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)
+                    )
+                    b_k = tl.load(p_k, boundary_check=(0, 1))
+                    b_kb = b_k * b_beta[:, None]
+                    if USE_G:
+                        b_kb *= b_g[:, None]
+                    if USE_GK:
+                        p_gk = tl.make_block_ptr(
+                            gk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0)
+                        )
+                        b_kb *= tl.exp(tl.load(p_gk, boundary_check=(0, 1)))
+                    b_w = tl.dot(b_A, b_kb.to(b_k.dtype))
+                    tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
+
+
+def recompute_w_u_fwd(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    beta: torch.Tensor,
+    A: torch.Tensor,
+    g: Optional[torch.Tensor] = None,
+    gk: Optional[torch.Tensor] = None,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    BT = A.shape[-1]
+    BK = 128
+    BV = 128
+
+    chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+    g = g.transpose(1, 2).contiguous() if g is not None else None
+    beta = beta.transpose(1, 2).contiguous()
+
+    w = torch.empty_like(k)
+    u = torch.empty_like(v)
+    cv_kernel_num = 24
+    recompute_w_u_fwd_kernel[(cv_kernel_num,)](
+        k=k,
+        v=v,
+        beta=beta,
+        w=w,
+        u=u,
+        A=A,
+        g=g,
+        gk=gk,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T_tmp=T,
+        B=B,
+        H=H,
+        K=K,
+        V=V,
+        NT=NT,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+    )
+    return w, u
+
+
+def prepare_wy_repr_bwd(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    g: torch.Tensor,
+    beta: torch.Tensor,
+    A: torch.Tensor,
+    dw: torch.Tensor,
+    du: torch.Tensor,
+    cu_seqlens: Optional[torch.LongTensor],
+    chunk_size: int = 64,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    B, T, H, K, V = *k.shape, v.shape[-1]
+    BT = chunk_size
+    chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+    BK = 128
+    BV = 128
+    beta = beta.transpose(1, 2).contiguous()
+    g = g.transpose(1, 2).contiguous()
+
+    dk = torch.empty_like(k)
+    dv = torch.empty_like(v)
+    dbeta = torch.empty_like(beta)
+    dg = torch.empty_like(g)
+
+    cv_kernel_num = 24
+    prepare_wy_repr_bwd_kernel[(cv_kernel_num,)](
+        k=k,
+        v=v,
+        beta=beta,
+        g=g,
+        A=A,
+        dw=dw,
+        du=du,
+        dk=dk,
+        dv=dv,
+        dbeta=dbeta,
+        dg=dg,
+        cu_seqlens=cu_seqlens,
+        chunk_indices=chunk_indices,
+        T=T,
+        B=B,
+        H=H,
+        K=K,
+        V=V,
+        NT=NT,
+        BT=BT,
+        BK=BK,
+        BV=BV,
+    )
+
+    dbeta = dbeta.transpose(1, 2).contiguous()
+    dg = dg.transpose(1, 2).contiguous()
+
+    return dk, dv, dbeta, dg
+
+
+bwd_prepare_wy_repr = prepare_wy_repr_bwd
+
+fwd_recompute_w_u = recompute_w_u_fwd

src/llamafactory/train/hyper_parallel/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .workflow import run_sft
+from .workflow import run_pt, run_sft
 
 
-__all__ = ["run_sft"]
+__all__ = ["run_pt", "run_sft"]

src/llamafactory/train/hyper_parallel/trainer.py

@@ -0,0 +1,409 @@
+# 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.
+
+"""HyperParallel distributed trainer for LlamaFactory."""
+
+import logging
+import os
+import types
+from contextlib import nullcontext
+from functools import partial
+from typing import Any, Optional
+
+import torch
+from hyper_parallel.integration.llamafactory import (
+    HSDPModule,
+    HyperParallelArguments,
+    export_to_hf_format,
+    fsdp2_prepare_model,
+    hsdp_sync_stream,
+    load_hsdp_model,
+    load_hsdp_optimizer_and_scheduler,
+    save_hsdp_checkpoint,
+    wrap_optimizer_with_skip_dtensor_dispatch,
+)
+from hyper_parallel.integration.llamafactory import (
+    clip_grad_norm_ as hp_clip_grad_norm_,
+)
+from hyper_parallel.integration.llamafactory.context_parallel import (
+    cp_prepare_model,
+    get_cp_rank,
+    get_dp_rank,
+    shard_inputs_for_cp,
+)
+from hyper_parallel.platform import get_platform
+from torch import nn
+
+from ..sft.trainer import CustomSeq2SeqTrainer
+
+
+logger = logging.getLogger(__name__)
+
+
+class _CPBatchRepeatedBatchSampler(torch.utils.data.BatchSampler):
+    """Repeat logical batches so Accelerate shards CP peers onto the same samples."""
+
+    def __init__(self, sampler, batch_size: int, drop_last: bool, repeat_factor: int, logical_group_size: int):
+        super().__init__(sampler, batch_size, drop_last)
+        self.repeat_factor = repeat_factor
+        self.logical_group_size = logical_group_size
+
+    def __len__(self):
+        logical_length = super().__len__()
+        if not self.drop_last and logical_length > 0:
+            logical_length = _ceil_div(logical_length, self.logical_group_size) * self.logical_group_size
+        return logical_length * self.repeat_factor
+
+    def __iter__(self):
+        initial_data = []
+        logical_count = 0
+        pad_cursor = 0
+        max_initial_data = self.batch_size * self.logical_group_size
+
+        def collect_initial_data(batch):
+            if len(initial_data) < max_initial_data:
+                initial_data.extend(batch[: max_initial_data - len(initial_data)])
+
+        def get_padding_item():
+            nonlocal pad_cursor
+            item = initial_data[pad_cursor % len(initial_data)]
+            pad_cursor += 1
+            return item
+
+        def pad_batch(batch):
+            batch = list(batch)
+            if self.drop_last or len(batch) == self.batch_size:
+                return batch
+
+            while len(batch) < self.batch_size:
+                batch.append(get_padding_item())
+            return batch
+
+        def make_padding_batch():
+            return [get_padding_item() for _ in range(self.batch_size)]
+
+        def repeat_batch(batch):
+            for _ in range(self.repeat_factor):
+                yield list(batch)
+
+        for batch in super().__iter__():
+            collect_initial_data(batch)
+            batch = pad_batch(batch)
+            logical_count += 1
+            yield from repeat_batch(batch)
+
+        if self.drop_last or logical_count == 0:
+            return
+
+        while logical_count % self.logical_group_size != 0:
+            logical_count += 1
+            yield from repeat_batch(make_padding_batch())
+
+
+class _CPDataLoaderLengthProxy:
+    """Keep baseline logical dataloader length while yielding CP-repeated batches."""
+
+    def __init__(self, dataloader, logical_length: int):
+        self._dataloader = dataloader
+        self._logical_length = logical_length
+
+    def __iter__(self):
+        return iter(self._dataloader)
+
+    def __len__(self):
+        return self._logical_length
+
+    def __getattr__(self, name):
+        return getattr(self._dataloader, name)
+
+
+def _ceil_div(numerator: int, denominator: int) -> int:
+    return (numerator + denominator - 1) // denominator
+
+
+class HyperParallelTrainer(CustomSeq2SeqTrainer):
+    """Trainer that replaces Accelerate FSDP2 with HyperParallel fully_shard.
+
+    Inherits CustomSeq2SeqTrainer for training algorithm logic (loss, metrics,
+    prediction, sampler, etc.) and only overrides HSDP-specific behavior.
+    """
+
+    def __init__(
+        self,
+        hp_args: HyperParallelArguments,
+        finetuning_args=None,
+        processor=None,
+        ref_model: Optional[nn.Module] = None,
+        **kwargs,
+    ):
+        self._hp_args = hp_args
+
+        # Let CustomSeq2SeqTrainer handle everything except ref_model —
+        # Custom would prepare it with accelerate's fsdp2_prepare_model,
+        # but we need HP's version instead.
+        super().__init__(
+            finetuning_args=finetuning_args,
+            processor=processor,
+            ref_model=None,
+            **kwargs,
+        )
+
+        if not getattr(self.accelerator, "is_fsdp2", False):
+            raise ValueError("HyperParallel trainer requires Accelerate FSDP2 mode to be enabled.")
+
+        self._cp_size = hp_args.cp_size
+        self._cp_rank = get_cp_rank(hp_args) if self._cp_size > 1 else 0
+        self._dp_rank = get_dp_rank(hp_args) if self._cp_size > 1 else get_platform().get_rank()
+
+        # Prepare ref_model with the same CP + HSDP path as the train model.
+        self.ref_model = ref_model
+        if self.ref_model is not None:
+            self.ref_model = self._prepare_model_for_hyper_parallel(self.ref_model)
+
+        self._orig_accelerator_clip_grad_norm = self.accelerator.clip_grad_norm_
+        self._orig_fsdp2_prepare_model = None
+        self._accelerator_patches_active = False
+
+    def _prepare_model_for_hyper_parallel(self, model: nn.Module) -> nn.Module:
+        """Apply CP runtime hooks before delegating to HyperParallel FSDP2 preparation."""
+        if self._cp_size > 1:
+            model = cp_prepare_model(model, self.accelerator, self._hp_args)
+        return fsdp2_prepare_model(self.accelerator, model, self._hp_args)
+
+    def _activate_accelerator_patches(self) -> None:
+        """Patch Accelerate to use HyperParallel fsdp2_prepare_model and clip_grad_norm_."""
+        if self._accelerator_patches_active:
+            return
+
+        import accelerate.accelerator as acc_module  # pylint: disable=C0415
+
+        self._orig_fsdp2_prepare_model = acc_module.fsdp2_prepare_model
+
+        def _hp_fsdp2_prepare_model(accelerator, model):
+            return self._prepare_model_for_hyper_parallel(model)
+
+        acc_module.fsdp2_prepare_model = _hp_fsdp2_prepare_model
+
+        def _hp_clip_grad_norm(accelerator, parameters, max_norm, norm_type=2):
+            if getattr(accelerator, "is_fsdp2", False):
+                accelerator.unscale_gradients()
+                parameter_list = list(parameters)
+                parameter_ids = {id(param) for param in parameter_list}
+                for model in accelerator._models:  # pylint: disable=protected-access
+                    if not isinstance(model, HSDPModule):
+                        continue
+                    model_param_ids = {id(param) for param in model.parameters()}
+                    if parameter_ids and parameter_ids.issubset(model_param_ids):
+                        return hp_clip_grad_norm_(parameter_list, max_norm, norm_type=norm_type)
+            return self._orig_accelerator_clip_grad_norm(parameters, max_norm, norm_type=norm_type)
+
+        self.accelerator.clip_grad_norm_ = types.MethodType(_hp_clip_grad_norm, self.accelerator)
+        self._accelerator_patches_active = True
+
+    def _restore_accelerator_patches(self) -> None:
+        """Restore original Accelerate methods."""
+        if not self._accelerator_patches_active:
+            return
+
+        import accelerate.accelerator as acc_module  # pylint: disable=C0415
+
+        if self._orig_fsdp2_prepare_model is not None:
+            acc_module.fsdp2_prepare_model = self._orig_fsdp2_prepare_model
+        self.accelerator.clip_grad_norm_ = self._orig_accelerator_clip_grad_norm
+        self._accelerator_patches_active = False
+
+    def _wrap_model(self, model: nn.Module, training: bool = True, dataloader=None) -> nn.Module:
+        """Let Accelerate own FSDP2/HSDP wrapping so optimizer remapping stays correct."""
+        del dataloader
+        if isinstance(model, HSDPModule):
+            return model
+        if training and getattr(self.accelerator, "is_fsdp2", False):
+            return model
+        return super()._wrap_model(model, training=training)
+
+    def _get_train_sampler(self, train_dataset=None):
+        """Match the no-CP baseline sampler semantics before CP repeats whole logical batches."""
+        if train_dataset is None:
+            train_dataset = self.train_dataset
+        if getattr(self.finetuning_args, "disable_shuffling", False):
+            return torch.utils.data.SequentialSampler(train_dataset)
+        return super()._get_train_sampler(train_dataset)
+
+    def _build_cp_batch_sampler(self, dataset, shuffle: bool, batch_size: int, drop_last: bool):
+        """Repeat complete logical batches so CP groups consume the same baseline batch."""
+        sampler = self._get_train_sampler(dataset) if shuffle else torch.utils.data.SequentialSampler(dataset)
+        return _CPBatchRepeatedBatchSampler(
+            sampler,
+            batch_size=batch_size,
+            drop_last=drop_last,
+            repeat_factor=self._cp_size,
+            logical_group_size=max(1, get_platform().get_world_size() // self._cp_size),
+        )
+
+    def _get_cp_dataloader(self, dataset, batch_size: int, shuffle: bool):
+        """Create a train dataloader whose logical batches are shared within each CP group."""
+        if isinstance(dataset, torch.utils.data.IterableDataset):
+            raise NotImplementedError(
+                "HyperParallel CP training requires a map-style dataset because iterable datasets cannot "
+                "repeat logical batches across CP ranks."
+            )
+
+        try:
+            import datasets  # pylint: disable=C0415
+        except ImportError:  # pragma: no cover
+            datasets = None
+
+        if datasets is not None and isinstance(dataset, datasets.Dataset):
+            dataset = self._remove_unused_columns(dataset, description="Training")
+            data_collator = self.data_collator
+        else:
+            data_collator = self._get_collator_with_removed_columns(self.data_collator, description="Training")
+
+        batch_sampler = self._build_cp_batch_sampler(
+            dataset,
+            shuffle=shuffle,
+            batch_size=batch_size,
+            drop_last=self.args.dataloader_drop_last,
+        )
+        logical_batches = len(batch_sampler) // self._cp_size
+        dp_size = max(1, get_platform().get_world_size() // self._cp_size)
+        logical_length = (
+            logical_batches // dp_size if self.args.dataloader_drop_last else _ceil_div(logical_batches, dp_size)
+        )
+
+        dataloader_params = {
+            "batch_sampler": batch_sampler,
+            "collate_fn": data_collator,
+            "num_workers": self.args.dataloader_num_workers,
+            "pin_memory": self.args.dataloader_pin_memory,
+            "persistent_workers": self.args.dataloader_persistent_workers
+            if self.args.dataloader_num_workers > 0
+            else False,
+        }
+        if self.args.dataloader_num_workers > 0:
+            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor
+
+        from transformers.trainer import seed_worker  # pylint: disable=C0415
+
+        dataloader_params["worker_init_fn"] = partial(
+            seed_worker,
+            num_workers=self.args.dataloader_num_workers,
+            rank=self.args.process_index,
+        )
+
+        dataloader = self.accelerator.prepare(torch.utils.data.DataLoader(dataset, **dataloader_params))
+        return _CPDataLoaderLengthProxy(dataloader, logical_length)
+
+    def get_train_dataloader(self):
+        """Keep the no-CP logical batch stream, then repeat each whole batch across CP peers."""
+        if self.train_dataset is None:
+            raise ValueError("Trainer: training requires a train_dataset.")
+        if self._cp_size <= 1:
+            return super().get_train_dataloader()
+
+        shuffle = not getattr(self.finetuning_args, "disable_shuffling", False)
+        return self._get_cp_dataloader(
+            dataset=self.train_dataset,
+            batch_size=self._train_batch_size,
+            shuffle=shuffle,
+        )
+
+    def _move_model_to_device(self, model: nn.Module, device: Optional[torch.device] = None):
+        """Skip redundant device moves for HSDP-wrapped models."""
+        if isinstance(model, HSDPModule):
+            return model
+        if device is None:
+            return model
+        return model.to(device)
+
+    def train(self, *args, **kwargs):
+        """Activate HP patches during training and restore afterwards."""
+        self._activate_accelerator_patches()
+        try:
+            return super().train(*args, **kwargs)
+        finally:
+            self._restore_accelerator_patches()
+
+    def training_step(
+        self,
+        model: nn.Module,
+        inputs: dict[str, Any],
+        num_items_in_batch: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Standard training step with HSDP sync plus optional CP input sharding."""
+        model.train()
+        inputs = self._prepare_inputs(inputs)
+
+        if self._cp_size > 1:
+            inputs = shard_inputs_for_cp(inputs, self._cp_rank, self._cp_size)
+
+        sync_gradients = getattr(self.accelerator, "sync_gradients", True)
+        if isinstance(model, HSDPModule):
+            model.set_is_last_backward(sync_gradients)
+            model.set_requires_gradient_sync(sync_gradients)
+
+        compute_loss_context_manager = getattr(self, "compute_loss_context_manager", nullcontext)
+        with compute_loss_context_manager():
+            loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
+
+        if self.args.n_gpu > 1:
+            loss = loss.mean()
+
+        if not getattr(self, "model_accepts_loss_kwargs", False) and getattr(self, "compute_loss_func", None) is None:
+            loss = loss / self.args.gradient_accumulation_steps
+
+        self.accelerator.backward(loss)
+
+        if isinstance(model, HSDPModule) and sync_gradients:
+            hsdp_sync_stream()
+
+        return loss.detach()
+
+    def create_optimizer(self):
+        """Create optimizer and wrap step with SkipDTensorDispatch."""
+        optimizer = super().create_optimizer()
+        wrap_optimizer_with_skip_dtensor_dispatch(optimizer)
+        return optimizer
+
+    def _save_optimizer_and_scheduler(self, output_dir: str) -> None:
+        """Save model/optimizer shards per-rank and scheduler."""
+        save_hsdp_checkpoint(
+            model=self.model,
+            optimizer=self.optimizer,
+            lr_scheduler=self.lr_scheduler,
+            output_dir=output_dir,
+            should_save_scheduler=self.args.should_save and self.lr_scheduler is not None,
+        )
+
+    def _load_from_checkpoint(self, resume_from_checkpoint: str, model: Optional[nn.Module] = None) -> None:
+        """Load model from HSDP sharded checkpoint."""
+        target = model if model is not None else self.model
+        loaded = load_hsdp_model(target, resume_from_checkpoint)
+        if not loaded:
+            return super()._load_from_checkpoint(resume_from_checkpoint, model=model)
+        self._pending_hsdp_checkpoint = resume_from_checkpoint
+        return None
+
+    def _load_optimizer_and_scheduler(self, checkpoint: Optional[str] = None) -> None:
+        """Load optimizer/scheduler from per-rank checkpoint files."""
+        ckpt_dir = getattr(self, "_pending_hsdp_checkpoint", None) or checkpoint
+        if ckpt_dir is None:
+            return
+        load_hsdp_optimizer_and_scheduler(self.optimizer, self.lr_scheduler, ckpt_dir)
+
+    def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = False):
+        """Save model weights in HuggingFace-compatible format."""
+        save_dir = output_dir or self.args.output_dir
+        os.makedirs(save_dir, exist_ok=True)
+        export_to_hf_format(self.model, getattr(self, "processing_class", None), save_dir)

src/llamafactory/train/hyper_parallel/workflow.py

@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 from typing import TYPE_CHECKING, Optional
 
+from transformers import DataCollatorForLanguageModeling
+
 from ...data import SFTDataCollatorWith4DAttentionMask, get_dataset, get_template_and_fix_tokenizer
 from ...extras.constants import IGNORE_INDEX
 from ...extras.logging import get_logger
@@ -21,9 +24,9 @@ from ...extras.misc import calculate_tps
 from ...extras.packages import is_hyper_parallel_available, is_transformers_version_greater_than
 from ...extras.ploting import plot_loss
 from ...model import load_model, load_tokenizer
-from ..callbacks import SaveProcessorCallback
 from ..sft.metric import ComputeAccuracy, ComputeSimilarity, eval_logit_processor
-from ..trainer_utils import asft_loss_func, create_modelcard_and_push, create_ref_model, dft_loss_func, eaft_loss_func
+from ..trainer_utils import create_modelcard_and_push, create_ref_model
+from .trainer import HyperParallelTrainer
 
 
 if TYPE_CHECKING:
@@ -35,6 +38,94 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
+def _prepare_hp_args(finetuning_args: "FinetuningArguments", model_args: "ModelArguments"):
+    r"""Load HyperParallel arguments and apply LlamaFactory-side overrides.
+
+    When activation optimization is enabled, skip native gradient checkpointing
+    so HP can install its own via ``setup_activation_optimization``.
+    """
+    if not is_hyper_parallel_available():
+        raise ImportError("hyper_parallel is not installed. Please install it with `pip install hyper_parallel`.")
+
+    from hyper_parallel.integration.llamafactory import HyperParallelArguments  # pylint: disable=C0415
+
+    hp_args = HyperParallelArguments.from_finetuning_args(finetuning_args)
+
+    if getattr(hp_args, "cp_size", None) != finetuning_args.hyper_parallel_cp_size:
+        setattr(hp_args, "cp_size", finetuning_args.hyper_parallel_cp_size)
+
+    if hp_args.activation_mode != "none":
+        model_args.disable_gradient_checkpointing = True
+    return hp_args
+
+
+def run_pt(
+    model_args: "ModelArguments",
+    data_args: "DataArguments",
+    training_args: "Seq2SeqTrainingArguments",
+    finetuning_args: "FinetuningArguments",
+    callbacks: Optional[list["TrainerCallback"]] = None,
+):
+    hp_args = _prepare_hp_args(finetuning_args, model_args)
+
+    tokenizer_module = load_tokenizer(model_args)
+    tokenizer = tokenizer_module["tokenizer"]
+    template = get_template_and_fix_tokenizer(tokenizer, data_args)
+    dataset_module = get_dataset(template, model_args, data_args, training_args, stage="pt", **tokenizer_module)
+    model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train)
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+
+    trainer = HyperParallelTrainer(
+        hp_args=hp_args,
+        model=model,
+        args=training_args,
+        finetuning_args=finetuning_args,
+        data_collator=data_collator,
+        callbacks=callbacks,
+        **dataset_module,
+        **tokenizer_module,
+    )
+
+    if training_args.do_train:
+        train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+        trainer.save_model()
+        trainer.log_metrics("train", train_result.metrics)
+        trainer.save_metrics("train", train_result.metrics)
+        trainer.save_state()
+        if trainer.is_world_process_zero() and finetuning_args.plot_loss:
+            keys = ["loss"]
+            if isinstance(dataset_module.get("eval_dataset"), dict):
+                keys += [f"eval_{key}_loss" for key in dataset_module["eval_dataset"].keys()]
+            else:
+                keys += ["eval_loss"]
+
+            plot_loss(training_args.output_dir, keys=keys)
+
+    if training_args.do_eval:
+        metrics = trainer.evaluate(metric_key_prefix="eval")
+
+        if isinstance(dataset_module.get("eval_dataset"), dict):
+            for key in dataset_module["eval_dataset"].keys():
+                try:
+                    perplexity = math.exp(metrics[f"eval_{key}_loss"])
+                except OverflowError:
+                    perplexity = float("inf")
+
+                metrics[f"eval_{key}_perplexity"] = perplexity
+        else:
+            try:
+                perplexity = math.exp(metrics["eval_loss"])
+            except OverflowError:
+                perplexity = float("inf")
+
+            metrics["eval_perplexity"] = perplexity
+
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+    create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)
+
+
 def run_sft(
     model_args: "ModelArguments",
     data_args: "DataArguments",
@@ -43,13 +134,7 @@ def run_sft(
     generating_args: "GeneratingArguments",
     callbacks: Optional[list["TrainerCallback"]] = None,
 ):
-    if not is_hyper_parallel_available():
-        raise ImportError("hyper_parallel is not installed. Please install it with `pip install hyper_parallel`.")
-
-    from hyper_parallel.integration.llamafactory import (  # pylint: disable=C0415
-        HyperParallelArguments,
-        HyperParallelTrainer,
-    )
+    hp_args = _prepare_hp_args(finetuning_args, model_args)
 
     tokenizer_module = load_tokenizer(model_args)
     tokenizer = tokenizer_module["tokenizer"]
@@ -94,25 +179,6 @@ def run_sft(
         gen_kwargs["eos_token_id"] = [tokenizer.eos_token_id] + tokenizer.additional_special_tokens_ids
     gen_kwargs["pad_token_id"] = tokenizer.pad_token_id
 
-    hp_args = HyperParallelArguments.from_finetuning_args(finetuning_args)
-
-    callbacks = list(callbacks or [])
-    processor = tokenizer_module.get("processor")
-    if processor is not None:
-        callbacks.append(SaveProcessorCallback(processor))
-
-    compute_loss_func = None
-    if finetuning_args.use_dft_loss:
-        compute_loss_func = dft_loss_func
-    elif finetuning_args.use_eaft_loss:
-        compute_loss_func = lambda outputs, labels, num_items_in_batch=None: eaft_loss_func(  # noqa: E731
-            outputs, labels, num_items_in_batch, finetuning_args.eaft_alpha
-        )
-    elif finetuning_args.use_asft_loss:
-        from functools import partial
-
-        compute_loss_func = partial(asft_loss_func, asft_alpha=finetuning_args.asft_alpha)
-
     trainer = HyperParallelTrainer(
         hp_args=hp_args,
         model=model,
@@ -122,20 +188,11 @@ def run_sft(
         callbacks=callbacks,
         gen_kwargs=gen_kwargs,
         ref_model=ref_model,
-        compute_loss_func=compute_loss_func,
         **dataset_module,
         **tokenizer_module,
         **metric_module,
     )
 
-    if finetuning_args.use_badam:
-        from types import MethodType
-
-        from badam import BAdamCallback, clip_grad_norm_old_version  # type: ignore[import]
-
-        trainer.accelerator.clip_grad_norm_ = MethodType(clip_grad_norm_old_version, trainer.accelerator)
-        trainer.add_callback(BAdamCallback)
-
     # Training
     if training_args.do_train:
         train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)

src/llamafactory/train/tuner.py

@@ -88,12 +88,17 @@ def _training_function(config: dict[str, Any]) -> None:
 
     callbacks.append(ReporterCallback(model_args, data_args, finetuning_args, generating_args))  # add to last
 
-    if finetuning_args.stage == "sft" and finetuning_args.use_hyper_parallel:
+    if finetuning_args.stage in ["pt", "sft"] and finetuning_args.use_hyper_parallel:
         if not is_hyper_parallel_available():
             raise ImportError("hyper_parallel is not installed. Please install it with `pip install hyper_parallel`.")
-        from .hyper_parallel import run_sft as run_sft_hp
+        if finetuning_args.stage == "pt":
+            from .hyper_parallel import run_pt as run_pt_hp
 
-        run_sft_hp(model_args, data_args, training_args, finetuning_args, generating_args, callbacks)
+            run_pt_hp(model_args, data_args, training_args, finetuning_args, callbacks)
+        else:
+            from .hyper_parallel import run_sft as run_sft_hp
+
+            run_sft_hp(model_args, data_args, training_args, finetuning_args, generating_args, callbacks)
 
     elif finetuning_args.stage in ["pt", "sft", "dpo"] and finetuning_args.use_mca:
         if not is_mcore_adapter_available():
@@ -127,7 +132,7 @@ def _training_function(config: dict[str, Any]) -> None:
         raise ValueError(f"Unknown task: {finetuning_args.stage}.")
 
     if is_ray_available() and ray.is_initialized():
-        return  # if ray is intialized it will destroy the process group on return
+        return  # if ray is initialized it will destroy the process group on return
 
     try:
         if dist.is_initialized():

src/llamafactory/v1/accelerator/interface.py

@@ -29,16 +29,20 @@ And data parallelism types:
 from dataclasses import dataclass
 from datetime import timedelta
 from enum import StrEnum
-from typing import Any, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 from torch.distributed import barrier, destroy_process_group, init_process_group
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 
 from ..utils import logging
-from ..utils.types import DistributedConfig, ProcessGroup, TensorLike
+from ..utils.types import ProcessGroup, TensorLike
 from . import helper
 
 
+if TYPE_CHECKING:
+    from ..config.training_args import TrainingArguments
+
+
 logger = logging.get_logger(__name__)
 
 
@@ -68,6 +72,11 @@ class DistributedStrategy:
         if not helper.is_distributed():
             self.mp_shard_size = 1
         elif self.mp_shard_size is None:
+            if helper.get_world_size() % self.mp_replicate_size != 0:
+                raise ValueError(
+                    f"world_size ({helper.get_world_size()}) must be divisible by "
+                    f"mp_replicate_size ({self.mp_replicate_size})."
+                )
             self.mp_shard_size = helper.get_world_size() // self.mp_replicate_size
         elif self.mp_replicate_size * self.mp_shard_size != helper.get_world_size():
             raise ValueError(
@@ -78,6 +87,10 @@ class DistributedStrategy:
         if not helper.is_distributed():
             self.dp_size = 1
         elif self.dp_size is None:
+            if helper.get_world_size() % self.cp_size != 0:
+                raise ValueError(
+                    f"world_size ({helper.get_world_size()}) must be divisible by cp_size ({self.cp_size})."
+                )
             self.dp_size = helper.get_world_size() // self.cp_size
         elif self.dp_size * self.cp_size != helper.get_world_size():
             raise ValueError(
@@ -119,12 +132,13 @@ class DistributedInterface:
 
         return cls._instance
 
-    def __init__(self, config: DistributedConfig | None = None) -> None:
+    def __init__(
+        self,
+        training_args: "TrainingArguments | None" = None,
+    ) -> None:
         if self._initialized:
             return
 
-        self.dist_config = config
-
         helper.set_device_index()
         self._is_distributed = helper.is_distributed()
         self._rank = helper.get_rank()
@@ -134,17 +148,17 @@ class DistributedInterface:
         self.current_device = helper.get_current_device()
         self.device_count = helper.get_device_count()
 
-        if config is None:
+        if training_args is None:
             self.strategy = DistributedStrategy()
             timeout = 18000
         else:
             self.strategy = DistributedStrategy(
-                mp_replicate_size=config.get("mp_replicate_size", 1),
-                mp_shard_size=config.get("mp_shard_size", None),
-                dp_size=config.get("dp_size", None),
-                cp_size=config.get("cp_size", 1),
+                mp_replicate_size=training_args.mp_replicate_size,
+                mp_shard_size=training_args.mp_shard_size,
+                dp_size=training_args.dp_size,
+                cp_size=training_args.cp_size,
             )
-            timeout = config.get("timeout", 18000)
+            timeout = training_args.dist_timeout
 
         if self._is_distributed:
             init_process_group(timeout=timedelta(seconds=timeout), backend=helper.get_process_group_backend())

src/llamafactory/v1/config/model_args.py

@@ -25,9 +25,9 @@ class ModelArguments:
         default="Qwen/Qwen3-4B-Instruct-2507",
         metadata={"help": "Path to the model or model identifier from Hugging Face."},
     )
-    template: str = field(
-        default="qwen3_nothink",
-        metadata={"help": "Template for the model."},
+    custom_chat_template: str | None = field(
+        default=None,
+        metadata={"help": "Custom Jinja2 chat template string. Overrides the model's built-in template."},
     )
     trust_remote_code: bool = field(
         default=False,

src/llamafactory/v1/config/training_args.py

@@ -14,11 +14,16 @@
 
 import os
 from dataclasses import dataclass, field
+from typing import Literal
 from uuid import uuid4
 
+from ..utils.logging import get_logger
 from .arg_utils import BatchingStrategy, PluginConfig, get_plugin_config
 
 
+logger = get_logger(__name__)
+
+
 @dataclass
 class TrainingArguments:
     output_dir: str = field(
@@ -71,7 +76,31 @@ class TrainingArguments:
     )
     dist_config: PluginConfig | None = field(
         default=None,
-        metadata={"help": "Distribution configuration for training."},
+        metadata={"help": "Distributed backend plugin configuration."},
+    )
+    dp_size: int | None = field(
+        default=None,
+        metadata={"help": "Data parallel size, default to world_size // cp_size."},
+    )
+    cp_size: int = field(
+        default=1,
+        metadata={"help": "Context parallel size."},
+    )
+    cp_mode: str = field(
+        default="ulysses",
+        metadata={"help": "Context parallel implementation."},
+    )
+    mp_replicate_size: int = field(
+        default=1,
+        metadata={"help": "Model parallel replicate size."},
+    )
+    mp_shard_size: int | None = field(
+        default=None,
+        metadata={"help": "Model parallel shard size, default to world_size // mp_replicate_size."},
+    )
+    dist_timeout: int = field(
+        default=18000,
+        metadata={"help": "Distributed process group initialization timeout in seconds."},
     )
     optim_config: PluginConfig | None = field(
         default=None,
@@ -115,11 +144,52 @@ class TrainingArguments:
         default=1,
         metadata={"help": "Log metrics every N optimizer steps."},
     )
+    pref_loss: Literal["sigmoid", "orpo", "simpo"] = field(
+        default="sigmoid",
+        metadata={"help": "The type of DPO loss to use."},
+    )
+    pref_beta: float = field(
+        default=0.1,
+        metadata={"help": "The beta parameter in the preference loss."},
+    )
+    pref_ftx: float = field(
+        default=0.0,
+        metadata={"help": "The supervised fine-tuning loss coefficient in DPO training."},
+    )
+    simpo_gamma: float = field(
+        default=0.5,
+        metadata={"help": "The target reward margin term in SimPO loss."},
+    )
+    dpo_label_smoothing: float = field(
+        default=0.0,
+        metadata={"help": "The robust DPO label smoothing parameter in cDPO that should be between 0 and 0.5."},
+    )
+    ld_alpha: float | None = field(
+        default=None,
+        metadata={"help": "Alpha parameter from LD-DPO, controls weighting of verbose token log-probabilities."},
+    )
 
     def __post_init__(self) -> None:
         self.dist_config = get_plugin_config(self.dist_config)
         self.optim_config = get_plugin_config(self.optim_config)
         self.lr_scheduler_config = get_plugin_config(self.lr_scheduler_config)
+        try:
+            from ..plugins.model_plugins.deepspeed_utils import register_deepspeed_dist_config
+
+            register_deepspeed_dist_config(self.dist_config)
+        except ImportError:
+            pass
+
+        # The optimizer learning rate has a single source of truth: ``learning_rate``.
+        # Propagate it into ``optim_config["lr"]`` so optimizer plugins (e.g. Muon) pick it up
+        # via ``optim_config.get("lr")`` without each plugin needing a separate ``learning_rate`` arg.
+        if self.optim_config is not None:
+            if "lr" in self.optim_config:
+                logger.warning_rank0(
+                    "`optim_config.lr` is overridden by `learning_rate`; set the learning rate via "
+                    "`learning_rate` instead and remove `lr` from `optim_config`."
+                )
+            self.optim_config["lr"] = self.learning_rate
 
         if str(self.batching_strategy) == str(BatchingStrategy.DYNAMIC_BATCHING):
             if self.max_steps is None or self.max_steps <= 0:

src/llamafactory/v1/core/base_sampler.py

@@ -16,8 +16,8 @@ from collections.abc import AsyncGenerator
 
 from ..config import ModelArguments, SampleArguments, SampleBackend
 from ..utils.types import HFModel, Message, Sample, TorchDataset
+from .rendering import Renderer
 from .utils.inference_engine import HuggingFaceEngine
-from .utils.rendering import Renderer
 
 
 class BaseSampler:

src/llamafactory/v1/core/base_trainer.py

@@ -31,6 +31,7 @@ from abc import abstractmethod
 
 import torch
 import torch.nn.functional as F
+from torch.distributed.tensor import DTensor
 
 from ..accelerator.helper import ReduceOp
 from ..accelerator.interface import Dim, DistributedInterface
@@ -44,9 +45,9 @@ from ..utils.callbacks import (
 )
 from ..utils.helper import compute_valid_tokens
 from ..utils.types import BatchInput, HFModel, ModelOutput, Tensor, TorchDataset
+from .rendering import Renderer
 from .utils.batching import BatchGenerator
 from .utils.checkpoint import TrainingCheckpointCoordinator
-from .utils.rendering import Renderer
 
 
 logger = logging.get_logger(__name__)
@@ -93,9 +94,12 @@ class BaseTrainer:
         dist_name = self.args.dist_config.name if self.args.dist_config is not None else None
 
         if dist_name == "deepspeed":
-            from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin
+            if self.args.cp_size > 1:
+                raise ValueError("Context parallelism currently requires `dist_config.name: fsdp2`.")
 
-            self._deepspeed_engine = DistributedPlugin("deepspeed")(
+            from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
+
+            self._deepspeed_engine = DistributedPlugin("deepspeed").shard_model(
                 self.model,
                 self.args.dist_config,
                 num_micro_batch=self.train_batch_generator.num_micro_batch,
@@ -138,7 +142,7 @@ class BaseTrainer:
         self.state.global_step = self.global_step
         self.state.epoch = self._resume_epoch
 
-        if self.args.dist_config is not None and self.args.dist_config.get("cp_size", 1) > 1:
+        if self.args.cp_size > 1:
             # qwen3.5 is not supported because of the different attention implementation, which will be supported in the future.
             if model.config.model_type == "qwen3_5":
                 raise RuntimeError(
@@ -151,7 +155,7 @@ class BaseTrainer:
                     "Sequence parallelism requires flash attention. Please set `flash_attn: flash_attention_2`."
                 )
 
-            SequenceParallelModelPlugin(self.args.dist_config.get("cp_mode", "ulysses"))(model, self.args.dist_config)
+            SequenceParallelModelPlugin(self.args.cp_mode)(model, self.args.cp_size)
 
     def _create_batch_generator(self) -> None:
         if (
@@ -182,9 +186,9 @@ class BaseTrainer:
                 device_ids = None if self.device.type == "cpu" else [self.device.index]
                 self.model = DDP(self.model, device_ids=device_ids)
         else:
-            from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin
+            from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
 
-            self.model = DistributedPlugin(self.args.dist_config.name)(
+            self.model = DistributedPlugin(self.args.dist_config.name).shard_model(
                 self.model,
                 self.args.dist_config,
                 bf16=self.args.bf16,
@@ -196,7 +200,7 @@ class BaseTrainer:
             _trainable_params = [p for p in self.model.parameters() if p.requires_grad]
             self.optimizer = torch.optim.AdamW(_trainable_params, lr=self.args.learning_rate)
         else:
-            from ..plugins.trainer_plugins.optimizer import OptimizerPlugin
+            from ..plugins.trainer_plugins.optimizers.optimizer import OptimizerPlugin
 
             self.optimizer = OptimizerPlugin(self.args.optim_config.name)(self.model, self.args.optim_config)
 
@@ -255,7 +259,7 @@ class BaseTrainer:
                 step_valid_tokens = DistributedInterface().all_reduce(step_valid_tokens, op=ReduceOp.SUM)
                 num_micro = len(micro_batches)
                 for i, micro_batch in enumerate(micro_batches):
-                    if self.args.dist_config and self.args.dist_config.get("cp_size", 1) > 1:
+                    if self.args.cp_size > 1:
                         from ..plugins.model_plugins.parallelization.sequence_parallel import (
                             SequenceParallelLossPlugin,
                         )
@@ -279,12 +283,19 @@ class BaseTrainer:
                     # deepspeed: engine.step() already ran inside backward at the sync boundary
                     grad_norm = self._deepspeed_engine.get_grad_norm()
                 else:
-                    grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.max_grad_norm).item()
-
-                    if self.args.dist_config and self.args.dist_config.get("cp_size", 1) > 1:
-                        grad_norm = grad_norm**2
-                        grad_norm = DistributedInterface().all_reduce(grad_norm, op=ReduceOp.SUM, dim=Dim.CP)
-                        grad_norm = grad_norm**0.5
+                    # 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()
 
                     if not torch.isfinite(torch.tensor(grad_norm)):  # type: ignore # pyright: ignore [reportUnknownReturnType]
                         logger.warning_rank0(f"Gradient norm is not finite: {grad_norm}")
@@ -318,6 +329,10 @@ class BaseTrainer:
                         "grad_norm": grad_norm,
                         "learning_rate": current_lr,
                     }
+                    # Merge per-step trainer metrics (e.g. DPO rewards/logps/logits)
+                    step_metrics = getattr(self, "_step_metrics", None)
+                    if step_metrics:
+                        logs.update(step_metrics)
                     self.callback_handler.on_log(self.args, self.state, logs)
 
                 if self.args.save_steps and self.global_step % self.args.save_steps == 0:
@@ -338,7 +353,7 @@ class BaseTrainer:
     def save_model(self) -> None:
         """Save the model."""
         if self.args.dist_config is not None and self.args.dist_config.name in ("deepspeed", "fsdp2"):
-            from ..plugins.trainer_plugins.distributed.hub import DistributedPlugin
+            from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
 
             DistributedPlugin(self.args.dist_config.name).save_model(
                 self.model, self.args.output_dir, self.renderer.processor

src/llamafactory/v1/core/data_engine.py

@@ -58,8 +58,9 @@ class DataEngine(Dataset):
         """Dict of (dataset_name, dataset)"""
         self.dataset_infos: dict[str, DatasetInfo] = {}
         """Dict of (dataset_name, dataset_info)"""
-        self.data_index: list[tuple[str, int]] = []
-        """List of (dataset_name, sample_index)"""
+        self.data_index: list[tuple[str, int, int | None]] = []
+        """List of (dataset_name, sample_index, cut). ``cut`` is the prefix length for multi-turn
+        split (messages[:cut], ending at one supervised assistant turn), or None for a whole sample."""
         self.streaming: bool = False
         """Whether dataset is streaming."""
         self._get_dataset_info()
@@ -98,12 +99,23 @@ class DataEngine(Dataset):
                 self.datasets[dataset_name] = DataLoaderPlugin(dataset_info["source"]).load(dataset_info)
 
     def _build_data_index(self) -> None:
-        """Build dataset index."""
+        """Build dataset index.
+
+        Multi-turn SFT conversations are prefix-expanded: one index entry per supervised assistant
+        turn, so ``len()`` reflects the true number of training samples (each trained on its last
+        turn). Entries are ``(dataset_name, sample_index, cut)``; ``cut`` is the prefix length
+        ``messages[:cut]`` ending at one supervised turn, or ``None`` for a whole sample (DPO,
+        streaming, or no supervised turn).
+        """
         for dataset_name, dataset in self.datasets.items():
-            if self.streaming:
-                data_index = [(dataset_name, -1) for _ in range(1000)]
+            if self.streaming:  # cannot pre-count turns -> keep whole, unsplit
+                data_index = [(dataset_name, -1, None) for _ in range(1000)]
             else:
-                data_index = [(dataset_name, sample_index) for sample_index in range(len(dataset))]
+                data_index = []
+                for sample_index in range(len(dataset)):
+                    sample = self._convert_data_sample(dataset[sample_index], dataset_name)
+                    for cut in self._prefix_cuts(sample):
+                        data_index.append((dataset_name, sample_index, cut))
 
             size = self.dataset_infos[dataset_name].get("size")
             weight = self.dataset_infos[dataset_name].get("weight")
@@ -114,6 +126,20 @@ class DataEngine(Dataset):
 
             self.data_index.extend(data_index)
 
+    @staticmethod
+    def _prefix_cuts(sample: Sample) -> list[int | None]:
+        """Prefix lengths to split a multi-turn conversation on: one per supervised assistant turn.
+
+        ``u1 a1 u2 a2`` -> ``[2, 4]`` (samples ``messages[:2]`` and ``messages[:4]``, each trained on
+        its last assistant turn). Non-SFT samples (no ``messages``) or those with no supervised turn
+        are kept whole (``[None]``).
+        """
+        messages = sample.get("messages")
+        if not messages:
+            return [None]
+        cuts = [i + 1 for i, m in enumerate(messages) if m["role"] == "assistant" and m.get("loss_weight", 1.0) > 1e-6]
+        return cuts or [None]
+
     def _convert_data_sample(self, raw_sample: dict[str, Any], dataset_name: str) -> Sample:
         """Convert dataset sample.
 
@@ -156,20 +182,22 @@ class DataEngine(Dataset):
             raise ValueError("Streaming dataset does not support index access.")
 
         if isinstance(index, int):
-            dataset_name, sample_index = self.data_index[index]
-            return self._convert_data_sample(self.datasets[dataset_name][sample_index], dataset_name)
+            return self._get(*self.data_index[index])
         else:  # data selector plugin
             from ..plugins.data_plugins.loader import select_data_sample
 
             selected_index = select_data_sample(self.data_index, index)
             if isinstance(selected_index, list):
-                return [
-                    self._convert_data_sample(self.datasets[dataset_name][sample_index], dataset_name)
-                    for dataset_name, sample_index in selected_index
-                ]
+                return [self._get(*entry) for entry in selected_index]
             else:
-                dataset_name, sample_index = selected_index
-                return self._convert_data_sample(self.datasets[dataset_name][sample_index], dataset_name)
+                return self._get(*selected_index)
+
+    def _get(self, dataset_name: str, sample_index: int, cut: int | None = None) -> Sample:
+        """Convert one raw row, truncating to a multi-turn prefix when ``cut`` is set."""
+        sample = self._convert_data_sample(self.datasets[dataset_name][sample_index], dataset_name)
+        if cut is not None and "messages" in sample:
+            sample = {**sample, "messages": sample["messages"][:cut]}
+        return sample
 
     def __iter__(self) -> Iterable[Sample]:
         """Get dataset iterator.

src/llamafactory/v1/core/model_engine.py

@@ -37,8 +37,9 @@ from ..accelerator.helper import DeviceType
 from ..accelerator.interface import DistributedInterface
 from ..config.model_args import ModelArguments, ModelClass
 from ..utils import logging
+from ..utils.helper import get_tokenizer, is_tokenizer
 from ..utils.types import HFConfig, HFModel, Processor
-from .utils.rendering import Renderer
+from .rendering import Renderer
 
 
 logger = logging.get_logger(__name__)
@@ -63,28 +64,30 @@ class ModelEngine:
         """Whether to train the model."""
         self.processor = self._init_processor()
         """Tokenizer or multi-modal processor."""
-        self.renderer = Renderer(self.args.template, self.processor)
-        """Renderer."""
+        self._sync_chat_template()
         self.model_config = self._init_model_config()
         """Model configuration."""
-        self._dist_config = DistributedInterface().dist_config
-        self._deepspeed_zero3_plugin = None
+        self.renderer = Renderer(self.processor)
+        """Renderer."""
         self._deepspeed_zero3_enabled = False
 
-        if self.is_train and self._dist_config is not None and self._dist_config.get("name") == "deepspeed":
+        try:
             from ..plugins.model_plugins.deepspeed_utils import (
+                is_deepspeed_zero3_enabled,
                 setup_deepspeed_zero3_model_loading,
                 teardown_deepspeed_zero3_model_loading,
             )
 
+            self._deepspeed_zero3_enabled = self.is_train and is_deepspeed_zero3_enabled()
+        except ImportError:
+            pass
+
+        if self._deepspeed_zero3_enabled:
+            plugin = setup_deepspeed_zero3_model_loading()
             try:
-                self._deepspeed_zero3_plugin = setup_deepspeed_zero3_model_loading(self.is_train, self._dist_config)
-                self._deepspeed_zero3_enabled = self._deepspeed_zero3_plugin is not None
                 self.model = self._init_model()
             finally:
-                teardown_deepspeed_zero3_model_loading(self._deepspeed_zero3_plugin)
-                self._deepspeed_zero3_plugin = None
-                self._deepspeed_zero3_enabled = False
+                teardown_deepspeed_zero3_model_loading(plugin)
         else:
             self.model = self._init_model()
 
@@ -99,6 +102,19 @@ class ModelEngine:
             trust_remote_code=self.args.trust_remote_code,
         )
 
+    def _sync_chat_template(self) -> None:
+        """Sync chat_template and inject custom_chat_template."""
+        tokenizer = get_tokenizer(self.processor)
+        if not is_tokenizer(self.processor) and not getattr(self.processor, "chat_template", None):
+            if getattr(tokenizer, "chat_template", None):
+                self.processor.chat_template = tokenizer.chat_template
+
+        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
+
     def _init_model_config(self) -> HFConfig:
         """Init model config."""
         return AutoConfig.from_pretrained(
@@ -127,9 +143,7 @@ class ModelEngine:
 
             init_kwargs = QuantizationPlugin(self.args.quant_config.name)(
                 init_kwargs=init_kwargs,
-                config=self.model_config,
-                tokenizer=self.processor,
-                model_args=self.args,
+                quant_config=self.args.quant_config,
                 is_trainable=self.is_train,
             )
 
@@ -185,17 +199,16 @@ class ModelEngine:
 
             from ..plugins.model_plugins.peft import PeftPlugin
 
-            model = PeftPlugin(self.args.peft_config.name)(model, self.args.peft_config, self.is_train)
+            model = PeftPlugin(self.args.peft_config.name)(
+                model,
+                peft_config=self.args.peft_config,
+                is_train=self.is_train,
+            )
 
         if self.args.kernel_config is not None:
-            from ..plugins.model_plugins.kernels.interface import KernelPlugin
+            from ..plugins.model_plugins.kernels.interface import apply_kernels
 
-            kernel_config = self.args.kernel_config
-            kernel_kwargs: dict = {"model": model, "include_kernels": kernel_config.get("include_kernels")}
-            if kernel_config.name == "liger_kernel":
-                # Fused linear CE omits logits; SFT stage needs logits for loss_weights.
-                kernel_kwargs["require_logits"] = self.is_train
-            model = KernelPlugin(kernel_config.name)(**kernel_kwargs)
+            model = apply_kernels(model, self.args.kernel_config, require_logits=self.is_train)
 
         return model
 

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

@@ -12,8 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from ...utils.plugin import BasePlugin
+"""Rendering: turn a v1 ``Sample`` into a tokenized ``ModelInput``.
+
+Public entry point is :class:`Renderer`. Internals are split by concern:
+``format`` (message<->HF conversion) and ``escape`` (special-token escaping). Assistant supervision
+is located by a prompt/full token diff rather than a per-model marker table.
+"""
+
+from .rendering import Renderer
 
 
-class OptimizerPlugin(BasePlugin):
-    pass
+__all__ = ["Renderer"]

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

@@ -0,0 +1,100 @@
+# 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.
+
+"""Special-token escaping (prompt-injection hardening).
+
+Neutralizes control-token strings (``<|im_start|>``, ``<|image_pad|>`` ...) that appear literally
+in user-controlled text, so a crafted dataset cannot inject role markers or media placeholders
+into the rendered stream. A no-op for normal data.
+"""
+
+import json
+
+from ...utils.types import Message
+
+
+def _special_token_strings(tokenizer) -> list[str]:
+    """Strings the tokenizer encodes to a reserved/special id.
+
+    Such strings must be neutralized if they appear literally in user text. Derived from
+    ``added_tokens_decoder`` so it covers every control token of the model (``<|im_start|>``,
+    ``<|image_pad|>``, ``<tts_pad>`` ...), not only ``<|...|>``-shaped ones. Sorted longest-first
+    so nested matches escape correctly.
+    """
+    specials = [str(t) for t in tokenizer.added_tokens_decoder.values() if getattr(t, "special", False)]
+    return sorted((s for s in specials if len(s) >= 2), key=len, reverse=True)
+
+
+def _escape_special(text: str, specials: list[str], special_ids: set[int], tokenizer) -> str:
+    """Break any special-token string in user text by inserting U+200B after its first char.
+
+    No-op (no tokenization cost) when the text contains no special-token string. When it does,
+    self-validate that the result no longer encodes to a special id -- some normalizers strip
+    zero-width chars and would resurrect the collision -- and raise if it does.
+    """
+    if not any(sp in text for sp in specials):
+        return text
+    out = text
+    for sp in specials:
+        if sp in out:
+            # Insert a zero-width space (U+200B) after the first char to break the exact
+            # special-token string match while keeping the text visually/semantically intact.
+            out = out.replace(sp, sp[0] + "\u200b" + sp[1:])
+    if special_ids.intersection(tokenizer(out, add_special_tokens=False)["input_ids"]):
+        raise ValueError(
+            "special-token escape failed: the tokenizer normalized away the break char; "
+            "user text contains a literal control token that cannot be safely neutralized."
+        )
+    return out
+
+
+def _escape_special_in_messages(
+    messages: list[Message], specials: list[str], special_ids: set[int], tokenizer
+) -> list[Message]:
+    """Return messages with special-token strings neutralized in user-controlled literal text.
+
+    Covers ``text``/``reasoning`` block values and string values inside ``tool_call`` arguments.
+    """
+    if not specials:
+        return messages
+    escaped: list[Message] = []
+    for message in messages:
+        new_content = []
+        for content in message["content"]:
+            if content["type"] in ("text", "reasoning"):
+                new_content.append(
+                    {**content, "value": _escape_special(content["value"], specials, special_ids, tokenizer)}
+                )
+            elif content["type"] == "tool_call":
+                try:
+                    tc = json.loads(content["value"])
+                except (json.JSONDecodeError, TypeError):
+                    new_content.append(content)
+                    continue
+                # A tool_call value that is valid JSON but not an object (list/str/int) carries no
+                # escapable argument strings -- pass it through untouched rather than crash on .get().
+                if isinstance(tc, dict):
+                    args = tc.get("arguments")
+                    if isinstance(args, dict):
+                        tc["arguments"] = {
+                            k: (_escape_special(v, specials, special_ids, tokenizer) if isinstance(v, str) else v)
+                            for k, v in args.items()
+                        }
+                    new_content.append({**content, "value": json.dumps(tc)})
+                else:
+                    new_content.append(content)
+            else:
+                new_content.append(content)
+        escaped.append({**message, "content": new_content})
+    return escaped

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

@@ -0,0 +1,65 @@
+# 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.
+
+"""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``.
+"""
+
+import json
+
+from ...utils.types import Message
+
+
+_FALLBACK_CHATML_JINJA = (
+    "{% for message in messages %}"
+    "{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}"
+    "{% endfor %}"
+    "{% if add_generation_prompt %}"
+    "{{'<|im_start|>assistant\n'}}"
+    "{% endif %}"
+)
+
+
+def _to_hf_messages(messages: list[Message]) -> 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 tool_calls:
+            hf_msg["tool_calls"] = tool_calls
+        if reasoning_content:
+            hf_msg["reasoning_content"] = reasoning_content
+
+        hf_messages.append(hf_msg)
+    return hf_messages

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

@@ -0,0 +1,203 @@
+# 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.
+
+"""Rendering: turn a v1 ``Sample`` into a tokenized ``ModelInput``.
+
+This module is the orchestration + public API (``Renderer``). The mechanical pieces live in
+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.
+"""
+
+import json
+
+from ...utils.constants import IGNORE_INDEX
+from ...utils.helper import get_tokenizer
+from ...utils.types import Message, ModelInput, Processor, Sample
+from .escape import _escape_special, _escape_special_in_messages, _special_token_strings
+from .format import _FALLBACK_CHATML_JINJA, _to_hf_messages
+
+
+def _render_messages(
+    processor: Processor,
+    messages: list[Message],
+    tools: str | None = None,
+    is_generate: bool = False,
+    **kwargs,
+) -> ModelInput:
+    r"""Render messages using the model's own chat template.
+
+    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
+
+    # 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)
+
+    tools_parsed = None
+    if tools:
+        tools = _escape_special(tools, specials, special_ids, tokenizer)  # E3: tools text is user-controlled
+        try:
+            tools_parsed = json.loads(tools)
+        except json.JSONDecodeError as e:
+            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
+        )
+        return tokenizer(text, add_special_tokens=False)["input_ids"]
+
+    # 1. Full sequence, used verbatim.
+    input_ids = _encode(hf_messages, add_generation_prompt=is_generate)
+    n = len(input_ids)
+
+    if is_generate:
+        # Generation prompt only -- nothing is supervised.
+        return ModelInput(
+            input_ids=input_ids,
+            attention_mask=[1] * n,
+            labels=[IGNORE_INDEX] * n,
+            loss_weights=[0.0] * n,
+        )
+
+    # 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)
+    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
+        # mislabel.
+        raise ValueError(
+            "prompt is not a token-prefix of the full sequence; the chat template is not "
+            "prefix-stable for this turn, so diff-based labeling is unsafe."
+        )
+
+    weight = messages[-1].get("loss_weight", 1.0)
+    supervised = weight > 1e-6
+    labels = [IGNORE_INDEX] * len(prompt_ids)
+    loss_weights = [0.0] * len(prompt_ids)
+    for tid in input_ids[len(prompt_ids) :]:
+        labels.append(tid if supervised else IGNORE_INDEX)
+        loss_weights.append(weight)
+
+    return ModelInput(
+        input_ids=input_ids,
+        attention_mask=[1] * n,
+        labels=labels,
+        loss_weights=loss_weights,
+    )
+
+
+class Renderer:
+    def __init__(self, processor: Processor) -> None:
+        self.processor = processor
+
+    def render_messages(
+        self,
+        messages: list[Message],
+        tools: str | None = None,
+        is_generate: bool = False,
+        **kwargs,
+    ) -> ModelInput:
+        """Render messages to model input using apply_chat_template.
+
+        Args:
+            messages: The messages to render. For training the last message must be the supervised
+                assistant turn (use ``process_samples`` to split multi-turn conversations).
+            tools: JSON string of tool definitions.
+            is_generate: Whether to render for generation (adds generation prompt, no supervision).
+            **kwargs: Extra chat-template kwargs (e.g. ``enable_thinking``) forwarded verbatim to
+                ``apply_chat_template``; unset ones fall back to the template's own defaults. A
+                supervised assistant turn carrying reasoning forces ``enable_thinking=True``.
+
+        Returns:
+            ModelInput with input_ids, attention_mask, labels, and loss_weights.
+        """
+        return _render_messages(self.processor, messages, tools, is_generate, **kwargs)
+
+    def process_samples(self, samples: list[Sample]) -> list[ModelInput]:
+        """Process samples to model input.
+
+        Multi-turn SFT conversations are already prefix-split in the data layer (DataEngine), so each
+        ``messages`` sample is rendered once -- the diff-based renderer supervises only its last
+        assistant turn.
+
+        Args:
+            samples: The samples to process.
+
+        Returns:
+            List of processed model inputs.
+        """
+        model_inputs = []
+        for sample in samples:
+            rendered: list[ModelInput] = []
+            if "messages" in sample:
+                model_input = self.render_messages(sample["messages"], sample.get("tools"))
+                model_input["position_ids"] = list(range(1, len(model_input["input_ids"]) + 1))
+                rendered.append(model_input)
+            elif "chosen_messages" in sample and "rejected_messages" in sample:
+                chosen_input = self.render_messages(sample["chosen_messages"], sample.get("tools"))
+                rejected_input = self.render_messages(sample["rejected_messages"], sample.get("tools"))
+                chosen_input["token_type_ids"] = [1] * len(chosen_input["input_ids"])
+                rejected_input["token_type_ids"] = [2] * len(rejected_input["input_ids"])
+                model_input = ModelInput(
+                    input_ids=chosen_input["input_ids"] + rejected_input["input_ids"],
+                    attention_mask=chosen_input["attention_mask"] + rejected_input["attention_mask"],
+                    labels=chosen_input["labels"] + rejected_input["labels"],
+                    loss_weights=chosen_input["loss_weights"] + rejected_input["loss_weights"],
+                    token_type_ids=chosen_input["token_type_ids"] + rejected_input["token_type_ids"],
+                )
+                # chosen and rejected are independent sequences; position ids must restart at 1 for
+                # each (a single continuous range would offset rejected's positional embeddings).
+                model_input["position_ids"] = list(range(1, len(chosen_input["input_ids"]) + 1)) + list(
+                    range(1, len(rejected_input["input_ids"]) + 1)
+                )
+                rendered.append(model_input)
+            else:
+                raise ValueError("No valid messages or chosen_messages/rejected_messages found in sample.")
+
+            for model_input in rendered:
+                if "extra_info" in sample:
+                    model_input["extra_info"] = sample["extra_info"]
+                if "_dataset_name" in sample:
+                    model_input["_dataset_name"] = sample["_dataset_name"]
+                model_inputs.append(model_input)
+
+        return model_inputs

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

@@ -37,7 +37,7 @@ from ...utils import logging
 from ...utils.helper import pad_and_truncate
 from ...utils.objects import StatefulBuffer
 from ...utils.types import BatchInfo, BatchInput, ModelInput, TorchDataset
-from .rendering import Renderer
+from ..rendering import Renderer
 
 
 logger = logging.get_logger(__name__)
@@ -87,6 +87,7 @@ class BatchGenerator(Iterator):
         self.pin_memory = pin_memory
         self.drop_last = drop_last
         self.seed = seed
+        self._warned_truncation = False  # warn once when dropping fully-truncated (zero-loss) samples
         # TODO: support length and infinity
         dp_size = DistributedInterface().get_world_size(Dim.DP)
 
@@ -185,6 +186,31 @@ class BatchGenerator(Iterator):
 
         return batch
 
+    def _drop_unsupervised(self, samples: list[ModelInput]) -> list[ModelInput]:
+        """Drop samples whose supervised span is entirely beyond ``cutoff_len``.
+
+        Prefix-split puts the supervised tokens at the tail, and truncation keeps ``[:cutoff_len]``,
+        so a sample longer than ``cutoff_len`` loses all supervision and would contribute a zero-loss
+        (wasted) step. Only such over-length samples are at risk -- samples that fit within
+        ``cutoff_len`` are never truncated and always keep supervision -- so they pass an O(1) length
+        test before any ``loss_weights`` scan. Drop the at-risk, fully-masked ones and warn once.
+        """
+        kept = []
+        for sample in samples:
+            if len(sample["input_ids"]) > self.cutoff_len and not any(
+                w > 1e-6 for w in sample["loss_weights"][: self.cutoff_len]
+            ):
+                if not self._warned_truncation:
+                    self._warned_truncation = True
+                    logger.warning_rank0(
+                        f"Dropping training sample(s) whose supervised tokens fall entirely beyond "
+                        f"cutoff_len={self.cutoff_len} (all loss masked after truncation). "
+                        "Increase cutoff_len to keep them."
+                    )
+                continue
+            kept.append(sample)
+        return kept
+
     def _fill_buffer(self) -> None:
         if self.batching_strategy == BatchingStrategy.NORMAL:
             while len(self._buffer) < self.micro_batch_size * self.num_micro_batch:
@@ -193,7 +219,7 @@ class BatchGenerator(Iterator):
                 except StopIteration:
                     break
 
-                self._buffer.put(samples)
+                self._buffer.put(self._drop_unsupervised(samples))
         else:
             from ...plugins.trainer_plugins.batching import BatchingPlugin
 

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

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

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

@@ -25,7 +25,7 @@ from ...accelerator.interface import DistributedInterface
 from ...config import ModelArguments, SampleArguments
 from ...utils.helper import get_tokenizer
 from ...utils.types import HFModel, Message, Sample, TorchDataset
-from .rendering import Renderer
+from ..rendering import Renderer
 
 
 class BaseEngine(ABC):

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

@@ -1,178 +0,0 @@
-# Copyright 2025 the LlamaFactory team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Rendering utils.
-
-How to use:
-renderer = Renderer(template, processor)
-renderer.render_messages(messages: list[Message], tools: str | None) -> ModelInputs
-renderer.parse_message(text: str) -> Message
-renderer.process_samples(samples: list[Sample]) -> list[ModelInput]
-"""
-
-import numpy as np
-
-from ...utils.constants import IGNORE_INDEX
-from ...utils.helper import get_tokenizer
-from ...utils.types import Message, ModelInput, Processor, Sample
-
-
-def render_chatml_messages(
-    processor: Processor,
-    messages: list[Message],
-    tools: str | None = None,
-    is_generate: bool = False,
-) -> ModelInput:
-    """Apply chatml template to messages and convert them to model input.
-
-    See https://huggingface.co/spaces/huggingfacejs/chat-template-playground?modelId=Qwen/Qwen2-7B-Instruct
-    """
-    tokenizer = get_tokenizer(processor)
-    input_ids, labels, loss_weights = [], [], []
-
-    for message in messages:
-        temp_str = "<|im_start|>" + message["role"] + "\n"
-        for content in message["content"]:
-            if content["type"] == "text":
-                temp_str += content["value"]
-            else:
-                raise ValueError(f"Unsupported content type: {content['type']}")
-
-        temp_str += "<|im_end|>\n"
-        temp_weight = message.get("loss_weight", 1.0 if message["role"] == "assistant" else 0.0)
-        temp_ids = tokenizer.encode(temp_str, add_special_tokens=False)
-        input_ids.extend(temp_ids)
-        loss_weights.extend([temp_weight] * len(temp_ids))
-        if temp_weight > 1e-6:
-            labels.extend(temp_ids)
-        else:
-            labels.extend([IGNORE_INDEX] * len(temp_ids))
-
-    if is_generate:
-        temp_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
-        input_ids.extend(temp_ids)
-        loss_weights.extend([0.0] * len(temp_ids))
-        labels.extend([IGNORE_INDEX] * len(temp_ids))
-
-    return ModelInput(
-        input_ids=input_ids,
-        attention_mask=[1] * len(input_ids),
-        labels=labels,
-        loss_weights=loss_weights,
-    )
-
-
-def parse_chatml_message(generated_text: str) -> Message:
-    """Parse a message in ChatML format.
-
-    Args:
-        generated_text (str): The generated text in ChatML format.
-
-    Returns:
-        Message: The parsed message.
-    """
-    return Message(role="assistant", content=[{"type": "text", "value": generated_text}])
-
-
-class Renderer:
-    def __init__(self, template: str, processor: Processor):
-        self.template = template
-        self.processor = processor
-
-    def render_messages(
-        self,
-        messages: list[Message],
-        tools: str | None = None,
-        is_generate: bool = False,
-        enable_thinking: bool = False,
-    ) -> ModelInput:
-        """Apply template to messages and convert them to model input.
-
-        Args:
-            messages (list[Message]): The messages to render.
-            tools (str | None, optional): The tools to use. Defaults to None.
-            is_generate (bool, optional): Whether to render for generation. Defaults to False.
-            enable_thinking (bool, optional): Whether to enable thinking mode for generation. Defaults to False.
-
-        Returns:
-            ModelInput: The rendered model input.
-        """
-        if self.template == "chatml":
-            return render_chatml_messages(self.processor, messages, tools, is_generate)
-        else:
-            from ...plugins.model_plugins.rendering import RenderingPlugin
-
-            return RenderingPlugin(self.template).render_messages(
-                self.processor, messages, tools, is_generate, enable_thinking
-            )
-
-    def parse_message(self, generated_text: str) -> Message:
-        """Parse a message in the template format.
-
-        Args:
-            generated_text (str): The generated text in the template format.
-
-        Returns:
-            Message: The parsed message.
-        """
-        if self.template == "chatml":
-            return parse_chatml_message(generated_text)
-        else:
-            from ...plugins.model_plugins.rendering import RenderingPlugin
-
-            return RenderingPlugin(self.template).parse_message(generated_text)
-
-    def process_samples(self, samples: list[Sample]) -> list[ModelInput]:
-        """Process samples to model input.
-
-        Args:
-            samples (list[Sample]): The samples to process.
-
-        Returns:
-            list[ModelInput]: The processed model inputs.
-        """
-        model_inputs = []
-        for sample in samples:
-            if "messages" in sample:
-                model_input = self.render_messages(sample["messages"], sample.get("tools"))
-                if "position_ids" not in model_input:
-                    model_input["position_ids"] = list(range(1, len(model_input["input_ids"]) + 1))
-            elif "chosen_messages" in sample and "rejected_messages" in sample:
-                chosen_input = self.render_messages(sample["chosen_messages"], sample.get("tools"))
-                rejected_input = self.render_messages(sample["rejected_messages"], sample.get("tools"))
-                chosen_input["token_type_ids"] = [1] * len(chosen_input["input_ids"])
-                rejected_input["token_type_ids"] = [2] * len(rejected_input["input_ids"])
-                model_input = ModelInput(
-                    input_ids=chosen_input["input_ids"] + rejected_input["input_ids"],
-                    attention_mask=chosen_input["attention_mask"] + rejected_input["attention_mask"],
-                    labels=chosen_input["labels"] + rejected_input["labels"],
-                    loss_weights=chosen_input["loss_weights"] + rejected_input["loss_weights"],
-                    token_type_ids=chosen_input["token_type_ids"] + rejected_input["token_type_ids"],
-                )
-                if "position_ids" in chosen_input:
-                    model_input["position_ids"] = np.concatenate(
-                        [chosen_input["position_ids"], rejected_input["position_ids"]], axis=-1
-                    )
-            else:
-                raise ValueError("No valid messages or chosen_messages/rejected_messages found in sample.")
-
-            if "extra_info" in sample:
-                model_input["extra_info"] = sample["extra_info"]
-
-            if "_dataset_name" in sample:
-                model_input["_dataset_name"] = sample["_dataset_name"]
-
-            model_inputs.append(model_input)
-
-        return model_inputs

src/llamafactory/v1/launcher.py

@@ -146,7 +146,9 @@ def launch():
 
             run_sft()
         elif command == "dpo":
-            raise NotImplementedError("DPO trainer is not implemented yet.")
+            from llamafactory.v1.trainers.dpo_trainer import run_dpo
+
+            run_dpo()
         elif command == "rm":
             from llamafactory.v1.trainers.rm_trainer import run_rm
 
@@ -173,9 +175,9 @@ def main():
 
         run_sft()
     elif command == "dpo":
-        # from llamafactory.v1.trainers.dpo_trainer import run_dpo
-        # run_dpo()
-        raise NotImplementedError("DPO trainer is not implemented yet.")
+        from llamafactory.v1.trainers.dpo_trainer import run_dpo
+
+        run_dpo()
     elif command == "rm":
         from llamafactory.v1.trainers.rm_trainer import run_rm
 

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

@@ -14,9 +14,32 @@
 
 import json
 from copy import deepcopy
+from functools import lru_cache
 from typing import Any
 
 
+_registered_dist_config: Any | None = None
+
+
+def register_deepspeed_dist_config(dist_config: Any | None) -> None:
+    """Register backend config before model loading without involving the accelerator."""
+    global _registered_dist_config
+    _registered_dist_config = dist_config
+    is_deepspeed_zero3_enabled.cache_clear()
+
+
+@lru_cache(maxsize=1)
+def is_deepspeed_zero3_enabled() -> bool:
+    dist_config = _registered_dist_config
+    if dist_config is None or getattr(dist_config, "name", None) != "deepspeed":
+        return False
+
+    config_file = dist_config.get("config_file")
+    if not config_file:
+        return False
+    return _load_deepspeed_config(config_file).get("zero_optimization", {}).get("stage") == 3
+
+
 def _normalize_precision_enabled(value: Any) -> bool | str:
     if isinstance(value, str):
         value_lower = value.lower()
@@ -69,18 +92,19 @@ def _load_deepspeed_config(config_file: str) -> dict[str, Any]:
         return json.load(f)
 
 
-def setup_deepspeed_zero3_model_loading(is_train: bool, dist_config: dict[str, Any] | None):
-    """Enable transformers' ZeRO-3-aware model loading for the current thread."""
-    config_file = dist_config.get("config_file")
+def setup_deepspeed_zero3_model_loading():
+    """Enable ZeRO-3-aware model loading for the registered backend config."""
+    dist_config = _registered_dist_config
+    config_file = dist_config.get("config_file") if dist_config is not None else None
     if not config_file:
         raise ValueError("DeepSpeed config_file is required in dist_config")
 
     from accelerate.utils import DeepSpeedPlugin
 
     try:
-        from transformers.integrations import is_deepspeed_zero3_enabled
+        from transformers.integrations import is_deepspeed_zero3_enabled as _hf_is_deepspeed_zero3_enabled
     except ImportError:
-        from transformers.deepspeed import is_deepspeed_zero3_enabled
+        from transformers.deepspeed import is_deepspeed_zero3_enabled as _hf_is_deepspeed_zero3_enabled
 
     # DeepSpeed configs often use "auto" placeholders that only make sense once
     # we know the current runtime batch settings and precision mode.
@@ -109,7 +133,7 @@ def setup_deepspeed_zero3_model_loading(is_train: bool, dist_config: dict[str, A
     plugin.set_mixed_precision(mixed_precision)
     plugin.set_deepspeed_weakref()
 
-    if not is_deepspeed_zero3_enabled():
+    if not _hf_is_deepspeed_zero3_enabled():
         raise RuntimeError(
             "DeepSpeed ZeRO-3 model-loading bootstrap failed: transformers still reports zero3 disabled "
             "after constructing HfDeepSpeedConfig. This usually means the runtime is using a different transformers "

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -37,8 +37,8 @@ logger = logging.get_logger(__name__)
 
 
 class SequenceParallelModelPlugin(BasePlugin):
-    def __call__(self, model, model_args):
-        return super().__call__(model, model_args)
+    def __call__(self, model, cp_size: int):
+        return super().__call__(model, cp_size)
 
 
 class SequenceParallelLossPlugin(BasePlugin):
@@ -82,15 +82,17 @@ def new_flash_attn_forward(
 
 
 @SequenceParallelModelPlugin("ulysses").register()
-def apply_sequence_parallel(model, model_args):
+def apply_sequence_parallel(model, cp_size: int):
     # Replace _flash_attention_forward with new_flash_attn_forward
     module = sys.modules[model.__module__]
-    cp_size = model_args.get("cp_size", 1)
 
     set_ulysses_sequence_parallel_group(DistributedInterface().get_group(Dim.CP))
 
     try:
-        num_attention_heads, num_key_value_heads = model.config.num_attention_heads, model.config.num_attention_heads
+        num_attention_heads, num_key_value_heads = (
+            model.config.num_attention_heads,
+            model.config.num_key_value_heads,
+        )
     except AttributeError:
         num_attention_heads, num_key_value_heads = (
             model.config.text_config.num_attention_heads,

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

@@ -81,7 +81,7 @@ class UlyssesAttention(torch.nn.Module):
         query: Tensor,
         key: Tensor,
         value: Tensor,
-        attention_mask: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
         query_length: int,
         dropout_p=0.0,
         softmax_scale=None,
@@ -122,25 +122,42 @@ class UlyssesAttention(torch.nn.Module):
         if softmax_scale is None:
             softmax_scale = q.shape[-1] ** -0.5
 
+        sp_world_size = get_ulysses_sequence_parallel_world_size(self.spg)
+        local_position_ids = position_ids
+
         if position_ids is not None:
-            global_position_ids = [
-                torch.empty_like(position_ids) for _ in range(get_ulysses_sequence_parallel_world_size(self.spg))
-            ]
+            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)
             position_ids = torch.cat(global_position_ids, dim=-1).contiguous()
-            attention_mask = None
-        else:
+
+        # HF may turn an all-ones local attention_mask into None before this
+        # function. Under CP, different ranks can then disagree: some local
+        # shards still contain padding and keep a mask, while others see None.
+        # Synchronize that boolean first so every rank takes the same collective
+        # path below.
+        has_attention_mask = torch.tensor([attention_mask is not None], dtype=torch.int64, device=query.device)
+        global_has_attention_mask = [torch.empty_like(has_attention_mask) for _ in range(sp_world_size)]
+        dist.all_gather(global_has_attention_mask, has_attention_mask, group=self.spg)
+
+        # Padded path: at least one shard has real padding, so rebuild the full
+        # sequence mask for all ranks. Ranks whose local mask was optimized away
+        # contribute an all-ones shard.
+        if torch.any(torch.stack(global_has_attention_mask)):
             if attention_mask is None:
-                attention_mask = torch.ones(q.shape[0], q.shape[1], dtype=torch.int64, device=q.device)
+                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)
             else:
                 attention_mask = attention_mask.to(torch.int64)
 
-            global_attention_mask = [
-                torch.empty_like(attention_mask) for _ in range(get_ulysses_sequence_parallel_world_size(self.spg))
-            ]
+            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)
+            attention_mask = torch.cat(global_attention_mask, dim=1).contiguous()
 
+        # Packed/dense path: no rank has a mask, so leave attention_mask as None.
+        # HF can then use position_ids for padding-free packed varlen attention,
+        # or dense flash attention when position_ids are monotonic.
         context_layer = self.attn_fn(
             q,
             k,

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

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

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

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

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

@@ -1,56 +0,0 @@
-# Copyright 2025 the LlamaFactory team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import importlib
-
-from ...utils import logging
-from ...utils.plugin import BasePlugin
-from ...utils.types import Message, ModelInput, Processor
-
-
-logger = logging.get_logger(__name__)
-
-
-class RenderingPlugin(BasePlugin):
-    _attempted_template_imports: set[str] = set()
-
-    def _ensure_template_imported(self) -> None:
-        if self.name is None or self.name in self._attempted_template_imports:
-            return
-
-        full_module_name = f"{__package__}.templates.{self.name}"
-        self._attempted_template_imports.add(self.name)
-        try:
-            importlib.import_module(full_module_name)
-        except Exception as exc:
-            logger.warning(f"[Template Registry] Failed to import {full_module_name}: {exc}")
-
-    def __getitem__(self, method_name: str):
-        self._ensure_template_imported()
-        return super().__getitem__(method_name)
-
-    def render_messages(
-        self,
-        processor: Processor,
-        messages: list[Message],
-        tools: str | None = None,
-        is_generate: bool = False,
-        enable_thinking: bool = False,
-    ) -> ModelInput:
-        """Render messages in the template format."""
-        return self["render_messages"](processor, messages, tools, is_generate, enable_thinking)
-
-    def parse_messages(self, generated_text: str) -> Message:
-        """Parse messages in the template format."""
-        return self["parse_messages"](generated_text)

src/llamafactory/v1/plugins/model_plugins/templates/__init__.py

@@ -1,13 +0,0 @@
-# Copyright 2025 the LlamaFactory team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.

src/llamafactory/v1/plugins/model_plugins/templates/qwen3.py

@@ -1,259 +0,0 @@
-# Copyright 2025 the LlamaFactory team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import json
-import re
-
-from ....utils.constants import IGNORE_INDEX
-from ....utils.helper import get_tokenizer
-from ....utils.types import Message, ModelInput, Processor, ToolCall
-from ..rendering import RenderingPlugin
-
-
-def _update_model_input(
-    processor: Processor,
-    input_ids: list[int],
-    labels: list[int],
-    loss_weights: list[int],
-    temp_str: str,
-    temp_weight: float,
-) -> str:
-    """Update model input with temporary string."""
-    if not temp_str:
-        return ""
-
-    tokenizer = get_tokenizer(processor)
-    temp_ids = tokenizer.encode(temp_str, add_special_tokens=False)
-    input_ids.extend(temp_ids)
-    loss_weights.extend([temp_weight] * len(temp_ids))
-    if temp_weight > 1e-6:
-        labels.extend(temp_ids)
-    else:
-        labels.extend([IGNORE_INDEX] * len(temp_ids))
-
-    return ""
-
-
-def _concat_text_content(message: Message) -> str:
-    """Concatenate text fields in a message."""
-    message_text = ""
-    for content in message["content"]:
-        if content["type"] == "text":
-            message_text += content["value"]
-        else:
-            raise ValueError(f"Unsupported content type: {content['type']}")
-
-    return message_text
-
-
-def _get_last_query_index(messages: list[Message]) -> int:
-    """Find the last user query index, excluding wrapped tool responses."""
-    last_query_index = len(messages) - 1
-    for idx in range(len(messages) - 1, -1, -1):
-        message = messages[idx]
-        if message["role"] != "user":
-            continue
-
-        user_text = ""
-        is_plain_text = True
-        for content in message["content"]:
-            if content["type"] != "text":
-                is_plain_text = False
-                break
-            user_text += content["value"]
-
-        if not is_plain_text:
-            continue
-
-        if not (user_text.startswith("<tool_response>") and user_text.endswith("</tool_response>")):
-            last_query_index = idx
-            break
-
-    return last_query_index
-
-
-def _split_assistant_content(message: Message) -> tuple[str, str, list[ToolCall]]:
-    """Split assistant message into text, reasoning and tool calls."""
-    text_content = ""
-    reasoning_content = ""
-    tool_calls: list[ToolCall] = []
-
-    for content in message["content"]:
-        if content["type"] == "text":
-            text_content += content["value"]
-        elif content["type"] == "reasoning":
-            reasoning_content += content["value"]
-        elif content["type"] == "tool_call":
-            try:
-                tool_call: ToolCall = json.loads(content["value"])
-            except json.JSONDecodeError:
-                raise ValueError(f"Invalid tool call format: {content['value']}.")
-
-            tool_calls.append(tool_call)
-        else:
-            raise ValueError(f"Unsupported content type: {content['type']}")
-
-    return text_content, reasoning_content, tool_calls
-
-
-@RenderingPlugin("qwen3").register("render_messages")
-def render_qwen3_messages(
-    processor: Processor,
-    messages: list[Message],
-    tools: str | None = None,
-    is_generate: bool = False,
-    enable_thinking: bool = False,
-) -> ModelInput:
-    """Render messages in the Qwen3 template format.
-
-    See https://huggingface.co/spaces/huggingfacejs/chat-template-playground?modelId=Qwen/Qwen3-8B
-    """
-    input_ids, labels, loss_weights = [], [], []
-    temp_str, temp_weight = "", 0.0
-    if tools:
-        temp_str += "<|im_start|>system\n"
-        if messages[0]["role"] == "system":
-            temp_str += _concat_text_content(messages[0]) + "\n\n"
-            temp_weight = messages[0].get("loss_weight", 0.0)
-
-        temp_str += (
-            "# Tools\n\nYou may call one or more functions to assist with the user query.\n\n"
-            "You are provided with function signatures within <tools></tools> XML tags:\n<tools>"
-        )
-        try:
-            tools = json.loads(tools)
-        except json.JSONDecodeError:
-            raise ValueError(f"Invalid tools format: {str(tools)}.")
-
-        if not isinstance(tools, list):
-            tools = [tools]
-
-        for tool in tools:
-            temp_str += "\n" + json.dumps(tool, ensure_ascii=False)
-
-        temp_str += (
-            "\n</tools>\n\nFor each function call, return a json object with function name "
-            'and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{"name": '
-            '<function-name>, "arguments": <args-json-object>}\n</tool_call><|im_end|>\n'
-        )
-    elif messages[0]["role"] == "system":
-        temp_str += "<|im_start|>system\n" + _concat_text_content(messages[0]) + "<|im_end|>\n"
-        temp_weight = messages[0].get("loss_weight", 0.0)
-
-    temp_str = _update_model_input(processor, input_ids, labels, loss_weights, temp_str, temp_weight)
-    last_query_index = _get_last_query_index(messages)
-
-    for turn_idx, message in enumerate(messages):
-        if message["role"] == "user" or (message["role"] == "system" and turn_idx != 0):
-            temp_str += "<|im_start|>" + message["role"] + "\n" + _concat_text_content(message) + "<|im_end|>\n"
-            temp_weight = message.get("loss_weight", 0.0)
-        elif message["role"] == "assistant":
-            temp_str += "<|im_start|>" + message["role"] + "\n"
-
-            text_content, reasoning_content, tool_calls = _split_assistant_content(message)
-            if turn_idx > last_query_index and (turn_idx == len(messages) - 1 or reasoning_content):
-                temp_str += "<think>\n" + reasoning_content.strip("\n") + "\n</think>\n\n" + text_content.lstrip("\n")
-            else:
-                temp_str += text_content
-
-            for tool_call_idx, tool_call in enumerate(tool_calls):
-                if (tool_call_idx == 0 and text_content) or tool_call_idx > 0:
-                    temp_str += "\n"
-
-                arguments = tool_call.get("arguments")
-                if isinstance(arguments, str):
-                    arguments_str = arguments
-                else:
-                    arguments_str = json.dumps(arguments, ensure_ascii=False)
-
-                temp_str += (
-                    '<tool_call>\n{"name": "'
-                    + tool_call["name"]
-                    + '", "arguments": '
-                    + arguments_str
-                    + "}\n</tool_call>"
-                )
-
-            temp_str += "<|im_end|>\n"
-            temp_weight = message.get("loss_weight", 1.0)
-        elif message["role"] == "tool":
-            if turn_idx == 0 or messages[turn_idx - 1]["role"] != "tool":
-                temp_str += "<|im_start|>user"
-
-            temp_str += "\n<tool_response>\n" + _concat_text_content(message) + "\n</tool_response>"
-            if turn_idx == len(messages) - 1 or messages[turn_idx + 1]["role"] != "tool":
-                temp_str += "<|im_end|>\n"
-
-            temp_weight = message.get("loss_weight", 0.0)
-
-        temp_str = _update_model_input(processor, input_ids, labels, loss_weights, temp_str, temp_weight)
-
-    if is_generate:
-        temp_str += "<|im_start|>assistant\n"
-        temp_weight = 0.0
-        if enable_thinking is False:
-            temp_str += "<think>\n\n</think>\n\n"
-
-    temp_str = _update_model_input(processor, input_ids, labels, loss_weights, temp_str, temp_weight)
-
-    attention_mask = [1] * len(input_ids)
-    return ModelInput(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
-        labels=labels,
-        loss_weights=loss_weights,
-    )
-
-
-@RenderingPlugin("qwen3").register("parse_message")
-def parse_qwen3_message(generated_text: str) -> Message:
-    """Parse a message in the Qwen3 template format. Supports interleaved reasoning and tool calls.
-
-    Args:
-        generated_text (str): The generated text in the Qwen3 template format.
-
-    Returns:
-        Message: The parsed message.
-    """
-    pattern = re.compile(r"<(think|tool_call)>\s*(.*?)\s*</\1>\s*", re.DOTALL)
-    content = []
-    last_end = 0
-
-    for match in pattern.finditer(generated_text):
-        start, end = match.span()
-        if start > last_end:
-            text = generated_text[last_end:start].strip()
-            if text:
-                content.append({"type": "text", "value": text})
-
-        tag_type = match.group(1)
-        tag_value = match.group(2).strip()
-        if tag_type == "think":
-            content.append({"type": "reasoning", "value": tag_value.strip()})
-        elif tag_type == "tool_call":
-            try:
-                json.loads(tag_value.strip())
-            except json.JSONDecodeError:
-                raise ValueError(f"Invalid tool call format: {tag_value.strip()}.")
-
-            content.append({"type": "tool_call", "value": tag_value.strip()})
-
-        last_end = end
-
-    if last_end < len(generated_text):
-        text = generated_text[last_end:].strip()
-        if text:
-            content.append({"type": "text", "value": text})
-
-    return Message(role="assistant", content=content)

src/llamafactory/v1/plugins/model_plugins/templates/qwen3_nothink.py

@@ -1,209 +0,0 @@
-# Copyright 2025 the LlamaFactory team.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import json
-import re
-
-from ....utils.constants import IGNORE_INDEX
-from ....utils.helper import get_tokenizer
-from ....utils.types import Message, ModelInput, Processor, ToolCall
-from ..rendering import RenderingPlugin
-
-
-def _update_model_input(
-    processor: Processor,
-    input_ids: list[int],
-    labels: list[int],
-    loss_weights: list[int],
-    temp_str: str,
-    temp_weight: float,
-) -> str:
-    """Update model input with temporary string."""
-    if not temp_str:
-        return ""
-
-    tokenizer = get_tokenizer(processor)
-    temp_ids = tokenizer.encode(temp_str, add_special_tokens=False)
-    input_ids.extend(temp_ids)
-    loss_weights.extend([temp_weight] * len(temp_ids))
-    if temp_weight > 1e-6:
-        labels.extend(temp_ids)
-    else:
-        labels.extend([IGNORE_INDEX] * len(temp_ids))
-
-    return ""
-
-
-def _concat_text_content(message: Message) -> str:
-    """Concatenate text fields in a message."""
-    message_text = ""
-    for content in message["content"]:
-        if content["type"] == "text":
-            message_text += content["value"]
-        else:
-            raise ValueError(f"Unsupported content type: {content['type']}")
-
-    return message_text
-
-
-@RenderingPlugin("qwen3_nothink").register("render_messages")
-def render_qwen3_nothink_messages(
-    processor: Processor,
-    messages: list[Message],
-    tools: str | None = None,
-    is_generate: bool = False,
-    enable_thinking: bool = False,
-) -> ModelInput:
-    """Render messages in the Qwen3 nothink template format.
-
-    See https://huggingface.co/spaces/huggingfacejs/chat-template-playground?modelId=Qwen/Qwen3-4B-Instruct-2507
-    """
-    input_ids, labels, loss_weights = [], [], []
-    temp_str, temp_weight = "", 0.0
-    if tools:
-        temp_str += "<|im_start|>system\n"
-        if messages[0]["role"] == "system":
-            temp_str += _concat_text_content(messages[0]) + "\n\n"
-            temp_weight = messages[0].get("loss_weight", 0.0)
-
-        temp_str += (
-            "# Tools\n\nYou may call one or more functions to assist with the user query.\n\n"
-            "You are provided with function signatures within <tools></tools> XML tags:\n<tools>"
-        )
-
-        try:
-            tools = json.loads(tools)
-        except json.JSONDecodeError:
-            raise ValueError(f"Invalid tools format: {str(tools)}.")
-
-        if not isinstance(tools, list):
-            tools = [tools]
-
-        for tool in tools:
-            temp_str += "\n" + json.dumps(tool, ensure_ascii=False)
-
-        temp_str += (
-            "\n</tools>\n\nFor each function call, return a json object with function name "
-            'and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{"name": '
-            '<function-name>, "arguments": <args-json-object>}\n</tool_call><|im_end|>\n'
-        )
-    elif messages[0]["role"] == "system":
-        temp_str += "<|im_start|>system\n" + _concat_text_content(messages[0]) + "<|im_end|>\n"
-        temp_weight = messages[0].get("loss_weight", 0.0)
-
-    temp_str = _update_model_input(processor, input_ids, labels, loss_weights, temp_str, temp_weight)
-
-    for turn_idx, message in enumerate(messages):
-        if message["role"] == "user" or (message["role"] == "system" and turn_idx != 0):
-            temp_str += "<|im_start|>" + message["role"] + "\n" + _concat_text_content(message) + "<|im_end|>\n"
-            temp_weight = message.get("loss_weight", 0.0)
-        elif message["role"] == "assistant":
-            temp_str += "<|im_start|>" + message["role"] + "\n"
-            for val_idx, content in enumerate(message["content"]):
-                if content["type"] == "text":
-                    temp_str += content["value"]
-                elif content["type"] == "reasoning":
-                    temp_str += "<thinking>\n" + content["value"] + "\n</thinking>\n\n"  # avoid using special tokens
-                elif content["type"] == "tool_call":
-                    if val_idx != 0 and message["content"][val_idx - 1]["type"] in ["text", "tool_call"]:
-                        temp_str += "\n"
-
-                    try:
-                        tool_call: ToolCall = json.loads(content["value"])
-                    except json.JSONDecodeError:
-                        raise ValueError(f"Invalid tool call format: {content['value']}.")
-
-                    temp_str += (
-                        '<tool_call>\n{"name": "'
-                        + tool_call["name"]
-                        + '", "arguments": '
-                        + json.dumps(tool_call["arguments"], ensure_ascii=False)
-                        + "}\n</tool_call>"
-                    )
-
-                else:
-                    raise ValueError(f"Unsupported content type: {content['type']}")
-
-            temp_str += "<|im_end|>\n"
-            temp_weight = message.get("loss_weight", 1.0)
-        elif message["role"] == "tool":
-            if turn_idx == 0 or messages[turn_idx - 1]["role"] != "tool":
-                temp_str += "<|im_start|>user"
-
-            temp_str += "\n<tool_response>\n" + _concat_text_content(message) + "\n</tool_response>"
-            if turn_idx == len(messages) - 1 or messages[turn_idx + 1]["role"] != "tool":
-                temp_str += "<|im_end|>\n"
-
-            temp_weight = message.get("loss_weight", 0.0)
-
-        temp_str = _update_model_input(processor, input_ids, labels, loss_weights, temp_str, temp_weight)
-
-    if is_generate:
-        temp_str += "<|im_start|>assistant\n"
-        temp_weight = 0.0
-        if enable_thinking:
-            raise ValueError("The qwen3_nothink template does not support thinking mode.")
-
-    temp_str = _update_model_input(processor, input_ids, labels, loss_weights, temp_str, temp_weight)
-
-    attention_mask = [1] * len(input_ids)
-    return ModelInput(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
-        labels=labels,
-        loss_weights=loss_weights,
-    )
-
-
-@RenderingPlugin("qwen3_nothink").register("parse_message")
-def parse_qwen3_nothink_message(generated_text: str) -> Message:
-    """Parse a message in the Qwen3 nothink template format. Supports interleaved reasoning and tool calls.
-
-    Args:
-        generated_text (str): The generated text in the Qwen3 nothink template format.
-
-    Returns:
-        Message: The parsed message.
-    """
-    pattern = re.compile(r"<(thinking|tool_call)>\s*(.*?)\s*</\1>\s*", re.DOTALL)
-    content = []
-    last_end = 0
-
-    for match in pattern.finditer(generated_text):
-        start, end = match.span()
-        if start > last_end:
-            text = generated_text[last_end:start].strip()
-            if text:
-                content.append({"type": "text", "value": text})
-
-        tag_type = match.group(1)
-        tag_value = match.group(2).strip()
-        if tag_type == "thinking":
-            content.append({"type": "reasoning", "value": tag_value.strip()})
-        elif tag_type == "tool_call":
-            try:
-                json.loads(tag_value.strip())
-            except json.JSONDecodeError:
-                raise ValueError(f"Invalid tool call format: {tag_value.strip()}.")
-
-            content.append({"type": "tool_call", "value": tag_value.strip()})
-
-        last_end = end
-
-    if last_end < len(generated_text):
-        text = generated_text[last_end:].strip()
-        if text:
-            content.append({"type": "text", "value": text})
-
-    return Message(role="assistant", content=content)

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,297 @@
+# 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.
+#
+# This module is vendored into v1 (independent of v0 / `llamafactory.third_party.muon`)
+# so that the v1 optimizer plugin does not depend on v0 code.
+#
+# Based on MoonshotAI's Moonlight library and Keller Jordan's Muon library:
+#   https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py
+#   https://github.com/KellerJordan/Muon/blob/master/muon.py
+# (originally MIT-licensed; re-distributed here under Apache 2.0).
+
+import math
+import os
+
+import torch
+import torch.distributed as dist
+
+
+def _dtensor_cls():
+    """Return the DTensor class if available, else None."""
+    try:
+        from torch.distributed.tensor import DTensor
+    except ImportError:  # pragma: no cover
+        try:
+            from torch.distributed._tensor import DTensor  # type: ignore[no-redef]
+        except ImportError:
+            return None
+    return DTensor
+
+
+def _is_dtensor(t) -> bool:
+    """True if ``t`` is a DTensor (i.e. sharded by FSDP2)."""
+    DT = _dtensor_cls()
+    return DT is not None and isinstance(t, DT)
+
+
+def _distribute(tensor, mesh, placements):
+    """Scatter a full (replicated) tensor into a DTensor with the given mesh/placements."""
+    try:
+        from torch.distributed.tensor import distribute_tensor
+    except ImportError:  # pragma: no cover
+        from torch.distributed._tensor import distribute_tensor  # type: ignore[no-redef]
+    return distribute_tensor(tensor, mesh, placements)
+
+
+def _is_rank0() -> bool:
+    """True on rank 0 (or when not distributed)."""
+    return not (dist.is_available() and dist.is_initialized()) or dist.get_rank() == 0
+
+
+def zeropower_via_newtonschulz5(G: "torch.Tensor", steps: int) -> "torch.Tensor":
+    """Newton-Schulz iteration to compute the zeroth power / orthogonalization of G.
+
+    We opt to use a quintic iteration whose coefficients are selected to maximize the slope at zero.
+    For the purpose of minimizing steps, it turns out to be empirically effective to keep increasing
+    the slope at zero even beyond the point where the iteration no longer converges all the way to
+    one everywhere on the interval. This iteration therefore does not produce UV^T but rather something
+    like US'V^T where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
+    performance at all relative to UV^T, where USV^T = G is the SVD.
+
+    Computation runs in ``bfloat16`` and the result is returned in ``bfloat16`` by design (NS is
+    stable in bf16, matching upstream Keller Jordan / Moonlight). The caller's in-place ``add_``
+    upcasts the operand to the parameter dtype, so no cast-back to ``G.dtype`` is needed.
+    """
+    assert len(G.shape) == 2
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    X = G.bfloat16()
+    if G.size(0) > G.size(1):
+        X = X.T
+    # Ensure spectral norm is at most 1
+    X = X / (X.norm() + 1e-7)
+    # Perform the NS iterations
+    for _ in range(steps):
+        A = X @ X.T
+        B = b * A + c * A @ A  # adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
+        X = a * X + B @ X
+
+    if G.size(0) > G.size(1):
+        X = X.T
+    return X
+
+
+class Muon(torch.optim.Optimizer):
+    """Muon - MomentUm Orthogonalized by Newton-schulz.
+
+    Muon internally runs standard SGD-momentum, and then performs an orthogonalization post-
+    processing step, in which each 2D parameter's update is replaced with the nearest orthogonal
+    matrix. To efficiently orthogonalize each update, we use a Newton-Schulz iteration, which has
+    the advantage that it can be stably run in bfloat16 on the GPU.
+
+    Some warnings:
+    - We believe this optimizer is unlikely to work well for training with small batch size.
+    - We believe it may not work well for finetuning pretrained models, but we haven't tested this.
+
+    Arguments:
+        muon_params: The parameters to be optimized by Muon.
+        lr: The learning rate. The updates will have spectral norm of `lr`. (0.02 is a good default)
+        wd: The weight decay.
+        momentum: The momentum used by the internal SGD. (0.95 is a good default)
+        nesterov: Whether to use Nesterov-style momentum in the internal SGD. (recommended)
+        ns_steps: The number of Newton-Schulz iterations to run. (6 is probably always enough)
+        adamw_params: The parameters to be optimized by AdamW. Any parameters in `muon_params` which are
+        {0, 1}-D or are detected as being the embed or lm_head will be optimized by AdamW as well.
+        adamw_betas: The betas for the internal AdamW.
+        adamw_eps: The epsilon for the internal AdamW.
+    """
+
+    def __init__(
+        self,
+        lr=1e-3,
+        wd=0.1,
+        muon_params=None,
+        momentum=0.95,
+        nesterov=True,
+        ns_steps=5,
+        adamw_params=None,
+        adamw_betas=(0.9, 0.95),
+        adamw_eps=1e-8,
+    ):
+        defaults = dict(
+            lr=lr,
+            wd=wd,
+            momentum=momentum,
+            nesterov=nesterov,
+            ns_steps=ns_steps,
+            adamw_betas=adamw_betas,
+            adamw_eps=adamw_eps,
+        )
+
+        params = list(muon_params)
+        adamw_params = list(adamw_params) if adamw_params is not None else []
+        params.extend(adamw_params)
+        super().__init__(params, defaults)
+        # Sort parameters into those for which we will use Muon, and those for which we will not
+        for p in muon_params:
+            # Use Muon for every parameter in muon_params which is >= 2D and doesn't look like an embedding or head layer
+            assert p.ndim == 2, p.ndim
+            self.state[p]["use_muon"] = True
+        for p in adamw_params:
+            # Do not use Muon for parameters in adamw_params
+            self.state[p]["use_muon"] = False
+
+        self._diag_done = False
+
+    def _v2_diag(self, p) -> None:
+        """Print (once, rank0) the param/grad/data types needed to implement the DTensor-aware v2.
+
+        Gate with env var LLAMAFACTORY_MUON_DIAG=1 so it is opt-in.
+        """
+        self._diag_done = True
+        if os.environ.get("LLAMAFACTORY_MUON_DIAG") != "1":
+            return
+        if not _is_rank0():
+            return
+        DT = _dtensor_cls()
+        g = p.grad
+        is_dt = (DT is not None) and isinstance(p, DT)
+        is_g_dt = (DT is not None) and isinstance(g, DT)
+        lines = ["[Muon v2-diag] === info for writing the DTensor-aware v2 ==="]
+        lines.append(f"  param: type={type(p).__name__} is_DT={is_dt} shape={tuple(p.shape)}")
+        if is_dt:
+            lines.append(f"    placements={p.placements} device_mesh={p.device_mesh}")
+            try:
+                lines.append(f"    p.to_local().shape={tuple(p.to_local().shape)}")
+            except Exception as e:  # noqa: BLE001
+                lines.append(f"    p.to_local() ERR={e!r}")
+        lines.append(f"  grad:  type={type(g).__name__} is_DT={is_g_dt} shape={tuple(g.shape)}")
+        lines.append(f"    grad.has_full_tensor={hasattr(g, 'full_tensor')}")
+        if is_g_dt:
+            lines.append(f"    grad.placements={g.placements} grad.device_mesh={g.device_mesh}")
+        lines.append(f"  p.data: type={type(p.data).__name__} shape={tuple(p.data.shape)}")
+        lines.append(
+            f"  compare: p.shape==p.data.shape ? {tuple(p.shape) == tuple(p.data.shape)} ; "
+            f"grad.shape==p.data.shape ? {tuple(g.shape) == tuple(p.data.shape)}"
+        )
+        print("\n".join(lines), flush=True)
+
+    def adjust_lr_for_muon(self, lr: float, param_shape: list[int]) -> float:
+        A, B = param_shape[:2]
+        # We adjust the learning rate and weight decay based on the size of the parameter matrix
+        # as described in the paper
+        adjusted_ratio = 0.2 * math.sqrt(max(A, B))
+        adjusted_lr = lr * adjusted_ratio
+        return adjusted_lr
+
+    def step(self, closure=None):
+        """Perform a single optimization step.
+
+        Args:
+            closure (Callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            # Muon loop
+            params = [p for p in group["params"] if self.state[p]["use_muon"]]
+            lr = group["lr"]
+            wd = group["wd"]
+            momentum = group["momentum"]
+
+            # generate weight updates in distributed fashion
+            for p in params:
+                # sanity check
+                g = p.grad
+                if g is None:
+                    continue
+                if not self._diag_done:
+                    self._v2_diag(p)
+
+                state = self.state[p]
+
+                # v2: under FSDP2, p.grad is a sharded DTensor. Newton-Schulz must run on
+                # the FULL 2D matrix (running it on the local shard computes a partial Gram
+                # matrix and the NS iteration diverges -> NaN). Momentum accumulation is
+                # elementwise, so the momentum buffer is kept sharded (mirroring g's
+                # placements -> 1/N memory and FSDP2-checkpoint-native); we all-gather only
+                # for the NS step, then scatter the update back to the local shard.
+                sharded = _is_dtensor(g)
+                if sharded:
+                    p_mesh, p_placements = p.device_mesh, p.placements
+                else:
+                    p_mesh = p_placements = None
+
+                # momentum buffer mirrors g's sharding (sharded DTensor under FSDP2, plain
+                # tensor otherwise); elementwise accumulation is correct on the local shard.
+                if "momentum_buffer" not in state:
+                    state["momentum_buffer"] = torch.zeros_like(g)
+                buf = state["momentum_buffer"]
+                buf.mul_(momentum).add_(g)
+                g_use = g.add(buf, alpha=momentum) if group["nesterov"] else buf
+
+                # all-gather ONLY here: NS needs the full 2D matrix (Gram matrix X @ X.T).
+                g_full = g_use.full_tensor() if sharded else g_use
+                if g_full.ndim > 2:
+                    g_full = g_full.view(g_full.size(0), -1)
+                u_full = zeropower_via_newtonschulz5(g_full, steps=group["ns_steps"])
+
+                # scale update (p.shape is the DTensor global shape -> correct A, B)
+                adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
+
+                # apply weight decay (in-place on the local shard; elementwise -> correct)
+                p.data.mul_(1 - lr * wd)
+
+                # apply update; scatter the full update back to the local shard under FSDP2
+                if sharded:
+                    u_dt = _distribute(u_full, p_mesh, p_placements)
+                    p.data.add_(u_dt, alpha=-adjusted_lr)
+                else:
+                    p.data.add_(u_full, alpha=-adjusted_lr)
+
+            # Adam backup
+            params = [p for p in group["params"] if not self.state[p]["use_muon"]]
+            lr = group["lr"]
+            beta1, beta2 = group["adamw_betas"]
+            eps = group["adamw_eps"]
+            weight_decay = group["wd"]
+
+            for p in params:
+                g = p.grad
+                if g is None:
+                    continue
+                state = self.state[p]
+                if "step" not in state:
+                    state["step"] = 0
+                    state["moment1"] = torch.zeros_like(g)
+                    state["moment2"] = torch.zeros_like(g)
+                state["step"] += 1
+                step = state["step"]
+                buf1 = state["moment1"]
+                buf2 = state["moment2"]
+                buf1.lerp_(g, 1 - beta1)
+                buf2.lerp_(g.square(), 1 - beta2)
+
+                g = buf1 / (eps + buf2.sqrt())
+
+                bias_correction1 = 1 - beta1**step
+                bias_correction2 = 1 - beta2**step
+                scale = bias_correction1 / bias_correction2**0.5
+                p.data.mul_(1 - lr * weight_decay)
+                p.data.add_(g, alpha=-lr / scale)
+
+        return loss

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

@@ -0,0 +1,81 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from ....utils import logging
+from ....utils.plugin import BasePlugin
+
+
+if TYPE_CHECKING:
+    from ....config.arg_utils import PluginConfig
+    from ....utils.types import HFModel
+
+
+logger = logging.get_logger(__name__)
+
+
+class OptimizerPlugin(BasePlugin):
+    pass
+
+
+@OptimizerPlugin("muon").register()
+def create_muon_optimizer(model: HFModel, optim_config: PluginConfig):
+    """Create a Muon optimizer.
+
+    Muon is used for 2D "hidden" weight matrices; the remaining parameters (1D bias/LayerNorm,
+    embeddings incl. GPT-2 ``wte``/``wpe``, the output ``lm_head``, and LoRA adapter factors) are
+    optimized by the built-in AdamW.
+
+    The Muon step is DTensor-aware: under FSDP2 it all-gathers the full gradient, runs Newton-Schulz
+    on the full 2D matrix, then scatters the update back to the local shard. So it is correct under
+    FSDP2 / sequence parallel (no longer approximate).
+    """
+    from .muon_optimizer import Muon
+
+    muon_params, adamw_params = [], []
+    for name, param in model.named_parameters():
+        if param.requires_grad:
+            # Muon is only appropriate for 2D "hidden" weight matrices. Route everything else to
+            # the internal AdamW: 1D bias/norm, embeddings ("embed", GPT-2 "wte"/"wpe"), the output
+            # head ("lm_head"), and LoRA adapter factors ("lora_A"/"lora_B"/"lora_embedding_*").
+            if (
+                param.ndim == 2
+                and "embed" not in name
+                and "lm_head" not in name
+                and "wte" not in name
+                and "wpe" not in name
+                and "lora" not in name
+            ):
+                muon_params.append(param)
+            else:
+                adamw_params.append(param)
+
+    optimizer = Muon(
+        lr=optim_config.get("lr", 1e-3),
+        wd=optim_config.get("wd", 0.1),
+        muon_params=muon_params,
+        momentum=optim_config.get("momentum", 0.95),
+        nesterov=optim_config.get("nesterov", True),
+        ns_steps=optim_config.get("ns_steps", 5),
+        adamw_params=adamw_params,
+        adamw_betas=tuple(optim_config.get("adamw_betas", [0.9, 0.95])),
+        adamw_eps=optim_config.get("adamw_eps", 1e-8),
+    )
+    logger.info_rank0(
+        f"Using Muon optimizer with {len(muon_params)} Muon params and {len(adamw_params)} AdamW params."
+    )
+    return optimizer

src/llamafactory/v1/samplers/cli_sampler.py

@@ -21,7 +21,7 @@ from ..config import InputArgument, ModelArguments, SampleArguments, SampleBacke
 from ..core.base_sampler import BaseSampler
 from ..core.data_engine import DataEngine
 from ..core.model_engine import ModelEngine
-from ..core.utils.rendering import Renderer
+from ..core.rendering import Renderer
 from ..utils.types import HFModel, Message, Sample, TorchDataset
 
 
@@ -118,7 +118,7 @@ def run_chat(args: InputArgument = None):
                 response += new_text
 
             print()
-            messages.append(model_engine.renderer.parse_message(response))
+            messages.append({"role": "assistant", "content": [{"type": "text", "value": response}]})
 
 
 if __name__ == "__main__":