diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index e3c3ba46..d7aaab02 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/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, ) diff --git a/miles/backends/fsdp_utils/arguments.py b/miles/backends/fsdp_utils/arguments.py index a84979a0..ca3b9fbd 100644 --- a/miles/backends/fsdp_utils/arguments.py +++ b/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: diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index 9f7495d4..62eb4e1b 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/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.""" diff --git a/miles/backends/fsdp_utils/parallel.py b/miles/backends/fsdp_utils/parallel.py index 5128ef35..8ac736ab 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/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 diff --git a/miles/backends/fsdp_utils/sequence_parallel/__init__.py b/miles/backends/fsdp_utils/sequence_parallel/__init__.py new file mode 100644 index 00000000..dabdf250 --- /dev/null +++ b/miles/backends/fsdp_utils/sequence_parallel/__init__.py @@ -0,0 +1 @@ +"""Sequence-parallel topology, operators, plans, and model integrations.""" diff --git a/miles/backends/fsdp_utils/sequence_parallel/attention.py b/miles/backends/fsdp_utils/sequence_parallel/attention.py new file mode 100644 index 00000000..cbbbd593 --- /dev/null +++ b/miles/backends/fsdp_utils/sequence_parallel/attention.py @@ -0,0 +1,198 @@ +"""Differentiable USP (Ulysses x Ring) attention operators, owned by the trainer. + +Layout convention matches sglang-diffusion's USP (heads sharded across the ulysses +group inside attention, sequence sharded outside), so training numerics stay +aligned with rollout; the collectives only move data. Local attention is torch +SDPA by default and may be injected by the model adapter; ring attention uses +torch's ring templates with the aten flash op. +""" + +import torch +import torch.distributed as dist + + +class _GatherSequence(torch.autograd.Function): + """All-gather local shards along dim; backward sums then returns each rank's slice. + + The backward all-reduces the incoming gradient over the sp group first: + downstream partial grads then carry an sp factor, so FSDP's + 1/(dp*sp) mean over a dp x sp shard mesh restores (1/dp) * sum_dp exactly. + """ + + @staticmethod + def forward(ctx, x, group, sp_rank, sp_size, dim): + ctx.group = group + ctx.sp_rank = sp_rank + ctx.dim = dim + ctx.local_size = x.shape[dim] + parts = [torch.empty_like(x) for _ in range(sp_size)] + dist.all_gather(parts, x.contiguous(), group=group) + return torch.cat(parts, dim=dim) + + @staticmethod + def backward(ctx, grad): + grad = grad.contiguous() + dist.all_reduce(grad, group=ctx.group) + start = ctx.sp_rank * ctx.local_size + return grad.narrow(ctx.dim, start, ctx.local_size), None, None, None, None + + +def shard_sequence(x, sp_rank, sp_size, dim=1): + s = x.shape[dim] + if s % sp_size: + raise ValueError(f"sequence length {s} is not divisible by sp_size {sp_size}") + s_local = s // sp_size + return x.narrow(dim, sp_rank * s_local, s_local) + + +def gather_sequence(x, group, sp_rank, sp_size, dim=1): + return _GatherSequence.apply(x, group, sp_rank, sp_size, dim) + + +class _AllToAllSingle(torch.autograd.Function): + """Even-split all-to-all; an involution, so the adjoint is the same collective.""" + + @staticmethod + def forward(ctx, x, group): + ctx.group = group + out = torch.empty_like(x) + dist.all_to_all_single(out, x, group=group) + return out + + @staticmethod + def backward(ctx, grad): + out = torch.empty_like(grad) + dist.all_to_all_single(out, grad.contiguous(), group=ctx.group) + return out, None + + +def _all_to_all_4d(x, group): + shape = x.shape + return _AllToAllSingle.apply(x.flatten(), group).reshape(shape) + + +def ulysses_input_all_to_all(x, group): + """[b, s_local, h, d] -> [b, s_local * world, h / world, d]: shard heads, gather sequence.""" + world_size = dist.get_world_size(group) + if world_size <= 1: + return x + b, s_local, h_global, d = x.shape + if h_global % world_size: + raise ValueError(f"num_heads({h_global}) is not divisible by ulysses world size({world_size})") + h_local, s_global = h_global // world_size, s_local * world_size + + x = x.permute(2, 0, 1, 3).contiguous() # [h_global, b, s_local, d] + x = _all_to_all_4d(x, group) + x = x.reshape(world_size, h_local, b, s_local, d) + return x.permute(2, 0, 3, 1, 4).contiguous().reshape(b, s_global, h_local, d) + + +def ulysses_output_all_to_all(x, group): + """[b, s_global, h_local, d] -> [b, s_global / world, h_local * world, d]: inverse of input.""" + world_size = dist.get_world_size(group) + if world_size <= 1: + return x + b, s_global, h_local, d = x.shape + if s_global % world_size: + raise ValueError(f"sequence({s_global}) is not divisible by ulysses world size({world_size})") + s_local, h_global = s_global // world_size, h_local * world_size + + x = x.permute(1, 0, 2, 3).contiguous() # [s_global, b, h_local, d] + x = _all_to_all_4d(x, group) + x = x.reshape(world_size, s_local, b, h_local, d) + return x.permute(2, 1, 0, 3, 4).contiguous().reshape(b, s_local, h_global, d) + + +class _RingFlashAttention(torch.autograd.Function): + """Ring attention via torch's ring templates (fwd + reverse-ring bwd), aten flash op. + + q/k/v: [B, H, S, D]. + """ + + @staticmethod + def forward(ctx, query, key, value, group, scale): + # torch's experimental (private) ring-attention templates; this home + # requires torch >= 2.11 (the CI image's pin) — before the 2.11 move + # they lived in torch.distributed.tensor.experimental._attention. + from torch.distributed.tensor.experimental._context_parallel._attention import _templated_ring_attention + + out, lse, cum_q, cum_k, max_q, max_k, philox_seed, philox_offset, _dbg = _templated_ring_attention( + group, + 2, + torch.ops.aten._scaled_dot_product_flash_attention, + query=query, + key=key, + value=value, + is_causal=False, + dropout_p=0.0, + scale=scale, + ) + out = out.to(query.dtype) + ctx.save_for_backward(query, key, value, out, lse, cum_q, cum_k, philox_seed, philox_offset) + ctx.group, ctx.scale, ctx.max_q, ctx.max_k = group, scale, max_q, max_k + return out + + @staticmethod + def backward(ctx, grad_out): + from torch.distributed.tensor.experimental._context_parallel._attention import ( + _templated_ring_attention_backward, + ) + + query, key, value, out, lse, cum_q, cum_k, philox_seed, philox_offset = ctx.saved_tensors + grad_q, grad_k, grad_v, *_ = _templated_ring_attention_backward( + ctx.group, + 2, + torch.ops.aten._scaled_dot_product_flash_attention_backward.default, + grad_out=grad_out.contiguous(), + grad_out_name="grad_out", + query=query, + key=key, + value=value, + out=out, + logsumexp=lse, + is_causal=False, + cum_seq_q=cum_q, + cum_seq_k=cum_k, + max_q=ctx.max_q, + max_k=ctx.max_k, + dropout_p=0.0, + philox_seed=philox_seed, + philox_offset=philox_offset, + scale=ctx.scale, + ) + return grad_q, grad_k, grad_v, None, None + + +def usp_attention(query, key, value, ulysses_group=None, ring_group=None, local_attention_fn=None): + """USP self-attention on [B, S_local, H, D] tensors; returns the same layout. + + Ulysses all-to-all temporarily gathers the sequence (sharding heads), ring + attention covers the remaining split. Without Ring, ``local_attention_fn`` may + route the gathered tensors through the model's configured attention backend; + the fallback is plain SDPA. + """ + if ulysses_group is not None: + query = ulysses_input_all_to_all(query, ulysses_group) + key = ulysses_input_all_to_all(key, ulysses_group) + value = ulysses_input_all_to_all(value, ulysses_group) + + if ring_group is not None: + scale = query.shape[-1] ** -0.5 + q = query.transpose(1, 2) # [B, H, S, D] + k = key.transpose(1, 2) + v = value.transpose(1, 2) + out = _RingFlashAttention.apply(q.contiguous(), k.contiguous(), v.contiguous(), ring_group, scale) + out = out.transpose(1, 2).contiguous() # [B, S, H, D] + elif local_attention_fn is not None: + out = local_attention_fn(query, key, value) + else: + scale = query.shape[-1] ** -0.5 + q = query.transpose(1, 2) + k = key.transpose(1, 2) + v = value.transpose(1, 2) + out = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False, scale=scale) + out = out.transpose(1, 2).contiguous() # [B, S, H, D] + + if ulysses_group is not None: + out = ulysses_output_all_to_all(out, ulysses_group) + return out diff --git a/miles/backends/fsdp_utils/sequence_parallel/diffusers_dispatch.py b/miles/backends/fsdp_utils/sequence_parallel/diffusers_dispatch.py new file mode 100644 index 00000000..7b53bd13 --- /dev/null +++ b/miles/backends/fsdp_utils/sequence_parallel/diffusers_dispatch.py @@ -0,0 +1,96 @@ +"""Diffusers dispatcher integration for USP self-attention. + +Wraps the modeling module's ``dispatch_attention_fn`` so self-attention call +sites (which pass ``_parallel_config`` per upstream convention) route through +``attention.usp_attention``; cross-attention and the model's own processors run +untouched. +""" + +import functools +import sys + +from .attention 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") + if parallel_config.ring_group is not None: + if kwargs.get("backend") is not None: + raise ValueError( + "An explicit attention backend is supported with pure Ulysses only, not Ring attention" + ) + return usp_attention(query, key, value, parallel_config.ulysses_group, parallel_config.ring_group) + + def local_attention_fn(local_query, local_key, local_value): + return original( + local_query, + local_key, + local_value, + *args, + parallel_config=None, + **kwargs, + ) + + return usp_attention( + query, + key, + value, + parallel_config.ulysses_group, + local_attention_fn=local_attention_fn, + ) + + 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, 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; its SequenceParallelPlan must provide a custom attention_installer" + ) + _wrap_dispatch(module) + config = _USPDispatchConfig(parallel_state) + for processor in base.attn_processors.values(): + processor._parallel_config = config diff --git a/miles/backends/fsdp_utils/sequence_parallel/plan.py b/miles/backends/fsdp_utils/sequence_parallel/plan.py new file mode 100644 index 00000000..aaef5eab --- /dev/null +++ b/miles/backends/fsdp_utils/sequence_parallel/plan.py @@ -0,0 +1,141 @@ +"""Sequence-parallel plan: the per-family declaration and its interpreter. + +A model family opts into SP through a ``SequenceParallelPlan`` (where the +sequence is sharded/gathered, how self-attention is installed, head count). +``apply_sequence_parallel`` consumes it: boundary hooks keep model outputs +full-sequence, so every parameter sees a partial gradient and loss/log_prob +code is untouched. +""" + +import inspect +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from diffusers.models._modeling_parallel import ContextParallelOutput + +from .attention import gather_sequence, shard_sequence + +MILES_SP_PLAN_ATTR = "_miles_sp_plan" + + +@dataclass(frozen=True) +class SequenceParallelPlan: + """What one transformer family declares to run under SP. + + ``boundaries``: fqn -> ``ContextParallelInput``/``ContextParallelOutput`` + (the diffusers ``_cp_plan`` vocabulary) — where the sequence dim is sharded + to S/sp and where full-sequence outputs are gathered back. + ``attention_installer``: called with (transformer, parallel_state); routes + the model's self-attention through ``attention.usp_attention``. + + Backends may attach one plan to a model instance as ``_miles_sp_plan``. + The plan is topology-independent; ranks and process groups remain in the + runtime parallel state passed to ``apply_sequence_parallel``. + """ + + boundaries: dict + attention_installer: Callable[[torch.nn.Module, object], None] + num_attention_heads: int + + def __post_init__(self) -> None: + wildcards = [path for path in self.boundaries if "*" in path] + if wildcards: + raise ValueError(f"SequenceParallelPlan does not support wildcard boundary paths: {wildcards}") + if self.num_attention_heads < 1: + raise ValueError(f"num_attention_heads must be positive, got {self.num_attention_heads}") + + +def _split_if_expected(x, spec, parallel_state): + if not isinstance(x, torch.Tensor): + return x + if spec.expected_dims is not None and x.ndim != spec.expected_dims: + return x + return shard_sequence(x, parallel_state.sp_rank, parallel_state.sp_size, dim=spec.split_dim) + + +def _resolve_submodule(root, path): + # getattr-based walk: unlike get_submodule, it also traverses attribute + # forwarding of wrappers such as peft's PeftModel. + module = root + for part in path.split("."): + module = module[int(part)] if part.isdigit() else getattr(module, part) + return module + + +def _install_boundary_hooks(transformer, boundaries, parallel_state): + """Install shard/gather hooks from the plan's boundary specs. + + ContextParallelInput entries split module inputs (or outputs when + split_output=True); ContextParallelOutput entries gather module outputs. + The gather's backward sums across SP, pairing with FSDP's 1/(dp*sp) mean. + """ + for path, spec in boundaries.items(): + module = _resolve_submodule(transformer, path) if path else transformer + + if isinstance(spec, ContextParallelOutput): + + def gather_output(mod, args, output, _spec=spec): + assert isinstance(output, torch.Tensor) + return gather_sequence( + output, + parallel_state.sp_group, + parallel_state.sp_rank, + parallel_state.sp_size, + dim=_spec.gather_dim, + ) + + module.register_forward_hook(gather_output) + continue + + input_specs = {k: v for k, v in spec.items() if not v.split_output} + output_specs = {k: v for k, v in spec.items() if v.split_output} + + if input_specs: + # Wrappers like PeftModel expose forward(*args, **kwargs); the real + # parameter names live on the base module. + sig_module = module.get_base_model() if hasattr(module, "get_base_model") else module + param_names = list(inspect.signature(sig_module.forward).parameters) + missing = [k for k in input_specs if isinstance(k, str) and k not in param_names] + if missing: + raise ValueError( + f"boundary keys {missing} at '{path}' are not parameters of " + f"{type(sig_module).__name__}.forward" + ) + + def split_inputs(mod, args, kwargs, _specs=input_specs, _names=param_names): + args = list(args) + for key, s in _specs.items(): + if isinstance(key, str) and key in kwargs: + kwargs[key] = _split_if_expected(kwargs[key], s, parallel_state) + continue + index = key if isinstance(key, int) else (_names.index(key) if key in _names else None) + if index is not None and index < len(args): + args[index] = _split_if_expected(args[index], s, parallel_state) + return tuple(args), kwargs + + module.register_forward_pre_hook(split_inputs, with_kwargs=True) + + if output_specs: + + def split_outputs(mod, args, kwargs, output, _specs=output_specs): + single = not isinstance(output, tuple) + out = [output] if single else list(output) + for index, s in _specs.items(): + out[index] = _split_if_expected(out[index], s, parallel_state) + return out[0] if single else tuple(out) + + module.register_forward_hook(split_outputs, with_kwargs=True) + + +def apply_sequence_parallel(transformer, parallel_state, plan): + """Wire SP into one transformer per its plan: install the family's SP + self-attention and the shard/gather boundary hooks. Call once per + transformer after FSDP wrapping.""" + if plan.num_attention_heads % parallel_state.ulysses_degree != 0: + raise ValueError( + f"num_attention_heads({plan.num_attention_heads}) is not divisible by " + f"ulysses_degree({parallel_state.ulysses_degree})" + ) + plan.attention_installer(transformer, parallel_state) + _install_boundary_hooks(transformer, plan.boundaries, parallel_state) diff --git a/miles/backends/fsdp_utils/sequence_parallel/topology.py b/miles/backends/fsdp_utils/sequence_parallel/topology.py new file mode 100644 index 00000000..8536ce04 --- /dev/null +++ b/miles/backends/fsdp_utils/sequence_parallel/topology.py @@ -0,0 +1,54 @@ +"""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.""" + if sequence_parallel_size < 1: + raise ValueError(f"sequence_parallel_size must be positive, got {sequence_parallel_size}") + if ulysses_degree < 0: + raise ValueError(f"ulysses_degree must be non-negative, got {ulysses_degree}") + sp = 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 = validate_sp_config(world_size, 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 diff --git a/miles/backends/training_utils/parallel.py b/miles/backends/training_utils/parallel.py index 4283e873..1ef33717 100644 --- a/miles/backends/training_utils/parallel.py +++ b/miles/backends/training_utils/parallel.py @@ -1,4 +1,5 @@ from dataclasses import dataclass + import torch.distributed as dist @@ -11,14 +12,19 @@ class ParallelState: dp_rank: int dp_src_rank: int dp_size: int - cp_rank: int - cp_size: int - dp_cp_rank: int - dp_cp_size: int dp_group: dist.ProcessGroup | None - dp_cp_group: dist.ProcessGroup | None - dp_cp_group_gloo: dist.ProcessGroup | None - cp_group: dist.ProcessGroup | None + # Sequence Parallelism (USP = Ulysses x Ring) + sp_rank: int + sp_size: int + sp_group: dist.ProcessGroup | None + ulysses_degree: int + ring_degree: int + ulysses_group: dist.ProcessGroup | None + ring_group: dist.ProcessGroup | None + # dp x sp spans every training rank + dp_sp_rank: int + dp_sp_size: int + dp_sp_group_gloo: dist.ProcessGroup | None tp_size: int tp_rank: int tp_group: dist.ProcessGroup | None diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index c123cf51..ab32b84d 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1276,7 +1276,6 @@ def parse_args(add_custom_arguments=None): args = load_fsdp_args(extra_args_provider=add_miles_arguments) args.rank = 0 # Primary process rank for wandb initialization args.world_size = args.actor_num_nodes * args.actor_num_gpus_per_node - assert args.context_parallel_size == 1, "Context parallelism is not supported for FSDP backend." miles_validate_args(args) sglang_validate_args(args) @@ -1447,6 +1446,11 @@ def miles_validate_args(args): "debug_rollout_only and debug_train_only cannot be set at the same time, " "please set only one of them." ) + if getattr(args, "diffusion_model", None): + from miles.backends.fsdp_utils.arguments import validate_sp_args + + validate_sp_args(args) + # always true on offload for colocate at the moment. if args.colocate: if args.offload_train is None: @@ -1496,7 +1500,9 @@ def miles_validate_args(args): ) args.global_batch_size = derived_gbs - dp_size = args.actor_num_gpus_per_node * args.actor_num_nodes + train_world_size = args.actor_num_gpus_per_node * args.actor_num_nodes + sp_size = args.sequence_parallel_size if getattr(args, "diffusion_model", None) else 1 + dp_size = train_world_size // sp_size if args.global_batch_size is not None: assert ( args.global_batch_size % dp_size == 0 diff --git a/miles/utils/train_data_utils.py b/miles/utils/train_data_utils.py index 18a95f5b..9a663317 100644 --- a/miles/utils/train_data_utils.py +++ b/miles/utils/train_data_utils.py @@ -429,21 +429,21 @@ def reorder_train_pairs_for_tiling( return [train_data[i] for step in schedule for micro_batch in step for i in micro_batch] -def validate_same_microbatch_counts_across_dp( +def validate_same_microbatch_counts_across_train_ranks( *, microbatch_schedule: list[list[tuple[int, int]]], parallel_state, ) -> None: - """Ensure every DP rank will run the same number of FSDP micro-batches.""" + """Ensure every rank in the flattened DP×SP FSDP mesh runs the same number of micro-batches.""" local_microbatch_counts = [len(step_ranges) for step_ranges in microbatch_schedule] - gathered_microbatch_counts = [None] * parallel_state.dp_cp_size + gathered_microbatch_counts = [None] * parallel_state.dp_sp_size dist.all_gather_object( gathered_microbatch_counts, local_microbatch_counts, - group=parallel_state.dp_cp_group_gloo, + group=parallel_state.dp_sp_group_gloo, ) if any(counts != local_microbatch_counts for counts in gathered_microbatch_counts): raise ValueError( - "Uneven train-pair counts would make DP ranks run different numbers of FSDP " + "Uneven train-pair counts would make training ranks run different numbers of FSDP " f"micro-batches per optimizer step: {gathered_microbatch_counts}" ) diff --git a/tests/fast-gpu/backends/fsdp_utils/sequence_parallel/_attention_parity_worker.py b/tests/fast-gpu/backends/fsdp_utils/sequence_parallel/_attention_parity_worker.py new file mode 100644 index 00000000..6680662d --- /dev/null +++ b/tests/fast-gpu/backends/fsdp_utils/sequence_parallel/_attention_parity_worker.py @@ -0,0 +1,213 @@ +"""Four-rank worker for USP parity against full-sequence SDPA. + +This is launched by ``test_attention_parity.py`` rather than discovered as a +standalone CI test. Every rank constructs the same full Q/K/V reference, then +runs USP from its sequence shard and compares its local output and input grads. +""" + +import argparse +import os + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + +from miles.backends.fsdp_utils.sequence_parallel.attention import usp_attention +from miles.backends.fsdp_utils.sequence_parallel.topology import sp_subgroups + + +SP_SIZE = 4 +SHAPE = (2, 128, 8, 64) # [batch, global sequence, heads, head dim] +DTYPE = torch.bfloat16 + +# These bounds compare production-dtype Ring attention against the same +# full-sequence flash-SDPA reference. The observed error is at most one bf16 +# quantization step for this input band; the bounds leave only a small margin. +# Pure Ulysses is a lossless permutation around per-head attention, so both its +# forward and dQ/dK/dV are required to remain bitwise identical. +TOLERANCES = { + 2: {"forward": (8e-3, 1e-3), "backward": (8e-3, 2.5e-4)}, + 1: {"forward": (8e-3, 1e-3), "backward": (8e-3, 2.5e-4)}, +} + + +def _sdpa(query, key, value): + scale = query.shape[-1] ** -0.5 + with sdpa_kernel(SDPBackend.FLASH_ATTENTION): + output = F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + dropout_p=0.0, + is_causal=False, + scale=scale, + ) + return output.transpose(1, 2).contiguous() + + +def _make_full_inputs(device): + generator = torch.Generator(device=device).manual_seed(20260721) + query = torch.randn(SHAPE, device=device, dtype=DTYPE, generator=generator) * 0.5 + key = torch.randn(SHAPE, device=device, dtype=DTYPE, generator=generator) * 0.5 + value = torch.randn(SHAPE, device=device, dtype=DTYPE, generator=generator) * 0.5 + grad_output = torch.randn(SHAPE, device=device, dtype=DTYPE, generator=generator) * 0.1 + for tensor in (query, key, value, grad_output): + dist.broadcast(tensor, src=0) + return query, key, value, grad_output + + +def _create_usp_groups(ulysses_degree): + rank = dist.get_rank() + _, _, _, ulysses_ranks, ring_ranks = sp_subgroups(SP_SIZE, SP_SIZE, ulysses_degree) + + ulysses_group = None + if ulysses_degree > 1: + for ranks in ulysses_ranks: + group = dist.new_group(ranks) + if rank in ranks: + ulysses_group = group + + ring_degree = SP_SIZE // ulysses_degree + ring_group = None + if ring_degree > 1: + for ranks in ring_ranks: + group = dist.new_group(ranks) + if rank in ranks: + ring_group = group + return ulysses_group, ring_group + + +def _run_reference(query, key, value, grad_output): + query = query.detach().clone().requires_grad_(True) + key = key.detach().clone().requires_grad_(True) + value = value.detach().clone().requires_grad_(True) + output = _sdpa(query, key, value) + output.backward(grad_output) + return output.detach(), (query.grad.detach(), key.grad.detach(), value.grad.detach()) + + +def _run_usp(query, key, value, grad_output, ulysses_group, ring_group): + rank = dist.get_rank() + local_sequence = SHAPE[1] // SP_SIZE + start = rank * local_sequence + query = query[:, start : start + local_sequence].detach().clone().requires_grad_(True) + key = key[:, start : start + local_sequence].detach().clone().requires_grad_(True) + value = value[:, start : start + local_sequence].detach().clone().requires_grad_(True) + local_grad_output = grad_output[:, start : start + local_sequence].contiguous() + + output = usp_attention( + query, + key, + value, + ulysses_group=ulysses_group, + ring_group=ring_group, + local_attention_fn=_sdpa, + ) + output.backward(local_grad_output) + return output.detach(), (query.grad.detach(), key.grad.detach(), value.grad.detach()) + + +def _local_shard(tensor): + local_sequence = SHAPE[1] // SP_SIZE + start = dist.get_rank() * local_sequence + return tensor[:, start : start + local_sequence].contiguous() + + +def _global_stats(actual, expected): + difference = (actual.float() - expected.float()).abs() + max_abs = difference.max() + normalized = max_abs / expected.float().abs().max().clamp_min(1e-12) + mismatches = torch.count_nonzero(actual != expected).to(torch.int64) + total = torch.tensor(actual.numel(), device=actual.device, dtype=torch.int64) + dist.all_reduce(max_abs, op=dist.ReduceOp.MAX) + dist.all_reduce(normalized, op=dist.ReduceOp.MAX) + dist.all_reduce(mismatches, op=dist.ReduceOp.SUM) + dist.all_reduce(total, op=dist.ReduceOp.SUM) + return max_abs.item(), normalized.item(), mismatches.item(), total.item() + + +def _assert_bitwise(name, actual, expected): + max_abs, normalized, mismatches, total = _global_stats(actual, expected) + if dist.get_rank() == 0: + print( + f"{name}: bitwise={mismatches == 0} mismatches={mismatches}/{total} " + f"max_abs={max_abs:.3e} normalized={normalized:.3e}", + flush=True, + ) + assert mismatches == 0, f"{name} must be bitwise identical; {mismatches}/{total} values differ" + + +def _assert_close(name, actual, expected, *, rtol, atol): + max_abs, normalized, mismatches, total = _global_stats(actual, expected) + if dist.get_rank() == 0: + print( + f"{name}: bitwise={mismatches == 0} mismatches={mismatches}/{total} " + f"max_abs={max_abs:.3e} normalized={normalized:.3e} rtol={rtol:.1e} atol={atol:.1e}", + flush=True, + ) + torch.testing.assert_close(actual, expected, rtol=rtol, atol=atol) + + +def _enable_deterministic_mode(): + assert os.environ.get("CUBLAS_WORKSPACE_CONFIG") == ":4096:8" + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + torch.use_deterministic_algorithms(True, warn_only=False) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--ulysses-degree", type=int, choices=(1, 2, 4), required=True) + parser.add_argument("--deterministic", action="store_true") + args = parser.parse_args() + + if args.deterministic: + _enable_deterministic_mode() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl", device_id=device) + assert dist.get_world_size() == SP_SIZE + + ulysses_group, ring_group = _create_usp_groups(args.ulysses_degree) + full_inputs = _make_full_inputs(device) + topology = f"sp4-u{args.ulysses_degree}r{SP_SIZE // args.ulysses_degree}" + + if args.deterministic: + output_1, grads_1 = _run_usp(*full_inputs, ulysses_group, ring_group) + output_2, grads_2 = _run_usp(*full_inputs, ulysses_group, ring_group) + _assert_bitwise(f"{topology} deterministic forward", output_1, output_2) + for name, actual, expected in zip(("dQ", "dK", "dV"), grads_1, grads_2, strict=True): + _assert_bitwise(f"{topology} deterministic {name}", actual, expected) + dist.barrier() + if dist.get_rank() == 0: + print(f"{topology} deterministic: PASS", flush=True) + dist.destroy_process_group() + return + + reference_output, reference_grads = _run_reference(*full_inputs) + usp_output, usp_grads = _run_usp(*full_inputs, ulysses_group, ring_group) + local_reference_output = _local_shard(reference_output) + if args.ulysses_degree == SP_SIZE: + _assert_bitwise(f"{topology} forward", usp_output, local_reference_output) + else: + rtol, atol = TOLERANCES[args.ulysses_degree]["forward"] + _assert_close(f"{topology} forward", usp_output, local_reference_output, rtol=rtol, atol=atol) + + for name, actual, expected in zip(("dQ", "dK", "dV"), usp_grads, reference_grads, strict=True): + local_expected = _local_shard(expected) + if args.ulysses_degree == SP_SIZE: + _assert_bitwise(f"{topology} {name}", actual, local_expected) + else: + rtol, atol = TOLERANCES[args.ulysses_degree]["backward"] + _assert_close(f"{topology} {name}", actual, local_expected, rtol=rtol, atol=atol) + + dist.barrier() + if dist.get_rank() == 0: + print(f"{topology}: PASS", flush=True) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/fast-gpu/backends/fsdp_utils/sequence_parallel/test_attention_parity.py b/tests/fast-gpu/backends/fsdp_utils/sequence_parallel/test_attention_parity.py new file mode 100644 index 00000000..c699b61f --- /dev/null +++ b/tests/fast-gpu/backends/fsdp_utils/sequence_parallel/test_attention_parity.py @@ -0,0 +1,63 @@ +from tests.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=240, + suite="stage-c-5-gpu-h200", + labels=["fsdp"], +) + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + + +_WORKER = Path(__file__).with_name("_attention_parity_worker.py") + + +def _run_worker(ulysses_degree, *, deterministic=False): + env = os.environ.copy() + env["PYTHONUNBUFFERED"] = "1" + command = [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nnodes=1", + "--nproc_per_node=4", + str(_WORKER), + "--ulysses-degree", + str(ulysses_degree), + ] + if deterministic: + env["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + command.append("--deterministic") + subprocess.run( + command, + check=True, + env=env, + ) + + +@pytest.mark.parametrize( + "ulysses_degree", + [4, 2, 1], + ids=["sp4-u4r1", "sp4-u2r2", "sp4-u1r4"], +) +def test_usp_forward_backward_matches_full_sequence_sdpa(ulysses_degree): + _run_worker(ulysses_degree) + + +@pytest.mark.parametrize( + "ulysses_degree", + [4, 2, 1], + ids=["sp4-u4r1", "sp4-u2r2", "sp4-u1r4"], +) +def test_usp_forward_backward_is_bitwise_repeatable_in_deterministic_mode(ulysses_degree): + _run_worker(ulysses_degree, deterministic=True) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))