Feat: USP sequence parallelism (Ulysses x Ring) via SequenceParallelPlan; Wan enabled

miles/backends/fsdp_utils/actor.py

@@ -31,6 +31,7 @@
 from .diffusion_update_weight_utils import DiffusionUpdateWeightFromTensor, DiffusionUpdateWeightFromTensorLoRA
 from .lr_scheduler import get_lr_scheduler
 from .parallel import create_fsdp_parallel_state
+from .sp_plan import apply_sequence_parallel
 
 logger = logging.getLogger(__name__)
 
@@ -131,12 +132,18 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
 
             model = apply_fsdp2(
                 model,
-                mesh=self.parallel_state.dp_mesh,
+                mesh=self.parallel_state.fsdp_mesh,
                 cpu_offload=self.args.fsdp_cpu_offload,
                 args=self.args,
                 no_split_modules=self.model_backend.fsdp_no_split_modules(model),
             )
             self.models[component] = model
+
+        if self.parallel_state.sp_size > 1:
+            for model in self.models.values():
+                plan = self.model_backend.sequence_parallel_plan(model)
+                apply_sequence_parallel(model, self.parallel_state, plan)
+
         # Force a sync to ensure sharding is complete and old memory is freed.
         torch.cuda.synchronize()
         clear_memory()

miles/backends/fsdp_utils/arguments.py

@@ -57,8 +57,9 @@ class FSDPArgs:
     # support matrix. Name kept identical to Megatron's.
     deterministic_mode: bool = False
 
-    # Context Parallelism
-    context_parallel_size: int = 1  # Context Parallelism size
+    # Sequence Parallelism (USP = Ulysses x Ring)
+    sequence_parallel_size: int = 1
+    ulysses_degree: int = 0  # 0=auto: ulysses fills sp; ring = sp // ulysses
 
     # YAML bookkeeping
     config: str | None = None

miles/backends/fsdp_utils/configs/train_pipeline_config.py

