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

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"