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

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