diff --git a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx index fe3c14808..164608275 100644 --- a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx +++ b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx @@ -842,7 +842,7 @@ the strict `quality="high"` deployment contract. - The released visual VAE quality recipe uses overlapping tiled decode. SGLang keeps that recipe by default and distributes complete tiles across the decode group; this changes scheduling, not the computation inside each tile. - H3 rejects `--vae-config.parallel-decode-mode spatial` and `spatial_shard`: validation found output mismatches. Use the default released tiled recipe. - Keep the default `--encoder-parallel auto`. With the server’s default `batching_max_size` of 1, single-node H100/H200/B200/B300 recipes with peer-to-peer access fold the Qwen text encoder over otherwise idle Ulysses ranks. This is separate from DiT tensor parallelism. A pure-TP recipe already shards the encoder over its TP group and does not add a world fold. -- For throughput-oriented serving, select **DP (batched throughput)**. The picker pairs `--encoder-parallel dp` with an editable `--batching-max-size` greater than 1; compatible requests are distributed across ranks, while every rank keeps a full encoder replica. Encoder DP requires TP1 and DiT DP1, so it is disabled for the H100 TP2 + Ulysses2 and RTX 5090 TP2 recipes. It provides no benefit for a batch of one and is not bitwise-identical to the folded deployment. +- For throughput-oriented serving, select **DP (batched throughput)**. The picker pairs `--encoder-parallel dp` with an editable `--batching-max-size` greater than 1. Encoder DP stays inside each DiT replica and composes with encoder TP: the H100 TP2 + Ulysses2 recipe has two TP-sharded encoder copies that can split a batch, while the RTX 5090 pure-TP2 recipe has one encoder copy and therefore no additional batch-DP degree. It provides no benefit for a batch of one and is not bitwise-identical to the unsplit deployment. - Use explicit **Fold** to prioritize single-request latency and encoder memory on a measured high-bandwidth single-node topology. Use **Replicate** as the compatibility path when folding or encoder DP is unsuitable. - `--use-fsdp-inference true` shards only the DiT. MiniMax-H3 preserves the original FP32 dtype of its patch, time, and output projections during FSDP all-gather, so this path does not trade numerical correctness for memory. On 4×H100, prefer TP2 + Ulysses2 for speed; use FSDP as an explicit capacity policy rather than assuming it is faster. - `speed` keeps model components resident, while `auto` applies the model-aware 120 GiB residency threshold. `memory` prioritizes avoiding OOM and includes the executable VAE decoder in its default layerwise set. A measured recipe with sufficient headroom can opt into `--component-residency vae=resident`; the 2×H100 CI recipe does this because the VAE's 4.8 GiB/GPU cost avoids repeated decoder transfers during tiled decode. DiT residency and prefetch knobs remain scoped to the DiT. Use `speed` only after confirming that the complete target workload fits. diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 78e46b81e..1e794cea8 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -92,7 +92,7 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis - `--ulysses-degree {N}` and `--ring-degree {N}`: USP parallelism controls - `--kv-gather-degree {N}`: sequence-parallel degree that splits rows inside attention and exchanges with one K/V all-gather (queries stay local) instead of Ulysses all-to-all. Non-causal attention only; does not compose with `--ulysses-degree`/`--ring-degree` yet. When no SP degree is set explicitly, `sp_degree=2` defaults to `kv_gather_degree=2` (its measured-win zone) and higher degrees default to Ulysses; under that auto assignment, attention calls the gather path cannot take fall back to the Ulysses exchange, while an explicit degree fails instead of degrading. - `--enable-cfg-parallel {true|false}`: enable or explicitly disable CFG parallelism -- `--encoder-parallel {auto|fold|dp|replicate}`: how the text/image encoders use the GPUs the DiT replica leaves idle during encoding. `auto` (the default for both `generate` and `serve`) TP-folds an encoder wide enough to pay for the per-layer all-reduce, selects DP for a server batch when it can engage, and otherwise replicates; `fold` forces the shard whenever the dims allow it; `dp` splits a batched encode across ranks and needs `--batching-max-size > 1` to engage; `replicate` encodes redundantly on every rank. `fold` and `replicate` are bitwise-identical to single-GPU encoding. See [Encoder Parallelism](/docs/sglang-diffusion/encoder_parallel). +- `--encoder-parallel {auto|fold|dp|replicate}`: how text/image encoders use the GPUs in each DiT replica. `auto` TP-folds an encoder wide enough to benefit, selects batch DP when it can engage, and otherwise keeps the existing encoder TP layout; `fold` shards across the full replica whenever dimensions allow; `dp` splits a batched encode across encoder copies and composes with encoder TP; `replicate` disables folding and batch DP. Encoder collectives never cross `--dp-size` replicas. See [Encoder Parallelism](/docs/sglang-diffusion/encoder_parallel). - `--warmup-mode {off|request|server}`: control startup warmup for `sglang serve`; `off` skips warmup, `request` primes the request path, and `server` runs a full synthetic server warmup before serving traffic - `--enable-torch-compile {true|false}`: compile native diffusion hot paths. When no warmup mode is configured, this also enables server warmup so first real requests do not pay compile latency. - `--offload-during-compile {true|false}`: when compile warmup is active, temporarily layerwise-offload DiT weights and move resident non-DiT components off-device so `max-autotune` fits on tighter-memory GPUs; the configured serving residency is restored before real traffic. Skipped under existing layerwise offload, Cache-DiT, or FSDP. diff --git a/docs/docs/sglang-diffusion/encoder_parallel.mdx b/docs/docs/sglang-diffusion/encoder_parallel.mdx index f8540a17c..610db498d 100644 --- a/docs/docs/sglang-diffusion/encoder_parallel.mdx +++ b/docs/docs/sglang-diffusion/encoder_parallel.mdx @@ -15,10 +15,10 @@ those otherwise-unused GPUs for the encoding stage. | Mode | What it does | Use when | | --- | --- | --- | -| `auto` | Picks `fold`, `dp`, or `replicate` per encoder from its width and the request's batch width | Default for `generate`; you want the decision made per encoder | +| `auto` | Picks `fold`, `dp`, or the existing encoder layout from its width and the request's batch width | Default; you want the decision made per encoder | | `fold` | TP-shards the encoder weights across the idle DiT replica | One wide encoder dominates a single-request encode | -| `dp` | Each rank encodes its slice of the prompt batch, then the outputs are all-gathered | Default for `serve`; needs `--batching-max-size > 1` to engage | -| `replicate` | Every rank encodes the whole batch redundantly | You want the encoding stage to match single-GPU numerics exactly | +| `dp` | Encoder copies split the prompt batch, then all-gather their outputs inside the replica | Throughput serving with `--batching-max-size > 1` | +| `replicate` | Keeps the encoder on its DiT TP group and encodes redundantly across the other replica ranks | You want to disable folding and batch DP | The two accelerated modes are mutually exclusive per encoder: folding shards the weights for the lifetime of the loaded model, so a folded encoder cannot also be @@ -44,32 +44,31 @@ configuration you measured yourself. ## Numerics -`fold` and `replicate` are bitwise-identical to single-GPU encoding: folding -shards a GEMM and reduces it, which is the same arithmetic the unsharded kernel -performs. +`replicate` matches single-GPU encoding bit-for-bit only when the encoder TP +degree is one. Both the existing DiT TP layout and `fold` can reorder parallel +reductions; they are mathematically equivalent but are not generally bitwise +identical to a single-GPU kernel. -`dp` is **not** bitwise-identical. Each rank runs the full unsharded encoder on -a smaller batch, so the GEMM tiling and reduction order differ from the batched -reference — the same floating-point reordering class as choosing a different -attention backend or parallelism strategy, not a precision loss. The gathered -result is mathematically equivalent, and per-request results stay deterministic -for a fixed batch shape, but embeddings will not match a `replicate` run -bit-for-bit, and long video sampling can amplify the difference into visible -frame differences. Use `replicate` (or `fold`) when you need bit-exact -reproducibility against a single-GPU reference, e.g. when refreshing consistency -baselines. +`dp` is also not bitwise-identical: each encoder copy sees a smaller batch, so +GEMM tiling can differ from the unsplit reference. It may compose with encoder +TP: all ranks in one TP group receive the same batch slice, and corresponding TP +ranks gather outputs across the orthogonal encoder-DP group. The result remains +mathematically equivalent and deterministic for a fixed topology and batch +shape, but long video sampling can amplify small floating-point differences. ## Recommended Commands -Throughput serving. `serve` already defaults to `dp`, but a single encode call -must carry more than one prompt for it to engage, so raise the batching ceiling -too — an encoder flag deliberately does not change DiT batching for you: +Throughput serving. A single encode call must carry more than one prompt for DP +to engage, so raise the batching ceiling too; an encoder flag deliberately does +not change DiT batching for you: ```bash sglang serve \ - --model-path Qwen/Qwen-Image-2512 \ + --model-path Wan-AI/Wan2.2-TI2V-5B-Diffusers \ --model-type diffusion \ --num-gpus 2 \ + --tp-size 1 \ + --ulysses-degree 2 \ --encoder-parallel dp \ --batching-max-size 2 ``` @@ -97,10 +96,16 @@ sglang serve \ ## Interaction With Other Flags -- **Tensor / data parallel**: `dp` requires a replicated encoder, so it is - skipped when `--tp-size > 1` or `--dp-size > 1`. -- **Dynamic batching**: `dp` only pays with a wide batch, so selecting it raises - the default batching ceiling. See [Inference Batching](./dynamic_batching). -- **Sequence parallelism**: independent — SP splits the DiT's latent sequence, - encoder parallelism splits the encoding stage. See +- **Tensor parallel**: encoder DP composes with TP. A TP group jointly encodes + one batch slice; the orthogonal ranks inside the same pipeline replica split + and gather the batch. A pure-TP replica has one encoder copy, so there is no + additional batch-DP degree. +- **Data parallel**: encoder collectives never cross pipeline replicas. With + `--dp-size > 1`, each replica independently uses its own TP/SP/CFG ranks. +- **Dynamic batching**: `dp` only pays with a wide batch. Selecting it does not + change `--batching-max-size`; configure that separately. See + [Inference Batching](./dynamic_batching). +- **Sequence parallelism**: SP splits the DiT latent sequence. During encoding, + those ranks either hold encoder copies for batch DP or join a folded encoder. + See [Sequence Parallelism](./ring_sp_performance). diff --git a/python/sglang/multimodal_gen/configs/models/encoders/base.py b/python/sglang/multimodal_gen/configs/models/encoders/base.py index cad6e764e..a43b713f8 100644 --- a/python/sglang/multimodal_gen/configs/models/encoders/base.py +++ b/python/sglang/multimodal_gen/configs/models/encoders/base.py @@ -78,9 +78,9 @@ class EncoderConfig(ModelConfig): # Parallel folding: during the encoding stage the whole DiT replica is idle, # so TP-shard the encoder across those otherwise-unused GPUs instead of - # running it on a single rank. None = replicated, else the group to fold - # over ("sp"|"world"); resolved by finalize_encoder_folding. - parallel_folding_mode: Literal["sp", "world"] | None = None + # leaving it on the DiT TP group. None keeps that TP group; the other modes + # override it with a wider group selected by finalize_encoder_folding. + parallel_folding_mode: Literal["sp", "world", "replica"] | None = None @dataclass diff --git a/python/sglang/multimodal_gen/runtime/distributed/__init__.py b/python/sglang/multimodal_gen/runtime/distributed/__init__.py index 276615cce..91f6fc8c7 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/__init__.py +++ b/python/sglang/multimodal_gen/runtime/distributed/__init__.py @@ -12,6 +12,8 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_dp_group, get_dp_rank, get_dp_world_size, + get_encoder_data_parallel_group, + get_replica_group, get_sp_group, get_sp_parallel_rank, get_sp_world_size, @@ -44,6 +46,8 @@ __all__ = [ "get_world_size", # Data parallel group "get_dp_group", + "get_replica_group", + "get_encoder_data_parallel_group", "get_dp_rank", "get_dp_world_size", # Sequence parallel group diff --git a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py index cd4efcec5..70e0a4cb2 100644 --- a/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py +++ b/python/sglang/multimodal_gen/runtime/distributed/parallel_state.py @@ -66,10 +66,18 @@ _SP: SequenceParallelGroupCoordinator | None = None _PP: PipelineGroupCoordinator | None = None _CFG: GroupCoordinator | None = None _DP: GroupCoordinator | None = None +# all ranks serving one pipeline replica (every dim except dp); with +# dp_size == 1 it coincides with the world group +_REPLICA: GroupCoordinator | None = None +# Corresponding TP lanes across the replicated encoder copies in one pipeline +# replica. None means the encoder has only one TP copy, so batch DP cannot run. +_ENCODER_DP: GroupCoordinator | None = None _VAE_DECODE: GroupCoordinator | None = None _DIT: ProcessGroup | None = None _VAE: ProcessGroup | None = None _VAE_DECODE_PARALLEL_AXES = "tp-sp-pp-cfg" +_REPLICA_PARALLEL_AXES = "tp-sp-pp-cfg" +_ENCODER_DP_PARALLEL_AXES = "sp-pp-cfg" TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) @@ -187,6 +195,8 @@ def init_parallel_group_coordinator( "sequence", "classifier_free_guidance", "vae_decode", + "replica", + "encoder_data", ], f"parallel_mode {parallel_mode} is not supported" if parallel_mode == "pipeline": return PipelineGroupCoordinator( @@ -204,19 +214,19 @@ def init_parallel_group_coordinator( **kwargs, ) else: + group_name = { + "tensor": "tp_group", + "vae_decode": "vae_decode_group", + "replica": "replica_group", + "encoder_data": "encoder_dp_group", + }.get(parallel_mode, "cfg_group") return GroupCoordinator( group_ranks=group_ranks, local_rank=local_rank, torch_distributed_backend=backend, use_device_communicator=parallel_mode != "tensor", use_srt_custom_allreduce=parallel_mode == "tensor", - group_name=( - "tp_group" - if parallel_mode == "tensor" - else ( - "vae_decode_group" if parallel_mode == "vae_decode" else "cfg_group" - ) - ), + group_name=group_name, ) @@ -228,6 +238,12 @@ def _get_vae_decode_group_ranks( return rank_generator.get_ranks(_VAE_DECODE_PARALLEL_AXES) +def _get_encoder_data_parallel_group_ranks( + rank_generator: RankGenerator, +) -> list[list[int]]: + return rank_generator.get_ranks(_ENCODER_DP_PARALLEL_AXES) + + def get_tp_group() -> GroupCoordinator: assert _TP is not None, "tensor model parallel group is not initialized" return _TP @@ -324,6 +340,35 @@ def get_dp_group() -> GroupCoordinator: return _DP +def get_replica_group() -> GroupCoordinator: + """The ranks that share this rank's request batch — its pipeline replica. + + Replica-wide encoder folding uses this group rather than the world group, + which spans replicas when dp_size > 1. The encoder batch-DP group is the + orthogonal non-TP subset of this replica. A process that only initialized + the world group has exactly one replica by construction, so the world group + doubles as it.""" + if _REPLICA is not None: + return _REPLICA + return get_world_group() + + +def get_encoder_data_parallel_group() -> GroupCoordinator | None: + """Encoder copies that may split a batch inside one pipeline replica. + + Native encoders use the DiT TP group when they are not folded. The batch-DP + group is therefore the orthogonal non-TP part of the replica: ranks with the + same TP coordinate receive different batch slices, while every rank in one + TP group receives the same slice. A pure-TP replica has only one encoder + copy and returns None. + """ + if _ENCODER_DP is not None: + return _ENCODER_DP + if not model_parallel_is_initialized(): + return get_world_group() + return None + + # xDiT def initialize_model_parallel( data_parallel_size: int = 1, @@ -486,6 +531,64 @@ def initialize_model_parallel( ) _sync_srt_tp_group() + global _REPLICA + assert _REPLICA is None, "replica group is already initialized" + if data_parallel_size == 1: + _REPLICA = get_world_group() + else: + replica_axis_groups = [ + group + for degree, group in ( + (tensor_parallel_degree, _TP), + (sequence_parallel_degree, _SP), + (pipeline_parallel_degree, _PP), + (classifier_free_guidance_degree, _CFG), + ) + if degree > 1 + ] + if len(replica_axis_groups) <= 1: + # One nontrivial axis already is the replica group. If every axis + # is size one, any singleton axis group represents the replica. + _REPLICA = replica_axis_groups[0] if replica_axis_groups else _TP + else: + _REPLICA = init_parallel_group_coordinator( + group_ranks=rank_generator.get_ranks(_REPLICA_PARALLEL_AXES), + local_rank=get_world_group().local_rank, + backend=backend, + parallel_mode="replica", + ) + + global _ENCODER_DP + assert _ENCODER_DP is None, "encoder data parallel group is already initialized" + encoder_dp_degree = ( + sequence_parallel_degree + * pipeline_parallel_degree + * classifier_free_guidance_degree + ) + if tensor_parallel_degree == 1: + # With no encoder TP sharding, every replica rank owns a full encoder. + _ENCODER_DP = _REPLICA + elif encoder_dp_degree > 1: + nontrivial_groups = [ + group + for degree, group in ( + (sequence_parallel_degree, _SP), + (pipeline_parallel_degree, _PP), + (classifier_free_guidance_degree, _CFG), + ) + if degree > 1 + ] + if len(nontrivial_groups) == 1: + # Reuse the existing axis group in the common TP x SP case. + _ENCODER_DP = nontrivial_groups[0] + else: + _ENCODER_DP = init_parallel_group_coordinator( + group_ranks=_get_encoder_data_parallel_group_ranks(rank_generator), + local_rank=get_world_group().local_rank, + backend=backend, + parallel_mode="encoder_data", + ) + global _VAE_DECODE assert _VAE_DECODE is None, "VAE decode parallel group is already initialized" _VAE_DECODE = init_parallel_group_coordinator( @@ -920,7 +1023,7 @@ def init_vae_group( def destroy_model_parallel() -> None: """Set the groups to none and destroy them.""" - global _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE + global _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE, _REPLICA, _ENCODER_DP _clear_srt_tp_group() # The IPC transport keeps CUDA mappings associated with the current @@ -930,9 +1033,25 @@ def destroy_model_parallel() -> None: IPC_A2A.reset() - for group in (_TP, _SP, _DP, _CFG, _PP, _VAE_DECODE): - if group is not None: + destroyed_groups = [] + for group in ( + _TP, + _SP, + _DP, + _CFG, + _PP, + _REPLICA, + _ENCODER_DP, + _VAE_DECODE, + ): + # Replica-derived groups may alias each other or the world group. + if ( + group is not None + and group is not _WORLD + and all(group is not destroyed for destroyed in destroyed_groups) + ): group.destroy() + destroyed_groups.append(group) # Ulysses and Ring groups are created separately from the SP coordinator, # so GroupCoordinator.destroy() does not own or release them. Explicitly @@ -954,4 +1073,6 @@ def destroy_model_parallel() -> None: if group is not None: torch.distributed.destroy_process_group(group) - _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE = (None,) * 8 + _TP, _SP, _DP, _CFG, _PP, _VAE_DECODE, _DIT, _VAE, _REPLICA, _ENCODER_DP = ( + None, + ) * 10 diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py index d79ce8cd2..1c395f2c6 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py @@ -15,6 +15,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImageEditPipelineConfig, ) from sglang.multimodal_gen.runtime.distributed import ( + get_encoder_data_parallel_group, get_local_torch_device, ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( @@ -420,17 +421,19 @@ class TextEncoderLoader(ComponentLoader): model_config, component_model_path, ) + encoder_dp_group = get_encoder_data_parallel_group() + prefer_dp = ( + server_args.batching_max_size > 1 + and encoder_dp_group is not None + and encoder_dp_group.world_size > 1 + and issubclass(model_cls, TextEncoder) + and model_cls.supports_dp_encode + ) # real dims are populated now; resolve fold vs replicate finalize_encoder_folding( encoder_config, server_args.encoder_parallel, - prefer_dp=( - server_args.batching_max_size > 1 - and (server_args.tp_size or 1) == 1 - and (server_args.dp_size or 1) == 1 - and issubclass(model_cls, TextEncoder) - and model_cls.supports_dp_encode - ), + prefer_dp=prefer_dp, ) encoder_dtype = server_args.pipeline_config.text_encoder_precisions[ encoder_index diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/base.py b/python/sglang/multimodal_gen/runtime/models/encoders/base.py index 0d8c34f73..fc8b486ed 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/base.py +++ b/python/sglang/multimodal_gen/runtime/models/encoders/base.py @@ -14,6 +14,7 @@ from sglang.multimodal_gen.configs.models.encoders import ( TextEncoderConfig, ) from sglang.multimodal_gen.runtime.distributed import ( + get_replica_group, get_sp_group, get_tp_group, get_world_group, @@ -36,6 +37,10 @@ def get_folding_tp_group(config: EncoderConfig): if mode == "world": # the whole single-replica DiT (all GPUs), regardless of tp/sp/cfg. return get_world_group() + if mode == "replica": + # the ranks serving this rank's pipeline replica (== world when + # dp_size is 1); the shape-independent choice for explicit folding + return get_replica_group() if mode is None: return get_tp_group() raise ValueError(f"Unsupported encoder folding mode: {mode!r}") diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py index c1bd476f8..025e8d4d9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py @@ -23,7 +23,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) -_MINIMAX_H3_SINGLE_RANK_TEXT_ENCODE_EXTRA_KEY = "minimax_h3_single_rank_text_encode" +_MINIMAX_H3_SINGLE_COPY_TEXT_ENCODE_EXTRA_KEY = "minimax_h3_single_copy_text_encode" class MiniMaxH3TextEncodingStage(TextEncodingStage): @@ -108,13 +108,14 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage): batches: list[Req], server_args: ServerArgs, ) -> list[Req]: - """Distribute independent H3 presentations over replicated encoders. + """Distribute independent H3 presentations over encoder copies. H3 presentations have variable multimodal layouts, so they cannot be stacked into the generic text batch without changing padding/kernels. - Assigning one complete request to a rank preserves the exact - single-request encoder path, then broadcasts that request's native - payload to the other ranks. + Assigning one complete request to a copy preserves the exact + single-request encoder path. A copy may itself span a TP group; the + orthogonal encoder-DP group then broadcasts that request's native + payload to the other copies. """ grouped = self._group_requests_by_fingerprint( batches, @@ -142,14 +143,14 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage): payload = None if dp_group.rank_in_group == owner: try: - first_batch.extra[_MINIMAX_H3_SINGLE_RANK_TEXT_ENCODE_EXTRA_KEY] = ( + first_batch.extra[_MINIMAX_H3_SINGLE_COPY_TEXT_ENCODE_EXTRA_KEY] = ( True ) try: first_result = self(first_batch, server_args) finally: first_batch.extra.pop( - _MINIMAX_H3_SINGLE_RANK_TEXT_ENCODE_EXTRA_KEY, None + _MINIMAX_H3_SINGLE_COPY_TEXT_ENCODE_EXTRA_KEY, None ) payload = first_result.extra.get( MINIMAX_H3_TEXT_EMBEDDINGS_EXTRA_KEY @@ -191,7 +192,7 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage): self._dp_choice_logged = True logger.info( "encoder_parallel: distributing %d independent MiniMax H3 " - "presentations over %d replicated encoder ranks", + "presentations over %d encoder copies", batch_size, world_size, ) @@ -381,7 +382,7 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage): share_across_replicas=( world > 1 and not bool( - batch.extra.get(_MINIMAX_H3_SINGLE_RANK_TEXT_ENCODE_EXTRA_KEY) + batch.extra.get(_MINIMAX_H3_SINGLE_COPY_TEXT_ENCODE_EXTRA_KEY) ) ), ) @@ -536,7 +537,7 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage): pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, ) - if batch.extra.get(_MINIMAX_H3_SINGLE_RANK_TEXT_ENCODE_EXTRA_KEY): + if batch.extra.get(_MINIMAX_H3_SINGLE_COPY_TEXT_ENCODE_EXTRA_KEY): batch.extra.pop(MINIMAX_H3_PREPARED_REFERENCE_VIDEO_EXTRA_KEY, None) return { "positive": { diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py index 04cb9dc38..509126830 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_encoding.py @@ -17,8 +17,8 @@ import torch from sglang.multimodal_gen.configs.models.encoders import BaseEncoderOutput from sglang.multimodal_gen.configs.pipeline_configs.base import TextConditioningOutput from sglang.multimodal_gen.runtime.distributed import ( + get_encoder_data_parallel_group, get_local_torch_device, - get_world_group, ) from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context from sglang.multimodal_gen.runtime.managers.memory_managers.component_manager import ( @@ -46,12 +46,13 @@ logger = init_logger(__name__) def _data_parallel_text_encode(forward_fn, forward_kwargs: dict, group): - """each rank encodes its 1/world_size batch slice, then all-gathers + """Each encoder copy encodes a batch slice, then all-gathers the outputs. - every rank runs the full unsharded encoder on its slice, so each row is - computed by the same kernels as the replicated forward; the batch is padded - to a multiple of world_size and padding rows are dropped after the gather. - Requires a TextEncoder (BaseEncoderOutput) -- see _text_encode_dp_group. + An encoder copy may itself be TP-sharded. Corresponding TP ranks use the + same batch-DP rank, so every rank in that TP group receives the same slice; + the orthogonal batch-DP groups then gather the replicated encoder outputs. + The batch is padded to a multiple of the DP degree and padding rows are + dropped after the gather. Requires a TextEncoder (BaseEncoderOutput). """ world = group.world_size rank = group.rank_in_group @@ -515,29 +516,32 @@ class TextEncodingStage(ConditionEncodingStage): ): """group to data-parallel a batched text-encode over, or None - requires a replicated encoder (tp==1, dp==1, not folded): each rank - would otherwise redundantly encode the whole batch. Also requires a - TextEncoder, whose forward returns BaseEncoderOutput -- the gather needs - to know which fields carry the batch, and a raw transformers encoder - returns its own output type (e.g. Qwen2_5_VLCausalLMOutputWithPast). + DP splits the request batch across encoder copies and all-gathers the + outputs. A non-folded native encoder uses the DiT TP group, so this uses + the orthogonal non-TP ranks inside the same pipeline replica. It never + mixes requests across pipeline replicas and composes with DiT TP. """ - policy = server_args.encoder_parallel - if ( - policy not in ("auto", "dp") - # isinstance first: the loader can return a raw transformers - # encoder, which carries no such attribute - or not isinstance(text_encoder, TextEncoder) - or not text_encoder.supports_dp_encode - or (server_args.tp_size or 1) != 1 - or (server_args.dp_size or 1) != 1 - or encoder_config.parallel_folding_mode is not None - ): + if server_args.encoder_parallel not in ("auto", "dp"): return None - group = get_world_group() - if group.world_size <= 1: + # A folded encoder has one TP copy spanning its folding group, so there + # are no independent copies over which to split the batch. + if encoder_config.parallel_folding_mode is not None: + return None + # the gather rebuilds a BaseEncoderOutput, which only a TextEncoder + # forward produces -- a raw transformers encoder returns its own output + # type (e.g. Qwen2_5_VLCausalLMOutputWithPast). isinstance first: the + # loader can return such an encoder, which carries no dp attribute. + if not isinstance(text_encoder, TextEncoder): + return None + if not text_encoder.supports_dp_encode: + return None + group = get_encoder_data_parallel_group() + if group is None or group.world_size <= 1: return None # explicit dp trusts the operator on an unmeasured topology; auto does not - measured = policy == "dp" or group_has_measured_topology(group) + measured = server_args.encoder_parallel == "dp" or group_has_measured_topology( + group + ) if not encoder_dp_worthwhile(encoder_config, batch_size, measured): return None self._log_dp_choice(batch_size, group.world_size) @@ -548,7 +552,7 @@ class TextEncodingStage(ConditionEncodingStage): return self._dp_choice_logged = True logger.info( - "encoder_parallel: data-parallel text encode over %d ranks " + "encoder_parallel: data-parallel text encode over %d encoder copies " "(batch %d). Measured 1.9x on the encode stage at batch 2/4/8 " "(2xH100, T5-XXL width) with max_abs_diff=0 against the replicated " "forward.", diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index f983f7455..7b709cb7e 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -710,8 +710,20 @@ class ServerArgs(DisaggServerArgsMixin): replica_size = (self.num_gpus or tp_size) // dp_size fold_world = dp_size == 1 and not self.disagg_mode and replica_size > tp_size + # An explicit fold policy is honored for any multi-rank replica — + # pure-TP (replica == tp) and dp>1 shapes included. Encoder layout is + # independent of the DiT's parallelism; the replica group is simply + # "the ranks that share this batch". `auto` keeps its conservative + # proposals below and never folds a pure-TP replica. + fold_replica = ( + self.encoder_parallel == "fold" + and not self.disagg_mode + and replica_size > 1 + ) - if fold_world: + if fold_replica: + mode = "replica" + elif fold_world: mode = "world" elif tp_size == 1 and sp_degree > 1: # Preserve prior behavior for dp>1 / disaggregated SP runs. @@ -1891,8 +1903,9 @@ class ServerArgs(DisaggServerArgsMixin): "`auto` folds encoders wide enough to benefit (best " "single-request latency) and data-parallels eligible native " "text encoders at batch>1; `fold` always tensor-parallels the " - "encoder weights; `dp` never folds and splits the batch across " - "ranks (best batched throughput; requires TP=1 and DP=1); " + "encoder weights across the replica; `dp` never folds and " + "splits the batch across encoder copies inside each replica, " + "composing with encoder TP (best batched throughput); " "`replicate` disables both. The default is `auto`." ), ) @@ -3171,10 +3184,6 @@ class ServerArgs(DisaggServerArgsMixin): raise ValueError("batching_max_size must be >= 1") if self.batching_delay_ms < 0: raise ValueError("batching_delay_ms must be >= 0") - if self.encoder_parallel == "dp" and ( - (self.tp_size or 1) != 1 or (self.dp_size or 1) != 1 - ): - raise ValueError("encoder_parallel=dp requires tp_size=1 and dp_size=1") def _set_default_attention_backend(self) -> None: """Configure ROCm defaults when users do not specify an attention backend.""" diff --git a/python/sglang/multimodal_gen/test/unit/test_encoder_world_folding.py b/python/sglang/multimodal_gen/test/unit/test_encoder_world_folding.py index 441ceaf67..10f3be5fa 100644 --- a/python/sglang/multimodal_gen/test/unit/test_encoder_world_folding.py +++ b/python/sglang/multimodal_gen/test/unit/test_encoder_world_folding.py @@ -68,10 +68,25 @@ def _proposed_mode(tp, sp, cfg, dp=1, disagg=False, num_gpus=None, policy="auto" def test_pure_tp_not_folded(): - # replica == tp: encoder already uses every replica GPU; nothing to fold. + # replica == tp under auto: each rank keeps a full encoder replica; auto + # never proposes folding here (the trade-off stays opt-in, see below). assert _proposed_mode(tp=2, sp=1, cfg=1) is None +def test_explicit_fold_proposes_replica_for_any_shape(): + # Encoder layout is independent of the DiT's parallelism: an operator's + # explicit fold is honored on every multi-rank replica instead of being + # silently ignored (pure TP) or narrowed to the SP group (dp > 1). + assert _proposed_mode(tp=2, sp=1, cfg=1, policy="fold") == "replica" + assert _proposed_mode(tp=1, sp=2, cfg=1, policy="fold") == "replica" + assert ( + _proposed_mode(tp=2, sp=2, cfg=1, dp=2, num_gpus=8, policy="fold") == "replica" + ) + # single-GPU replica: nothing to shard over + assert _proposed_mode(tp=1, sp=1, cfg=1, policy="fold") is None + assert _proposed_mode(tp=1, sp=1, cfg=1, dp=2, num_gpus=2, policy="fold") is None + + def test_single_gpu_not_folded(): assert _proposed_mode(tp=1, sp=1, cfg=1) is None @@ -121,10 +136,14 @@ def test_all_encoders_get_the_same_proposed_mode(): assert img.parallel_folding_mode == "world" -def test_adjust_proposes_regardless_of_policy(): - # adjust reads the parallelism only; finalize owns the policy decision. - for policy in ("auto", "fold", "dp", "replicate"): +def test_adjust_proposal_policy_dependence(): + # adjust reads the parallelism only for auto/dp/replicate; finalize owns + # those policy decisions. An explicit fold is the one exception: it widens + # the proposal to the whole replica, because finalize cannot conjure a + # group that was never proposed (pure-TP replicas had none at all). + for policy in ("auto", "dp", "replicate"): assert _proposed_mode(tp=1, sp=2, cfg=1, policy=policy) == "world", policy + assert _proposed_mode(tp=1, sp=2, cfg=1, policy="fold") == "replica" def test_no_policy_touches_batching_max_size(): diff --git a/python/sglang/multimodal_gen/test/unit/test_text_encode_dp_gate.py b/python/sglang/multimodal_gen/test/unit/test_text_encode_dp_gate.py new file mode 100644 index 000000000..be020137e --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_text_encode_dp_gate.py @@ -0,0 +1,119 @@ +"""Unit test for the batched text-encode data-parallel gate. + +_text_encode_dp_group decides whether a batched encode may split across +encoder copies. It composes with DiT TP through the orthogonal encoder-DP +group and stays off for folded encoders, single-copy replicas, and non-DP +policies. Pure logic, no GPU / distributed init. +""" + +from types import SimpleNamespace +from unittest.mock import Mock + +from sglang.multimodal_gen.configs.models.encoders import TextEncoderConfig +from sglang.multimodal_gen.runtime.distributed.parallel_state import ( + _get_encoder_data_parallel_group_ranks, +) +from sglang.multimodal_gen.runtime.models.encoders.base import TextEncoder +from sglang.multimodal_gen.runtime.pipelines_core.stages import text_encoding as _te_mod +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.distributed import RankGenerator + + +def _enc(hidden=4096, heads=64, inter=10240, folding_mode=None): + enc = TextEncoderConfig() + enc.hidden_size = hidden + enc.num_attention_heads = heads + enc.intermediate_size = inter + enc.parallel_folding_mode = folding_mode + return enc + + +def _gate( + monkeypatch, + policy="dp", + tp=1, + dp=1, + folding_mode=None, + batch_size=4, + world_size=4, + group_available=True, + measured=True, + supports_dp=True, +): + group = SimpleNamespace(world_size=world_size) if group_available else None + monkeypatch.setattr(_te_mod, "get_encoder_data_parallel_group", lambda: group) + monkeypatch.setattr(_te_mod, "group_has_measured_topology", lambda group: measured) + encoder = Mock(spec=TextEncoder) + encoder.supports_dp_encode = supports_dp + server_args = SimpleNamespace(encoder_parallel=policy, tp_size=tp, dp_size=dp) + self = SimpleNamespace(_log_dp_choice=lambda batch, world: None) + return _te_mod.TextEncodingStage._text_encode_dp_group( + self, server_args, _enc(folding_mode=folding_mode), batch_size, encoder + ) + + +def test_dp_engages_for_replicated_encoder(monkeypatch): + assert _gate(monkeypatch) is not None + + +def test_dit_tp_composes_with_encoder_dp_group(monkeypatch): + # TP ranks jointly run one encoder copy; the orthogonal group distributes + # batch slices across those TP-sharded copies. + assert _gate(monkeypatch, tp=4) is not None + assert _gate(monkeypatch, policy="auto", tp=4) is not None + + +def test_pure_tp_has_no_independent_encoder_copy(monkeypatch): + # A pure-TP replica already uses every GPU for one encoder copy. + assert _gate(monkeypatch, tp=4, group_available=False) is None + + +def test_folded_encoder_blocks_dp(monkeypatch): + # folding shards the weights over the folding group; a rank cannot + # encode alone, whatever the policy says + assert _gate(monkeypatch, folding_mode="world") is None + assert _gate(monkeypatch, tp=4, folding_mode="world") is None + + +def test_multi_replica_uses_per_replica_encoder_group(monkeypatch): + # dp>1 no longer blocks: the encoder-DP group stays within one pipeline + # replica. A replica with one encoder copy uses its normal TP forward. + assert _gate(monkeypatch, dp=2) is not None + assert _gate(monkeypatch, dp=2, world_size=1) is None + + +def test_policy_and_encoder_gates(monkeypatch): + assert _gate(monkeypatch, policy="replicate") is None + assert _gate(monkeypatch, policy="fold") is None + assert _gate(monkeypatch, supports_dp=False) is None + assert _gate(monkeypatch, world_size=1) is None + assert _gate(monkeypatch, batch_size=1) is None # unbatched: not worthwhile + # explicit dp trusts the operator on an unmeasured topology; auto does not + assert _gate(monkeypatch, policy="dp", measured=False) is not None + assert _gate(monkeypatch, policy="auto", measured=False) is None + + +def _validate_batching(encoder_parallel, tp, dp): + args = SimpleNamespace( + batching_mode="dynamic", + batching_max_size=1, + batching_delay_ms=0, + encoder_parallel=encoder_parallel, + tp_size=tp, + dp_size=dp, + ) + ServerArgs._validate_batching(args) + + +def test_server_args_accepts_dp_for_any_parallel_shape(): + # Encoder DP composes with TP and stays within each pipeline replica. + _validate_batching("dp", tp=4, dp=1) + _validate_batching("dp", tp=1, dp=2) + _validate_batching("dp", tp=4, dp=2) + + +def test_encoder_dp_groups_are_tp_orthogonal_and_replica_local(): + ranks = _get_encoder_data_parallel_group_ranks( + RankGenerator(tp=2, sp=2, pp=1, cfg=1, dp=2, order="tp-sp-pp-cfg-dp") + ) + assert ranks == [[0, 2], [1, 3], [4, 6], [5, 7]]