[diffusion] fix: serve requests that turn cfg off on a cfg-parallel server (#39303)
Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
This commit is contained in:
@@ -394,34 +394,30 @@ class InputValidationStage(PipelineStage):
|
||||
f"Guidance scale must be positive, but got {batch.guidance_scale}"
|
||||
)
|
||||
|
||||
# Reject requests that do not enable CFG on a server launched with
|
||||
# --enable-cfg-parallel. CFG-parallel splits cond/uncond across ranks,
|
||||
# so rank 1 has no work and returns None for noise_pred, which crashes
|
||||
# scheduler.step() ~30 minutes later under a gloo broadcast timeout.
|
||||
# Earlier, field-specific checks above (negative_prompt missing,
|
||||
# guidance_scale < 0) fire first and produce better messages for those
|
||||
# cases; this is the catch-all for any combination that still leaves
|
||||
# do_classifier_free_guidance=False under cfg-parallel.
|
||||
# A request that leaves CFG off is servable under CFG parallelism: the
|
||||
# dispatcher gives branch 0 to rank 0, and every other rank runs branch 0
|
||||
# too so the all-gather has shapes to work with. Both ranks then read the
|
||||
# owner's prediction, so the answer is the single-branch answer and the
|
||||
# extra ranks are only redundant.
|
||||
#
|
||||
# This used to raise. That guard was added for a warmup hang (#23198)
|
||||
# two weeks BEFORE the multi-branch refactor (#23736) taught the
|
||||
# dispatcher to handle a single branch, and the warmup path has since
|
||||
# grown its own fix -- the warmup builder forces CFG on whenever
|
||||
# cfg-parallel is enabled. What was left was a server refusing traffic
|
||||
# it could serve, and the runtime AUTO-enables cfg-parallel from the
|
||||
# model's default sampling params, so `sglang serve --num-gpus 2` on a
|
||||
# CFG-defaulting model rejected every guidance_scale=1.0 request while
|
||||
# blaming a flag the user never passed.
|
||||
if server_args.enable_cfg_parallel and not batch.do_classifier_free_guidance:
|
||||
neg_prompt_state = (
|
||||
"not set"
|
||||
if batch.negative_prompt is None
|
||||
else "empty"
|
||||
if batch.negative_prompt == ""
|
||||
else "set"
|
||||
)
|
||||
raise ValueError(
|
||||
f"Server was launched with --enable-cfg-parallel but this "
|
||||
f"request does not use classifier-free guidance "
|
||||
f"(do_classifier_free_guidance={batch.do_classifier_free_guidance}, "
|
||||
f"guidance_scale={batch.guidance_scale}, "
|
||||
f"true_cfg_scale={batch.true_cfg_scale}, "
|
||||
f"negative_prompt={neg_prompt_state}). "
|
||||
f"CFG-parallel splits cond/uncond across ranks and requires "
|
||||
f"both to be active. Either disable --enable-cfg-parallel or "
|
||||
f"ensure the request enables CFG (set guidance_scale > 1.0 or "
|
||||
f"true_cfg_scale > 1.0, with a non-empty negative_prompt or "
|
||||
f"negative_prompt_embeds)."
|
||||
logger.warning_once(
|
||||
"CFG parallelism is enabled but this request does not use "
|
||||
"classifier-free guidance (guidance_scale=%s, true_cfg_scale=%s), "
|
||||
"so it has one branch and the other CFG rank(s) recompute it "
|
||||
"redundantly. Pass --cfg-parallel-size 1 to spend those GPUs on "
|
||||
"another parallelism instead.",
|
||||
batch.guidance_scale,
|
||||
batch.true_cfg_scale,
|
||||
)
|
||||
|
||||
# for i2v, get image from image_path
|
||||
|
||||
@@ -1466,9 +1466,11 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
self.enable_cfg_parallel = auto_cfg_parallel_degree > 1
|
||||
if self.enable_cfg_parallel:
|
||||
logger.info(
|
||||
"Automatically enabled CFG parallel at degree %d for %d GPUs. "
|
||||
"Use --sp-degree / --ulysses-degree to use sequence "
|
||||
"parallelism instead.",
|
||||
"Automatically enabled CFG parallel at degree %d for %d GPUs "
|
||||
"because this model uses classifier-free guidance by default. "
|
||||
"A request that turns CFG off still runs, but it has one branch, "
|
||||
"so the other CFG rank(s) recompute it redundantly. Override with "
|
||||
"--cfg-parallel-size 1, --tp-size, or --sp-degree / --ulysses-degree.",
|
||||
self.cfg_parallel_degree,
|
||||
self.num_gpus,
|
||||
)
|
||||
|
||||
@@ -100,16 +100,53 @@ class SortedHelpFormatter(argparse.HelpFormatter):
|
||||
super().add_arguments(actions)
|
||||
|
||||
|
||||
@lru_cache
|
||||
def _print_info_once(logger: Logger, msg: str) -> None:
|
||||
# Set the stacklevel to 2 to print the original caller's line info
|
||||
logger.info(msg, stacklevel=2)
|
||||
# `logger.warning_once(msg, *args)` is bound as MethodType(_print_warning_once,
|
||||
# logger), so there is exactly ONE frame between the caller and logger.warning --
|
||||
# and stacklevel=2 is part of the observable contract, asserted literally by
|
||||
# test_diffusion_bcg_padding. Any helper in between pushes the record's filename
|
||||
# to this file and breaks that assertion, so the dedup cannot be an lru_cache on
|
||||
# a second function.
|
||||
#
|
||||
# It also cannot be an lru_cache on THIS function: keyed on the arguments it would
|
||||
# hold a strong reference to each one for the life of the process, and callers
|
||||
# here pass tensors. Hence a set of formatted text, which stores only strings.
|
||||
#
|
||||
# The args themselves are new: these helpers used to take the message alone, so a
|
||||
# caller that formatted lazily -- the way the standard contract implies -- raised
|
||||
# TypeError instead of logging, always on a branch too rare to have been seen.
|
||||
_logged_once: set[tuple[str, int, str]] = set()
|
||||
|
||||
|
||||
@lru_cache
|
||||
def _print_warning_once(logger: Logger, msg: str) -> None:
|
||||
# Set the stacklevel to 2 to print the original caller's line info
|
||||
logger.warning(msg, stacklevel=2)
|
||||
def _log_once_guard(logger: Logger, level: int, msg: str, *args: Any) -> str | None:
|
||||
"""The text to log, or None when this message has already been logged."""
|
||||
text = msg % args if args else msg
|
||||
key = (logger.name, level, text)
|
||||
if key in _logged_once:
|
||||
return None
|
||||
_logged_once.add(key)
|
||||
return text
|
||||
|
||||
|
||||
def _print_info_once(logger: Logger, msg: str, *args: Any) -> None:
|
||||
text = _log_once_guard(logger, logging.INFO, msg, *args)
|
||||
# stacklevel=2 is asserted literally by test_diffusion_bcg_padding, so it is
|
||||
# contract rather than a tuning knob. It does NOT reach the caller: init_logger
|
||||
# also patches `warning` into a forwarder to `logger.log`, adding a frame, so
|
||||
# the record names this module. That was true before these helpers too.
|
||||
if text is not None:
|
||||
logger.info(text, stacklevel=2)
|
||||
|
||||
|
||||
def _print_warning_once(logger: Logger, msg: str, *args: Any) -> None:
|
||||
text = _log_once_guard(logger, logging.WARNING, msg, *args)
|
||||
if text is not None:
|
||||
logger.warning(text, stacklevel=2)
|
||||
|
||||
|
||||
# These were lru_cache objects, so `.cache_clear()` was part of their surface and
|
||||
# a test resets the dedup through it.
|
||||
_print_info_once.cache_clear = _logged_once.clear
|
||||
_print_warning_once.cache_clear = _logged_once.clear
|
||||
|
||||
|
||||
def get_is_main_process():
|
||||
@@ -167,19 +204,19 @@ class _SGLDiffusionLogger(Logger):
|
||||
`intel_extension_for_pytorch.utils._logger`.
|
||||
"""
|
||||
|
||||
def info_once(self, msg: str) -> None:
|
||||
def info_once(self, msg: str, *args: Any) -> None:
|
||||
"""
|
||||
As :meth:`info`, but subsequent calls with the same message
|
||||
are silently dropped.
|
||||
and args are silently dropped.
|
||||
"""
|
||||
_print_info_once(self, msg)
|
||||
_print_info_once(self, msg, *args)
|
||||
|
||||
def warning_once(self, msg: str) -> None:
|
||||
def warning_once(self, msg: str, *args: Any) -> None:
|
||||
"""
|
||||
As :meth:`warning`, but subsequent calls with the same message
|
||||
are silently dropped.
|
||||
and args are silently dropped.
|
||||
"""
|
||||
_print_warning_once(self, msg)
|
||||
_print_warning_once(self, msg, *args)
|
||||
|
||||
def info( # type: ignore[override]
|
||||
self,
|
||||
|
||||
@@ -4,8 +4,8 @@ Covers warmup and cfg-parallel guard paths introduced alongside this file:
|
||||
- build_warmup_reqs synthesizes warmup Reqs that actually enable
|
||||
classifier-free guidance when cfg-parallel is on.
|
||||
- DiffGenerator sends explicit warmup resolutions through the scheduler client.
|
||||
- InputValidationStage.forward rejects non-CFG requests when the server
|
||||
has cfg-parallel on.
|
||||
- InputValidationStage.forward ACCEPTS non-CFG requests when the server
|
||||
has cfg-parallel on, and the branch dispatcher serves them.
|
||||
- Server-based warmup can opt into model-default negative prompts so warmup
|
||||
populates the negative text embedding cache.
|
||||
- Req-based warmup remains available only through the lazy legacy path.
|
||||
@@ -1047,26 +1047,28 @@ class TestImageVaeEncodingLatentRetrieval(unittest.TestCase):
|
||||
)
|
||||
|
||||
|
||||
class TestInputValidationCfgParallelGuard(unittest.TestCase):
|
||||
"""Commit 2: per-request cfg-parallel check.
|
||||
class TestInputValidationCfgParallelSingleBranch(unittest.TestCase):
|
||||
"""A request that turns CFG off must still be served under cfg-parallel.
|
||||
|
||||
This used to raise. The guard came from a warmup hang (#23198, 2026-04-23);
|
||||
two weeks later the multi-branch refactor (#23736) taught the dispatcher to
|
||||
handle a single branch, and the warmup builder grew its own fix (it forces
|
||||
CFG on whenever cfg-parallel is enabled). What the guard still did was refuse
|
||||
live traffic the runtime could serve -- and because cfg-parallel is
|
||||
AUTO-enabled from the model's default sampling params, a plain
|
||||
`sglang serve --num-gpus 2` on a CFG-defaulting model rejected every
|
||||
guidance_scale=1.0 request, citing a flag the user never passed.
|
||||
|
||||
Both tests patch _generate_seeds (the first statement of
|
||||
InputValidationStage.forward, input_validation.py:274) to sidestep
|
||||
its device-lookup / generator-creation code which pulls in torch
|
||||
CUDA bindings — keeps the suite strictly CPU-only. We still need
|
||||
num_inference_steps on the Req because the stage's
|
||||
"num_inference_steps <= 0" check at L305-308 raises TypeError on
|
||||
None before the new commit-2 check is reached.
|
||||
InputValidationStage.forward) to sidestep its device-lookup / generator
|
||||
creation, keeping the suite CPU-only. num_inference_steps must be set because
|
||||
the "num_inference_steps <= 0" check raises TypeError on None first.
|
||||
"""
|
||||
|
||||
def test_input_validation_rejects_cfg_parallel_without_cfg(self):
|
||||
# negative_prompt="" (non-None) ensures the existing
|
||||
# negative_prompt-is-None check at input_validation.py:295-298
|
||||
# does NOT fire first — this isolates the new commit-2 check.
|
||||
# width/height/num_outputs_per_prompt pre-set so the stage's
|
||||
# default-dimension block at L352-361 doesn't mutate the Req
|
||||
# in a way that obscures the assertion target.
|
||||
req = Req(
|
||||
def _single_branch_req(self) -> Req:
|
||||
# negative_prompt="" (non-None) keeps the negative_prompt-is-None check
|
||||
# from firing first, so this isolates the cfg-parallel path.
|
||||
return Req(
|
||||
prompt="test",
|
||||
negative_prompt="",
|
||||
guidance_scale=1.0,
|
||||
@@ -1076,29 +1078,26 @@ class TestInputValidationCfgParallelGuard(unittest.TestCase):
|
||||
width=512,
|
||||
height=512,
|
||||
)
|
||||
|
||||
def test_input_validation_accepts_cfg_parallel_without_cfg(self):
|
||||
req = self._single_branch_req()
|
||||
self.assertIs(
|
||||
req.do_classifier_free_guidance,
|
||||
False,
|
||||
"Sanity: test setup must leave do_cfg=False so the "
|
||||
"commit-2 check is the one that fires, not an upstream check.",
|
||||
"Sanity: the setup must leave do_cfg=False, or this tests nothing.",
|
||||
)
|
||||
|
||||
stage = _make_input_validation_stage()
|
||||
server_args = _make_validation_server_args(enable_cfg_parallel=True)
|
||||
|
||||
with patch.object(InputValidationStage, "_generate_seeds"):
|
||||
with self.assertRaises(ValueError) as ctx:
|
||||
try:
|
||||
stage.forward(req, server_args)
|
||||
|
||||
msg = str(ctx.exception).lower()
|
||||
self.assertIn("cfg-parallel", msg)
|
||||
for field in (
|
||||
"do_classifier_free_guidance",
|
||||
"guidance_scale",
|
||||
"true_cfg_scale",
|
||||
"negative_prompt",
|
||||
):
|
||||
self.assertIn(field, str(ctx.exception))
|
||||
except ValueError as e:
|
||||
self.fail(
|
||||
"forward() rejected a single-branch request under "
|
||||
f"cfg-parallel; the dispatcher can serve it: {e}"
|
||||
)
|
||||
|
||||
def test_input_validation_passes_cfg_parallel_with_cfg(self):
|
||||
req = Req(
|
||||
@@ -1127,5 +1126,78 @@ class TestInputValidationCfgParallelGuard(unittest.TestCase):
|
||||
self.fail(f"forward() raised ValueError on a valid CFG request: {e}")
|
||||
|
||||
|
||||
class TestCfgParallelServesOneBranch(unittest.TestCase):
|
||||
"""The property that makes accepting a single-branch request safe.
|
||||
|
||||
Dropping the validation guard is only correct because the dispatcher already
|
||||
handles n_branches=1 on a 2-rank CFG group: rank 0 owns the branch, every
|
||||
other rank runs it too so the all-gather has shapes, and the reorder step
|
||||
hands both ranks the owner's prediction. Pin it from the rank that owns
|
||||
nothing -- that is the rank the old comment said returned None and hung a
|
||||
gloo broadcast for half an hour.
|
||||
"""
|
||||
|
||||
def _run_on_rank(self, cfg_rank: int, n_branches: int = 1, world_size: int = 2):
|
||||
from sglang.multimodal_gen.runtime.distributed.cfg_policy import (
|
||||
CFGBranch,
|
||||
CFGPolicy,
|
||||
)
|
||||
|
||||
mod = "sglang.multimodal_gen.runtime.distributed.cfg_parallel_utils"
|
||||
branches = [CFGBranch(f"b{i}", i == 0, {"tag": i}) for i in range(n_branches)]
|
||||
policy = CFGPolicy(branches=branches)
|
||||
seen: list[int] = []
|
||||
|
||||
def predict_fn(branch):
|
||||
seen.append(branch.kwargs["tag"])
|
||||
return torch.full((1, 2), float(branch.kwargs["tag"]))
|
||||
|
||||
# A real 2-rank gather returns one tensor per rank. Both ranks ran the
|
||||
# same branch here, so both contributions carry the same values.
|
||||
def fake_all_gather(t, dim=0, separate_tensors=False):
|
||||
return [t.clone() for _ in range(world_size)]
|
||||
|
||||
with (
|
||||
patch(f"{mod}.get_classifier_free_guidance_rank", return_value=cfg_rank),
|
||||
patch(
|
||||
f"{mod}.get_classifier_free_guidance_world_size",
|
||||
return_value=world_size,
|
||||
),
|
||||
patch(f"{mod}.get_local_torch_device", return_value=torch.device("cpu")),
|
||||
patch(f"{mod}.cfg_model_parallel_all_gather", side_effect=fake_all_gather),
|
||||
):
|
||||
from sglang.multimodal_gen.runtime.distributed.cfg_parallel_utils import (
|
||||
run_cfg_parallel,
|
||||
)
|
||||
|
||||
return run_cfg_parallel(policy, predict_fn), seen
|
||||
|
||||
def test_branch_owner_gets_the_single_prediction(self):
|
||||
preds, seen = self._run_on_rank(cfg_rank=0)
|
||||
self.assertEqual(len(preds), 1)
|
||||
self.assertEqual(seen, [0], "the owning rank runs branch 0 once")
|
||||
self.assertTrue(torch.equal(preds[0], torch.zeros(1, 2)))
|
||||
|
||||
def test_rank_without_a_branch_still_returns_the_owners_prediction(self):
|
||||
preds, seen = self._run_on_rank(cfg_rank=1)
|
||||
self.assertEqual(
|
||||
seen,
|
||||
[0],
|
||||
"the rank that owns no branch must still run one, or the "
|
||||
"all-gather has no shapes to work with",
|
||||
)
|
||||
self.assertEqual(len(preds), 1)
|
||||
self.assertIsNotNone(preds[0])
|
||||
self.assertTrue(torch.equal(preds[0], torch.zeros(1, 2)))
|
||||
|
||||
def test_two_branches_still_split_across_the_ranks(self):
|
||||
from sglang.multimodal_gen.runtime.distributed.cfg_parallel_utils import (
|
||||
dispatch_branches,
|
||||
)
|
||||
|
||||
self.assertEqual(dispatch_branches(1, 2), [[0], []])
|
||||
self.assertEqual(dispatch_branches(2, 2), [[0], [1]])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -1,10 +1,13 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import gc
|
||||
import logging
|
||||
import unittest
|
||||
import weakref
|
||||
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import (
|
||||
globally_suppress_loggers,
|
||||
init_logger,
|
||||
)
|
||||
|
||||
|
||||
@@ -27,5 +30,75 @@ class TestSuppressNoisyDependencyLogs(unittest.TestCase):
|
||||
)
|
||||
|
||||
|
||||
class TestLogOnceTakesFormatArgs(unittest.TestCase):
|
||||
"""`warning_once(msg, *args)` must format, not raise.
|
||||
|
||||
The helpers took only the message, so every caller that formatted lazily --
|
||||
the way logger.warning wants -- raised TypeError instead of logging. Three
|
||||
call sites did, and each sat on a branch that rarely runs, so the bug was
|
||||
invisible: cfg_parallel_utils only reaches its call when a CFG-parallel
|
||||
group has more ranks than branches, which input validation used to reject
|
||||
outright. Removing that rejection turned the latent TypeError into a crash
|
||||
on the first single-branch request.
|
||||
"""
|
||||
|
||||
def test_warning_once_formats_lazy_args(self):
|
||||
logger = init_logger("sglang.test.logonce.warning")
|
||||
with self.assertLogs(logger, level=logging.WARNING) as captured:
|
||||
logger.warning_once("cfg_parallel_size=%d > n_branches=%d", 2, 1)
|
||||
self.assertIn("cfg_parallel_size=2 > n_branches=1", captured.output[0])
|
||||
|
||||
def test_info_once_formats_lazy_args(self):
|
||||
logger = init_logger("sglang.test.logonce.info")
|
||||
with self.assertLogs(logger, level=logging.INFO) as captured:
|
||||
logger.info_once("degree %d on %d GPUs", 2, 2)
|
||||
self.assertIn("degree 2 on 2 GPUs", captured.output[0])
|
||||
|
||||
def test_record_does_not_name_the_caller(self):
|
||||
"""Documents a wart, so nobody "fixes" it and breaks the bcg assertion.
|
||||
|
||||
init_logger also replaces `logger.warning` with a patched method that
|
||||
forwards to `logger.log`, so there is one more frame than the stacklevel
|
||||
accounts for and the record names this module rather than the caller.
|
||||
That predates these helpers -- the original passed stacklevel=2 through
|
||||
the same patched method -- and raising the number would contradict
|
||||
test_diffusion_bcg_padding, which asserts the literal `stacklevel=2`.
|
||||
"""
|
||||
logger = init_logger("sglang.test.logonce.stacklevel")
|
||||
with self.assertLogs(logger, level=logging.WARNING) as captured:
|
||||
logger.warning_once("from the caller %d", 1)
|
||||
self.assertEqual(captured.records[0].filename, "logging_utils.py")
|
||||
|
||||
def test_arguments_are_not_retained(self):
|
||||
"""The once-cache must key on text, not on the arguments.
|
||||
|
||||
An lru_cache keyed on the arguments holds a strong reference to each of
|
||||
them for the life of the process, and callers in this package pass
|
||||
tensors. Formatting first and caching the result keeps only strings.
|
||||
"""
|
||||
logger = init_logger("sglang.test.logonce.retain")
|
||||
|
||||
class _Heavy:
|
||||
def __repr__(self):
|
||||
return "<heavy>"
|
||||
|
||||
obj = _Heavy()
|
||||
ref = weakref.ref(obj)
|
||||
with self.assertLogs(logger, level=logging.WARNING) as captured:
|
||||
logger.warning_once("holding %s", obj)
|
||||
self.assertIn("<heavy>", captured.output[0])
|
||||
del obj
|
||||
gc.collect()
|
||||
self.assertIsNone(ref(), "the once-cache kept the argument alive")
|
||||
|
||||
def test_same_message_and_args_logs_once(self):
|
||||
logger = init_logger("sglang.test.logonce.dedup")
|
||||
with self.assertLogs(logger, level=logging.WARNING) as captured:
|
||||
logger.warning_once("idle ranks: %d", 1)
|
||||
logger.warning_once("idle ranks: %d", 1)
|
||||
logger.warning_once("idle ranks: %d", 2) # different args, new line
|
||||
self.assertEqual(len(captured.output), 2, captured.output)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user