@@ -179,3 +179,14 @@ def cfg_combine(
     @abc.abstractmethod
     def preprocess_model_before_fsdp(self, model: torch.nn.Module) -> None:
         """Preprocess the model before FSDP."""
+
+    def apply_sp_attention(self, model: torch.nn.Module, parallel_state) -> None:
+        """Route this family's self-attention through ``sp_ops.usp_attention``.
+
+        Default = intercept at diffusers' attention dispatch (covers every model
+        whose processors call ``dispatch_attention_fn``); families with a
+        non-dispatch attention path override with their own installer.
+        """
+        from ..sp_attention import apply_dispatch_sp_attention
+
+        apply_dispatch_sp_attention(model, parallel_state)

miles/backends/fsdp_utils/model_backend.py

@@ -1,12 +1,13 @@
 """Model backend: owns model-side behavior for the FSDP trainer.
 
 Selected via ``--model-backend-path`` (miles custom-function style); the
-family config declares the default. Three concerns, all properties of the
+family config declares the default. Four concerns, all properties of the
 concrete modeling rather than of the training loop:
 
   - ``load_models_and_scheduler``: checkpoint -> ``({component: model}, scheduler)``
   - ``enable_gradient_checkpointing``: how this model turns on grad ckpt
   - ``fsdp_no_split_modules``: which block classes FSDP wraps
+  - ``sequence_parallel_plan``: the model's SP boundaries/attention declaration
 
 Defaults implement the diffusers protocol (see ``models/__init__.py``); a
 native model overrides methods here instead of retrofitting its instances.
@@ -20,6 +21,8 @@
 import torch
 from diffusers import DiffusionPipeline
 
+from .sp_plan import SequenceParallelPlan
+
 
 class ModelBackend(abc.ABC):
     def __init__(self, train_pipeline_config):
@@ -46,6 +49,10 @@ def set_attention_backend(self, model: torch.nn.Module, backend: str) -> None:
         """Select the DiT attention backend; default = the diffusers protocol method."""
         model.set_attention_backend(backend)
 
+    def sequence_parallel_plan(self, model: torch.nn.Module) -> SequenceParallelPlan:
+        """The model's SequenceParallelPlan (boundaries + attention installer)."""
+        raise NotImplementedError(f"{type(self).__name__} does not support sequence parallelism")
+
 
 class DiffusersModelBackend(ModelBackend):
     """Load trainable components from a diffusers pipeline checkpoint."""
@@ -75,3 +82,41 @@ def load_models_and_scheduler(
         scheduler = pipeline.scheduler
         del pipeline
         return raw_models, scheduler
+
+    def sequence_parallel_plan(self, model: torch.nn.Module) -> SequenceParallelPlan:
+        base = model.get_base_model() if hasattr(model, "get_base_model") else model
+        boundaries = getattr(base, "_cp_plan", None)
+        if not boundaries:
+            raise ValueError(f"{base.__class__.__name__} declares no _cp_plan; sequence parallelism unavailable")
+        wildcards = [k for k in boundaries if "*" in k]
+        if wildcards:
+            raise ValueError(
+                f"{base.__class__.__name__}._cp_plan uses wildcard boundaries {wildcards}, "
+                "which the boundary-hook installer does not support yet"
+            )
+        return SequenceParallelPlan(
+            boundaries=boundaries,
+            attention=self.config.apply_sp_attention,
+            num_attention_heads=base.config.num_attention_heads,
+        )
+
+
+# to validate and fail earlier if there's invalid condition
+def validate_sp_support(args, cfg_cls) -> None:
+    """Driver-side, before any actor launches: reject launches whose SP config
+    cannot work, without loading weights.
+
+    Models that lack a _cp_plan or don't route attention through diffusers'
+    dispatch fail later (plan construction / attention install) with a clear
+    message — checking those here would require loading the model class.
+    """
+    from miles.utils.misc import load_function
+
+    if args.fsdp_attention_backend is not None:
+        raise ValueError(
+            "--fsdp-attention-backend has no effect with sequence parallelism: "
+            "USP owns the self-attention kernel choice"
+        )
+    backend_cls = load_function(args.model_backend_path)
+    if backend_cls.sequence_parallel_plan is ModelBackend.sequence_parallel_plan:
+        raise ValueError(f"{backend_cls.__name__} does not support sequence parallelism")

miles/backends/fsdp_utils/parallel.py

@@ -7,55 +7,78 @@
 from miles.utils.distributed_utils import get_gloo_group
 
 from ..training_utils.parallel import ParallelState
+from .sp_mesh import locate_rank, sp_subgroups, validate_sp_config
 
 logger = logging.getLogger(__name__)
 
 
 def create_fsdp_parallel_state(args: Namespace) -> ParallelState:
-    """Create a ParallelState instance for FSDP configuration."""
+    """ParallelState for FSDP + optional sequence parallelism.
+
+    SP gets its own process groups. FSDP shards parameters over every mesh
+    axis (dp x sp flattened) — should a replicate axis (HSDP) ever be added,
+    flatten only the shard axes, not the whole world. Data dispatch is by
+    dp_rank; sp peers share samples.
+    """
     world_size = dist.get_world_size()
     rank = dist.get_rank()
 
-    cp_size = args.context_parallel_size
-    dp_rank = rank // cp_size
-    cp_rank = rank % cp_size
-
-    mesh = init_device_mesh("cuda", mesh_shape=(world_size // cp_size, cp_size), mesh_dim_names=("dp", "cp"))
+    sp_size, ulysses_degree, ring_degree = validate_sp_config(
+        world_size, args.sequence_parallel_size, args.ulysses_degree
+    )
+    dp_rank, sp_rank, _, _ = locate_rank(rank, sp_size, ulysses_degree)
+    dp_size = world_size // sp_size
 
+    mesh = init_device_mesh("cuda", mesh_shape=(dp_size, sp_size), mesh_dim_names=("dp", "sp"))
+    dp_group = mesh.get_group("dp")
+    sp_group = mesh.get_group("sp")
     logger.info(
-        f"[Rank {rank}] Device mesh (2D): world_size={world_size}, "
-        f"cp_size={cp_size}, dp_size={world_size // cp_size}"
+        f"[Rank {rank}] mesh dp={dp_size} sp={sp_size} (ulysses={ulysses_degree} ring={ring_degree}), "
+        f"dp_rank={dp_rank} sp_rank={sp_rank}"
     )
-    logger.info(f"[Rank {rank}] Mesh shape: {mesh.shape}, " f"dp_rank={dp_rank}, cp_rank={cp_rank}")
-
-    # Setup Ring Flash Attention with CP group from mesh (only when cp_size > 1).
-    # Lazy import to avoid a hard dependency on ring_flash_attn at module load;
-    # the package is incompatible with transformers>=5.4 for pure-DP runs.
-    if cp_size > 1:
-        from ring_flash_attn import substitute_hf_flash_attn
 
-        substitute_hf_flash_attn(mesh.get_group("cp"), heads_k_stride=1)
-        logger.info(f"[Rank {rank}] CP initialized via device mesh")
-    else:
-        logger.info(f"[Rank {rank}] Pure DP mode (cp_size=1)")
+    # dist.new_group is collective: every rank must create every group.
+    # Degree-1 dimensions stay None (usp_attention treats None as local).
+    ulysses_group = ring_group = None
+    if sp_size > 1:
+        _, _, _, ulysses_groups, ring_groups = sp_subgroups(world_size, sp_size, ulysses_degree)
+        if ulysses_degree > 1:
+            for ranks in ulysses_groups:
+                group = dist.new_group(ranks)
+                if rank in ranks:
+                    ulysses_group = group
+        if ring_degree > 1:
+            for ranks in ring_groups:
+                group = dist.new_group(ranks)
+                if rank in ranks:
+                    ring_group = group
 
     parallel_state = ParallelState(
         dp_rank=dp_rank,
         dp_src_rank=dp_rank // world_size,
-        dp_size=world_size // cp_size,
-        cp_rank=cp_rank,
-        cp_size=cp_size,
+        dp_size=dp_size,
+        cp_rank=sp_rank,
+        cp_size=sp_size,
         dp_cp_rank=rank,
         dp_cp_size=world_size,
-        dp_group=mesh.get_group("dp"),
+        dp_group=dp_group,
         dp_cp_group=dist.group.WORLD,
         dp_cp_group_gloo=get_gloo_group(),
-        cp_group=mesh.get_group("cp"),
+        cp_group=sp_group,
         tp_size=1,
         tp_rank=0,
-        tp_group=dist.new_group([rank]),
+        tp_group=None,
+        sp_rank=sp_rank,
+        sp_size=sp_size,
+        sp_group=sp_group,
+        ulysses_degree=ulysses_degree,
+        ring_degree=ring_degree,
+        ulysses_group=ulysses_group,
+        ring_group=ring_group,
     )
-
     parallel_state.dp_mesh = mesh["dp"]
-
+    if sp_size > 1:
+        parallel_state.fsdp_mesh = mesh[("dp", "sp")]._flatten("dp_sp")
+    else:
+        parallel_state.fsdp_mesh = mesh["dp"]
     return parallel_state

miles/backends/fsdp_utils/sp_attention.py

@@ -0,0 +1,73 @@
+"""Dispatch-level USP attention: the default installer for diffusers models.
+
+Wraps the modeling module's ``dispatch_attention_fn`` so self-attention call
+sites (which pass ``_parallel_config`` per upstream convention) route through
+``sp_ops.usp_attention``; cross-attention and the model's own processors run
+untouched.
+"""
+
+import functools
+import sys
+
+from .sp_ops import usp_attention
+
+
+class _USPDispatchConfig:
+    """Marker consumed by the wrapped dispatch_attention_fn; models pass it
+    through ``_parallel_config`` for self-attention call sites only."""
+
+    def __init__(self, parallel_state):
+        self.ulysses_group = parallel_state.ulysses_group
+        self.ring_group = parallel_state.ring_group
+
+
+def _wrap_dispatch(module):
+    original = module.dispatch_attention_fn
+    if getattr(original, "_miles_usp_wrapped", False):
+        return
+
+    @functools.wraps(original)
+    def dispatch(query, key, value, *args, parallel_config=None, **kwargs):
+        if not isinstance(parallel_config, _USPDispatchConfig):
+            return original(query, key, value, *args, parallel_config=parallel_config, **kwargs)
+        attn_mask = args[0] if args else kwargs.get("attn_mask")
+        if attn_mask is not None:
+            raise ValueError("USP self-attention does not support attention masks")
+        if (len(args) > 1 and args[1]) or kwargs.get("dropout_p"):
+            raise ValueError("USP self-attention requires dropout_p=0 (per-rank RNG streams diverge)")
+        if (len(args) > 2 and args[2]) or kwargs.get("is_causal"):
+            raise ValueError("USP self-attention does not support is_causal yet")
+        if (len(args) > 3 and args[3] is not None) or kwargs.get("scale") is not None:
+            raise ValueError("USP self-attention uses the default 1/sqrt(d) scale")
+        return usp_attention(query, key, value, parallel_config.ulysses_group, parallel_config.ring_group)
+
+    dispatch._miles_usp_wrapped = True
+    module.dispatch_attention_fn = dispatch
+
+
+def apply_dispatch_sp_attention(transformer, parallel_state):
+    """Default SP attention: intercept the model module's dispatch_attention_fn
+    so self-attention call sites (which pass ``_parallel_config`` per upstream
+    convention) route through usp_attention; the model's own processors and
+    cross-attention stay untouched."""
+    base = transformer.get_base_model() if hasattr(transformer, "get_base_model") else transformer
+    # fully_shard swizzles the class (FSDP<Name>, defined in torch's fsdp
+    # module); the modeling module that imported dispatch_attention_fn is
+    # found through the MRO.
+    module = next(
+        (
+            mod
+            for cls in type(base).__mro__
+            if (mod := sys.modules.get(cls.__module__)) is not None and hasattr(mod, "dispatch_attention_fn")
+        ),
+        None,
+    )
+    if module is None:
+        raise ValueError(
+            f"{type(base).__name__} does not route attention through diffusers' "
+            "dispatch_attention_fn; the family must override apply_sp_attention"
+        )
+    _wrap_dispatch(module)
+    config = _USPDispatchConfig(parallel_state)
+    for processor in base.attn_processors.values():
+        processor._parallel_config = config

miles/backends/fsdp_utils/sp_mesh.py

@@ -0,0 +1,50 @@
+"""Sequence-parallel rank layout: sp = ulysses * ring, global rank = dp_rank * sp + sp_rank.
+
+Pure functions, no distributed init required. Group layout matches sglang's USP
+(Ulysses ranks contiguous within an SP group, Ring ranks strided).
+"""
+
+
+def resolve_sp_degrees(sequence_parallel_size, ulysses_degree=0):
+    """Normalize to (sp, ulysses, ring); ring = sp // ulysses. 0 means auto: ulysses fills sp."""
+    sp = max(1, sequence_parallel_size)
+    u = ulysses_degree or sp
+    if sp % u:
+        raise ValueError(f"sequence_parallel_size({sp}) is not divisible by ulysses_degree({u})")
+    return sp, u, sp // u
+
+
+def validate_sp_config(world_size, sequence_parallel_size, ulysses_degree=0):
+    """Validate at startup. Returns (sp, ulysses, ring).
+
+    The num_heads % ulysses check lives in apply_sequence_parallel, where the
+    real model config is available.
+    """
+    sp, u, r = resolve_sp_degrees(sequence_parallel_size, ulysses_degree)
+    if world_size % sp != 0:
+        raise ValueError(f"world_size({world_size}) is not divisible by sequence_parallel_size({sp})")
+    return sp, u, r
+
+
+def sp_subgroups(world_size, sequence_parallel_size, ulysses_degree=0):
+    """Return (dp_size, sp_size, sp_groups, ulysses_groups, ring_groups); groups are global-rank lists.
+
+    Inverse of locate_rank: coordinates -> full member list per group.
+    """
+    sp, u, r = resolve_sp_degrees(sequence_parallel_size, ulysses_degree)
+    dp_size = world_size // sp
+    sp_groups = [list(range(d * sp, (d + 1) * sp)) for d in range(dp_size)]
+    ulysses_groups = [g[i * u : (i + 1) * u] for g in sp_groups for i in range(r)]
+    ring_groups = [g[i::u] for g in sp_groups for i in range(u)]
+    return dp_size, sp, sp_groups, ulysses_groups, ring_groups
+
+
+def locate_rank(rank, sequence_parallel_size, ulysses_degree=0):
+    """Locate a global rank's (dp_rank, sp_rank, ulysses_rank, ring_rank).
+
+    Inverse of sp_subgroups: global rank -> its index on each axis.
+    """
+    sp, u, _ = resolve_sp_degrees(sequence_parallel_size, ulysses_degree)
+    dp_rank, sp_rank = divmod(rank, sp)
+    ring_rank, ulysses_rank = divmod(sp_rank, u)
+    return dp_rank, sp_rank, ulysses_rank, ring_rank