[v1] add FSDPTurbo EP/EFSDP plugin for MoE training (#10676)

.github/workflows/docs.yml

@@ -66,6 +66,7 @@ jobs:
           path: docs/_build/html
 
   deploy:
+    if: github.event_name == 'push' && github.ref == 'refs/heads/main'
     environment:
       name: github-pages
       url: ${{ steps.deployment.outputs.page_url }}

.github/workflows/tests_npu.yml

@@ -8,6 +8,7 @@ on:
     paths:
       - "**/*.py"
       - "pyproject.toml"
+      - "requirements/fsdpturbo.txt"
       - "Makefile"
       - ".github/workflows/*.yml"
   pull_request:
@@ -16,6 +17,7 @@ on:
     paths:
       - "**/*.py"
       - "pyproject.toml"
+      - "requirements/fsdpturbo.txt"
       - "Makefile"
       - ".github/workflows/*.yml"
 
@@ -68,6 +70,7 @@ jobs:
           uv pip install -e .
           uv pip install -r requirements/npu.txt
           uv pip install -r requirements/dev.txt
+          uv pip install --no-deps -r requirements/fsdpturbo.txt
 
       - name: Install node
         run: |

docs/en/advanced/distributed/fsdpturbo-ep-efsdp.md

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

docs/en/advanced/ktransformers.md

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

docs/en/index.rst

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

docs/zh/advanced/distributed/fsdpturbo-ep-efsdp.md

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

docs/zh/advanced/ktransformers.md

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

docs/zh/index.rst

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

examples/ktransformers/accelerate/fsdp2_kt_bf16.yaml

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

examples/ktransformers/accelerate/fsdp2_kt_int4.yaml

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

examples/ktransformers/accelerate/fsdp2_kt_int8.yaml

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

examples/ktransformers/accelerate/fsdp2_kt_int8_1gpu.yaml

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

examples/ktransformers/accelerate/fsdp2_kt_int8_8gpu.yaml

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

examples/ktransformers/train_lora/deepseek_v3_int8_lora_sft_kt.yaml

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

examples/ktransformers/train_lora/deepseek_v3_lora_sft_kt.yaml

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

examples/ktransformers/train_lora/qwen3_5moe_lora_sft_kt.yaml

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

examples/v1/train_full/train_full_qwen3_moe_fsdpturbo_ep_fsdp.yaml

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

requirements/fsdpturbo.txt

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

src/llamafactory/hparams/model_args.py

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

src/llamafactory/hparams/parser.py

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

src/llamafactory/model/adapter.py

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

src/llamafactory/model/model_utils/attention.py

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

src/llamafactory/model/model_utils/checkpointing.py

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

src/llamafactory/model/model_utils/quantization.py

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

src/llamafactory/v1/core/base_trainer.py

@@ -294,19 +294,24 @@ class BaseTrainer:
                     # deepspeed: engine.step() already ran inside backward at the sync boundary
                     grad_norm = self._deepspeed_engine.get_grad_norm()
                 else:
-                    # FSDP2 shards params/grads across the fsdp mesh, so clip_grad_norm_ returns a
-                    # per-rank local shard norm (global / sqrt(shard_size)): reported grad_norm then
-                    # scales as 1/sqrt(dp_size) and the clip coefficient is applied per-shard. Reduce
-                    # to the true global norm first, then clip with it.
-                    grads = [p.grad for p in self.model.parameters() if p.grad is not None]
-                    total_norm = torch.nn.utils.get_total_norm(grads)
-                    if isinstance(total_norm, DTensor):
-                        # full_tensor all-reduces across the fsdp mesh (spans CP under default
-                        # mp_shard=world); a separate CP reduce would over-count by sqrt(cp_size).
-                        total_norm = total_norm.full_tensor()
-                    # pass a Tensor: clip_grads_with_norm_ clamps max_norm / (total_norm + 1e-6).
-                    torch.nn.utils.clip_grads_with_norm_(self.model.parameters(), self.args.max_grad_norm, total_norm)
-                    grad_norm = total_norm.item()
+                    dist_name = self.args.dist_config.name if self.args.dist_config else None
+                    if dist_name == "fsdpturbo":
+                        from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
+
+                        grad_norm = DistributedPlugin(dist_name).clip_grad_norm(self.model, self.args.max_grad_norm)
+                    else:
+                        # FSDP2 shards params/grads across the fsdp mesh, so clip_grad_norm_ returns a
+                        # per-rank local shard norm. Materialize the true global norm before clipping.
+                        grads = [p.grad for p in self.model.parameters() if p.grad is not None]
+                        total_norm = torch.nn.utils.get_total_norm(grads)
+                        if isinstance(total_norm, DTensor):
+                            # full_tensor all-reduces across the fsdp mesh (spans CP under default
+                            # mp_shard=world); a separate CP reduce would over-count by sqrt(cp_size).
+                            total_norm = total_norm.full_tensor()
+                        torch.nn.utils.clip_grads_with_norm_(
+                            self.model.parameters(), self.args.max_grad_norm, total_norm
+                        )
+                        grad_norm = total_norm.item()
 
                     if not torch.isfinite(torch.tensor(grad_norm)):  # type: ignore # pyright: ignore [reportUnknownReturnType]
                         logger.warning_rank0(f"Gradient norm is not finite: {grad_norm}")
@@ -363,7 +368,7 @@ class BaseTrainer:
 
     def save_model(self) -> None:
         """Save the model."""
-        if self.args.dist_config is not None and self.args.dist_config.name in ("deepspeed", "fsdp2"):
+        if self.args.dist_config is not None and self.args.dist_config.name in ("deepspeed", "fsdp2", "fsdpturbo"):
             from ..plugins.trainer_plugins.distributed.interface import DistributedPlugin
 
             DistributedPlugin(self.args.dist_config.name).save_model(

src/llamafactory/v1/core/model_engine.py

@@ -112,8 +112,7 @@ class ModelEngine:
         if self.args.custom_chat_template:
             if not is_tokenizer(self.processor):
                 self.processor.chat_template = self.args.custom_chat_template
-            else:
-                tokenizer.chat_template = self.args.custom_chat_template
+            tokenizer.chat_template = self.args.custom_chat_template
 
     def _init_model_config(self) -> HFConfig:
         """Init model config."""
@@ -184,7 +183,7 @@ class ModelEngine:
         if init_device.type == DeviceType.META:
             assert self.args.quant_config is None, "Quantization is not supported with meta device."
             with init_empty_weights():
-                model = AutoClass.from_config(self.model_config)
+                model = AutoClass.from_config(self.model_config, attn_implementation=self.args.flash_attn)
         else:
             model = AutoClass.from_pretrained(
                 self.args.model,

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

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

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

@@ -20,6 +20,7 @@ from .base import KernelPlugin
 
 # Import built-in implementations so their class decorators populate the registry.
 from .liger_kernel_ops import LigerKernel  # noqa: F401
+from .ops.linear_attention.fla import FlashLinearAttentionKernel  # noqa: F401
 from .ops.mlp.cuda_fused_moe import CudaFusedMoEKernel  # noqa: F401
 from .ops.mlp.npu_fused_moe import NpuFusedMoEKernel  # noqa: F401
 from .ops.mlp.npu_swiglu import NpuSwiGluKernel  # noqa: F401

src/llamafactory/v1/plugins/model_plugins/kernels/ops/linear_attention/fla.py

@@ -0,0 +1,103 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Flash Linear Attention kernel plugin backed by FSDPTurbo's operator registry."""
+
+from functools import partial
+
+from ......accelerator.helper import DeviceType, get_current_accelerator
+from ......utils import logging
+from ......utils.types import HFModel
+from ...base import BaseKernel, KernelPlugin
+
+
+logger = logging.get_logger(__name__)
+
+CHUNK_GATED_DELTA_RULE = "chunk_gated_delta_rule"
+FUSED_RECURRENT_GATED_DELTA_RULE = "fused_recurrent_gated_delta_rule"
+FLASH_LINEAR_ATTENTION_KERNELS = (
+    CHUNK_GATED_DELTA_RULE,
+    FUSED_RECURRENT_GATED_DELTA_RULE,
+)
+FLA_MODULE_ATTRIBUTES = {
+    CHUNK_GATED_DELTA_RULE: "chunk_gated_delta_rule",
+    FUSED_RECURRENT_GATED_DELTA_RULE: "recurrent_gated_delta_rule",
+}
+SUPPORTED_CHUNK_SIZES = (16, 32, 64)
+
+
+@KernelPlugin("flash-linear-attention").register()
+class FlashLinearAttentionKernel(BaseKernel):
+    """Install selected FLA callables through FSDPTurbo's device operator registry."""
+
+    @staticmethod
+    def check_device() -> None:
+        current = get_current_accelerator().type
+        if current not in (DeviceType.CUDA, DeviceType.NPU):
+            raise RuntimeError(f"FlashLinearAttentionKernel requires CUDA or NPU, current accelerator is {current}.")
+
+    @staticmethod
+    def check_deps() -> None:
+        try:
+            import fla.ops.gated_delta_rule  # noqa: F401
+            import fsdp_turbo.ops.fla  # noqa: F401
+            from fsdp_turbo.ops.registry import get_op  # noqa: F401
+            from fsdp_turbo.utils.patch import patch_model_members  # noqa: F401
+        except ImportError as exc:
+            raise RuntimeError("Flash Linear Attention and FSDPTurbo are required for this kernel.") from exc
+
+    @staticmethod
+    def _apply(**kwargs) -> HFModel:
+        model = kwargs["model"]
+        config = kwargs.get("config") or {}
+        include_kernels = config.get("include_kernels", "auto")
+        chunk_size = config.get("chunk_size", 64)
+
+        if include_kernels == "auto" or include_kernels is True:
+            selected = list(FLASH_LINEAR_ATTENTION_KERNELS)
+        elif isinstance(include_kernels, str):
+            selected = [name.strip() for name in include_kernels.split(",") if name.strip()]
+        else:
+            raise TypeError("kernel_config.include_kernels must be 'auto' or a comma-separated string.")
+
+        if not selected:
+            raise ValueError("kernel_config.include_kernels must select at least one FLA kernel.")
+
+        unsupported = set(selected).difference(FLASH_LINEAR_ATTENTION_KERNELS)
+        if unsupported:
+            raise ValueError(f"Unsupported Flash Linear Attention kernels: {sorted(unsupported)}")
+        if isinstance(chunk_size, bool) or not isinstance(chunk_size, int) or chunk_size not in SUPPORTED_CHUNK_SIZES:
+            raise ValueError(f"chunk_size must be one of {SUPPORTED_CHUNK_SIZES}, got {chunk_size!r}.")
+
+        from fsdp_turbo.ops.registry import get_op
+        from fsdp_turbo.utils.patch import patch_model_members
+
+        patched = 0
+        named_modules = tuple(model.named_modules())
+        for op_name in selected:
+            module_attribute = FLA_MODULE_ATTRIBUTES[op_name]
+            op = get_op(op_name)
+            configured_op = partial(op, chunk_size=chunk_size) if op_name == CHUNK_GATED_DELTA_RULE else op
+            targets = {
+                f"{type(module).__module__}.{type(module).__name__}.{module_attribute}"
+                for _, module in named_modules
+                if callable(getattr(module, module_attribute, None))
+            }
+            matched = patch_model_members(model, sorted(targets), configured_op) if targets else 0
+            if matched == 0:
+                raise RuntimeError(f"FLA operator `{op_name}` did not match any model module attributes.")
+            patched += matched
+
+        logger.info_rank0(f"Flash Linear Attention kernels updated {patched} module callables: {selected}.")
+        return model

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

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

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

@@ -220,7 +220,7 @@ class FSDP2Engine:
     def is_lora_module_wrap(self, model) -> bool:
         return any(isinstance(module, LoraLayer) for module in model.modules())
 
-    def prepare_model(self, model: HFModel) -> HFModel:
+    def prepare_model(self, model: HFModel, ignored_params: set[nn.Parameter] | None = None) -> HFModel:
         if self.fsdp_mesh is None:
             logger.warning("No FSDP Mesh available, skipping FSDP wrapping.")
             return model
@@ -236,6 +236,11 @@ class FSDP2Engine:
             names = ", ".join(cls.__name__ for cls in transformer_layer_cls_to_wrap)
             logger.info(f"Applying per-layer FSDP to: {names}")
 
+        def _ignored_params_for(module: nn.Module) -> set[nn.Parameter] | None:
+            if not ignored_params:
+                return None
+            return ignored_params.intersection(module.parameters()) or None
+
         if self.is_lora_module_wrap(model):
             lora_modules = []
             for module in model.modules():
@@ -251,6 +256,7 @@ class FSDP2Engine:
                     reshard_after_forward=self.reshard_after_forward,
                     mp_policy=mp_policy,
                     offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None,
+                    ignored_params=_ignored_params_for(module),
                 )
 
             logger.info("Applying FSDP wrap for LoRA layer separately.")
@@ -271,6 +277,7 @@ class FSDP2Engine:
                     reshard_after_forward=self.reshard_after_forward,
                     mp_policy=mp_policy,
                     offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None,
+                    ignored_params=_ignored_params_for(module),
                 )
 
         # BaseTrainer is the single source of truth for gradient checkpointing.
@@ -299,6 +306,7 @@ class FSDP2Engine:
             reshard_after_forward=self.reshard_after_forward,
             mp_policy=mp_policy,
             offload_policy=CPUOffloadPolicy(pin_memory=self.pin_memory) if self.offload_params else None,
+            ignored_params=_ignored_params_for(model),
         )
 
         return model

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

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

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

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

