diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 561a9a4e..0ecf29f3 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/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() diff --git a/miles/backends/fsdp_utils/arguments.py b/miles/backends/fsdp_utils/arguments.py index 6184d39e..4e84fc97 100644 --- a/miles/backends/fsdp_utils/arguments.py +++ b/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 diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 04079032..8bd0b90e 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/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) diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index 9bb8e500..c4ace19f 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/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") diff --git a/miles/backends/fsdp_utils/parallel.py b/miles/backends/fsdp_utils/parallel.py index 5128ef35..2ac732ce 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/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 diff --git a/miles/backends/fsdp_utils/sp_attention.py b/miles/backends/fsdp_utils/sp_attention.py new file mode 100644 index 00000000..d1eda772 --- /dev/null +++ b/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, 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 diff --git a/miles/backends/fsdp_utils/sp_mesh.py b/miles/backends/fsdp_utils/sp_mesh.py new file mode 100644 index 00000000..ce61eabd --- /dev/null +++ b/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 diff --git a/miles/backends/fsdp_utils/sp_ops.py b/miles/backends/fsdp_utils/sp_ops.py new file mode 100644 index 00000000..59d798e1 --- /dev/null +++ b/miles/backends/fsdp_utils/sp_ops.py @@ -0,0 +1,186 @@ +"""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; 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 returns each rank's slice. + + With sum_grad 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, sum_grad): + ctx.group = group + ctx.sp_rank = sp_rank + ctx.dim = dim + ctx.local_size = x.shape[dim] + ctx.sum_grad = sum_grad + 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): + if ctx.sum_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, 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, sum_grad=False): + return _GatherSequence.apply(x, group, sp_rank, sp_size, dim, sum_grad) + + +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, is_causal): + from torch.distributed.tensor.experimental._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=is_causal, + 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.is_causal, ctx.max_q, ctx.max_k = group, scale, is_causal, max_q, max_k + return out + + @staticmethod + def backward(ctx, grad_out): + from torch.distributed.tensor.experimental._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=ctx.is_causal, + 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, None + + +def usp_attention(query, key, value, ulysses_group=None, ring_group=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; with no groups this is plain SDPA. + """ + scale = query.shape[-1] ** -0.5 + + 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) + + q = query.transpose(1, 2) # [B, H, S, D] + k = key.transpose(1, 2) + v = value.transpose(1, 2) + if ring_group is not None: + out = _RingFlashAttention.apply(q.contiguous(), k.contiguous(), v.contiguous(), ring_group, scale, False) + else: + 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/sp_plan.py b/miles/backends/fsdp_utils/sp_plan.py new file mode 100644 index 00000000..b6eee22a --- /dev/null +++ b/miles/backends/fsdp_utils/sp_plan.py @@ -0,0 +1,131 @@ +"""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 .sp_ops import gather_sequence, shard_sequence + + +@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``: called with (transformer, parallel_state); routes the model's + self-attention through ``sp_ops.usp_attention``. + """ + + boundaries: dict + attention: Callable[[torch.nn.Module, object], None] + num_attention_heads: int + + +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, sum_grad=True): + """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 is a differentiable all-gather; its sum_grad backward pairs + with FSDP's 1/(dp*sp) mean (sum_grad=False is the sp-replicated-parameter + variant, used by tests as a validation anchor). + """ + 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, + sum_grad=sum_grad, + ) + + 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, sum_grad=True): + """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(transformer, parallel_state) + _install_boundary_hooks(transformer, plan.boundaries, parallel_state, sum_grad=sum_grad) diff --git a/miles/backends/training_utils/parallel.py b/miles/backends/training_utils/parallel.py index 4283e873..4378be87 100644 --- a/miles/backends/training_utils/parallel.py +++ b/miles/backends/training_utils/parallel.py @@ -25,3 +25,11 @@ class ParallelState: is_pp_last_stage: bool = True vpp_size: int | None = 1 microbatch_group_size_per_vp_stage: int | None = None + # Sequence Parallelism (USP = Ulysses x Ring). cp_* fields above alias sp_*. + sp_rank: int = 0 + sp_size: int = 1 + sp_group: dist.ProcessGroup | None = None + ulysses_degree: int = 1 + ring_degree: int = 1 + ulysses_group: dist.ProcessGroup | None = None + ring_group: dist.ProcessGroup | None = None diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 0ada321d..73b0f506 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1244,8 +1244,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) @@ -1360,6 +1358,10 @@ def miles_validate_args(args): if args.model_backend_path is None: args.model_backend_path = cfg_cls.model_backend_path cfg_cls.validate_args(args) + if getattr(args, "sequence_parallel_size", 1) > 1: + from miles.backends.fsdp_utils.model_backend import validate_sp_support + + validate_sp_support(args, cfg_cls) if args.dump_details is not None: args.save_debug_rollout_data = f"{args.dump_details}/rollout_data/{{rollout_id}}.pt" diff --git a/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh b/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh index bb2d7904..c542fd59 100755 --- a/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh +++ b/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh @@ -19,9 +19,8 @@ # --micro-batch-size 2: the one-step-per-rollout schedule keeps every # micro-batch phase-pure (one DiT, one CFG scale); mbs=4 OOMs on H200, 2 fits. # -# NOTE: gradient checkpointing stays OFF. Wan2.2 under FSDP2 mixed precision hits -# torch.utils.checkpoint CheckpointError (fp32 RoPE freqs buffers; fix pending -# in a separate PR). If you OOM, lower --rollout-batch-size, +# NOTE: gradient checkpointing stays OFF by default (not needed at 5 frames). +# If you OOM, enable --gradient-checkpointing or lower --rollout-batch-size, # --n-samples-per-prompt, or --diffusion-microgroup-size. # # Layout: first 4 GPUs in CUDA_VISIBLE_DEVICES = train+sgld colocate, @@ -49,6 +48,11 @@ fi PYTHON_BIN="${PYTHON_BIN:-python}" +# Sequence parallelism (USP). Default off; sp = ulysses x ring must divide the +# train world size. 0 = auto (ulysses fills sp); ring = sp / ulysses. +SP_SIZE="${SP_SIZE:-1}" +ULYSSES_DEGREE="${ULYSSES_DEGREE:-0}" + DATASETS_DIR="/root/datasets/miles-diffusion-datasets" if [[ ! -f "${DATASETS_DIR}/flowgrpo_pickscore/train.jsonl" ]]; then hf download --repo-type dataset rockdu/miles-diffusion-datasets \ @@ -77,6 +81,8 @@ WAN_LORA_TARGET_MODULES=( --diffusion-microgroup-size 8 \ --micro-batch-size 2 \ --actor-num-gpus-per-node 4 \ + --sequence-parallel-size "${SP_SIZE}" \ + --ulysses-degree "${ULYSSES_DEGREE}" \ --rollout-num-gpus 4 \ --rollout-num-gpus-per-engine 1 \ --num-gpus-per-node 5 \