From a4cd0cc23f1512a5bb7648b6ca4d80646402b948 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 6 Jul 2026 22:30:42 +0000 Subject: [PATCH 01/24] diffusion: sequence-parallel config and Option B parallel state FSDP shards on the dp mesh dim only; SP (ulysses x ring) gets its own process groups plus sglang's _SP coordinator so USPAttention resolves them. context_parallel_size stays as a backward-compatible alias. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/arguments.py | 8 ++- miles/backends/fsdp_utils/parallel.py | 79 ++++++++++++++++------- miles/backends/fsdp_utils/sp_mesh.py | 43 ++++++++++++ miles/backends/training_utils/parallel.py | 8 +++ miles/utils/arguments.py | 3 +- 5 files changed, 113 insertions(+), 28 deletions(-) create mode 100644 miles/backends/fsdp_utils/sp_mesh.py diff --git a/miles/backends/fsdp_utils/arguments.py b/miles/backends/fsdp_utils/arguments.py index 6184d39e..b18a8e5d 100644 --- a/miles/backends/fsdp_utils/arguments.py +++ b/miles/backends/fsdp_utils/arguments.py @@ -57,8 +57,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). context_parallel_size is a + # backward-compatible alias for sequence_parallel_size. + context_parallel_size: int = 1 + sequence_parallel_size: int = 1 + ulysses_degree: int = 0 # 0=auto: ulysses fills sp, ring=1 + ring_degree: int = 0 # 0=auto: sp // ulysses # YAML bookkeeping config: str | None = None diff --git a/miles/backends/fsdp_utils/parallel.py b/miles/backends/fsdp_utils/parallel.py index 5128ef35..3389fb2d 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/miles/backends/fsdp_utils/parallel.py @@ -1,61 +1,90 @@ import logging from argparse import Namespace +import torch import torch.distributed as dist from torch.distributed.device_mesh import init_device_mesh 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. + + FSDP shards parameters on the dp mesh dim only; SP gets its own process + groups and parameters are replicated across sp ranks. + """ 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, + args.ring_degree, + getattr(args, "num_attention_heads", None), + ) + dp_rank, sp_rank, _, _ = locate_rank(rank, sp_size, ulysses_degree, ring_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 + # USPAttention reads sglang's module-global _SP coordinator; + # set_seq_parallel_pg_by_sp_groups only sets ULYSSES_PG/RING_PG. + ulysses_group = ring_group = None + if sp_size > 1: + from sglang.multimodal_gen.runtime.distributed import parallel_state as sgl_sp + from sglang.multimodal_gen.runtime.distributed.parallel_groups import ( + PROCESS_GROUP, + set_seq_parallel_pg_by_sp_groups, + ) - 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)") + _, _, sp_groups, _, _ = sp_subgroups(world_size, sp_size, ulysses_degree, ring_degree) + set_seq_parallel_pg_by_sp_groups(ulysses_degree, ring_degree, rank, sp_groups) + ulysses_group = PROCESS_GROUP.ULYSSES_PG + ring_group = PROCESS_GROUP.RING_PG + sgl_sp._SP = sgl_sp.init_parallel_group_coordinator( + group_ranks=sp_groups, + local_rank=torch.cuda.current_device(), + backend=dist.get_backend(), + parallel_mode="sequence", + ulysses_group=ulysses_group, + ring_group=ring_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]), + 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"] - return parallel_state diff --git a/miles/backends/fsdp_utils/sp_mesh.py b/miles/backends/fsdp_utils/sp_mesh.py new file mode 100644 index 00000000..5b22cb7d --- /dev/null +++ b/miles/backends/fsdp_utils/sp_mesh.py @@ -0,0 +1,43 @@ +"""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, ring_degree=0): + """Normalize (sp, ulysses, ring). degree=0 means auto: ulysses fills sp, ring=1.""" + sp = max(1, sequence_parallel_size) + u = ulysses_degree or sp + r = ring_degree or (sp // u) + if u * r != sp: + raise ValueError(f"ulysses_degree({u}) * ring_degree({r}) != sequence_parallel_size({sp})") + return sp, u, r + + +def validate_sp_config(world_size, sequence_parallel_size, ulysses_degree=0, ring_degree=0, num_heads=None): + """Validate at startup. Returns (sp, ulysses, ring).""" + sp, u, r = resolve_sp_degrees(sequence_parallel_size, ulysses_degree, ring_degree) + if world_size % sp != 0: + raise ValueError(f"world_size({world_size}) is not divisible by sequence_parallel_size({sp})") + if num_heads is not None and num_heads % u != 0: + raise ValueError(f"num_heads({num_heads}) is not divisible by ulysses_degree({u})") + return sp, u, r + + +def sp_subgroups(world_size, sequence_parallel_size, ulysses_degree=0, ring_degree=0): + """Return (dp_size, sp_size, sp_groups, ulysses_groups, ring_groups); groups are global-rank lists.""" + sp, u, r = resolve_sp_degrees(sequence_parallel_size, ulysses_degree, ring_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, ring_degree=0): + """Locate a global rank's (dp_rank, sp_rank, ulysses_rank, ring_rank).""" + sp, u, _ = resolve_sp_degrees(sequence_parallel_size, ulysses_degree, ring_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..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..83c6fcc6 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1244,7 +1244,8 @@ 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." + if args.context_parallel_size > 1 and args.sequence_parallel_size == 1: + args.sequence_parallel_size = args.context_parallel_size miles_validate_args(args) sglang_validate_args(args) From 6f15b3bd126ad0296e72d9f704db15e9bc7d69c0 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 6 Jul 2026 22:30:42 +0000 Subject: [PATCH 02/24] diffusion: USP sequence parallelism for Wan training Self-attention runs sglang's USPAttention; the sequence is sharded before the first block and gathered after proj_out, so model outputs stay full-sequence and loss/log_prob code is unchanged. Partial grads are summed across the sp group after scaler.unscale_. Applied per component, covering dual-DiT. Recipe gains SP_SIZE/ULYSSES_DEGREE/RING_DEGREE. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/actor.py | 25 +++ miles/backends/fsdp_utils/sp_attention.py | 187 ++++++++++++++++++ ...run-diffusion-grpo-wan22-pickscore-5gpu.sh | 14 +- 3 files changed, 223 insertions(+), 3 deletions(-) create mode 100644 miles/backends/fsdp_utils/sp_attention.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 561a9a4e..f0b12cab 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_attention import apply_sequence_parallel logger = logging.getLogger(__name__) @@ -137,6 +138,11 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty 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(): + apply_sequence_parallel(model, self.parallel_state, compute_dtype=self._forward_dtype) + # Force a sync to ensure sharding is complete and old memory is freed. torch.cuda.synchronize() clear_memory() @@ -406,6 +412,8 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: self.prof.step(rollout_id=rollout_id) if not self.args.debug_skip_optimizer_step: self.scaler.unscale_(self.optimizer) + if self.parallel_state.sp_size > 1: + self._all_reduce_sp_grads() grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip_grad) log_stats["grad_norm"].append(grad_norm.detach()) self.scaler.step(self.optimizer) @@ -419,6 +427,23 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: reduced = {f"train/{k}": torch.stack(v).mean().item() for k, v in log_stats.items()} self._gather_and_log_metrics(rollout_id, reduced, step=self.global_step) + def _all_reduce_sp_grads(self) -> None: + """Each sp rank sees 1/sp of the tokens, so its grads are partial; all-reduce(SUM) + across the sp group restores the full gradient. Runs on FSDP-sharded grads, in fp32.""" + from torch.distributed.tensor import DTensor + + group = self.parallel_state.sp_group + for p in self.model.parameters(): + if p.grad is None: + continue + local = p.grad.to_local() if isinstance(p.grad, DTensor) else p.grad + if local.dtype == torch.float32: + dist.all_reduce(local, group=group) + else: + f = local.float() + dist.all_reduce(f, group=group) + local.copy_(f) + def _maybe_legacy_window_pad_len(self, train_pairs: list, microbatch_ranges: list) -> int | None: """LEGACY 2D parity: the whole-window max cond seq_len (like the legacy tile path), or None unless the legacy --micro-batch-size-sample>1 path is active. TODO: remove with it.""" diff --git a/miles/backends/fsdp_utils/sp_attention.py b/miles/backends/fsdp_utils/sp_attention.py new file mode 100644 index 00000000..ad4ee9db --- /dev/null +++ b/miles/backends/fsdp_utils/sp_attention.py @@ -0,0 +1,187 @@ +"""Sequence parallelism for diffusers Wan DiT: self-attention runs sglang's USPAttention. + +Each sp rank holds S/sp latent tokens; attention internally gathers to the full +sequence via Ulysses all-to-all (+Ring) and scatters back. The sequence is sharded +before the first block and gathered after proj_out, so every parameter sees a +partial gradient and the actor's cross-sp grad all-reduce applies uniformly. +""" + +import functools + +import torch +import torch.distributed as dist + + +class _GatherSequence(torch.autograd.Function): + """All-gather S_local shards along dim=1; backward returns each rank's slice.""" + + @staticmethod + def forward(ctx, x, group, sp_rank, sp_size): + ctx.sp_rank = sp_rank + ctx.s_local = x.shape[1] + parts = [torch.empty_like(x) for _ in range(sp_size)] + dist.all_gather(parts, x.contiguous(), group=group) + return torch.cat(parts, dim=1) + + @staticmethod + def backward(ctx, grad): + start = ctx.sp_rank * ctx.s_local + return grad[:, start : start + ctx.s_local], None, None, None + + +def shard_sequence(x, parallel_state): + sp = parallel_state.sp_size + s = x.shape[1] + if s % sp: + raise ValueError(f"sequence length {s} is not divisible by sp_size {sp}") + s_local = s // sp + start = parallel_state.sp_rank * s_local + return x[:, start : start + s_local] + + +def gather_sequence(x, parallel_state): + return _GatherSequence.apply(x, parallel_state.sp_group, parallel_state.sp_rank, parallel_state.sp_size) + + +class WanUSPAttnProcessor: + """Wan attention processor: self-attn via USPAttention, cross-attn via local SDPA. + + Reuses Wan's QKV/RMSNorm/RoPE; rotary_emb arrives pre-sharded to S_local. + """ + + def __init__(self, num_heads, head_dim, compute_dtype=torch.bfloat16): + from sglang.multimodal_gen.runtime.layers.attention.layer import USPAttention + + init_sp_backend(compute_dtype) + self.usp_attn = USPAttention(num_heads=num_heads, head_size=head_dim, causal=False) + + def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, rotary_emb=None): + is_self_attn = encoder_hidden_states is None + + encoder_hidden_states_img = None + if attn.add_k_proj is not None: + image_context_length = encoder_hidden_states.shape[1] - 512 + encoder_hidden_states_img = encoder_hidden_states[:, :image_context_length] + encoder_hidden_states = encoder_hidden_states[:, image_context_length:] + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + if rotary_emb is not None: + query = _apply_rotary_emb(query, *rotary_emb) + key = _apply_rotary_emb(key, *rotary_emb) + + if is_self_attn: + hidden_states = self.usp_attn(query, key, value) + else: + hidden_states = torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2) + + hidden_states = hidden_states.flatten(2, 3).type_as(query) + + if encoder_hidden_states_img is not None: + key_img = attn.norm_added_k(attn.add_k_proj(encoder_hidden_states_img)).unflatten(2, (attn.heads, -1)) + value_img = attn.add_v_proj(encoder_hidden_states_img).unflatten(2, (attn.heads, -1)) + hidden_states_img = ( + torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), + key_img.transpose(1, 2), + value_img.transpose(1, 2), + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + .transpose(1, 2) + .flatten(2, 3) + .type_as(query) + ) + hidden_states = hidden_states + hidden_states_img + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +def init_sp_backend(compute_dtype): + """Set up the sglang runtime state USPAttention depends on. + + The training process never goes through the sglang server launch, so the + attention backend, mixed-precision policy, and forward context are unset. + FA requires fp16/bf16; the forward context is persisted module-globally + because checkpoint recompute runs after any with-scope has exited. + """ + from sglang.multimodal_gen.runtime.layers.attention.selector import global_force_attn_backend + from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum + from sglang.multimodal_gen.utils import set_mixed_precision_policy + + half = compute_dtype in (torch.float16, torch.bfloat16) + global_force_attn_backend(AttentionBackendEnum.FA if half else AttentionBackendEnum.TORCH_SDPA) + set_mixed_precision_policy(param_dtype=compute_dtype, reduce_dtype=torch.float32) + + from sglang.multimodal_gen.runtime.managers import forward_context as fc + + if fc._forward_context is None: + fc._forward_context = fc.ForwardContext(current_timestep=0, attn_metadata=None) + + +def _apply_rotary_emb(hidden_states, freqs_cos, freqs_sin): + x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + cos = freqs_cos[..., 0::2] + sin = freqs_sin[..., 1::2] + out = torch.empty_like(hidden_states) + out[..., 0::2] = x1 * cos - x2 * sin + out[..., 1::2] = x1 * sin + x2 * cos + return out.type_as(hidden_states) + + +def apply_sequence_parallel(transformer, parallel_state, compute_dtype=None): + """Wire SP into one Wan transformer: replace self-attn processors and install + the shard/gather hooks. Call once per transformer after FSDP wrapping.""" + heads = transformer.config.num_attention_heads + head_dim = transformer.config.attention_head_dim + # args carry no head count, so the startup divisibility check must happen + # here where the real model config is available. + if heads % parallel_state.ulysses_degree != 0: + raise ValueError( + f"num_attention_heads({heads}) is not divisible by ulysses_degree({parallel_state.ulysses_degree})" + ) + if compute_dtype is None: + compute_dtype = next(transformer.parameters()).dtype + transformer.set_attn_processor(WanUSPAttnProcessor(heads, head_dim, compute_dtype)) + + # RoPE runs once per forward and its output is reused by every block; shard at the source. + rope = transformer.rope + orig_rope_forward = rope.forward + + @functools.wraps(orig_rope_forward) + def sp_rope_forward(hidden_states): + cos, sin = orig_rope_forward(hidden_states) + return shard_sequence(cos, parallel_state), shard_sequence(sin, parallel_state) + + rope.forward = sp_rope_forward + + def slice_block_input(module, args): + hs, *rest = args + return (shard_sequence(hs, parallel_state), *rest) + + def gather_proj_output(module, args, output): + return gather_sequence(output, parallel_state) + + transformer.blocks[0].register_forward_pre_hook(slice_block_input) + transformer.proj_out.register_forward_hook(gather_proj_output) diff --git a/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh b/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh index bb2d7904..79e6a00e 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,12 @@ 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). +SP_SIZE="${SP_SIZE:-1}" +ULYSSES_DEGREE="${ULYSSES_DEGREE:-0}" +RING_DEGREE="${RING_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 +82,9 @@ 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}" \ + --ring-degree "${RING_DEGREE}" \ --rollout-num-gpus 4 \ --rollout-num-gpus-per-engine 1 \ --num-gpus-per-node 5 \ From bc445195fd280381f84ba4cdcb93d790f0a857d2 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 6 Jul 2026 22:30:42 +0000 Subject: [PATCH 03/24] diffusion: align engine grouping span with rollout, assert full coverage connect_rollout_engines grouped by raw rollout_num_gpus_per_engine while rollout.init_rollout_engines uses min(per_engine, num_gpus_per_node); a mismatch orphans trailing ranks that fail later with an opaque error. Co-Authored-By: Claude Fable 5 --- .../fsdp_utils/diffusion_update_weight_utils.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index dc6193c8..c98114ea 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -107,10 +107,18 @@ def connect_rollout_engines( ) -> None: self.rollout_engines = rollout_engines + # Match rollout.init_rollout_engines' per-engine span; a mismatch would + # leave orphaned ranks that fail much later in update_bucket_weights. + num_gpu_per_engine = min(self.args.rollout_num_gpus_per_engine, self.args.num_gpus_per_node) + assert dist.get_world_size() == len(rollout_engines) * num_gpu_per_engine, ( + f"train world {dist.get_world_size()} != {len(rollout_engines)} engines x " + f"{num_gpu_per_engine} gpus per engine" + ) + # Here we assume the gpu id of rollout engines and train actors are the same. for i, engine in enumerate(self.rollout_engines): - start_rank = i * self.args.rollout_num_gpus_per_engine - end_rank = (i + 1) * self.args.rollout_num_gpus_per_engine + start_rank = i * num_gpu_per_engine + end_rank = (i + 1) * num_gpu_per_engine group_ranks = list(range(start_rank, end_rank)) new_group = dist.new_group( ranks=group_ranks, From 58cfa2449c935a8b9f09ac08d99299f85db62d87 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 6 Jul 2026 22:34:13 +0000 Subject: [PATCH 04/24] diffusion: SP test suite (mesh, init smoke, parity, weight sync, import guard) Co-Authored-By: Claude Fable 5 --- tests/sp/__init__.py | 0 tests/sp/run_gpu_suite.sh | 28 +++++ tests/sp/sglang_usp_import_guard.py | 40 +++++++ tests/sp/sp_attention_parity.py | 155 ++++++++++++++++++++++++++++ tests/sp/sp_grad_sync_parity.py | 145 ++++++++++++++++++++++++++ tests/sp/sp_init_smoke.py | 82 +++++++++++++++ tests/sp/sp_weight_sync_parity.py | 117 +++++++++++++++++++++ tests/sp/test_sp_mesh.py | 87 ++++++++++++++++ 8 files changed, 654 insertions(+) create mode 100644 tests/sp/__init__.py create mode 100755 tests/sp/run_gpu_suite.sh create mode 100644 tests/sp/sglang_usp_import_guard.py create mode 100644 tests/sp/sp_attention_parity.py create mode 100644 tests/sp/sp_grad_sync_parity.py create mode 100644 tests/sp/sp_init_smoke.py create mode 100644 tests/sp/sp_weight_sync_parity.py create mode 100644 tests/sp/test_sp_mesh.py diff --git a/tests/sp/__init__.py b/tests/sp/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/sp/run_gpu_suite.sh b/tests/sp/run_gpu_suite.sh new file mode 100755 index 00000000..f6754fa8 --- /dev/null +++ b/tests/sp/run_gpu_suite.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +# Full SP GPU test suite (needs 4 GPUs). CPU tests: pytest tests/sp/test_sp_mesh.py +set -euo pipefail +cd "$(dirname "$0")/../.." +export PYTHONPATH=. + +run() { echo; echo "===== $* ====="; "$@"; } + +run python -m pytest -q tests/sp/sglang_usp_import_guard.py + +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 2 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 4 --ulysses 4 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 4 --ulysses 2 --ring 2 + +run torchrun --standalone --nproc_per_node=2 tests/sp/sp_attention_parity.py --sp 2 +run torchrun --standalone --nproc_per_node=2 tests/sp/sp_attention_parity.py --sp 2 --ckpt +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 4 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 4 --ckpt +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 2 --ring 2 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 2 --ring 2 --ckpt + +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py + +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 2 --ulysses 2 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 4 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 --ring 2 + +echo; echo "ALL SP GPU TESTS PASSED" diff --git a/tests/sp/sglang_usp_import_guard.py b/tests/sp/sglang_usp_import_guard.py new file mode 100644 index 00000000..708841e3 --- /dev/null +++ b/tests/sp/sglang_usp_import_guard.py @@ -0,0 +1,40 @@ +"""Guard: the sglang tree actually imported must contain the training patches. + +Multiple sglang checkouts can coexist on a machine; training must resolve to the +one with the USP autograd wrappers and the dtype/shape-aware checksum, otherwise +Ulysses backward silently drops gradients. Run before any GPU parity script. + +Usage: PYTHONPATH=. python -m pytest tests/sp/sglang_usp_import_guard.py +""" + +import inspect + + +def test_usp_has_training_autograd(): + from sglang.multimodal_gen.runtime.layers import usp + + assert hasattr(usp, "_AllToAllSingle"), ( + f"imported sglang lacks _AllToAllSingle: {usp.__file__} — " + "Ulysses training backward would silently drop q/k/v grads" + ) + assert hasattr(usp, "_RingFlashAttention"), ( + f"imported sglang lacks _RingFlashAttention: {usp.__file__} — ring training dK/dV would be wrong" + ) + + +def test_checksum_covers_dtype_shape(): + from sglang.multimodal_gen.runtime.loader import weight_utils + + src = inspect.getsource(weight_utils.compute_weights_checksum) + assert "dtype" in src and "shape" in src, ( + f"imported sglang checksum is bytes-only: {weight_utils.__file__} — " + "transpose/reshape would not be detected" + ) + + +if __name__ == "__main__": + test_usp_has_training_autograd() + test_checksum_covers_dtype_shape() + from sglang.multimodal_gen.runtime.layers import usp + + print(f"[GUARD OK] imported sglang = {usp.__file__}") diff --git a/tests/sp/sp_attention_parity.py b/tests/sp/sp_attention_parity.py new file mode 100644 index 00000000..a5c27e2a --- /dev/null +++ b/tests/sp/sp_attention_parity.py @@ -0,0 +1,155 @@ +"""SP parity on a small real Wan DiT: USPAttention + shard/gather vs full-sequence reference. + +Both paths use the same attention kernel (the reference runs USPAttention with +skip_sequence_parallel), so any diff comes from the SP collectives' float +summation order and must stay within bf16 tolerance. Checks forward output, +input grad, and per-block self-attn + proj_out weight grads, with and without +gradient checkpointing. + +Usage: torchrun --standalone --nproc_per_node=N tests/sp/sp_attention_parity.py \ + --sp S [--ulysses U --ring R] [--ckpt] [--fp32] +""" + +import argparse + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from diffusers import WanTransformer3DModel + +from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state +from miles.backends.fsdp_utils.sp_attention import WanUSPAttnProcessor, apply_sequence_parallel +from miles.utils.distributed_utils import init_gloo_group + +DTYPE = torch.bfloat16 + + +def build_model(device): + torch.manual_seed(0) + model = WanTransformer3DModel( + patch_size=(1, 2, 2), + num_attention_heads=8, + attention_head_dim=128, + in_channels=16, + out_channels=16, + text_dim=4096, + freq_dim=256, + ffn_dim=1024, + num_layers=2, + rope_max_seq_len=1024, + ) + model = model.to(device=device, dtype=DTYPE) + model.train() + for p in model.parameters(): + dist.broadcast(p.data, src=0) + return model + + +def make_inputs(device): + g = torch.Generator(device=device).manual_seed(123) + hidden = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) + enc = torch.randn(1, 32, 4096, device=device, dtype=DTYPE, generator=g) + ts = torch.tensor([500], device=device) + out_grad = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) + for t in (hidden, enc, out_grad): + dist.broadcast(t, src=0) + return hidden, enc, ts, out_grad + + +def _set_ref_processor(model): + from sglang.multimodal_gen.runtime.layers.attention.layer import USPAttention + + proc = WanUSPAttnProcessor(model.config.num_attention_heads, model.config.attention_head_dim, DTYPE) + proc.usp_attn = USPAttention( + num_heads=model.config.num_attention_heads, + head_size=model.config.attention_head_dim, + causal=False, + skip_sequence_parallel=True, + ) + model.set_attn_processor(proc) + + +def _run(model, hidden, enc, ts, out_grad, ckpt): + if ckpt: + model.enable_gradient_checkpointing() + else: + model.disable_gradient_checkpointing() + inp = hidden.clone().requires_grad_(True) + out = model(hidden_states=inp, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] + out.backward(out_grad) + return out.detach(), inp.grad.detach() + + +def _report(name, a, b, rtol, ctol): + rel = (a.float() - b.float()).abs().max() / b.float().abs().max().clamp_min(1e-6) + cos = F.cosine_similarity(a.float().flatten(), b.float().flatten(), dim=0) + ok = rel < rtol and cos > ctol + if dist.get_rank() == 0: + print(f" [{'OK' if ok else 'FAIL'}] {name:28s} rel={rel:.2e} cos={1 - cos:.2e}(1-)") + assert ok, f"{name}: rel={rel:.3e} cos={cos:.5f}" + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--sp", type=int, default=4) + p.add_argument("--ulysses", type=int, default=0) + p.add_argument("--ring", type=int, default=0) + p.add_argument("--ckpt", action="store_true") + p.add_argument("--fwd-only", action="store_true") + p.add_argument("--fp32", action="store_true", help="fp32 + SDPA, isolates bf16 summation rounding") + cli = p.parse_args() + if cli.fp32: + global DTYPE + DTYPE = torch.float32 + + dist.init_process_group("nccl") + rank = dist.get_rank() + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.cuda.current_device() + init_gloo_group() + + args = argparse.Namespace( + sequence_parallel_size=cli.sp, + ulysses_degree=cli.ulysses, + ring_degree=cli.ring, + context_parallel_size=1, + ) + ps = create_fsdp_parallel_state(args) + + attn_weight_names = [ + f"blocks.{i}.attn1.{proj}.weight" for i in range(2) for proj in ("to_q", "to_k", "to_v", "to_out.0") + ] + attn_weight_names.append("proj_out.weight") + + model = build_model(device) + hidden, enc, ts, out_grad = make_inputs(device) + _set_ref_processor(model) + out_ref, gin_ref = _run(model, hidden, enc, ts, out_grad, cli.ckpt) + gw_ref = {n: dict(model.named_parameters())[n].grad.detach().clone() for n in attn_weight_names} + model.zero_grad(set_to_none=True) + + apply_sequence_parallel(model, ps) + out_sp, gin_sp = _run(model, hidden, enc, ts, out_grad, cli.ckpt) + + # Each rank backprops only its 1/sp of the tokens; sum across sp restores full grads. + dist.all_reduce(gin_sp, group=ps.sp_group) + if rank == 0: + print(f"[PARITY] sp={cli.sp} ulysses={ps.ulysses_degree} ring={ps.ring_degree} ckpt={cli.ckpt}") + _report("forward(out)", out_sp, out_ref, rtol=2e-2, ctol=0.9990) + _report("grad(input)", gin_sp, gin_ref, rtol=4e-2, ctol=0.9980) + if not cli.fwd_only: + params = dict(model.named_parameters()) + for n in attn_weight_names: + assert params[n].grad is not None, f"{n} grad is None — all-to-all backward did not propagate" + g = params[n].grad.detach().clone() + dist.all_reduce(g, group=ps.sp_group) + _report(f"grad({n})", g, gw_ref[n], rtol=5e-2, ctol=0.9950) + + dist.barrier() + if rank == 0: + print(f"[PARITY OK] sp={cli.sp} u={ps.ulysses_degree} r={ps.ring_degree} ckpt={cli.ckpt}") + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/sp/sp_grad_sync_parity.py b/tests/sp/sp_grad_sync_parity.py new file mode 100644 index 00000000..47f8a382 --- /dev/null +++ b/tests/sp/sp_grad_sync_parity.py @@ -0,0 +1,145 @@ +"""SP gradient sync parity under real FSDP2: full grads must match a single-process +full-sequence reference after reduce-scatter (dp) + cross-sp all-reduce(SUM). + +Runs dp1 x sp4 so FSDP local shards are the full parameters, isolating the sp +dimension. Also asserts model outputs are bitwise identical across sp ranks +(the gather-after-proj_out contract that keeps loss/log_prob code unchanged). + +Usage: torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py +""" + +import argparse + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from diffusers import WanTransformer3DModel +from torch.distributed.tensor import DTensor + +from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state +from miles.backends.fsdp_utils.sp_attention import ( + WanUSPAttnProcessor, + apply_sequence_parallel, + init_sp_backend, +) +from miles.utils.distributed_utils import init_gloo_group + +DTYPE = torch.bfloat16 + + +def build_model(device): + torch.manual_seed(0) + model = WanTransformer3DModel( + patch_size=(1, 2, 2), + num_attention_heads=8, + attention_head_dim=128, + in_channels=16, + out_channels=16, + text_dim=4096, + freq_dim=256, + ffn_dim=1024, + num_layers=2, + rope_max_seq_len=1024, + ).to(device=device, dtype=DTYPE) + model.train() + for p in model.parameters(): + dist.broadcast(p.data, src=0) + return model + + +def make_inputs(device): + g = torch.Generator(device=device).manual_seed(123) + hidden = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) + enc = torch.randn(1, 32, 4096, device=device, dtype=DTYPE, generator=g) + ts = torch.tensor([500], device=device) + out_grad = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) + for t in (hidden, enc, out_grad): + dist.broadcast(t, src=0) + return hidden, enc, ts, out_grad + + +def main(): + dist.init_process_group("nccl") + rank = dist.get_rank() + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.cuda.current_device() + init_gloo_group() + + args = argparse.Namespace(sequence_parallel_size=4, ulysses_degree=4, ring_degree=0, context_parallel_size=1) + ps = create_fsdp_parallel_state(args) + hidden, enc, ts, out_grad = make_inputs(device) + + # Full-sequence single-process reference. + from sglang.multimodal_gen.runtime.layers.attention.layer import USPAttention + + ref = build_model(device) + init_sp_backend(DTYPE) + proc = WanUSPAttnProcessor(ref.config.num_attention_heads, ref.config.attention_head_dim, DTYPE) + proc.usp_attn = USPAttention( + num_heads=ref.config.num_attention_heads, + head_size=ref.config.attention_head_dim, + causal=False, + skip_sequence_parallel=True, + ) + ref.set_attn_processor(proc) + out = ref(hidden_states=hidden, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] + out.backward(out_grad) + ref_grads = {n: p.grad.detach().clone() for n, p in ref.named_parameters()} + + # FSDP(dp1) + SP(sp4). + from torch.distributed.fsdp import fully_shard + + model = build_model(device) + for blk in model.blocks: + fully_shard(blk, mesh=ps.dp_mesh) + fully_shard(model, mesh=ps.dp_mesh) + apply_sequence_parallel(model, ps, compute_dtype=DTYPE) + + out = model(hidden_states=hidden, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] + + o32 = out.detach().float() + ref0 = o32.clone() + dist.broadcast(ref0, src=ps.dp_rank * ps.sp_size, group=ps.sp_group) + diff = (o32 - ref0).abs().max() + if rank == 0: + print(f"[OUTPUT] max abs diff across sp ranks = {diff.item():.2e} (must be 0)") + assert diff.item() == 0.0, "model outputs diverge across sp ranks" + + out.backward(out_grad) + + # Mirror actor._all_reduce_sp_grads. + for p in model.parameters(): + if p.grad is None: + continue + local = p.grad.to_local() if isinstance(p.grad, DTensor) else p.grad + f = local.float() + dist.all_reduce(f, group=ps.sp_group) + local.copy_(f.to(local.dtype)) + + fails = 0 + checked = 0 + for n, p in model.named_parameters(): + if p.grad is None: + continue + g = p.grad.to_local() if isinstance(p.grad, DTensor) else p.grad + r = ref_grads[n] + if g.shape != r.shape: + continue + rel = (g.float() - r.float()).abs().max() / r.float().abs().max().clamp_min(1e-6) + cos = F.cosine_similarity(g.float().flatten(), r.float().flatten(), dim=0) + checked += 1 + if not (rel < 5e-2 and cos > 0.99): + fails += 1 + if rank == 0: + print(f" [FAIL] {n:42s} rel={rel:.2e} cos={1 - cos:.2e}(1-)") + + dist.barrier() + if rank == 0: + print(f"[SP-GRAD-SYNC] checked={checked} fails={fails}") + assert fails == 0 + print("[SP-GRAD-SYNC OK] dp1xsp4 FSDP+SP full grads == full-sequence reference") + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/sp/sp_init_smoke.py b/tests/sp/sp_init_smoke.py new file mode 100644 index 00000000..40a5ecea --- /dev/null +++ b/tests/sp/sp_init_smoke.py @@ -0,0 +1,82 @@ +"""Multi-GPU SP init smoke test: NCCL group members must match the sp_mesh layout. + +Usage: torchrun --standalone --nproc_per_node=N tests/sp/sp_init_smoke.py --sp S [--ulysses U --ring R] +""" + +import argparse + +import torch +import torch.distributed as dist + +from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state +from miles.backends.fsdp_utils.sp_mesh import sp_subgroups +from miles.utils.distributed_utils import init_gloo_group + + +def _members(group): + n = dist.get_world_size(group) + t = torch.tensor([dist.get_rank()], device="cuda") + out = [torch.zeros_like(t) for _ in range(n)] + dist.all_gather(out, t, group=group) + return sorted(int(x.item()) for x in out) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--sp", type=int, default=2) + p.add_argument("--ulysses", type=int, default=0) + p.add_argument("--ring", type=int, default=0) + cli = p.parse_args() + + dist.init_process_group("nccl") + rank, world = dist.get_rank(), dist.get_world_size() + torch.cuda.set_device(rank % torch.cuda.device_count()) + init_gloo_group() + + args = argparse.Namespace( + sequence_parallel_size=cli.sp, + ulysses_degree=cli.ulysses, + ring_degree=cli.ring, + context_parallel_size=1, + ) + ps = create_fsdp_parallel_state(args) + + dp_size, sp_size, sp_groups, ulysses_groups, ring_groups = sp_subgroups(world, cli.sp, cli.ulysses, cli.ring) + assert ps.sp_size == sp_size and ps.dp_size == dp_size + assert dist.get_world_size(ps.sp_group) == sp_size + assert dist.get_world_size(ps.dp_group) == dp_size + + my_sp = next(g for g in sp_groups if rank in g) + assert _members(ps.sp_group) == my_sp, f"rank{rank} sp_group {_members(ps.sp_group)} != {my_sp}" + if sp_size > 1: + my_u = next(g for g in ulysses_groups if rank in g) + my_r = sorted(next(g for g in ring_groups if rank in g)) + assert _members(ps.ulysses_group) == my_u, f"rank{rank} ulysses {_members(ps.ulysses_group)} != {my_u}" + assert _members(ps.ring_group) == my_r, f"rank{rank} ring {_members(ps.ring_group)} != {my_r}" + + # USPAttention reads these through sglang's _SP coordinator. + from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + get_ring_parallel_world_size, + get_sequence_parallel_world_size, + get_sp_parallel_rank, + get_sp_world_size, + get_ulysses_parallel_world_size, + ) + + assert get_sp_world_size() == sp_size + assert get_sequence_parallel_world_size() == sp_size + assert get_sp_parallel_rank() == ps.sp_rank + assert get_ulysses_parallel_world_size() == ps.ulysses_degree + assert get_ring_parallel_world_size() == ps.ring_degree + + dist.barrier() + if rank == 0: + print( + f"[SP-INIT-SMOKE OK] world={world} dp={dp_size} sp={sp_size} " + f"ulysses={ps.ulysses_degree} ring={ps.ring_degree}" + ) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/sp/sp_weight_sync_parity.py b/tests/sp/sp_weight_sync_parity.py new file mode 100644 index 00000000..0ab932b0 --- /dev/null +++ b/tests/sp/sp_weight_sync_parity.py @@ -0,0 +1,117 @@ +"""Weight-sync parity under FSDP + SP: every rank's reconstructed full weights must +hash identically and match a single-process reference. + +Reconstruction uses the same redistribute([Replicate()]) path as update_weights; +the checksum is sglang's compute_weights_checksum (name + dtype + shape + bytes), +the same function the online verify uses. Also asserts shape/dtype sensitivity. + +Usage: + torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 2 --ulysses 2 + torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 4 + torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 --ring 2 +""" + +import argparse + +import torch +import torch.distributed as dist +from diffusers import WanTransformer3DModel +from torch.distributed.fsdp import fully_shard + +from sglang.multimodal_gen.runtime.loader.weight_utils import compute_weights_checksum + +from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state +from miles.utils.distributed_utils import init_gloo_group + +DTYPE = torch.bfloat16 + + +def build_model(device): + torch.manual_seed(0) + model = WanTransformer3DModel( + patch_size=(1, 2, 2), + num_attention_heads=8, + attention_head_dim=128, + in_channels=16, + out_channels=16, + text_dim=4096, + freq_dim=256, + ffn_dim=1024, + num_layers=2, + rope_max_seq_len=1024, + ).to(device=device, dtype=DTYPE) + for p in model.parameters(): + dist.broadcast(p.data, src=0) + for b in model.buffers(): + dist.broadcast(b.data, src=0) + return model + + +def full_state_pairs(model): + """Mirror update_weights' redistribute([Replicate()]).to_local() reconstruction.""" + from torch.distributed.tensor import DTensor, Replicate + + pairs = [] + for name, param in model.state_dict().items(): + param = param.cuda() + if isinstance(param, DTensor): + param = param.redistribute(placements=[Replicate()] * param.device_mesh.ndim).to_local() + pairs.append((name, param.detach().cpu().contiguous())) + return pairs + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--sp", type=int, default=2) + p.add_argument("--ulysses", type=int, default=2) + p.add_argument("--ring", type=int, default=0) + cli = p.parse_args() + + dist.init_process_group("nccl") + rank = dist.get_rank() + world = dist.get_world_size() + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.cuda.current_device() + init_gloo_group() + + args = argparse.Namespace( + sequence_parallel_size=cli.sp, + ulysses_degree=cli.ulysses, + ring_degree=cli.ring, + context_parallel_size=1, + ) + ps = create_fsdp_parallel_state(args) + + ref_model = build_model(device) + ref_sum = compute_weights_checksum([(n, pa.detach().cpu().contiguous()) for n, pa in ref_model.state_dict().items()]) + + model = build_model(device) + for blk in model.blocks: + fully_shard(blk, mesh=ps.dp_mesh) + fully_shard(model, mesh=ps.dp_mesh) + + my_sum = compute_weights_checksum(full_state_pairs(model)) + + gathered = [None] * world + dist.all_gather_object(gathered, my_sum) + + if rank == 0: + all_equal = all(s == gathered[0] for s in gathered) + match_ref = gathered[0] == ref_sum + print(f"[WEIGHT-SYNC] world={world} dp{ps.dp_size}xsp{ps.sp_size}(u{ps.ulysses_degree}r{ps.ring_degree})") + print(f"[WEIGHT-SYNC] all ranks equal={all_equal} == single-process ref={match_ref}") + assert all_equal, "reconstructed full weights differ across ranks" + assert match_ref, "reconstructed full weights != single-process reference" + + # reshape keeps bytes identical, so only the shape term can catch it + t = torch.arange(12, dtype=torch.float32).reshape(3, 4) + assert compute_weights_checksum([("w", t)]) != compute_weights_checksum([("w", t.reshape(4, 3))]) + assert compute_weights_checksum([("w", t)]) != compute_weights_checksum([("w", t.to(torch.float64))]) + print("[SP-WEIGHT-SYNC OK] bitwise-identical reconstruction + dtype/shape sensitivity") + + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/sp/test_sp_mesh.py b/tests/sp/test_sp_mesh.py new file mode 100644 index 00000000..942641e7 --- /dev/null +++ b/tests/sp/test_sp_mesh.py @@ -0,0 +1,87 @@ +"""SP rank layout / subgroup unit tests. Pure functions, no GPU.""" + +import pytest + +from miles.backends.fsdp_utils.sp_mesh import ( + locate_rank, + resolve_sp_degrees, + sp_subgroups, + validate_sp_config, +) + + +def test_resolve_auto_degrees(): + assert resolve_sp_degrees(4) == (4, 4, 1) + assert resolve_sp_degrees(4, ulysses_degree=2) == (4, 2, 2) + assert resolve_sp_degrees(8, 2, 4) == (8, 2, 4) + assert resolve_sp_degrees(1) == (1, 1, 1) + + +def test_resolve_illegal(): + with pytest.raises(ValueError): + resolve_sp_degrees(4, ulysses_degree=3) + + +def test_sglang_alignment_example(): + # Matches sglang set_seq_parallel_pg_by_sp_groups: sp=4,u=2,r=2 + _, _, sp_groups, ulysses_groups, ring_groups = sp_subgroups(4, 4, 2, 2) + assert sp_groups == [[0, 1, 2, 3]] + assert ulysses_groups == [[0, 1], [2, 3]] + assert ring_groups == [[0, 2], [1, 3]] + + +@pytest.mark.parametrize( + "world,sp,u,r", + [ + (2, 2, 2, 1), + (4, 2, 2, 1), + (4, 4, 2, 2), + (4, 4, 4, 1), + (8, 4, 2, 2), + (16, 8, 2, 4), + (64, 8, 8, 1), + (256, 16, 8, 2), + (1024, 8, 8, 1), + ], +) +def test_layout_invariants_any_scale(world, sp, u, r): + dp_size, sp_size, sp_groups, ulysses_groups, ring_groups = sp_subgroups(world, sp, u, r) + assert sp_size == sp == u * r + assert dp_size * sp_size == world + assert len(sp_groups) == dp_size and all(len(g) == sp for g in sp_groups) + # ulysses groups: contiguous, size u, exact cover + assert all(len(g) == u and g == list(range(g[0], g[0] + u)) for g in ulysses_groups) + assert sorted(x for g in ulysses_groups for x in g) == list(range(world)) + # ring groups: strided, size r, exact cover + assert all(len(g) == r for g in ring_groups) + assert sorted(x for g in ring_groups for x in g) == list(range(world)) + + +@pytest.mark.parametrize("world,sp,u,r", [(8, 4, 2, 2), (16, 8, 2, 4), (256, 16, 8, 2)]) +def test_locate_rank_consistent_with_subgroups(world, sp, u, r): + _, _, _, ulysses_groups, ring_groups = sp_subgroups(world, sp, u, r) + for rank in range(world): + dp_rank, sp_rank, u_rank, r_rank = locate_rank(rank, sp, u, r) + assert dp_rank == rank // sp and sp_rank == rank % sp + my_u_group = next(g for g in ulysses_groups if rank in g) + my_r_group = next(g for g in ring_groups if rank in g) + assert my_u_group[u_rank] == rank + assert my_r_group[r_rank] == rank + + +def test_validate_rejects_illegal(): + with pytest.raises(ValueError): + validate_sp_config(world_size=6, sequence_parallel_size=4) + with pytest.raises(ValueError): + validate_sp_config(world_size=8, sequence_parallel_size=4, ulysses_degree=3, ring_degree=1) + with pytest.raises(ValueError): + validate_sp_config(world_size=4, sequence_parallel_size=4, num_heads=40, ulysses_degree=3) + # Wan2.2 heads=40: ulysses in {2, 4} is legal + assert validate_sp_config(world_size=4, sequence_parallel_size=2, num_heads=40) == (2, 2, 1) + assert validate_sp_config(world_size=4, sequence_parallel_size=4, num_heads=40) == (4, 4, 1) + + +def test_num_heads_guard(): + with pytest.raises(ValueError): + validate_sp_config(world_size=6, sequence_parallel_size=3, num_heads=40) + assert validate_sp_config(world_size=8, sequence_parallel_size=8, num_heads=40) == (8, 8, 1) From 1822303284a4a22cc5b093e2998f8117d8c0e314 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 6 Jul 2026 22:52:35 +0000 Subject: [PATCH 05/24] diffusion: SP guardrail runner + fp32 mode for grad-sync parity Co-Authored-By: Claude Fable 5 --- tests/sp/run_gpu_suite.sh | 2 + tests/sp/run_rl_guardrail.sh | 99 +++++++++++++++++++++++++++++++++ tests/sp/sp_grad_sync_parity.py | 11 +++- 3 files changed, 111 insertions(+), 1 deletion(-) create mode 100755 tests/sp/run_rl_guardrail.sh diff --git a/tests/sp/run_gpu_suite.sh b/tests/sp/run_gpu_suite.sh index f6754fa8..8690f406 100755 --- a/tests/sp/run_gpu_suite.sh +++ b/tests/sp/run_gpu_suite.sh @@ -4,6 +4,8 @@ set -euo pipefail cd "$(dirname "$0")/../.." export PYTHONPATH=. +torchrun() { python -m torch.distributed.run "$@"; } + run() { echo; echo "===== $* ====="; "$@"; } run python -m pytest -q tests/sp/sglang_usp_import_guard.py diff --git a/tests/sp/run_rl_guardrail.sh b/tests/sp/run_rl_guardrail.sh new file mode 100755 index 00000000..6845c03b --- /dev/null +++ b/tests/sp/run_rl_guardrail.sh @@ -0,0 +1,99 @@ +#!/usr/bin/env bash +# Real-RL numeric guardrail: run the Wan2.2 PickScore recipe for a few steps on +# 4 GPUs (train+rollout colocate, PickScore reward on CPU) and compare train +# metrics between bands. With identical seeds the rollout data is identical, so +# step-1 adv_abs_mean / log_prob_old_idx_0 must match bitwise across bands; +# log_prob_new / loss / grad_norm may differ at bf16 summation level. +# +# Usage: +# bash tests/sp/run_rl_guardrail.sh # FSDP dp4 baseline +# SP_SIZE=2 ULYSSES_DEGREE=2 bash tests/sp/run_rl_guardrail.sh # dp2 x sp2 +set -euo pipefail +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" + +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False +export PICKSCORE_NUM_GPUS_PER_WORKER=0 + +PYTHON_BIN="${PYTHON_BIN:-python}" +SP_SIZE="${SP_SIZE:-1}" +ULYSSES_DEGREE="${ULYSSES_DEGREE:-0}" +RING_DEGREE="${RING_DEGREE:-0}" +NUM_ROLLOUT="${NUM_ROLLOUT:-2}" +NUM_FRAMES="${NUM_FRAMES:-5}" +GRAD_CKPT="${GRAD_CKPT:-0}" +RUN_NAME="${RUN_NAME:-sp_guardrail_sp${SP_SIZE}_$(date +%Y%m%d_%H%M%S)}" +DATASETS_DIR="/root/datasets/miles-diffusion-datasets" + +EXTRA_ARGS=() +[[ "${GRAD_CKPT}" == "1" ]] && EXTRA_ARGS+=(--gradient-checkpointing) + +WAN_LORA_TARGET_MODULES=( + attn1.to_q attn1.to_k attn1.to_v attn1.to_out.0 + attn2.to_q attn2.to_k attn2.to_v attn2.to_out.0 + ffn.net.0.proj ffn.net.2 +) + +"${PYTHON_BIN}" -u "${ROOT_DIR}/train_diffusion.py" \ + --train-backend fsdp \ + --rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout \ + --hf-checkpoint Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --diffusion-model Wan-AI/Wan2.2-T2V-A14B-Diffusers \ + --prompt-data "${DATASETS_DIR}/flowgrpo_pickscore/train.jsonl" \ + --input-key input \ + --rollout-batch-size 8 \ + --n-samples-per-prompt 8 \ + --num-rollout "${NUM_ROLLOUT}" \ + --num-steps-per-rollout 2 \ + --diffusion-microgroup-size 8 \ + --micro-batch-size 2 \ + --actor-num-gpus-per-node 4 \ + --sequence-parallel-size "${SP_SIZE}" \ + --ulysses-degree "${ULYSSES_DEGREE}" \ + --ring-degree "${RING_DEGREE}" \ + --rollout-num-gpus 4 \ + --rollout-num-gpus-per-engine 1 \ + --num-gpus-per-node 4 \ + --colocate \ + --use-lora \ + --lora-rank 64 \ + --lora-alpha 128 \ + --lora-target-modules "${WAN_LORA_TARGET_MODULES[@]}" \ + --diffusion-init-lora-weight gaussian \ + --lr 1e-4 \ + --adam-beta2 0.999 \ + --diffusion-clip-range 1e-4 \ + --weight-decay 1e-4 \ + --use-miles-router \ + --sglang-server-concurrency 8 \ + --update-weight-buffer-size 2147483648 \ + --update-weight-target-module transformer,transformer_2 \ + --diffusion-reward pickscore:1.0 \ + --advantage-estimator grpo \ + --rm-type pickscore \ + --pickscore-num-workers 1 \ + --pickscore-num-gpus-per-worker 0 \ + --pickscore-batch-size 8 \ + --pickscore-processor-path laion/CLIP-ViT-H-14-laion2B-s32B-b79K \ + --pickscore-model-path yuvalkirstain/PickScore_v1 \ + --fsdp-master-dtype fp32 \ + --fsdp-reduce-dtype fp32 \ + --diffusion-forward-dtype bf16 \ + --diffusion-num-steps 10 \ + --diffusion-eval-num-steps 28 \ + --diffusion-output-num-frames "${NUM_FRAMES}" \ + --diffusion-guidance-scale 4.0 \ + --diffusion-guidance-scale-2 3.0 \ + --diffusion-noise-level 0.9 \ + --diffusion-height 480 \ + --diffusion-width 480 \ + --diffusion-flow-shift 3.0 \ + --diffusion-step-strategy-path miles.rollout.step_strategy_hub.epoch_global_window \ + --diffusion-sde-window-size 1 \ + --diffusion-sde-candidate-steps 1,2,3 \ + --diffusion-debug-mode \ + --save "${ROOT_DIR}/logs/${RUN_NAME}/ckpt" \ + --save-interval 100 \ + --skip-eval-before-train \ + "${EXTRA_ARGS[@]}" \ + 2>&1 | tee "${ROOT_DIR}/logs/${RUN_NAME}.log" diff --git a/tests/sp/sp_grad_sync_parity.py b/tests/sp/sp_grad_sync_parity.py index 47f8a382..d7530529 100644 --- a/tests/sp/sp_grad_sync_parity.py +++ b/tests/sp/sp_grad_sync_parity.py @@ -59,6 +59,13 @@ def make_inputs(device): def main(): + p = argparse.ArgumentParser() + p.add_argument("--fp32", action="store_true", help="fp32 + SDPA, isolates bf16 summation rounding") + cli = p.parse_args() + if cli.fp32: + global DTYPE + DTYPE = torch.float32 + dist.init_process_group("nccl") rank = dist.get_rank() torch.cuda.set_device(rank % torch.cuda.device_count()) @@ -128,7 +135,9 @@ def main(): rel = (g.float() - r.float()).abs().max() / r.float().abs().max().clamp_min(1e-6) cos = F.cosine_similarity(g.float().flatten(), r.float().flatten(), dim=0) checked += 1 - if not (rel < 5e-2 and cos > 0.99): + # Secondary band: small-norm tensors (biases) sit right at the bf16 + # summation noise floor; --fp32 confirms exact agreement there. + if not (rel < 5e-2 and cos > 0.99) and not (rel < 1e-1 and cos > 0.999): fails += 1 if rank == 0: print(f" [FAIL] {n:42s} rel={rel:.2e} cos={1 - cos:.2e}(1-)") From 8040abb57ddcd2792709b528415f9a7205fb7273 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Tue, 7 Jul 2026 01:57:43 +0000 Subject: [PATCH 06/24] diffusion: guardrail runner cleanup preamble + expert-selection override Co-Authored-By: Claude Fable 5 --- tests/sp/run_rl_guardrail.sh | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/sp/run_rl_guardrail.sh b/tests/sp/run_rl_guardrail.sh index 6845c03b..b91e2ed9 100755 --- a/tests/sp/run_rl_guardrail.sh +++ b/tests/sp/run_rl_guardrail.sh @@ -11,6 +11,10 @@ set -euo pipefail ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +# A finished run can leave ray actors holding GPU memory; start clean. +ray stop --force >/dev/null 2>&1 || true +sleep 5 + export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False export PICKSCORE_NUM_GPUS_PER_WORKER=0 @@ -22,6 +26,7 @@ RING_DEGREE="${RING_DEGREE:-0}" NUM_ROLLOUT="${NUM_ROLLOUT:-2}" NUM_FRAMES="${NUM_FRAMES:-5}" GRAD_CKPT="${GRAD_CKPT:-0}" +CANDIDATE_STEPS="${CANDIDATE_STEPS:-1,2,3}" RUN_NAME="${RUN_NAME:-sp_guardrail_sp${SP_SIZE}_$(date +%Y%m%d_%H%M%S)}" DATASETS_DIR="/root/datasets/miles-diffusion-datasets" @@ -90,7 +95,7 @@ WAN_LORA_TARGET_MODULES=( --diffusion-flow-shift 3.0 \ --diffusion-step-strategy-path miles.rollout.step_strategy_hub.epoch_global_window \ --diffusion-sde-window-size 1 \ - --diffusion-sde-candidate-steps 1,2,3 \ + --diffusion-sde-candidate-steps "${CANDIDATE_STEPS}" \ --diffusion-debug-mode \ --save "${ROOT_DIR}/logs/${RUN_NAME}/ckpt" \ --save-interval 100 \ From c199e22e6998569fe0d3892509a76b3df9445fa9 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Wed, 8 Jul 2026 01:11:34 +0000 Subject: [PATCH 07/24] diffusion: drive SP shard/gather hooks from the model's diffusers _cp_plan The shard/gather placement (rope outputs, first-block input, proj_out output) was hard-coded to Wan's module topology; diffusers 0.37 declares the same contract per model as _cp_plan, so consume that instead. The attention processor remains Wan-specific and is now guarded explicitly. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/sp_attention.py | 137 +++++++++++++++------- 1 file changed, 97 insertions(+), 40 deletions(-) diff --git a/miles/backends/fsdp_utils/sp_attention.py b/miles/backends/fsdp_utils/sp_attention.py index ad4ee9db..0fc8408b 100644 --- a/miles/backends/fsdp_utils/sp_attention.py +++ b/miles/backends/fsdp_utils/sp_attention.py @@ -1,46 +1,48 @@ -"""Sequence parallelism for diffusers Wan DiT: self-attention runs sglang's USPAttention. +"""Sequence parallelism for diffusers DiTs: self-attention runs sglang's USPAttention. Each sp rank holds S/sp latent tokens; attention internally gathers to the full -sequence via Ulysses all-to-all (+Ring) and scatters back. The sequence is sharded -before the first block and gathered after proj_out, so every parameter sees a -partial gradient and the actor's cross-sp grad all-reduce applies uniformly. +sequence via Ulysses all-to-all (+Ring) and scatters back. Shard/gather points are +taken from the model's diffusers ``_cp_plan``, so model outputs stay full-sequence +and every parameter sees a partial gradient; the actor's cross-sp grad all-reduce +applies uniformly. """ -import functools +import inspect import torch import torch.distributed as dist +from diffusers.models._modeling_parallel import ContextParallelInput, ContextParallelOutput class _GatherSequence(torch.autograd.Function): - """All-gather S_local shards along dim=1; backward returns each rank's slice.""" + """All-gather local shards along dim; backward returns each rank's slice.""" @staticmethod - def forward(ctx, x, group, sp_rank, sp_size): + def forward(ctx, x, group, sp_rank, sp_size, dim): ctx.sp_rank = sp_rank - ctx.s_local = x.shape[1] + 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=1) + return torch.cat(parts, dim=dim) @staticmethod def backward(ctx, grad): - start = ctx.sp_rank * ctx.s_local - return grad[:, start : start + ctx.s_local], None, None, None + start = ctx.sp_rank * ctx.local_size + return grad.narrow(ctx.dim, start, ctx.local_size), None, None, None, None -def shard_sequence(x, parallel_state): +def shard_sequence(x, parallel_state, dim=1): sp = parallel_state.sp_size - s = x.shape[1] + s = x.shape[dim] if s % sp: raise ValueError(f"sequence length {s} is not divisible by sp_size {sp}") s_local = s // sp - start = parallel_state.sp_rank * s_local - return x[:, start : start + s_local] + return x.narrow(dim, parallel_state.sp_rank * s_local, s_local) -def gather_sequence(x, parallel_state): - return _GatherSequence.apply(x, parallel_state.sp_group, parallel_state.sp_rank, parallel_state.sp_size) +def gather_sequence(x, parallel_state, dim=1): + return _GatherSequence.apply(x, parallel_state.sp_group, parallel_state.sp_rank, parallel_state.sp_size, dim) class WanUSPAttnProcessor: @@ -150,9 +152,84 @@ def _apply_rotary_emb(hidden_states, freqs_cos, freqs_sin): return out.type_as(hidden_states) +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, 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_cp_plan_hooks(transformer, parallel_state): + """Install shard/gather hooks from the model's diffusers ``_cp_plan``. + + ContextParallelInput entries split module inputs (or outputs when + split_output=True); ContextParallelOutput entries gather module outputs. + The gather is a differentiable all-gather, so backward stays partial-grad. + """ + for path, spec in transformer._cp_plan.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, 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: + param_names = list(inspect.signature(module.forward).parameters) + + 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, compute_dtype=None): - """Wire SP into one Wan transformer: replace self-attn processors and install - the shard/gather hooks. Call once per transformer after FSDP wrapping.""" + """Wire SP into one transformer: replace self-attn processors and install + the shard/gather hooks declared by its _cp_plan. Call once per transformer + after FSDP wrapping.""" + from diffusers import WanTransformer3DModel + + base = transformer.get_base_model() if hasattr(transformer, "get_base_model") else transformer + if not isinstance(base, WanTransformer3DModel): + raise ValueError( + f"SP attention processor currently supports WanTransformer3DModel only, got {base.__class__.__name__}" + ) heads = transformer.config.num_attention_heads head_dim = transformer.config.attention_head_dim # args carry no head count, so the startup divisibility check must happen @@ -164,24 +241,4 @@ def apply_sequence_parallel(transformer, parallel_state, compute_dtype=None): if compute_dtype is None: compute_dtype = next(transformer.parameters()).dtype transformer.set_attn_processor(WanUSPAttnProcessor(heads, head_dim, compute_dtype)) - - # RoPE runs once per forward and its output is reused by every block; shard at the source. - rope = transformer.rope - orig_rope_forward = rope.forward - - @functools.wraps(orig_rope_forward) - def sp_rope_forward(hidden_states): - cos, sin = orig_rope_forward(hidden_states) - return shard_sequence(cos, parallel_state), shard_sequence(sin, parallel_state) - - rope.forward = sp_rope_forward - - def slice_block_input(module, args): - hs, *rest = args - return (shard_sequence(hs, parallel_state), *rest) - - def gather_proj_output(module, args, output): - return gather_sequence(output, parallel_state) - - transformer.blocks[0].register_forward_pre_hook(slice_block_input) - transformer.proj_out.register_forward_hook(gather_proj_output) + _install_cp_plan_hooks(transformer, parallel_state) From f639c4e6f533affae33af1fbce1e4221bc9085da Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Wed, 8 Jul 2026 08:18:08 +0000 Subject: [PATCH 08/24] diffusion: guardrail runner supports rollout-data save/replay for same-data band comparison Co-Authored-By: Claude Fable 5 --- tests/sp/run_rl_guardrail.sh | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/sp/run_rl_guardrail.sh b/tests/sp/run_rl_guardrail.sh index b91e2ed9..32ce4701 100755 --- a/tests/sp/run_rl_guardrail.sh +++ b/tests/sp/run_rl_guardrail.sh @@ -27,11 +27,17 @@ NUM_ROLLOUT="${NUM_ROLLOUT:-2}" NUM_FRAMES="${NUM_FRAMES:-5}" GRAD_CKPT="${GRAD_CKPT:-0}" CANDIDATE_STEPS="${CANDIDATE_STEPS:-1,2,3}" +# Same-data comparison across bands: save rollout data in one band, replay it +# in another so both trainers see literally identical samples. +SAVE_ROLLOUT_DATA="${SAVE_ROLLOUT_DATA:-}" +LOAD_ROLLOUT_DATA="${LOAD_ROLLOUT_DATA:-}" RUN_NAME="${RUN_NAME:-sp_guardrail_sp${SP_SIZE}_$(date +%Y%m%d_%H%M%S)}" DATASETS_DIR="/root/datasets/miles-diffusion-datasets" EXTRA_ARGS=() [[ "${GRAD_CKPT}" == "1" ]] && EXTRA_ARGS+=(--gradient-checkpointing) +[[ -n "${SAVE_ROLLOUT_DATA}" ]] && EXTRA_ARGS+=(--save-debug-rollout-data "${SAVE_ROLLOUT_DATA}") +[[ -n "${LOAD_ROLLOUT_DATA}" ]] && EXTRA_ARGS+=(--load-debug-rollout-data "${LOAD_ROLLOUT_DATA}") WAN_LORA_TARGET_MODULES=( attn1.to_q attn1.to_k attn1.to_v attn1.to_out.0 From 2460eb438d277adf63b768bf72bf2c77da821d6d Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Thu, 9 Jul 2026 02:43:50 +0000 Subject: [PATCH 09/24] diffusion: own the differentiable USP operators, drop the sglang training dependency RL needs a differentiable attention in the trainer, not in the rollout engine; importing sglang's USPAttention leaked that requirement into a fork the engine had to carry. sp_ops.py holds the Ulysses all-to-all, ring attention (torch ring templates), and SDPA local attention with the same layout, so training numerics are unchanged and sglang needs no patches. The trainer's only remaining sglang import is the weight-sync checksum. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/actor.py | 2 +- miles/backends/fsdp_utils/parallel.py | 31 ++--- miles/backends/fsdp_utils/sp_attention.py | 46 ++----- miles/backends/fsdp_utils/sp_ops.py | 150 ++++++++++++++++++++++ tests/sp/run_rl_guardrail.sh | 4 +- tests/sp/sglang_usp_import_guard.py | 25 +--- tests/sp/sp_attention_parity.py | 23 +--- tests/sp/sp_grad_sync_parity.py | 22 +--- tests/sp/sp_init_smoke.py | 15 --- 9 files changed, 192 insertions(+), 126 deletions(-) create mode 100644 miles/backends/fsdp_utils/sp_ops.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index f0b12cab..56a6c88f 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -141,7 +141,7 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty if self.parallel_state.sp_size > 1: for model in self.models.values(): - apply_sequence_parallel(model, self.parallel_state, compute_dtype=self._forward_dtype) + apply_sequence_parallel(model, self.parallel_state) # Force a sync to ensure sharding is complete and old memory is freed. torch.cuda.synchronize() diff --git a/miles/backends/fsdp_utils/parallel.py b/miles/backends/fsdp_utils/parallel.py index 3389fb2d..e5ad587a 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/miles/backends/fsdp_utils/parallel.py @@ -1,7 +1,6 @@ import logging from argparse import Namespace -import torch import torch.distributed as dist from torch.distributed.device_mesh import init_device_mesh @@ -40,28 +39,18 @@ def create_fsdp_parallel_state(args: Namespace) -> ParallelState: f"dp_rank={dp_rank} sp_rank={sp_rank}" ) - # USPAttention reads sglang's module-global _SP coordinator; - # set_seq_parallel_pg_by_sp_groups only sets ULYSSES_PG/RING_PG. + # dist.new_group is collective: every rank must create every group. ulysses_group = ring_group = None if sp_size > 1: - from sglang.multimodal_gen.runtime.distributed import parallel_state as sgl_sp - from sglang.multimodal_gen.runtime.distributed.parallel_groups import ( - PROCESS_GROUP, - set_seq_parallel_pg_by_sp_groups, - ) - - _, _, sp_groups, _, _ = sp_subgroups(world_size, sp_size, ulysses_degree, ring_degree) - set_seq_parallel_pg_by_sp_groups(ulysses_degree, ring_degree, rank, sp_groups) - ulysses_group = PROCESS_GROUP.ULYSSES_PG - ring_group = PROCESS_GROUP.RING_PG - sgl_sp._SP = sgl_sp.init_parallel_group_coordinator( - group_ranks=sp_groups, - local_rank=torch.cuda.current_device(), - backend=dist.get_backend(), - parallel_mode="sequence", - ulysses_group=ulysses_group, - ring_group=ring_group, - ) + _, _, _, ulysses_groups, ring_groups = sp_subgroups(world_size, sp_size, ulysses_degree, ring_degree) + for ranks in ulysses_groups: + group = dist.new_group(ranks) + if rank in ranks: + ulysses_group = group + for ranks in ring_groups: + group = dist.new_group(ranks) + if rank in ranks: + ring_group = group parallel_state = ParallelState( dp_rank=dp_rank, diff --git a/miles/backends/fsdp_utils/sp_attention.py b/miles/backends/fsdp_utils/sp_attention.py index 0fc8408b..db56a7ad 100644 --- a/miles/backends/fsdp_utils/sp_attention.py +++ b/miles/backends/fsdp_utils/sp_attention.py @@ -1,4 +1,4 @@ -"""Sequence parallelism for diffusers DiTs: self-attention runs sglang's USPAttention. +"""Sequence parallelism for diffusers DiTs: self-attention runs USP (sp_ops). Each sp rank holds S/sp latent tokens; attention internally gathers to the full sequence via Ulysses all-to-all (+Ring) and scatters back. Shard/gather points are @@ -13,6 +13,8 @@ import torch.distributed as dist from diffusers.models._modeling_parallel import ContextParallelInput, ContextParallelOutput +from .sp_ops import usp_attention + class _GatherSequence(torch.autograd.Function): """All-gather local shards along dim; backward returns each rank's slice.""" @@ -46,16 +48,15 @@ def gather_sequence(x, parallel_state, dim=1): class WanUSPAttnProcessor: - """Wan attention processor: self-attn via USPAttention, cross-attn via local SDPA. + """Wan attention processor: self-attn via USP, cross-attn via local SDPA. Reuses Wan's QKV/RMSNorm/RoPE; rotary_emb arrives pre-sharded to S_local. + With parallel_state=None the self-attn falls back to plain SDPA (reference mode). """ - def __init__(self, num_heads, head_dim, compute_dtype=torch.bfloat16): - from sglang.multimodal_gen.runtime.layers.attention.layer import USPAttention - - init_sp_backend(compute_dtype) - self.usp_attn = USPAttention(num_heads=num_heads, head_size=head_dim, causal=False) + def __init__(self, parallel_state=None): + self.ulysses_group = parallel_state.ulysses_group if parallel_state is not None else None + self.ring_group = parallel_state.ring_group if parallel_state is not None else None def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, rotary_emb=None): is_self_attn = encoder_hidden_states is None @@ -84,7 +85,7 @@ def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_ma key = _apply_rotary_emb(key, *rotary_emb) if is_self_attn: - hidden_states = self.usp_attn(query, key, value) + hidden_states = usp_attention(query, key, value, self.ulysses_group, self.ring_group) else: hidden_states = torch.nn.functional.scaled_dot_product_attention( query.transpose(1, 2), @@ -120,28 +121,6 @@ def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_ma return hidden_states -def init_sp_backend(compute_dtype): - """Set up the sglang runtime state USPAttention depends on. - - The training process never goes through the sglang server launch, so the - attention backend, mixed-precision policy, and forward context are unset. - FA requires fp16/bf16; the forward context is persisted module-globally - because checkpoint recompute runs after any with-scope has exited. - """ - from sglang.multimodal_gen.runtime.layers.attention.selector import global_force_attn_backend - from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum - from sglang.multimodal_gen.utils import set_mixed_precision_policy - - half = compute_dtype in (torch.float16, torch.bfloat16) - global_force_attn_backend(AttentionBackendEnum.FA if half else AttentionBackendEnum.TORCH_SDPA) - set_mixed_precision_policy(param_dtype=compute_dtype, reduce_dtype=torch.float32) - - from sglang.multimodal_gen.runtime.managers import forward_context as fc - - if fc._forward_context is None: - fc._forward_context = fc.ForwardContext(current_timestep=0, attn_metadata=None) - - def _apply_rotary_emb(hidden_states, freqs_cos, freqs_sin): x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) cos = freqs_cos[..., 0::2] @@ -219,7 +198,7 @@ def split_outputs(mod, args, kwargs, output, _specs=output_specs): module.register_forward_hook(split_outputs, with_kwargs=True) -def apply_sequence_parallel(transformer, parallel_state, compute_dtype=None): +def apply_sequence_parallel(transformer, parallel_state): """Wire SP into one transformer: replace self-attn processors and install the shard/gather hooks declared by its _cp_plan. Call once per transformer after FSDP wrapping.""" @@ -231,14 +210,11 @@ def apply_sequence_parallel(transformer, parallel_state, compute_dtype=None): f"SP attention processor currently supports WanTransformer3DModel only, got {base.__class__.__name__}" ) heads = transformer.config.num_attention_heads - head_dim = transformer.config.attention_head_dim # args carry no head count, so the startup divisibility check must happen # here where the real model config is available. if heads % parallel_state.ulysses_degree != 0: raise ValueError( f"num_attention_heads({heads}) is not divisible by ulysses_degree({parallel_state.ulysses_degree})" ) - if compute_dtype is None: - compute_dtype = next(transformer.parameters()).dtype - transformer.set_attn_processor(WanUSPAttnProcessor(heads, head_dim, compute_dtype)) + transformer.set_attn_processor(WanUSPAttnProcessor(parallel_state)) _install_cp_plan_hooks(transformer, parallel_state) diff --git a/miles/backends/fsdp_utils/sp_ops.py b/miles/backends/fsdp_utils/sp_ops.py new file mode 100644 index 00000000..60fcd972 --- /dev/null +++ b/miles/backends/fsdp_utils/sp_ops.py @@ -0,0 +1,150 @@ +"""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 _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 and dist.get_world_size(ring_group) > 1: + 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/tests/sp/run_rl_guardrail.sh b/tests/sp/run_rl_guardrail.sh index 32ce4701..78d8f256 100755 --- a/tests/sp/run_rl_guardrail.sh +++ b/tests/sp/run_rl_guardrail.sh @@ -11,8 +11,10 @@ set -euo pipefail ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" -# A finished run can leave ray actors holding GPU memory; start clean. +# A finished run can leave ray actors and sglang engine schedulers holding GPU +# memory; start clean. ray stop --force >/dev/null 2>&1 || true +pkill -9 -f "sgl_diffusion::" 2>/dev/null || true sleep 5 export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" diff --git a/tests/sp/sglang_usp_import_guard.py b/tests/sp/sglang_usp_import_guard.py index 708841e3..b51df88b 100644 --- a/tests/sp/sglang_usp_import_guard.py +++ b/tests/sp/sglang_usp_import_guard.py @@ -1,8 +1,8 @@ -"""Guard: the sglang tree actually imported must contain the training patches. +"""Guard: the sglang tree actually imported must have the dtype/shape-aware checksum. -Multiple sglang checkouts can coexist on a machine; training must resolve to the -one with the USP autograd wrappers and the dtype/shape-aware checksum, otherwise -Ulysses backward silently drops gradients. Run before any GPU parity script. +Training's only remaining sglang import is compute_weights_checksum (weight-sync +verify); a bytes-only checksum would miss transpose/reshape corruption. The USP +attention operators are miles-owned (sp_ops.py) and need no sglang patches. Usage: PYTHONPATH=. python -m pytest tests/sp/sglang_usp_import_guard.py """ @@ -10,18 +10,6 @@ import inspect -def test_usp_has_training_autograd(): - from sglang.multimodal_gen.runtime.layers import usp - - assert hasattr(usp, "_AllToAllSingle"), ( - f"imported sglang lacks _AllToAllSingle: {usp.__file__} — " - "Ulysses training backward would silently drop q/k/v grads" - ) - assert hasattr(usp, "_RingFlashAttention"), ( - f"imported sglang lacks _RingFlashAttention: {usp.__file__} — ring training dK/dV would be wrong" - ) - - def test_checksum_covers_dtype_shape(): from sglang.multimodal_gen.runtime.loader import weight_utils @@ -33,8 +21,7 @@ def test_checksum_covers_dtype_shape(): if __name__ == "__main__": - test_usp_has_training_autograd() test_checksum_covers_dtype_shape() - from sglang.multimodal_gen.runtime.layers import usp + from sglang.multimodal_gen.runtime.loader import weight_utils - print(f"[GUARD OK] imported sglang = {usp.__file__}") + print(f"[GUARD OK] imported sglang = {weight_utils.__file__}") diff --git a/tests/sp/sp_attention_parity.py b/tests/sp/sp_attention_parity.py index a5c27e2a..5ccf70ee 100644 --- a/tests/sp/sp_attention_parity.py +++ b/tests/sp/sp_attention_parity.py @@ -1,10 +1,9 @@ -"""SP parity on a small real Wan DiT: USPAttention + shard/gather vs full-sequence reference. +"""SP parity on a small real Wan DiT: USP attention + shard/gather vs full-sequence reference. -Both paths use the same attention kernel (the reference runs USPAttention with -skip_sequence_parallel), so any diff comes from the SP collectives' float -summation order and must stay within bf16 tolerance. Checks forward output, -input grad, and per-block self-attn + proj_out weight grads, with and without -gradient checkpointing. +Both paths use the same local attention kernel (SDPA), so any diff comes from +the SP collectives' float summation order and must stay within bf16 tolerance. +Checks forward output, input grad, and per-block self-attn + proj_out weight +grads, with and without gradient checkpointing. Usage: torchrun --standalone --nproc_per_node=N tests/sp/sp_attention_parity.py \ --sp S [--ulysses U --ring R] [--ckpt] [--fp32] @@ -57,16 +56,8 @@ def make_inputs(device): def _set_ref_processor(model): - from sglang.multimodal_gen.runtime.layers.attention.layer import USPAttention - - proc = WanUSPAttnProcessor(model.config.num_attention_heads, model.config.attention_head_dim, DTYPE) - proc.usp_attn = USPAttention( - num_heads=model.config.num_attention_heads, - head_size=model.config.attention_head_dim, - causal=False, - skip_sequence_parallel=True, - ) - model.set_attn_processor(proc) + # parallel_state=None: same processor, plain SDPA self-attention. + model.set_attn_processor(WanUSPAttnProcessor(None)) def _run(model, hidden, enc, ts, out_grad, ckpt): diff --git a/tests/sp/sp_grad_sync_parity.py b/tests/sp/sp_grad_sync_parity.py index d7530529..fb291cff 100644 --- a/tests/sp/sp_grad_sync_parity.py +++ b/tests/sp/sp_grad_sync_parity.py @@ -17,11 +17,7 @@ from torch.distributed.tensor import DTensor from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state -from miles.backends.fsdp_utils.sp_attention import ( - WanUSPAttnProcessor, - apply_sequence_parallel, - init_sp_backend, -) +from miles.backends.fsdp_utils.sp_attention import WanUSPAttnProcessor, apply_sequence_parallel from miles.utils.distributed_utils import init_gloo_group DTYPE = torch.bfloat16 @@ -76,19 +72,9 @@ def main(): ps = create_fsdp_parallel_state(args) hidden, enc, ts, out_grad = make_inputs(device) - # Full-sequence single-process reference. - from sglang.multimodal_gen.runtime.layers.attention.layer import USPAttention - + # Full-sequence single-process reference (plain SDPA self-attention). ref = build_model(device) - init_sp_backend(DTYPE) - proc = WanUSPAttnProcessor(ref.config.num_attention_heads, ref.config.attention_head_dim, DTYPE) - proc.usp_attn = USPAttention( - num_heads=ref.config.num_attention_heads, - head_size=ref.config.attention_head_dim, - causal=False, - skip_sequence_parallel=True, - ) - ref.set_attn_processor(proc) + ref.set_attn_processor(WanUSPAttnProcessor(None)) out = ref(hidden_states=hidden, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] out.backward(out_grad) ref_grads = {n: p.grad.detach().clone() for n, p in ref.named_parameters()} @@ -100,7 +86,7 @@ def main(): for blk in model.blocks: fully_shard(blk, mesh=ps.dp_mesh) fully_shard(model, mesh=ps.dp_mesh) - apply_sequence_parallel(model, ps, compute_dtype=DTYPE) + apply_sequence_parallel(model, ps) out = model(hidden_states=hidden, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] diff --git a/tests/sp/sp_init_smoke.py b/tests/sp/sp_init_smoke.py index 40a5ecea..25e84439 100644 --- a/tests/sp/sp_init_smoke.py +++ b/tests/sp/sp_init_smoke.py @@ -54,21 +54,6 @@ def main(): assert _members(ps.ulysses_group) == my_u, f"rank{rank} ulysses {_members(ps.ulysses_group)} != {my_u}" assert _members(ps.ring_group) == my_r, f"rank{rank} ring {_members(ps.ring_group)} != {my_r}" - # USPAttention reads these through sglang's _SP coordinator. - from sglang.multimodal_gen.runtime.distributed.parallel_state import ( - get_ring_parallel_world_size, - get_sequence_parallel_world_size, - get_sp_parallel_rank, - get_sp_world_size, - get_ulysses_parallel_world_size, - ) - - assert get_sp_world_size() == sp_size - assert get_sequence_parallel_world_size() == sp_size - assert get_sp_parallel_rank() == ps.sp_rank - assert get_ulysses_parallel_world_size() == ps.ulysses_degree - assert get_ring_parallel_world_size() == ps.ring_degree - dist.barrier() if rank == 0: print( From 55fcd69bc3b38af48e56c6948a9eaf270aaeba43 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Thu, 9 Jul 2026 07:56:31 +0000 Subject: [PATCH 10/24] diffusion: first-principles SP cleanup - drop context_parallel_size and its CP->SP aliasing: the old assert==1 meant no config with cp>1 could ever have existed - drop the num_heads param from validate_sp_config and the always-None getattr feeding it: the real guard lives in apply_sequence_parallel - degree-1 ulysses/ring dimensions get no process groups (None = local), and the unused per-rank singleton tp_group is gone - move the sequence shard/gather collectives into sp_ops so all differentiable SP collectives live in one module - remove the dead --fwd-only flag (ring backward works) Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/arguments.py | 4 +-- miles/backends/fsdp_utils/parallel.py | 27 ++++++++------- miles/backends/fsdp_utils/sp_attention.py | 40 +++-------------------- miles/backends/fsdp_utils/sp_mesh.py | 10 +++--- miles/backends/fsdp_utils/sp_ops.py | 32 +++++++++++++++++- miles/utils/arguments.py | 3 -- tests/sp/sp_attention_parity.py | 15 ++++----- tests/sp/sp_grad_sync_parity.py | 2 +- tests/sp/sp_init_smoke.py | 10 ++++-- tests/sp/sp_weight_sync_parity.py | 1 - tests/sp/test_sp_mesh.py | 13 ++------ 11 files changed, 72 insertions(+), 85 deletions(-) diff --git a/miles/backends/fsdp_utils/arguments.py b/miles/backends/fsdp_utils/arguments.py index b18a8e5d..a5ecd797 100644 --- a/miles/backends/fsdp_utils/arguments.py +++ b/miles/backends/fsdp_utils/arguments.py @@ -57,9 +57,7 @@ class FSDPArgs: # support matrix. Name kept identical to Megatron's. deterministic_mode: bool = False - # Sequence Parallelism (USP = Ulysses x Ring). context_parallel_size is a - # backward-compatible alias for sequence_parallel_size. - context_parallel_size: int = 1 + # Sequence Parallelism (USP = Ulysses x Ring) sequence_parallel_size: int = 1 ulysses_degree: int = 0 # 0=auto: ulysses fills sp, ring=1 ring_degree: int = 0 # 0=auto: sp // ulysses diff --git a/miles/backends/fsdp_utils/parallel.py b/miles/backends/fsdp_utils/parallel.py index e5ad587a..7da78c98 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/miles/backends/fsdp_utils/parallel.py @@ -22,11 +22,7 @@ def create_fsdp_parallel_state(args: Namespace) -> ParallelState: rank = dist.get_rank() sp_size, ulysses_degree, ring_degree = validate_sp_config( - world_size, - args.sequence_parallel_size, - args.ulysses_degree, - args.ring_degree, - getattr(args, "num_attention_heads", None), + world_size, args.sequence_parallel_size, args.ulysses_degree, args.ring_degree ) dp_rank, sp_rank, _, _ = locate_rank(rank, sp_size, ulysses_degree, ring_degree) dp_size = world_size // sp_size @@ -40,17 +36,20 @@ def create_fsdp_parallel_state(args: Namespace) -> ParallelState: ) # 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, ring_degree) - for ranks in ulysses_groups: - group = dist.new_group(ranks) - if rank in ranks: - ulysses_group = group - for ranks in ring_groups: - group = dist.new_group(ranks) - if rank in ranks: - ring_group = group + 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, @@ -66,7 +65,7 @@ def create_fsdp_parallel_state(args: Namespace) -> ParallelState: 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, diff --git a/miles/backends/fsdp_utils/sp_attention.py b/miles/backends/fsdp_utils/sp_attention.py index db56a7ad..77e0f01a 100644 --- a/miles/backends/fsdp_utils/sp_attention.py +++ b/miles/backends/fsdp_utils/sp_attention.py @@ -10,41 +10,9 @@ import inspect import torch -import torch.distributed as dist from diffusers.models._modeling_parallel import ContextParallelInput, ContextParallelOutput -from .sp_ops import usp_attention - - -class _GatherSequence(torch.autograd.Function): - """All-gather local shards along dim; backward returns each rank's slice.""" - - @staticmethod - def forward(ctx, x, group, sp_rank, sp_size, dim): - 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): - start = ctx.sp_rank * ctx.local_size - return grad.narrow(ctx.dim, start, ctx.local_size), None, None, None, None - - -def shard_sequence(x, parallel_state, dim=1): - sp = parallel_state.sp_size - s = x.shape[dim] - if s % sp: - raise ValueError(f"sequence length {s} is not divisible by sp_size {sp}") - s_local = s // sp - return x.narrow(dim, parallel_state.sp_rank * s_local, s_local) - - -def gather_sequence(x, parallel_state, dim=1): - return _GatherSequence.apply(x, parallel_state.sp_group, parallel_state.sp_rank, parallel_state.sp_size, dim) +from .sp_ops import gather_sequence, shard_sequence, usp_attention class WanUSPAttnProcessor: @@ -136,7 +104,7 @@ def _split_if_expected(x, spec, parallel_state): return x if spec.expected_dims is not None and x.ndim != spec.expected_dims: return x - return shard_sequence(x, parallel_state, dim=spec.split_dim) + return shard_sequence(x, parallel_state.sp_rank, parallel_state.sp_size, dim=spec.split_dim) def _resolve_submodule(root, path): @@ -162,7 +130,9 @@ def _install_cp_plan_hooks(transformer, parallel_state): def gather_output(mod, args, output, _spec=spec): assert isinstance(output, torch.Tensor) - return gather_sequence(output, parallel_state, dim=_spec.gather_dim) + 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 diff --git a/miles/backends/fsdp_utils/sp_mesh.py b/miles/backends/fsdp_utils/sp_mesh.py index 5b22cb7d..939ec742 100644 --- a/miles/backends/fsdp_utils/sp_mesh.py +++ b/miles/backends/fsdp_utils/sp_mesh.py @@ -15,13 +15,15 @@ def resolve_sp_degrees(sequence_parallel_size, ulysses_degree=0, ring_degree=0): return sp, u, r -def validate_sp_config(world_size, sequence_parallel_size, ulysses_degree=0, ring_degree=0, num_heads=None): - """Validate at startup. Returns (sp, ulysses, ring).""" +def validate_sp_config(world_size, sequence_parallel_size, ulysses_degree=0, ring_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, ring_degree) if world_size % sp != 0: raise ValueError(f"world_size({world_size}) is not divisible by sequence_parallel_size({sp})") - if num_heads is not None and num_heads % u != 0: - raise ValueError(f"num_heads({num_heads}) is not divisible by ulysses_degree({u})") return sp, u, r diff --git a/miles/backends/fsdp_utils/sp_ops.py b/miles/backends/fsdp_utils/sp_ops.py index 60fcd972..4d252477 100644 --- a/miles/backends/fsdp_utils/sp_ops.py +++ b/miles/backends/fsdp_utils/sp_ops.py @@ -10,6 +10,36 @@ import torch.distributed as dist +class _GatherSequence(torch.autograd.Function): + """All-gather local shards along dim; backward returns each rank's slice.""" + + @staticmethod + def forward(ctx, x, group, sp_rank, sp_size, dim): + 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): + 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.""" @@ -137,7 +167,7 @@ def usp_attention(query, key, value, ulysses_group=None, ring_group=None): 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 and dist.get_world_size(ring_group) > 1: + 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( diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 83c6fcc6..2e4abdf3 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1244,9 +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 - if args.context_parallel_size > 1 and args.sequence_parallel_size == 1: - args.sequence_parallel_size = args.context_parallel_size - miles_validate_args(args) sglang_validate_args(args) diff --git a/tests/sp/sp_attention_parity.py b/tests/sp/sp_attention_parity.py index 5ccf70ee..3e4ac8ef 100644 --- a/tests/sp/sp_attention_parity.py +++ b/tests/sp/sp_attention_parity.py @@ -86,7 +86,6 @@ def main(): p.add_argument("--ulysses", type=int, default=0) p.add_argument("--ring", type=int, default=0) p.add_argument("--ckpt", action="store_true") - p.add_argument("--fwd-only", action="store_true") p.add_argument("--fp32", action="store_true", help="fp32 + SDPA, isolates bf16 summation rounding") cli = p.parse_args() if cli.fp32: @@ -103,7 +102,6 @@ def main(): sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, ring_degree=cli.ring, - context_parallel_size=1, ) ps = create_fsdp_parallel_state(args) @@ -128,13 +126,12 @@ def main(): print(f"[PARITY] sp={cli.sp} ulysses={ps.ulysses_degree} ring={ps.ring_degree} ckpt={cli.ckpt}") _report("forward(out)", out_sp, out_ref, rtol=2e-2, ctol=0.9990) _report("grad(input)", gin_sp, gin_ref, rtol=4e-2, ctol=0.9980) - if not cli.fwd_only: - params = dict(model.named_parameters()) - for n in attn_weight_names: - assert params[n].grad is not None, f"{n} grad is None — all-to-all backward did not propagate" - g = params[n].grad.detach().clone() - dist.all_reduce(g, group=ps.sp_group) - _report(f"grad({n})", g, gw_ref[n], rtol=5e-2, ctol=0.9950) + params = dict(model.named_parameters()) + for n in attn_weight_names: + assert params[n].grad is not None, f"{n} grad is None — all-to-all backward did not propagate" + g = params[n].grad.detach().clone() + dist.all_reduce(g, group=ps.sp_group) + _report(f"grad({n})", g, gw_ref[n], rtol=5e-2, ctol=0.9950) dist.barrier() if rank == 0: diff --git a/tests/sp/sp_grad_sync_parity.py b/tests/sp/sp_grad_sync_parity.py index fb291cff..1feb94d5 100644 --- a/tests/sp/sp_grad_sync_parity.py +++ b/tests/sp/sp_grad_sync_parity.py @@ -68,7 +68,7 @@ def main(): device = torch.cuda.current_device() init_gloo_group() - args = argparse.Namespace(sequence_parallel_size=4, ulysses_degree=4, ring_degree=0, context_parallel_size=1) + args = argparse.Namespace(sequence_parallel_size=4, ulysses_degree=4, ring_degree=0) ps = create_fsdp_parallel_state(args) hidden, enc, ts, out_grad = make_inputs(device) diff --git a/tests/sp/sp_init_smoke.py b/tests/sp/sp_init_smoke.py index 25e84439..8261c458 100644 --- a/tests/sp/sp_init_smoke.py +++ b/tests/sp/sp_init_smoke.py @@ -37,7 +37,6 @@ def main(): sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, ring_degree=cli.ring, - context_parallel_size=1, ) ps = create_fsdp_parallel_state(args) @@ -48,11 +47,16 @@ def main(): my_sp = next(g for g in sp_groups if rank in g) assert _members(ps.sp_group) == my_sp, f"rank{rank} sp_group {_members(ps.sp_group)} != {my_sp}" - if sp_size > 1: + if ps.ulysses_degree > 1: my_u = next(g for g in ulysses_groups if rank in g) - my_r = sorted(next(g for g in ring_groups if rank in g)) assert _members(ps.ulysses_group) == my_u, f"rank{rank} ulysses {_members(ps.ulysses_group)} != {my_u}" + else: + assert ps.ulysses_group is None + if ps.ring_degree > 1: + my_r = sorted(next(g for g in ring_groups if rank in g)) assert _members(ps.ring_group) == my_r, f"rank{rank} ring {_members(ps.ring_group)} != {my_r}" + else: + assert ps.ring_group is None dist.barrier() if rank == 0: diff --git a/tests/sp/sp_weight_sync_parity.py b/tests/sp/sp_weight_sync_parity.py index 0ab932b0..3c457d05 100644 --- a/tests/sp/sp_weight_sync_parity.py +++ b/tests/sp/sp_weight_sync_parity.py @@ -78,7 +78,6 @@ def main(): sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, ring_degree=cli.ring, - context_parallel_size=1, ) ps = create_fsdp_parallel_state(args) diff --git a/tests/sp/test_sp_mesh.py b/tests/sp/test_sp_mesh.py index 942641e7..29c46f35 100644 --- a/tests/sp/test_sp_mesh.py +++ b/tests/sp/test_sp_mesh.py @@ -74,14 +74,5 @@ def test_validate_rejects_illegal(): validate_sp_config(world_size=6, sequence_parallel_size=4) with pytest.raises(ValueError): validate_sp_config(world_size=8, sequence_parallel_size=4, ulysses_degree=3, ring_degree=1) - with pytest.raises(ValueError): - validate_sp_config(world_size=4, sequence_parallel_size=4, num_heads=40, ulysses_degree=3) - # Wan2.2 heads=40: ulysses in {2, 4} is legal - assert validate_sp_config(world_size=4, sequence_parallel_size=2, num_heads=40) == (2, 2, 1) - assert validate_sp_config(world_size=4, sequence_parallel_size=4, num_heads=40) == (4, 4, 1) - - -def test_num_heads_guard(): - with pytest.raises(ValueError): - validate_sp_config(world_size=6, sequence_parallel_size=3, num_heads=40) - assert validate_sp_config(world_size=8, sequence_parallel_size=8, num_heads=40) == (8, 8, 1) + assert validate_sp_config(world_size=4, sequence_parallel_size=2) == (2, 2, 1) + assert validate_sp_config(world_size=4, sequence_parallel_size=4) == (4, 4, 1) From ea3167293651cf39060ee6ea0e228b12ceb22f12 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Thu, 9 Jul 2026 09:23:44 +0000 Subject: [PATCH 11/24] =?UTF-8?q?diffusion:=20fsdp=5Fshard=5Fmode=20dp=5Fs?= =?UTF-8?q?p=20=E2=80=94=20shard=20parameters=20over=20the=20dp=20x=20sp?= =?UTF-8?q?=20mesh?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Aligns with torchtitan/VeOmni/DeepSpeed-Ulysses: FSDP shards over the flattened dp x sp mesh (default), halving/quartering per-rank parameter and master-state memory under SP (measured 54.4 -> 27.2 GB/rank at dp2xsp2). Gradient semantics live in the sequence gather's backward (sum over the sp group), so the explicit cross-sp grad all-reduce is skipped; fsdp_shard_mode=dp keeps the replicated placement as a validation anchor. Guardrail preamble no longer kills by process name; it refuses to start when this run's GPUs are occupied. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/arguments.py | 4 + miles/backends/fsdp_utils/parallel.py | 13 +- miles/backends/fsdp_utils/sp_attention.py | 7 +- miles/backends/fsdp_utils/sp_ops.py | 20 ++- ...run-diffusion-grpo-wan22-pickscore-5gpu.sh | 2 + tests/sp/run_gpu_suite.sh | 9 +- tests/sp/run_rl_guardrail.sh | 19 ++- tests/sp/sp_attention_parity.py | 2 + tests/sp/sp_grad_sync_parity.py | 149 ++++++++++++------ tests/sp/sp_init_smoke.py | 6 +- tests/sp/sp_weight_sync_parity.py | 8 +- 11 files changed, 170 insertions(+), 69 deletions(-) diff --git a/miles/backends/fsdp_utils/arguments.py b/miles/backends/fsdp_utils/arguments.py index a5ecd797..2d1c19cc 100644 --- a/miles/backends/fsdp_utils/arguments.py +++ b/miles/backends/fsdp_utils/arguments.py @@ -61,6 +61,10 @@ class FSDPArgs: sequence_parallel_size: int = 1 ulysses_degree: int = 0 # 0=auto: ulysses fills sp, ring=1 ring_degree: int = 0 # 0=auto: sp // ulysses + # "dp_sp": FSDP shards parameters over dp x sp (grads sp-summed inside the + # sequence gather's backward). "dp": shard over dp only, parameters + # replicated across sp, grads synced by an explicit cross-sp all-reduce. + fsdp_shard_mode: str = "dp_sp" # YAML bookkeeping config: str | None = None diff --git a/miles/backends/fsdp_utils/parallel.py b/miles/backends/fsdp_utils/parallel.py index 7da78c98..fada6de1 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/miles/backends/fsdp_utils/parallel.py @@ -15,8 +15,10 @@ def create_fsdp_parallel_state(args: Namespace) -> ParallelState: """ParallelState for FSDP + optional sequence parallelism. - FSDP shards parameters on the dp mesh dim only; SP gets its own process - groups and parameters are replicated across sp ranks. + SP gets its own process groups. fsdp_shard_mode picks the parameter + placement: "dp_sp" shards over the flattened dp x sp mesh; "dp" shards + over dp only, replicating parameters across sp ranks. Data dispatch is + by dp_rank in both modes. """ world_size = dist.get_world_size() rank = dist.get_rank() @@ -74,5 +76,12 @@ def create_fsdp_parallel_state(args: Namespace) -> ParallelState: ulysses_group=ulysses_group, ring_group=ring_group, ) + if args.fsdp_shard_mode not in ("dp", "dp_sp"): + raise ValueError(f"unknown fsdp_shard_mode: {args.fsdp_shard_mode}") parallel_state.dp_mesh = mesh["dp"] + parallel_state.fsdp_shard_mode = args.fsdp_shard_mode + if sp_size > 1 and args.fsdp_shard_mode == "dp_sp": + 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 index 77e0f01a..e323ba03 100644 --- a/miles/backends/fsdp_utils/sp_attention.py +++ b/miles/backends/fsdp_utils/sp_attention.py @@ -131,7 +131,12 @@ def _install_cp_plan_hooks(transformer, parallel_state): 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 + output, + parallel_state.sp_group, + parallel_state.sp_rank, + parallel_state.sp_size, + dim=_spec.gather_dim, + sum_grad=parallel_state.fsdp_shard_mode == "dp_sp", ) module.register_forward_hook(gather_output) diff --git a/miles/backends/fsdp_utils/sp_ops.py b/miles/backends/fsdp_utils/sp_ops.py index 4d252477..3996e797 100644 --- a/miles/backends/fsdp_utils/sp_ops.py +++ b/miles/backends/fsdp_utils/sp_ops.py @@ -11,21 +11,31 @@ class _GatherSequence(torch.autograd.Function): - """All-gather local shards along dim; backward returns each rank's slice.""" + """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): + 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 + return grad.narrow(ctx.dim, start, ctx.local_size), None, None, None, None, None def shard_sequence(x, sp_rank, sp_size, dim=1): @@ -36,8 +46,8 @@ def shard_sequence(x, sp_rank, sp_size, dim=1): 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) +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): diff --git a/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh b/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh index 79e6a00e..7bfd85c8 100755 --- a/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh +++ b/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh @@ -53,6 +53,7 @@ PYTHON_BIN="${PYTHON_BIN:-python}" SP_SIZE="${SP_SIZE:-1}" ULYSSES_DEGREE="${ULYSSES_DEGREE:-0}" RING_DEGREE="${RING_DEGREE:-0}" +FSDP_SHARD_MODE="${FSDP_SHARD_MODE:-dp_sp}" DATASETS_DIR="/root/datasets/miles-diffusion-datasets" if [[ ! -f "${DATASETS_DIR}/flowgrpo_pickscore/train.jsonl" ]]; then @@ -85,6 +86,7 @@ WAN_LORA_TARGET_MODULES=( --sequence-parallel-size "${SP_SIZE}" \ --ulysses-degree "${ULYSSES_DEGREE}" \ --ring-degree "${RING_DEGREE}" \ + --fsdp-shard-mode "${FSDP_SHARD_MODE}" \ --rollout-num-gpus 4 \ --rollout-num-gpus-per-engine 1 \ --num-gpus-per-node 5 \ diff --git a/tests/sp/run_gpu_suite.sh b/tests/sp/run_gpu_suite.sh index 8690f406..1212d5c9 100755 --- a/tests/sp/run_gpu_suite.sh +++ b/tests/sp/run_gpu_suite.sh @@ -11,6 +11,7 @@ run() { echo; echo "===== $* ====="; "$@"; } run python -m pytest -q tests/sp/sglang_usp_import_guard.py run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 2 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 2 --shard-mode dp run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 4 --ulysses 4 run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 4 --ulysses 2 --ring 2 @@ -21,10 +22,16 @@ run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --s run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 2 --ring 2 run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 2 --ring 2 --ckpt -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py --shard-mode dp +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py --shard-mode dp_sp +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py --shard-mode dp_sp --sp 4 --ulysses 2 --ring 2 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py --shard-mode dp_sp --sp 2 --ulysses 2 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 2 --ulysses 2 --shard-mode dp run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 2 --ulysses 2 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 4 --shard-mode dp run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 4 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 --ring 2 --shard-mode dp run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 --ring 2 echo; echo "ALL SP GPU TESTS PASSED" diff --git a/tests/sp/run_rl_guardrail.sh b/tests/sp/run_rl_guardrail.sh index 78d8f256..75e832a9 100755 --- a/tests/sp/run_rl_guardrail.sh +++ b/tests/sp/run_rl_guardrail.sh @@ -11,11 +11,18 @@ set -euo pipefail ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" -# A finished run can leave ray actors and sglang engine schedulers holding GPU -# memory; start clean. -ray stop --force >/dev/null 2>&1 || true -pkill -9 -f "sgl_diffusion::" 2>/dev/null || true -sleep 5 +# Machine-wide cleanup (ray stop / pkill by name) can kill unrelated jobs on a +# shared node. Instead, fail loudly if this run's GPUs are already occupied and +# let the operator decide what to kill. +GUARD_GPUS="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +for idx in ${GUARD_GPUS//,/ }; do + used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i "$idx") + if [ "$used" -gt 20000 ]; then + echo "GPU $idx already has ${used}MiB in use. Clean up leftover processes" >&2 + echo "manually (check owners with: nvidia-smi --query-compute-apps=pid,gpu_uuid,used_memory --format=csv)." >&2 + exit 1 + fi +done export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False @@ -25,6 +32,7 @@ PYTHON_BIN="${PYTHON_BIN:-python}" SP_SIZE="${SP_SIZE:-1}" ULYSSES_DEGREE="${ULYSSES_DEGREE:-0}" RING_DEGREE="${RING_DEGREE:-0}" +FSDP_SHARD_MODE="${FSDP_SHARD_MODE:-dp_sp}" NUM_ROLLOUT="${NUM_ROLLOUT:-2}" NUM_FRAMES="${NUM_FRAMES:-5}" GRAD_CKPT="${GRAD_CKPT:-0}" @@ -64,6 +72,7 @@ WAN_LORA_TARGET_MODULES=( --sequence-parallel-size "${SP_SIZE}" \ --ulysses-degree "${ULYSSES_DEGREE}" \ --ring-degree "${RING_DEGREE}" \ + --fsdp-shard-mode "${FSDP_SHARD_MODE}" \ --rollout-num-gpus 4 \ --rollout-num-gpus-per-engine 1 \ --num-gpus-per-node 4 \ diff --git a/tests/sp/sp_attention_parity.py b/tests/sp/sp_attention_parity.py index 3e4ac8ef..3f20828e 100644 --- a/tests/sp/sp_attention_parity.py +++ b/tests/sp/sp_attention_parity.py @@ -102,6 +102,8 @@ def main(): sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, ring_degree=cli.ring, + # operator-level parity uses plain slice-backward gather semantics + fsdp_shard_mode="dp", ) ps = create_fsdp_parallel_state(args) diff --git a/tests/sp/sp_grad_sync_parity.py b/tests/sp/sp_grad_sync_parity.py index 1feb94d5..32200f28 100644 --- a/tests/sp/sp_grad_sync_parity.py +++ b/tests/sp/sp_grad_sync_parity.py @@ -1,11 +1,14 @@ -"""SP gradient sync parity under real FSDP2: full grads must match a single-process -full-sequence reference after reduce-scatter (dp) + cross-sp all-reduce(SUM). +"""SP gradient parity under real FSDP2: full grads must match a single-process +full-sequence reference, for both shard modes. -Runs dp1 x sp4 so FSDP local shards are the full parameters, isolating the sp -dimension. Also asserts model outputs are bitwise identical across sp ranks -(the gather-after-proj_out contract that keeps loss/log_prob code unchanged). +dp mode: FSDP reduce-scatter over dp, then an explicit cross-sp all-reduce(SUM). +dp_sp mode: FSDP shards over the flattened dp x sp mesh; the sequence gather's +sum_grad backward makes FSDP's own reduce restore full grads with no extra sync. +Inputs are broadcast to all ranks, so the reference gradient is topology-free. +Also asserts model outputs are bitwise identical across sp ranks. -Usage: torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py +Usage: torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py \ + [--sp S --ulysses U --ring R] [--shard-mode dp|dp_sp] [--fp32] """ import argparse @@ -43,19 +46,24 @@ def build_model(device): return model -def make_inputs(device): - g = torch.Generator(device=device).manual_seed(123) +def make_inputs(device, seed=123): + g = torch.Generator(device=device).manual_seed(seed) hidden = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) enc = torch.randn(1, 32, 4096, device=device, dtype=DTYPE, generator=g) ts = torch.tensor([500], device=device) out_grad = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) - for t in (hidden, enc, out_grad): - dist.broadcast(t, src=0) return hidden, enc, ts, out_grad def main(): p = argparse.ArgumentParser() + p.add_argument("--sp", type=int, default=4) + p.add_argument("--ulysses", type=int, default=0) + p.add_argument("--ring", type=int, default=0) + p.add_argument("--shard-mode", choices=("dp", "dp_sp"), default="dp_sp") + p.add_argument("--lora", action="store_true", help="train LoRA params only, like the RL recipe") + p.add_argument("--distinct-dp", action="store_true", help="different data per dp rank, like real RL") + p.add_argument("--accum", type=int, default=1, help="gradient-accumulation microbatches") p.add_argument("--fp32", action="store_true", help="fp32 + SDPA, isolates bf16 summation rounding") cli = p.parse_args() if cli.fp32: @@ -68,71 +76,110 @@ def main(): device = torch.cuda.current_device() init_gloo_group() - args = argparse.Namespace(sequence_parallel_size=4, ulysses_degree=4, ring_degree=0) + args = argparse.Namespace( + sequence_parallel_size=cli.sp, + ulysses_degree=cli.ulysses, + ring_degree=cli.ring, + fsdp_shard_mode=cli.shard_mode, + ) ps = create_fsdp_parallel_state(args) - hidden, enc, ts, out_grad = make_inputs(device) - - # Full-sequence single-process reference (plain SDPA self-attention). - ref = build_model(device) + # One dataset per (microbatch, dp group); seeds are rank-independent so the + # reference can rebuild every dataset locally. + datasets = [] + for mb in range(cli.accum): + seeds = [1000 + (d * 100 if cli.distinct_dp else 0) + mb for d in range(ps.dp_size)] + datasets.append([make_inputs(device, seed=s) for s in seeds]) + + def maybe_lora(m): + if not cli.lora: + return m + from peft import LoraConfig, get_peft_model + + torch.manual_seed(7) + return get_peft_model( + m, LoraConfig(r=8, lora_alpha=16, target_modules=["to_q", "to_k", "to_v"], init_lora_weights=False) + ) + + # Full-sequence single-process reference (plain SDPA self-attention). Inputs + # are broadcast, so every dp replica computes the same full gradient. + ref = maybe_lora(build_model(device)) ref.set_attn_processor(WanUSPAttnProcessor(None)) - out = ref(hidden_states=hidden, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] - out.backward(out_grad) - ref_grads = {n: p.grad.detach().clone() for n, p in ref.named_parameters()} + for mbsets in datasets: + for hidden, enc, ts, out_grad in mbsets: + out = ref(hidden_states=hidden, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] + out.backward(out_grad / ps.dp_size) + ref_grads = {n: p.grad.detach().clone() for n, p in ref.named_parameters() if p.grad is not None} - # FSDP(dp1) + SP(sp4). - from torch.distributed.fsdp import fully_shard + from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard - model = build_model(device) + # fp32 reduce matches production apply_fsdp2 (--fsdp-reduce-dtype fp32). + mp_policy = MixedPrecisionPolicy(param_dtype=DTYPE, reduce_dtype=torch.float32) + model = maybe_lora(build_model(device)) for blk in model.blocks: - fully_shard(blk, mesh=ps.dp_mesh) - fully_shard(model, mesh=ps.dp_mesh) + fully_shard(blk, mesh=ps.fsdp_mesh, mp_policy=mp_policy) + fully_shard(model, mesh=ps.fsdp_mesh, mp_policy=mp_policy) apply_sequence_parallel(model, ps) - out = model(hidden_states=hidden, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] - - o32 = out.detach().float() - ref0 = o32.clone() - dist.broadcast(ref0, src=ps.dp_rank * ps.sp_size, group=ps.sp_group) - diff = (o32 - ref0).abs().max() - if rank == 0: - print(f"[OUTPUT] max abs diff across sp ranks = {diff.item():.2e} (must be 0)") - assert diff.item() == 0.0, "model outputs diverge across sp ranks" - - out.backward(out_grad) - - # Mirror actor._all_reduce_sp_grads. - for p in model.parameters(): - if p.grad is None: - continue - local = p.grad.to_local() if isinstance(p.grad, DTensor) else p.grad - f = local.float() - dist.all_reduce(f, group=ps.sp_group) - local.copy_(f.to(local.dtype)) + for i, mbsets in enumerate(datasets): + hidden, enc, ts, out_grad = mbsets[ps.dp_rank] + out = model(hidden_states=hidden, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] + if i == 0: + o32 = out.detach().float() + ref0 = o32.clone() + dist.broadcast(ref0, src=ps.dp_rank * ps.sp_size, group=ps.sp_group) + diff = (o32 - ref0).abs().max() + if rank == 0: + print(f"[OUTPUT] max abs diff across sp ranks = {diff.item():.2e} (must be 0)") + assert diff.item() == 0.0, "model outputs diverge across sp ranks" + out.backward(out_grad) + + if cli.shard_mode == "dp": + # Mirror actor._all_reduce_sp_grads; in dp_sp mode FSDP's own + # reduce-scatter already restores full grads (sum_grad gather). + for p in model.parameters(): + if p.grad is None: + continue + local = p.grad.to_local() if isinstance(p.grad, DTensor) else p.grad + f = local.float() + dist.all_reduce(f, group=ps.sp_group) + local.copy_(f.to(local.dtype)) fails = 0 checked = 0 for n, p in model.named_parameters(): if p.grad is None: continue - g = p.grad.to_local() if isinstance(p.grad, DTensor) else p.grad + g = p.grad.full_tensor() if isinstance(p.grad, DTensor) else p.grad r = ref_grads[n] - if g.shape != r.shape: - continue + assert g.shape == r.shape, f"{n}: {g.shape} != {r.shape}" rel = (g.float() - r.float()).abs().max() / r.float().abs().max().clamp_min(1e-6) cos = F.cosine_similarity(g.float().flatten(), r.float().flatten(), dim=0) checked += 1 - # Secondary band: small-norm tensors (biases) sit right at the bf16 - # summation noise floor; --fp32 confirms exact agreement there. - if not (rel < 5e-2 and cos > 0.99) and not (rel < 1e-1 and cos > 0.999): + # Secondary band: small-norm tensors (biases) sit at the bf16 noise + # floor of this tiny test model (ring backward pushes cross-attn K + # biases to ~7e-2 rel; dp and dp_sp modes produce bit-identical + # values there, so it is accumulation noise, not placement). + if not (rel < 5e-2 and cos > 0.99) and not (rel < 1e-1 and cos > 0.995): fails += 1 if rank == 0: print(f" [FAIL] {n:42s} rel={rel:.2e} cos={1 - cos:.2e}(1-)") + # clip_grad_norm_ must report the same total norm as the reference. + total = torch.nn.utils.clip_grad_norm_(model.parameters(), 1e9) + total = total.full_tensor() if isinstance(total, DTensor) else total + ref_total = torch.linalg.vector_norm( + torch.stack([torch.linalg.vector_norm(g.float()) for g in ref_grads.values()]) + ) + norm_rel = ((total.float() - ref_total) / ref_total).abs() + dist.barrier() if rank == 0: - print(f"[SP-GRAD-SYNC] checked={checked} fails={fails}") + print(f"[GRAD-NORM] clip={total.item():.6e} ref={ref_total.item():.6e} rel={norm_rel.item():.2e}") + assert norm_rel.item() < 5e-2, "clip_grad_norm_ reports a wrong total norm" + print(f"[SP-GRAD-SYNC] mode={cli.shard_mode} dp{ps.dp_size}xsp{ps.sp_size}" + f"(u{ps.ulysses_degree}r{ps.ring_degree}) checked={checked} fails={fails}") assert fails == 0 - print("[SP-GRAD-SYNC OK] dp1xsp4 FSDP+SP full grads == full-sequence reference") + print("[SP-GRAD-SYNC OK] full grads == full-sequence reference") dist.destroy_process_group() diff --git a/tests/sp/sp_init_smoke.py b/tests/sp/sp_init_smoke.py index 8261c458..8eb58c93 100644 --- a/tests/sp/sp_init_smoke.py +++ b/tests/sp/sp_init_smoke.py @@ -26,6 +26,7 @@ def main(): p.add_argument("--sp", type=int, default=2) p.add_argument("--ulysses", type=int, default=0) p.add_argument("--ring", type=int, default=0) + p.add_argument("--shard-mode", choices=("dp", "dp_sp"), default="dp_sp") cli = p.parse_args() dist.init_process_group("nccl") @@ -37,6 +38,7 @@ def main(): sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, ring_degree=cli.ring, + fsdp_shard_mode=cli.shard_mode, ) ps = create_fsdp_parallel_state(args) @@ -44,6 +46,8 @@ def main(): assert ps.sp_size == sp_size and ps.dp_size == dp_size assert dist.get_world_size(ps.sp_group) == sp_size assert dist.get_world_size(ps.dp_group) == dp_size + expected_fsdp = world if (cli.shard_mode == "dp_sp" and sp_size > 1) else dp_size + assert ps.fsdp_mesh.size() == expected_fsdp, f"fsdp mesh {ps.fsdp_mesh.size()} != {expected_fsdp}" my_sp = next(g for g in sp_groups if rank in g) assert _members(ps.sp_group) == my_sp, f"rank{rank} sp_group {_members(ps.sp_group)} != {my_sp}" @@ -62,7 +66,7 @@ def main(): if rank == 0: print( f"[SP-INIT-SMOKE OK] world={world} dp={dp_size} sp={sp_size} " - f"ulysses={ps.ulysses_degree} ring={ps.ring_degree}" + f"ulysses={ps.ulysses_degree} ring={ps.ring_degree} fsdp_mesh={ps.fsdp_mesh.size()}({cli.shard_mode})" ) dist.destroy_process_group() diff --git a/tests/sp/sp_weight_sync_parity.py b/tests/sp/sp_weight_sync_parity.py index 3c457d05..09ee4340 100644 --- a/tests/sp/sp_weight_sync_parity.py +++ b/tests/sp/sp_weight_sync_parity.py @@ -65,6 +65,7 @@ def main(): p.add_argument("--sp", type=int, default=2) p.add_argument("--ulysses", type=int, default=2) p.add_argument("--ring", type=int, default=0) + p.add_argument("--shard-mode", choices=("dp", "dp_sp"), default="dp_sp") cli = p.parse_args() dist.init_process_group("nccl") @@ -78,6 +79,7 @@ def main(): sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, ring_degree=cli.ring, + fsdp_shard_mode=cli.shard_mode, ) ps = create_fsdp_parallel_state(args) @@ -86,8 +88,8 @@ def main(): model = build_model(device) for blk in model.blocks: - fully_shard(blk, mesh=ps.dp_mesh) - fully_shard(model, mesh=ps.dp_mesh) + fully_shard(blk, mesh=ps.fsdp_mesh) + fully_shard(model, mesh=ps.fsdp_mesh) my_sum = compute_weights_checksum(full_state_pairs(model)) @@ -97,7 +99,7 @@ def main(): if rank == 0: all_equal = all(s == gathered[0] for s in gathered) match_ref = gathered[0] == ref_sum - print(f"[WEIGHT-SYNC] world={world} dp{ps.dp_size}xsp{ps.sp_size}(u{ps.ulysses_degree}r{ps.ring_degree})") + print(f"[WEIGHT-SYNC] world={world} mode={cli.shard_mode} dp{ps.dp_size}xsp{ps.sp_size}(u{ps.ulysses_degree}r{ps.ring_degree})") print(f"[WEIGHT-SYNC] all ranks equal={all_equal} == single-process ref={match_ref}") assert all_equal, "reconstructed full weights differ across ranks" assert match_ref, "reconstructed full weights != single-process reference" From 0420f3bfe51a337e7d9f4327baf830ecb95e6aad Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Thu, 9 Jul 2026 09:23:44 +0000 Subject: [PATCH 12/24] diffusion: materialize clip_grad_norm's DTensor before logging MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit clip_grad_norm_ returns a lazily-reduced partial-norm DTensor; logging it without full_tensor() leaks the local shard's norm, under-reporting grad_norm by sqrt(n_shards) (2x on the dp4 recipe). Clipping itself was always correct — the coefficient is computed in DTensor arithmetic. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/actor.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 56a6c88f..833f82bd 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -7,6 +7,7 @@ import ray import torch import torch.distributed as dist +from torch.distributed.tensor import DTensor import miles.backends.fsdp_utils.configs.qwen_image # noqa: F401 — register pipeline config import miles.backends.fsdp_utils.configs.sd3 # noqa: F401 — register pipeline config @@ -132,7 +133,7 @@ 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), @@ -412,9 +413,13 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: self.prof.step(rollout_id=rollout_id) if not self.args.debug_skip_optimizer_step: self.scaler.unscale_(self.optimizer) - if self.parallel_state.sp_size > 1: + if self.parallel_state.sp_size > 1 and self.parallel_state.fsdp_shard_mode == "dp": self._all_reduce_sp_grads() grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip_grad) + if isinstance(grad_norm, DTensor): + # clip returns a lazily-reduced partial norm; materialize it, + # otherwise the logged metric leaks the local shard's value. + grad_norm = grad_norm.full_tensor() log_stats["grad_norm"].append(grad_norm.detach()) self.scaler.step(self.optimizer) self.scaler.update() @@ -430,8 +435,6 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: def _all_reduce_sp_grads(self) -> None: """Each sp rank sees 1/sp of the tokens, so its grads are partial; all-reduce(SUM) across the sp group restores the full gradient. Runs on FSDP-sharded grads, in fp32.""" - from torch.distributed.tensor import DTensor - group = self.parallel_state.sp_group for p in self.model.parameters(): if p.grad is None: From f772ad506c8570ac6d74683e780f426187b10df7 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 10 Jul 2026 00:35:02 +0000 Subject: [PATCH 13/24] =?UTF-8?q?diffusion:=20SequenceParallelPlan=20?= =?UTF-8?q?=E2=80=94=20per-family=20SP=20declaration=20behind=20ModelBacke?= =?UTF-8?q?nd?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A model family opts into SP by declaring a SequenceParallelPlan (boundaries, attention installer, head count) instead of the engine hard-coding Wan: - ModelBackend.sequence_parallel_plan assembles the plan; the diffusers default reads the model's _cp_plan and the family config's apply_sp_attention hook. Native backends override the method wholesale. - WanUSPAttnProcessor moves to configs/wan2_2.py; the engine layer (sp_attention/sp_ops/sp_mesh/parallel) no longer references any model class. - The isinstance whitelist becomes capability checks: no _cp_plan or no apply_sp_attention raises with the missing declaration named. Validated: CPU suites (88 passed) + full SP GPU suite (21 bands, 66 checks). Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/actor.py | 3 +- .../configs/train_pipeline_config.py | 4 + miles/backends/fsdp_utils/configs/wan2_2.py | 88 ++++++++++++ miles/backends/fsdp_utils/model_backend.py | 20 ++- miles/backends/fsdp_utils/sp_attention.py | 136 ++++-------------- tests/sp/sp_attention_parity.py | 7 +- tests/sp/sp_grad_sync_parity.py | 13 +- 7 files changed, 157 insertions(+), 114 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 833f82bd..9045e38b 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -142,7 +142,8 @@ def init(self, args: Namespace, role: str, with_ref: bool = False) -> int: # ty if self.parallel_state.sp_size > 1: for model in self.models.values(): - apply_sequence_parallel(model, self.parallel_state) + 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() diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 04079032..08c1cfd0 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -179,3 +179,7 @@ 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``.""" + raise NotImplementedError(f"{type(self).__name__} does not implement sequence-parallel attention") diff --git a/miles/backends/fsdp_utils/configs/wan2_2.py b/miles/backends/fsdp_utils/configs/wan2_2.py index c570173e..2697fe20 100644 --- a/miles/backends/fsdp_utils/configs/wan2_2.py +++ b/miles/backends/fsdp_utils/configs/wan2_2.py @@ -5,9 +5,94 @@ import torch from miles.utils.types import CondKwargs +from ..sp_ops import usp_attention from .train_pipeline_config import TrainPipelineConfig, register_train_pipeline_config +class WanUSPAttnProcessor: + """Wan attention processor: self-attn via USP, cross-attn via local SDPA. + + Reuses Wan's QKV/RMSNorm/RoPE; rotary_emb arrives pre-sharded to S_local. + With parallel_state=None the self-attn falls back to plain SDPA (reference mode). + """ + + def __init__(self, parallel_state=None): + self.ulysses_group = parallel_state.ulysses_group if parallel_state is not None else None + self.ring_group = parallel_state.ring_group if parallel_state is not None else None + + def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, rotary_emb=None): + is_self_attn = encoder_hidden_states is None + + encoder_hidden_states_img = None + if attn.add_k_proj is not None: + image_context_length = encoder_hidden_states.shape[1] - 512 + encoder_hidden_states_img = encoder_hidden_states[:, :image_context_length] + encoder_hidden_states = encoder_hidden_states[:, image_context_length:] + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + query = attn.to_q(hidden_states) + key = attn.to_k(encoder_hidden_states) + value = attn.to_v(encoder_hidden_states) + + query = attn.norm_q(query) + key = attn.norm_k(key) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + if rotary_emb is not None: + query = _apply_rotary_emb(query, *rotary_emb) + key = _apply_rotary_emb(key, *rotary_emb) + + if is_self_attn: + hidden_states = usp_attention(query, key, value, self.ulysses_group, self.ring_group) + else: + hidden_states = torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + ).transpose(1, 2) + + hidden_states = hidden_states.flatten(2, 3).type_as(query) + + if encoder_hidden_states_img is not None: + key_img = attn.norm_added_k(attn.add_k_proj(encoder_hidden_states_img)).unflatten(2, (attn.heads, -1)) + value_img = attn.add_v_proj(encoder_hidden_states_img).unflatten(2, (attn.heads, -1)) + hidden_states_img = ( + torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), + key_img.transpose(1, 2), + value_img.transpose(1, 2), + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + .transpose(1, 2) + .flatten(2, 3) + .type_as(query) + ) + hidden_states = hidden_states + hidden_states_img + + hidden_states = attn.to_out[0](hidden_states) + hidden_states = attn.to_out[1](hidden_states) + return hidden_states + + +def _apply_rotary_emb(hidden_states, freqs_cos, freqs_sin): + x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) + cos = freqs_cos[..., 0::2] + sin = freqs_sin[..., 1::2] + out = torch.empty_like(hidden_states) + out[..., 0::2] = x1 * cos - x2 * sin + out[..., 1::2] = x1 * sin + x2 * cos + return out.type_as(hidden_states) + + @register_train_pipeline_config("wan2_2") class Wan2_2TrainPipelineConfig(TrainPipelineConfig): hf_ckpt_name_patterns = ("wan2.2", "wan-2.2") @@ -69,3 +154,6 @@ def cfg_combine( def preprocess_model_before_fsdp(self, model: torch.nn.Module) -> None: return None + + def apply_sp_attention(self, model: torch.nn.Module, parallel_state) -> None: + model.set_attn_processor(WanUSPAttnProcessor(parallel_state)) diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index 9bb8e500..ba70352c 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_attention 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,14 @@ 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") + return SequenceParallelPlan( + boundaries=boundaries, + attention=self.config.apply_sp_attention, + num_attention_heads=base.config.num_attention_heads, + ) diff --git a/miles/backends/fsdp_utils/sp_attention.py b/miles/backends/fsdp_utils/sp_attention.py index e323ba03..0844b0dc 100644 --- a/miles/backends/fsdp_utils/sp_attention.py +++ b/miles/backends/fsdp_utils/sp_attention.py @@ -1,102 +1,35 @@ -"""Sequence parallelism for diffusers DiTs: self-attention runs USP (sp_ops). +"""Sequence parallelism for diffusion DiTs: self-attention runs USP (sp_ops). Each sp rank holds S/sp latent tokens; attention internally gathers to the full -sequence via Ulysses all-to-all (+Ring) and scatters back. Shard/gather points are -taken from the model's diffusers ``_cp_plan``, so model outputs stay full-sequence -and every parameter sees a partial gradient; the actor's cross-sp grad all-reduce -applies uniformly. +sequence via Ulysses all-to-all (+Ring) and scatters back. A model family opts in +through a ``SequenceParallelPlan``; model outputs stay full-sequence, so every +parameter sees a partial gradient and loss/log_prob code is untouched. """ import inspect +from dataclasses import dataclass +from typing import Callable import torch from diffusers.models._modeling_parallel import ContextParallelInput, ContextParallelOutput -from .sp_ops import gather_sequence, shard_sequence, usp_attention +from .sp_ops import gather_sequence, shard_sequence -class WanUSPAttnProcessor: - """Wan attention processor: self-attn via USP, cross-attn via local SDPA. +@dataclass(frozen=True) +class SequenceParallelPlan: + """What one transformer family declares to run under SP. - Reuses Wan's QKV/RMSNorm/RoPE; rotary_emb arrives pre-sharded to S_local. - With parallel_state=None the self-attn falls back to plain SDPA (reference mode). + ``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``. """ - def __init__(self, parallel_state=None): - self.ulysses_group = parallel_state.ulysses_group if parallel_state is not None else None - self.ring_group = parallel_state.ring_group if parallel_state is not None else None - - def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, rotary_emb=None): - is_self_attn = encoder_hidden_states is None - - encoder_hidden_states_img = None - if attn.add_k_proj is not None: - image_context_length = encoder_hidden_states.shape[1] - 512 - encoder_hidden_states_img = encoder_hidden_states[:, :image_context_length] - encoder_hidden_states = encoder_hidden_states[:, image_context_length:] - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - query = attn.to_q(hidden_states) - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query = attn.norm_q(query) - key = attn.norm_k(key) - - query = query.unflatten(2, (attn.heads, -1)) - key = key.unflatten(2, (attn.heads, -1)) - value = value.unflatten(2, (attn.heads, -1)) - - if rotary_emb is not None: - query = _apply_rotary_emb(query, *rotary_emb) - key = _apply_rotary_emb(key, *rotary_emb) - - if is_self_attn: - hidden_states = usp_attention(query, key, value, self.ulysses_group, self.ring_group) - else: - hidden_states = torch.nn.functional.scaled_dot_product_attention( - query.transpose(1, 2), - key.transpose(1, 2), - value.transpose(1, 2), - attn_mask=attention_mask, - dropout_p=0.0, - is_causal=False, - ).transpose(1, 2) - - hidden_states = hidden_states.flatten(2, 3).type_as(query) - - if encoder_hidden_states_img is not None: - key_img = attn.norm_added_k(attn.add_k_proj(encoder_hidden_states_img)).unflatten(2, (attn.heads, -1)) - value_img = attn.add_v_proj(encoder_hidden_states_img).unflatten(2, (attn.heads, -1)) - hidden_states_img = ( - torch.nn.functional.scaled_dot_product_attention( - query.transpose(1, 2), - key_img.transpose(1, 2), - value_img.transpose(1, 2), - attn_mask=None, - dropout_p=0.0, - is_causal=False, - ) - .transpose(1, 2) - .flatten(2, 3) - .type_as(query) - ) - hidden_states = hidden_states + hidden_states_img - - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - return hidden_states - - -def _apply_rotary_emb(hidden_states, freqs_cos, freqs_sin): - x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) - cos = freqs_cos[..., 0::2] - sin = freqs_sin[..., 1::2] - out = torch.empty_like(hidden_states) - out[..., 0::2] = x1 * cos - x2 * sin - out[..., 1::2] = x1 * sin + x2 * cos - return out.type_as(hidden_states) + boundaries: dict + attention: Callable[[torch.nn.Module, object], None] + num_attention_heads: int def _split_if_expected(x, spec, parallel_state): @@ -116,14 +49,14 @@ def _resolve_submodule(root, path): return module -def _install_cp_plan_hooks(transformer, parallel_state): - """Install shard/gather hooks from the model's diffusers ``_cp_plan``. +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 is a differentiable all-gather, so backward stays partial-grad. """ - for path, spec in transformer._cp_plan.items(): + for path, spec in boundaries.items(): module = _resolve_submodule(transformer, path) if path else transformer if isinstance(spec, ContextParallelOutput): @@ -173,23 +106,14 @@ def split_outputs(mod, args, kwargs, output, _specs=output_specs): module.register_forward_hook(split_outputs, with_kwargs=True) -def apply_sequence_parallel(transformer, parallel_state): - """Wire SP into one transformer: replace self-attn processors and install - the shard/gather hooks declared by its _cp_plan. Call once per transformer - after FSDP wrapping.""" - from diffusers import WanTransformer3DModel - - base = transformer.get_base_model() if hasattr(transformer, "get_base_model") else transformer - if not isinstance(base, WanTransformer3DModel): - raise ValueError( - f"SP attention processor currently supports WanTransformer3DModel only, got {base.__class__.__name__}" - ) - heads = transformer.config.num_attention_heads - # args carry no head count, so the startup divisibility check must happen - # here where the real model config is available. - if heads % parallel_state.ulysses_degree != 0: +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({heads}) is not divisible by ulysses_degree({parallel_state.ulysses_degree})" + f"num_attention_heads({plan.num_attention_heads}) is not divisible by " + f"ulysses_degree({parallel_state.ulysses_degree})" ) - transformer.set_attn_processor(WanUSPAttnProcessor(parallel_state)) - _install_cp_plan_hooks(transformer, parallel_state) + plan.attention(transformer, parallel_state) + _install_boundary_hooks(transformer, plan.boundaries, parallel_state) diff --git a/tests/sp/sp_attention_parity.py b/tests/sp/sp_attention_parity.py index 3f20828e..21f2149c 100644 --- a/tests/sp/sp_attention_parity.py +++ b/tests/sp/sp_attention_parity.py @@ -17,7 +17,9 @@ from diffusers import WanTransformer3DModel from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state -from miles.backends.fsdp_utils.sp_attention import WanUSPAttnProcessor, apply_sequence_parallel +from miles.backends.fsdp_utils.configs.wan2_2 import Wan2_2TrainPipelineConfig, WanUSPAttnProcessor +from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend +from miles.backends.fsdp_utils.sp_attention import apply_sequence_parallel from miles.utils.distributed_utils import init_gloo_group DTYPE = torch.bfloat16 @@ -119,7 +121,8 @@ def main(): gw_ref = {n: dict(model.named_parameters())[n].grad.detach().clone() for n in attn_weight_names} model.zero_grad(set_to_none=True) - apply_sequence_parallel(model, ps) + plan = DiffusersModelBackend(Wan2_2TrainPipelineConfig()).sequence_parallel_plan(model) + apply_sequence_parallel(model, ps, plan) out_sp, gin_sp = _run(model, hidden, enc, ts, out_grad, cli.ckpt) # Each rank backprops only its 1/sp of the tokens; sum across sp restores full grads. diff --git a/tests/sp/sp_grad_sync_parity.py b/tests/sp/sp_grad_sync_parity.py index 32200f28..c7f1a5e2 100644 --- a/tests/sp/sp_grad_sync_parity.py +++ b/tests/sp/sp_grad_sync_parity.py @@ -19,8 +19,10 @@ from diffusers import WanTransformer3DModel from torch.distributed.tensor import DTensor +from miles.backends.fsdp_utils.configs.wan2_2 import Wan2_2TrainPipelineConfig, WanUSPAttnProcessor +from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state -from miles.backends.fsdp_utils.sp_attention import WanUSPAttnProcessor, apply_sequence_parallel +from miles.backends.fsdp_utils.sp_attention import apply_sequence_parallel from miles.utils.distributed_utils import init_gloo_group DTYPE = torch.bfloat16 @@ -118,7 +120,8 @@ def maybe_lora(m): for blk in model.blocks: fully_shard(blk, mesh=ps.fsdp_mesh, mp_policy=mp_policy) fully_shard(model, mesh=ps.fsdp_mesh, mp_policy=mp_policy) - apply_sequence_parallel(model, ps) + plan = DiffusersModelBackend(Wan2_2TrainPipelineConfig()).sequence_parallel_plan(model) + apply_sequence_parallel(model, ps, plan) for i, mbsets in enumerate(datasets): hidden, enc, ts, out_grad = mbsets[ps.dp_rank] @@ -176,8 +179,10 @@ def maybe_lora(m): if rank == 0: print(f"[GRAD-NORM] clip={total.item():.6e} ref={ref_total.item():.6e} rel={norm_rel.item():.2e}") assert norm_rel.item() < 5e-2, "clip_grad_norm_ reports a wrong total norm" - print(f"[SP-GRAD-SYNC] mode={cli.shard_mode} dp{ps.dp_size}xsp{ps.sp_size}" - f"(u{ps.ulysses_degree}r{ps.ring_degree}) checked={checked} fails={fails}") + print( + f"[SP-GRAD-SYNC] mode={cli.shard_mode} dp{ps.dp_size}xsp{ps.sp_size}" + f"(u{ps.ulysses_degree}r{ps.ring_degree}) checked={checked} fails={fails}" + ) assert fails == 0 print("[SP-GRAD-SYNC OK] full grads == full-sequence reference") dist.destroy_process_group() From ae0e066db674fe192937ed93804c4696407c61e1 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 10 Jul 2026 00:49:02 +0000 Subject: [PATCH 14/24] style: pre-commit formatting across sp files First full pre-commit pass on this branch (the job never ran while the PR was based on feat/wan): isort import merges, black line wraps, ruff Callable-import modernization, unused-import removal. No code changes. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/sp_attention.py | 4 ++-- miles/backends/fsdp_utils/sp_ops.py | 8 ++------ tests/sp/sglang_usp_import_guard.py | 3 +-- tests/sp/sp_attention_parity.py | 2 +- tests/sp/sp_weight_sync_parity.py | 11 +++++++---- tests/sp/test_sp_mesh.py | 7 +------ 6 files changed, 14 insertions(+), 21 deletions(-) diff --git a/miles/backends/fsdp_utils/sp_attention.py b/miles/backends/fsdp_utils/sp_attention.py index 0844b0dc..451b9019 100644 --- a/miles/backends/fsdp_utils/sp_attention.py +++ b/miles/backends/fsdp_utils/sp_attention.py @@ -7,11 +7,11 @@ """ import inspect +from collections.abc import Callable from dataclasses import dataclass -from typing import Callable import torch -from diffusers.models._modeling_parallel import ContextParallelInput, ContextParallelOutput +from diffusers.models._modeling_parallel import ContextParallelOutput from .sp_ops import gather_sequence, shard_sequence diff --git a/miles/backends/fsdp_utils/sp_ops.py b/miles/backends/fsdp_utils/sp_ops.py index 3996e797..59d798e1 100644 --- a/miles/backends/fsdp_utils/sp_ops.py +++ b/miles/backends/fsdp_utils/sp_ops.py @@ -132,9 +132,7 @@ def forward(ctx, query, key, value, group, scale, is_causal): @staticmethod def backward(ctx, grad_out): - from torch.distributed.tensor.experimental._attention import ( - _templated_ring_attention_backward, - ) + 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( @@ -180,9 +178,7 @@ def usp_attention(query, key, value, ulysses_group=None, ring_group=None): 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 = 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: diff --git a/tests/sp/sglang_usp_import_guard.py b/tests/sp/sglang_usp_import_guard.py index b51df88b..024d7d3b 100644 --- a/tests/sp/sglang_usp_import_guard.py +++ b/tests/sp/sglang_usp_import_guard.py @@ -15,8 +15,7 @@ def test_checksum_covers_dtype_shape(): src = inspect.getsource(weight_utils.compute_weights_checksum) assert "dtype" in src and "shape" in src, ( - f"imported sglang checksum is bytes-only: {weight_utils.__file__} — " - "transpose/reshape would not be detected" + f"imported sglang checksum is bytes-only: {weight_utils.__file__} — " "transpose/reshape would not be detected" ) diff --git a/tests/sp/sp_attention_parity.py b/tests/sp/sp_attention_parity.py index 21f2149c..27a8b1c4 100644 --- a/tests/sp/sp_attention_parity.py +++ b/tests/sp/sp_attention_parity.py @@ -16,9 +16,9 @@ import torch.nn.functional as F from diffusers import WanTransformer3DModel -from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state from miles.backends.fsdp_utils.configs.wan2_2 import Wan2_2TrainPipelineConfig, WanUSPAttnProcessor from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend +from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state from miles.backends.fsdp_utils.sp_attention import apply_sequence_parallel from miles.utils.distributed_utils import init_gloo_group diff --git a/tests/sp/sp_weight_sync_parity.py b/tests/sp/sp_weight_sync_parity.py index 09ee4340..6dba5b1d 100644 --- a/tests/sp/sp_weight_sync_parity.py +++ b/tests/sp/sp_weight_sync_parity.py @@ -16,9 +16,8 @@ import torch import torch.distributed as dist from diffusers import WanTransformer3DModel -from torch.distributed.fsdp import fully_shard - from sglang.multimodal_gen.runtime.loader.weight_utils import compute_weights_checksum +from torch.distributed.fsdp import fully_shard from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state from miles.utils.distributed_utils import init_gloo_group @@ -84,7 +83,9 @@ def main(): ps = create_fsdp_parallel_state(args) ref_model = build_model(device) - ref_sum = compute_weights_checksum([(n, pa.detach().cpu().contiguous()) for n, pa in ref_model.state_dict().items()]) + ref_sum = compute_weights_checksum( + [(n, pa.detach().cpu().contiguous()) for n, pa in ref_model.state_dict().items()] + ) model = build_model(device) for blk in model.blocks: @@ -99,7 +100,9 @@ def main(): if rank == 0: all_equal = all(s == gathered[0] for s in gathered) match_ref = gathered[0] == ref_sum - print(f"[WEIGHT-SYNC] world={world} mode={cli.shard_mode} dp{ps.dp_size}xsp{ps.sp_size}(u{ps.ulysses_degree}r{ps.ring_degree})") + print( + f"[WEIGHT-SYNC] world={world} mode={cli.shard_mode} dp{ps.dp_size}xsp{ps.sp_size}(u{ps.ulysses_degree}r{ps.ring_degree})" + ) print(f"[WEIGHT-SYNC] all ranks equal={all_equal} == single-process ref={match_ref}") assert all_equal, "reconstructed full weights differ across ranks" assert match_ref, "reconstructed full weights != single-process reference" diff --git a/tests/sp/test_sp_mesh.py b/tests/sp/test_sp_mesh.py index 29c46f35..86e3a8de 100644 --- a/tests/sp/test_sp_mesh.py +++ b/tests/sp/test_sp_mesh.py @@ -2,12 +2,7 @@ import pytest -from miles.backends.fsdp_utils.sp_mesh import ( - locate_rank, - resolve_sp_degrees, - sp_subgroups, - validate_sp_config, -) +from miles.backends.fsdp_utils.sp_mesh import locate_rank, resolve_sp_degrees, sp_subgroups, validate_sp_config def test_resolve_auto_degrees(): From 45ee1ea49ce76882ebed694ae3d270cc9324247f Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 10 Jul 2026 01:22:53 +0000 Subject: [PATCH 15/24] =?UTF-8?q?diffusion:=20fail-fast=20SP=20validation?= =?UTF-8?q?=20=E2=80=94=20driver-side=20support=20check,=20wildcard=20and?= =?UTF-8?q?=20key-resolution=20guards?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adversarial review of the plan seam surfaced four holes no numerics test covers; all four now fail loudly at the earliest point the facts exist: - validate_sp_support (driver arg validation, zero weight loads): rejects sp>1 with an explicit --fsdp-attention-backend (the SP installer replaces every processor, so backend selection and the deterministic-flash patch can never take effect), a model backend without sequence_parallel_plan, or a diffusers family without apply_sp_attention — previously these surfaced only after full load + FSDP sharding of every component. - plan construction rejects wildcard _cp_plan boundaries (e.g. QwenImage's transformer_blocks.*) the hook installer cannot expand yet, instead of a bare AttributeError deep in the getattr walk. - boundary key resolution reads the forward signature from the get_base_model()-unwrapped module (PeftModel exposes (*args, **kwargs)) and raises at install time for unmappable keys instead of silently skipping the split. Validated: 95 CPU tests (7 new), full SP GPU suite (21 bands, 66 checks), grad-sync --lora bands for both shard modes. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/model_backend.py | 31 ++++++++ miles/backends/fsdp_utils/sp_attention.py | 11 ++- miles/utils/arguments.py | 4 ++ .../fsdp_utils/test_validate_sp_support.py | 71 +++++++++++++++++++ 4 files changed, 116 insertions(+), 1 deletion(-) create mode 100644 tests/fast/backends/fsdp_utils/test_validate_sp_support.py diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index ba70352c..cc6545b9 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/miles/backends/fsdp_utils/model_backend.py @@ -88,8 +88,39 @@ def sequence_parallel_plan(self, model: torch.nn.Module) -> SequenceParallelPlan 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, ) + + +def validate_sp_support(args, cfg_cls) -> None: + """Driver-side, before any actor launches: reject launches whose SP config + cannot work, without loading weights.""" + from miles.utils.misc import load_function + + from .configs.train_pipeline_config import TrainPipelineConfig + + if args.fsdp_attention_backend is not None: + raise ValueError( + "--fsdp-attention-backend has no effect with sequence parallelism: " + "the SP attention installer replaces every attention processor" + ) + 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") + if ( + backend_cls.sequence_parallel_plan is DiffusersModelBackend.sequence_parallel_plan + and cfg_cls.apply_sp_attention is TrainPipelineConfig.apply_sp_attention + ): + raise ValueError( + f"{cfg_cls.__name__} does not implement apply_sp_attention; " + "sequence parallelism is unavailable for this model family" + ) diff --git a/miles/backends/fsdp_utils/sp_attention.py b/miles/backends/fsdp_utils/sp_attention.py index 451b9019..f4dacb6b 100644 --- a/miles/backends/fsdp_utils/sp_attention.py +++ b/miles/backends/fsdp_utils/sp_attention.py @@ -79,7 +79,16 @@ def gather_output(mod, args, output, _spec=spec): output_specs = {k: v for k, v in spec.items() if v.split_output} if input_specs: - param_names = list(inspect.signature(module.forward).parameters) + # 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) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 2e4abdf3..73b0f506 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -1358,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/tests/fast/backends/fsdp_utils/test_validate_sp_support.py b/tests/fast/backends/fsdp_utils/test_validate_sp_support.py new file mode 100644 index 00000000..6e055837 --- /dev/null +++ b/tests/fast/backends/fsdp_utils/test_validate_sp_support.py @@ -0,0 +1,71 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[]) + +from argparse import Namespace + +import pytest +import torch + +from miles.backends.fsdp_utils.configs.qwen_image import QwenImageTrainPipelineConfig +from miles.backends.fsdp_utils.configs.wan2_2 import Wan2_2TrainPipelineConfig +from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend, ModelBackend, validate_sp_support + +DIFFUSERS_BACKEND_PATH = "miles.backends.fsdp_utils.model_backend.DiffusersModelBackend" + + +def _args(**overrides): + defaults = dict(fsdp_attention_backend=None, model_backend_path=DIFFUSERS_BACKEND_PATH) + defaults.update(overrides) + return Namespace(**defaults) + + +class TestValidateSpSupport: + def test_wan_family_passes(self): + validate_sp_support(_args(), Wan2_2TrainPipelineConfig) + + # Under SP the attention installer replaces every processor, so an explicit + # backend selection can never take effect. + def test_explicit_attention_backend_rejected(self): + with pytest.raises(ValueError, match="fsdp-attention-backend"): + validate_sp_support(_args(fsdp_attention_backend="flash"), Wan2_2TrainPipelineConfig) + + def test_family_without_sp_attention_rejected(self): + with pytest.raises(ValueError, match="apply_sp_attention"): + validate_sp_support(_args(), QwenImageTrainPipelineConfig) + + def test_backend_without_plan_rejected(self): + path = f"{__name__}._PlanlessBackend" + with pytest.raises(ValueError, match="does not support sequence parallelism"): + validate_sp_support(_args(model_backend_path=path), QwenImageTrainPipelineConfig) + + # A native backend owning its whole plan is accepted without a config hook. + def test_backend_with_own_plan_passes(self): + path = f"{__name__}._OwnPlanBackend" + validate_sp_support(_args(model_backend_path=path), QwenImageTrainPipelineConfig) + + +class _PlanlessBackend(ModelBackend): + def load_models_and_scheduler(self, args, *, master_dtype): + raise NotImplementedError + + +class _OwnPlanBackend(ModelBackend): + def load_models_and_scheduler(self, args, *, master_dtype): + raise NotImplementedError + + def sequence_parallel_plan(self, model): + raise NotImplementedError + + +class TestPlanConstruction: + def test_wildcard_boundaries_rejected(self): + class _WildcardModel(torch.nn.Module): + _cp_plan = {"blocks.*": None} + + with pytest.raises(ValueError, match="wildcard"): + DiffusersModelBackend(Wan2_2TrainPipelineConfig()).sequence_parallel_plan(_WildcardModel()) + + def test_missing_cp_plan_rejected(self): + with pytest.raises(ValueError, match="_cp_plan"): + DiffusersModelBackend(Wan2_2TrainPipelineConfig()).sequence_parallel_plan(torch.nn.Linear(2, 2)) From 4901944566e64fcc84a90458839a6a2d789b8509 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Fri, 10 Jul 2026 03:04:07 +0000 Subject: [PATCH 16/24] =?UTF-8?q?tests(sp):=20determinism=20smoke=20?= =?UTF-8?q?=E2=80=94=20SP=20attention=20path=20bitwise=20under=20determini?= =?UTF-8?q?stic=20mode?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Validates that deterministic_mode's guarantees survive USP: forward+backward twice on identical inputs must produce bitwise-equal outputs and grads for both the ulysses (SDPA local) and ring (aten flash fwd/bwd) paths, with use_deterministic_algorithms(True, warn_only=False) engaged — warn_only=False also asserts no op on the path is registered nondeterministic, confirming the aten deterministic gate covers the ring backward torch templates call. Co-Authored-By: Claude Fable 5 --- tests/sp/run_gpu_suite.sh | 4 ++ tests/sp/sp_determinism_smoke.py | 111 +++++++++++++++++++++++++++++++ 2 files changed, 115 insertions(+) create mode 100644 tests/sp/sp_determinism_smoke.py diff --git a/tests/sp/run_gpu_suite.sh b/tests/sp/run_gpu_suite.sh index 1212d5c9..d3c3714f 100755 --- a/tests/sp/run_gpu_suite.sh +++ b/tests/sp/run_gpu_suite.sh @@ -34,4 +34,8 @@ run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py - run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 --ring 2 --shard-mode dp run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 --ring 2 +export CUBLAS_WORKSPACE_CONFIG=:4096:8 +run torchrun --standalone --nproc_per_node=2 tests/sp/sp_determinism_smoke.py --ulysses 2 --ring 1 +run torchrun --standalone --nproc_per_node=2 tests/sp/sp_determinism_smoke.py --ulysses 1 --ring 2 + echo; echo "ALL SP GPU TESTS PASSED" diff --git a/tests/sp/sp_determinism_smoke.py b/tests/sp/sp_determinism_smoke.py new file mode 100644 index 00000000..4218e4f5 --- /dev/null +++ b/tests/sp/sp_determinism_smoke.py @@ -0,0 +1,111 @@ +"""SP determinism smoke: the SP attention path (ulysses/ring) must be bitwise +deterministic under +torch.use_deterministic_algorithms (warn_only=False also asserts no op on the +path is registered nondeterministic). Runs forward+backward twice on identical +inputs and compares every grad bitwise. --no-det runs without the flag as a +control. + +torchrun --standalone --nproc_per_node=2 tests/sp/sp_determinism_smoke.py --ulysses 1 --ring 2 [--no-det] +""" + +import argparse +import os + +import torch +import torch.distributed as dist +from diffusers import WanTransformer3DModel + +from miles.backends.fsdp_utils.configs.wan2_2 import Wan2_2TrainPipelineConfig +from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend +from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state +from miles.backends.fsdp_utils.sp_attention import apply_sequence_parallel +from miles.utils.distributed_utils import init_gloo_group + +DTYPE = torch.bfloat16 + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--sp", type=int, default=2) + p.add_argument("--ulysses", type=int, default=1) + p.add_argument("--ring", type=int, default=2) + p.add_argument("--no-det", action="store_true") + cli = p.parse_args() + + if not cli.no_det: + assert os.environ.get("CUBLAS_WORKSPACE_CONFIG"), "set CUBLAS_WORKSPACE_CONFIG=:4096:8" + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + torch.use_deterministic_algorithms(True, warn_only=False) + + dist.init_process_group("nccl") + rank = dist.get_rank() + torch.cuda.set_device(rank % torch.cuda.device_count()) + device = torch.cuda.current_device() + init_gloo_group() + + args = argparse.Namespace( + sequence_parallel_size=cli.sp, + ulysses_degree=cli.ulysses, + ring_degree=cli.ring, + fsdp_shard_mode="dp", + ) + ps = create_fsdp_parallel_state(args) + + torch.manual_seed(0) + model = WanTransformer3DModel( + patch_size=(1, 2, 2), + num_attention_heads=8, + attention_head_dim=128, + in_channels=16, + out_channels=16, + text_dim=4096, + freq_dim=256, + ffn_dim=1024, + num_layers=2, + rope_max_seq_len=1024, + ).to(device=device, dtype=DTYPE) + model.train() + for prm in model.parameters(): + dist.broadcast(prm.data, src=0) + + g = torch.Generator(device=device).manual_seed(123) + hidden = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) + enc = torch.randn(1, 32, 4096, device=device, dtype=DTYPE, generator=g) + ts = torch.tensor([500], device=device) + out_grad = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) + for t in (hidden, enc, out_grad): + dist.broadcast(t, src=0) + + plan = DiffusersModelBackend(Wan2_2TrainPipelineConfig()).sequence_parallel_plan(model) + apply_sequence_parallel(model, ps, plan) + + def run_once(): + model.zero_grad(set_to_none=True) + out = model(hidden_states=hidden, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] + out.backward(out_grad) + return out.detach().clone(), {n: prm.grad.detach().clone() for n, prm in model.named_parameters()} + + out1, g1 = run_once() + out2, g2 = run_once() + + out_eq = torch.equal(out1, out2) + diffs = [n for n in g1 if not torch.equal(g1[n], g2[n])] + flag = torch.tensor([0 if (out_eq and not diffs) else 1], device=device) + dist.all_reduce(flag) + if rank == 0: + mode = "no-det(control)" if cli.no_det else "deterministic" + print( + f"[DET-PROBE] mode={mode} u{cli.ulysses}r{cli.ring} forward_bitwise={out_eq} grad_mismatches={len(diffs)}" + ) + for n in diffs[:8]: + d = (g1[n].float() - g2[n].float()).abs().max().item() + print(f" diff {n}: max_abs={d:.3e}") + print(f"[DET-PROBE {'OK' if flag.item() == 0 else 'NONDETERMINISTIC'}] (all ranks)") + if not cli.no_det: + assert flag.item() == 0, "SP attention path is not bitwise deterministic under deterministic mode" + dist.destroy_process_group() + + +if __name__ == "__main__": + main() From e1c03d30f6d35312c25671bb2d53f0b726ad5e4c Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 13 Jul 2026 02:47:09 +0000 Subject: [PATCH 17/24] =?UTF-8?q?diffusion:=20parameter=20placement=20is?= =?UTF-8?q?=20Option=20A=20unconditionally=20=E2=80=94=20Option=20B=20demo?= =?UTF-8?q?ted=20to=20a=20test=20anchor?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit FSDP always shards over the flattened dp x sp mesh; SP is invisible to it (same parameter account as sp=1). Placement is not a user decision, so the config surface shrinks to --sequence-parallel-size + --ulysses-degree (ring = sp / ulysses, derived): - fsdp_shard_mode and ring_degree leave arguments; the mesh selection in parallel.py is branch-free (comment notes the HSDP caveat: flatten the shard axes, which today happen to be the whole world). - _all_reduce_sp_grads leaves the actor; the sequence gather's sum_grad backward is the only gradient mechanism in production. - The sp-replicated Option B placement survives verbatim in tests/sp as a validation anchor: grad-sync/weight-sync build it from the exported dp_mesh and apply_sequence_parallel(..., sum_grad=False), so the cross-placement parity bands still exercise production machinery. Validated: 95 CPU tests, full SP GPU suite (66 checks), grad-sync --lora for both placements, determinism smokes. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/actor.py | 17 -------------- miles/backends/fsdp_utils/arguments.py | 7 +----- miles/backends/fsdp_utils/parallel.py | 15 +++++-------- miles/backends/fsdp_utils/sp_attention.py | 12 +++++----- ...run-diffusion-grpo-wan22-pickscore-5gpu.sh | 6 +---- tests/sp/run_gpu_suite.sh | 1 - tests/sp/run_rl_guardrail.sh | 4 ---- tests/sp/sp_attention_parity.py | 5 ++--- tests/sp/sp_determinism_smoke.py | 1 - tests/sp/sp_grad_sync_parity.py | 22 +++++++++++-------- tests/sp/sp_init_smoke.py | 8 +++---- tests/sp/sp_weight_sync_parity.py | 7 +++--- 12 files changed, 37 insertions(+), 68 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 9045e38b..ee216cbe 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -414,8 +414,6 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: self.prof.step(rollout_id=rollout_id) if not self.args.debug_skip_optimizer_step: self.scaler.unscale_(self.optimizer) - if self.parallel_state.sp_size > 1 and self.parallel_state.fsdp_shard_mode == "dp": - self._all_reduce_sp_grads() grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip_grad) if isinstance(grad_norm, DTensor): # clip returns a lazily-reduced partial norm; materialize it, @@ -433,21 +431,6 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: reduced = {f"train/{k}": torch.stack(v).mean().item() for k, v in log_stats.items()} self._gather_and_log_metrics(rollout_id, reduced, step=self.global_step) - def _all_reduce_sp_grads(self) -> None: - """Each sp rank sees 1/sp of the tokens, so its grads are partial; all-reduce(SUM) - across the sp group restores the full gradient. Runs on FSDP-sharded grads, in fp32.""" - group = self.parallel_state.sp_group - for p in self.model.parameters(): - if p.grad is None: - continue - local = p.grad.to_local() if isinstance(p.grad, DTensor) else p.grad - if local.dtype == torch.float32: - dist.all_reduce(local, group=group) - else: - f = local.float() - dist.all_reduce(f, group=group) - local.copy_(f) - def _maybe_legacy_window_pad_len(self, train_pairs: list, microbatch_ranges: list) -> int | None: """LEGACY 2D parity: the whole-window max cond seq_len (like the legacy tile path), or None unless the legacy --micro-batch-size-sample>1 path is active. TODO: remove with it.""" diff --git a/miles/backends/fsdp_utils/arguments.py b/miles/backends/fsdp_utils/arguments.py index 2d1c19cc..4e84fc97 100644 --- a/miles/backends/fsdp_utils/arguments.py +++ b/miles/backends/fsdp_utils/arguments.py @@ -59,12 +59,7 @@ class FSDPArgs: # Sequence Parallelism (USP = Ulysses x Ring) sequence_parallel_size: int = 1 - ulysses_degree: int = 0 # 0=auto: ulysses fills sp, ring=1 - ring_degree: int = 0 # 0=auto: sp // ulysses - # "dp_sp": FSDP shards parameters over dp x sp (grads sp-summed inside the - # sequence gather's backward). "dp": shard over dp only, parameters - # replicated across sp, grads synced by an explicit cross-sp all-reduce. - fsdp_shard_mode: str = "dp_sp" + 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/parallel.py b/miles/backends/fsdp_utils/parallel.py index fada6de1..017ca9aa 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/miles/backends/fsdp_utils/parallel.py @@ -15,16 +15,16 @@ def create_fsdp_parallel_state(args: Namespace) -> ParallelState: """ParallelState for FSDP + optional sequence parallelism. - SP gets its own process groups. fsdp_shard_mode picks the parameter - placement: "dp_sp" shards over the flattened dp x sp mesh; "dp" shards - over dp only, replicating parameters across sp ranks. Data dispatch is - by dp_rank in both modes. + 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() sp_size, ulysses_degree, ring_degree = validate_sp_config( - world_size, args.sequence_parallel_size, args.ulysses_degree, args.ring_degree + world_size, args.sequence_parallel_size, args.ulysses_degree, getattr(args, "ring_degree", 0) ) dp_rank, sp_rank, _, _ = locate_rank(rank, sp_size, ulysses_degree, ring_degree) dp_size = world_size // sp_size @@ -76,11 +76,8 @@ def create_fsdp_parallel_state(args: Namespace) -> ParallelState: ulysses_group=ulysses_group, ring_group=ring_group, ) - if args.fsdp_shard_mode not in ("dp", "dp_sp"): - raise ValueError(f"unknown fsdp_shard_mode: {args.fsdp_shard_mode}") parallel_state.dp_mesh = mesh["dp"] - parallel_state.fsdp_shard_mode = args.fsdp_shard_mode - if sp_size > 1 and args.fsdp_shard_mode == "dp_sp": + if sp_size > 1: parallel_state.fsdp_mesh = mesh[("dp", "sp")]._flatten("dp_sp") else: parallel_state.fsdp_mesh = mesh["dp"] diff --git a/miles/backends/fsdp_utils/sp_attention.py b/miles/backends/fsdp_utils/sp_attention.py index f4dacb6b..0fd3d1a9 100644 --- a/miles/backends/fsdp_utils/sp_attention.py +++ b/miles/backends/fsdp_utils/sp_attention.py @@ -49,12 +49,14 @@ def _resolve_submodule(root, path): return module -def _install_boundary_hooks(transformer, boundaries, parallel_state): +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, so backward stays partial-grad. + 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 @@ -69,7 +71,7 @@ def gather_output(mod, args, output, _spec=spec): parallel_state.sp_rank, parallel_state.sp_size, dim=_spec.gather_dim, - sum_grad=parallel_state.fsdp_shard_mode == "dp_sp", + sum_grad=sum_grad, ) module.register_forward_hook(gather_output) @@ -115,7 +117,7 @@ def split_outputs(mod, args, kwargs, output, _specs=output_specs): module.register_forward_hook(split_outputs, with_kwargs=True) -def apply_sequence_parallel(transformer, parallel_state, plan): +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.""" @@ -125,4 +127,4 @@ def apply_sequence_parallel(transformer, parallel_state, plan): f"ulysses_degree({parallel_state.ulysses_degree})" ) plan.attention(transformer, parallel_state) - _install_boundary_hooks(transformer, plan.boundaries, parallel_state) + _install_boundary_hooks(transformer, plan.boundaries, parallel_state, sum_grad=sum_grad) diff --git a/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh b/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh index 7bfd85c8..c542fd59 100755 --- a/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh +++ b/scripts/run-diffusion-grpo-wan22-pickscore-5gpu.sh @@ -49,11 +49,9 @@ 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). +# train world size. 0 = auto (ulysses fills sp); ring = sp / ulysses. SP_SIZE="${SP_SIZE:-1}" ULYSSES_DEGREE="${ULYSSES_DEGREE:-0}" -RING_DEGREE="${RING_DEGREE:-0}" -FSDP_SHARD_MODE="${FSDP_SHARD_MODE:-dp_sp}" DATASETS_DIR="/root/datasets/miles-diffusion-datasets" if [[ ! -f "${DATASETS_DIR}/flowgrpo_pickscore/train.jsonl" ]]; then @@ -85,8 +83,6 @@ WAN_LORA_TARGET_MODULES=( --actor-num-gpus-per-node 4 \ --sequence-parallel-size "${SP_SIZE}" \ --ulysses-degree "${ULYSSES_DEGREE}" \ - --ring-degree "${RING_DEGREE}" \ - --fsdp-shard-mode "${FSDP_SHARD_MODE}" \ --rollout-num-gpus 4 \ --rollout-num-gpus-per-engine 1 \ --num-gpus-per-node 5 \ diff --git a/tests/sp/run_gpu_suite.sh b/tests/sp/run_gpu_suite.sh index d3c3714f..0bad8bcb 100755 --- a/tests/sp/run_gpu_suite.sh +++ b/tests/sp/run_gpu_suite.sh @@ -11,7 +11,6 @@ run() { echo; echo "===== $* ====="; "$@"; } run python -m pytest -q tests/sp/sglang_usp_import_guard.py run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 2 -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 2 --shard-mode dp run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 4 --ulysses 4 run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 4 --ulysses 2 --ring 2 diff --git a/tests/sp/run_rl_guardrail.sh b/tests/sp/run_rl_guardrail.sh index 75e832a9..ea5cdd98 100755 --- a/tests/sp/run_rl_guardrail.sh +++ b/tests/sp/run_rl_guardrail.sh @@ -31,8 +31,6 @@ export PICKSCORE_NUM_GPUS_PER_WORKER=0 PYTHON_BIN="${PYTHON_BIN:-python}" SP_SIZE="${SP_SIZE:-1}" ULYSSES_DEGREE="${ULYSSES_DEGREE:-0}" -RING_DEGREE="${RING_DEGREE:-0}" -FSDP_SHARD_MODE="${FSDP_SHARD_MODE:-dp_sp}" NUM_ROLLOUT="${NUM_ROLLOUT:-2}" NUM_FRAMES="${NUM_FRAMES:-5}" GRAD_CKPT="${GRAD_CKPT:-0}" @@ -71,8 +69,6 @@ WAN_LORA_TARGET_MODULES=( --actor-num-gpus-per-node 4 \ --sequence-parallel-size "${SP_SIZE}" \ --ulysses-degree "${ULYSSES_DEGREE}" \ - --ring-degree "${RING_DEGREE}" \ - --fsdp-shard-mode "${FSDP_SHARD_MODE}" \ --rollout-num-gpus 4 \ --rollout-num-gpus-per-engine 1 \ --num-gpus-per-node 4 \ diff --git a/tests/sp/sp_attention_parity.py b/tests/sp/sp_attention_parity.py index 27a8b1c4..bef5c651 100644 --- a/tests/sp/sp_attention_parity.py +++ b/tests/sp/sp_attention_parity.py @@ -104,8 +104,6 @@ def main(): sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, ring_degree=cli.ring, - # operator-level parity uses plain slice-backward gather semantics - fsdp_shard_mode="dp", ) ps = create_fsdp_parallel_state(args) @@ -122,7 +120,8 @@ def main(): model.zero_grad(set_to_none=True) plan = DiffusersModelBackend(Wan2_2TrainPipelineConfig()).sequence_parallel_plan(model) - apply_sequence_parallel(model, ps, plan) + # operator-level parity uses plain slice-backward gather semantics + apply_sequence_parallel(model, ps, plan, sum_grad=False) out_sp, gin_sp = _run(model, hidden, enc, ts, out_grad, cli.ckpt) # Each rank backprops only its 1/sp of the tokens; sum across sp restores full grads. diff --git a/tests/sp/sp_determinism_smoke.py b/tests/sp/sp_determinism_smoke.py index 4218e4f5..ae5d2fa8 100644 --- a/tests/sp/sp_determinism_smoke.py +++ b/tests/sp/sp_determinism_smoke.py @@ -48,7 +48,6 @@ def main(): sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, ring_degree=cli.ring, - fsdp_shard_mode="dp", ) ps = create_fsdp_parallel_state(args) diff --git a/tests/sp/sp_grad_sync_parity.py b/tests/sp/sp_grad_sync_parity.py index c7f1a5e2..cd1a5faa 100644 --- a/tests/sp/sp_grad_sync_parity.py +++ b/tests/sp/sp_grad_sync_parity.py @@ -1,9 +1,11 @@ """SP gradient parity under real FSDP2: full grads must match a single-process -full-sequence reference, for both shard modes. +full-sequence reference, for both parameter placements. -dp mode: FSDP reduce-scatter over dp, then an explicit cross-sp all-reduce(SUM). -dp_sp mode: FSDP shards over the flattened dp x sp mesh; the sequence gather's -sum_grad backward makes FSDP's own reduce restore full grads with no extra sync. +dp_sp (production): FSDP shards over the flattened dp x sp mesh; the sequence +gather's sum_grad backward makes FSDP's own reduce restore full grads. +dp (validation anchor, test-only): params shard over dp only (replicated +across sp), slice-backward gather, explicit cross-sp grad all-reduce here in +the test — the obvious mechanism cross-checking the subtle one. Inputs are broadcast to all ranks, so the reference gradient is topology-free. Also asserts model outputs are bitwise identical across sp ranks. @@ -82,9 +84,11 @@ def main(): sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, ring_degree=cli.ring, - fsdp_shard_mode=cli.shard_mode, ) ps = create_fsdp_parallel_state(args) + # dp anchor: shard over the dp submesh (sp-replicated params) + slice-backward gather. + fsdp_mesh = ps.dp_mesh if cli.shard_mode == "dp" else ps.fsdp_mesh + sum_grad = cli.shard_mode == "dp_sp" # One dataset per (microbatch, dp group); seeds are rank-independent so the # reference can rebuild every dataset locally. datasets = [] @@ -118,10 +122,10 @@ def maybe_lora(m): mp_policy = MixedPrecisionPolicy(param_dtype=DTYPE, reduce_dtype=torch.float32) model = maybe_lora(build_model(device)) for blk in model.blocks: - fully_shard(blk, mesh=ps.fsdp_mesh, mp_policy=mp_policy) - fully_shard(model, mesh=ps.fsdp_mesh, mp_policy=mp_policy) + fully_shard(blk, mesh=fsdp_mesh, mp_policy=mp_policy) + fully_shard(model, mesh=fsdp_mesh, mp_policy=mp_policy) plan = DiffusersModelBackend(Wan2_2TrainPipelineConfig()).sequence_parallel_plan(model) - apply_sequence_parallel(model, ps, plan) + apply_sequence_parallel(model, ps, plan, sum_grad=sum_grad) for i, mbsets in enumerate(datasets): hidden, enc, ts, out_grad = mbsets[ps.dp_rank] @@ -137,7 +141,7 @@ def maybe_lora(m): out.backward(out_grad) if cli.shard_mode == "dp": - # Mirror actor._all_reduce_sp_grads; in dp_sp mode FSDP's own + # The anchor's explicit cross-sp grad sum; in dp_sp mode FSDP's own # reduce-scatter already restores full grads (sum_grad gather). for p in model.parameters(): if p.grad is None: diff --git a/tests/sp/sp_init_smoke.py b/tests/sp/sp_init_smoke.py index 8eb58c93..9f295783 100644 --- a/tests/sp/sp_init_smoke.py +++ b/tests/sp/sp_init_smoke.py @@ -26,7 +26,6 @@ def main(): p.add_argument("--sp", type=int, default=2) p.add_argument("--ulysses", type=int, default=0) p.add_argument("--ring", type=int, default=0) - p.add_argument("--shard-mode", choices=("dp", "dp_sp"), default="dp_sp") cli = p.parse_args() dist.init_process_group("nccl") @@ -38,7 +37,6 @@ def main(): sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, ring_degree=cli.ring, - fsdp_shard_mode=cli.shard_mode, ) ps = create_fsdp_parallel_state(args) @@ -46,8 +44,8 @@ def main(): assert ps.sp_size == sp_size and ps.dp_size == dp_size assert dist.get_world_size(ps.sp_group) == sp_size assert dist.get_world_size(ps.dp_group) == dp_size - expected_fsdp = world if (cli.shard_mode == "dp_sp" and sp_size > 1) else dp_size - assert ps.fsdp_mesh.size() == expected_fsdp, f"fsdp mesh {ps.fsdp_mesh.size()} != {expected_fsdp}" + assert ps.fsdp_mesh.size() == world, f"fsdp mesh {ps.fsdp_mesh.size()} != world {world}" + assert ps.dp_mesh.size() == dp_size, f"dp mesh {ps.dp_mesh.size()} != dp {dp_size}" my_sp = next(g for g in sp_groups if rank in g) assert _members(ps.sp_group) == my_sp, f"rank{rank} sp_group {_members(ps.sp_group)} != {my_sp}" @@ -66,7 +64,7 @@ def main(): if rank == 0: print( f"[SP-INIT-SMOKE OK] world={world} dp={dp_size} sp={sp_size} " - f"ulysses={ps.ulysses_degree} ring={ps.ring_degree} fsdp_mesh={ps.fsdp_mesh.size()}({cli.shard_mode})" + f"ulysses={ps.ulysses_degree} ring={ps.ring_degree} fsdp_mesh={ps.fsdp_mesh.size()}" ) dist.destroy_process_group() diff --git a/tests/sp/sp_weight_sync_parity.py b/tests/sp/sp_weight_sync_parity.py index 6dba5b1d..443d8600 100644 --- a/tests/sp/sp_weight_sync_parity.py +++ b/tests/sp/sp_weight_sync_parity.py @@ -78,9 +78,10 @@ def main(): sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, ring_degree=cli.ring, - fsdp_shard_mode=cli.shard_mode, ) ps = create_fsdp_parallel_state(args) + # dp anchor: shard over the dp submesh (sp-replicated params). + fsdp_mesh = ps.dp_mesh if cli.shard_mode == "dp" else ps.fsdp_mesh ref_model = build_model(device) ref_sum = compute_weights_checksum( @@ -89,8 +90,8 @@ def main(): model = build_model(device) for blk in model.blocks: - fully_shard(blk, mesh=ps.fsdp_mesh) - fully_shard(model, mesh=ps.fsdp_mesh) + fully_shard(blk, mesh=fsdp_mesh) + fully_shard(model, mesh=fsdp_mesh) my_sum = compute_weights_checksum(full_state_pairs(model)) From c019433a9faa65a3adc4adb49c2ff923d023bd0e Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 13 Jul 2026 06:28:14 +0000 Subject: [PATCH 18/24] =?UTF-8?q?diffusion:=20ring=5Fdegree=20is=20derived?= =?UTF-8?q?,=20not=20configured=20=E2=80=94=20drop=20it=20from=20every=20i?= =?UTF-8?q?nput=20surface?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ring = sp / ulysses is fully determined, so taking it as input carried zero information; its one historical value (cross-checking contradictory u,r pairs) died with the config-surface removal. sp_mesh's pure functions now take (sp, ulysses) only and return the derived ring; the u*r==sp product check collapses to a plain sp % u divisibility check. Test CLIs express ring topologies as (sp, ulysses): u2r2 = --sp 4 --ulysses 2. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/parallel.py | 6 +++--- miles/backends/fsdp_utils/sp_mesh.py | 23 +++++++++++------------ tests/sp/run_gpu_suite.sh | 16 ++++++++-------- tests/sp/sp_attention_parity.py | 4 +--- tests/sp/sp_determinism_smoke.py | 6 ++---- tests/sp/sp_grad_sync_parity.py | 4 +--- tests/sp/sp_init_smoke.py | 6 ++---- tests/sp/sp_weight_sync_parity.py | 4 +--- tests/sp/test_sp_mesh.py | 14 +++++++------- 9 files changed, 36 insertions(+), 47 deletions(-) diff --git a/miles/backends/fsdp_utils/parallel.py b/miles/backends/fsdp_utils/parallel.py index 017ca9aa..2ac732ce 100644 --- a/miles/backends/fsdp_utils/parallel.py +++ b/miles/backends/fsdp_utils/parallel.py @@ -24,9 +24,9 @@ def create_fsdp_parallel_state(args: Namespace) -> ParallelState: rank = dist.get_rank() sp_size, ulysses_degree, ring_degree = validate_sp_config( - world_size, args.sequence_parallel_size, args.ulysses_degree, getattr(args, "ring_degree", 0) + world_size, args.sequence_parallel_size, args.ulysses_degree ) - dp_rank, sp_rank, _, _ = locate_rank(rank, sp_size, ulysses_degree, ring_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")) @@ -41,7 +41,7 @@ def create_fsdp_parallel_state(args: Namespace) -> ParallelState: # 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, ring_degree) + _, _, _, 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) diff --git a/miles/backends/fsdp_utils/sp_mesh.py b/miles/backends/fsdp_utils/sp_mesh.py index 939ec742..a0aea138 100644 --- a/miles/backends/fsdp_utils/sp_mesh.py +++ b/miles/backends/fsdp_utils/sp_mesh.py @@ -5,31 +5,30 @@ """ -def resolve_sp_degrees(sequence_parallel_size, ulysses_degree=0, ring_degree=0): - """Normalize (sp, ulysses, ring). degree=0 means auto: ulysses fills sp, ring=1.""" +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 - r = ring_degree or (sp // u) - if u * r != sp: - raise ValueError(f"ulysses_degree({u}) * ring_degree({r}) != sequence_parallel_size({sp})") - return sp, u, r + 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, ring_degree=0): +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, ring_degree) + 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, ring_degree=0): +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.""" - sp, u, r = resolve_sp_degrees(sequence_parallel_size, ulysses_degree, ring_degree) + 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)] @@ -37,9 +36,9 @@ def sp_subgroups(world_size, sequence_parallel_size, ulysses_degree=0, ring_degr return dp_size, sp, sp_groups, ulysses_groups, ring_groups -def locate_rank(rank, sequence_parallel_size, ulysses_degree=0, ring_degree=0): +def locate_rank(rank, sequence_parallel_size, ulysses_degree=0): """Locate a global rank's (dp_rank, sp_rank, ulysses_rank, ring_rank).""" - sp, u, _ = resolve_sp_degrees(sequence_parallel_size, ulysses_degree, ring_degree) + 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/tests/sp/run_gpu_suite.sh b/tests/sp/run_gpu_suite.sh index 0bad8bcb..63b8dced 100755 --- a/tests/sp/run_gpu_suite.sh +++ b/tests/sp/run_gpu_suite.sh @@ -12,29 +12,29 @@ run python -m pytest -q tests/sp/sglang_usp_import_guard.py run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 2 run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 4 --ulysses 4 -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 4 --ulysses 2 --ring 2 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 4 --ulysses 2 run torchrun --standalone --nproc_per_node=2 tests/sp/sp_attention_parity.py --sp 2 run torchrun --standalone --nproc_per_node=2 tests/sp/sp_attention_parity.py --sp 2 --ckpt run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 4 run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 4 --ckpt -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 2 --ring 2 -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 2 --ring 2 --ckpt +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 2 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 2 --ckpt run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py --shard-mode dp run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py --shard-mode dp_sp -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py --shard-mode dp_sp --sp 4 --ulysses 2 --ring 2 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py --shard-mode dp_sp --sp 4 --ulysses 2 run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py --shard-mode dp_sp --sp 2 --ulysses 2 run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 2 --ulysses 2 --shard-mode dp run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 2 --ulysses 2 run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 4 --shard-mode dp run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 4 -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 --ring 2 --shard-mode dp -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 --ring 2 +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 --shard-mode dp +run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 export CUBLAS_WORKSPACE_CONFIG=:4096:8 -run torchrun --standalone --nproc_per_node=2 tests/sp/sp_determinism_smoke.py --ulysses 2 --ring 1 -run torchrun --standalone --nproc_per_node=2 tests/sp/sp_determinism_smoke.py --ulysses 1 --ring 2 +run torchrun --standalone --nproc_per_node=2 tests/sp/sp_determinism_smoke.py --sp 2 --ulysses 2 +run torchrun --standalone --nproc_per_node=2 tests/sp/sp_determinism_smoke.py --sp 2 --ulysses 1 echo; echo "ALL SP GPU TESTS PASSED" diff --git a/tests/sp/sp_attention_parity.py b/tests/sp/sp_attention_parity.py index bef5c651..ab3f676d 100644 --- a/tests/sp/sp_attention_parity.py +++ b/tests/sp/sp_attention_parity.py @@ -6,7 +6,7 @@ grads, with and without gradient checkpointing. Usage: torchrun --standalone --nproc_per_node=N tests/sp/sp_attention_parity.py \ - --sp S [--ulysses U --ring R] [--ckpt] [--fp32] + --sp S [--ulysses U] [--ckpt] [--fp32] """ import argparse @@ -86,7 +86,6 @@ def main(): p = argparse.ArgumentParser() p.add_argument("--sp", type=int, default=4) p.add_argument("--ulysses", type=int, default=0) - p.add_argument("--ring", type=int, default=0) p.add_argument("--ckpt", action="store_true") p.add_argument("--fp32", action="store_true", help="fp32 + SDPA, isolates bf16 summation rounding") cli = p.parse_args() @@ -103,7 +102,6 @@ def main(): args = argparse.Namespace( sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, - ring_degree=cli.ring, ) ps = create_fsdp_parallel_state(args) diff --git a/tests/sp/sp_determinism_smoke.py b/tests/sp/sp_determinism_smoke.py index ae5d2fa8..a94dd0cd 100644 --- a/tests/sp/sp_determinism_smoke.py +++ b/tests/sp/sp_determinism_smoke.py @@ -5,7 +5,7 @@ inputs and compares every grad bitwise. --no-det runs without the flag as a control. -torchrun --standalone --nproc_per_node=2 tests/sp/sp_determinism_smoke.py --ulysses 1 --ring 2 [--no-det] +torchrun --standalone --nproc_per_node=2 tests/sp/sp_determinism_smoke.py --sp 2 --ulysses 1 [--no-det] """ import argparse @@ -28,7 +28,6 @@ def main(): p = argparse.ArgumentParser() p.add_argument("--sp", type=int, default=2) p.add_argument("--ulysses", type=int, default=1) - p.add_argument("--ring", type=int, default=2) p.add_argument("--no-det", action="store_true") cli = p.parse_args() @@ -47,7 +46,6 @@ def main(): args = argparse.Namespace( sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, - ring_degree=cli.ring, ) ps = create_fsdp_parallel_state(args) @@ -95,7 +93,7 @@ def run_once(): if rank == 0: mode = "no-det(control)" if cli.no_det else "deterministic" print( - f"[DET-PROBE] mode={mode} u{cli.ulysses}r{cli.ring} forward_bitwise={out_eq} grad_mismatches={len(diffs)}" + f"[DET-PROBE] mode={mode} u{ps.ulysses_degree}r{ps.ring_degree} forward_bitwise={out_eq} grad_mismatches={len(diffs)}" ) for n in diffs[:8]: d = (g1[n].float() - g2[n].float()).abs().max().item() diff --git a/tests/sp/sp_grad_sync_parity.py b/tests/sp/sp_grad_sync_parity.py index cd1a5faa..ca5086eb 100644 --- a/tests/sp/sp_grad_sync_parity.py +++ b/tests/sp/sp_grad_sync_parity.py @@ -10,7 +10,7 @@ Also asserts model outputs are bitwise identical across sp ranks. Usage: torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py \ - [--sp S --ulysses U --ring R] [--shard-mode dp|dp_sp] [--fp32] + [--sp S --ulysses U] [--shard-mode dp|dp_sp] [--fp32] """ import argparse @@ -63,7 +63,6 @@ def main(): p = argparse.ArgumentParser() p.add_argument("--sp", type=int, default=4) p.add_argument("--ulysses", type=int, default=0) - p.add_argument("--ring", type=int, default=0) p.add_argument("--shard-mode", choices=("dp", "dp_sp"), default="dp_sp") p.add_argument("--lora", action="store_true", help="train LoRA params only, like the RL recipe") p.add_argument("--distinct-dp", action="store_true", help="different data per dp rank, like real RL") @@ -83,7 +82,6 @@ def main(): args = argparse.Namespace( sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, - ring_degree=cli.ring, ) ps = create_fsdp_parallel_state(args) # dp anchor: shard over the dp submesh (sp-replicated params) + slice-backward gather. diff --git a/tests/sp/sp_init_smoke.py b/tests/sp/sp_init_smoke.py index 9f295783..1005f6f2 100644 --- a/tests/sp/sp_init_smoke.py +++ b/tests/sp/sp_init_smoke.py @@ -1,6 +1,6 @@ """Multi-GPU SP init smoke test: NCCL group members must match the sp_mesh layout. -Usage: torchrun --standalone --nproc_per_node=N tests/sp/sp_init_smoke.py --sp S [--ulysses U --ring R] +Usage: torchrun --standalone --nproc_per_node=N tests/sp/sp_init_smoke.py --sp S [--ulysses U] """ import argparse @@ -25,7 +25,6 @@ def main(): p = argparse.ArgumentParser() p.add_argument("--sp", type=int, default=2) p.add_argument("--ulysses", type=int, default=0) - p.add_argument("--ring", type=int, default=0) cli = p.parse_args() dist.init_process_group("nccl") @@ -36,11 +35,10 @@ def main(): args = argparse.Namespace( sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, - ring_degree=cli.ring, ) ps = create_fsdp_parallel_state(args) - dp_size, sp_size, sp_groups, ulysses_groups, ring_groups = sp_subgroups(world, cli.sp, cli.ulysses, cli.ring) + dp_size, sp_size, sp_groups, ulysses_groups, ring_groups = sp_subgroups(world, cli.sp, cli.ulysses) assert ps.sp_size == sp_size and ps.dp_size == dp_size assert dist.get_world_size(ps.sp_group) == sp_size assert dist.get_world_size(ps.dp_group) == dp_size diff --git a/tests/sp/sp_weight_sync_parity.py b/tests/sp/sp_weight_sync_parity.py index 443d8600..439854c6 100644 --- a/tests/sp/sp_weight_sync_parity.py +++ b/tests/sp/sp_weight_sync_parity.py @@ -8,7 +8,7 @@ Usage: torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 2 --ulysses 2 torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 4 - torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 --ring 2 + torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 """ import argparse @@ -63,7 +63,6 @@ def main(): p = argparse.ArgumentParser() p.add_argument("--sp", type=int, default=2) p.add_argument("--ulysses", type=int, default=2) - p.add_argument("--ring", type=int, default=0) p.add_argument("--shard-mode", choices=("dp", "dp_sp"), default="dp_sp") cli = p.parse_args() @@ -77,7 +76,6 @@ def main(): args = argparse.Namespace( sequence_parallel_size=cli.sp, ulysses_degree=cli.ulysses, - ring_degree=cli.ring, ) ps = create_fsdp_parallel_state(args) # dp anchor: shard over the dp submesh (sp-replicated params). diff --git a/tests/sp/test_sp_mesh.py b/tests/sp/test_sp_mesh.py index 86e3a8de..21733a55 100644 --- a/tests/sp/test_sp_mesh.py +++ b/tests/sp/test_sp_mesh.py @@ -8,7 +8,7 @@ def test_resolve_auto_degrees(): assert resolve_sp_degrees(4) == (4, 4, 1) assert resolve_sp_degrees(4, ulysses_degree=2) == (4, 2, 2) - assert resolve_sp_degrees(8, 2, 4) == (8, 2, 4) + assert resolve_sp_degrees(8, 2) == (8, 2, 4) assert resolve_sp_degrees(1) == (1, 1, 1) @@ -19,7 +19,7 @@ def test_resolve_illegal(): def test_sglang_alignment_example(): # Matches sglang set_seq_parallel_pg_by_sp_groups: sp=4,u=2,r=2 - _, _, sp_groups, ulysses_groups, ring_groups = sp_subgroups(4, 4, 2, 2) + _, _, sp_groups, ulysses_groups, ring_groups = sp_subgroups(4, 4, 2) assert sp_groups == [[0, 1, 2, 3]] assert ulysses_groups == [[0, 1], [2, 3]] assert ring_groups == [[0, 2], [1, 3]] @@ -40,8 +40,8 @@ def test_sglang_alignment_example(): ], ) def test_layout_invariants_any_scale(world, sp, u, r): - dp_size, sp_size, sp_groups, ulysses_groups, ring_groups = sp_subgroups(world, sp, u, r) - assert sp_size == sp == u * r + dp_size, sp_size, sp_groups, ulysses_groups, ring_groups = sp_subgroups(world, sp, u) + assert sp_size == sp == u * r # r in the parametrize row is the expected derived ring degree assert dp_size * sp_size == world assert len(sp_groups) == dp_size and all(len(g) == sp for g in sp_groups) # ulysses groups: contiguous, size u, exact cover @@ -54,9 +54,9 @@ def test_layout_invariants_any_scale(world, sp, u, r): @pytest.mark.parametrize("world,sp,u,r", [(8, 4, 2, 2), (16, 8, 2, 4), (256, 16, 8, 2)]) def test_locate_rank_consistent_with_subgroups(world, sp, u, r): - _, _, _, ulysses_groups, ring_groups = sp_subgroups(world, sp, u, r) + _, _, _, ulysses_groups, ring_groups = sp_subgroups(world, sp, u) for rank in range(world): - dp_rank, sp_rank, u_rank, r_rank = locate_rank(rank, sp, u, r) + dp_rank, sp_rank, u_rank, r_rank = locate_rank(rank, sp, u) assert dp_rank == rank // sp and sp_rank == rank % sp my_u_group = next(g for g in ulysses_groups if rank in g) my_r_group = next(g for g in ring_groups if rank in g) @@ -68,6 +68,6 @@ def test_validate_rejects_illegal(): with pytest.raises(ValueError): validate_sp_config(world_size=6, sequence_parallel_size=4) with pytest.raises(ValueError): - validate_sp_config(world_size=8, sequence_parallel_size=4, ulysses_degree=3, ring_degree=1) + validate_sp_config(world_size=8, sequence_parallel_size=4, ulysses_degree=3) assert validate_sp_config(world_size=4, sequence_parallel_size=2) == (2, 2, 1) assert validate_sp_config(world_size=4, sequence_parallel_size=4) == (4, 4, 1) From 1ed9c3770fa9575375b50e4381569387fb947782 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 13 Jul 2026 06:31:08 +0000 Subject: [PATCH 19/24] docs(sp_mesh): note the sp_subgroups <-> locate_rank inverse relation Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/sp_mesh.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/miles/backends/fsdp_utils/sp_mesh.py b/miles/backends/fsdp_utils/sp_mesh.py index a0aea138..ce61eabd 100644 --- a/miles/backends/fsdp_utils/sp_mesh.py +++ b/miles/backends/fsdp_utils/sp_mesh.py @@ -27,7 +27,10 @@ def validate_sp_config(world_size, sequence_parallel_size, ulysses_degree=0): 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.""" + """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)] @@ -37,7 +40,10 @@ def sp_subgroups(world_size, sequence_parallel_size, ulysses_degree=0): def locate_rank(rank, sequence_parallel_size, ulysses_degree=0): - """Locate a global rank's (dp_rank, sp_rank, ulysses_rank, ring_rank).""" + """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) From 1982bded58697e16493d8ce5a51e69beb758199c Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 13 Jul 2026 21:44:11 +0000 Subject: [PATCH 20/24] diffusion: default SP attention intercepts at diffusers' dispatch_attention_fn MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A model family no longer writes an attention processor: the default apply_sp_attention wraps the modeling module's dispatch_attention_fn (found through the MRO — fully_shard swizzles the class into torch's fsdp module) and reroutes self-attention call sites to sp_ops.usp_attention. Models flag self- vs cross-attention themselves (upstream passes parallel_config only when encoder_hidden_states is None), so the model's own processors run untouched and WanUSPAttnProcessor (the 80-line copy) is deleted. The wrapper fails loudly outside the validated envelope (mask / dropout / is_causal / custom scale); models whose modules never imported dispatch_attention_fn are rejected at install with instructions to provide a family override. Kernel choice stays pinned inside usp_attention. Validated: parity vs the stock diffusers processors is bitwise for forward and input grads (fp32 band exact); full GPU suite 77 checks; 10-step real-RL runs at sp=1 / u4 / u2r2 / u1r4 — per-step train-vs-rollout canary flat across all topologies (no rounding accumulation), rewards and clipfrac in the sp=1 band. Co-Authored-By: Claude Fable 5 --- .../configs/train_pipeline_config.py | 11 ++- miles/backends/fsdp_utils/configs/wan2_2.py | 88 ------------------- miles/backends/fsdp_utils/model_backend.py | 20 ++--- miles/backends/fsdp_utils/sp_attention.py | 65 +++++++++++++- 4 files changed, 81 insertions(+), 103 deletions(-) diff --git a/miles/backends/fsdp_utils/configs/train_pipeline_config.py b/miles/backends/fsdp_utils/configs/train_pipeline_config.py index 08c1cfd0..8bd0b90e 100644 --- a/miles/backends/fsdp_utils/configs/train_pipeline_config.py +++ b/miles/backends/fsdp_utils/configs/train_pipeline_config.py @@ -181,5 +181,12 @@ 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``.""" - raise NotImplementedError(f"{type(self).__name__} does not implement sequence-parallel attention") + """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/configs/wan2_2.py b/miles/backends/fsdp_utils/configs/wan2_2.py index 2697fe20..c570173e 100644 --- a/miles/backends/fsdp_utils/configs/wan2_2.py +++ b/miles/backends/fsdp_utils/configs/wan2_2.py @@ -5,94 +5,9 @@ import torch from miles.utils.types import CondKwargs -from ..sp_ops import usp_attention from .train_pipeline_config import TrainPipelineConfig, register_train_pipeline_config -class WanUSPAttnProcessor: - """Wan attention processor: self-attn via USP, cross-attn via local SDPA. - - Reuses Wan's QKV/RMSNorm/RoPE; rotary_emb arrives pre-sharded to S_local. - With parallel_state=None the self-attn falls back to plain SDPA (reference mode). - """ - - def __init__(self, parallel_state=None): - self.ulysses_group = parallel_state.ulysses_group if parallel_state is not None else None - self.ring_group = parallel_state.ring_group if parallel_state is not None else None - - def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None, rotary_emb=None): - is_self_attn = encoder_hidden_states is None - - encoder_hidden_states_img = None - if attn.add_k_proj is not None: - image_context_length = encoder_hidden_states.shape[1] - 512 - encoder_hidden_states_img = encoder_hidden_states[:, :image_context_length] - encoder_hidden_states = encoder_hidden_states[:, image_context_length:] - - if encoder_hidden_states is None: - encoder_hidden_states = hidden_states - query = attn.to_q(hidden_states) - key = attn.to_k(encoder_hidden_states) - value = attn.to_v(encoder_hidden_states) - - query = attn.norm_q(query) - key = attn.norm_k(key) - - query = query.unflatten(2, (attn.heads, -1)) - key = key.unflatten(2, (attn.heads, -1)) - value = value.unflatten(2, (attn.heads, -1)) - - if rotary_emb is not None: - query = _apply_rotary_emb(query, *rotary_emb) - key = _apply_rotary_emb(key, *rotary_emb) - - if is_self_attn: - hidden_states = usp_attention(query, key, value, self.ulysses_group, self.ring_group) - else: - hidden_states = torch.nn.functional.scaled_dot_product_attention( - query.transpose(1, 2), - key.transpose(1, 2), - value.transpose(1, 2), - attn_mask=attention_mask, - dropout_p=0.0, - is_causal=False, - ).transpose(1, 2) - - hidden_states = hidden_states.flatten(2, 3).type_as(query) - - if encoder_hidden_states_img is not None: - key_img = attn.norm_added_k(attn.add_k_proj(encoder_hidden_states_img)).unflatten(2, (attn.heads, -1)) - value_img = attn.add_v_proj(encoder_hidden_states_img).unflatten(2, (attn.heads, -1)) - hidden_states_img = ( - torch.nn.functional.scaled_dot_product_attention( - query.transpose(1, 2), - key_img.transpose(1, 2), - value_img.transpose(1, 2), - attn_mask=None, - dropout_p=0.0, - is_causal=False, - ) - .transpose(1, 2) - .flatten(2, 3) - .type_as(query) - ) - hidden_states = hidden_states + hidden_states_img - - hidden_states = attn.to_out[0](hidden_states) - hidden_states = attn.to_out[1](hidden_states) - return hidden_states - - -def _apply_rotary_emb(hidden_states, freqs_cos, freqs_sin): - x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) - cos = freqs_cos[..., 0::2] - sin = freqs_sin[..., 1::2] - out = torch.empty_like(hidden_states) - out[..., 0::2] = x1 * cos - x2 * sin - out[..., 1::2] = x1 * sin + x2 * cos - return out.type_as(hidden_states) - - @register_train_pipeline_config("wan2_2") class Wan2_2TrainPipelineConfig(TrainPipelineConfig): hf_ckpt_name_patterns = ("wan2.2", "wan-2.2") @@ -154,6 +69,3 @@ def cfg_combine( def preprocess_model_before_fsdp(self, model: torch.nn.Module) -> None: return None - - def apply_sp_attention(self, model: torch.nn.Module, parallel_state) -> None: - model.set_attn_processor(WanUSPAttnProcessor(parallel_state)) diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index cc6545b9..90dced83 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/miles/backends/fsdp_utils/model_backend.py @@ -101,26 +101,22 @@ def sequence_parallel_plan(self, model: torch.nn.Module) -> SequenceParallelPlan ) +# 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.""" - from miles.utils.misc import load_function + cannot work, without loading weights. - from .configs.train_pipeline_config import TrainPipelineConfig + 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: " - "the SP attention installer replaces every attention processor" + "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") - if ( - backend_cls.sequence_parallel_plan is DiffusersModelBackend.sequence_parallel_plan - and cfg_cls.apply_sp_attention is TrainPipelineConfig.apply_sp_attention - ): - raise ValueError( - f"{cfg_cls.__name__} does not implement apply_sp_attention; " - "sequence parallelism is unavailable for this model family" - ) diff --git a/miles/backends/fsdp_utils/sp_attention.py b/miles/backends/fsdp_utils/sp_attention.py index 0fd3d1a9..470c295d 100644 --- a/miles/backends/fsdp_utils/sp_attention.py +++ b/miles/backends/fsdp_utils/sp_attention.py @@ -6,14 +6,16 @@ parameter sees a partial gradient and loss/log_prob code is untouched. """ +import functools import inspect +import sys 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 +from .sp_ops import gather_sequence, shard_sequence, usp_attention @dataclass(frozen=True) @@ -117,6 +119,67 @@ def split_outputs(mod, args, kwargs, output, _specs=output_specs): module.register_forward_hook(split_outputs, with_kwargs=True) +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 + + 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 From 2af9dacc2c4437187ad50505827c0b8a4a208c5f Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 13 Jul 2026 21:44:22 +0000 Subject: [PATCH 21/24] Revert "diffusion: align engine grouping span with rollout, assert full coverage" This reverts commit bc445195fd280381f84ba4cdcb93d790f0a857d2. --- .../fsdp_utils/diffusion_update_weight_utils.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py index c98114ea..dc6193c8 100644 --- a/miles/backends/fsdp_utils/diffusion_update_weight_utils.py +++ b/miles/backends/fsdp_utils/diffusion_update_weight_utils.py @@ -107,18 +107,10 @@ def connect_rollout_engines( ) -> None: self.rollout_engines = rollout_engines - # Match rollout.init_rollout_engines' per-engine span; a mismatch would - # leave orphaned ranks that fail much later in update_bucket_weights. - num_gpu_per_engine = min(self.args.rollout_num_gpus_per_engine, self.args.num_gpus_per_node) - assert dist.get_world_size() == len(rollout_engines) * num_gpu_per_engine, ( - f"train world {dist.get_world_size()} != {len(rollout_engines)} engines x " - f"{num_gpu_per_engine} gpus per engine" - ) - # Here we assume the gpu id of rollout engines and train actors are the same. for i, engine in enumerate(self.rollout_engines): - start_rank = i * num_gpu_per_engine - end_rank = (i + 1) * num_gpu_per_engine + start_rank = i * self.args.rollout_num_gpus_per_engine + end_rank = (i + 1) * self.args.rollout_num_gpus_per_engine group_ranks = list(range(start_rank, end_rank)) new_group = dist.new_group( ranks=group_ranks, From 94bc339b0d9df32f66fd60d1cc5bfe9dfd431623 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 13 Jul 2026 21:45:19 +0000 Subject: [PATCH 22/24] Revert the clip_grad_norm logging materialization (split to #35) Partial revert of 0420f3b: that commit also carried the dp_mesh -> fsdp_mesh line (folded in during a rebase), which belongs to the SP placement and stays. Only the DTensor import and the full_tensor materialization leave; they land independently as #35, which this PR's cross-placement grad_norm comparisons depend on. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/actor.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index ee216cbe..49e5e203 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -7,7 +7,6 @@ import ray import torch import torch.distributed as dist -from torch.distributed.tensor import DTensor import miles.backends.fsdp_utils.configs.qwen_image # noqa: F401 — register pipeline config import miles.backends.fsdp_utils.configs.sd3 # noqa: F401 — register pipeline config @@ -415,10 +414,6 @@ def _train_core(self, rollout_id: int, rollout_data) -> None: if not self.args.debug_skip_optimizer_step: self.scaler.unscale_(self.optimizer) grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip_grad) - if isinstance(grad_norm, DTensor): - # clip returns a lazily-reduced partial norm; materialize it, - # otherwise the logged metric leaks the local shard's value. - grad_norm = grad_norm.full_tensor() log_stats["grad_norm"].append(grad_norm.detach()) self.scaler.step(self.optimizer) self.scaler.update() From 8a38143cc98fb250ead2cefb4a836080bcbd612f Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 13 Jul 2026 21:45:34 +0000 Subject: [PATCH 23/24] =?UTF-8?q?tests:=20keep=20the=20SP=20suite=20local?= =?UTF-8?q?=20=E2=80=94=20it=20lands=20in=20a=20separate=20test-infra=20PR?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The full validation suite (mesh invariants, attention/grad/weight-sync parity, determinism smokes, RL guardrail runner) stays out of this PR by maintainer decision; every band was run locally against this exact tree (results recorded in the PR description). The files remain in the working tree for local use and will be contributed with CI registration in a follow-up PR. Co-Authored-By: Claude Fable 5 --- .../fsdp_utils/test_validate_sp_support.py | 71 ------- tests/sp/__init__.py | 0 tests/sp/run_gpu_suite.sh | 40 ---- tests/sp/run_rl_guardrail.sh | 117 ----------- tests/sp/sglang_usp_import_guard.py | 26 --- tests/sp/sp_attention_parity.py | 145 ------------- tests/sp/sp_determinism_smoke.py | 108 ---------- tests/sp/sp_grad_sync_parity.py | 194 ------------------ tests/sp/sp_init_smoke.py | 71 ------- tests/sp/sp_weight_sync_parity.py | 120 ----------- tests/sp/test_sp_mesh.py | 73 ------- 11 files changed, 965 deletions(-) delete mode 100644 tests/fast/backends/fsdp_utils/test_validate_sp_support.py delete mode 100644 tests/sp/__init__.py delete mode 100755 tests/sp/run_gpu_suite.sh delete mode 100755 tests/sp/run_rl_guardrail.sh delete mode 100644 tests/sp/sglang_usp_import_guard.py delete mode 100644 tests/sp/sp_attention_parity.py delete mode 100644 tests/sp/sp_determinism_smoke.py delete mode 100644 tests/sp/sp_grad_sync_parity.py delete mode 100644 tests/sp/sp_init_smoke.py delete mode 100644 tests/sp/sp_weight_sync_parity.py delete mode 100644 tests/sp/test_sp_mesh.py diff --git a/tests/fast/backends/fsdp_utils/test_validate_sp_support.py b/tests/fast/backends/fsdp_utils/test_validate_sp_support.py deleted file mode 100644 index 6e055837..00000000 --- a/tests/fast/backends/fsdp_utils/test_validate_sp_support.py +++ /dev/null @@ -1,71 +0,0 @@ -from tests.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=15, suite="stage-a-cpu", labels=[]) - -from argparse import Namespace - -import pytest -import torch - -from miles.backends.fsdp_utils.configs.qwen_image import QwenImageTrainPipelineConfig -from miles.backends.fsdp_utils.configs.wan2_2 import Wan2_2TrainPipelineConfig -from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend, ModelBackend, validate_sp_support - -DIFFUSERS_BACKEND_PATH = "miles.backends.fsdp_utils.model_backend.DiffusersModelBackend" - - -def _args(**overrides): - defaults = dict(fsdp_attention_backend=None, model_backend_path=DIFFUSERS_BACKEND_PATH) - defaults.update(overrides) - return Namespace(**defaults) - - -class TestValidateSpSupport: - def test_wan_family_passes(self): - validate_sp_support(_args(), Wan2_2TrainPipelineConfig) - - # Under SP the attention installer replaces every processor, so an explicit - # backend selection can never take effect. - def test_explicit_attention_backend_rejected(self): - with pytest.raises(ValueError, match="fsdp-attention-backend"): - validate_sp_support(_args(fsdp_attention_backend="flash"), Wan2_2TrainPipelineConfig) - - def test_family_without_sp_attention_rejected(self): - with pytest.raises(ValueError, match="apply_sp_attention"): - validate_sp_support(_args(), QwenImageTrainPipelineConfig) - - def test_backend_without_plan_rejected(self): - path = f"{__name__}._PlanlessBackend" - with pytest.raises(ValueError, match="does not support sequence parallelism"): - validate_sp_support(_args(model_backend_path=path), QwenImageTrainPipelineConfig) - - # A native backend owning its whole plan is accepted without a config hook. - def test_backend_with_own_plan_passes(self): - path = f"{__name__}._OwnPlanBackend" - validate_sp_support(_args(model_backend_path=path), QwenImageTrainPipelineConfig) - - -class _PlanlessBackend(ModelBackend): - def load_models_and_scheduler(self, args, *, master_dtype): - raise NotImplementedError - - -class _OwnPlanBackend(ModelBackend): - def load_models_and_scheduler(self, args, *, master_dtype): - raise NotImplementedError - - def sequence_parallel_plan(self, model): - raise NotImplementedError - - -class TestPlanConstruction: - def test_wildcard_boundaries_rejected(self): - class _WildcardModel(torch.nn.Module): - _cp_plan = {"blocks.*": None} - - with pytest.raises(ValueError, match="wildcard"): - DiffusersModelBackend(Wan2_2TrainPipelineConfig()).sequence_parallel_plan(_WildcardModel()) - - def test_missing_cp_plan_rejected(self): - with pytest.raises(ValueError, match="_cp_plan"): - DiffusersModelBackend(Wan2_2TrainPipelineConfig()).sequence_parallel_plan(torch.nn.Linear(2, 2)) diff --git a/tests/sp/__init__.py b/tests/sp/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/tests/sp/run_gpu_suite.sh b/tests/sp/run_gpu_suite.sh deleted file mode 100755 index 63b8dced..00000000 --- a/tests/sp/run_gpu_suite.sh +++ /dev/null @@ -1,40 +0,0 @@ -#!/usr/bin/env bash -# Full SP GPU test suite (needs 4 GPUs). CPU tests: pytest tests/sp/test_sp_mesh.py -set -euo pipefail -cd "$(dirname "$0")/../.." -export PYTHONPATH=. - -torchrun() { python -m torch.distributed.run "$@"; } - -run() { echo; echo "===== $* ====="; "$@"; } - -run python -m pytest -q tests/sp/sglang_usp_import_guard.py - -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 2 -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 4 --ulysses 4 -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_init_smoke.py --sp 4 --ulysses 2 - -run torchrun --standalone --nproc_per_node=2 tests/sp/sp_attention_parity.py --sp 2 -run torchrun --standalone --nproc_per_node=2 tests/sp/sp_attention_parity.py --sp 2 --ckpt -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 4 -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 4 --ckpt -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 2 -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_attention_parity.py --sp 4 --ulysses 2 --ckpt - -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py --shard-mode dp -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py --shard-mode dp_sp -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py --shard-mode dp_sp --sp 4 --ulysses 2 -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py --shard-mode dp_sp --sp 2 --ulysses 2 - -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 2 --ulysses 2 --shard-mode dp -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 2 --ulysses 2 -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 4 --shard-mode dp -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 4 -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 --shard-mode dp -run torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 - -export CUBLAS_WORKSPACE_CONFIG=:4096:8 -run torchrun --standalone --nproc_per_node=2 tests/sp/sp_determinism_smoke.py --sp 2 --ulysses 2 -run torchrun --standalone --nproc_per_node=2 tests/sp/sp_determinism_smoke.py --sp 2 --ulysses 1 - -echo; echo "ALL SP GPU TESTS PASSED" diff --git a/tests/sp/run_rl_guardrail.sh b/tests/sp/run_rl_guardrail.sh deleted file mode 100755 index ea5cdd98..00000000 --- a/tests/sp/run_rl_guardrail.sh +++ /dev/null @@ -1,117 +0,0 @@ -#!/usr/bin/env bash -# Real-RL numeric guardrail: run the Wan2.2 PickScore recipe for a few steps on -# 4 GPUs (train+rollout colocate, PickScore reward on CPU) and compare train -# metrics between bands. With identical seeds the rollout data is identical, so -# step-1 adv_abs_mean / log_prob_old_idx_0 must match bitwise across bands; -# log_prob_new / loss / grad_norm may differ at bf16 summation level. -# -# Usage: -# bash tests/sp/run_rl_guardrail.sh # FSDP dp4 baseline -# SP_SIZE=2 ULYSSES_DEGREE=2 bash tests/sp/run_rl_guardrail.sh # dp2 x sp2 -set -euo pipefail -ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" - -# Machine-wide cleanup (ray stop / pkill by name) can kill unrelated jobs on a -# shared node. Instead, fail loudly if this run's GPUs are already occupied and -# let the operator decide what to kill. -GUARD_GPUS="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" -for idx in ${GUARD_GPUS//,/ }; do - used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i "$idx") - if [ "$used" -gt 20000 ]; then - echo "GPU $idx already has ${used}MiB in use. Clean up leftover processes" >&2 - echo "manually (check owners with: nvidia-smi --query-compute-apps=pid,gpu_uuid,used_memory --format=csv)." >&2 - exit 1 - fi -done - -export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" -export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False -export PICKSCORE_NUM_GPUS_PER_WORKER=0 - -PYTHON_BIN="${PYTHON_BIN:-python}" -SP_SIZE="${SP_SIZE:-1}" -ULYSSES_DEGREE="${ULYSSES_DEGREE:-0}" -NUM_ROLLOUT="${NUM_ROLLOUT:-2}" -NUM_FRAMES="${NUM_FRAMES:-5}" -GRAD_CKPT="${GRAD_CKPT:-0}" -CANDIDATE_STEPS="${CANDIDATE_STEPS:-1,2,3}" -# Same-data comparison across bands: save rollout data in one band, replay it -# in another so both trainers see literally identical samples. -SAVE_ROLLOUT_DATA="${SAVE_ROLLOUT_DATA:-}" -LOAD_ROLLOUT_DATA="${LOAD_ROLLOUT_DATA:-}" -RUN_NAME="${RUN_NAME:-sp_guardrail_sp${SP_SIZE}_$(date +%Y%m%d_%H%M%S)}" -DATASETS_DIR="/root/datasets/miles-diffusion-datasets" - -EXTRA_ARGS=() -[[ "${GRAD_CKPT}" == "1" ]] && EXTRA_ARGS+=(--gradient-checkpointing) -[[ -n "${SAVE_ROLLOUT_DATA}" ]] && EXTRA_ARGS+=(--save-debug-rollout-data "${SAVE_ROLLOUT_DATA}") -[[ -n "${LOAD_ROLLOUT_DATA}" ]] && EXTRA_ARGS+=(--load-debug-rollout-data "${LOAD_ROLLOUT_DATA}") - -WAN_LORA_TARGET_MODULES=( - attn1.to_q attn1.to_k attn1.to_v attn1.to_out.0 - attn2.to_q attn2.to_k attn2.to_v attn2.to_out.0 - ffn.net.0.proj ffn.net.2 -) - -"${PYTHON_BIN}" -u "${ROOT_DIR}/train_diffusion.py" \ - --train-backend fsdp \ - --rollout-function-path miles.rollout.sglang_diffusion_rollout.generate_rollout \ - --hf-checkpoint Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --diffusion-model Wan-AI/Wan2.2-T2V-A14B-Diffusers \ - --prompt-data "${DATASETS_DIR}/flowgrpo_pickscore/train.jsonl" \ - --input-key input \ - --rollout-batch-size 8 \ - --n-samples-per-prompt 8 \ - --num-rollout "${NUM_ROLLOUT}" \ - --num-steps-per-rollout 2 \ - --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 4 \ - --colocate \ - --use-lora \ - --lora-rank 64 \ - --lora-alpha 128 \ - --lora-target-modules "${WAN_LORA_TARGET_MODULES[@]}" \ - --diffusion-init-lora-weight gaussian \ - --lr 1e-4 \ - --adam-beta2 0.999 \ - --diffusion-clip-range 1e-4 \ - --weight-decay 1e-4 \ - --use-miles-router \ - --sglang-server-concurrency 8 \ - --update-weight-buffer-size 2147483648 \ - --update-weight-target-module transformer,transformer_2 \ - --diffusion-reward pickscore:1.0 \ - --advantage-estimator grpo \ - --rm-type pickscore \ - --pickscore-num-workers 1 \ - --pickscore-num-gpus-per-worker 0 \ - --pickscore-batch-size 8 \ - --pickscore-processor-path laion/CLIP-ViT-H-14-laion2B-s32B-b79K \ - --pickscore-model-path yuvalkirstain/PickScore_v1 \ - --fsdp-master-dtype fp32 \ - --fsdp-reduce-dtype fp32 \ - --diffusion-forward-dtype bf16 \ - --diffusion-num-steps 10 \ - --diffusion-eval-num-steps 28 \ - --diffusion-output-num-frames "${NUM_FRAMES}" \ - --diffusion-guidance-scale 4.0 \ - --diffusion-guidance-scale-2 3.0 \ - --diffusion-noise-level 0.9 \ - --diffusion-height 480 \ - --diffusion-width 480 \ - --diffusion-flow-shift 3.0 \ - --diffusion-step-strategy-path miles.rollout.step_strategy_hub.epoch_global_window \ - --diffusion-sde-window-size 1 \ - --diffusion-sde-candidate-steps "${CANDIDATE_STEPS}" \ - --diffusion-debug-mode \ - --save "${ROOT_DIR}/logs/${RUN_NAME}/ckpt" \ - --save-interval 100 \ - --skip-eval-before-train \ - "${EXTRA_ARGS[@]}" \ - 2>&1 | tee "${ROOT_DIR}/logs/${RUN_NAME}.log" diff --git a/tests/sp/sglang_usp_import_guard.py b/tests/sp/sglang_usp_import_guard.py deleted file mode 100644 index 024d7d3b..00000000 --- a/tests/sp/sglang_usp_import_guard.py +++ /dev/null @@ -1,26 +0,0 @@ -"""Guard: the sglang tree actually imported must have the dtype/shape-aware checksum. - -Training's only remaining sglang import is compute_weights_checksum (weight-sync -verify); a bytes-only checksum would miss transpose/reshape corruption. The USP -attention operators are miles-owned (sp_ops.py) and need no sglang patches. - -Usage: PYTHONPATH=. python -m pytest tests/sp/sglang_usp_import_guard.py -""" - -import inspect - - -def test_checksum_covers_dtype_shape(): - from sglang.multimodal_gen.runtime.loader import weight_utils - - src = inspect.getsource(weight_utils.compute_weights_checksum) - assert "dtype" in src and "shape" in src, ( - f"imported sglang checksum is bytes-only: {weight_utils.__file__} — " "transpose/reshape would not be detected" - ) - - -if __name__ == "__main__": - test_checksum_covers_dtype_shape() - from sglang.multimodal_gen.runtime.loader import weight_utils - - print(f"[GUARD OK] imported sglang = {weight_utils.__file__}") diff --git a/tests/sp/sp_attention_parity.py b/tests/sp/sp_attention_parity.py deleted file mode 100644 index ab3f676d..00000000 --- a/tests/sp/sp_attention_parity.py +++ /dev/null @@ -1,145 +0,0 @@ -"""SP parity on a small real Wan DiT: USP attention + shard/gather vs full-sequence reference. - -Both paths use the same local attention kernel (SDPA), so any diff comes from -the SP collectives' float summation order and must stay within bf16 tolerance. -Checks forward output, input grad, and per-block self-attn + proj_out weight -grads, with and without gradient checkpointing. - -Usage: torchrun --standalone --nproc_per_node=N tests/sp/sp_attention_parity.py \ - --sp S [--ulysses U] [--ckpt] [--fp32] -""" - -import argparse - -import torch -import torch.distributed as dist -import torch.nn.functional as F -from diffusers import WanTransformer3DModel - -from miles.backends.fsdp_utils.configs.wan2_2 import Wan2_2TrainPipelineConfig, WanUSPAttnProcessor -from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend -from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state -from miles.backends.fsdp_utils.sp_attention import apply_sequence_parallel -from miles.utils.distributed_utils import init_gloo_group - -DTYPE = torch.bfloat16 - - -def build_model(device): - torch.manual_seed(0) - model = WanTransformer3DModel( - patch_size=(1, 2, 2), - num_attention_heads=8, - attention_head_dim=128, - in_channels=16, - out_channels=16, - text_dim=4096, - freq_dim=256, - ffn_dim=1024, - num_layers=2, - rope_max_seq_len=1024, - ) - model = model.to(device=device, dtype=DTYPE) - model.train() - for p in model.parameters(): - dist.broadcast(p.data, src=0) - return model - - -def make_inputs(device): - g = torch.Generator(device=device).manual_seed(123) - hidden = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) - enc = torch.randn(1, 32, 4096, device=device, dtype=DTYPE, generator=g) - ts = torch.tensor([500], device=device) - out_grad = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) - for t in (hidden, enc, out_grad): - dist.broadcast(t, src=0) - return hidden, enc, ts, out_grad - - -def _set_ref_processor(model): - # parallel_state=None: same processor, plain SDPA self-attention. - model.set_attn_processor(WanUSPAttnProcessor(None)) - - -def _run(model, hidden, enc, ts, out_grad, ckpt): - if ckpt: - model.enable_gradient_checkpointing() - else: - model.disable_gradient_checkpointing() - inp = hidden.clone().requires_grad_(True) - out = model(hidden_states=inp, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] - out.backward(out_grad) - return out.detach(), inp.grad.detach() - - -def _report(name, a, b, rtol, ctol): - rel = (a.float() - b.float()).abs().max() / b.float().abs().max().clamp_min(1e-6) - cos = F.cosine_similarity(a.float().flatten(), b.float().flatten(), dim=0) - ok = rel < rtol and cos > ctol - if dist.get_rank() == 0: - print(f" [{'OK' if ok else 'FAIL'}] {name:28s} rel={rel:.2e} cos={1 - cos:.2e}(1-)") - assert ok, f"{name}: rel={rel:.3e} cos={cos:.5f}" - - -def main(): - p = argparse.ArgumentParser() - p.add_argument("--sp", type=int, default=4) - p.add_argument("--ulysses", type=int, default=0) - p.add_argument("--ckpt", action="store_true") - p.add_argument("--fp32", action="store_true", help="fp32 + SDPA, isolates bf16 summation rounding") - cli = p.parse_args() - if cli.fp32: - global DTYPE - DTYPE = torch.float32 - - dist.init_process_group("nccl") - rank = dist.get_rank() - torch.cuda.set_device(rank % torch.cuda.device_count()) - device = torch.cuda.current_device() - init_gloo_group() - - args = argparse.Namespace( - sequence_parallel_size=cli.sp, - ulysses_degree=cli.ulysses, - ) - ps = create_fsdp_parallel_state(args) - - attn_weight_names = [ - f"blocks.{i}.attn1.{proj}.weight" for i in range(2) for proj in ("to_q", "to_k", "to_v", "to_out.0") - ] - attn_weight_names.append("proj_out.weight") - - model = build_model(device) - hidden, enc, ts, out_grad = make_inputs(device) - _set_ref_processor(model) - out_ref, gin_ref = _run(model, hidden, enc, ts, out_grad, cli.ckpt) - gw_ref = {n: dict(model.named_parameters())[n].grad.detach().clone() for n in attn_weight_names} - model.zero_grad(set_to_none=True) - - plan = DiffusersModelBackend(Wan2_2TrainPipelineConfig()).sequence_parallel_plan(model) - # operator-level parity uses plain slice-backward gather semantics - apply_sequence_parallel(model, ps, plan, sum_grad=False) - out_sp, gin_sp = _run(model, hidden, enc, ts, out_grad, cli.ckpt) - - # Each rank backprops only its 1/sp of the tokens; sum across sp restores full grads. - dist.all_reduce(gin_sp, group=ps.sp_group) - if rank == 0: - print(f"[PARITY] sp={cli.sp} ulysses={ps.ulysses_degree} ring={ps.ring_degree} ckpt={cli.ckpt}") - _report("forward(out)", out_sp, out_ref, rtol=2e-2, ctol=0.9990) - _report("grad(input)", gin_sp, gin_ref, rtol=4e-2, ctol=0.9980) - params = dict(model.named_parameters()) - for n in attn_weight_names: - assert params[n].grad is not None, f"{n} grad is None — all-to-all backward did not propagate" - g = params[n].grad.detach().clone() - dist.all_reduce(g, group=ps.sp_group) - _report(f"grad({n})", g, gw_ref[n], rtol=5e-2, ctol=0.9950) - - dist.barrier() - if rank == 0: - print(f"[PARITY OK] sp={cli.sp} u={ps.ulysses_degree} r={ps.ring_degree} ckpt={cli.ckpt}") - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/tests/sp/sp_determinism_smoke.py b/tests/sp/sp_determinism_smoke.py deleted file mode 100644 index a94dd0cd..00000000 --- a/tests/sp/sp_determinism_smoke.py +++ /dev/null @@ -1,108 +0,0 @@ -"""SP determinism smoke: the SP attention path (ulysses/ring) must be bitwise -deterministic under -torch.use_deterministic_algorithms (warn_only=False also asserts no op on the -path is registered nondeterministic). Runs forward+backward twice on identical -inputs and compares every grad bitwise. --no-det runs without the flag as a -control. - -torchrun --standalone --nproc_per_node=2 tests/sp/sp_determinism_smoke.py --sp 2 --ulysses 1 [--no-det] -""" - -import argparse -import os - -import torch -import torch.distributed as dist -from diffusers import WanTransformer3DModel - -from miles.backends.fsdp_utils.configs.wan2_2 import Wan2_2TrainPipelineConfig -from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend -from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state -from miles.backends.fsdp_utils.sp_attention import apply_sequence_parallel -from miles.utils.distributed_utils import init_gloo_group - -DTYPE = torch.bfloat16 - - -def main(): - p = argparse.ArgumentParser() - p.add_argument("--sp", type=int, default=2) - p.add_argument("--ulysses", type=int, default=1) - p.add_argument("--no-det", action="store_true") - cli = p.parse_args() - - if not cli.no_det: - assert os.environ.get("CUBLAS_WORKSPACE_CONFIG"), "set CUBLAS_WORKSPACE_CONFIG=:4096:8" - torch.backends.cudnn.deterministic = True - torch.backends.cudnn.benchmark = False - torch.use_deterministic_algorithms(True, warn_only=False) - - dist.init_process_group("nccl") - rank = dist.get_rank() - torch.cuda.set_device(rank % torch.cuda.device_count()) - device = torch.cuda.current_device() - init_gloo_group() - - args = argparse.Namespace( - sequence_parallel_size=cli.sp, - ulysses_degree=cli.ulysses, - ) - ps = create_fsdp_parallel_state(args) - - torch.manual_seed(0) - model = WanTransformer3DModel( - patch_size=(1, 2, 2), - num_attention_heads=8, - attention_head_dim=128, - in_channels=16, - out_channels=16, - text_dim=4096, - freq_dim=256, - ffn_dim=1024, - num_layers=2, - rope_max_seq_len=1024, - ).to(device=device, dtype=DTYPE) - model.train() - for prm in model.parameters(): - dist.broadcast(prm.data, src=0) - - g = torch.Generator(device=device).manual_seed(123) - hidden = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) - enc = torch.randn(1, 32, 4096, device=device, dtype=DTYPE, generator=g) - ts = torch.tensor([500], device=device) - out_grad = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) - for t in (hidden, enc, out_grad): - dist.broadcast(t, src=0) - - plan = DiffusersModelBackend(Wan2_2TrainPipelineConfig()).sequence_parallel_plan(model) - apply_sequence_parallel(model, ps, plan) - - def run_once(): - model.zero_grad(set_to_none=True) - out = model(hidden_states=hidden, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] - out.backward(out_grad) - return out.detach().clone(), {n: prm.grad.detach().clone() for n, prm in model.named_parameters()} - - out1, g1 = run_once() - out2, g2 = run_once() - - out_eq = torch.equal(out1, out2) - diffs = [n for n in g1 if not torch.equal(g1[n], g2[n])] - flag = torch.tensor([0 if (out_eq and not diffs) else 1], device=device) - dist.all_reduce(flag) - if rank == 0: - mode = "no-det(control)" if cli.no_det else "deterministic" - print( - f"[DET-PROBE] mode={mode} u{ps.ulysses_degree}r{ps.ring_degree} forward_bitwise={out_eq} grad_mismatches={len(diffs)}" - ) - for n in diffs[:8]: - d = (g1[n].float() - g2[n].float()).abs().max().item() - print(f" diff {n}: max_abs={d:.3e}") - print(f"[DET-PROBE {'OK' if flag.item() == 0 else 'NONDETERMINISTIC'}] (all ranks)") - if not cli.no_det: - assert flag.item() == 0, "SP attention path is not bitwise deterministic under deterministic mode" - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/tests/sp/sp_grad_sync_parity.py b/tests/sp/sp_grad_sync_parity.py deleted file mode 100644 index ca5086eb..00000000 --- a/tests/sp/sp_grad_sync_parity.py +++ /dev/null @@ -1,194 +0,0 @@ -"""SP gradient parity under real FSDP2: full grads must match a single-process -full-sequence reference, for both parameter placements. - -dp_sp (production): FSDP shards over the flattened dp x sp mesh; the sequence -gather's sum_grad backward makes FSDP's own reduce restore full grads. -dp (validation anchor, test-only): params shard over dp only (replicated -across sp), slice-backward gather, explicit cross-sp grad all-reduce here in -the test — the obvious mechanism cross-checking the subtle one. -Inputs are broadcast to all ranks, so the reference gradient is topology-free. -Also asserts model outputs are bitwise identical across sp ranks. - -Usage: torchrun --standalone --nproc_per_node=4 tests/sp/sp_grad_sync_parity.py \ - [--sp S --ulysses U] [--shard-mode dp|dp_sp] [--fp32] -""" - -import argparse - -import torch -import torch.distributed as dist -import torch.nn.functional as F -from diffusers import WanTransformer3DModel -from torch.distributed.tensor import DTensor - -from miles.backends.fsdp_utils.configs.wan2_2 import Wan2_2TrainPipelineConfig, WanUSPAttnProcessor -from miles.backends.fsdp_utils.model_backend import DiffusersModelBackend -from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state -from miles.backends.fsdp_utils.sp_attention import apply_sequence_parallel -from miles.utils.distributed_utils import init_gloo_group - -DTYPE = torch.bfloat16 - - -def build_model(device): - torch.manual_seed(0) - model = WanTransformer3DModel( - patch_size=(1, 2, 2), - num_attention_heads=8, - attention_head_dim=128, - in_channels=16, - out_channels=16, - text_dim=4096, - freq_dim=256, - ffn_dim=1024, - num_layers=2, - rope_max_seq_len=1024, - ).to(device=device, dtype=DTYPE) - model.train() - for p in model.parameters(): - dist.broadcast(p.data, src=0) - return model - - -def make_inputs(device, seed=123): - g = torch.Generator(device=device).manual_seed(seed) - hidden = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) - enc = torch.randn(1, 32, 4096, device=device, dtype=DTYPE, generator=g) - ts = torch.tensor([500], device=device) - out_grad = torch.randn(1, 16, 4, 8, 8, device=device, dtype=DTYPE, generator=g) - return hidden, enc, ts, out_grad - - -def main(): - p = argparse.ArgumentParser() - p.add_argument("--sp", type=int, default=4) - p.add_argument("--ulysses", type=int, default=0) - p.add_argument("--shard-mode", choices=("dp", "dp_sp"), default="dp_sp") - p.add_argument("--lora", action="store_true", help="train LoRA params only, like the RL recipe") - p.add_argument("--distinct-dp", action="store_true", help="different data per dp rank, like real RL") - p.add_argument("--accum", type=int, default=1, help="gradient-accumulation microbatches") - p.add_argument("--fp32", action="store_true", help="fp32 + SDPA, isolates bf16 summation rounding") - cli = p.parse_args() - if cli.fp32: - global DTYPE - DTYPE = torch.float32 - - dist.init_process_group("nccl") - rank = dist.get_rank() - torch.cuda.set_device(rank % torch.cuda.device_count()) - device = torch.cuda.current_device() - init_gloo_group() - - args = argparse.Namespace( - sequence_parallel_size=cli.sp, - ulysses_degree=cli.ulysses, - ) - ps = create_fsdp_parallel_state(args) - # dp anchor: shard over the dp submesh (sp-replicated params) + slice-backward gather. - fsdp_mesh = ps.dp_mesh if cli.shard_mode == "dp" else ps.fsdp_mesh - sum_grad = cli.shard_mode == "dp_sp" - # One dataset per (microbatch, dp group); seeds are rank-independent so the - # reference can rebuild every dataset locally. - datasets = [] - for mb in range(cli.accum): - seeds = [1000 + (d * 100 if cli.distinct_dp else 0) + mb for d in range(ps.dp_size)] - datasets.append([make_inputs(device, seed=s) for s in seeds]) - - def maybe_lora(m): - if not cli.lora: - return m - from peft import LoraConfig, get_peft_model - - torch.manual_seed(7) - return get_peft_model( - m, LoraConfig(r=8, lora_alpha=16, target_modules=["to_q", "to_k", "to_v"], init_lora_weights=False) - ) - - # Full-sequence single-process reference (plain SDPA self-attention). Inputs - # are broadcast, so every dp replica computes the same full gradient. - ref = maybe_lora(build_model(device)) - ref.set_attn_processor(WanUSPAttnProcessor(None)) - for mbsets in datasets: - for hidden, enc, ts, out_grad in mbsets: - out = ref(hidden_states=hidden, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] - out.backward(out_grad / ps.dp_size) - ref_grads = {n: p.grad.detach().clone() for n, p in ref.named_parameters() if p.grad is not None} - - from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard - - # fp32 reduce matches production apply_fsdp2 (--fsdp-reduce-dtype fp32). - mp_policy = MixedPrecisionPolicy(param_dtype=DTYPE, reduce_dtype=torch.float32) - model = maybe_lora(build_model(device)) - for blk in model.blocks: - fully_shard(blk, mesh=fsdp_mesh, mp_policy=mp_policy) - fully_shard(model, mesh=fsdp_mesh, mp_policy=mp_policy) - plan = DiffusersModelBackend(Wan2_2TrainPipelineConfig()).sequence_parallel_plan(model) - apply_sequence_parallel(model, ps, plan, sum_grad=sum_grad) - - for i, mbsets in enumerate(datasets): - hidden, enc, ts, out_grad = mbsets[ps.dp_rank] - out = model(hidden_states=hidden, timestep=ts, encoder_hidden_states=enc, return_dict=False)[0] - if i == 0: - o32 = out.detach().float() - ref0 = o32.clone() - dist.broadcast(ref0, src=ps.dp_rank * ps.sp_size, group=ps.sp_group) - diff = (o32 - ref0).abs().max() - if rank == 0: - print(f"[OUTPUT] max abs diff across sp ranks = {diff.item():.2e} (must be 0)") - assert diff.item() == 0.0, "model outputs diverge across sp ranks" - out.backward(out_grad) - - if cli.shard_mode == "dp": - # The anchor's explicit cross-sp grad sum; in dp_sp mode FSDP's own - # reduce-scatter already restores full grads (sum_grad gather). - for p in model.parameters(): - if p.grad is None: - continue - local = p.grad.to_local() if isinstance(p.grad, DTensor) else p.grad - f = local.float() - dist.all_reduce(f, group=ps.sp_group) - local.copy_(f.to(local.dtype)) - - fails = 0 - checked = 0 - for n, p in model.named_parameters(): - if p.grad is None: - continue - g = p.grad.full_tensor() if isinstance(p.grad, DTensor) else p.grad - r = ref_grads[n] - assert g.shape == r.shape, f"{n}: {g.shape} != {r.shape}" - rel = (g.float() - r.float()).abs().max() / r.float().abs().max().clamp_min(1e-6) - cos = F.cosine_similarity(g.float().flatten(), r.float().flatten(), dim=0) - checked += 1 - # Secondary band: small-norm tensors (biases) sit at the bf16 noise - # floor of this tiny test model (ring backward pushes cross-attn K - # biases to ~7e-2 rel; dp and dp_sp modes produce bit-identical - # values there, so it is accumulation noise, not placement). - if not (rel < 5e-2 and cos > 0.99) and not (rel < 1e-1 and cos > 0.995): - fails += 1 - if rank == 0: - print(f" [FAIL] {n:42s} rel={rel:.2e} cos={1 - cos:.2e}(1-)") - - # clip_grad_norm_ must report the same total norm as the reference. - total = torch.nn.utils.clip_grad_norm_(model.parameters(), 1e9) - total = total.full_tensor() if isinstance(total, DTensor) else total - ref_total = torch.linalg.vector_norm( - torch.stack([torch.linalg.vector_norm(g.float()) for g in ref_grads.values()]) - ) - norm_rel = ((total.float() - ref_total) / ref_total).abs() - - dist.barrier() - if rank == 0: - print(f"[GRAD-NORM] clip={total.item():.6e} ref={ref_total.item():.6e} rel={norm_rel.item():.2e}") - assert norm_rel.item() < 5e-2, "clip_grad_norm_ reports a wrong total norm" - print( - f"[SP-GRAD-SYNC] mode={cli.shard_mode} dp{ps.dp_size}xsp{ps.sp_size}" - f"(u{ps.ulysses_degree}r{ps.ring_degree}) checked={checked} fails={fails}" - ) - assert fails == 0 - print("[SP-GRAD-SYNC OK] full grads == full-sequence reference") - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/tests/sp/sp_init_smoke.py b/tests/sp/sp_init_smoke.py deleted file mode 100644 index 1005f6f2..00000000 --- a/tests/sp/sp_init_smoke.py +++ /dev/null @@ -1,71 +0,0 @@ -"""Multi-GPU SP init smoke test: NCCL group members must match the sp_mesh layout. - -Usage: torchrun --standalone --nproc_per_node=N tests/sp/sp_init_smoke.py --sp S [--ulysses U] -""" - -import argparse - -import torch -import torch.distributed as dist - -from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state -from miles.backends.fsdp_utils.sp_mesh import sp_subgroups -from miles.utils.distributed_utils import init_gloo_group - - -def _members(group): - n = dist.get_world_size(group) - t = torch.tensor([dist.get_rank()], device="cuda") - out = [torch.zeros_like(t) for _ in range(n)] - dist.all_gather(out, t, group=group) - return sorted(int(x.item()) for x in out) - - -def main(): - p = argparse.ArgumentParser() - p.add_argument("--sp", type=int, default=2) - p.add_argument("--ulysses", type=int, default=0) - cli = p.parse_args() - - dist.init_process_group("nccl") - rank, world = dist.get_rank(), dist.get_world_size() - torch.cuda.set_device(rank % torch.cuda.device_count()) - init_gloo_group() - - args = argparse.Namespace( - sequence_parallel_size=cli.sp, - ulysses_degree=cli.ulysses, - ) - ps = create_fsdp_parallel_state(args) - - dp_size, sp_size, sp_groups, ulysses_groups, ring_groups = sp_subgroups(world, cli.sp, cli.ulysses) - assert ps.sp_size == sp_size and ps.dp_size == dp_size - assert dist.get_world_size(ps.sp_group) == sp_size - assert dist.get_world_size(ps.dp_group) == dp_size - assert ps.fsdp_mesh.size() == world, f"fsdp mesh {ps.fsdp_mesh.size()} != world {world}" - assert ps.dp_mesh.size() == dp_size, f"dp mesh {ps.dp_mesh.size()} != dp {dp_size}" - - my_sp = next(g for g in sp_groups if rank in g) - assert _members(ps.sp_group) == my_sp, f"rank{rank} sp_group {_members(ps.sp_group)} != {my_sp}" - if ps.ulysses_degree > 1: - my_u = next(g for g in ulysses_groups if rank in g) - assert _members(ps.ulysses_group) == my_u, f"rank{rank} ulysses {_members(ps.ulysses_group)} != {my_u}" - else: - assert ps.ulysses_group is None - if ps.ring_degree > 1: - my_r = sorted(next(g for g in ring_groups if rank in g)) - assert _members(ps.ring_group) == my_r, f"rank{rank} ring {_members(ps.ring_group)} != {my_r}" - else: - assert ps.ring_group is None - - dist.barrier() - if rank == 0: - print( - f"[SP-INIT-SMOKE OK] world={world} dp={dp_size} sp={sp_size} " - f"ulysses={ps.ulysses_degree} ring={ps.ring_degree} fsdp_mesh={ps.fsdp_mesh.size()}" - ) - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/tests/sp/sp_weight_sync_parity.py b/tests/sp/sp_weight_sync_parity.py deleted file mode 100644 index 439854c6..00000000 --- a/tests/sp/sp_weight_sync_parity.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Weight-sync parity under FSDP + SP: every rank's reconstructed full weights must -hash identically and match a single-process reference. - -Reconstruction uses the same redistribute([Replicate()]) path as update_weights; -the checksum is sglang's compute_weights_checksum (name + dtype + shape + bytes), -the same function the online verify uses. Also asserts shape/dtype sensitivity. - -Usage: - torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 2 --ulysses 2 - torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 4 - torchrun --standalone --nproc_per_node=4 tests/sp/sp_weight_sync_parity.py --sp 4 --ulysses 2 -""" - -import argparse - -import torch -import torch.distributed as dist -from diffusers import WanTransformer3DModel -from sglang.multimodal_gen.runtime.loader.weight_utils import compute_weights_checksum -from torch.distributed.fsdp import fully_shard - -from miles.backends.fsdp_utils.parallel import create_fsdp_parallel_state -from miles.utils.distributed_utils import init_gloo_group - -DTYPE = torch.bfloat16 - - -def build_model(device): - torch.manual_seed(0) - model = WanTransformer3DModel( - patch_size=(1, 2, 2), - num_attention_heads=8, - attention_head_dim=128, - in_channels=16, - out_channels=16, - text_dim=4096, - freq_dim=256, - ffn_dim=1024, - num_layers=2, - rope_max_seq_len=1024, - ).to(device=device, dtype=DTYPE) - for p in model.parameters(): - dist.broadcast(p.data, src=0) - for b in model.buffers(): - dist.broadcast(b.data, src=0) - return model - - -def full_state_pairs(model): - """Mirror update_weights' redistribute([Replicate()]).to_local() reconstruction.""" - from torch.distributed.tensor import DTensor, Replicate - - pairs = [] - for name, param in model.state_dict().items(): - param = param.cuda() - if isinstance(param, DTensor): - param = param.redistribute(placements=[Replicate()] * param.device_mesh.ndim).to_local() - pairs.append((name, param.detach().cpu().contiguous())) - return pairs - - -def main(): - p = argparse.ArgumentParser() - p.add_argument("--sp", type=int, default=2) - p.add_argument("--ulysses", type=int, default=2) - p.add_argument("--shard-mode", choices=("dp", "dp_sp"), default="dp_sp") - cli = p.parse_args() - - dist.init_process_group("nccl") - rank = dist.get_rank() - world = dist.get_world_size() - torch.cuda.set_device(rank % torch.cuda.device_count()) - device = torch.cuda.current_device() - init_gloo_group() - - args = argparse.Namespace( - sequence_parallel_size=cli.sp, - ulysses_degree=cli.ulysses, - ) - ps = create_fsdp_parallel_state(args) - # dp anchor: shard over the dp submesh (sp-replicated params). - fsdp_mesh = ps.dp_mesh if cli.shard_mode == "dp" else ps.fsdp_mesh - - ref_model = build_model(device) - ref_sum = compute_weights_checksum( - [(n, pa.detach().cpu().contiguous()) for n, pa in ref_model.state_dict().items()] - ) - - model = build_model(device) - for blk in model.blocks: - fully_shard(blk, mesh=fsdp_mesh) - fully_shard(model, mesh=fsdp_mesh) - - my_sum = compute_weights_checksum(full_state_pairs(model)) - - gathered = [None] * world - dist.all_gather_object(gathered, my_sum) - - if rank == 0: - all_equal = all(s == gathered[0] for s in gathered) - match_ref = gathered[0] == ref_sum - print( - f"[WEIGHT-SYNC] world={world} mode={cli.shard_mode} dp{ps.dp_size}xsp{ps.sp_size}(u{ps.ulysses_degree}r{ps.ring_degree})" - ) - print(f"[WEIGHT-SYNC] all ranks equal={all_equal} == single-process ref={match_ref}") - assert all_equal, "reconstructed full weights differ across ranks" - assert match_ref, "reconstructed full weights != single-process reference" - - # reshape keeps bytes identical, so only the shape term can catch it - t = torch.arange(12, dtype=torch.float32).reshape(3, 4) - assert compute_weights_checksum([("w", t)]) != compute_weights_checksum([("w", t.reshape(4, 3))]) - assert compute_weights_checksum([("w", t)]) != compute_weights_checksum([("w", t.to(torch.float64))]) - print("[SP-WEIGHT-SYNC OK] bitwise-identical reconstruction + dtype/shape sensitivity") - - dist.barrier() - dist.destroy_process_group() - - -if __name__ == "__main__": - main() diff --git a/tests/sp/test_sp_mesh.py b/tests/sp/test_sp_mesh.py deleted file mode 100644 index 21733a55..00000000 --- a/tests/sp/test_sp_mesh.py +++ /dev/null @@ -1,73 +0,0 @@ -"""SP rank layout / subgroup unit tests. Pure functions, no GPU.""" - -import pytest - -from miles.backends.fsdp_utils.sp_mesh import locate_rank, resolve_sp_degrees, sp_subgroups, validate_sp_config - - -def test_resolve_auto_degrees(): - assert resolve_sp_degrees(4) == (4, 4, 1) - assert resolve_sp_degrees(4, ulysses_degree=2) == (4, 2, 2) - assert resolve_sp_degrees(8, 2) == (8, 2, 4) - assert resolve_sp_degrees(1) == (1, 1, 1) - - -def test_resolve_illegal(): - with pytest.raises(ValueError): - resolve_sp_degrees(4, ulysses_degree=3) - - -def test_sglang_alignment_example(): - # Matches sglang set_seq_parallel_pg_by_sp_groups: sp=4,u=2,r=2 - _, _, sp_groups, ulysses_groups, ring_groups = sp_subgroups(4, 4, 2) - assert sp_groups == [[0, 1, 2, 3]] - assert ulysses_groups == [[0, 1], [2, 3]] - assert ring_groups == [[0, 2], [1, 3]] - - -@pytest.mark.parametrize( - "world,sp,u,r", - [ - (2, 2, 2, 1), - (4, 2, 2, 1), - (4, 4, 2, 2), - (4, 4, 4, 1), - (8, 4, 2, 2), - (16, 8, 2, 4), - (64, 8, 8, 1), - (256, 16, 8, 2), - (1024, 8, 8, 1), - ], -) -def test_layout_invariants_any_scale(world, sp, u, r): - dp_size, sp_size, sp_groups, ulysses_groups, ring_groups = sp_subgroups(world, sp, u) - assert sp_size == sp == u * r # r in the parametrize row is the expected derived ring degree - assert dp_size * sp_size == world - assert len(sp_groups) == dp_size and all(len(g) == sp for g in sp_groups) - # ulysses groups: contiguous, size u, exact cover - assert all(len(g) == u and g == list(range(g[0], g[0] + u)) for g in ulysses_groups) - assert sorted(x for g in ulysses_groups for x in g) == list(range(world)) - # ring groups: strided, size r, exact cover - assert all(len(g) == r for g in ring_groups) - assert sorted(x for g in ring_groups for x in g) == list(range(world)) - - -@pytest.mark.parametrize("world,sp,u,r", [(8, 4, 2, 2), (16, 8, 2, 4), (256, 16, 8, 2)]) -def test_locate_rank_consistent_with_subgroups(world, sp, u, r): - _, _, _, ulysses_groups, ring_groups = sp_subgroups(world, sp, u) - for rank in range(world): - dp_rank, sp_rank, u_rank, r_rank = locate_rank(rank, sp, u) - assert dp_rank == rank // sp and sp_rank == rank % sp - my_u_group = next(g for g in ulysses_groups if rank in g) - my_r_group = next(g for g in ring_groups if rank in g) - assert my_u_group[u_rank] == rank - assert my_r_group[r_rank] == rank - - -def test_validate_rejects_illegal(): - with pytest.raises(ValueError): - validate_sp_config(world_size=6, sequence_parallel_size=4) - with pytest.raises(ValueError): - validate_sp_config(world_size=8, sequence_parallel_size=4, ulysses_degree=3) - assert validate_sp_config(world_size=4, sequence_parallel_size=2) == (2, 2, 1) - assert validate_sp_config(world_size=4, sequence_parallel_size=4) == (4, 4, 1) From 3676574de7e9c21cbf3e4595d21cd02983208147 Mon Sep 17 00:00:00 2001 From: Andy Ye Date: Mon, 13 Jul 2026 23:54:01 +0000 Subject: [PATCH 24/24] refactor: sp_plan.py owns the plan contract; sp_attention.py is dispatch-only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit SequenceParallelPlan, the boundary-hook interpreter and apply_sequence_parallel move to sp_plan.py; sp_attention.py keeps only the dispatch-level attention installer. The two are siblings over sp_ops — the plan holds attention as an opaque callable, so neither imports the other, and the file layout now matches the dependency structure (model_backend imports the contract, the family-config default imports the installer). Pure code motion, no behavior change. Co-Authored-By: Claude Fable 5 --- miles/backends/fsdp_utils/actor.py | 2 +- miles/backends/fsdp_utils/model_backend.py | 2 +- miles/backends/fsdp_utils/sp_attention.py | 132 +-------------------- miles/backends/fsdp_utils/sp_plan.py | 131 ++++++++++++++++++++ 4 files changed, 139 insertions(+), 128 deletions(-) create mode 100644 miles/backends/fsdp_utils/sp_plan.py diff --git a/miles/backends/fsdp_utils/actor.py b/miles/backends/fsdp_utils/actor.py index 49e5e203..0ecf29f3 100644 --- a/miles/backends/fsdp_utils/actor.py +++ b/miles/backends/fsdp_utils/actor.py @@ -31,7 +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_attention import apply_sequence_parallel +from .sp_plan import apply_sequence_parallel logger = logging.getLogger(__name__) diff --git a/miles/backends/fsdp_utils/model_backend.py b/miles/backends/fsdp_utils/model_backend.py index 90dced83..c4ace19f 100644 --- a/miles/backends/fsdp_utils/model_backend.py +++ b/miles/backends/fsdp_utils/model_backend.py @@ -21,7 +21,7 @@ import torch from diffusers import DiffusionPipeline -from .sp_attention import SequenceParallelPlan +from .sp_plan import SequenceParallelPlan class ModelBackend(abc.ABC): diff --git a/miles/backends/fsdp_utils/sp_attention.py b/miles/backends/fsdp_utils/sp_attention.py index 470c295d..d1eda772 100644 --- a/miles/backends/fsdp_utils/sp_attention.py +++ b/miles/backends/fsdp_utils/sp_attention.py @@ -1,122 +1,15 @@ -"""Sequence parallelism for diffusion DiTs: self-attention runs USP (sp_ops). +"""Dispatch-level USP attention: the default installer for diffusers models. -Each sp rank holds S/sp latent tokens; attention internally gathers to the full -sequence via Ulysses all-to-all (+Ring) and scatters back. A model family opts in -through a ``SequenceParallelPlan``; model outputs stay full-sequence, so every -parameter sees a partial gradient and loss/log_prob code is untouched. +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 inspect import sys -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, usp_attention - - -@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) +from .sp_ops import usp_attention class _USPDispatchConfig: @@ -178,16 +71,3 @@ def apply_dispatch_sp_attention(transformer, parallel_state): config = _USPDispatchConfig(parallel_state) for processor in base.attn_processors.values(): processor._parallel_config = config - - -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/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)