tests_v1/config/test_args_parser.py

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

tests_v1/plugins/model_plugins/test_kernel_plugin.py

@@ -13,12 +13,32 @@
 # limitations under the License.
 
 import sys
+from functools import partial
 from unittest.mock import MagicMock, patch
 
+import pytest
 import torch.multiprocessing as mp
+from torch import nn
 from transformers import AutoModelForCausalLM
 
 
+def _original_fla_op(*args, **kwargs):
+    return args, kwargs
+
+
+class _LinearAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.chunk_gated_delta_rule = _original_fla_op
+        self.recurrent_gated_delta_rule = _original_fla_op
+
+
+class _FLAModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.linear_attn = _LinearAttention()
+
+
 def _apply_kernel(rank) -> None:
     with patch("torch.accelerator.current_accelerator") as mock_get_accelerator:
         mock_device = MagicMock()
@@ -73,3 +93,62 @@ def test_apply_kernel():
 
 def test_apply_all_kernels():
     mp.spawn(_apply_all_kernels)
+
+
+@pytest.mark.runs_on(["npu"])
+def test_flash_linear_attention_kernels_compose_with_auto(monkeypatch):
+    import fsdp_turbo.ops.fla  # noqa: F401
+    from fsdp_turbo.ops import get_op
+
+    from llamafactory.v1.plugins.model_plugins.kernels import interface
+    from llamafactory.v1.plugins.model_plugins.kernels.ops.linear_attention.fla import (
+        FlashLinearAttentionKernel,
+    )
+
+    model = _FLAModel()
+    auto_calls = []
+
+    monkeypatch.setattr(
+        interface,
+        "_apply_auto_kernels",
+        lambda model, **kwargs: auto_calls.append((model, kwargs)) or model,
+    )
+    # FLA execution is outside this bridge test; its external runtime is not required.
+    monkeypatch.setattr(FlashLinearAttentionKernel, "check_deps", staticmethod(lambda: None))
+
+    config = {
+        "name": "auto, flash-linear-attention",
+        "include_kernels": "fused_recurrent_gated_delta_rule, chunk_gated_delta_rule",
+        "chunk_size": 32,
+    }
+    assert interface.apply_kernels(model, config) is model
+    assert auto_calls == [(model, {"config": config, "require_logits": False})]
+    assert get_op("chunk_gated_delta_rule").__module__ == "fsdp_turbo.ops.fla"
+
+    chunk_op = model.linear_attn.chunk_gated_delta_rule
+    assert isinstance(chunk_op, partial)
+    assert chunk_op.func.__module__ == "fsdp_turbo.ops.fla"
+    assert chunk_op.keywords == {"chunk_size": 32}
+    assert model.linear_attn.recurrent_gated_delta_rule.__module__ == "fsdp_turbo.ops.fla"
+
+    with pytest.raises(RuntimeError, match="did not match any model module attributes"):
+        FlashLinearAttentionKernel.apply(
+            model=nn.Linear(2, 2),
+            config={"include_kernels": "chunk_gated_delta_rule", "chunk_size": 32},
+        )
+
+
+def test_flash_linear_attention_kernel_validates_config(monkeypatch):
+    from llamafactory.v1.plugins.model_plugins.kernels.ops.linear_attention.fla import (
+        FlashLinearAttentionKernel,
+    )
+
+    model = nn.Sequential(nn.Linear(2, 2))
+    monkeypatch.setattr(FlashLinearAttentionKernel, "check_device", staticmethod(lambda: None))
+    monkeypatch.setattr(FlashLinearAttentionKernel, "check_deps", staticmethod(lambda: None))
+
+    with pytest.raises(ValueError, match="chunk_size"):
+        FlashLinearAttentionKernel.apply(model=model, config={"include_kernels": "auto", "chunk_size": 48})
+
+    with pytest.raises(ValueError, match="Unsupported Flash Linear Attention kernels"):
+        FlashLinearAttentionKernel.apply(model=model, config={"include_kernels": "not_a_kernel"})

