anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[diffusion][feature] Add component-selective offload policies (#5929)

Browse Source 
Signed-off-by: lishunyang12 <lishunyang12@163.com>
This commit is contained in:
SYLAR authored and GitHub committed 2026-09-06 05:40:09 +08:00
1 parent 44d3ae100a
commit ea3cf99894
74 files changed
+5121 -1584
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/test_multiproc_engine_concurrency.py
+93 -21
View File
@@ -254,8 +254,11 @@ class TestRequestModeDispatch:
executor.od_config = SimpleNamespace(
step_execution=False,
parallel_config=SimpleNamespace(data_parallel_size=2),
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
diffusion_offload_config={
"mode": "layer",
"components": ["dit"],
"layer_options": {"dit": {"weight_transfer": "allgather"}},
},
)
executor.execute_request = Mock(return_value="dlo-dp")
executor.collective_rpc = Mock()
@@ -358,7 +361,7 @@ class TestRequestModeDispatch:
for rank, result_mq in enumerate(executor._result_mqs):
assert result_mq.dequeue.call_count == (1 if rank in expected_primary_ranks else 0)
@pytest.mark.parametrize("empty_prompt", ["", {"prompt": ""}])
@pytest.mark.parametrize("empty_prompt", ["", [], {}, {"prompt": ""}])
def test_dlo_dp_rejects_empty_prompt_before_worker_dispatch(self, empty_prompt):
executor, _, _ = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
@@ -376,6 +379,71 @@ class TestRequestModeDispatch:
executor.collective_rpc.assert_not_called()
def test_dlo_dp_accepts_precomputed_embeddings_without_prompt_text(self):
executor, _, _ = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
step_execution=False,
parallel_config=SimpleNamespace(data_parallel_size=2),
diffusion_offload_config={
"mode": "layer",
"components": ["text_encoder"],
"layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
},
)
executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
scheduler_output = _make_sched_output("A", "B")
for new_req in scheduler_output.scheduled_new_reqs:
new_req.req.prompt = {"prompt": "", "prompt_embeds": object()}
result = executor.execute_request(scheduler_output)
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
@pytest.mark.parametrize(
"prompt",
[
[1, 2],
{"type": "token", "prompt_token_ids": [1, 2]},
{"prompt_ids": [1, 2], "prompt_mask": object()},
],
)
def test_dlo_dp_accepts_tokenized_prompts(self, prompt):
executor, _, _ = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
step_execution=False,
parallel_config=SimpleNamespace(data_parallel_size=2),
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
)
executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
scheduler_output = _make_sched_output("A", "B")
for new_req in scheduler_output.scheduled_new_reqs:
new_req.req.prompt = prompt
result = executor.execute_request(scheduler_output)
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
def test_text_encoder_allgather_rejects_mismatched_embedding_paths(self):
executor, _, _ = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
step_execution=False,
parallel_config=SimpleNamespace(data_parallel_size=2),
diffusion_offload_config={
"mode": "layer",
"components": ["text_encoder"],
"layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
},
)
executor.collective_rpc = Mock()
scheduler_output = _make_sched_output("A", "B")
scheduler_output.scheduled_new_reqs[1].req.prompt = {"prompt_embeds": object()}
with pytest.raises(ValueError, match="same positive/negative prompt embedding fields"):
executor.execute_request(scheduler_output)
executor.collective_rpc.assert_not_called()
def test_dlo_dp_valid_wave_still_runs_after_rejected_wave(self):
executor, _, _ = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
@@ -396,24 +464,6 @@ class TestRequestModeDispatch:
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
executor.collective_rpc.assert_called_once()
def test_dlo_dp_allows_shared_default_denoise_steps(self):
executor, _, _ = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
step_execution=False,
parallel_config=SimpleNamespace(data_parallel_size=2),
enable_distributed_layerwise_offload=True,
dlo_use_allgather=True,
)
executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
scheduler_output = _make_sched_output("A", "B")
for new_req in scheduler_output.scheduled_new_reqs:
new_req.req.sampling_params.num_inference_steps = None
result = executor.execute_request(scheduler_output)
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
executor.collective_rpc.assert_called_once()
def test_dlo_dp_forwards_request_metadata_envelopes_as_one_wave(self):
executor, _, _ = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
@@ -496,6 +546,28 @@ class TestRequestModeDispatch:
executor._fail_closed_on_dp_wave_timeout.assert_called_once()
assert isinstance(executor._fail_closed_on_dp_wave_timeout.call_args.args[0], TimeoutError)
def test_single_allgather_request_times_out_and_fails_closed(self):
from vllm_omni.diffusion.executor import multiproc_executor as executor_module
executor, _, _ = _make_executor(num_gpus=2)
executor.od_config = SimpleNamespace(
step_execution=False,
parallel_config=SimpleNamespace(data_parallel_size=2),
diffusion_offload_config={
"mode": "layer",
"components": ["dit"],
"layer_options": {"dit": {"weight_transfer": "allgather"}},
},
)
executor.collective_rpc = Mock(side_effect=TimeoutError("timed out"))
executor._fail_closed_on_dp_wave_timeout = Mock()
result = executor.execute_request(_make_sched_output("A"))
assert result.runner_outputs[0].result.error == "timed out"
assert executor.collective_rpc.call_args.kwargs["timeout"] == executor_module._DLO_DP_WAVE_TIMEOUT_S
executor._fail_closed_on_dp_wave_timeout.assert_called_once()
# ───────────────── concurrent collective RPC ─────────────────