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

miles/backends/fsdp_utils/actor.py

@@ -26,7 +26,7 @@
     build_microbatch_schedule,
     scheduler_meta_from_rollout,
     stack_train_pair_rollout_debug,
-    validate_same_microbatch_counts_across_dp,
+    validate_same_microbatch_counts_across_train_ranks,
 )
 from . import checkpoint
 from .diffusion_update_weight_utils import (
@@ -36,6 +36,7 @@
 )
 from .lr_scheduler import get_lr_scheduler
 from .parallel import create_fsdp_parallel_state
+from .sequence_parallel.plan import apply_sequence_parallel
 
 logger = logging.getLogger(__name__)
 
@@ -121,7 +122,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
             full_state = model.state_dict() if rank == 0 else {}
             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),
@@ -130,6 +131,12 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int:  # ty
             del full_state
             self.train_pipeline_config.postprocess_model_after_materialize(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()
@@ -268,14 +275,14 @@ def _gather_and_log_metrics(self, rollout_id: int, log_dict: dict[str, float], s
                 log_dict["train/lr"] = float(self.optimizer.param_groups[0]["lr"])
             except Exception:
                 pass
-        if self.parallel_state.dp_cp_rank == 0:
-            dp_size = self.parallel_state.dp_cp_size
+        if self.parallel_state.dp_sp_rank == 0:
+            dp_size = self.parallel_state.dp_sp_size
             gathered = [None] * dp_size
             dist.gather_object(
                 log_dict,
                 gathered,
                 dst=self.parallel_state.dp_src_rank,
-                group=self.parallel_state.dp_cp_group_gloo,
+                group=self.parallel_state.dp_sp_group_gloo,
             )
             reduced = {k: sum(d[k] for d in gathered) / dp_size for k in log_dict}
             reduced["train/epoch"] = float(rollout_id)
@@ -296,7 +303,7 @@ def _gather_and_log_metrics(self, rollout_id: int, log_dict: dict[str, float], s
                 log_dict,
                 None,
                 dst=self.parallel_state.dp_src_rank,
-                group=self.parallel_state.dp_cp_group_gloo,
+                group=self.parallel_state.dp_sp_group_gloo,
             )
 
     def train(self, rollout_id: int, rollout_data_ref) -> None:  # type: ignore[override]
@@ -376,7 +383,7 @@ def _train_core(self, rollout_id: int, rollout_data) -> None:
             num_optim_steps_per_rollout=num_optim_steps_per_rollout,
             micro_batch_size=micro_bs,
         )
-        validate_same_microbatch_counts_across_dp(
+        validate_same_microbatch_counts_across_train_ranks(
             microbatch_schedule=microbatch_schedule,
             parallel_state=self.parallel_state,
         )

miles/backends/fsdp_utils/arguments.py

@@ -4,6 +4,8 @@
 
 import yaml
 
+from .sequence_parallel.topology import validate_sp_config
+
 
 @dataclass
 class FSDPArgs:
@@ -33,7 +35,8 @@ class FSDPArgs:
     attn_implementation: str = "flash_attention_2"
 
     # DiT attention backend, passed to diffusers set_attention_backend (e.g.
-    # "flash", "sage", "native"). None keeps the diffusers default.
+    # "flash", "sage", "native"). None keeps the diffusers default. Under SP,
+    # explicit backends are supported for pure Ulysses; Ring owns its kernel.
     fsdp_attention_backend: str | None = None
 
     # Logging
@@ -57,8 +60,12 @@ 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
+    # 0=auto: ulysses fills sp; ring = sp // ulysses. Ring degrees > 1 run on
+    # torch's experimental (private) ring-attention implementation and require
+    # torch >= 2.11 (the CI image's pin).
+    ulysses_degree: int = 0
 
     # YAML bookkeeping
     config: str | None = None
@@ -153,6 +160,31 @@ def validate_attention_args(args):
     )
 
 
+def validate_sp_args(args) -> None:
+    """Validate the finalized train topology and SP/backend combination on the driver.
+
+    Model-instance constraints such as ``_cp_plan`` availability and attention
+    dispatch compatibility are checked later, after the model is loaded.
+    """
+    from miles.utils.misc import load_function
+
+    sp_size, _, ring_degree = validate_sp_config(
+        args.actor_num_gpus_per_node * args.actor_num_nodes,
+        args.sequence_parallel_size,
+        args.ulysses_degree,
+    )
+    if sp_size == 1:
+        return
+    if args.fsdp_attention_backend is not None and ring_degree > 1:
+        raise ValueError(
+            "--fsdp-attention-backend is supported with pure Ulysses only: "
+            f"the configured SP topology has ring_degree={ring_degree}, whose ring attention owns the kernel choice"
+        )
+    backend_cls = load_function(args.model_backend_path)
+    if not backend_cls.supports_sequence_parallelism():
+        raise ValueError(f"{backend_cls.__name__} does not support sequence parallelism")
+
+
 def load_fsdp_args(extra_args_provider=None):
     args = parse_fsdp_cli(extra_args_provider)
     if args.config:

miles/backends/fsdp_utils/model_backend.py

@@ -1,15 +1,17 @@
 """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_component`` / ``load_scheduler``: checkpoint -> model components and 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.
+Defaults adapt the diffusers protocol (see ``models/__init__.py``); a native
+backend overrides the model-side seams and provides one plan for each model it
+supports under sequence parallelism.
 """
 
 from __future__ import annotations
@@ -24,6 +26,9 @@
 import torch.distributed as dist
 from diffusers import DiffusionPipeline
 
+from .sequence_parallel.diffusers_dispatch import apply_dispatch_sp_attention
+from .sequence_parallel.plan import MILES_SP_PLAN_ATTR, SequenceParallelPlan
+
 logger = logging.getLogger(__name__)
 
 
@@ -68,6 +73,15 @@ def fsdp_no_split_modules(self, model: torch.nn.Module) -> list[str]:
     def set_attention_backend(self, model: torch.nn.Module, backend: str) -> None:
         raise NotImplementedError
 
+    @classmethod
+    def supports_sequence_parallelism(cls) -> bool:
+        """Whether the backend declares a model-specific SP plan resolver."""
+        return cls.sequence_parallel_plan is not ModelBackend.sequence_parallel_plan
+
+    def sequence_parallel_plan(self, model: torch.nn.Module) -> SequenceParallelPlan:
+        """Return 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."""
@@ -157,6 +171,28 @@ def _component_class(spec):
                 return None
         return getattr(module, class_name, None)
 
+    def sequence_parallel_plan(self, model: torch.nn.Module) -> SequenceParallelPlan:
+        base = model.get_base_model() if hasattr(model, "get_base_model") else model
+        plan = getattr(base, MILES_SP_PLAN_ATTR, None)
+        if plan is not None:
+            if not isinstance(plan, SequenceParallelPlan):
+                raise TypeError(
+                    f"{base.__class__.__name__}.{MILES_SP_PLAN_ATTR} must be a SequenceParallelPlan, "
+                    f"got {type(plan).__name__}"
+                )
+            return plan
+
+        boundaries = getattr(base, "_cp_plan", None)
+        if not boundaries:
+            raise ValueError(f"{base.__class__.__name__} declares no _cp_plan; sequence parallelism unavailable")
+        plan = SequenceParallelPlan(
+            boundaries=boundaries,
+            attention_installer=apply_dispatch_sp_attention,
+            num_attention_heads=base.config.num_attention_heads,
+        )
+        setattr(base, MILES_SP_PLAN_ATTR, plan)
+        return plan
+
 
 class LTXModelBackend(ModelBackend):
     """Native LTX-2 loading via ltx_core; model instances stay unmodified."""

miles/backends/fsdp_utils/parallel.py

@@ -7,55 +7,70 @@
 from miles.utils.distributed_utils import get_gloo_group
 
 from ..training_utils.parallel import ParallelState
+from .sequence_parallel.topology 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_cp_rank=rank,
-        dp_cp_size=world_size,
-        dp_group=mesh.get_group("dp"),
-        dp_cp_group=dist.group.WORLD,
-        dp_cp_group_gloo=get_gloo_group(),
-        cp_group=mesh.get_group("cp"),
+        dp_src_rank=0,
+        dp_size=dp_size,
+        dp_group=dp_group,
+        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,
+        dp_sp_rank=rank,
+        dp_sp_size=world_size,
+        dp_sp_group_gloo=get_gloo_group(),
         tp_size=1,
         tp_rank=0,
-        tp_group=dist.new_group([rank]),
+        tp_group=None,
     )
-
-    parallel_state.dp_mesh = mesh["dp"]
-
+    parallel_state.fsdp_mesh = mesh[("dp", "sp")]._flatten("dp_sp") if sp_size > 1 else mesh["dp"]
     return parallel_state

miles/backends/fsdp_utils/sequence_parallel/__init__.py

@@ -0,0 +1 @@
+"""Sequence-parallel topology, operators, plans, and model integrations."""