tests_v1/plugins/model_plugins/test_ulysses_cp.py

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

tests_v1/plugins/trainer_plugins/distributed/test_fsdpturbo_ep.py

@@ -0,0 +1,134 @@
+# Copyright 2025 the LlamaFactory team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from llamafactory.v1.plugins.trainer_plugins.distributed import fsdpturbo as fsdpturbo_module
+from llamafactory.v1.plugins.trainer_plugins.distributed.fsdpturbo import (
+    FSDPTurboEPModelSpec,
+    FSDPTurboFSDP2Engine,
+    FSDPTurboParallelState,
+)
+from llamafactory.v1.plugins.trainer_plugins.distributed.interface import (
+    DistributedPlugin,
+    FSDPTurboParams,
+)
+
+
+class _Model(torch.nn.Module):
+    def __init__(self, model_type: str):
+        super().__init__()
+        self.config = SimpleNamespace(model_type=model_type)
+
+
+def test_qwen35_ep_model_spec():
+    spec = FSDPTurboEPModelSpec.get(_Model("qwen3_5_moe"))
+
+    assert spec is not None
+    assert spec.ep_modules == ["model.language_model.layers.{*}.mlp.experts"]
+    assert spec.ep_fsdp_modules == ["model.language_model.layers.{*}.mlp"]
+
+
+def test_fsdpturbo_uses_class_plugin_and_strict_backend_params():
+    plugin = DistributedPlugin("fsdpturbo")
+    params = plugin.parse_params({"name": "fsdpturbo", "ep_size": 4}, FSDPTurboParams)
+
+    assert params.ep_size == 4
+    assert callable(plugin.shard_model)
+    assert callable(plugin.clip_grad_norm)
+    with pytest.raises(ValueError, match="Unknown params"):
+        plugin.parse_params({"name": "fsdpturbo", "cp_size": 2}, FSDPTurboParams)
+    for key in ("ep_modules", "ep_fsdp_modules"):
+        with pytest.raises(ValueError, match="Unknown params"):
+            plugin.parse_params({"name": "fsdpturbo", key: ["model.layers.*.mlp"]}, FSDPTurboParams)
+
+
+def test_fsdpturbo_sets_storage_dtype_inside_backend(monkeypatch):
+    from llamafactory.v1.plugins.trainer_plugins.distributed.fsdp2 import FSDP2Engine
+
+    monkeypatch.setattr(FSDP2Engine, "shard_model", lambda self, model: model)
+    engine = object.__new__(FSDPTurboFSDP2Engine)
+    engine.mixed_precision = "bf16"
+    model = torch.nn.Linear(2, 2, dtype=torch.float32)
+
+    assert engine.shard_model(model).weight.dtype == torch.bfloat16
+
+
+def test_fsdpturbo_sets_public_efsdp_gradient_divide_factor(monkeypatch):
+    expert_parallel_module = pytest.importorskip("fsdp_turbo.distributed.expert_parallel.expert_parallel")
+    expert_fully_shard_module = pytest.importorskip(
+        "fsdp_turbo.distributed.expert_parallel.expert_fully_shard_parallel"
+    )
+    captured = {}
+    monkeypatch.setattr(expert_parallel_module, "expert_parallelize_modules", lambda model, mesh, plan: model)
+
+    def _expert_fully_shard_modules(model, mesh, ep_plan, fsdp_plan):
+        captured["gradient_divide_factor"] = ep_plan.gradient_divide_factor
+        return model
+
+    monkeypatch.setattr(expert_fully_shard_module, "expert_fully_shard_modules", _expert_fully_shard_modules)
+
+    engine = object.__new__(FSDPTurboFSDP2Engine)
+    engine.dist_config = {"ep_dispatcher": "eager"}
+    engine.ep_size = 4
+    engine.ep_fsdp_size = 2
+    engine.parallel_state = SimpleNamespace(efsdp_size=2, ep_mesh=object(), efsdp_mesh=object())
+    engine.rank = 0
+
+    engine.prepare_model_ep(_Model("qwen3_5_moe"))
+
+    assert captured["gradient_divide_factor"] == 8.0
+
+
+def test_fsdpturbo_owns_expert_mesh_topology(monkeypatch):
+    calls = []
+
+    class _Mesh:
+        def __init__(self, name="expert"):
+            self.name = name
+
+        def __getitem__(self, name):
+            return _Mesh(name)
+
+    def _init_device_mesh(**kwargs):
+        calls.append(kwargs)
+        return _Mesh()
+
+    class _DistributedInterface:
+        current_device = torch.device("cpu")
+        strategy = SimpleNamespace(cp_size=1)
+
+        def get_world_size(self, dim):
+            return 16
+
+        def get_device_mesh(self, dim):
+            return _Mesh("dp")
+
+    monkeypatch.setattr(fsdpturbo_module, "init_device_mesh", _init_device_mesh)
+    state = FSDPTurboParallelState()
+    state.initialize(_DistributedInterface(), {"ep_size": 8})
+
+    assert calls == [
+        {
+            "device_type": "cpu",
+            "mesh_shape": (1, 2, 8, 1),
+            "mesh_dim_names": ("edp", "efsdp", "ep", "expert_cp"),
+        }
+    ]
+    assert state.ep_mesh.name == "ep"
+    assert state.efsdp_mesh.name == "efsdp"
+    assert state.expert_cp_mesh.name == "expert_cp"