From da4b2f9aef769151c019023bade228864c8b6eb7 Mon Sep 17 00:00:00 2001 From: Linyu Wu Date: Tue, 1 Sep 2026 10:05:34 +0800 Subject: [PATCH 01/23] feat(rollout): move the vllm-omni engine to 0.28.0rc1 on transformers 5.12.1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit vllm-omni 0.28.0rc1 requires transformers >=5.10.1,<5.15, above the repo's 5.6.x baseline. The base pin widens to admit both engine forks rather than moving to the new floor, because sglang hard-pins 5.6.0 and raising the base leaves that fork unsolvable; the vllm extra instead narrows itself to an exact 5.12.1. That version is deliberate — sglang 0.5.15+ pins it too, so unifying the two forks later is a bump rather than a second zoo-wide migration. The exact pin also matters because the repo carries no uv.lock, so a range here would follow the next patch release silently. torch is unchanged: vllm 0.28.0 declares the same 2.13.0 as 0.27.0, so the CUDA-13 stack carries over. Core vLLM 0.28 re-slots the reusable MoE workspace by lane within ubatch (ubatch_id * num_lanes + lane) rather than by ubatch alone, so the workspace pool patch now mirrors that arithmetic instead of indexing _current_workspaces directly. Reading num_lanes off the manager and the lane off the ContextVar keeps the old flat layout when either is absent. Nothing exercises lanes today — the diffusion rollout never enables DSpark speculative decoding — but both failure modes here are silent: a stale index writes None over the wrong slot and orphans a tensor that still has live views. --- pyproject.toml | 16 ++++--- .../vllm_omni/patches/compat_moe_workspace.py | 45 ++++++++++++------- 2 files changed, 39 insertions(+), 22 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b43622db7..fa2d8a066 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,8 +22,9 @@ dependencies = [ "diffusers>=0.38.0", # LTX-2 needs 0.38: transformer.forward gained sigma/audio_sigma/isolate_modalities/use_cross_timestep, connector gained padding_side, calculate_shift now pins image_seq_len to max_image_seq_len. (>=0.37 was already required for Qwen-Image RoPE text-len-from-mask; 0.38 verified to keep that + every other model API unchanged.) "hydra-core>=1.3", "omegaconf>=2.3", - # Match the engine stacks' 5.6.x baseline; BAGEL vendor smoke passed on 5.6.0. - "transformers>=5.6,<5.7", + # Wide enough to admit both engine forks — sglang hard-pins 5.6.0 while the + # vllm extra narrows itself to 5.12.1, so raising this floor unsolves sglang. + "transformers>=5.6,<5.13", "huggingface-hub>=0.34", # Vendored Boogu-Image DiT (unirl/models/boogu_image/vendor/) imports einops # in transformer_boogu.py / attention_processor.py / rope.py. Pure-python, @@ -82,8 +83,12 @@ sglang = [ # +cu129 pin would resolve anyway, linking a CUDA-12 torch under a CUDA-13 vllm. # Prerelease, so installs need --prerelease=allow, plus sglang's compat layer. vllm = [ - "vllm==0.27.0 ; sys_platform == 'linux'", - "vllm-omni==0.27.0rc1 ; sys_platform == 'linux'", + "vllm==0.28.0 ; sys_platform == 'linux'", + "vllm-omni==0.28.0rc1 ; sys_platform == 'linux'", + # vllm-omni 0.28 needs >=5.10.1,<5.15 (5.10.0 is yanked, 5.15 has a known + # construction regression); exact because there is no uv.lock to hold it, and + # 5.12.1 is what sglang 0.5.15+ pins, so unifying later is a bump not a port. + "transformers==5.12.1 ; sys_platform == 'linux'", "torch==2.13.0+cu130 ; sys_platform == 'linux'", "torchvision==0.28.0+cu130 ; sys_platform == 'linux'", # torchaudio 2.11.0+cu130 is ABI-built against torch 2.11; there is no @@ -152,7 +157,8 @@ index-strategy = "unsafe-best-match" # sglang 0.5.12.post1 pins transformers==5.6.0 but leaves its `kernels` # integration unpinned; kernels >=0.15 changed LayerRepository's signature and # breaks transformers 5.6.0 at import time. Hold it to transformers' own -# declared bound. +# declared bound. Scoped to that fork in practice: nothing the vllm extra +# resolves requires kernels, so this does not cap its transformers at <5.13. # # transformers 5.6.0 accepts tokenizers<=0.23.0, but tokenizers 0.23.0(rc) # dropped the `cls` kwarg from RobertaProcessing.__new__, so deserializing a diff --git a/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py b/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py index 3fa1d336d..b8dd75c21 100644 --- a/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py +++ b/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py @@ -10,16 +10,16 @@ _WORKSPACE_MANAGER_MARKER = "_diffrl_moe_workspace_pool_installed" -def _ubatch_pool_key(ubatch_id: int) -> str: - return f"{MOE_WORKSPACE_TAG}:{ubatch_id}" +def _workspace_pool_key(workspace_id: int) -> str: + return f"{MOE_WORKSPACE_TAG}:{workspace_id}" @contextmanager -def _workspace_pool_context(allocator: Any, ubatch_id: int) -> Iterator[None]: - """Create one replaceable pool per ubatch while keeping one sleep tag.""" - # The dictionary key is ubatch-specific so DBO slots do not replace one - # another, while AllocationData keeps one model-level workspace tag. - pool_key = _ubatch_pool_key(ubatch_id) +def _workspace_pool_context(allocator: Any, workspace_id: int) -> Iterator[None]: + """Create one replaceable pool per workspace while keeping one sleep tag.""" + # The dictionary key is per-workspace so DBO and lane slots do not replace + # one another, while AllocationData keeps one model-level workspace tag. + pool_key = _workspace_pool_key(workspace_id) previous_tag = allocator.current_tag with allocator.use_memory_pool(tag=pool_key): allocator.current_tag = MOE_WORKSPACE_TAG @@ -30,9 +30,9 @@ def _workspace_pool_context(allocator: Any, ubatch_id: int) -> Iterator[None]: allocator.current_tag = previous_tag -def _release_ubatch_pool(allocator: Any, ubatch_id: int) -> None: +def _release_workspace_pool(allocator: Any, workspace_id: int) -> None: """Release a retained MemPool before its pluggable allocator.""" - pool_state = allocator.allocator_and_pools.pop(_ubatch_pool_key(ubatch_id), None) + pool_state = allocator.allocator_and_pools.pop(_workspace_pool_key(workspace_id), None) if pool_state is None: return memory_pool, pluggable_allocator = pool_state @@ -46,7 +46,7 @@ def _patch_workspace_manager_class( workspace_manager_class: type, *, allocator_provider: Callable[[], Any], - ubatch_provider: Callable[[], int], + workspace_id_provider: Callable[[Any], int], ) -> None: """Route ``_ensure_workspace_size`` growth through the dedicated CuMem tag.""" if getattr(workspace_manager_class, _WORKSPACE_MANAGER_MARKER, False): @@ -55,9 +55,9 @@ def _patch_workspace_manager_class( original_ensure_workspace_size = workspace_manager_class._ensure_workspace_size def ensure_workspace_size(manager: Any, required_bytes: int) -> Any: - ubatch_id = ubatch_provider() + workspace_id = workspace_id_provider(manager) try: - current_workspace = manager._current_workspaces[ubatch_id] + current_workspace = manager._current_workspaces[workspace_id] current_size = manager._workspace_size_bytes(current_workspace) workspace_locked = manager.is_locked() except AttributeError: @@ -77,21 +77,21 @@ def ensure_workspace_size(manager: Any, required_bytes: int) -> Any: import torch # ``use_memory_pool`` replaces the retained MemPool for a key. Release - # this ubatch's old tensor before replacing its pool; reversing that + # this slot's old tensor before replacing its pool; reversing that # order aborts in MemPool::~MemPool on torch 2.11. replacing_workspace = current_workspace is not None if replacing_workspace: # vLLM 0.20's CuMem free callback can unmap immediately, so a # device-wide sync must drain every stream first. torch.accelerator.synchronize() - manager._current_workspaces[ubatch_id] = None + manager._current_workspaces[workspace_id] = None del current_workspace gc.collect() - _release_ubatch_pool(allocator, ubatch_id) + _release_workspace_pool(allocator, workspace_id) if replacing_workspace: torch.accelerator.empty_cache() - with _workspace_pool_context(allocator, ubatch_id): + with _workspace_pool_context(allocator, workspace_id): return original_ensure_workspace_size(manager, required_bytes) workspace_manager_class._ensure_workspace_size = ensure_workspace_size @@ -109,10 +109,21 @@ def patch_moe_workspace_pool() -> None: except (ImportError, AttributeError): return + try: + from vllm.v1.worker.workspace import _workspace_lane + except ImportError: + _workspace_lane = None + + def workspace_id_provider(manager: Any) -> int: + # vLLM 0.28 slots workspaces by lane within ubatch; these defaults + # reproduce the flat per-ubatch layout that earlier versions used. + lane = 0 if _workspace_lane is None else _workspace_lane.get() + return dbo_current_ubatch_id() * getattr(manager, "_num_lanes", 1) + lane + _patch_workspace_manager_class( WorkspaceManager, allocator_provider=lambda: CuMemAllocator.instance, - ubatch_provider=dbo_current_ubatch_id, + workspace_id_provider=workspace_id_provider, ) From db9c9871aa15c1c87a3d46332677212c6db269b0 Mon Sep 17 00:00:00 2001 From: Linyu Wu Date: Wed, 2 Sep 2026 19:17:45 +0800 Subject: [PATCH 02/23] fix(rollout): let vLLM report an unconfigured MoE workspace lane itself --- .../rollout/engine/vllm_omni/patches/compat_moe_workspace.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py b/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py index b8dd75c21..b536d44e0 100644 --- a/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py +++ b/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py @@ -60,7 +60,9 @@ def ensure_workspace_size(manager: Any, required_bytes: int) -> Any: current_workspace = manager._current_workspaces[workspace_id] current_size = manager._workspace_size_bytes(current_workspace) workspace_locked = manager.is_locked() - except AttributeError: + except (AttributeError, IndexError): + # An out-of-range slot means an unconfigured lane, which vLLM itself + # reports with a precise error once it reaches its own bounds check. return original_ensure_workspace_size(manager, required_bytes) if current_size >= required_bytes: return original_ensure_workspace_size(manager, required_bytes) From e076e858fb6a2aeed8226164ecb8fc0de3a5d32b Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Mon, 14 Sep 2026 23:59:12 +0800 Subject: [PATCH 03/23] feat(rollout): adapt vllm-omni to stable 0.28 --- INSTALL.md | 6 +- datasets/ucf101/README.md | 2 +- ...i_audio_dcase_gspo_lora_vllm_omni_1x4.yaml | 2 +- ...i_audio_video_gspo_lora_vllm_omni_1x4.yaml | 2 +- ...i_audio_video_gspo_lora_vllm_omni_1x8.yaml | 2 +- ...mage_video_r1_gspo_lora_vllm_omni_1x4.yaml | 2 +- ...omni_video_r1_gspo_lora_vllm_omni_1x4.yaml | 2 +- ...omni_video_r1_gspo_lora_vllm_omni_1x8.yaml | 2 +- .../bagel/bagel_it2i_managed_editscore.yaml | 2 +- .../diffusion/bagel/bagel_it2i_vllmomni.yaml | 2 +- examples/diffusion/bagel/bagel_vllmomni.yaml | 2 +- .../diffusion/bagel/bagel_vllmomni_async.yaml | 2 +- pyproject.toml | 31 ++-- unirl/rollout/README.md | 4 +- .../engine/vllm_omni/adapters/bagel.py | 2 +- .../rollout/engine/vllm_omni/adapters/base.py | 6 +- .../rollout/engine/vllm_omni/adapters/hi3.py | 18 +- .../rollout/engine/vllm_omni/adapters/hv15.py | 6 +- .../engine/vllm_omni/adapters/qwen3_omni.py | 5 +- .../engine/vllm_omni/adapters/qwen_image.py | 4 +- .../rollout/engine/vllm_omni/adapters/sd3.py | 4 +- .../rollout/engine/vllm_omni/backends/base.py | 2 +- .../engine/vllm_omni/backends/native.py | 154 +++++++++++------- unirl/rollout/engine/vllm_omni/config.py | 6 +- .../deploy_configs/bagel_t2i_rl.yaml | 35 ++++ .../hunyuan_image3_ar_recaption_rl.yaml | 31 ++++ .../deploy_configs/hunyuan_image3_ar_rl.yaml | 26 +++ .../hunyuan_image3_dit_recaption_rl.yaml | 25 +++ .../deploy_configs/hunyuan_image3_i2t_rl.yaml | 26 +++ .../hunyuan_image3_it2i_rl.yaml | 48 ++++++ .../deploy_configs/hunyuan_image3_t2i_rl.yaml | 69 ++++++++ .../hunyuan_video15_t2v_rl.yaml | 34 ++++ .../qwen3_omni_thinker_only_rl_1x4.yaml | 30 ++++ ..._omni_thinker_only_rl_audio_video_1x4.yaml | 30 ++++ .../deploy_configs/qwen_image_t2i_rl.yaml | 33 ++++ .../vllm_omni/deploy_configs/sd35_t2i_rl.yaml | 33 ++++ .../engine/vllm_omni/patches/README.md | 5 +- .../vllm_omni/patches/compat_hi3_lora.py | 55 ------- .../vllm_omni/patches/compat_moe_workspace.py | 19 ++- .../vllm_omni/patches/compat_tokenizer.py | 2 - .../engine/vllm_omni/patches/runtime.py | 77 +++------ .../engine/vllm_omni/pipeline_configs.py | 80 +++++++++ unirl/rollout/engine/vllm_omni/plugin.py | 11 ++ .../vllm_omni/stage_configs/bagel_t2i_rl.yaml | 103 ------------ .../hunyuan_image3_ar_recaption_rl.yaml | 77 --------- .../hunyuan_image3_dit_recaption_rl.yaml | 68 -------- .../stage_configs/hunyuan_image3_i2t.yaml | 40 ----- .../stage_configs/hunyuan_image3_it2i_rl.yaml | 87 ---------- .../stage_configs/hunyuan_image3_t2i_rl.yaml | 126 -------------- .../stage_configs/hunyuan_image3_t2t.yaml | 41 ----- .../stage_configs/hunyuan_video15_t2v_rl.yaml | 79 --------- .../qwen3_omni_thinker_only_rl_1x4.yaml | 51 ------ ..._omni_thinker_only_rl_audio_video_1x4.yaml | 52 ------ .../stage_configs/qwen_image_t2i_rl.yaml | 85 ---------- .../vllm_omni/stage_configs/sd35_t2i_rl.yaml | 73 --------- .../rollout/engine/vllm_omni/utils/tracks.py | 13 +- 56 files changed, 700 insertions(+), 1134 deletions(-) create mode 100644 unirl/rollout/engine/vllm_omni/deploy_configs/bagel_t2i_rl.yaml create mode 100644 unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_recaption_rl.yaml create mode 100644 unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml create mode 100644 unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_dit_recaption_rl.yaml create mode 100644 unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml create mode 100644 unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_it2i_rl.yaml create mode 100644 unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_t2i_rl.yaml create mode 100644 unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_video15_t2v_rl.yaml create mode 100644 unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_1x4.yaml create mode 100644 unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_audio_video_1x4.yaml create mode 100644 unirl/rollout/engine/vllm_omni/deploy_configs/qwen_image_t2i_rl.yaml create mode 100644 unirl/rollout/engine/vllm_omni/deploy_configs/sd35_t2i_rl.yaml delete mode 100644 unirl/rollout/engine/vllm_omni/patches/compat_hi3_lora.py create mode 100644 unirl/rollout/engine/vllm_omni/pipeline_configs.py delete mode 100644 unirl/rollout/engine/vllm_omni/stage_configs/bagel_t2i_rl.yaml delete mode 100644 unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_ar_recaption_rl.yaml delete mode 100644 unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_dit_recaption_rl.yaml delete mode 100644 unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_i2t.yaml delete mode 100644 unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_it2i_rl.yaml delete mode 100644 unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_t2i_rl.yaml delete mode 100644 unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_t2t.yaml delete mode 100644 unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_video15_t2v_rl.yaml delete mode 100644 unirl/rollout/engine/vllm_omni/stage_configs/qwen3_omni_thinker_only_rl_1x4.yaml delete mode 100644 unirl/rollout/engine/vllm_omni/stage_configs/qwen3_omni_thinker_only_rl_audio_video_1x4.yaml delete mode 100644 unirl/rollout/engine/vllm_omni/stage_configs/qwen_image_t2i_rl.yaml delete mode 100644 unirl/rollout/engine/vllm_omni/stage_configs/sd35_t2i_rl.yaml diff --git a/INSTALL.md b/INSTALL.md index d76d4f3af..fe9c46ddb 100644 --- a/INSTALL.md +++ b/INSTALL.md @@ -16,14 +16,14 @@ explicit `CUDA_COMPAT_DIR`. ```bash uv venv --python 3.12 --seed .venv && source .venv/bin/activate -uv pip install -e ".[vllm,train,infer]" --prerelease=allow +uv pip install -e ".[vllm,train,infer]" ``` ## sglang ```bash uv venv --python 3.12 --seed .venv-sglang && source .venv-sglang/bin/activate -uv pip install -e ".[sglang,train,infer]" --prerelease=allow +uv pip install -e ".[sglang,train,infer]" ``` This extra reaches `causal-conv1d` through `flash-linear-attention[conv1d]`, @@ -69,7 +69,7 @@ For development tools (lint and tests): ```bash uv pip install -e ".[vllm,train,infer,eval,dev]" # or, for the sglang engine: -uv pip install -e ".[sglang,train,infer,eval,dev]" --prerelease=allow +uv pip install -e ".[sglang,train,infer,eval,dev]" ``` ## Environment diff --git a/datasets/ucf101/README.md b/datasets/ucf101/README.md index 324eb9105..07be49266 100644 --- a/datasets/ucf101/README.md +++ b/datasets/ucf101/README.md @@ -18,7 +18,7 @@ training environment can be installed with: ```bash uv venv --python 3.12 --seed .venv source .venv/bin/activate -uv pip install -e ".[vllm,train,infer]" --prerelease=allow +uv pip install -e ".[vllm,train,infer]" ``` Both supported engine extras (`vllm` and `sglang`) install PyAV for raw video diff --git a/examples/ar/qwen3_omni_audio_dcase_gspo_lora_vllm_omni_1x4.yaml b/examples/ar/qwen3_omni_audio_dcase_gspo_lora_vllm_omni_1x4.yaml index d87a40663..05896b3e4 100644 --- a/examples/ar/qwen3_omni_audio_dcase_gspo_lora_vllm_omni_1x4.yaml +++ b/examples/ar/qwen3_omni_audio_dcase_gspo_lora_vllm_omni_1x4.yaml @@ -114,7 +114,7 @@ rollout: _target_: unirl.rollout.engine.vllm_omni.config.VLLMOmniEngineConfig model_path: ${oc.env:QWEN3_OMNI_PATH,/path/to/Qwen3-Omni-30B-A3B-Instruct} modality: qwen3_omni_thinker - stage_yaml_override: qwen3_omni_thinker_only_rl_1x4.yaml + deploy_config_override: qwen3_omni_thinker_only_rl_1x4.yaml enable_sleep_mode: true max_prompt_length: 16384 video_fps: 2.0 diff --git a/examples/ar/qwen3_omni_audio_video_gspo_lora_vllm_omni_1x4.yaml b/examples/ar/qwen3_omni_audio_video_gspo_lora_vllm_omni_1x4.yaml index 9d80b2584..a2ece8e31 100644 --- a/examples/ar/qwen3_omni_audio_video_gspo_lora_vllm_omni_1x4.yaml +++ b/examples/ar/qwen3_omni_audio_video_gspo_lora_vllm_omni_1x4.yaml @@ -124,7 +124,7 @@ rollout: model_path: ${oc.env:QWEN3_OMNI_PATH,/path/to/Qwen3-Omni-30B-A3B-Instruct} modality: qwen3_omni_thinker # Audio-in-video must serialize packed forwards to avoid elevated rollout/replay K3. - stage_yaml_override: qwen3_omni_thinker_only_rl_audio_video_1x4.yaml + deploy_config_override: qwen3_omni_thinker_only_rl_audio_video_1x4.yaml enable_sleep_mode: true max_prompt_length: 16384 video_fps: 2.0 diff --git a/examples/ar/qwen3_omni_audio_video_gspo_lora_vllm_omni_1x8.yaml b/examples/ar/qwen3_omni_audio_video_gspo_lora_vllm_omni_1x8.yaml index 8d3d0d99a..f42dc3b34 100644 --- a/examples/ar/qwen3_omni_audio_video_gspo_lora_vllm_omni_1x8.yaml +++ b/examples/ar/qwen3_omni_audio_video_gspo_lora_vllm_omni_1x8.yaml @@ -114,7 +114,7 @@ rollout: model_path: ${oc.env:QWEN3_OMNI_PATH,/path/to/Qwen3-Omni-30B-A3B-Instruct} modality: qwen3_omni_thinker # Audio-in-video must serialize packed forwards to avoid elevated rollout/replay K3. - stage_yaml_override: qwen3_omni_thinker_only_rl_audio_video_1x4.yaml + deploy_config_override: qwen3_omni_thinker_only_rl_audio_video_1x4.yaml enable_sleep_mode: true max_prompt_length: 16384 video_fps: 2.0 diff --git a/examples/ar/qwen3_omni_image_video_r1_gspo_lora_vllm_omni_1x4.yaml b/examples/ar/qwen3_omni_image_video_r1_gspo_lora_vllm_omni_1x4.yaml index 1131021c6..944ef2741 100644 --- a/examples/ar/qwen3_omni_image_video_r1_gspo_lora_vllm_omni_1x4.yaml +++ b/examples/ar/qwen3_omni_image_video_r1_gspo_lora_vllm_omni_1x4.yaml @@ -109,7 +109,7 @@ rollout: _target_: unirl.rollout.engine.vllm_omni.config.VLLMOmniEngineConfig model_path: ${oc.env:QWEN3_OMNI_PATH,/path/to/Qwen3-Omni-30B-A3B-Instruct} modality: qwen3_omni_thinker - stage_yaml_override: qwen3_omni_thinker_only_rl_1x4.yaml + deploy_config_override: qwen3_omni_thinker_only_rl_1x4.yaml enable_sleep_mode: true max_prompt_length: 16384 image_max_pixels: ${bundle.config.image_max_pixels} diff --git a/examples/ar/qwen3_omni_video_r1_gspo_lora_vllm_omni_1x4.yaml b/examples/ar/qwen3_omni_video_r1_gspo_lora_vllm_omni_1x4.yaml index 17e123817..96391067d 100644 --- a/examples/ar/qwen3_omni_video_r1_gspo_lora_vllm_omni_1x4.yaml +++ b/examples/ar/qwen3_omni_video_r1_gspo_lora_vllm_omni_1x4.yaml @@ -122,7 +122,7 @@ rollout: _target_: unirl.rollout.engine.vllm_omni.config.VLLMOmniEngineConfig model_path: ${oc.env:QWEN3_OMNI_PATH,/path/to/Qwen3-Omni-30B-A3B-Instruct} modality: qwen3_omni_thinker - stage_yaml_override: qwen3_omni_thinker_only_rl_1x4.yaml # TP=4 (audio tower 20 heads) + deploy_config_override: qwen3_omni_thinker_only_rl_1x4.yaml # TP=4 (audio tower 20 heads) enable_sleep_mode: true omni_extra: stage_init_timeout: 1200 diff --git a/examples/ar/qwen3_omni_video_r1_gspo_lora_vllm_omni_1x8.yaml b/examples/ar/qwen3_omni_video_r1_gspo_lora_vllm_omni_1x8.yaml index b924219fd..69da1bb14 100644 --- a/examples/ar/qwen3_omni_video_r1_gspo_lora_vllm_omni_1x8.yaml +++ b/examples/ar/qwen3_omni_video_r1_gspo_lora_vllm_omni_1x8.yaml @@ -109,7 +109,7 @@ rollout: _target_: unirl.rollout.engine.vllm_omni.config.VLLMOmniEngineConfig model_path: ${oc.env:QWEN3_OMNI_PATH,/path/to/Qwen3-Omni-30B-A3B-Instruct} modality: qwen3_omni_thinker - stage_yaml_override: qwen3_omni_thinker_only_rl_1x4.yaml # TP=4 (audio tower 20 heads) + deploy_config_override: qwen3_omni_thinker_only_rl_1x4.yaml # TP=4 (audio tower 20 heads) enable_sleep_mode: true omni_extra: stage_init_timeout: 1200 diff --git a/examples/diffusion/bagel/bagel_it2i_managed_editscore.yaml b/examples/diffusion/bagel/bagel_it2i_managed_editscore.yaml index 33e574747..addb8555d 100644 --- a/examples/diffusion/bagel/bagel_it2i_managed_editscore.yaml +++ b/examples/diffusion/bagel/bagel_it2i_managed_editscore.yaml @@ -161,7 +161,7 @@ rollout: # Required; same checkpoint the bundle loads. model_path: ${oc.env:BAGEL_PATH,ByteDance-Seed/BAGEL-7B-MoT} # BAGEL single-stage editing modality (registers BagelIt2iAdapter + boots - # stage_configs/bagel_t2i_rl.yaml — one YAML serves both image-out modalities — + # deploy_configs/bagel_t2i_rl.yaml — one YAML serves both image-out modalities — # with RLBagelPipeline). modality: bagel_it2i # Required for colocate: the trainer calls sleep()/wake_up() around the train diff --git a/examples/diffusion/bagel/bagel_it2i_vllmomni.yaml b/examples/diffusion/bagel/bagel_it2i_vllmomni.yaml index 1c6cf72a6..27fa23bee 100644 --- a/examples/diffusion/bagel/bagel_it2i_vllmomni.yaml +++ b/examples/diffusion/bagel/bagel_it2i_vllmomni.yaml @@ -159,7 +159,7 @@ rollout: # Required; same checkpoint the bundle loads. model_path: ${oc.env:BAGEL_PATH,ByteDance-Seed/BAGEL-7B-MoT} # BAGEL single-stage editing modality (registers BagelIt2iAdapter + boots - # stage_configs/bagel_t2i_rl.yaml — one YAML serves both image-out modalities — + # deploy_configs/bagel_t2i_rl.yaml — one YAML serves both image-out modalities — # with RLBagelPipeline). modality: bagel_it2i # Required for colocate: the trainer calls sleep()/wake_up() around the train diff --git a/examples/diffusion/bagel/bagel_vllmomni.yaml b/examples/diffusion/bagel/bagel_vllmomni.yaml index 50cfcbaa0..364f9195a 100644 --- a/examples/diffusion/bagel/bagel_vllmomni.yaml +++ b/examples/diffusion/bagel/bagel_vllmomni.yaml @@ -119,7 +119,7 @@ rollout: # Required; same checkpoint the bundle loads. model_path: ${oc.env:BAGEL_PATH,/root/hf_model/BAGEL-7B-MoT} # BAGEL single-stage T2I diffusion modality (registers BagelT2iAdapter + - # boots stage_configs/bagel_t2i_rl.yaml with RLBagelPipeline). + # boots deploy_configs/bagel_t2i_rl.yaml with RLBagelPipeline). modality: bagel_t2i # Required for colocate: trainer calls sleep()/wake_up() around the train # phase, a no-op unless CuMemAllocator is enabled at vLLM-Omni init time. diff --git a/examples/diffusion/bagel/bagel_vllmomni_async.yaml b/examples/diffusion/bagel/bagel_vllmomni_async.yaml index 0e1f34495..9f6e85de5 100644 --- a/examples/diffusion/bagel/bagel_vllmomni_async.yaml +++ b/examples/diffusion/bagel/bagel_vllmomni_async.yaml @@ -108,7 +108,7 @@ rollout: # Required; same checkpoint the bundle loads. model_path: ${oc.env:BAGEL_PATH,/root/hf_model/BAGEL-7B-MoT} # BAGEL single-stage T2I diffusion modality (registers BagelT2iAdapter + - # boots stage_configs/bagel_t2i_rl.yaml with RLBagelPipeline). + # boots deploy_configs/bagel_t2i_rl.yaml with RLBagelPipeline). modality: bagel_t2i # Separate slabs don't time-share GPUs, so sleep/wake is unnecessary. enable_sleep_mode: false diff --git a/pyproject.toml b/pyproject.toml index eb4fff5b6..fb22a97f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,10 +41,9 @@ dependencies = [ "tensordict>=0.5", ] -# Parent/inline engines load this. Multi-stage StageDiffusionProc does not — -# VLLMOmniHijack.hijack() reinstalls the same hook in spawn children. +# vllm-omni 0.28 loads this in parent and spawned stage processes. [project.entry-points."vllm_omni.general_plugins"] -unirl_capture_flush = "unirl.rollout.engine.vllm_omni.plugin:register_capture_flush" +unirl_runtime = "unirl.rollout.engine.vllm_omni.plugin:register_unirl_runtime" [project.optional-dependencies] # sglang pinned to the upstream release the UniRL sglang patch package @@ -60,6 +59,9 @@ unirl_capture_flush = "unirl.rollout.engine.vllm_omni.plugin:register_capture_fl # (upgrade pending; verified meanwhile via a 2.34-prefixed python). sglang = [ "sglang[diffusion]==0.5.12.post1 ; sys_platform == 'linux'", + # sglang leaves kernels unpinned; transformers 5.6.0 requires its 0.12 API. + # The vllm extra conflicts with this extra and resolves its own 0.14.1 pin. + "kernels>=0.12,<0.13 ; sys_platform == 'linux'", # sglang pins a prerelease; a first-party declaration makes it "explicitly # requested" under uv's default prerelease policy (scoped, not global). "flash-attn-4>=4.0.0b9,<5 ; sys_platform == 'linux'", @@ -81,10 +83,9 @@ sglang = [ ] # CUDA 13: vllm >=0.26 wants a bare `torch==2.13.0` whose wheel is cu130, and a # +cu129 pin would resolve anyway, linking a CUDA-12 torch under a CUDA-13 vllm. -# Prerelease, so installs need --prerelease=allow, plus sglang's compat layer. vllm = [ "vllm==0.28.0 ; sys_platform == 'linux'", - "vllm-omni==0.28.0rc1 ; sys_platform == 'linux'", + "vllm-omni==0.28.0 ; sys_platform == 'linux'", # vllm-omni 0.28 needs >=5.10.1,<5.15 (5.10.0 is yanked, 5.15 has a known # construction regression); exact because there is no uv.lock to hold it, and # 5.12.1 is what sglang 0.5.15+ pins, so unifying later is a bump not a port. @@ -165,26 +166,16 @@ conflicts = [ # other package take its newest from PyPI while +cuXXX outranks same-release # PyPI builds where it matters. index-strategy = "unsafe-best-match" -# sglang 0.5.12.post1 pins transformers==5.6.0 but leaves its `kernels` -# integration unpinned; kernels >=0.15 changed LayerRepository's signature and -# breaks transformers 5.6.0 at import time. Hold it to transformers' own -# declared bound. Scoped to that fork in practice: nothing the vllm extra -# resolves requires kernels, so this does not cap its transformers at <5.13. -# # transformers 5.6.0 accepts tokenizers<=0.23.0, but tokenizers 0.23.0(rc) # dropped the `cls` kwarg from RobertaProcessing.__new__, so deserializing a # saved CLIP tokenizer.json raises TypeError at rollout init (the SD3 text -# encoders, e2e-confirmed on H20). The sglang install's --prerelease=allow flag -# (see INSTALL.md) otherwise selects 0.23.0rc*; this override outranks it AND -# guards against 0.23.0 shipping stable. tokenizers 0.22.2 verified on both -# engine venvs. +# encoders, e2e-confirmed on H20). Keep both engine stacks on the verified +# tokenizers 0.22 family. override-dependencies = [ - "kernels>=0.12,<0.13", "tokenizers>=0.22,<0.23", - # sglang[diffusion] hard-pins diffusers==0.37.0 (still pinned on sglang main), - # colliding with the base diffusers>=0.38.0 floor (#91); override to the floor - # so ".[sglang,...]" stays solvable (same tactic as kernels/tokenizers above). - "diffusers>=0.38.0", + # Override sglang[diffusion]'s stale 0.37 pin with the exact version required + # by vllm-omni 0.28.0. Both mutually exclusive engine stacks use this version. + "diffusers==0.40.0", ] environments = [ "sys_platform == 'linux' and platform_machine == 'x86_64'", diff --git a/unirl/rollout/README.md b/unirl/rollout/README.md index 726087287..dba2650b4 100644 --- a/unirl/rollout/README.md +++ b/unirl/rollout/README.md @@ -90,7 +90,7 @@ Engine dirs use two layouts. The compact engines (`trainside`, `fastvideo`, management), `utils/`, `weight_sync.py`, and a runtime-patch dir for the pinned upstream (`sglang_diffusion/_patches/`, `vllm_omni/patches/`). `vllm_omni` additionally carries worker-subprocess code -(`pipelines/`, `worker/`) and stage boot configs (`stage_configs/`). +(`pipelines/`, `worker/`) and deployment configs (`deploy_configs/`). Model onboarding is per-engine, and the adapter file is usually **not** the whole change surface: @@ -107,7 +107,7 @@ change surface: model needs a new upstream patch. - **`vllm_omni`:** add an `adapters/.py` binder (keyed by modality), register it, import it in `adapters/__init__.py`, and add the appropriate boot - YAML under `stage_configs/`. DiT families additionally need a worker-side + YAML under `deploy_configs/`. DiT families additionally need a worker-side `pipelines//pipeline.py`; if the AR/DiT worker needs new behavior, add a `worker/` extension or `patches/compat_.py`. diff --git a/unirl/rollout/engine/vllm_omni/adapters/bagel.py b/unirl/rollout/engine/vllm_omni/adapters/bagel.py index d88221941..6a4e8b95b 100644 --- a/unirl/rollout/engine/vllm_omni/adapters/bagel.py +++ b/unirl/rollout/engine/vllm_omni/adapters/bagel.py @@ -245,7 +245,7 @@ def build_conditions(self, sample: Sample, per_request: List[List[OmniRawResult] class BagelAdapter(ModelAdapter): """Bind BAGEL t2i and it2i to one single-stage DiT worker.""" - stage_yaml = "bagel_t2i_rl.yaml" + deploy_config = "bagel_t2i_rl.yaml" omni_mode = "text-to-image" needs_driver_tokenizer = False image_input: bool = False # Whether the modality requires an edit-source image. diff --git a/unirl/rollout/engine/vllm_omni/adapters/base.py b/unirl/rollout/engine/vllm_omni/adapters/base.py index a76a15eae..26ffc0ccd 100644 --- a/unirl/rollout/engine/vllm_omni/adapters/base.py +++ b/unirl/rollout/engine/vllm_omni/adapters/base.py @@ -45,7 +45,7 @@ class ModelAdapter(ABC): modality: str = "" - stage_yaml: str = "" + deploy_config: str = "" omni_mode: Optional[str] = None needs_sigmas: bool = True needs_driver_tokenizer: bool = True @@ -75,9 +75,9 @@ def resolve_sde_label(strategy: Any) -> Optional[str]: def boot_kwargs(self) -> Dict[str, Any]: """Model-specific boot intent beyond the generic config spelling.""" - require(bool(self.stage_yaml), f"{type(self).__name__} must set stage_yaml") + require(bool(self.deploy_config), f"{type(self).__name__} must set deploy_config") kwargs: Dict[str, Any] = { - "stage_yaml": self.stage_yaml, + "deploy_config": self.deploy_config, "needs_driver_tokenizer": bool(self.needs_driver_tokenizer), "clear_cuda_visible": bool(self.clear_cuda_visible), } diff --git a/unirl/rollout/engine/vllm_omni/adapters/hi3.py b/unirl/rollout/engine/vllm_omni/adapters/hi3.py index 5c9730c2f..6a497406b 100644 --- a/unirl/rollout/engine/vllm_omni/adapters/hi3.py +++ b/unirl/rollout/engine/vllm_omni/adapters/hi3.py @@ -84,7 +84,7 @@ def hi3_fused_conditions(diff_outputs: List[OmniRawResult], *, modality: str) -> "Check that RLHunyuanImage3Pipeline.prepare_inputs_for_generation " "hook ran in every DiT worker — the subclass swap may not have " "taken effect (verify custom_pipeline_args.pipeline_class in " - "the stage YAML)." + "the deploy config)." ) sequence_lengths = [int(c["input_ids"].shape[-1]) for c in captures] @@ -512,7 +512,7 @@ def build_conditions(self, sample: Sample, per_request: List[List[OmniRawResult] class Hi3T2iAdapter(ModelAdapter): """HI3 text → AR think → DiT image.""" - stage_yaml = "hunyuan_image3_t2i_rl.yaml" + deploy_config = "hunyuan_image3_t2i_rl.yaml" omni_mode = "text-to-image" ar_lora_passthrough = True clear_cuda_visible = True @@ -548,7 +548,7 @@ def build_response(self, sample: Sample, per_request: List[List[OmniRawResult]]) class Hi3It2iAdapter(ModelAdapter): """HI3 image+text → AR recaption → DiT edited image.""" - stage_yaml = "hunyuan_image3_it2i_rl.yaml" + deploy_config = "hunyuan_image3_it2i_rl.yaml" omni_mode = "text-to-image" ar_lora_passthrough = True clear_cuda_visible = True @@ -582,9 +582,9 @@ def build_response(self, sample: Sample, per_request: List[List[OmniRawResult]]) @register_adapter("hi3_i2t") class Hi3I2tAdapter(ModelAdapter): - """HI3 image+text → AR text (vendored comprehension YAML).""" + """HI3 image+text → AR text.""" - stage_yaml = "hunyuan_image3_i2t.yaml" + deploy_config = "hunyuan_image3_i2t_rl.yaml" needs_sigmas = False ar_lora_passthrough = True clear_cuda_visible = True @@ -617,9 +617,9 @@ def build_response(self, sample: Sample, per_request: List[List[OmniRawResult]]) @register_adapter("hi3_t2t") class Hi3T2tAdapter(ModelAdapter): - """HI3 text → AR text (vendored comprehension YAML).""" + """HI3 text → AR text.""" - stage_yaml = "hunyuan_image3_t2t.yaml" + deploy_config = "hunyuan_image3_ar_rl.yaml" needs_sigmas = False ar_lora_passthrough = True clear_cuda_visible = True @@ -652,7 +652,7 @@ def build_response(self, sample: Sample, per_request: List[List[OmniRawResult]]) class Hi3ArRecaptionAdapter(ModelAdapter): """Two-engine trainer's AR think/recaption producer.""" - stage_yaml = "hunyuan_image3_ar_recaption_rl.yaml" + deploy_config = "hunyuan_image3_ar_recaption_rl.yaml" needs_sigmas = False ar_lora_passthrough = True clear_cuda_visible = True @@ -681,7 +681,7 @@ def build_response(self, sample: Sample, per_request: List[List[OmniRawResult]]) class Hi3DitRecaptionAdapter(ModelAdapter): """Standalone HI3 DiT — the two-engine trainer's image half.""" - stage_yaml = "hunyuan_image3_dit_recaption_rl.yaml" + deploy_config = "hunyuan_image3_dit_recaption_rl.yaml" omni_mode = "text-to-image" clear_cuda_visible = True # HI3 two-engine stages are TP>1 — wake-time LoRA re-push must use the byte-copy transport. diff --git a/unirl/rollout/engine/vllm_omni/adapters/hv15.py b/unirl/rollout/engine/vllm_omni/adapters/hv15.py index 079d61867..61118c216 100644 --- a/unirl/rollout/engine/vllm_omni/adapters/hv15.py +++ b/unirl/rollout/engine/vllm_omni/adapters/hv15.py @@ -54,8 +54,8 @@ class Hv15VideoOutputAdapter(DitOutputAdapter): "on the output envelope's unirl metadata (or it lacked the dual-stream " "text_mllm/text_glyph embeds). Check that " "RLHunyuanVideo15Pipeline's encode_prompt hook ran in every DiT " - "worker — verify custom_pipeline_args.pipeline_class in the stage " - "YAML." + "worker — verify custom_pipeline_args.pipeline_class in the deploy " + "config." ) def build_decoded(self, sample: Sample, per_request: List[List[OmniRawResult]]) -> Any: @@ -116,7 +116,7 @@ def _cat_field(field_name: str) -> Optional[torch.Tensor]: class Hv15T2vAdapter(ModelAdapter): """HunyuanVideo-1.5 text → video (single diffusion stage, TP=1).""" - stage_yaml = "hunyuan_video15_t2v_rl.yaml" + deploy_config = "hunyuan_video15_t2v_rl.yaml" needs_driver_tokenizer = False def __init__(self, config: Any, model_config: Any, *, strategy: Any = None, tokenize_fn: Any = None) -> None: diff --git a/unirl/rollout/engine/vllm_omni/adapters/qwen3_omni.py b/unirl/rollout/engine/vllm_omni/adapters/qwen3_omni.py index 1df55989c..2ad3cddfb 100644 --- a/unirl/rollout/engine/vllm_omni/adapters/qwen3_omni.py +++ b/unirl/rollout/engine/vllm_omni/adapters/qwen3_omni.py @@ -457,8 +457,7 @@ def _status(cls, per_request: List[List[OmniRawResult]]) -> torch.Tensor: values: List[int] = [] for group in per_request: output = cls._stage0(group) - request_output = getattr(output, "request_output", None) - completions = getattr(request_output, "outputs", None) or [] + completions = getattr(output, "outputs", None) or [] finish_reason = getattr(completions[0], "finish_reason", None) if completions else None values.append(int(mapping.get(str(finish_reason), SegmentStatus.PENDING))) return torch.tensor(values, dtype=torch.long) @@ -495,7 +494,7 @@ def build(self, sample: Sample, per_request: List[List[OmniRawResult]]) -> Sampl class Qwen3OmniThinkerAdapter(ModelAdapter): """Qwen3-Omni Thinker — text/video → AR text (single stage, TP>1, LoRA).""" - stage_yaml = "qwen3_omni_thinker_only_rl_1x4.yaml" + deploy_config = "qwen3_omni_thinker_only_rl_1x4.yaml" omni_mode = None needs_sigmas = False needs_driver_tokenizer = False diff --git a/unirl/rollout/engine/vllm_omni/adapters/qwen_image.py b/unirl/rollout/engine/vllm_omni/adapters/qwen_image.py index 0d2e0abe3..cfe6b4f60 100644 --- a/unirl/rollout/engine/vllm_omni/adapters/qwen_image.py +++ b/unirl/rollout/engine/vllm_omni/adapters/qwen_image.py @@ -97,7 +97,7 @@ class QwenImageOutputAdapter(DitOutputAdapter): "the output envelope's unirl metadata. Check that RLQwenImagePipeline's " "encode_prompt tap ran in every DiT worker — the subclass swap may " "not have taken effect (verify custom_pipeline_args.pipeline_class " - "in the stage YAML)." + "in the deploy config)." ) def build_conditions(self, sample: Sample, per_request: List[List[OmniRawResult]]) -> Dict[str, Any]: @@ -138,7 +138,7 @@ def build_conditions(self, sample: Sample, per_request: List[List[OmniRawResult] class QwenImageT2iAdapter(ModelAdapter): """Qwen-Image text → image (single diffusion stage, TP=1).""" - stage_yaml = "qwen_image_t2i_rl.yaml" + deploy_config = "qwen_image_t2i_rl.yaml" omni_mode = "text-to-image" needs_driver_tokenizer = False diff --git a/unirl/rollout/engine/vllm_omni/adapters/sd3.py b/unirl/rollout/engine/vllm_omni/adapters/sd3.py index c4e06e935..44f264838 100644 --- a/unirl/rollout/engine/vllm_omni/adapters/sd3.py +++ b/unirl/rollout/engine/vllm_omni/adapters/sd3.py @@ -60,7 +60,7 @@ def build_conditions(self, sample: Sample, per_request: List[List[OmniRawResult] "RLStableDiffusion3Pipeline._install_encode_prompt_hook ran " "in every DiT worker — the subclass swap may not have taken " "effect (verify custom_pipeline_args.pipeline_class in the " - "stage YAML)." + "deploy config)." ) embeds = torch.cat([c["prompt_embeds"] for c in captures], dim=0) @@ -93,7 +93,7 @@ def build_conditions(self, sample: Sample, per_request: List[List[OmniRawResult] class Sd3T2iAdapter(ModelAdapter): """SD3.5-medium text → image (single diffusion stage, TP=1).""" - stage_yaml = "sd35_t2i_rl.yaml" + deploy_config = "sd35_t2i_rl.yaml" omni_mode = "text-to-image" needs_driver_tokenizer = False diff --git a/unirl/rollout/engine/vllm_omni/backends/base.py b/unirl/rollout/engine/vllm_omni/backends/base.py index c7f724ad8..4397c9dff 100644 --- a/unirl/rollout/engine/vllm_omni/backends/base.py +++ b/unirl/rollout/engine/vllm_omni/backends/base.py @@ -61,7 +61,7 @@ class OmniRawResult(Protocol): request_id: str stage_id: Optional[int] final_output_type: Optional[str] - request_output: Optional[Any] + outputs: Sequence[Any] prompt_token_ids: Optional[Sequence[int]] images: Optional[Sequence[Any]] trajectory_latents: Optional["torch.Tensor"] diff --git a/unirl/rollout/engine/vllm_omni/backends/native.py b/unirl/rollout/engine/vllm_omni/backends/native.py index 5acdd47ac..49d6126de 100644 --- a/unirl/rollout/engine/vllm_omni/backends/native.py +++ b/unirl/rollout/engine/vllm_omni/backends/native.py @@ -54,12 +54,12 @@ def _import_omni_runtime() -> Dict[str, Any]: } -def _resolve_stage_yaml(name: str) -> str: - """Return the absolute path of the local stage-config YAML asset.""" +def _resolve_deploy_config(name: str) -> str: + """Return the absolute path of the local deploy-config YAML asset.""" here = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) - path = os.path.join(here, "stage_configs", name) + path = os.path.join(here, "deploy_configs", name) if not os.path.exists(path): - raise FileNotFoundError(f"_resolve_stage_yaml: YAML not found at {path}") + raise FileNotFoundError(f"_resolve_deploy_config: YAML not found at {path}") return path @@ -141,9 +141,11 @@ def boot(cls, intent: Dict[str, Any]) -> "VLLMOmniBackend": # colocated train actor's state; spawned workers initialize their own. with _preserve_cudnn_sdp_state(): from unirl.rollout.engine.vllm_omni.patches import install as install_patches + from unirl.rollout.engine.vllm_omni.plugin import register_unirl_runtime install_patches() rt = _import_omni_runtime() + register_unirl_runtime() if intent.get("clear_cuda_visible"): os.environ.pop("CUDA_VISIBLE_DEVICES", None) @@ -159,7 +161,7 @@ def boot(cls, intent: Dict[str, Any]) -> "VLLMOmniBackend": except Exception: # noqa: BLE001 - belt and braces; never block a boot pass - yaml_path = _resolve_stage_yaml(str(intent["stage_yaml"])) + deploy_config_path = _resolve_deploy_config(str(intent["deploy_config"])) omni_kwargs = _assemble_omni_kwargs(intent) ports = intent.get("ports") boot_master_port = int(ports.master_port) if ports is not None else None @@ -168,7 +170,7 @@ def boot(cls, intent: Dict[str, Any]) -> "VLLMOmniBackend": pformat( { **intent, - "stage_yaml_path": yaml_path, + "deploy_config_path": deploy_config_path, "assembled_omni_kwargs": omni_kwargs, "boot_master_port": boot_master_port, }, @@ -185,7 +187,7 @@ def boot(cls, intent: Dict[str, Any]) -> "VLLMOmniBackend": with _master_port_env(boot_master_port): omni = rt["Omni"]( model=str(intent["model_path"]), - stage_configs_path=yaml_path, + deploy_config=deploy_config_path, **omni_kwargs, ) finally: @@ -196,12 +198,13 @@ def boot(cls, intent: Dict[str, Any]) -> "VLLMOmniBackend": try: from omegaconf import OmegaConf + stage_configs = omni.engine.stage_configs resolved_stage_configs = OmegaConf.to_container( - OmegaConf.create(omni.stage_configs), + OmegaConf.create(stage_configs), resolve=True, ) except Exception: # noqa: BLE001 - config logging must never block boot - resolved_stage_configs = omni.stage_configs + resolved_stage_configs = omni.engine.stage_configs logger.info( "VLLM-Omni resolved runtime stage configs (after all overrides):\n%s", pformat(resolved_stage_configs, sort_dicts=True), @@ -215,7 +218,7 @@ def boot(cls, intent: Dict[str, Any]) -> "VLLMOmniBackend": omni, rt, tokenizer=tokenizer, - tp_per_stage=_tp_from_stage_configs(omni.stage_configs), + tp_per_stage=_tp_from_stage_configs(omni.engine.stage_configs), ) except BaseException: logger.exception("VLLM-Omni boot failed; tearing down any engine processes") @@ -301,6 +304,49 @@ def tp_per_stage(self) -> Dict[int, int]: def _stage_ids(self) -> List[int]: return list(range(self.num_stages())) + def _stage_type(self, stage_id: int) -> str: + metadata = self._require_omni().engine.get_stage_metadata(stage_id) + return str(getattr(metadata, "stage_type", "llm")) + + @staticmethod + def _require_rpc_success(action: str, stage_id: int, results: object) -> None: + """Fail closed when a stage or replica reports an unsupported/failed RPC.""" + + def validate(result: object) -> None: + if result is None: + return + if isinstance(result, (list, tuple)): + for item in result: + validate(item) + return + if not isinstance(result, Mapping): + return + + status = result.get("status") + error = result.get("error") + if result.get("supported") is False or error or (status is not None and status != "SUCCESS"): + reason = result.get("reason") or result.get("error_msg") or error or status + raise RuntimeError(f"vllm-omni {action} failed on stage {stage_id}: {reason!r}") + + validate(results) + + def _collective_rpc( + self, + stage_id: int, + method: str, + *, + args: tuple = (), + kwargs: Optional[dict] = None, + ) -> List[Any]: + results = self._require_omni().engine.collective_rpc( + method=method, + args=args, + kwargs=kwargs, + stage_ids=[stage_id], + ) + self._require_rpc_success(method, stage_id, results) + return results + @staticmethod def _require_ack_success(action: str, stage_id: int, task_id: str, acks: object) -> None: # Worker handlers catch their own exceptions and answer @@ -351,11 +397,14 @@ def validate_result(result: object) -> None: raise RuntimeError(f"vllm-omni {action} returned no successful ACK for stage {stage_id}") def sleep_task(self) -> None: - """Fan ``handle_sleep_task`` to every stage's workers (level 1).""" + """Sleep AR stages through EngineCore and diffusion stages through worker RPC.""" import uuid omni = self._require_omni() for sid in self._stage_ids(): + if self._stage_type(sid) != "diffusion": + self._collective_rpc(sid, "sleep", args=(1, "abort")) + continue task_id = str(uuid.uuid4()) acks = omni.engine.collective_rpc( method="handle_sleep_task", @@ -365,13 +414,17 @@ def sleep_task(self) -> None: self._require_ack_success("sleep", int(sid), task_id, acks) def wake_task(self) -> None: - """Fan ``handle_wake_task`` to every stage's workers + sync CUDA.""" + """Wake AR stages through EngineCore and diffusion stages through worker RPC.""" import uuid import torch omni = self._require_omni() for sid in self._stage_ids(): + if self._stage_type(sid) != "diffusion": + # ``None`` wakes every tag, including UniRL's MoE workspace tag. + self._collective_rpc(sid, "wake_up", kwargs={"tags": None}) + continue task_id = str(uuid.uuid4()) acks = omni.engine.collective_rpc( method="handle_wake_task", @@ -407,7 +460,6 @@ def update_from_ipc( replica_rank: Optional[int], ) -> None: """Fan a bucketed CUDA-IPC state-dict update out to per-stage workers.""" - omni = self._require_omni() kwargs = { "peft_config": peft_config, "base_sync_done": base_sync_done, @@ -415,11 +467,10 @@ def update_from_ipc( "replica_rank": replica_rank, } for sid in self._stage_ids(): - omni.engine.collective_rpc( - method="update_weights_from_ipc", - args=(), + self._collective_rpc( + sid, + "update_weights_from_ipc", kwargs={**kwargs, "stage_id": int(sid)}, - stage_ids=[int(sid)], ) def init_weights_group( @@ -432,7 +483,6 @@ def init_weights_group( group_name: str, backend: str, ) -> None: - omni = self._require_omni() kwargs = { "master_address": str(master_address), "master_port": int(master_port), @@ -442,11 +492,10 @@ def init_weights_group( "backend": str(backend), } for sid in self._stage_ids(): - omni.engine.collective_rpc( - method="init_weights_update_group", - args=(), + self._collective_rpc( + sid, + "init_weights_update_group", kwargs=kwargs, - stage_ids=[int(sid)], ) def update_from_distributed( @@ -459,7 +508,6 @@ def update_from_distributed( target_modules: Optional[List[str]], flush_cache: bool, ) -> None: - omni = self._require_omni() kwargs = { "names": list(names), "dtypes": list(dtypes), @@ -469,22 +517,20 @@ def update_from_distributed( "flush_cache": bool(flush_cache), } for sid in self._stage_ids(): - omni.engine.collective_rpc( - method="update_weights_from_distributed", - args=(), + self._collective_rpc( + sid, + "update_weights_from_distributed", kwargs=kwargs, - stage_ids=[int(sid)], ) def destroy_weights_group(self, *, group_name: str) -> None: if self._omni is None: return for sid in self._stage_ids(): - self._omni.engine.collective_rpc( - method="destroy_weights_update_group", - args=(), + self._collective_rpc( + sid, + "destroy_weights_update_group", kwargs={"group_name": str(group_name)}, - stage_ids=[int(sid)], ) def update_from_tensor( @@ -496,7 +542,6 @@ def update_from_tensor( flush_cache: bool, ) -> None: """Fan a SGLang-shape tensor payload to per-stage workers.""" - omni = self._require_omni() kwargs = { "serialized_named_tensors": list(serialized_named_tensors), "target_modules": list(target_modules) if target_modules else None, @@ -504,11 +549,10 @@ def update_from_tensor( "flush_cache": bool(flush_cache), } for sid in self._stage_ids(): - omni.engine.collective_rpc( - method="update_weights_from_tensor", - args=(), + self._collective_rpc( + sid, + "update_weights_from_tensor", kwargs=kwargs, - stage_ids=[int(sid)], ) def set_lora_handle( @@ -527,7 +571,6 @@ def set_lora_handle( DIFFRL_LORA_PATH, ) - omni = self._require_omni() lora_tensors = self._wrap_peft_envelope(lora_tensors) self._remove_existing_lora(int(DIFFRL_LORA_INT_ID)) @@ -540,8 +583,9 @@ def set_lora_handle( name: t.detach().clone() if isinstance(t, torch.Tensor) else t for name, t in lora_tensors.items() } serialized = MultiprocessingSerializer.serialize(cloned, output_str=True) - omni.engine.collective_rpc( - method="set_lora_from_tensor_dict", + self._collective_rpc( + sid, + "set_lora_from_tensor_dict", args=( str(adapter_name) or DIFFRL_LORA_NAME, int(DIFFRL_LORA_INT_ID), @@ -549,7 +593,6 @@ def set_lora_handle( dict(peft_config or {}), serialized, ), - stage_ids=[int(sid)], ) def set_lora_copy( @@ -571,7 +614,6 @@ def set_lora_copy( DIFFRL_LORA_PATH, ) - omni = self._require_omni() lora_tensors = self._wrap_peft_envelope(lora_tensors) self._remove_existing_lora(int(DIFFRL_LORA_INT_ID)) @@ -583,8 +625,9 @@ def set_lora_copy( serialized = base64.b64encode(buf.getvalue()).decode("ascii") for sid in self._stage_ids(): - omni.engine.collective_rpc( - method="set_lora_from_tensor_dict_copy", + self._collective_rpc( + sid, + "set_lora_from_tensor_dict_copy", args=( str(adapter_name) or DIFFRL_LORA_NAME, int(DIFFRL_LORA_INT_ID), @@ -592,7 +635,6 @@ def set_lora_copy( dict(peft_config or {}), serialized, ), - stage_ids=[int(sid)], ) @staticmethod @@ -607,39 +649,29 @@ def _wrap_peft_envelope(lora_tensors: Dict[str, Any]) -> Dict[str, Any]: def _remove_existing_lora(self, adapter_id: int) -> None: """Drop the existing adapter on every stage before re-adding.""" - omni = self._require_omni() for sid in self._stage_ids(): - try: - omni.engine.collective_rpc( - method="remove_lora", - args=(int(adapter_id),), - stage_ids=[int(sid)], - ) - except Exception: - pass + self._collective_rpc(sid, "remove_lora", args=(int(adapter_id),)) def param_checksums(self, *, names: List[str]) -> dict: """Fan ``_diffrl_loaded_param_checksums`` across stages and ranks.""" - omni = self._require_omni() out: dict = {} for sid in self._stage_ids(): - results = omni.engine.collective_rpc( - method="_diffrl_loaded_param_checksums", + results = self._collective_rpc( + sid, + "_diffrl_loaded_param_checksums", args=(list(names),), - stage_ids=[int(sid)], ) out[int(sid)] = results[0] if isinstance(results, list) and results else results return out def lora_checksums(self, *, adapter_id: int, names: Optional[List[str]]) -> dict: """Fan ``_diffrl_loaded_lora_checksums`` across stages and ranks.""" - omni = self._require_omni() out: dict = {} for sid in self._stage_ids(): - results = omni.engine.collective_rpc( - method="_diffrl_loaded_lora_checksums", + results = self._collective_rpc( + sid, + "_diffrl_loaded_lora_checksums", args=(int(adapter_id), list(names) if names else None), - stage_ids=[int(sid)], ) out[int(sid)] = results[0] if isinstance(results, list) and results else results return out diff --git a/unirl/rollout/engine/vllm_omni/config.py b/unirl/rollout/engine/vllm_omni/config.py index ed4520524..7bd98562b 100644 --- a/unirl/rollout/engine/vllm_omni/config.py +++ b/unirl/rollout/engine/vllm_omni/config.py @@ -31,7 +31,7 @@ def make_engine(self, **deps: Any): enable_sleep_mode: bool = True - stage_yaml_override: Optional[str] = None + deploy_config_override: Optional[str] = None omni_extra: Dict[str, Any] = field(default_factory=dict) @@ -71,8 +71,8 @@ def server_intent( } intent.update(extra) - if self.stage_yaml_override: - intent["stage_yaml"] = str(self.stage_yaml_override) + if self.deploy_config_override: + intent["deploy_config"] = str(self.deploy_config_override) omni_kwargs: Dict[str, Any] = dict( stage_init_timeout=1200, diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/bagel_t2i_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/bagel_t2i_rl.yaml new file mode 100644 index 000000000..2ef9de384 --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/bagel_t2i_rl.yaml @@ -0,0 +1,35 @@ +# BAGEL-7B-MoT single-stage deployment shared by t2i and it2i rollout. +pipeline: bagel_single_stage +async_chunk: false +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + vae_use_slicing: false + vae_use_tiling: false + cache_backend: null + cache_config: null + enable_cache_dit_summary: false + enable_lora: true + max_lora_rank: 64 + parallel_config: + pipeline_parallel_size: 1 + data_parallel_size: 1 + tensor_parallel_size: 1 + enable_expert_parallel: false + sequence_parallel_size: 1 + ulysses_degree: 1 + ring_degree: 1 + cfg_parallel_size: 1 + vae_patch_parallel_size: 1 + use_hsdp: false + hsdp_shard_size: -1 + hsdp_replicate_size: 1 + custom_pipeline_args: + pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.bagel.pipeline.RLBagelPipeline + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_recaption_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_recaption_rl.yaml new file mode 100644 index 000000000..ce0653938 --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_recaption_rl.yaml @@ -0,0 +1,31 @@ +# HunyuanImage-3 terminal AR recaption deployment for the two-engine trainer. +pipeline: unirl_hi3_ar_text +async_chunk: false +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "0,1,2,3" + max_num_seqs: 1 + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.ar_extension.HI3ARWeightSyncExtension + enable_lora: true + fully_sharded_loras: true + max_model_len: 8192 + gpu_memory_utilization: 0.7 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + max_num_batched_tokens: 8192 + tensor_parallel_size: 4 + pipeline_parallel_size: 1 + hf_overrides: + rope_parameters: + mrope_section: [0, 32, 32] + rope_type: default + default_sampling_params: + temperature: 0.7 + top_p: 0.95 + top_k: 1024 + max_tokens: 4096 + detokenize: true + repetition_penalty: 1.1 diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml new file mode 100644 index 000000000..be521e8cc --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml @@ -0,0 +1,26 @@ +# HunyuanImage-3 text-only AR deployment for the t2t adapter. +pipeline: unirl_hi3_ar_text +async_chunk: false + +stages: + - stage_id: 0 + devices: "0,1,2,3" + max_num_seqs: 1 + gpu_memory_utilization: 0.95 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + max_num_batched_tokens: 32768 + tensor_parallel_size: 4 + pipeline_parallel_size: 1 + hf_overrides: + rope_parameters: + mrope_section: [0, 32, 32] + rope_type: default + default_sampling_params: + temperature: 0.0 + top_p: 0.95 + top_k: 1024 + max_tokens: 2048 + stop_token_ids: [127957, 128024, 128026] + detokenize: true diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_dit_recaption_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_dit_recaption_rl.yaml new file mode 100644 index 000000000..442163bbc --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_dit_recaption_rl.yaml @@ -0,0 +1,25 @@ +# HunyuanImage-3 standalone DiT deployment for the two-engine trainer. +pipeline: hunyuan_image3_dit +async_chunk: false +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "4,5,6,7" + max_num_seqs: 1 + gpu_memory_utilization: 0.7 + enforce_eager: true + trust_remote_code: true + vae_use_slicing: false + vae_use_tiling: false + cache_backend: null + cache_config: null + enable_cache_dit_summary: false + omni_kv_config: + need_recv_cache: true + parallel_config: + tensor_parallel_size: 4 + enable_expert_parallel: true + custom_pipeline_args: + pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.hi3.pipeline.RLHunyuanImage3Pipeline + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml new file mode 100644 index 000000000..cb650b959 --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml @@ -0,0 +1,26 @@ +# HunyuanImage-3 multimodal AR deployment for the i2t adapter. +pipeline: unirl_hi3_ar_multimodal_text +async_chunk: false + +stages: + - stage_id: 0 + devices: "0,1,2,3" + max_num_seqs: 1 + gpu_memory_utilization: 0.95 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + max_num_batched_tokens: 32768 + tensor_parallel_size: 4 + pipeline_parallel_size: 1 + hf_overrides: + rope_parameters: + mrope_section: [0, 32, 32] + rope_type: default + default_sampling_params: + temperature: 0.0 + top_p: 0.95 + top_k: 1024 + max_tokens: 2048 + stop_token_ids: [127957, 128024, 128026] + detokenize: true diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_it2i_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_it2i_rl.yaml new file mode 100644 index 000000000..9caadb965 --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_it2i_rl.yaml @@ -0,0 +1,48 @@ +# HunyuanImage-3 image-to-image deployment with UniRL worker extensions. +pipeline: hunyuan_image_3_moe +async_chunk: false + +stages: + - stage_id: 0 + devices: "0,1,2,3" + max_num_seqs: 1 + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.ar_extension.HI3ARWeightSyncExtension + enable_lora: true + gpu_memory_utilization: 0.95 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + max_num_batched_tokens: 32768 + tensor_parallel_size: 4 + pipeline_parallel_size: 1 + hf_overrides: + rope_parameters: + mrope_section: [0, 32, 32] + rope_type: default + default_sampling_params: + temperature: 0.6 + top_p: 0.95 + top_k: 1024 + max_tokens: 4096 + stop_token_ids: [127957] + detokenize: true + + - stage_id: 1 + devices: "4,5,6,7" + max_num_seqs: 1 + enforce_eager: true + trust_remote_code: true + distributed_executor_backend: mp + parallel_config: + tensor_parallel_size: 4 + enable_expert_parallel: true + custom_pipeline_args: + pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.hi3.pipeline.RLHunyuanImage3Pipeline + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension + default_sampling_params: + num_inference_steps: 50 + guidance_scale: 2.5 + +edges: + - from: 0 + to: 1 diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_t2i_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_t2i_rl.yaml new file mode 100644 index 000000000..ecaf64e79 --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_t2i_rl.yaml @@ -0,0 +1,69 @@ +# HunyuanImage-3 text-to-image deployment with UniRL worker extensions. +pipeline: hunyuan_image_3_moe +async_chunk: false + +stages: + - stage_id: 0 + devices: "0,1,2,3" + max_num_seqs: 1 + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.ar_extension.HI3ARWeightSyncExtension + enable_lora: true + gpu_memory_utilization: 0.9 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + max_num_batched_tokens: 32768 + tensor_parallel_size: 4 + pipeline_parallel_size: 1 + distributed_executor_backend: mp + hf_overrides: + rope_parameters: + mrope_section: [0, 32, 32] + rope_type: default + omni_kv_config: + need_send_cache: true + kv_transfer_criteria: + type: prefill_finished + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 2048 + seed: 42 + detokenize: true + repetition_penalty: 1.1 + + - stage_id: 1 + devices: "4,5,6,7" + max_num_seqs: 1 + enforce_eager: true + distributed_executor_backend: mp + vae_use_slicing: false + vae_use_tiling: false + cache_backend: null + cache_config: null + enable_cache_dit_summary: false + omni_kv_config: + need_recv_cache: true + parallel_config: + pipeline_parallel_size: 1 + data_parallel_size: 1 + tensor_parallel_size: 4 + enable_expert_parallel: false + sequence_parallel_size: 1 + ulysses_degree: 1 + ring_degree: 1 + cfg_parallel_size: 1 + vae_patch_parallel_size: 1 + use_hsdp: false + hsdp_shard_size: -1 + hsdp_replicate_size: 1 + custom_pipeline_args: + pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.hi3.pipeline.RLHunyuanImage3Pipeline + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension + +edges: + - from: 0 + to: 1 + window_size: -1 + max_inflight: 1 diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_video15_t2v_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_video15_t2v_rl.yaml new file mode 100644 index 000000000..4e73f3dc0 --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_video15_t2v_rl.yaml @@ -0,0 +1,34 @@ +# HunyuanVideo-1.5 single-stage text-to-video deployment for RL rollout. +pipeline: hunyuan_video_15 +async_chunk: false +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + enforce_eager: true + vae_use_slicing: false + vae_use_tiling: false + cache_backend: null + cache_config: null + enable_cache_dit_summary: false + diffusion_attention_backend: TORCH_SDPA + enable_lora: true + max_lora_rank: 64 + parallel_config: + pipeline_parallel_size: 1 + data_parallel_size: 1 + tensor_parallel_size: 1 + enable_expert_parallel: false + sequence_parallel_size: 1 + ulysses_degree: 1 + ring_degree: 1 + cfg_parallel_size: 1 + vae_patch_parallel_size: 1 + use_hsdp: false + hsdp_shard_size: -1 + hsdp_replicate_size: 1 + custom_pipeline_args: + pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.hv15.pipeline.RLHunyuanVideo15Pipeline + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_1x4.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_1x4.yaml new file mode 100644 index 000000000..375d90aad --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_1x4.yaml @@ -0,0 +1,30 @@ +# Qwen3-Omni Thinker-only deployment, TP=4. +pipeline: qwen3_omni_moe_thinker_only +async_chunk: false +distributed_executor_backend: mp +dtype: bfloat16 +trust_remote_code: true +enable_prefix_caching: false +enable_chunked_prefill: true + +stages: + - stage_id: 0 + devices: "0,1,2,3" + tensor_parallel_size: 4 + gpu_memory_utilization: 0.35 + max_num_seqs: 64 + max_num_batched_tokens: 32768 + limit_mm_per_prompt: + audio: 1 + video: 1 + image: 1 + mm_processor_cache_gb: 0 + load_format: safetensors + enforce_eager: false + enable_lora: true + max_lora_rank: 64 + max_loras: 1 + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.qwen3_omni_ar_extension.Qwen3OmniARWeightSyncExtension + enable_sleep_mode: true + logprobs_mode: processed_logprobs + disable_log_stats: true diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_audio_video_1x4.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_audio_video_1x4.yaml new file mode 100644 index 000000000..01cea9016 --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_audio_video_1x4.yaml @@ -0,0 +1,30 @@ +# Qwen3-Omni Thinker-only audio-in-video deployment, TP=4. +pipeline: qwen3_omni_moe_thinker_only +async_chunk: false +distributed_executor_backend: mp +dtype: bfloat16 +trust_remote_code: true +enable_prefix_caching: false +enable_chunked_prefill: true + +stages: + - stage_id: 0 + devices: "0,1,2,3" + tensor_parallel_size: 4 + gpu_memory_utilization: 0.35 + max_num_seqs: 1 + max_num_batched_tokens: 32768 + limit_mm_per_prompt: + audio: 1 + video: 1 + image: 0 + mm_processor_cache_gb: 0 + load_format: safetensors + enforce_eager: false + enable_lora: true + max_lora_rank: 64 + max_loras: 1 + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.qwen3_omni_ar_extension.Qwen3OmniARWeightSyncExtension + enable_sleep_mode: true + logprobs_mode: processed_logprobs + disable_log_stats: true diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/qwen_image_t2i_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/qwen_image_t2i_rl.yaml new file mode 100644 index 000000000..fb089da30 --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/qwen_image_t2i_rl.yaml @@ -0,0 +1,33 @@ +# Single-stage Qwen-Image deployment for RL rollout. +pipeline: unirl_rl_image_diffusion +async_chunk: false +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + enforce_eager: true + vae_use_slicing: false + vae_use_tiling: false + cache_backend: null + cache_config: null + enable_cache_dit_summary: false + enable_lora: true + max_lora_rank: 64 + parallel_config: + pipeline_parallel_size: 1 + data_parallel_size: 1 + tensor_parallel_size: 1 + enable_expert_parallel: false + sequence_parallel_size: 1 + ulysses_degree: 1 + ring_degree: 1 + cfg_parallel_size: 1 + vae_patch_parallel_size: 1 + use_hsdp: false + hsdp_shard_size: -1 + hsdp_replicate_size: 1 + custom_pipeline_args: + pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.qwen_image.pipeline.RLQwenImagePipeline + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/sd35_t2i_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/sd35_t2i_rl.yaml new file mode 100644 index 000000000..168f76c8b --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/sd35_t2i_rl.yaml @@ -0,0 +1,33 @@ +# Single-stage SD3.5-medium deployment for RL rollout. +pipeline: unirl_rl_image_diffusion +async_chunk: false +distributed_executor_backend: mp + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + enforce_eager: true + vae_use_slicing: false + vae_use_tiling: false + cache_backend: null + cache_config: null + enable_cache_dit_summary: false + enable_lora: true + max_lora_rank: 32 + parallel_config: + pipeline_parallel_size: 1 + data_parallel_size: 1 + tensor_parallel_size: 1 + enable_expert_parallel: false + sequence_parallel_size: 1 + ulysses_degree: 1 + ring_degree: 1 + cfg_parallel_size: 1 + vae_patch_parallel_size: 1 + use_hsdp: false + hsdp_shard_size: -1 + hsdp_replicate_size: 1 + custom_pipeline_args: + pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.sd3.pipeline.RLStableDiffusion3Pipeline + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension diff --git a/unirl/rollout/engine/vllm_omni/patches/README.md b/unirl/rollout/engine/vllm_omni/patches/README.md index e318c649f..2633a29f2 100644 --- a/unirl/rollout/engine/vllm_omni/patches/README.md +++ b/unirl/rollout/engine/vllm_omni/patches/README.md @@ -39,7 +39,7 @@ All in `runtime.py` unless noted. | Patch | Why | DELETE-WHEN | | --- | --- | --- | | `wrap_mp_process_for_children` | Re-installs the bundle inside every spawn child. **Must run first** | the rest of the bundle is empty | -| `register_capture_flush` (`plugin.py`) | Flushes `_unirl_captures` into formatter `metadata["unirl"]` after postprocess. The `vllm_omni.general_plugins` entry point is not loaded in `StageDiffusionProc` (HI3's 2-stage path), so hijack installs it in spawn children | vllm-omni loads general plugins in every stage process including `StageDiffusionProc` | +| `register_capture_flush` (`plugin.py`) | Flushes `_unirl_captures` into formatter `metadata["unirl"]` after postprocess. Stable 0.28 loads the general plugin in every stage process, so no child-process monkeypatch is needed | vllm-omni exposes the captured RL metadata natively | | `patch_dit_lora_loader` / `patch_ar_lora_loader` | Stock `DiffusionLoRAManager._load_adapter` loads only from a file path; RL pushes freshly-trained adapter tensors without a disk round-trip (`OmniTensorLoRARequest`). Lifted verbatim from verl-omni | vllm-omni's LoRA managers accept tensor-bag requests natively | | `patch_dit_hi3_lora_weights` | Resolves HI3 DiT `transformer.layers.*` wrappers against PEFT `model.layers.*` keys and converts GQA-interleaved fused-QKV LoRA-B rows to vLLM's packed `[q, k, v]` slices | [vllm-omni #6411](https://github.com/vllm-project/vllm-omni/issues/6411) is fixed and the pinned release includes it | | `patch_fp32_skip` | Punica kernels hard-assert dtype; HI3's MoE router gate is fp32, so non-fp16/bf16 layers must be skipped for LoRA wrapping | vllm's `from_layer` skips unsupported dtypes itself | @@ -50,8 +50,7 @@ All in `runtime.py` unless noted. | `patch_sigmas_passthrough` | HI3's DiT `scheduler.set_timesteps` never receives `sampling_params.sigmas` | upstream forwards `sigmas` itself | | `install_fate_sharing` | `PR_SET_PDEATHSIG` is bound by Linux to the **specific creating thread**, so arming it for children of short-lived init threads kills healthy workers; and a worker inside a CUDA/NCCL call never observes vLLM's `death_pipe` EOF | vllm's own child-reaping is thread-safe | | `compat_tokenizer` (module) | HI3's `__init__` looks up `` and computes `ratio_36 + 1`; the Base checkpoint ships ratio tokens 0-32 only → `TypeError: … 'NoneType' and 'int'`. Both the slow **and** fast tokenizer classes must be patched, not the shared base. The module import *is* the install trigger (it is the `HI3ARWorkerExtension` qualname target). Upstream ≥ v0.20.0 raises a clean `ValueError` instead — a better error, but the Base ckpt still needs this 0-fallback to work | Base-ckpt support is dropped (Instruct ships the tokens) | -| `compat_hi3_lora` (module) | vllm 0.20 expects a flat `list[tuple[str, str, int, str]]` from `get_expert_mapping`; HI3 returns a 2-tuple, so `process_packed_modules_mapping` trips `ValueError: too many values to unpack` at boot under `enable_lora` | vllm handles the 2-tuple / HI3 returns the flat list | -| `patch_moe_workspace_pool` (`compat_moe_workspace`) | vLLM's reusable MoE `WorkspaceManager` scratch buffer is allocated outside the `weights`/`kv_cache` CuMem pools, so worker sleep leaves it resident through the colocated train phase. Gives each DBO ubatch an independently replaceable CuMem pool while tagging all allocations `moe_workspace`; worker sleep then discards the scratch pages with the other non-weight pools and UniRL's untagged wake remaps them. Self-skips when the allocator is absent (sleep mode off → `CuMemAllocator.instance is None`) | vLLM allocates the MoE workspace inside a sleep-managed pool itself | +| `patch_moe_workspace_pool` (`compat_moe_workspace`) | vLLM's reusable MoE `WorkspaceManager` scratch buffer is allocated outside the `weights`/`kv_cache` CuMem pools, so worker sleep leaves it resident through the colocated train phase. Gives each DBO×lane workspace an independently replaceable CuMem pool while tagging all allocations `moe_workspace`; worker sleep then discards the scratch pages and UniRL's untagged wake remaps them. Self-skips when the allocator is absent | vLLM allocates the MoE workspace inside a sleep-managed pool itself | ## HI3 DiT LoRA namespace and layout diff --git a/unirl/rollout/engine/vllm_omni/patches/compat_hi3_lora.py b/unirl/rollout/engine/vllm_omni/patches/compat_hi3_lora.py deleted file mode 100644 index e735a6335..000000000 --- a/unirl/rollout/engine/vllm_omni/patches/compat_hi3_lora.py +++ /dev/null @@ -1,55 +0,0 @@ -"""Runtime patch making ``get_moe_expert_mapping`` tolerate HI3's 2-tuple ``get_expert_mapping`` shape.""" - -from __future__ import annotations - -_INSTALLED = False - - -def install() -> None: - """Patch ``get_moe_expert_mapping`` everywhere it's used, unwrapping HI3's 2-tuple to vllm 0.20's flat list.""" - global _INSTALLED - if _INSTALLED: - return - try: - from vllm.model_executor import utils as vllm_mu - except ImportError: - return - - original = getattr(vllm_mu, "get_moe_expert_mapping", None) - if original is None: - _INSTALLED = True - return - if getattr(original, "_diffrl_hi3_unwrap", False): - _INSTALLED = True - return - - def _patched(model, _orig=original): - result = _orig(model) - if ( - isinstance(result, tuple) - and len(result) == 2 - and not isinstance(result[0], (str, int, float)) - and hasattr(result[0], "__iter__") - and isinstance(result[1], dict) - ): - return result[0] - return result - - _patched._diffrl_hi3_unwrap = True # type: ignore[attr-defined] - vllm_mu.get_moe_expert_mapping = _patched - - try: - from vllm.lora import utils as vllm_lora_utils - - if hasattr(vllm_lora_utils, "get_moe_expert_mapping"): - vllm_lora_utils.get_moe_expert_mapping = _patched - except ImportError: - pass - - _INSTALLED = True - - -install() - - -__all__ = ["install"] diff --git a/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py b/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py index b536d44e0..4c5c8f6bc 100644 --- a/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py +++ b/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py @@ -46,7 +46,7 @@ def _patch_workspace_manager_class( workspace_manager_class: type, *, allocator_provider: Callable[[], Any], - workspace_id_provider: Callable[[Any], int], + workspace_id_provider: Callable[[Any], int | None], ) -> None: """Route ``_ensure_workspace_size`` growth through the dedicated CuMem tag.""" if getattr(workspace_manager_class, _WORKSPACE_MANAGER_MARKER, False): @@ -55,14 +55,14 @@ def _patch_workspace_manager_class( original_ensure_workspace_size = workspace_manager_class._ensure_workspace_size def ensure_workspace_size(manager: Any, required_bytes: int) -> Any: - workspace_id = workspace_id_provider(manager) try: + workspace_id = workspace_id_provider(manager) + if workspace_id is None: + return original_ensure_workspace_size(manager, required_bytes) current_workspace = manager._current_workspaces[workspace_id] current_size = manager._workspace_size_bytes(current_workspace) workspace_locked = manager.is_locked() except (AttributeError, IndexError): - # An out-of-range slot means an unconfigured lane, which vLLM itself - # reports with a precise error once it reaches its own bounds check. return original_ensure_workspace_size(manager, required_bytes) if current_size >= required_bytes: return original_ensure_workspace_size(manager, required_bytes) @@ -116,11 +116,14 @@ def patch_moe_workspace_pool() -> None: except ImportError: _workspace_lane = None - def workspace_id_provider(manager: Any) -> int: - # vLLM 0.28 slots workspaces by lane within ubatch; these defaults - # reproduce the flat per-ubatch layout that earlier versions used. + def workspace_id_provider(manager: Any) -> int | None: + # Delegate invalid lanes before indexing. A lane can otherwise flatten + # onto the next ubatch's valid slot and release an unrelated workspace. lane = 0 if _workspace_lane is None else _workspace_lane.get() - return dbo_current_ubatch_id() * getattr(manager, "_num_lanes", 1) + lane + num_lanes = getattr(manager, "_num_lanes", 1) + if lane < 0 or lane >= num_lanes: + return None + return dbo_current_ubatch_id() * num_lanes + lane _patch_workspace_manager_class( WorkspaceManager, diff --git a/unirl/rollout/engine/vllm_omni/patches/compat_tokenizer.py b/unirl/rollout/engine/vllm_omni/patches/compat_tokenizer.py index a6578b208..ccedfafed 100644 --- a/unirl/rollout/engine/vllm_omni/patches/compat_tokenizer.py +++ b/unirl/rollout/engine/vllm_omni/patches/compat_tokenizer.py @@ -48,8 +48,6 @@ def _filtered(self, tokens, *args, _orig=original, **kwargs): install() -from unirl.rollout.engine.vllm_omni.patches import compat_hi3_lora as _hi3_lora_compat # noqa: F401, E402 - class HI3ARWorkerExtension: """vllm-omni ``worker_extension_cls`` qualname target for HI3 AR.""" diff --git a/unirl/rollout/engine/vllm_omni/patches/runtime.py b/unirl/rollout/engine/vllm_omni/patches/runtime.py index e47a49acf..3631fe324 100644 --- a/unirl/rollout/engine/vllm_omni/patches/runtime.py +++ b/unirl/rollout/engine/vllm_omni/patches/runtime.py @@ -11,15 +11,9 @@ import torch from msgspec import field - -try: - from vllm.lora.lora_model import LoRAModel -except ImportError: - from vllm.lora.models import LoRAModel # type: ignore[no-redef] - +from vllm.lora.lora_model import LoRAModel from vllm.lora.lora_weights import LoRALayerWeights, PackedLoRALayerWeights from vllm.lora.peft_helper import PEFTHelper -from vllm.lora.utils import get_adapter_absolute_path from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager, logger from vllm_omni.lora.request import LoRARequest as OmniLoRARequest @@ -137,28 +131,23 @@ def start(self): def patch_dit_lora_loader() -> None: """Patch ``DiffusionLoRAManager._load_adapter`` (DiT stage) to support in-memory tensors.""" + original = DiffusionLoRAManager._load_adapter + if getattr(original, "_diffrl_tensor_lora_loader", False): + return - def hijack__load_adapter(self, lora_request: OmniTensorLoRARequest) -> tuple[LoRAModel, PEFTHelper]: + def hijack__load_adapter( + self, + lora_request: OmniTensorLoRARequest, + _orig=original, + ) -> tuple[LoRAModel, PEFTHelper]: + if not isinstance(lora_request, OmniTensorLoRARequest): + return _orig(self, lora_request) if not self._expected_lora_modules: raise ValueError("No supported LoRA modules found in the diffusion pipeline.") logger.debug("Supported LoRA modules: %s", self._expected_lora_modules) - lora_tensors = None - - if isinstance(lora_request, OmniTensorLoRARequest): - peft_config = lora_request.peft_config - lora_tensors = lora_request.lora_tensors - peft_helper = PEFTHelper.from_dict(peft_config) - else: - lora_path = get_adapter_absolute_path(lora_request.lora_path) - logger.debug("Resolved LoRA path: %s", lora_path) - - peft_helper = PEFTHelper.from_local_dir( - lora_path, - max_position_embeddings=None, - tensorizer_config_dict=lora_request.tensorizer_config_dict, - ) + peft_helper = PEFTHelper.from_dict(lora_request.peft_config or {}) logger.info( "Loaded PEFT config: r=%d, lora_alpha=%d, target_modules=%s", @@ -167,28 +156,15 @@ def hijack__load_adapter(self, lora_request: OmniTensorLoRARequest) -> tuple[LoR peft_helper.target_modules, ) - if isinstance(lora_request, OmniTensorLoRARequest): - lora_model = LoRAModel.from_lora_tensors( - tensors=lora_tensors, - peft_helper=peft_helper, - lora_model_id=lora_request.lora_int_id, - device="cpu", - dtype=self.dtype, - model_vocab_size=None, - weights_mapper=None, - ) - else: - lora_model = LoRAModel.from_local_checkpoint( - lora_path, - expected_lora_modules=self._expected_lora_modules, - peft_helper=peft_helper, - lora_model_id=lora_request.lora_int_id, - device="cpu", - dtype=self.dtype, - model_vocab_size=None, - tensorizer_config_dict=lora_request.tensorizer_config_dict, - weights_mapper=None, - ) + lora_model = LoRAModel.from_lora_tensors( + tensors=lora_request.lora_tensors or {}, + peft_helper=peft_helper, + lora_model_id=lora_request.lora_int_id, + device="cpu", + dtype=self.dtype, + model_vocab_size=None, + weights_mapper=None, + ) logger.info( "Loaded LoRA model: id=%d, num_modules=%d, modules=%s", @@ -202,6 +178,7 @@ def hijack__load_adapter(self, lora_request: OmniTensorLoRARequest) -> tuple[LoR return lora_model, peft_helper + hijack__load_adapter._diffrl_tensor_lora_loader = True # type: ignore[attr-defined] setattr(DiffusionLoRAManager, "_load_adapter", hijack__load_adapter) @@ -317,6 +294,9 @@ def hijack_ar__load_adapter(self, lora_request, _orig=_orig_ar_load_adapter) -> model = self._adapter_manager.model hf_to_vllm_mapper = getattr(model, "hf_to_vllm_mapper", None) + if hf_to_vllm_mapper is not None: + hf_to_vllm_mapper = hf_to_vllm_mapper.get_unstacked_mapper() + lora_skip_prefixes = getattr(model, "lora_skip_prefixes", None) lora = self._lora_model_cls.from_lora_tensors( tensors=lora_request.lora_tensors or {}, peft_helper=peft_helper, @@ -325,6 +305,7 @@ def hijack_ar__load_adapter(self, lora_request, _orig=_orig_ar_load_adapter) -> dtype=self.lora_config.lora_dtype, model_vocab_size=self.vocab_size, weights_mapper=hf_to_vllm_mapper, + skip_prefixes=lora_skip_prefixes, ) return lora @@ -401,7 +382,6 @@ def _patched_from_layer( for _modname in ( "vllm.lora.lora_model", - "vllm.lora.models", "vllm.lora.model_manager", "vllm.lora.worker_manager", ): @@ -678,11 +658,6 @@ class VLLMOmniHijack: def hijack() -> None: wrap_mp_process_for_children() - # StageDiffusionProc never loads vllm_omni.general_plugins, so spawn children get the flush only via wrap_mp. - from unirl.rollout.engine.vllm_omni.plugin import register_capture_flush - - register_capture_flush() - patch_dit_lora_loader() patch_dit_hi3_lora_weights() patch_ar_lora_loader() diff --git a/unirl/rollout/engine/vllm_omni/pipeline_configs.py b/unirl/rollout/engine/vllm_omni/pipeline_configs.py new file mode 100644 index 000000000..13f95d582 --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/pipeline_configs.py @@ -0,0 +1,80 @@ +"""Out-of-tree vLLM-Omni pipeline topology registrations.""" + +from __future__ import annotations + +from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +UNIRL_RL_IMAGE_DIFFUSION = PipelineConfig( + model_type="unirl_rl_image_diffusion", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + final_output=True, + final_output_type="image", + ), + ), +) + +_HI3_MODEL_ARCH = "HunyuanImage3ForCausalMM" + + +def _hi3_text_pipeline(model_type: str, *, requires_multimodal_data: bool) -> PipelineConfig: + return PipelineConfig( + model_type=model_type, + model_arch=_HI3_MODEL_ARCH, + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="AR", + execution_type=StageExecutionType.LLM_AR, + final_output=True, + final_output_type="text", + owns_tokenizer=True, + requires_multimodal_data=requires_multimodal_data, + model_arch=_HI3_MODEL_ARCH, + engine_output_type="text", + ), + ), + ) + + +UNIRL_HI3_AR_TEXT = _hi3_text_pipeline( + "unirl_hi3_ar_text", + requires_multimodal_data=False, +) +UNIRL_HI3_AR_MULTIMODAL_TEXT = _hi3_text_pipeline( + "unirl_hi3_ar_multimodal_text", + requires_multimodal_data=True, +) + + +def register_unirl_pipeline_configs() -> None: + """Register UniRL's custom topologies once, rejecting key collisions.""" + for pipeline in ( + UNIRL_RL_IMAGE_DIFFUSION, + UNIRL_HI3_AR_TEXT, + UNIRL_HI3_AR_MULTIMODAL_TEXT, + ): + existing = OMNI_PIPELINES.get(pipeline.model_type) + if existing is None: + register_pipeline(pipeline) + elif existing != pipeline: + raise RuntimeError( + f"vLLM-Omni pipeline key {pipeline.model_type!r} is already " + f"registered to an incompatible pipeline: {existing!r}" + ) + + +__all__ = [ + "UNIRL_HI3_AR_MULTIMODAL_TEXT", + "UNIRL_HI3_AR_TEXT", + "UNIRL_RL_IMAGE_DIFFUSION", + "register_unirl_pipeline_configs", +] diff --git a/unirl/rollout/engine/vllm_omni/plugin.py b/unirl/rollout/engine/vllm_omni/plugin.py index 93b47b49a..d4539cafe 100644 --- a/unirl/rollout/engine/vllm_omni/plugin.py +++ b/unirl/rollout/engine/vllm_omni/plugin.py @@ -9,6 +9,14 @@ _PATCH_FLAG = "_unirl_capture_flush" +def register_unirl_runtime() -> None: + """Register UniRL topology and output hooks in this vllm-omni process.""" + from unirl.rollout.engine.vllm_omni.pipeline_configs import register_unirl_pipeline_configs + + register_unirl_pipeline_configs() + register_capture_flush() + + def register_capture_flush() -> None: """Flush pipeline captures into formatter metadata after postprocess sees raw media.""" from vllm_omni.diffusion import diffusion_engine, output_formatter @@ -33,3 +41,6 @@ def patched(*args, **kwargs): # diffusion_engine took a from-import, so it holds its own binding. diffusion_engine.format_diffusion_outputs = patched logger.info("unirl: diffusion formatter flushes capture metadata after postprocess") + + +__all__ = ["register_capture_flush", "register_unirl_runtime"] diff --git a/unirl/rollout/engine/vllm_omni/stage_configs/bagel_t2i_rl.yaml b/unirl/rollout/engine/vllm_omni/stage_configs/bagel_t2i_rl.yaml deleted file mode 100644 index 38f82e89f..000000000 --- a/unirl/rollout/engine/vllm_omni/stage_configs/bagel_t2i_rl.yaml +++ /dev/null @@ -1,103 +0,0 @@ -# Single-stage BAGEL-7B-MoT image-out config for RL rollout. -# -# BAGEL is the single-stage topology where the DiT worker owns a full -# Qwen2-MoT LLM + ViT + VAE + tokenizer and handles text2img end-to-end (no -# separate AR/Thinker prelude — think mode is off for the GRPO path). -# -# Shared by BOTH image-out modalities — ``bagel_t2i`` and ``bagel_it2i`` -# (editing). The custom worker pipeline owns the ViT and editing prefill; the -# source image travels on the request (``multi_modal_data["image"]``), not in this -# config, so the two modalities need no stage-level difference. -# -# Diff vs upstream's bagel_single_stage deploy yaml: -# - custom_pipeline_args installs our RLBagelPipeline subclass. Unlike SD3/Qwen -# (which swap a diffusers scheduler), RLBagelPipeline sets self.scheduler to -# the BAGEL-specific BagelFlowSDEScheduler that upstream generate_image -# already threads through its hand-rolled denoise loop — giving per-step SDE -# transitions + dense trajectory capture (latents, reconstructed sigmas, -# sparse log_probs) whose math matches the trainside FlowSDEStrategy. It also -# taps bagel.prepare_vae_latent to inject the driver-authored x_T. -# - worker_extension_cls installs DiTWeightSyncExtension so the worker exposes -# update_weights_from_ipc (bucketed CUDA-IPC) and the NCCL update_weights -# paths the trainer's LoRA sync needs. -# -# Contract notes: -# - The adapter translates canonical guidance_scale to BAGEL's -# extra_args.cfg_text_scale; guidance_scale is not forwarded as a top-level -# worker argument. It also always sends cfg_img_scale because upstream -# BAGEL defaults a missing cfg_text_scale to 4.0 (CFG ON); trainside cfg=1 -# must remain a single-forward path. -# - num_inference_steps is sent as trainside_steps + 1: BAGEL builds -# linspace(1, 0, num_timesteps) and loops num_timesteps-1 steps. BAGEL builds -# its sigma schedule internally (it ignores sampling_params.sigmas), with a -# hardwired timestep_shift=3.0 == the trainside shift, so the worker schedule -# equals the diffusion Part's pinned sigmas; response sigma-echo verification asserts it. -# - trust_remote_code: BAGEL ships custom configs (llm_config.json / vit_config.json). -# -# NB: no ``pipeline:`` key — the legacy stage_configs loader reads only -# ``stage_args``, and ``custom_pipeline_args.pipeline_class`` bypasses model-type -# detection entirely (the worker imports RLBagelPipeline by qualname, which loads -# the BAGEL checkpoint itself), exactly like the SD3 / Qwen-Image RL configs. -# -# Sized for BAGEL-7B-MoT on TP=1 (navit bs=1; one sample per generate call). -# Single-GPU rollout matches the trainer's full-forward math more closely than -# tensor-parallel rollout and avoids the systematic bf16 drift seen in replay -# log-prob parity checks. - -async_chunk: false - -stage_args: - - stage_id: 0 - stage_type: diffusion - runtime: - process: true - devices: "0" - engine_args: - model_stage: dit - max_num_seqs: 1 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - distributed_executor_backend: "mp" - vae_use_slicing: false - vae_use_tiling: false - cache_backend: null - cache_config: null - enable_cache_dit_summary: false - # Required for vLLM to wrap DiT Linear modules in LoRA-aware variants. - # Without ``enable_lora``, ``add_lora`` still registers the adapter (so - # checksum probes pass) but no Linear consults it and the forward silently - # runs base weights. ``max_lora_rank`` must be >= the recipe's LoRA rank - # (64 in the BAGEL LoRA recipes); vLLM's default of 16 rejects larger ranks. - enable_lora: true - max_lora_rank: 64 - parallel_config: - pipeline_parallel_size: 1 - data_parallel_size: 1 - tensor_parallel_size: 1 - enable_expert_parallel: false - sequence_parallel_size: 1 - ulysses_degree: 1 - ring_degree: 1 - cfg_parallel_size: 1 - vae_patch_parallel_size: 1 - use_hsdp: false - hsdp_shard_size: -1 - hsdp_replicate_size: 1 - # unirl-specific: install RLBagelPipeline so SDE-mode rollouts get the - # trajectory-capturing scheduler + driver-noise injection. - custom_pipeline_args: - pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.bagel.pipeline.RLBagelPipeline - # unirl-specific: install the weight-sync extension so the DiT worker - # exposes update_weights_from_ipc (bucketed CUDA-IPC) and later - # init_weight_transfer_engine / update_weights (NCCL). Coexists with the - # custom pipeline above; vllm-omni composes them via multiple inheritance. - worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension - final_output: true - final_output_type: image - -runtime: - enabled: true - defaults: - window_size: -1 - max_inflight: 1 diff --git a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_ar_recaption_rl.yaml b/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_ar_recaption_rl.yaml deleted file mode 100644 index 1bac171e7..000000000 --- a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_ar_recaption_rl.yaml +++ /dev/null @@ -1,77 +0,0 @@ -# AR-ONLY think_recaption stage config (two-engine v2 trainer). -# -# A single, TERMINAL AR (LLM) stage that, in is_comprehension:false mode, -# generates ......... -# and stops at EOS. The full AR text is the recaption the DiT engine consumes as -# cot_text (extra["ar_generated_text"]); per-token logp (logprobs=1) feeds -# ar_capture.extract_ar_segment for ARGRPO replay. -# -# Differs from upstream i2t/t2t (which are is_comprehension:true comprehension -# answers, NOT think/recaption) — see hi3_v2_IMPL_spec.md plan-correction #1. -# -# Terminal AR-stage specifics: -# - final_output_type stays text (terminal; no DiT downstream). -# - omni_kv_config.need_send_cache DROPPED (no DiT stage to receive KV). -# -# Devices 0-3 (AR engine half of the 8-GPU colocate pool); the DiT engine -# (hunyuan_image3_dit_recaption_rl.yaml) takes 4-7. - -stage_args: - - stage_id: 0 - stage_type: llm - runtime: - process: true - devices: "0,1,2,3" - engine_args: - model_stage: AR - max_num_seqs: 1 - model_arch: HunyuanImage3ForCausalMM - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - # Side-effect import installs the Base-ckpt tokenizer fix AND composes - # the bucketed CUDA-IPC weight-sync mixin (LoRA sync target). See - # t2i_rl.yaml for the full rationale. - worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.ar_extension.HI3ARWeightSyncExtension - enable_lora: true - # Routes MoE LoRA off vLLM 0.27's fused one-shot kernel, which faults here. - fully_sharded_loras: true - # Colocate-memory trim: the AR stage only needs prompt + max_new_tokens - # (~4096), not the model's full 22800 ctx. Capping max_model_len shrinks - # the KV-cache reservation (~21 GiB → a few GiB) and gpu_memory_util 0.9→0.7 - # lowers the resident/slept footprint — both free room on the shared train - # cards (engine sleep keeps a few GiB on GPU 0-3, which the FSDP backward - # competes with). max_num_batched_tokens 32768→8192 also clears the - # "exceeds max_num_seqs*max_model_len" boot warning. - max_model_len: 8192 - gpu_memory_utilization: 0.7 - enforce_eager: true - trust_remote_code: true - engine_output_type: latent - enable_prefix_caching: false - max_num_batched_tokens: 8192 - tensor_parallel_size: 4 - pipeline_parallel_size: 1 - distributed_executor_backend: "mp" - hf_overrides: - rope_parameters: - mrope_section: [0, 32, 32] - rope_type: default - # is_comprehension:false → AR generates think/recaption (not a plain - # comprehension answer). No omni_kv_config: this stage is terminal. - is_comprehension: false - final_output: true - final_output_type: text - # No default seed on purpose: a fixed seed + temperature would make the N - # temperature-sampled recaptions of a prompt identical (→ AR GRPO advantage - # 0). Omitting it lets vLLM draw a fresh seed per request so the N recaptions - # actually differ. - default_sampling_params: - temperature: 0.7 - top_p: 0.95 - top_k: 1024 - max_tokens: 4096 - detokenize: True - repetition_penalty: 1.1 - -runtime: - enabled: true diff --git a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_dit_recaption_rl.yaml b/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_dit_recaption_rl.yaml deleted file mode 100644 index 0e0312905..000000000 --- a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_dit_recaption_rl.yaml +++ /dev/null @@ -1,68 +0,0 @@ -# DiT-ONLY stage config (two-engine v2 trainer) — critical gotcha. -# -# A single standalone diffusion stage that consumes an EXTERNAL recaption as -# cot_text (injected per-prompt via extra["ar_generated_text"] by request.py's -# dit_recaption branch) and renders M images at distinct noise. No AR stage, -# no KV transfer — the DiT pipeline self-prefills the text context from -# prompt + cot_text (RLHunyuanImage3Pipeline.forward → upstream forward → -# prepare_model_inputs(mode="gen_image")). -# -# Base = upstream hunyuan_image3_t2i.yaml, the ONLY HI3 DiT config verified -# standalone (single diffusion stage, 4x H20): enable_expert_parallel:true, -# need_recv_cache:true (no upstream edge feeds it, so it self-prefills — same -# as the verified standalone). Plus our RL bits from t2i_rl.yaml's DiT stage: -# - custom_pipeline_args → RLHunyuanImage3Pipeline (SDE scheduler swap + -# dense trajectory capture + fused_mm_capture for DiffusionGRPO replay). -# - worker_extension_cls → DiTWeightSyncExtension (LoRA receive target). -# - vae_use_slicing/tiling:false, cache_*:null (trajectory-capture parity). -# -# Devices 4-7 (DiT engine half of the 8-GPU colocate pool); AR engine -# (hunyuan_image3_ar_recaption_rl.yaml) takes 0-3. - -stage_args: - - stage_id: 0 - stage_type: diffusion - runtime: - devices: "4,5,6,7" - engine_args: - max_num_seqs: 1 - model_stage: dit - # Colocate-memory trim: lower the DiT engine's resident/slept footprint on - # GPU 4-7 so the FSDP backward has more room (mirrors the AR stage trim). - gpu_memory_utilization: 0.7 - enforce_eager: true - trust_remote_code: true - distributed_executor_backend: "mp" - vae_use_slicing: false - vae_use_tiling: false - cache_backend: null - cache_config: null - enable_cache_dit_summary: false - # No upstream sender; the recv side is set up but no transfer arrives, - # so the pipeline prefills its own text context (matches the verified - # standalone upstream t2i.yaml, which is also need_recv_cache:true with - # no edge feeding it). - omni_kv_config: - need_recv_cache: true - parallel_config: - tensor_parallel_size: 4 - enable_expert_parallel: true - # RL pipeline subclass: SDE scheduler + per-step trajectory + fused_mm - # capture (DiffusionGRPO replay needs conditions["fused"]). - custom_pipeline_args: - pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.hi3.pipeline.RLHunyuanImage3Pipeline - # Weight-sync extension: DiT worker exposes update_weights_from_ipc / - # the LoRA receive path so set_lora_from_tensors lands the adapter. - worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension - final_output: true - final_output_type: image - is_comprehension: false - # No default seed on purpose. Per-image distinct seeds are set by the engine - # (one generate() per prompt, seed = seed_from_sample_id) — NOT here and NOT - # via the request's sampling params. A fixed default (e.g. seed:42) here, OR - # relying on vllm-omni's random-seed fallback, makes every per-recaption - # request share ONE seed → byte-identical images → diffusion advantage 0. - # See engine.seed_from_sample_id / request._to_omni_dit_recaption. - -runtime: - enabled: true diff --git a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_i2t.yaml b/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_i2t.yaml deleted file mode 100644 index ad0be96ef..000000000 --- a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_i2t.yaml +++ /dev/null @@ -1,40 +0,0 @@ -# Vendored at its last upstream revision, since vllm-omni deleted stage_configs/ in #3172. -# HunyuanImage-3.0 image-to-text: one AR stage reads image + prompt and emits text. - -stage_args: - - stage_id: 0 - stage_type: llm - runtime: - process: true - devices: "0,1,2,3" - requires_multimodal_data: true - engine_args: - model_stage: AR - max_num_seqs: 1 - model_arch: HunyuanImage3ForCausalMM - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.95 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - max_num_batched_tokens: 32768 - tensor_parallel_size: 4 - pipeline_parallel_size: 1 - hf_overrides: - rope_parameters: - mrope_section: [0, 32, 32] - rope_type: default - is_comprehension: true - final_output: true - final_output_type: text - default_sampling_params: - temperature: 0.0 - top_p: 0.95 - top_k: 1024 - max_tokens: 2048 - stop_token_ids: [127957, 128024, 128026] # <|endoftext|>, , - detokenize: True - -runtime: - enabled: true diff --git a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_it2i_rl.yaml b/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_it2i_rl.yaml deleted file mode 100644 index 0124650df..000000000 --- a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_it2i_rl.yaml +++ /dev/null @@ -1,87 +0,0 @@ -# RL-aware copy of upstream hunyuan_image3_it2i.yaml. -# Diff vs upstream: -# 1. Stage 0 final_output: true (was false). Surfaces the AR text + per-token -# logp our ar_capture.py reads to fill the AR generation Part's segment. The KV transfer -# to Stage 1 is independent of this flag, so the AR→DiT bridge is -# unaffected. -# 2. Stage 1 engine_args.custom_pipeline_args injected to install our -# RLHunyuanImage3Pipeline subclass (scheduler swap + trajectory drain -# for SDE-mode rollouts). - -stage_args: - - stage_id: 0 - stage_type: llm - runtime: - process: true - devices: "0,1,2,3" - requires_multimodal_data: true - engine_args: - max_num_seqs: 1 - model_stage: AR - model_arch: HunyuanImage3ForCausalMM - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - # See t2i YAML for rationale; same tokenizer patch + weight-sync mixin. - worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.ar_extension.HI3ARWeightSyncExtension - # Enable LoRA on the AR stage; mirrors the t2i YAML. - enable_lora: true - gpu_memory_utilization: 0.95 - enforce_eager: true - trust_remote_code: true - engine_output_type: latent - enable_prefix_caching: false - max_num_batched_tokens: 32768 - tensor_parallel_size: 4 - pipeline_parallel_size: 1 - hf_overrides: - rope_parameters: - mrope_section: [0, 32, 32] - rope_type: default - is_comprehension: false - # unirl-specific: was `false` upstream — flipped so AR tokens + - # per-token logp surface to the orchestrator's caller for RL. - final_output: true - final_output_type: text - default_sampling_params: - temperature: 0.6 - top_p: 0.95 - top_k: 1024 - max_tokens: 4096 - stop_token_ids: [127957] - detokenize: true - - - stage_id: 1 - stage_type: diffusion - runtime: - process: true - devices: "4,5,6,7" - requires_multimodal_data: true - engine_args: - max_num_seqs: 1 - model_stage: dit - model_arch: HunyuanImage3ForCausalMM - enforce_eager: true - trust_remote_code: true - distributed_executor_backend: "mp" - parallel_config: - tensor_parallel_size: 4 - enable_expert_parallel: true - # unirl-specific: install RLHunyuanImage3Pipeline. - custom_pipeline_args: - pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.hi3.pipeline.RLHunyuanImage3Pipeline - # unirl-specific: install the weight-sync extension on the DiT - # worker (mirrors t2i YAML — see that file for rationale). - worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension - engine_input_source: [0] - custom_process_input_func: vllm_omni.model_executor.stage_input_processors.hunyuan_image3.ar2diffusion - final_output: true - final_output_type: image - default_sampling_params: - num_inference_steps: 50 - guidance_scale: 2.5 - -runtime: - enabled: true - edges: - - from: 0 - to: 1 diff --git a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_t2i_rl.yaml b/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_t2i_rl.yaml deleted file mode 100644 index 44a5a3642..000000000 --- a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_t2i_rl.yaml +++ /dev/null @@ -1,126 +0,0 @@ -# RL-aware copy of upstream hunyuan_image3_moe.yaml. -# Diff vs upstream: -# - Stage 1 engine_args.custom_pipeline_args injected to install our -# RLHunyuanImage3Pipeline subclass (scheduler swap + trajectory drain -# for SDE-mode rollouts). -# Stage 0 (AR) is unchanged: upstream already ships final_output: true, -# which surfaces the AR text + per-token logp our ar_capture.py reads. -# -# Verified on 8x L40S-48G GPU (4 for AR + 4 for DiT) per upstream docs. - -stage_args: - - stage_id: 0 - stage_type: llm - runtime: - process: true - devices: "0,1,2,3" - engine_args: - model_stage: AR - max_num_seqs: 1 - model_arch: HunyuanImage3ForCausalMM - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - # Side-effect import: ``HI3ARWeightSyncExtension`` subclasses - # ``HI3ARWorkerExtension`` (in ``compat/tokenizer``), so resolving - # the qualname imports that module and triggers ``install()`` — - # patching ``convert_tokens_to_ids`` -> 0-on-None for the Base - # HunyuanImage-3 ckpt (its tokenizer is missing ````; - # vllm-omni's HI3 model does ratio_36 + 1 unconditionally and would - # otherwise TypeError). No-op on Instruct since the lookup returns - # a real id there. Composing the weight-sync mixin on top adds - # ``update_weights_from_ipc`` (bucketed CUDA-IPC) without losing - # the tokenizer fix. - worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.ar_extension.HI3ARWeightSyncExtension - # unirl-specific: enable LoRA on the AR stage so - # ``set_lora_from_tensors`` can install in-memory adapters via - # ``add_lora`` (the LoRA tensor-bag transport). Defaults from - # ``vllm.config.lora.LoRAConfig`` (max_loras=1, max_lora_rank=16) - # are sufficient for our single-adapter trainer-side LoRA flow. - enable_lora: true - gpu_memory_utilization: 0.9 - enforce_eager: true - trust_remote_code: true - engine_output_type: latent - enable_prefix_caching: false - max_num_batched_tokens: 32768 - tensor_parallel_size: 4 - pipeline_parallel_size: 1 - # When running inside a Ray actor (e.g. RolloutActor), vllm's auto- - # detection picks the ray executor backend, but the StageEngineCoreProc - # subprocess loses the placement-group handle (KeyError: 'bundles'). - # Force mp executor so each stage spawns workers via multiprocessing, - # matching the legacy in-process smoke and the DiT stage below. - distributed_executor_backend: "mp" - hf_overrides: - rope_parameters: - mrope_section: [0, 32, 32] - rope_type: default - omni_kv_config: - need_send_cache: true - kv_transfer_criteria: - type: prefill_finished - is_comprehension: true - final_output: true - final_output_type: text - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 2048 - seed: 42 - detokenize: True - repetition_penalty: 1.1 - - - stage_id: 1 - stage_type: diffusion - runtime: - process: true - devices: "4,5,6,7" - engine_args: - model_stage: diffusion - max_num_seqs: 1 - enforce_eager: true - distributed_executor_backend: "mp" - vae_use_slicing: false - vae_use_tiling: false - cache_backend: null - cache_config: null - enable_cache_dit_summary: false - omni_kv_config: - need_recv_cache: true - parallel_config: - pipeline_parallel_size: 1 - data_parallel_size: 1 - tensor_parallel_size: 4 - enable_expert_parallel: false - sequence_parallel_size: 1 - ulysses_degree: 1 - ring_degree: 1 - cfg_parallel_size: 1 - vae_patch_parallel_size: 1 - use_hsdp: false - hsdp_shard_size: -1 - hsdp_replicate_size: 1 - # unirl-specific: install RLHunyuanImage3Pipeline so SDE-mode - # rollouts get a swappable scheduler + per-step trajectory capture. - custom_pipeline_args: - pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.hi3.pipeline.RLHunyuanImage3Pipeline - # unirl-specific: install the weight-sync extension so the DiT - # worker exposes ``update_weights_from_ipc`` (bucketed CUDA-IPC) and - # later ``init_weight_transfer_engine`` / ``update_weights`` (NCCL). - # Must coexist with the custom pipeline above; vllm-omni's worker - # composes these via multiple inheritance at instantiation time. - worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension - engine_input_source: [0] - final_output: true - final_output_type: image - -runtime: - enabled: true - defaults: - window_size: -1 - max_inflight: 1 - edges: - - from: 0 - to: 1 - window_size: -1 diff --git a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_t2t.yaml b/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_t2t.yaml deleted file mode 100644 index 9edae5c24..000000000 --- a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_image3_t2t.yaml +++ /dev/null @@ -1,41 +0,0 @@ -# Vendored at its last upstream revision, since vllm-omni deleted stage_configs/ in #3172. -# HunyuanImage-3.0 text-to-text: one AR stage reads a prompt and emits text. -# Sampling params match the official generation_config.json. - -stage_args: - - stage_id: 0 - stage_type: llm - runtime: - process: true - devices: "0,1,2,3" - requires_multimodal_data: false - engine_args: - model_stage: AR - max_num_seqs: 1 - model_arch: HunyuanImage3ForCausalMM - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.95 - enforce_eager: true - trust_remote_code: true - enable_prefix_caching: false - max_num_batched_tokens: 32768 - tensor_parallel_size: 4 - pipeline_parallel_size: 1 - hf_overrides: - rope_parameters: - mrope_section: [0, 32, 32] - rope_type: default - is_comprehension: true - final_output: true - final_output_type: text - default_sampling_params: - temperature: 0.0 - top_p: 0.95 - top_k: 1024 - max_tokens: 2048 - stop_token_ids: [127957, 128024, 128026] # <|endoftext|>, , - detokenize: True - -runtime: - enabled: true diff --git a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_video15_t2v_rl.yaml b/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_video15_t2v_rl.yaml deleted file mode 100644 index a215511bb..000000000 --- a/unirl/rollout/engine/vllm_omni/stage_configs/hunyuan_video15_t2v_rl.yaml +++ /dev/null @@ -1,79 +0,0 @@ -# Single-stage HunyuanVideo-1.5 text-to-video config for RL rollout. -# -# Diff vs SD3.5 stage config: -# - Same single-stage diffusion architecture (no AR prelude — HV1.5 is pure -# generative text-to-video, not multi-modal like HunyuanImage-3). -# - custom_pipeline_args installs our RLHunyuanVideo15Pipeline subclass -# so SDE-mode rollouts get trajectory capture (latents, sigmas, log_probs) -# and dual-text-encoder condition extraction (text_mllm + text_glyph). -# - worker_extension_cls installs DiTWeightSyncExtension (same as SD3, -# reused for HV1.5 DiT backbone). -# - Dual-text-encoder condition fields (Qwen2.5-VL MLLM + ByT5 glyph) -# are captured by the pipeline subclass and routed via the output envelope -# for the response handler to unpack into HunyuanVideo15Conditions. -# -# Sized for HunyuanVideo-1.5 on TP=1. Single-GPU rollout matches the trainer's -# full-forward math more closely than tensor-parallel rollout. - -stage_args: - - stage_id: 0 - stage_type: diffusion - runtime: - process: true - devices: "0" - engine_args: - model_stage: diffusion - # HV1.5 is non-batching upstream, so the engine requires this. - max_num_seqs: 1 - enforce_eager: true - distributed_executor_backend: "mp" - vae_use_slicing: false - vae_use_tiling: false - cache_backend: null - cache_config: null - enable_cache_dit_summary: false - # Match the trainer's PyTorch SDPA forward path. - # 0.27 reads ``diffusion_attention_backend`` (not the AR field ``attention_backend``). - diffusion_attention_backend: TORCH_SDPA - # Required for vLLM to wrap DiT Linear modules in LoRA-aware variants. - # Without ``enable_lora``, ``add_lora`` still registers the adapter in - # ``DiffusionLoRAManager._registered_adapters`` (so checksum probes - # pass), but no Linear consults it and the forward silently runs base - # weights. ``max_lora_rank`` must be >= the recipe's LoRAPolicy rank; - # the vllm.config.lora default of 16 silently rejects larger adapters - # at add time. - enable_lora: true - max_lora_rank: 64 - parallel_config: - pipeline_parallel_size: 1 - data_parallel_size: 1 - tensor_parallel_size: 1 - enable_expert_parallel: false - sequence_parallel_size: 1 - ulysses_degree: 1 - ring_degree: 1 - cfg_parallel_size: 1 - vae_patch_parallel_size: 1 - use_hsdp: false - hsdp_shard_size: -1 - hsdp_replicate_size: 1 - # unirl-specific: install RLHunyuanVideo15Pipeline so SDE-mode - # rollouts get a swappable scheduler, per-step trajectory capture - # (latents, sigmas, log_probs), and dual-text-encoder condition - # extraction (text_mllm + text_glyph from the dual-encoder architecture). - custom_pipeline_args: - pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.hv15.pipeline.RLHunyuanVideo15Pipeline - # unirl-specific: install the weight-sync extension so the DiT - # worker exposes update_weights_from_ipc (bucketed CUDA-IPC) and - # later init_weight_transfer_engine / update_weights (NCCL). - # Must coexist with the custom pipeline above; vllm-omni's worker - # composes these via multiple inheritance at instantiation time. - worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension - final_output: true - final_output_type: video - -runtime: - enabled: true - defaults: - window_size: -1 - max_inflight: 1 diff --git a/unirl/rollout/engine/vllm_omni/stage_configs/qwen3_omni_thinker_only_rl_1x4.yaml b/unirl/rollout/engine/vllm_omni/stage_configs/qwen3_omni_thinker_only_rl_1x4.yaml deleted file mode 100644 index a728f248b..000000000 --- a/unirl/rollout/engine/vllm_omni/stage_configs/qwen3_omni_thinker_only_rl_1x4.yaml +++ /dev/null @@ -1,51 +0,0 @@ -# Qwen3-Omni Thinker-only stage, TP=4. -stage_args: - - stage_id: 0 - runtime: - devices: "0,1,2,3" - engine_args: - model_stage: thinker - model_arch: Qwen3OmniMoeThinkerForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - engine_output_type: text - hf_config_name: thinker_config - - tensor_parallel_size: 4 - distributed_executor_backend: "mp" - - gpu_memory_utilization: 0.35 - # Standalone image, audio, and video prompts are safe to batch. - # Audio-in-video uses the dedicated serialized stage config instead. - max_num_seqs: 64 - max_num_batched_tokens: 32768 - enable_chunked_prefill: true - - # Audio fusion is enabled per request; setting use_audio_in_video globally - # breaks vLLM's video-only dummy profiling during startup. - limit_mm_per_prompt: - audio: 1 - video: 1 - image: 1 - - mm_processor_cache_gb: 0 - - dtype: bfloat16 - load_format: safetensors - trust_remote_code: true - enable_prefix_caching: false - enforce_eager: false - - enable_lora: true - max_lora_rank: 64 - max_loras: 1 - - # Add the weight-sync receivers to the AR worker. - worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.qwen3_omni_ar_extension.Qwen3OmniARWeightSyncExtension - - enable_sleep_mode: true - logprobs_mode: processed_logprobs - disable_log_stats: true - final_output: true - final_output_type: text - is_comprehension: true diff --git a/unirl/rollout/engine/vllm_omni/stage_configs/qwen3_omni_thinker_only_rl_audio_video_1x4.yaml b/unirl/rollout/engine/vllm_omni/stage_configs/qwen3_omni_thinker_only_rl_audio_video_1x4.yaml deleted file mode 100644 index b79dacfac..000000000 --- a/unirl/rollout/engine/vllm_omni/stage_configs/qwen3_omni_thinker_only_rl_audio_video_1x4.yaml +++ /dev/null @@ -1,52 +0,0 @@ -# Qwen3-Omni Thinker-only audio-in-video stage, TP=4. -stage_args: - - stage_id: 0 - runtime: - devices: "0,1,2,3" - engine_args: - model_stage: thinker - model_arch: Qwen3OmniMoeThinkerForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - engine_output_type: text - hf_config_name: thinker_config - - tensor_parallel_size: 4 - distributed_executor_backend: "mp" - - gpu_memory_utilization: 0.35 - # Audio-in-video decoder numerics become batch-sensitive when several - # long multimodal requests share one packed forward. Serialize requests - # to keep rollout/replay K3 near the bf16 numerical floor. - max_num_seqs: 1 - max_num_batched_tokens: 32768 - enable_chunked_prefill: true - - # Audio fusion is enabled per request; setting use_audio_in_video globally - # breaks vLLM's video-only dummy profiling during startup. - limit_mm_per_prompt: - audio: 1 - video: 1 - image: 0 - - mm_processor_cache_gb: 0 - - dtype: bfloat16 - load_format: safetensors - trust_remote_code: true - enable_prefix_caching: false - enforce_eager: false - - enable_lora: true - max_lora_rank: 64 - max_loras: 1 - - # Add the weight-sync receivers to the AR worker. - worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.qwen3_omni_ar_extension.Qwen3OmniARWeightSyncExtension - - enable_sleep_mode: true - logprobs_mode: processed_logprobs - disable_log_stats: true - final_output: true - final_output_type: text - is_comprehension: true diff --git a/unirl/rollout/engine/vllm_omni/stage_configs/qwen_image_t2i_rl.yaml b/unirl/rollout/engine/vllm_omni/stage_configs/qwen_image_t2i_rl.yaml deleted file mode 100644 index 5468c503f..000000000 --- a/unirl/rollout/engine/vllm_omni/stage_configs/qwen_image_t2i_rl.yaml +++ /dev/null @@ -1,85 +0,0 @@ -# Single-stage Qwen-Image text-to-image config for RL rollout. -# -# Diff vs the (non-existent) upstream Qwen-Image stage config (upstream boots -# it through the default single-diffusion-stage factory): -# - custom_pipeline_args installs our RLQwenImagePipeline subclass so -# SDE-mode rollouts get a swappable scheduler + per-step trajectory -# capture (latents, sigmas, log_probs). The subclass also owns the -# packed-latent boundary: x_T injection packs the driver's spatial -# [B, C, H, W] noise; trajectory harvest unpacks back from the -# transformer's [B, S, C*4] patch layout. -# - worker_extension_cls installs DiTWeightSyncExtension so the worker -# exposes update_weights_from_ipc (bucketed CUDA-IPC) and the NCCL -# update_weights paths the trainer needs. -# -# Contract notes: -# - CFG is off by default (recipes run guidance_scale 1.0 = true_cfg 1.0); -# the driver-side adapter omits negative_prompt entirely when CFG is off -# so upstream's "" -> has_neg_prompt + true_cfg 4.0 default never fires. -# - height/width must be multiples of 16 (vae_scale_factor * packing); -# otherwise upstream normalizes them and the driver-pinned dynamic-shift -# sigma schedule desyncs — the response layer's sigma-echo verify then -# fails loudly. -# -# Sized for Qwen-Image (20B MMDiT + co-resident Qwen2.5-VL text encoder) on -# TP=1. Single-GPU rollout matches the trainer's full-forward math more -# closely than tensor-parallel rollout and avoids the systematic bf16 drift -# seen in replay log-prob parity checks. - -stage_args: - - stage_id: 0 - stage_type: diffusion - runtime: - process: true - devices: "0" - engine_args: - model_stage: diffusion - # The pipeline sets supports_request_batch=False, which requires this. - max_num_seqs: 1 - enforce_eager: true - distributed_executor_backend: "mp" - vae_use_slicing: false - vae_use_tiling: false - cache_backend: null - cache_config: null - enable_cache_dit_summary: false - # Required for vLLM to wrap DiT Linear modules in LoRA-aware variants. - # Without ``enable_lora``, ``add_lora`` still registers the adapter in - # ``DiffusionLoRAManager._registered_adapters`` (so checksum probes - # pass), but no Linear consults it and the forward silently runs base - # weights. ``max_lora_rank`` must be >= the recipe's LoRAPolicy rank - # (64 in qwen_image_dancegrpo); the vllm.config.lora default of 16 - # silently rejects larger adapters at add time. - enable_lora: true - max_lora_rank: 64 - parallel_config: - pipeline_parallel_size: 1 - data_parallel_size: 1 - tensor_parallel_size: 1 - enable_expert_parallel: false - sequence_parallel_size: 1 - ulysses_degree: 1 - ring_degree: 1 - cfg_parallel_size: 1 - vae_patch_parallel_size: 1 - use_hsdp: false - hsdp_shard_size: -1 - hsdp_replicate_size: 1 - # unirl-specific: install RLQwenImagePipeline so SDE-mode rollouts get - # a swappable scheduler + per-step trajectory capture. - custom_pipeline_args: - pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.qwen_image.pipeline.RLQwenImagePipeline - # unirl-specific: install the weight-sync extension so the DiT - # worker exposes update_weights_from_ipc (bucketed CUDA-IPC) and - # later init_weight_transfer_engine / update_weights (NCCL). - # Must coexist with the custom pipeline above; vllm-omni's worker - # composes these via multiple inheritance at instantiation time. - worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension - final_output: true - final_output_type: image - -runtime: - enabled: true - defaults: - window_size: -1 - max_inflight: 1 diff --git a/unirl/rollout/engine/vllm_omni/stage_configs/sd35_t2i_rl.yaml b/unirl/rollout/engine/vllm_omni/stage_configs/sd35_t2i_rl.yaml deleted file mode 100644 index 2c932a1b3..000000000 --- a/unirl/rollout/engine/vllm_omni/stage_configs/sd35_t2i_rl.yaml +++ /dev/null @@ -1,73 +0,0 @@ -# Single-stage SD3.5-medium text-to-image config for RL rollout. -# -# Diff vs the (non-existent) upstream SD3 stage config: -# - Single diffusion stage (no AR prelude — SD3.5 has no chain-of-thought -# stage like HunyuanImage-3). -# - custom_pipeline_args installs our RLStableDiffusion3Pipeline subclass -# so SDE-mode rollouts get a swappable scheduler + per-step trajectory -# capture (latents, sigmas, log_probs). -# - worker_extension_cls installs DiTWeightSyncExtension so the worker -# exposes update_weights_from_ipc (bucketed CUDA-IPC) and the NCCL -# update_weights paths the trainer needs. -# -# Sized for SD3.5-medium on TP=1. Single-GPU rollout matches the trainer's -# full-forward math more closely than tensor-parallel rollout and avoids the -# systematic bf16 drift seen in replay log-prob parity checks. - -stage_args: - - stage_id: 0 - stage_type: diffusion - runtime: - process: true - devices: "0" - engine_args: - model_stage: diffusion - # The pipeline sets supports_request_batch=False, which requires this. - max_num_seqs: 1 - enforce_eager: true - distributed_executor_backend: "mp" - vae_use_slicing: false - vae_use_tiling: false - cache_backend: null - cache_config: null - enable_cache_dit_summary: false - # Required for vLLM to wrap DiT Linear modules in LoRA-aware variants. - # Without ``enable_lora``, ``add_lora`` still registers the adapter in - # ``DiffusionLoRAManager._registered_adapters`` (so checksum probes - # pass), but no Linear consults it and the forward silently runs base - # weights. ``max_lora_rank`` must be >= the recipe's LoRAPolicy rank - # (32 in flowgrpo_fast_sd3); the vllm.config.lora default of 16 - # silently rejects larger adapters at add time. - enable_lora: true - max_lora_rank: 32 - parallel_config: - pipeline_parallel_size: 1 - data_parallel_size: 1 - tensor_parallel_size: 1 - enable_expert_parallel: false - sequence_parallel_size: 1 - ulysses_degree: 1 - ring_degree: 1 - cfg_parallel_size: 1 - vae_patch_parallel_size: 1 - use_hsdp: false - hsdp_shard_size: -1 - hsdp_replicate_size: 1 - # unirl-specific: install RLStableDiffusion3Pipeline so SDE-mode - # rollouts get a swappable scheduler + per-step trajectory capture. - custom_pipeline_args: - pipeline_class: unirl.rollout.engine.vllm_omni.pipelines.sd3.pipeline.RLStableDiffusion3Pipeline - # unirl-specific: install the weight-sync extension so the DiT - # worker exposes update_weights_from_ipc (bucketed CUDA-IPC) and - # later init_weight_transfer_engine / update_weights (NCCL). - # Must coexist with the custom pipeline above; vllm-omni's worker - # composes these via multiple inheritance at instantiation time. - worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.dit_extension.DiTWeightSyncExtension - final_output: true - final_output_type: image - -runtime: - enabled: true - defaults: - window_size: -1 - max_inflight: 1 diff --git a/unirl/rollout/engine/vllm_omni/utils/tracks.py b/unirl/rollout/engine/vllm_omni/utils/tracks.py index 7ef35b4c9..19e6de3bb 100644 --- a/unirl/rollout/engine/vllm_omni/utils/tracks.py +++ b/unirl/rollout/engine/vllm_omni/utils/tracks.py @@ -189,11 +189,9 @@ def decoded_text_from_ar(per_request: Sequence[Sequence[Any]]) -> Texts: ar = pick_stage_output(outputs, final_output_type="text", stage_id=0) text_str = "" if ar is not None: - ro = getattr(ar, "request_output", None) - if ro is not None: - completions = getattr(ro, "outputs", None) or [] - if completions: - text_str = getattr(completions[0], "text", "") or "" + completions = getattr(ar, "outputs", None) or [] + if completions: + text_str = getattr(completions[0], "text", "") or "" texts.append(Text(text=text_str)) return Texts.from_list(texts) @@ -229,10 +227,7 @@ def _flatten_logprobs(logprobs: Any, fallback_len: int) -> Optional[torch.Tensor def _extract_completion(out: Any) -> Tuple[List[int], Optional[torch.Tensor]]: """Pull ``(token_ids, per_token_logp)`` out of a Stage-0 result.""" - request_output = getattr(out, "request_output", None) - if request_output is None: - return [], None - completions = getattr(request_output, "outputs", None) or [] + completions = getattr(out, "outputs", None) or [] if not completions: return [], None completion = completions[0] From 5378e0f6b6fae5590fceaa03813d90e081508e62 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Mon, 14 Sep 2026 23:59:22 +0800 Subject: [PATCH 04/23] test(rollout): cover vllm-omni 0.28 contracts --- .github/workflows/vllm-omni-contracts.yml | 24 +++ .../vllm_omni/test_compat_moe_workspace.py | 135 +++++++++++++++++ .../vllm_omni/test_dependency_contracts.py | 22 +++ .../rollout/vllm_omni/test_deploy_configs.py | 100 +++++++++++++ .../test_stable_installed_contracts.py | 92 ++++++++++++ .../test_stable_runtime_contracts.py | 140 ++++++++++++++++++ 6 files changed, 513 insertions(+) create mode 100644 .github/workflows/vllm-omni-contracts.yml create mode 100644 tests/rollout/vllm_omni/test_compat_moe_workspace.py create mode 100644 tests/rollout/vllm_omni/test_dependency_contracts.py create mode 100644 tests/rollout/vllm_omni/test_deploy_configs.py create mode 100644 tests/rollout/vllm_omni/test_stable_installed_contracts.py create mode 100644 tests/rollout/vllm_omni/test_stable_runtime_contracts.py diff --git a/.github/workflows/vllm-omni-contracts.yml b/.github/workflows/vllm-omni-contracts.yml new file mode 100644 index 000000000..f893da775 --- /dev/null +++ b/.github/workflows/vllm-omni-contracts.yml @@ -0,0 +1,24 @@ +name: vLLM-Omni Contracts + +on: + pull_request: + branches: [main] + +jobs: + contracts: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install focused test dependencies + run: | + python -m pip install pytest pyyaml pillow psutil hydra-core omegaconf + python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + + - name: Run vLLM-Omni contracts + run: python -m pytest -q tests/rollout/vllm_omni diff --git a/tests/rollout/vllm_omni/test_compat_moe_workspace.py b/tests/rollout/vllm_omni/test_compat_moe_workspace.py new file mode 100644 index 000000000..8851c1090 --- /dev/null +++ b/tests/rollout/vllm_omni/test_compat_moe_workspace.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +import sys +from contextlib import contextmanager +from contextvars import ContextVar +from types import ModuleType + +import pytest + +from unirl.rollout.engine.vllm_omni.patches.compat_moe_workspace import patch_moe_workspace_pool + + +class _Allocator: + def __init__(self) -> None: + self.current_tag = None + self.allocator_and_pools: dict[str, tuple[object, object]] = {} + self.entered_pool_keys: list[str] = [] + + @contextmanager + def use_memory_pool(self, tag: str): + self.entered_pool_keys.append(tag) + self.allocator_and_pools[tag] = (object(), object()) + yield + + +def _install_fake_vllm( + monkeypatch: pytest.MonkeyPatch, + workspace_manager: type, + allocator: _Allocator, + ubatch_id: ContextVar[int], + lane: ContextVar[int], +) -> None: + modules = { + "vllm": ModuleType("vllm"), + "vllm.device_allocator": ModuleType("vllm.device_allocator"), + "vllm.device_allocator.cumem": ModuleType("vllm.device_allocator.cumem"), + "vllm.v1": ModuleType("vllm.v1"), + "vllm.v1.worker": ModuleType("vllm.v1.worker"), + "vllm.v1.worker.ubatching": ModuleType("vllm.v1.worker.ubatching"), + "vllm.v1.worker.workspace": ModuleType("vllm.v1.worker.workspace"), + } + + class CuMemAllocator: + instance = allocator + + modules["vllm.device_allocator.cumem"].CuMemAllocator = CuMemAllocator + modules["vllm.v1.worker.ubatching"].dbo_current_ubatch_id = ubatch_id.get + modules["vllm.v1.worker.workspace"].WorkspaceManager = workspace_manager + modules["vllm.v1.worker.workspace"]._workspace_lane = lane + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + +def test_invalid_lane_does_not_release_adjacent_ubatch(monkeypatch: pytest.MonkeyPatch) -> None: + ubatch_id = ContextVar("test_ubatch", default=0) + lane = ContextVar("test_lane", default=0) + other_workspace = object() + other_pool = (object(), object()) + + class WorkspaceManager: + def __init__(self) -> None: + self._num_lanes = 1 + self._current_workspaces = [None, other_workspace] + + def _workspace_size_bytes(self, workspace: object | None) -> int: + return int(workspace is not None) + + def is_locked(self) -> bool: + return False + + def _ensure_workspace_size(self, required_bytes: int): + current_lane = lane.get() + if current_lane >= self._num_lanes: + raise RuntimeError(f"Workspace lane {current_lane} is not configured") + return self._current_workspaces[ubatch_id.get() * self._num_lanes + current_lane] + + allocator = _Allocator() + allocator.allocator_and_pools["moe_workspace:1"] = other_pool + _install_fake_vllm(monkeypatch, WorkspaceManager, allocator, ubatch_id, lane) + patch_moe_workspace_pool() + + manager = WorkspaceManager() + token = lane.set(1) + try: + with pytest.raises(RuntimeError, match="lane 1 is not configured"): + manager._ensure_workspace_size(2) + finally: + lane.reset(token) + + assert manager._current_workspaces[1] is other_workspace + assert allocator.allocator_and_pools["moe_workspace:1"] is other_pool + + +def test_two_ubatches_by_two_lanes_use_distinct_pool_keys(monkeypatch: pytest.MonkeyPatch) -> None: + ubatch_id = ContextVar("test_ubatch", default=0) + lane = ContextVar("test_lane", default=0) + + class WorkspaceManager: + def __init__(self) -> None: + self._num_lanes = 2 + self._current_workspaces = [None] * 4 + + def _workspace_size_bytes(self, workspace: object | None) -> int: + return int(workspace is not None) + + def is_locked(self) -> bool: + return False + + def _ensure_workspace_size(self, required_bytes: int): + workspace_id = ubatch_id.get() * self._num_lanes + lane.get() + self._current_workspaces[workspace_id] = object() + return self._current_workspaces[workspace_id] + + allocator = _Allocator() + _install_fake_vllm(monkeypatch, WorkspaceManager, allocator, ubatch_id, lane) + patch_moe_workspace_pool() + + manager = WorkspaceManager() + for ubatch in range(2): + for workspace_lane in range(2): + ubatch_token = ubatch_id.set(ubatch) + lane_token = lane.set(workspace_lane) + try: + manager._ensure_workspace_size(1) + finally: + lane.reset(lane_token) + ubatch_id.reset(ubatch_token) + + assert allocator.entered_pool_keys == [ + "moe_workspace:0", + "moe_workspace:1", + "moe_workspace:2", + "moe_workspace:3", + ] + assert all(workspace is not None for workspace in manager._current_workspaces) diff --git a/tests/rollout/vllm_omni/test_dependency_contracts.py b/tests/rollout/vllm_omni/test_dependency_contracts.py new file mode 100644 index 000000000..80de303a0 --- /dev/null +++ b/tests/rollout/vllm_omni/test_dependency_contracts.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +import tomllib +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[3] + + +def test_stable_vllm_omni_dependency_contract() -> None: + with (ROOT / "pyproject.toml").open("rb") as file: + project = tomllib.load(file) + + extras = project["project"]["optional-dependencies"] + assert "vllm==0.28.0 ; sys_platform == 'linux'" in extras["vllm"] + assert "vllm-omni==0.28.0 ; sys_platform == 'linux'" in extras["vllm"] + assert "transformers==5.12.1 ; sys_platform == 'linux'" in extras["vllm"] + assert "kernels>=0.12,<0.13 ; sys_platform == 'linux'" in extras["sglang"] + + overrides = project["tool"]["uv"]["override-dependencies"] + assert "diffusers==0.40.0" in overrides + assert "tokenizers>=0.22,<0.23" in overrides + assert all("kernels" not in requirement for requirement in overrides) diff --git a/tests/rollout/vllm_omni/test_deploy_configs.py b/tests/rollout/vllm_omni/test_deploy_configs.py new file mode 100644 index 000000000..05dab3515 --- /dev/null +++ b/tests/rollout/vllm_omni/test_deploy_configs.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +import yaml + +ROOT = Path(__file__).resolve().parents[3] +DEPLOY_DIR = ROOT / "unirl/rollout/engine/vllm_omni/deploy_configs" + + +def _walk(value): + if isinstance(value, dict): + yield value + for child in value.values(): + yield from _walk(child) + elif isinstance(value, list): + for child in value: + yield from _walk(child) + + +def test_all_deploy_configs_use_stable_schema() -> None: + paths = sorted(DEPLOY_DIR.glob("*.yaml")) + assert len(paths) == 12 + + for path in paths: + config = yaml.safe_load(path.read_text()) + assert "stage_args" not in config, path + assert isinstance(config.get("pipeline"), str), path + assert isinstance(config.get("stages"), list), path + assert config["stages"], path + + +def test_all_vllm_omni_recipes_use_deploy_config_vocabulary() -> None: + examples = ROOT / "examples" + paths = { + *examples.glob("**/*vllmomni*.yaml"), + *examples.glob("**/*vllm_omni*.yaml"), + examples / "diffusion/bagel/bagel_it2i_managed_editscore.yaml", + } + assert len(paths) == 22 + + for path in sorted(paths): + config = yaml.safe_load(path.read_text()) + for mapping in _walk(config): + assert "stage_yaml_override" not in mapping, path + assert "stage_yaml" not in mapping, path + override = mapping.get("deploy_config_override") + if override is not None: + assert (DEPLOY_DIR / override).is_file(), (path, override) + + +def test_every_adapter_deploy_config_exists() -> None: + from unirl.rollout.engine.vllm_omni.adapters import registered_adapters + from unirl.rollout.engine.vllm_omni.adapters.base import get_adapter + + for name in registered_adapters(): + filename = get_adapter(name).deploy_config + assert filename + assert (DEPLOY_DIR / filename).is_file(), (name, filename) + + +def test_stable_factory_resolves_every_deploy_config() -> None: + pytest.importorskip("vllm_omni") + from vllm_omni.config.config_factory import StageConfigFactory + from vllm_omni.config.pipeline_registry import resolve_pipeline_config + from vllm_omni.config.stage_config import load_deploy_config + + from unirl.rollout.engine.vllm_omni.pipeline_configs import register_unirl_pipeline_configs + + register_unirl_pipeline_configs() + for path in sorted(DEPLOY_DIR.glob("*.yaml")): + deploy = load_deploy_config(path) + pipeline = resolve_pipeline_config(deploy.pipeline) + assert pipeline is not None, path + stages, _ = StageConfigFactory._create_legacy_from_registry( + pipeline, + {}, + str(path), + deploy, + ) + assert stages, path + assert [stage.stage_id for stage in stages] == list(range(len(stages))), path + + +def test_hi3_text_topologies_preserve_legacy_output_contract() -> None: + pytest.importorskip("vllm_omni") + from unirl.rollout.engine.vllm_omni.pipeline_configs import ( + UNIRL_HI3_AR_MULTIMODAL_TEXT, + UNIRL_HI3_AR_TEXT, + ) + + text_stage = UNIRL_HI3_AR_TEXT.stages[0] + multimodal_stage = UNIRL_HI3_AR_MULTIMODAL_TEXT.stages[0] + for stage in (text_stage, multimodal_stage): + assert stage.final_output_type == "text" + assert stage.engine_output_type == "text" + assert stage.owns_tokenizer is True + assert text_stage.requires_multimodal_data is False + assert multimodal_stage.requires_multimodal_data is True diff --git a/tests/rollout/vllm_omni/test_stable_installed_contracts.py b/tests/rollout/vllm_omni/test_stable_installed_contracts.py new file mode 100644 index 000000000..b9fd83fe7 --- /dev/null +++ b/tests/rollout/vllm_omni/test_stable_installed_contracts.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +import importlib +from importlib.metadata import version + +import pytest + +pytest.importorskip("vllm") +pytest.importorskip("vllm_omni") + + +def test_exact_stable_runtime_is_installed() -> None: + assert version("vllm") == "0.28.0" + assert version("vllm-omni") == "0.28.0" + assert version("transformers") == "5.12.1" + assert version("kernels") == "0.14.1" + assert version("diffusers") == "0.40.0" + + +def test_real_stable_output_uses_flat_completion_fields() -> None: + from vllm.outputs import CompletionOutput + from vllm_omni.outputs import OmniRequestOutput + + from unirl.rollout.engine.vllm_omni.utils.tracks import _extract_completion, decoded_text_from_ar + + completion = CompletionOutput( + index=0, + text="answer", + token_ids=[1, 2], + cumulative_logprob=None, + logprobs=None, + finish_reason="stop", + ) + output = OmniRequestOutput( + request_id="0", + stage_id=0, + final_output_type="text", + outputs=[completion], + ) + + assert decoded_text_from_ar([[output]]).texts == ["answer"] + assert _extract_completion(output) == ([1, 2], None) + + +@pytest.mark.parametrize( + "module_name", + [ + "unirl.rollout.engine.vllm_omni.pipelines.sd3.pipeline", + "unirl.rollout.engine.vllm_omni.pipelines.qwen_image.pipeline", + "unirl.rollout.engine.vllm_omni.pipelines.hv15.pipeline", + "unirl.rollout.engine.vllm_omni.pipelines.hi3.pipeline", + "unirl.rollout.engine.vllm_omni.pipelines.bagel.pipeline", + "unirl.rollout.engine.vllm_omni.worker.dit_extension", + "unirl.rollout.engine.vllm_omni.worker.ar_extension", + ], +) +def test_custom_pipeline_and_worker_modules_import(module_name: str) -> None: + importlib.import_module(module_name) + + +def test_capture_plugin_patches_both_formatter_bindings() -> None: + from vllm_omni.diffusion import diffusion_engine, output_formatter + + from unirl.rollout.engine.vllm_omni.plugin import register_capture_flush + + register_capture_flush() + patched = output_formatter.format_diffusion_outputs + assert patched is diffusion_engine.format_diffusion_outputs + + register_capture_flush() + assert output_formatter.format_diffusion_outputs is patched + + +def test_diffusion_lora_loader_delegates_non_tensor_requests(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager + + from unirl.rollout.engine.vllm_omni.patches.runtime import patch_dit_lora_loader + + request = object() + expected = object() + calls = [] + + def original(manager, received): + calls.append((manager, received)) + return expected + + monkeypatch.setattr(DiffusionLoRAManager, "_load_adapter", original) + patch_dit_lora_loader() + + manager = object() + assert DiffusionLoRAManager._load_adapter(manager, request) is expected + assert calls == [(manager, request)] diff --git a/tests/rollout/vllm_omni/test_stable_runtime_contracts.py b/tests/rollout/vllm_omni/test_stable_runtime_contracts.py new file mode 100644 index 000000000..6e6b0c7f3 --- /dev/null +++ b/tests/rollout/vllm_omni/test_stable_runtime_contracts.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from unirl.rollout.engine.vllm_omni import patches, plugin +from unirl.rollout.engine.vllm_omni.backends import native +from unirl.rollout.engine.vllm_omni.backends.native import VLLMOmniBackend +from unirl.rollout.engine.vllm_omni.utils.tracks import _extract_completion, decoded_text_from_ar + + +class _Task: + def __init__(self, *, task_id: str, level: int | None = None, tags=None) -> None: + self.task_id = task_id + self.level = level + self.tags = tags + + +class _Engine: + num_stages = 2 + + def __init__(self) -> None: + self.calls: list[tuple[int, str, tuple, dict | None]] = [] + + def get_stage_metadata(self, stage_id: int): + return SimpleNamespace(stage_type="llm" if stage_id == 0 else "diffusion") + + def collective_rpc(self, *, method: str, args=(), kwargs=None, stage_ids=None): + stage_id = stage_ids[0] + self.calls.append((stage_id, method, args, kwargs)) + if method in {"handle_sleep_task", "handle_wake_task"}: + task = args[0] + return [ + { + "status": "SUCCESS", + "stage_id": stage_id, + "task_id": task.task_id, + "rank": 0, + } + ] + return [True] + + +def _backend(engine: _Engine) -> VLLMOmniBackend: + runtime = { + "OmniSleepTask": _Task, + "OmniWakeTask": _Task, + } + return VLLMOmniBackend( + SimpleNamespace(engine=engine), + runtime, + tokenizer=None, + tp_per_stage={0: 1, 1: 1}, + ) + + +def test_stable_outputs_are_read_directly() -> None: + completion = SimpleNamespace(text="answer", token_ids=[1, 2], logprobs=None) + output = SimpleNamespace(stage_id=0, final_output_type="text", outputs=[completion]) + + texts = decoded_text_from_ar([[output]]) + tokens, logprobs = _extract_completion(output) + + assert texts.texts == ["answer"] + assert tokens == [1, 2] + assert logprobs is None + + +def test_sleep_and_wake_route_by_stage_type(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + engine = _Engine() + backend = _backend(engine) + + backend.sleep_task() + backend.wake_task() + + assert [(stage, method) for stage, method, _args, _kwargs in engine.calls] == [ + (0, "sleep"), + (1, "handle_sleep_task"), + (0, "wake_up"), + (1, "handle_wake_task"), + ] + assert engine.calls[0][2] == (1, "abort") + assert engine.calls[2][3] == {"tags": None} + + +def test_collective_rpc_rejects_unsupported_stage() -> None: + backend = _backend(_Engine()) + with pytest.raises(RuntimeError, match="not installed"): + backend._require_rpc_success( + "update_weights", + 1, + [{"supported": False, "error": "worker extension is not installed"}], + ) + + +def test_boot_passes_stable_deploy_config_and_reads_engine_stages(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict = {} + + class Omni: + def __init__(self, **kwargs) -> None: + captured.update(kwargs) + self.engine = SimpleNamespace( + num_stages=1, + stage_configs=[ + { + "stage_id": 0, + "engine_args": {"tensor_parallel_size": 2}, + } + ], + ) + + monkeypatch.setattr(patches, "install", lambda: None) + monkeypatch.setattr(plugin, "register_unirl_runtime", lambda: None) + monkeypatch.setattr( + native, + "_import_omni_runtime", + lambda: { + "Omni": Omni, + "OmniSleepTask": _Task, + "OmniWakeTask": _Task, + }, + ) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + backend = VLLMOmniBackend.boot( + { + "model_path": "unused", + "deploy_config": "sd35_t2i_rl.yaml", + "enable_sleep_mode": False, + "ports": None, + } + ) + + assert captured["model"] == "unused" + assert captured["deploy_config"].endswith("deploy_configs/sd35_t2i_rl.yaml") + assert "stage_configs_path" not in captured + assert backend.tp_per_stage() == {0: 2} From 1e0ff2f8d1a66453b915599fa2e8a1717f6b301f Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 15 Sep 2026 00:05:32 +0800 Subject: [PATCH 05/23] test(rollout): add real CuMem workspace smoke --- .../rollout/vllm_omni/gpu_cumem_workspace.py | 74 +++++++++++++++++++ 1 file changed, 74 insertions(+) create mode 100644 tests/rollout/vllm_omni/gpu_cumem_workspace.py diff --git a/tests/rollout/vllm_omni/gpu_cumem_workspace.py b/tests/rollout/vllm_omni/gpu_cumem_workspace.py new file mode 100644 index 000000000..4dd244787 --- /dev/null +++ b/tests/rollout/vllm_omni/gpu_cumem_workspace.py @@ -0,0 +1,74 @@ +"""Real-GPU CuMem workspace isolation and sleep/wake smoke.""" + +from __future__ import annotations + +import gc + +import torch + + +def main() -> None: + assert torch.cuda.is_available() + + from vllm.device_allocator.cumem import CuMemAllocator, cumem_available + from vllm.v1.worker import ubatching, workspace + from vllm.v1.worker.workspace import WorkspaceManager, use_workspace_lane + + from unirl.rollout.engine.vllm_omni.patches.compat_moe_workspace import ( + MOE_WORKSPACE_TAG, + patch_moe_workspace_pool, + ) + + assert cumem_available + current_ubatch = 0 + + def ubatch_id() -> int: + return current_ubatch + + ubatching.dbo_current_ubatch_id = ubatch_id + workspace.dbo_current_ubatch_id = ubatch_id + + allocator = CuMemAllocator.get_instance() + patch_moe_workspace_pool() + manager = WorkspaceManager(torch.device("cuda"), num_ubatches=2, num_lanes=2) + + for ubatch in range(2): + current_ubatch = ubatch + for lane in range(2): + with use_workspace_lane(lane): + tensor = manager._ensure_workspace_size(1024 * 1024 + ubatch * 4096 + lane * 1024) + tensor.fill_(ubatch * 2 + lane) + + expected_keys = {f"{MOE_WORKSPACE_TAG}:{index}" for index in range(4)} + assert expected_keys.issubset(allocator.allocator_and_pools) + assert len({id(tensor) for tensor in manager._current_workspaces}) == 4 + + current_ubatch = 0 + saved_workspace = manager._current_workspaces[2] + saved_pool = allocator.allocator_and_pools[f"{MOE_WORKSPACE_TAG}:2"] + with use_workspace_lane(2): + try: + manager._ensure_workspace_size(2 * 1024 * 1024) + except RuntimeError as error: + assert "not configured" in str(error) + else: + raise AssertionError("invalid workspace lane was accepted") + assert manager._current_workspaces[2] is saved_workspace + assert allocator.allocator_and_pools[f"{MOE_WORKSPACE_TAG}:2"] is saved_pool + + for _ in range(2): + allocator.sleep(offload_tags=("weights",)) + allocator.wake_up() + for tensor in manager._current_workspaces: + tensor.zero_() + torch.accelerator.synchronize() + + manager._current_workspaces = [None] * 4 + gc.collect() + allocator.release_pools() + CuMemAllocator.instance = None + print("gpu_cumem_workspace_ok") + + +if __name__ == "__main__": + main() From c08fc96b9f46b67c4e34738f09b69f7dd6d645da Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 15 Sep 2026 09:24:20 +0800 Subject: [PATCH 06/23] ci: install vllm-omni contract dependencies --- .github/workflows/vllm-omni-contracts.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/vllm-omni-contracts.yml b/.github/workflows/vllm-omni-contracts.yml index f893da775..e90702c27 100644 --- a/.github/workflows/vllm-omni-contracts.yml +++ b/.github/workflows/vllm-omni-contracts.yml @@ -17,7 +17,7 @@ jobs: - name: Install focused test dependencies run: | - python -m pip install pytest pyyaml pillow psutil hydra-core omegaconf + python -m pip install pytest pyyaml pillow psutil hydra-core omegaconf numpy ray python -m pip install torch --index-url https://download.pytorch.org/whl/cpu - name: Run vLLM-Omni contracts From caf6195ed7e65a0d72b577815d9da12b42539f9c Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 15 Sep 2026 09:34:37 +0800 Subject: [PATCH 07/23] ci: install diffusion contract dependency --- .github/workflows/vllm-omni-contracts.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/vllm-omni-contracts.yml b/.github/workflows/vllm-omni-contracts.yml index e90702c27..7ce53c46f 100644 --- a/.github/workflows/vllm-omni-contracts.yml +++ b/.github/workflows/vllm-omni-contracts.yml @@ -17,7 +17,7 @@ jobs: - name: Install focused test dependencies run: | - python -m pip install pytest pyyaml pillow psutil hydra-core omegaconf numpy ray + python -m pip install pytest pyyaml pillow psutil hydra-core omegaconf numpy ray diffusers==0.40.0 python -m pip install torch --index-url https://download.pytorch.org/whl/cpu - name: Run vLLM-Omni contracts From 144fdbb6b3ad2858a5fd09a00a26134176f73d3b Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 15 Sep 2026 10:00:22 +0800 Subject: [PATCH 08/23] fix(rollout): keep Qwen3-Omni thinker LoRA-capable --- .../rollout/vllm_omni/test_deploy_configs.py | 5 +++++ .../qwen3_omni_thinker_only_rl_1x4.yaml | 2 +- ..._omni_thinker_only_rl_audio_video_1x4.yaml | 2 +- .../engine/vllm_omni/pipeline_configs.py | 22 +++++++++++++++++++ 4 files changed, 29 insertions(+), 2 deletions(-) diff --git a/tests/rollout/vllm_omni/test_deploy_configs.py b/tests/rollout/vllm_omni/test_deploy_configs.py index 05dab3515..b5d7deaed 100644 --- a/tests/rollout/vllm_omni/test_deploy_configs.py +++ b/tests/rollout/vllm_omni/test_deploy_configs.py @@ -88,6 +88,7 @@ def test_hi3_text_topologies_preserve_legacy_output_contract() -> None: from unirl.rollout.engine.vllm_omni.pipeline_configs import ( UNIRL_HI3_AR_MULTIMODAL_TEXT, UNIRL_HI3_AR_TEXT, + UNIRL_QWEN3_OMNI_THINKER, ) text_stage = UNIRL_HI3_AR_TEXT.stages[0] @@ -98,3 +99,7 @@ def test_hi3_text_topologies_preserve_legacy_output_contract() -> None: assert stage.owns_tokenizer is True assert text_stage.requires_multimodal_data is False assert multimodal_stage.requires_multimodal_data is True + + qwen_stage = UNIRL_QWEN3_OMNI_THINKER.stages[0] + assert qwen_stage.model_arch == "Qwen3OmniMoeThinkerForConditionalGeneration" + assert qwen_stage.engine_output_type == "text" diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_1x4.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_1x4.yaml index 375d90aad..984d20fa1 100644 --- a/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_1x4.yaml +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_1x4.yaml @@ -1,5 +1,5 @@ # Qwen3-Omni Thinker-only deployment, TP=4. -pipeline: qwen3_omni_moe_thinker_only +pipeline: unirl_qwen3_omni_thinker async_chunk: false distributed_executor_backend: mp dtype: bfloat16 diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_audio_video_1x4.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_audio_video_1x4.yaml index 01cea9016..75b73c381 100644 --- a/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_audio_video_1x4.yaml +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/qwen3_omni_thinker_only_rl_audio_video_1x4.yaml @@ -1,5 +1,5 @@ # Qwen3-Omni Thinker-only audio-in-video deployment, TP=4. -pipeline: qwen3_omni_moe_thinker_only +pipeline: unirl_qwen3_omni_thinker async_chunk: false distributed_executor_backend: mp dtype: bfloat16 diff --git a/unirl/rollout/engine/vllm_omni/pipeline_configs.py b/unirl/rollout/engine/vllm_omni/pipeline_configs.py index 13f95d582..f2522b7c4 100644 --- a/unirl/rollout/engine/vllm_omni/pipeline_configs.py +++ b/unirl/rollout/engine/vllm_omni/pipeline_configs.py @@ -54,6 +54,26 @@ def _hi3_text_pipeline(model_type: str, *, requires_multimodal_data: bool) -> Pi requires_multimodal_data=True, ) +UNIRL_QWEN3_OMNI_THINKER = PipelineConfig( + model_type="unirl_qwen3_omni_thinker", + model_arch="Qwen3OmniMoeThinkerForConditionalGeneration", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="thinker", + execution_type=StageExecutionType.LLM_AR, + final_output=True, + final_output_type="text", + owns_tokenizer=True, + requires_multimodal_data=True, + hf_config_name="thinker_config", + model_arch="Qwen3OmniMoeThinkerForConditionalGeneration", + engine_output_type="text", + sampling_constraints={"detokenize": True}, + ), + ), +) + def register_unirl_pipeline_configs() -> None: """Register UniRL's custom topologies once, rejecting key collisions.""" @@ -61,6 +81,7 @@ def register_unirl_pipeline_configs() -> None: UNIRL_RL_IMAGE_DIFFUSION, UNIRL_HI3_AR_TEXT, UNIRL_HI3_AR_MULTIMODAL_TEXT, + UNIRL_QWEN3_OMNI_THINKER, ): existing = OMNI_PIPELINES.get(pipeline.model_type) if existing is None: @@ -75,6 +96,7 @@ def register_unirl_pipeline_configs() -> None: __all__ = [ "UNIRL_HI3_AR_MULTIMODAL_TEXT", "UNIRL_HI3_AR_TEXT", + "UNIRL_QWEN3_OMNI_THINKER", "UNIRL_RL_IMAGE_DIFFUSION", "register_unirl_pipeline_configs", ] From bb029ec40281f1d2cf06e6656d36a87c795e4c79 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 15 Sep 2026 11:12:55 +0800 Subject: [PATCH 09/23] fix(rollout): retain HI3 expert mapping compatibility --- .../test_stable_installed_contracts.py | 14 +++++ .../engine/vllm_omni/patches/README.md | 1 + .../vllm_omni/patches/compat_hi3_lora.py | 55 +++++++++++++++++++ .../vllm_omni/patches/compat_tokenizer.py | 2 + 4 files changed, 72 insertions(+) create mode 100644 unirl/rollout/engine/vllm_omni/patches/compat_hi3_lora.py diff --git a/tests/rollout/vllm_omni/test_stable_installed_contracts.py b/tests/rollout/vllm_omni/test_stable_installed_contracts.py index b9fd83fe7..32336524b 100644 --- a/tests/rollout/vllm_omni/test_stable_installed_contracts.py +++ b/tests/rollout/vllm_omni/test_stable_installed_contracts.py @@ -90,3 +90,17 @@ def original(manager, received): manager = object() assert DiffusionLoRAManager._load_adapter(manager, request) is expected assert calls == [(manager, request)] + + +def test_hi3_expert_mapping_tuple_is_unwrapped(monkeypatch: pytest.MonkeyPatch) -> None: + from vllm.model_executor import utils as vllm_utils + + from unirl.rollout.engine.vllm_omni.patches import compat_hi3_lora + + mapping = [("experts.0", "weight", 0, "w")] + monkeypatch.setattr(vllm_utils, "get_moe_expert_mapping", lambda _model: (mapping, {"old": "new"})) + monkeypatch.setattr(compat_hi3_lora, "_INSTALLED", False) + + compat_hi3_lora.install() + + assert vllm_utils.get_moe_expert_mapping(object()) == mapping diff --git a/unirl/rollout/engine/vllm_omni/patches/README.md b/unirl/rollout/engine/vllm_omni/patches/README.md index 2633a29f2..5d9445faa 100644 --- a/unirl/rollout/engine/vllm_omni/patches/README.md +++ b/unirl/rollout/engine/vllm_omni/patches/README.md @@ -50,6 +50,7 @@ All in `runtime.py` unless noted. | `patch_sigmas_passthrough` | HI3's DiT `scheduler.set_timesteps` never receives `sampling_params.sigmas` | upstream forwards `sigmas` itself | | `install_fate_sharing` | `PR_SET_PDEATHSIG` is bound by Linux to the **specific creating thread**, so arming it for children of short-lived init threads kills healthy workers; and a worker inside a CUDA/NCCL call never observes vLLM's `death_pipe` EOF | vllm's own child-reaping is thread-safe | | `compat_tokenizer` (module) | HI3's `__init__` looks up `` and computes `ratio_36 + 1`; the Base checkpoint ships ratio tokens 0-32 only → `TypeError: … 'NoneType' and 'int'`. Both the slow **and** fast tokenizer classes must be patched, not the shared base. The module import *is* the install trigger (it is the `HI3ARWorkerExtension` qualname target). Upstream ≥ v0.20.0 raises a clean `ValueError` instead — a better error, but the Base ckpt still needs this 0-fallback to work | Base-ckpt support is dropped (Instruct ships the tokens) | +| `compat_hi3_lora` (module) | HI3's AR model returns `(expert_mapping, replacements)`, while vLLM 0.28's LoRA manager still iterates the result as a flat four-field mapping and fails at boot | vLLM consumes the mapping-plus-replacements tuple natively | | `patch_moe_workspace_pool` (`compat_moe_workspace`) | vLLM's reusable MoE `WorkspaceManager` scratch buffer is allocated outside the `weights`/`kv_cache` CuMem pools, so worker sleep leaves it resident through the colocated train phase. Gives each DBO×lane workspace an independently replaceable CuMem pool while tagging all allocations `moe_workspace`; worker sleep then discards the scratch pages and UniRL's untagged wake remaps them. Self-skips when the allocator is absent | vLLM allocates the MoE workspace inside a sleep-managed pool itself | ## HI3 DiT LoRA namespace and layout diff --git a/unirl/rollout/engine/vllm_omni/patches/compat_hi3_lora.py b/unirl/rollout/engine/vllm_omni/patches/compat_hi3_lora.py new file mode 100644 index 000000000..8120a4c8f --- /dev/null +++ b/unirl/rollout/engine/vllm_omni/patches/compat_hi3_lora.py @@ -0,0 +1,55 @@ +"""Make vLLM's MoE LoRA mapping accept HI3's mapping-plus-replacements tuple.""" + +from __future__ import annotations + +_INSTALLED = False + + +def install() -> None: + """Unwrap HI3's ``(expert_mapping, replacements)`` result for vLLM LoRA.""" + global _INSTALLED + if _INSTALLED: + return + try: + from vllm.model_executor import utils as vllm_mu + except ImportError: + return + + original = getattr(vllm_mu, "get_moe_expert_mapping", None) + if original is None: + _INSTALLED = True + return + if getattr(original, "_diffrl_hi3_unwrap", False): + _INSTALLED = True + return + + def _patched(model, _orig=original): + result = _orig(model) + if ( + isinstance(result, tuple) + and len(result) == 2 + and not isinstance(result[0], (str, int, float)) + and hasattr(result[0], "__iter__") + and isinstance(result[1], dict) + ): + return result[0] + return result + + _patched._diffrl_hi3_unwrap = True # type: ignore[attr-defined] + vllm_mu.get_moe_expert_mapping = _patched + + try: + from vllm.lora import utils as vllm_lora_utils + + if hasattr(vllm_lora_utils, "get_moe_expert_mapping"): + vllm_lora_utils.get_moe_expert_mapping = _patched + except ImportError: + pass + + _INSTALLED = True + + +install() + + +__all__ = ["install"] diff --git a/unirl/rollout/engine/vllm_omni/patches/compat_tokenizer.py b/unirl/rollout/engine/vllm_omni/patches/compat_tokenizer.py index ccedfafed..a6578b208 100644 --- a/unirl/rollout/engine/vllm_omni/patches/compat_tokenizer.py +++ b/unirl/rollout/engine/vllm_omni/patches/compat_tokenizer.py @@ -48,6 +48,8 @@ def _filtered(self, tokens, *args, _orig=original, **kwargs): install() +from unirl.rollout.engine.vllm_omni.patches import compat_hi3_lora as _hi3_lora_compat # noqa: F401, E402 + class HI3ARWorkerExtension: """vllm-omni ``worker_extension_cls`` qualname target for HI3 AR.""" From e73fea5eb93bb4322aba4660c77f12e6669e618d Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 15 Sep 2026 11:27:23 +0800 Subject: [PATCH 10/23] fix(rollout): defer HI3 diffusion handoff until decode --- tests/rollout/vllm_omni/test_deploy_configs.py | 6 ++++++ .../vllm_omni/deploy_configs/hunyuan_image3_t2i_rl.yaml | 2 -- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/rollout/vllm_omni/test_deploy_configs.py b/tests/rollout/vllm_omni/test_deploy_configs.py index b5d7deaed..2e03a8637 100644 --- a/tests/rollout/vllm_omni/test_deploy_configs.py +++ b/tests/rollout/vllm_omni/test_deploy_configs.py @@ -103,3 +103,9 @@ def test_hi3_text_topologies_preserve_legacy_output_contract() -> None: qwen_stage = UNIRL_QWEN3_OMNI_THINKER.stages[0] assert qwen_stage.model_arch == "Qwen3OmniMoeThinkerForConditionalGeneration" assert qwen_stage.engine_output_type == "text" + + +def test_hi3_t2i_does_not_forward_at_raw_prefill_output() -> None: + config = yaml.safe_load((DEPLOY_DIR / "hunyuan_image3_t2i_rl.yaml").read_text()) + stage_zero = config["stages"][0] + assert stage_zero["omni_kv_config"] == {"need_send_cache": True} diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_t2i_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_t2i_rl.yaml index ecaf64e79..7f7bf433c 100644 --- a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_t2i_rl.yaml +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_t2i_rl.yaml @@ -22,8 +22,6 @@ stages: rope_type: default omni_kv_config: need_send_cache: true - kv_transfer_criteria: - type: prefill_finished default_sampling_params: temperature: 0.0 top_p: 1.0 From d2e807628442acc8026e8714cc911ee7ad9f6e2c Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 15 Sep 2026 12:28:02 +0800 Subject: [PATCH 11/23] fix(rollout): probe wrapped diffusion parameters --- .../test_stable_installed_contracts.py | 14 ++++++ .../vllm_omni/worker/ipc_receive_mixin.py | 47 ++++++++----------- 2 files changed, 34 insertions(+), 27 deletions(-) diff --git a/tests/rollout/vllm_omni/test_stable_installed_contracts.py b/tests/rollout/vllm_omni/test_stable_installed_contracts.py index 32336524b..a852788d6 100644 --- a/tests/rollout/vllm_omni/test_stable_installed_contracts.py +++ b/tests/rollout/vllm_omni/test_stable_installed_contracts.py @@ -104,3 +104,17 @@ def test_hi3_expert_mapping_tuple_is_unwrapped(monkeypatch: pytest.MonkeyPatch) compat_hi3_lora.install() assert vllm_utils.get_moe_expert_mapping(object()) == mapping + + +def test_diffusion_parameter_probe_reaches_wrapped_transformer() -> None: + import torch + + from unirl.rollout.engine.vllm_omni.worker.ipc_receive_mixin import BucketedIPCReceiveMixin + + worker = object.__new__(BucketedIPCReceiveMixin) + transformer = torch.nn.Linear(4, 3) + worker.model_runner = type("Runner", (), {"pipeline": type("Pipeline", (), {"transformer": transformer})()})() + + descriptions = worker._diffrl_describe_params() + + assert descriptions["weight"] == ((3, 4), "torch.float32") diff --git a/unirl/rollout/engine/vllm_omni/worker/ipc_receive_mixin.py b/unirl/rollout/engine/vllm_omni/worker/ipc_receive_mixin.py index 4d10273b7..3fc2af8be 100644 --- a/unirl/rollout/engine/vllm_omni/worker/ipc_receive_mixin.py +++ b/unirl/rollout/engine/vllm_omni/worker/ipc_receive_mixin.py @@ -226,20 +226,29 @@ def set_lora_from_tensor_dict_copy( ) return self.add_lora(request) + def _diffrl_parameter_source(self): + """Return the loaded module behind AR or diffusion runner wrappers.""" + runner = getattr(self, "model_runner", None) + if runner is None: + return None + queue = [getattr(runner, attr, None) for attr in ("pipeline", "model")] + seen: set[int] = set() + while queue: + obj = queue.pop(0) + if obj is None or id(obj) in seen: + continue + seen.add(id(obj)) + if callable(getattr(obj, "named_parameters", None)): + return obj + queue.extend(getattr(obj, attr, None) for attr in ("transformer", "model", "bagel")) + return None + def _diffrl_describe_params( self, names: Optional[list] = None, ) -> dict: """Return ``{name: (shape_tuple, dtype_str)}`` for the worker's loaded model.""" - runner = getattr(self, "model_runner", None) - if runner is None: - return {} - param_source = None - for attr in ("pipeline", "model"): - obj = getattr(runner, attr, None) - if obj is not None and hasattr(obj, "named_parameters"): - param_source = obj - break + param_source = self._diffrl_parameter_source() if param_source is None: return {} @@ -258,15 +267,7 @@ def _diffrl_param_checksums( """Return ``{name: short_sha256_hex}`` for the worker's loaded model.""" import hashlib - runner = getattr(self, "model_runner", None) - if runner is None: - return {} - param_source = None - for attr in ("pipeline", "model"): - obj = getattr(runner, attr, None) - if obj is not None and hasattr(obj, "named_parameters"): - param_source = obj - break + param_source = self._diffrl_parameter_source() if param_source is None: return {} @@ -298,15 +299,7 @@ def _diffrl_loaded_param_checksums( fingerprint_tensor, ) - runner = getattr(self, "model_runner", None) - if runner is None: - return {} - param_source = None - for attr in ("pipeline", "model"): - obj = getattr(runner, attr, None) - if obj is not None and hasattr(obj, "named_parameters"): - param_source = obj - break + param_source = self._diffrl_parameter_source() if param_source is None: return {} From 4668869bee3207e65be18cdf32514dfb8a848f2e Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 15 Sep 2026 12:33:26 +0800 Subject: [PATCH 12/23] fix(rollout): traverse stable diffusion worker wrappers --- .../vllm_omni/worker/ipc_receive_mixin.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/unirl/rollout/engine/vllm_omni/worker/ipc_receive_mixin.py b/unirl/rollout/engine/vllm_omni/worker/ipc_receive_mixin.py index 3fc2af8be..f1a51e769 100644 --- a/unirl/rollout/engine/vllm_omni/worker/ipc_receive_mixin.py +++ b/unirl/rollout/engine/vllm_omni/worker/ipc_receive_mixin.py @@ -228,10 +228,7 @@ def set_lora_from_tensor_dict_copy( def _diffrl_parameter_source(self): """Return the loaded module behind AR or diffusion runner wrappers.""" - runner = getattr(self, "model_runner", None) - if runner is None: - return None - queue = [getattr(runner, attr, None) for attr in ("pipeline", "model")] + queue = [self] seen: set[int] = set() while queue: obj = queue.pop(0) @@ -240,7 +237,18 @@ def _diffrl_parameter_source(self): seen.add(id(obj)) if callable(getattr(obj, "named_parameters", None)): return obj - queue.extend(getattr(obj, attr, None) for attr in ("transformer", "model", "bagel")) + queue.extend( + getattr(obj, attr, None) + for attr in ( + "worker", + "model_runner", + "runner", + "pipeline", + "transformer", + "model", + "bagel", + ) + ) return None def _diffrl_describe_params( From 279c37ef30965dcbc268c54bbcb277d2ea636067 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 15 Sep 2026 12:38:47 +0800 Subject: [PATCH 13/23] fix(rollout): skip empty parameter wrappers --- .../vllm_omni/test_stable_installed_contracts.py | 1 + .../engine/vllm_omni/worker/ipc_receive_mixin.py | 10 ++++++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/rollout/vllm_omni/test_stable_installed_contracts.py b/tests/rollout/vllm_omni/test_stable_installed_contracts.py index a852788d6..bcda1b4e1 100644 --- a/tests/rollout/vllm_omni/test_stable_installed_contracts.py +++ b/tests/rollout/vllm_omni/test_stable_installed_contracts.py @@ -112,6 +112,7 @@ def test_diffusion_parameter_probe_reaches_wrapped_transformer() -> None: from unirl.rollout.engine.vllm_omni.worker.ipc_receive_mixin import BucketedIPCReceiveMixin worker = object.__new__(BucketedIPCReceiveMixin) + worker.worker = torch.nn.Module() transformer = torch.nn.Linear(4, 3) worker.model_runner = type("Runner", (), {"pipeline": type("Pipeline", (), {"transformer": transformer})()})() diff --git a/unirl/rollout/engine/vllm_omni/worker/ipc_receive_mixin.py b/unirl/rollout/engine/vllm_omni/worker/ipc_receive_mixin.py index f1a51e769..08ff0904f 100644 --- a/unirl/rollout/engine/vllm_omni/worker/ipc_receive_mixin.py +++ b/unirl/rollout/engine/vllm_omni/worker/ipc_receive_mixin.py @@ -235,8 +235,14 @@ def _diffrl_parameter_source(self): if obj is None or id(obj) in seen: continue seen.add(id(obj)) - if callable(getattr(obj, "named_parameters", None)): - return obj + named_parameters = getattr(obj, "named_parameters", None) + if callable(named_parameters): + try: + next(iter(named_parameters())) + except StopIteration: + pass + else: + return obj queue.extend( getattr(obj, attr, None) for attr in ( From 73e563b8f8e05061c2e2d8045eb755e64e503421 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 15 Sep 2026 13:52:11 +0800 Subject: [PATCH 14/23] fix(rollout): require stable stage metadata --- unirl/rollout/engine/vllm_omni/backends/native.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unirl/rollout/engine/vllm_omni/backends/native.py b/unirl/rollout/engine/vllm_omni/backends/native.py index 49d6126de..3f5358698 100644 --- a/unirl/rollout/engine/vllm_omni/backends/native.py +++ b/unirl/rollout/engine/vllm_omni/backends/native.py @@ -306,7 +306,7 @@ def _stage_ids(self) -> List[int]: def _stage_type(self, stage_id: int) -> str: metadata = self._require_omni().engine.get_stage_metadata(stage_id) - return str(getattr(metadata, "stage_type", "llm")) + return str(metadata.stage_type) @staticmethod def _require_rpc_success(action: str, stage_id: int, results: object) -> None: From d03c1b79f7cd0ecca20709b1fa2c9eb2408a4a47 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 15 Sep 2026 13:58:03 +0800 Subject: [PATCH 15/23] fix(rollout): preserve HI3 AR generation mode --- .../rollout/vllm_omni/test_deploy_configs.py | 15 ++++++++++++ .../test_stable_runtime_contracts.py | 10 ++++++++ .../engine/vllm_omni/backends/native.py | 24 +++++++++++++++---- .../hunyuan_image3_ar_recaption_rl.yaml | 2 +- .../deploy_configs/hunyuan_image3_ar_rl.yaml | 2 ++ .../deploy_configs/hunyuan_image3_i2t_rl.yaml | 2 ++ .../engine/vllm_omni/pipeline_configs.py | 22 +++++++++++++---- 7 files changed, 68 insertions(+), 9 deletions(-) diff --git a/tests/rollout/vllm_omni/test_deploy_configs.py b/tests/rollout/vllm_omni/test_deploy_configs.py index 2e03a8637..0579cd722 100644 --- a/tests/rollout/vllm_omni/test_deploy_configs.py +++ b/tests/rollout/vllm_omni/test_deploy_configs.py @@ -87,6 +87,7 @@ def test_hi3_text_topologies_preserve_legacy_output_contract() -> None: pytest.importorskip("vllm_omni") from unirl.rollout.engine.vllm_omni.pipeline_configs import ( UNIRL_HI3_AR_MULTIMODAL_TEXT, + UNIRL_HI3_AR_RECAPTION, UNIRL_HI3_AR_TEXT, UNIRL_QWEN3_OMNI_THINKER, ) @@ -100,6 +101,12 @@ def test_hi3_text_topologies_preserve_legacy_output_contract() -> None: assert text_stage.requires_multimodal_data is False assert multimodal_stage.requires_multimodal_data is True + recaption_stage = UNIRL_HI3_AR_RECAPTION.stages[0] + assert recaption_stage.final_output_type == "text" + assert recaption_stage.engine_output_type == "latent" + assert recaption_stage.owns_tokenizer is True + assert recaption_stage.requires_multimodal_data is False + qwen_stage = UNIRL_QWEN3_OMNI_THINKER.stages[0] assert qwen_stage.model_arch == "Qwen3OmniMoeThinkerForConditionalGeneration" assert qwen_stage.engine_output_type == "text" @@ -109,3 +116,11 @@ def test_hi3_t2i_does_not_forward_at_raw_prefill_output() -> None: config = yaml.safe_load((DEPLOY_DIR / "hunyuan_image3_t2i_rl.yaml").read_text()) stage_zero = config["stages"][0] assert stage_zero["omni_kv_config"] == {"need_send_cache": True} + + +@pytest.mark.parametrize("name", ["hunyuan_image3_ar_rl.yaml", "hunyuan_image3_i2t_rl.yaml"]) +def test_hi3_text_lora_adapters_install_weight_sync_extension(name: str) -> None: + config = yaml.safe_load((DEPLOY_DIR / name).read_text()) + stage_zero = config["stages"][0] + assert stage_zero["enable_lora"] is True + assert stage_zero["worker_extension_cls"].endswith(".HI3ARWeightSyncExtension") diff --git a/tests/rollout/vllm_omni/test_stable_runtime_contracts.py b/tests/rollout/vllm_omni/test_stable_runtime_contracts.py index 6e6b0c7f3..b9741a118 100644 --- a/tests/rollout/vllm_omni/test_stable_runtime_contracts.py +++ b/tests/rollout/vllm_omni/test_stable_runtime_contracts.py @@ -96,6 +96,16 @@ def test_collective_rpc_rejects_unsupported_stage() -> None: ) +def test_collective_rpc_rejects_false_except_idempotent_lora_removal() -> None: + backend = _backend(_Engine()) + with pytest.raises(RuntimeError, match="returned False"): + backend._require_rpc_success("set_lora_from_tensor_dict", 0, [False]) + with pytest.raises(RuntimeError, match="unexpected result"): + backend._require_rpc_success("update_weights", 0, ["failed"]) + + backend._require_rpc_success("remove_lora", 0, [False], allow_false=True) + + def test_boot_passes_stable_deploy_config_and_reads_engine_stages(monkeypatch: pytest.MonkeyPatch) -> None: captured: dict = {} diff --git a/unirl/rollout/engine/vllm_omni/backends/native.py b/unirl/rollout/engine/vllm_omni/backends/native.py index 3f5358698..ae91723c5 100644 --- a/unirl/rollout/engine/vllm_omni/backends/native.py +++ b/unirl/rollout/engine/vllm_omni/backends/native.py @@ -309,7 +309,13 @@ def _stage_type(self, stage_id: int) -> str: return str(metadata.stage_type) @staticmethod - def _require_rpc_success(action: str, stage_id: int, results: object) -> None: + def _require_rpc_success( + action: str, + stage_id: int, + results: object, + *, + allow_false: bool = False, + ) -> None: """Fail closed when a stage or replica reports an unsupported/failed RPC.""" def validate(result: object) -> None: @@ -319,8 +325,12 @@ def validate(result: object) -> None: for item in result: validate(item) return - if not isinstance(result, Mapping): + if isinstance(result, bool): + if not result and not allow_false: + raise RuntimeError(f"vllm-omni {action} returned False on stage {stage_id}") return + if not isinstance(result, Mapping): + raise RuntimeError(f"vllm-omni {action} returned an unexpected result on stage {stage_id}: {result!r}") status = result.get("status") error = result.get("error") @@ -337,6 +347,7 @@ def _collective_rpc( *, args: tuple = (), kwargs: Optional[dict] = None, + allow_false: bool = False, ) -> List[Any]: results = self._require_omni().engine.collective_rpc( method=method, @@ -344,7 +355,7 @@ def _collective_rpc( kwargs=kwargs, stage_ids=[stage_id], ) - self._require_rpc_success(method, stage_id, results) + self._require_rpc_success(method, stage_id, results, allow_false=allow_false) return results @staticmethod @@ -650,7 +661,12 @@ def _wrap_peft_envelope(lora_tensors: Dict[str, Any]) -> Dict[str, Any]: def _remove_existing_lora(self, adapter_id: int) -> None: """Drop the existing adapter on every stage before re-adding.""" for sid in self._stage_ids(): - self._collective_rpc(sid, "remove_lora", args=(int(adapter_id),)) + self._collective_rpc( + sid, + "remove_lora", + args=(int(adapter_id),), + allow_false=True, + ) def param_checksums(self, *, names: List[str]) -> dict: """Fan ``_diffrl_loaded_param_checksums`` across stages and ranks.""" diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_recaption_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_recaption_rl.yaml index ce0653938..0b1b48337 100644 --- a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_recaption_rl.yaml +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_recaption_rl.yaml @@ -1,5 +1,5 @@ # HunyuanImage-3 terminal AR recaption deployment for the two-engine trainer. -pipeline: unirl_hi3_ar_text +pipeline: unirl_hi3_ar_recaption async_chunk: false distributed_executor_backend: mp diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml index be521e8cc..fca277bbe 100644 --- a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml @@ -6,6 +6,8 @@ stages: - stage_id: 0 devices: "0,1,2,3" max_num_seqs: 1 + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.ar_extension.HI3ARWeightSyncExtension + enable_lora: true gpu_memory_utilization: 0.95 enforce_eager: true trust_remote_code: true diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml index cb650b959..8fd674866 100644 --- a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml @@ -6,6 +6,8 @@ stages: - stage_id: 0 devices: "0,1,2,3" max_num_seqs: 1 + worker_extension_cls: unirl.rollout.engine.vllm_omni.worker.ar_extension.HI3ARWeightSyncExtension + enable_lora: true gpu_memory_utilization: 0.95 enforce_eager: true trust_remote_code: true diff --git a/unirl/rollout/engine/vllm_omni/pipeline_configs.py b/unirl/rollout/engine/vllm_omni/pipeline_configs.py index f2522b7c4..e1f986a24 100644 --- a/unirl/rollout/engine/vllm_omni/pipeline_configs.py +++ b/unirl/rollout/engine/vllm_omni/pipeline_configs.py @@ -25,7 +25,12 @@ _HI3_MODEL_ARCH = "HunyuanImage3ForCausalMM" -def _hi3_text_pipeline(model_type: str, *, requires_multimodal_data: bool) -> PipelineConfig: +def _hi3_ar_pipeline( + model_type: str, + *, + requires_multimodal_data: bool, + engine_output_type: str, +) -> PipelineConfig: return PipelineConfig( model_type=model_type, model_arch=_HI3_MODEL_ARCH, @@ -39,19 +44,26 @@ def _hi3_text_pipeline(model_type: str, *, requires_multimodal_data: bool) -> Pi owns_tokenizer=True, requires_multimodal_data=requires_multimodal_data, model_arch=_HI3_MODEL_ARCH, - engine_output_type="text", + engine_output_type=engine_output_type, ), ), ) -UNIRL_HI3_AR_TEXT = _hi3_text_pipeline( +UNIRL_HI3_AR_TEXT = _hi3_ar_pipeline( "unirl_hi3_ar_text", requires_multimodal_data=False, + engine_output_type="text", ) -UNIRL_HI3_AR_MULTIMODAL_TEXT = _hi3_text_pipeline( +UNIRL_HI3_AR_MULTIMODAL_TEXT = _hi3_ar_pipeline( "unirl_hi3_ar_multimodal_text", requires_multimodal_data=True, + engine_output_type="text", +) +UNIRL_HI3_AR_RECAPTION = _hi3_ar_pipeline( + "unirl_hi3_ar_recaption", + requires_multimodal_data=False, + engine_output_type="latent", ) UNIRL_QWEN3_OMNI_THINKER = PipelineConfig( @@ -81,6 +93,7 @@ def register_unirl_pipeline_configs() -> None: UNIRL_RL_IMAGE_DIFFUSION, UNIRL_HI3_AR_TEXT, UNIRL_HI3_AR_MULTIMODAL_TEXT, + UNIRL_HI3_AR_RECAPTION, UNIRL_QWEN3_OMNI_THINKER, ): existing = OMNI_PIPELINES.get(pipeline.model_type) @@ -95,6 +108,7 @@ def register_unirl_pipeline_configs() -> None: __all__ = [ "UNIRL_HI3_AR_MULTIMODAL_TEXT", + "UNIRL_HI3_AR_RECAPTION", "UNIRL_HI3_AR_TEXT", "UNIRL_QWEN3_OMNI_THINKER", "UNIRL_RL_IMAGE_DIFFUSION", From acad292813f4c29ab80c747ea061c08918f2723f Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 15 Sep 2026 14:10:50 +0800 Subject: [PATCH 16/23] ci: exercise installed vllm-omni contracts --- .github/workflows/vllm-omni-contracts.yml | 16 ++++++++++++++-- tests/rollout/vllm_omni/test_deploy_configs.py | 2 +- .../vllm_omni/test_stable_installed_contracts.py | 9 +++++++-- .../rollout/engine/vllm_omni/backends/native.py | 2 +- unirl/rollout/engine/vllm_omni/patches/README.md | 6 +++--- 5 files changed, 26 insertions(+), 9 deletions(-) diff --git a/.github/workflows/vllm-omni-contracts.yml b/.github/workflows/vllm-omni-contracts.yml index 7ce53c46f..905a35043 100644 --- a/.github/workflows/vllm-omni-contracts.yml +++ b/.github/workflows/vllm-omni-contracts.yml @@ -17,8 +17,20 @@ jobs: - name: Install focused test dependencies run: | - python -m pip install pytest pyyaml pillow psutil hydra-core omegaconf numpy ray diffusers==0.40.0 - python -m pip install torch --index-url https://download.pytorch.org/whl/cpu + python -m pip install torch==2.13.0 torchvision==0.28.0 \ + --index-url https://download.pytorch.org/whl/cpu + python -m pip install \ + pytest pyyaml pillow psutil hydra-core omegaconf numpy ray \ + diffusers==0.40.0 transformers==5.12.1 kernels==0.14.1 \ + aenum==3.1.16 pydantic cbor2 pyzmq janus aiohttp \ + openai-harmony openai msgspec pybase64 blake3 setproctitle \ + prometheus-client cloudpickle py-cpuinfo cachetools uvloop \ + llguidance tiktoken xgrammar loguru fastapi prettytable \ + accelerate==1.12.0 soundfile torchsde "imageio[ffmpeg]" \ + x-transformers einops gguf pydantic-settings cache-dit==1.5.0 + python -m pip install --no-deps vllm==0.28.0 vllm-omni==0.28.0 - name: Run vLLM-Omni contracts + env: + VLLM_OMNI_REQUIRE_INSTALLED: "1" run: python -m pytest -q tests/rollout/vllm_omni diff --git a/tests/rollout/vllm_omni/test_deploy_configs.py b/tests/rollout/vllm_omni/test_deploy_configs.py index 0579cd722..41fea98d5 100644 --- a/tests/rollout/vllm_omni/test_deploy_configs.py +++ b/tests/rollout/vllm_omni/test_deploy_configs.py @@ -38,7 +38,7 @@ def test_all_vllm_omni_recipes_use_deploy_config_vocabulary() -> None: *examples.glob("**/*vllm_omni*.yaml"), examples / "diffusion/bagel/bagel_it2i_managed_editscore.yaml", } - assert len(paths) == 22 + assert paths for path in sorted(paths): config = yaml.safe_load(path.read_text()) diff --git a/tests/rollout/vllm_omni/test_stable_installed_contracts.py b/tests/rollout/vllm_omni/test_stable_installed_contracts.py index bcda1b4e1..e1e1e91c9 100644 --- a/tests/rollout/vllm_omni/test_stable_installed_contracts.py +++ b/tests/rollout/vllm_omni/test_stable_installed_contracts.py @@ -1,12 +1,17 @@ from __future__ import annotations import importlib +import os from importlib.metadata import version import pytest -pytest.importorskip("vllm") -pytest.importorskip("vllm_omni") +if os.environ.get("VLLM_OMNI_REQUIRE_INSTALLED") == "1": + importlib.import_module("vllm") + importlib.import_module("vllm_omni") +else: + pytest.importorskip("vllm") + pytest.importorskip("vllm_omni") def test_exact_stable_runtime_is_installed() -> None: diff --git a/unirl/rollout/engine/vllm_omni/backends/native.py b/unirl/rollout/engine/vllm_omni/backends/native.py index ae91723c5..c1fb2ba73 100644 --- a/unirl/rollout/engine/vllm_omni/backends/native.py +++ b/unirl/rollout/engine/vllm_omni/backends/native.py @@ -292,7 +292,7 @@ def tokenize_prompt(self, text: str, *, task: str, sys_type: str) -> List[int]: build_prompt_tokens, ) - # vllm-omni 0.27 returns a PromptTokensResult here, not a bare id list. + # vLLM-Omni returns a PromptTokensResult here, not a bare id list. return list(build_prompt_tokens(text, self._tokenizer, task=task, sys_type=sys_type).token_ids) def num_stages(self) -> int: diff --git a/unirl/rollout/engine/vllm_omni/patches/README.md b/unirl/rollout/engine/vllm_omni/patches/README.md index 5d9445faa..2a8f3184e 100644 --- a/unirl/rollout/engine/vllm_omni/patches/README.md +++ b/unirl/rollout/engine/vllm_omni/patches/README.md @@ -80,6 +80,6 @@ known-invalid adapter. submodules, loaded lazily — `import unirl.rollout.engine.vllm_omni.patches` must not pull vllm. - **Every patch needs a DELETE-WHEN row.** Without one it is permanent by default. -- **0.27 honors `engine_args.diffusion_attention_backend`, not `attention_backend`.** - The old AR field is silently dropped by `OmniDiffusionConfig.from_kwargs`; - HV1.5 pins the diffusion backend in its stage YAML. +- **0.28 honors `diffusion_attention_backend`, not `attention_backend`.** + The AR field is silently dropped by `OmniDiffusionConfig.from_kwargs`; + HV1.5 pins the diffusion backend in its deploy config. From 44441e0840525fb9bd8c9ab88626f92484b9f9c8 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Tue, 15 Sep 2026 16:45:35 +0800 Subject: [PATCH 17/23] fix(rollout): pin HI3 AR stages to mp --- tests/rollout/vllm_omni/test_deploy_configs.py | 12 ++++++++++-- .../deploy_configs/hunyuan_image3_ar_rl.yaml | 1 + .../deploy_configs/hunyuan_image3_i2t_rl.yaml | 1 + .../deploy_configs/hunyuan_image3_it2i_rl.yaml | 1 + 4 files changed, 13 insertions(+), 2 deletions(-) diff --git a/tests/rollout/vllm_omni/test_deploy_configs.py b/tests/rollout/vllm_omni/test_deploy_configs.py index 41fea98d5..1cbe6693b 100644 --- a/tests/rollout/vllm_omni/test_deploy_configs.py +++ b/tests/rollout/vllm_omni/test_deploy_configs.py @@ -118,9 +118,17 @@ def test_hi3_t2i_does_not_forward_at_raw_prefill_output() -> None: assert stage_zero["omni_kv_config"] == {"need_send_cache": True} -@pytest.mark.parametrize("name", ["hunyuan_image3_ar_rl.yaml", "hunyuan_image3_i2t_rl.yaml"]) -def test_hi3_text_lora_adapters_install_weight_sync_extension(name: str) -> None: +@pytest.mark.parametrize( + "name", + [ + "hunyuan_image3_ar_rl.yaml", + "hunyuan_image3_i2t_rl.yaml", + "hunyuan_image3_it2i_rl.yaml", + ], +) +def test_hi3_ar_lora_stages_force_mp_weight_sync_workers(name: str) -> None: config = yaml.safe_load((DEPLOY_DIR / name).read_text()) stage_zero = config["stages"][0] assert stage_zero["enable_lora"] is True assert stage_zero["worker_extension_cls"].endswith(".HI3ARWeightSyncExtension") + assert stage_zero["distributed_executor_backend"] == "mp" diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml index fca277bbe..9fbd4a294 100644 --- a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml @@ -15,6 +15,7 @@ stages: max_num_batched_tokens: 32768 tensor_parallel_size: 4 pipeline_parallel_size: 1 + distributed_executor_backend: mp hf_overrides: rope_parameters: mrope_section: [0, 32, 32] diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml index 8fd674866..8b0098343 100644 --- a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml @@ -15,6 +15,7 @@ stages: max_num_batched_tokens: 32768 tensor_parallel_size: 4 pipeline_parallel_size: 1 + distributed_executor_backend: mp hf_overrides: rope_parameters: mrope_section: [0, 32, 32] diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_it2i_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_it2i_rl.yaml index 9caadb965..634e08989 100644 --- a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_it2i_rl.yaml +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_it2i_rl.yaml @@ -15,6 +15,7 @@ stages: max_num_batched_tokens: 32768 tensor_parallel_size: 4 pipeline_parallel_size: 1 + distributed_executor_backend: mp hf_overrides: rope_parameters: mrope_section: [0, 32, 32] From aa11992f9f8a69c25774d4fe18d55810fe0d5b48 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Fri, 18 Sep 2026 09:24:35 +0800 Subject: [PATCH 18/23] test(rollout): remove vllm-omni PR tests --- .github/workflows/vllm-omni-contracts.yml | 36 ----- .../rollout/vllm_omni/gpu_cumem_workspace.py | 74 --------- .../vllm_omni/test_compat_moe_workspace.py | 135 ---------------- .../vllm_omni/test_dependency_contracts.py | 22 --- .../rollout/vllm_omni/test_deploy_configs.py | 134 ---------------- .../test_stable_installed_contracts.py | 126 --------------- .../test_stable_runtime_contracts.py | 150 ------------------ 7 files changed, 677 deletions(-) delete mode 100644 .github/workflows/vllm-omni-contracts.yml delete mode 100644 tests/rollout/vllm_omni/gpu_cumem_workspace.py delete mode 100644 tests/rollout/vllm_omni/test_compat_moe_workspace.py delete mode 100644 tests/rollout/vllm_omni/test_dependency_contracts.py delete mode 100644 tests/rollout/vllm_omni/test_deploy_configs.py delete mode 100644 tests/rollout/vllm_omni/test_stable_installed_contracts.py delete mode 100644 tests/rollout/vllm_omni/test_stable_runtime_contracts.py diff --git a/.github/workflows/vllm-omni-contracts.yml b/.github/workflows/vllm-omni-contracts.yml deleted file mode 100644 index 905a35043..000000000 --- a/.github/workflows/vllm-omni-contracts.yml +++ /dev/null @@ -1,36 +0,0 @@ -name: vLLM-Omni Contracts - -on: - pull_request: - branches: [main] - -jobs: - contracts: - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - - name: Set up Python - uses: actions/setup-python@v5 - with: - python-version: "3.12" - - - name: Install focused test dependencies - run: | - python -m pip install torch==2.13.0 torchvision==0.28.0 \ - --index-url https://download.pytorch.org/whl/cpu - python -m pip install \ - pytest pyyaml pillow psutil hydra-core omegaconf numpy ray \ - diffusers==0.40.0 transformers==5.12.1 kernels==0.14.1 \ - aenum==3.1.16 pydantic cbor2 pyzmq janus aiohttp \ - openai-harmony openai msgspec pybase64 blake3 setproctitle \ - prometheus-client cloudpickle py-cpuinfo cachetools uvloop \ - llguidance tiktoken xgrammar loguru fastapi prettytable \ - accelerate==1.12.0 soundfile torchsde "imageio[ffmpeg]" \ - x-transformers einops gguf pydantic-settings cache-dit==1.5.0 - python -m pip install --no-deps vllm==0.28.0 vllm-omni==0.28.0 - - - name: Run vLLM-Omni contracts - env: - VLLM_OMNI_REQUIRE_INSTALLED: "1" - run: python -m pytest -q tests/rollout/vllm_omni diff --git a/tests/rollout/vllm_omni/gpu_cumem_workspace.py b/tests/rollout/vllm_omni/gpu_cumem_workspace.py deleted file mode 100644 index 4dd244787..000000000 --- a/tests/rollout/vllm_omni/gpu_cumem_workspace.py +++ /dev/null @@ -1,74 +0,0 @@ -"""Real-GPU CuMem workspace isolation and sleep/wake smoke.""" - -from __future__ import annotations - -import gc - -import torch - - -def main() -> None: - assert torch.cuda.is_available() - - from vllm.device_allocator.cumem import CuMemAllocator, cumem_available - from vllm.v1.worker import ubatching, workspace - from vllm.v1.worker.workspace import WorkspaceManager, use_workspace_lane - - from unirl.rollout.engine.vllm_omni.patches.compat_moe_workspace import ( - MOE_WORKSPACE_TAG, - patch_moe_workspace_pool, - ) - - assert cumem_available - current_ubatch = 0 - - def ubatch_id() -> int: - return current_ubatch - - ubatching.dbo_current_ubatch_id = ubatch_id - workspace.dbo_current_ubatch_id = ubatch_id - - allocator = CuMemAllocator.get_instance() - patch_moe_workspace_pool() - manager = WorkspaceManager(torch.device("cuda"), num_ubatches=2, num_lanes=2) - - for ubatch in range(2): - current_ubatch = ubatch - for lane in range(2): - with use_workspace_lane(lane): - tensor = manager._ensure_workspace_size(1024 * 1024 + ubatch * 4096 + lane * 1024) - tensor.fill_(ubatch * 2 + lane) - - expected_keys = {f"{MOE_WORKSPACE_TAG}:{index}" for index in range(4)} - assert expected_keys.issubset(allocator.allocator_and_pools) - assert len({id(tensor) for tensor in manager._current_workspaces}) == 4 - - current_ubatch = 0 - saved_workspace = manager._current_workspaces[2] - saved_pool = allocator.allocator_and_pools[f"{MOE_WORKSPACE_TAG}:2"] - with use_workspace_lane(2): - try: - manager._ensure_workspace_size(2 * 1024 * 1024) - except RuntimeError as error: - assert "not configured" in str(error) - else: - raise AssertionError("invalid workspace lane was accepted") - assert manager._current_workspaces[2] is saved_workspace - assert allocator.allocator_and_pools[f"{MOE_WORKSPACE_TAG}:2"] is saved_pool - - for _ in range(2): - allocator.sleep(offload_tags=("weights",)) - allocator.wake_up() - for tensor in manager._current_workspaces: - tensor.zero_() - torch.accelerator.synchronize() - - manager._current_workspaces = [None] * 4 - gc.collect() - allocator.release_pools() - CuMemAllocator.instance = None - print("gpu_cumem_workspace_ok") - - -if __name__ == "__main__": - main() diff --git a/tests/rollout/vllm_omni/test_compat_moe_workspace.py b/tests/rollout/vllm_omni/test_compat_moe_workspace.py deleted file mode 100644 index 8851c1090..000000000 --- a/tests/rollout/vllm_omni/test_compat_moe_workspace.py +++ /dev/null @@ -1,135 +0,0 @@ -from __future__ import annotations - -import sys -from contextlib import contextmanager -from contextvars import ContextVar -from types import ModuleType - -import pytest - -from unirl.rollout.engine.vllm_omni.patches.compat_moe_workspace import patch_moe_workspace_pool - - -class _Allocator: - def __init__(self) -> None: - self.current_tag = None - self.allocator_and_pools: dict[str, tuple[object, object]] = {} - self.entered_pool_keys: list[str] = [] - - @contextmanager - def use_memory_pool(self, tag: str): - self.entered_pool_keys.append(tag) - self.allocator_and_pools[tag] = (object(), object()) - yield - - -def _install_fake_vllm( - monkeypatch: pytest.MonkeyPatch, - workspace_manager: type, - allocator: _Allocator, - ubatch_id: ContextVar[int], - lane: ContextVar[int], -) -> None: - modules = { - "vllm": ModuleType("vllm"), - "vllm.device_allocator": ModuleType("vllm.device_allocator"), - "vllm.device_allocator.cumem": ModuleType("vllm.device_allocator.cumem"), - "vllm.v1": ModuleType("vllm.v1"), - "vllm.v1.worker": ModuleType("vllm.v1.worker"), - "vllm.v1.worker.ubatching": ModuleType("vllm.v1.worker.ubatching"), - "vllm.v1.worker.workspace": ModuleType("vllm.v1.worker.workspace"), - } - - class CuMemAllocator: - instance = allocator - - modules["vllm.device_allocator.cumem"].CuMemAllocator = CuMemAllocator - modules["vllm.v1.worker.ubatching"].dbo_current_ubatch_id = ubatch_id.get - modules["vllm.v1.worker.workspace"].WorkspaceManager = workspace_manager - modules["vllm.v1.worker.workspace"]._workspace_lane = lane - for name, module in modules.items(): - monkeypatch.setitem(sys.modules, name, module) - - -def test_invalid_lane_does_not_release_adjacent_ubatch(monkeypatch: pytest.MonkeyPatch) -> None: - ubatch_id = ContextVar("test_ubatch", default=0) - lane = ContextVar("test_lane", default=0) - other_workspace = object() - other_pool = (object(), object()) - - class WorkspaceManager: - def __init__(self) -> None: - self._num_lanes = 1 - self._current_workspaces = [None, other_workspace] - - def _workspace_size_bytes(self, workspace: object | None) -> int: - return int(workspace is not None) - - def is_locked(self) -> bool: - return False - - def _ensure_workspace_size(self, required_bytes: int): - current_lane = lane.get() - if current_lane >= self._num_lanes: - raise RuntimeError(f"Workspace lane {current_lane} is not configured") - return self._current_workspaces[ubatch_id.get() * self._num_lanes + current_lane] - - allocator = _Allocator() - allocator.allocator_and_pools["moe_workspace:1"] = other_pool - _install_fake_vllm(monkeypatch, WorkspaceManager, allocator, ubatch_id, lane) - patch_moe_workspace_pool() - - manager = WorkspaceManager() - token = lane.set(1) - try: - with pytest.raises(RuntimeError, match="lane 1 is not configured"): - manager._ensure_workspace_size(2) - finally: - lane.reset(token) - - assert manager._current_workspaces[1] is other_workspace - assert allocator.allocator_and_pools["moe_workspace:1"] is other_pool - - -def test_two_ubatches_by_two_lanes_use_distinct_pool_keys(monkeypatch: pytest.MonkeyPatch) -> None: - ubatch_id = ContextVar("test_ubatch", default=0) - lane = ContextVar("test_lane", default=0) - - class WorkspaceManager: - def __init__(self) -> None: - self._num_lanes = 2 - self._current_workspaces = [None] * 4 - - def _workspace_size_bytes(self, workspace: object | None) -> int: - return int(workspace is not None) - - def is_locked(self) -> bool: - return False - - def _ensure_workspace_size(self, required_bytes: int): - workspace_id = ubatch_id.get() * self._num_lanes + lane.get() - self._current_workspaces[workspace_id] = object() - return self._current_workspaces[workspace_id] - - allocator = _Allocator() - _install_fake_vllm(monkeypatch, WorkspaceManager, allocator, ubatch_id, lane) - patch_moe_workspace_pool() - - manager = WorkspaceManager() - for ubatch in range(2): - for workspace_lane in range(2): - ubatch_token = ubatch_id.set(ubatch) - lane_token = lane.set(workspace_lane) - try: - manager._ensure_workspace_size(1) - finally: - lane.reset(lane_token) - ubatch_id.reset(ubatch_token) - - assert allocator.entered_pool_keys == [ - "moe_workspace:0", - "moe_workspace:1", - "moe_workspace:2", - "moe_workspace:3", - ] - assert all(workspace is not None for workspace in manager._current_workspaces) diff --git a/tests/rollout/vllm_omni/test_dependency_contracts.py b/tests/rollout/vllm_omni/test_dependency_contracts.py deleted file mode 100644 index 80de303a0..000000000 --- a/tests/rollout/vllm_omni/test_dependency_contracts.py +++ /dev/null @@ -1,22 +0,0 @@ -from __future__ import annotations - -import tomllib -from pathlib import Path - -ROOT = Path(__file__).resolve().parents[3] - - -def test_stable_vllm_omni_dependency_contract() -> None: - with (ROOT / "pyproject.toml").open("rb") as file: - project = tomllib.load(file) - - extras = project["project"]["optional-dependencies"] - assert "vllm==0.28.0 ; sys_platform == 'linux'" in extras["vllm"] - assert "vllm-omni==0.28.0 ; sys_platform == 'linux'" in extras["vllm"] - assert "transformers==5.12.1 ; sys_platform == 'linux'" in extras["vllm"] - assert "kernels>=0.12,<0.13 ; sys_platform == 'linux'" in extras["sglang"] - - overrides = project["tool"]["uv"]["override-dependencies"] - assert "diffusers==0.40.0" in overrides - assert "tokenizers>=0.22,<0.23" in overrides - assert all("kernels" not in requirement for requirement in overrides) diff --git a/tests/rollout/vllm_omni/test_deploy_configs.py b/tests/rollout/vllm_omni/test_deploy_configs.py deleted file mode 100644 index 1cbe6693b..000000000 --- a/tests/rollout/vllm_omni/test_deploy_configs.py +++ /dev/null @@ -1,134 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import pytest -import yaml - -ROOT = Path(__file__).resolve().parents[3] -DEPLOY_DIR = ROOT / "unirl/rollout/engine/vllm_omni/deploy_configs" - - -def _walk(value): - if isinstance(value, dict): - yield value - for child in value.values(): - yield from _walk(child) - elif isinstance(value, list): - for child in value: - yield from _walk(child) - - -def test_all_deploy_configs_use_stable_schema() -> None: - paths = sorted(DEPLOY_DIR.glob("*.yaml")) - assert len(paths) == 12 - - for path in paths: - config = yaml.safe_load(path.read_text()) - assert "stage_args" not in config, path - assert isinstance(config.get("pipeline"), str), path - assert isinstance(config.get("stages"), list), path - assert config["stages"], path - - -def test_all_vllm_omni_recipes_use_deploy_config_vocabulary() -> None: - examples = ROOT / "examples" - paths = { - *examples.glob("**/*vllmomni*.yaml"), - *examples.glob("**/*vllm_omni*.yaml"), - examples / "diffusion/bagel/bagel_it2i_managed_editscore.yaml", - } - assert paths - - for path in sorted(paths): - config = yaml.safe_load(path.read_text()) - for mapping in _walk(config): - assert "stage_yaml_override" not in mapping, path - assert "stage_yaml" not in mapping, path - override = mapping.get("deploy_config_override") - if override is not None: - assert (DEPLOY_DIR / override).is_file(), (path, override) - - -def test_every_adapter_deploy_config_exists() -> None: - from unirl.rollout.engine.vllm_omni.adapters import registered_adapters - from unirl.rollout.engine.vllm_omni.adapters.base import get_adapter - - for name in registered_adapters(): - filename = get_adapter(name).deploy_config - assert filename - assert (DEPLOY_DIR / filename).is_file(), (name, filename) - - -def test_stable_factory_resolves_every_deploy_config() -> None: - pytest.importorskip("vllm_omni") - from vllm_omni.config.config_factory import StageConfigFactory - from vllm_omni.config.pipeline_registry import resolve_pipeline_config - from vllm_omni.config.stage_config import load_deploy_config - - from unirl.rollout.engine.vllm_omni.pipeline_configs import register_unirl_pipeline_configs - - register_unirl_pipeline_configs() - for path in sorted(DEPLOY_DIR.glob("*.yaml")): - deploy = load_deploy_config(path) - pipeline = resolve_pipeline_config(deploy.pipeline) - assert pipeline is not None, path - stages, _ = StageConfigFactory._create_legacy_from_registry( - pipeline, - {}, - str(path), - deploy, - ) - assert stages, path - assert [stage.stage_id for stage in stages] == list(range(len(stages))), path - - -def test_hi3_text_topologies_preserve_legacy_output_contract() -> None: - pytest.importorskip("vllm_omni") - from unirl.rollout.engine.vllm_omni.pipeline_configs import ( - UNIRL_HI3_AR_MULTIMODAL_TEXT, - UNIRL_HI3_AR_RECAPTION, - UNIRL_HI3_AR_TEXT, - UNIRL_QWEN3_OMNI_THINKER, - ) - - text_stage = UNIRL_HI3_AR_TEXT.stages[0] - multimodal_stage = UNIRL_HI3_AR_MULTIMODAL_TEXT.stages[0] - for stage in (text_stage, multimodal_stage): - assert stage.final_output_type == "text" - assert stage.engine_output_type == "text" - assert stage.owns_tokenizer is True - assert text_stage.requires_multimodal_data is False - assert multimodal_stage.requires_multimodal_data is True - - recaption_stage = UNIRL_HI3_AR_RECAPTION.stages[0] - assert recaption_stage.final_output_type == "text" - assert recaption_stage.engine_output_type == "latent" - assert recaption_stage.owns_tokenizer is True - assert recaption_stage.requires_multimodal_data is False - - qwen_stage = UNIRL_QWEN3_OMNI_THINKER.stages[0] - assert qwen_stage.model_arch == "Qwen3OmniMoeThinkerForConditionalGeneration" - assert qwen_stage.engine_output_type == "text" - - -def test_hi3_t2i_does_not_forward_at_raw_prefill_output() -> None: - config = yaml.safe_load((DEPLOY_DIR / "hunyuan_image3_t2i_rl.yaml").read_text()) - stage_zero = config["stages"][0] - assert stage_zero["omni_kv_config"] == {"need_send_cache": True} - - -@pytest.mark.parametrize( - "name", - [ - "hunyuan_image3_ar_rl.yaml", - "hunyuan_image3_i2t_rl.yaml", - "hunyuan_image3_it2i_rl.yaml", - ], -) -def test_hi3_ar_lora_stages_force_mp_weight_sync_workers(name: str) -> None: - config = yaml.safe_load((DEPLOY_DIR / name).read_text()) - stage_zero = config["stages"][0] - assert stage_zero["enable_lora"] is True - assert stage_zero["worker_extension_cls"].endswith(".HI3ARWeightSyncExtension") - assert stage_zero["distributed_executor_backend"] == "mp" diff --git a/tests/rollout/vllm_omni/test_stable_installed_contracts.py b/tests/rollout/vllm_omni/test_stable_installed_contracts.py deleted file mode 100644 index e1e1e91c9..000000000 --- a/tests/rollout/vllm_omni/test_stable_installed_contracts.py +++ /dev/null @@ -1,126 +0,0 @@ -from __future__ import annotations - -import importlib -import os -from importlib.metadata import version - -import pytest - -if os.environ.get("VLLM_OMNI_REQUIRE_INSTALLED") == "1": - importlib.import_module("vllm") - importlib.import_module("vllm_omni") -else: - pytest.importorskip("vllm") - pytest.importorskip("vllm_omni") - - -def test_exact_stable_runtime_is_installed() -> None: - assert version("vllm") == "0.28.0" - assert version("vllm-omni") == "0.28.0" - assert version("transformers") == "5.12.1" - assert version("kernels") == "0.14.1" - assert version("diffusers") == "0.40.0" - - -def test_real_stable_output_uses_flat_completion_fields() -> None: - from vllm.outputs import CompletionOutput - from vllm_omni.outputs import OmniRequestOutput - - from unirl.rollout.engine.vllm_omni.utils.tracks import _extract_completion, decoded_text_from_ar - - completion = CompletionOutput( - index=0, - text="answer", - token_ids=[1, 2], - cumulative_logprob=None, - logprobs=None, - finish_reason="stop", - ) - output = OmniRequestOutput( - request_id="0", - stage_id=0, - final_output_type="text", - outputs=[completion], - ) - - assert decoded_text_from_ar([[output]]).texts == ["answer"] - assert _extract_completion(output) == ([1, 2], None) - - -@pytest.mark.parametrize( - "module_name", - [ - "unirl.rollout.engine.vllm_omni.pipelines.sd3.pipeline", - "unirl.rollout.engine.vllm_omni.pipelines.qwen_image.pipeline", - "unirl.rollout.engine.vllm_omni.pipelines.hv15.pipeline", - "unirl.rollout.engine.vllm_omni.pipelines.hi3.pipeline", - "unirl.rollout.engine.vllm_omni.pipelines.bagel.pipeline", - "unirl.rollout.engine.vllm_omni.worker.dit_extension", - "unirl.rollout.engine.vllm_omni.worker.ar_extension", - ], -) -def test_custom_pipeline_and_worker_modules_import(module_name: str) -> None: - importlib.import_module(module_name) - - -def test_capture_plugin_patches_both_formatter_bindings() -> None: - from vllm_omni.diffusion import diffusion_engine, output_formatter - - from unirl.rollout.engine.vllm_omni.plugin import register_capture_flush - - register_capture_flush() - patched = output_formatter.format_diffusion_outputs - assert patched is diffusion_engine.format_diffusion_outputs - - register_capture_flush() - assert output_formatter.format_diffusion_outputs is patched - - -def test_diffusion_lora_loader_delegates_non_tensor_requests(monkeypatch: pytest.MonkeyPatch) -> None: - from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager - - from unirl.rollout.engine.vllm_omni.patches.runtime import patch_dit_lora_loader - - request = object() - expected = object() - calls = [] - - def original(manager, received): - calls.append((manager, received)) - return expected - - monkeypatch.setattr(DiffusionLoRAManager, "_load_adapter", original) - patch_dit_lora_loader() - - manager = object() - assert DiffusionLoRAManager._load_adapter(manager, request) is expected - assert calls == [(manager, request)] - - -def test_hi3_expert_mapping_tuple_is_unwrapped(monkeypatch: pytest.MonkeyPatch) -> None: - from vllm.model_executor import utils as vllm_utils - - from unirl.rollout.engine.vllm_omni.patches import compat_hi3_lora - - mapping = [("experts.0", "weight", 0, "w")] - monkeypatch.setattr(vllm_utils, "get_moe_expert_mapping", lambda _model: (mapping, {"old": "new"})) - monkeypatch.setattr(compat_hi3_lora, "_INSTALLED", False) - - compat_hi3_lora.install() - - assert vllm_utils.get_moe_expert_mapping(object()) == mapping - - -def test_diffusion_parameter_probe_reaches_wrapped_transformer() -> None: - import torch - - from unirl.rollout.engine.vllm_omni.worker.ipc_receive_mixin import BucketedIPCReceiveMixin - - worker = object.__new__(BucketedIPCReceiveMixin) - worker.worker = torch.nn.Module() - transformer = torch.nn.Linear(4, 3) - worker.model_runner = type("Runner", (), {"pipeline": type("Pipeline", (), {"transformer": transformer})()})() - - descriptions = worker._diffrl_describe_params() - - assert descriptions["weight"] == ((3, 4), "torch.float32") diff --git a/tests/rollout/vllm_omni/test_stable_runtime_contracts.py b/tests/rollout/vllm_omni/test_stable_runtime_contracts.py deleted file mode 100644 index b9741a118..000000000 --- a/tests/rollout/vllm_omni/test_stable_runtime_contracts.py +++ /dev/null @@ -1,150 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace - -import pytest -import torch - -from unirl.rollout.engine.vllm_omni import patches, plugin -from unirl.rollout.engine.vllm_omni.backends import native -from unirl.rollout.engine.vllm_omni.backends.native import VLLMOmniBackend -from unirl.rollout.engine.vllm_omni.utils.tracks import _extract_completion, decoded_text_from_ar - - -class _Task: - def __init__(self, *, task_id: str, level: int | None = None, tags=None) -> None: - self.task_id = task_id - self.level = level - self.tags = tags - - -class _Engine: - num_stages = 2 - - def __init__(self) -> None: - self.calls: list[tuple[int, str, tuple, dict | None]] = [] - - def get_stage_metadata(self, stage_id: int): - return SimpleNamespace(stage_type="llm" if stage_id == 0 else "diffusion") - - def collective_rpc(self, *, method: str, args=(), kwargs=None, stage_ids=None): - stage_id = stage_ids[0] - self.calls.append((stage_id, method, args, kwargs)) - if method in {"handle_sleep_task", "handle_wake_task"}: - task = args[0] - return [ - { - "status": "SUCCESS", - "stage_id": stage_id, - "task_id": task.task_id, - "rank": 0, - } - ] - return [True] - - -def _backend(engine: _Engine) -> VLLMOmniBackend: - runtime = { - "OmniSleepTask": _Task, - "OmniWakeTask": _Task, - } - return VLLMOmniBackend( - SimpleNamespace(engine=engine), - runtime, - tokenizer=None, - tp_per_stage={0: 1, 1: 1}, - ) - - -def test_stable_outputs_are_read_directly() -> None: - completion = SimpleNamespace(text="answer", token_ids=[1, 2], logprobs=None) - output = SimpleNamespace(stage_id=0, final_output_type="text", outputs=[completion]) - - texts = decoded_text_from_ar([[output]]) - tokens, logprobs = _extract_completion(output) - - assert texts.texts == ["answer"] - assert tokens == [1, 2] - assert logprobs is None - - -def test_sleep_and_wake_route_by_stage_type(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(torch.cuda, "is_available", lambda: False) - engine = _Engine() - backend = _backend(engine) - - backend.sleep_task() - backend.wake_task() - - assert [(stage, method) for stage, method, _args, _kwargs in engine.calls] == [ - (0, "sleep"), - (1, "handle_sleep_task"), - (0, "wake_up"), - (1, "handle_wake_task"), - ] - assert engine.calls[0][2] == (1, "abort") - assert engine.calls[2][3] == {"tags": None} - - -def test_collective_rpc_rejects_unsupported_stage() -> None: - backend = _backend(_Engine()) - with pytest.raises(RuntimeError, match="not installed"): - backend._require_rpc_success( - "update_weights", - 1, - [{"supported": False, "error": "worker extension is not installed"}], - ) - - -def test_collective_rpc_rejects_false_except_idempotent_lora_removal() -> None: - backend = _backend(_Engine()) - with pytest.raises(RuntimeError, match="returned False"): - backend._require_rpc_success("set_lora_from_tensor_dict", 0, [False]) - with pytest.raises(RuntimeError, match="unexpected result"): - backend._require_rpc_success("update_weights", 0, ["failed"]) - - backend._require_rpc_success("remove_lora", 0, [False], allow_false=True) - - -def test_boot_passes_stable_deploy_config_and_reads_engine_stages(monkeypatch: pytest.MonkeyPatch) -> None: - captured: dict = {} - - class Omni: - def __init__(self, **kwargs) -> None: - captured.update(kwargs) - self.engine = SimpleNamespace( - num_stages=1, - stage_configs=[ - { - "stage_id": 0, - "engine_args": {"tensor_parallel_size": 2}, - } - ], - ) - - monkeypatch.setattr(patches, "install", lambda: None) - monkeypatch.setattr(plugin, "register_unirl_runtime", lambda: None) - monkeypatch.setattr( - native, - "_import_omni_runtime", - lambda: { - "Omni": Omni, - "OmniSleepTask": _Task, - "OmniWakeTask": _Task, - }, - ) - monkeypatch.setattr(torch.cuda, "is_available", lambda: False) - - backend = VLLMOmniBackend.boot( - { - "model_path": "unused", - "deploy_config": "sd35_t2i_rl.yaml", - "enable_sleep_mode": False, - "ports": None, - } - ) - - assert captured["model"] == "unused" - assert captured["deploy_config"].endswith("deploy_configs/sd35_t2i_rl.yaml") - assert "stage_configs_path" not in captured - assert backend.tp_per_stage() == {0: 2} From 15b801f9f19636a2bc0271da14a41e576e912116 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Mon, 21 Sep 2026 11:53:25 +0800 Subject: [PATCH 19/23] fix(rollout): support direct vllm on 0.28 --- unirl/rollout/engine/vllm/runtime.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unirl/rollout/engine/vllm/runtime.py b/unirl/rollout/engine/vllm/runtime.py index 3a4726d37..05e965242 100644 --- a/unirl/rollout/engine/vllm/runtime.py +++ b/unirl/rollout/engine/vllm/runtime.py @@ -8,7 +8,7 @@ from typing import Any, Dict, List _PROTOCOL_VERSION = 1 -_SUPPORTED_VLLM_VERSION = "0.27.0" +_SUPPORTED_VLLM_VERSION = "0.28.0" class _ProtocolError(RuntimeError): From 78014c6ff05bcb23e398fd694aaa2113fb3a1218 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Mon, 21 Sep 2026 13:49:52 +0800 Subject: [PATCH 20/23] refactor(rollout): simplify vllm-omni backend helpers --- .../engine/vllm_omni/backends/native.py | 38 ++++++------------- 1 file changed, 11 insertions(+), 27 deletions(-) diff --git a/unirl/rollout/engine/vllm_omni/backends/native.py b/unirl/rollout/engine/vllm_omni/backends/native.py index c1fb2ba73..b79f94b67 100644 --- a/unirl/rollout/engine/vllm_omni/backends/native.py +++ b/unirl/rollout/engine/vllm_omni/backends/native.py @@ -63,36 +63,20 @@ def _resolve_deploy_config(name: str) -> str: return path -def _cfg_get(cfg: Any, key: str, default: Any = None) -> Any: - """Mapping/OmegaConf/attr-tolerant getter (``None`` coerces to default).""" - if cfg is None: - return default - getter = getattr(cfg, "get", None) - value = getter(key, default) if callable(getter) else getattr(cfg, key, default) - return default if value is None else value - - def _tp_from_stage_configs(stage_configs: Sequence[Any]) -> Dict[int, int]: """Extract ``{stage_id: tensor_parallel_size}`` from the runtime's configs.""" tp_map: Dict[int, int] = {} for entry in stage_configs: - sid = int(_cfg_get(entry, "stage_id", len(tp_map))) - ea = _cfg_get(entry, "engine_args", {}) - tp = _cfg_get(ea, "tensor_parallel_size") + engine_args = entry.engine_args + tp = engine_args.get("tensor_parallel_size") if tp is None: - tp = _cfg_get(_cfg_get(ea, "parallel_config", {}), "tensor_parallel_size") - tp_map[sid] = int(tp) if tp is not None else 1 + parallel_config = engine_args.get("parallel_config") + if parallel_config is not None: + tp = parallel_config.get("tensor_parallel_size") + tp_map[int(entry.stage_id)] = int(tp) if tp is not None else 1 return tp_map -def _assemble_omni_kwargs(intent: Dict[str, Any]) -> Dict[str, Any]: - """Spell the boot intent into ``Omni`` ctor kwargs.""" - omni_kwargs = dict(intent.get("omni_kwargs") or {}) - if intent.get("enable_sleep_mode"): - omni_kwargs["enable_sleep_mode"] = True - return omni_kwargs - - @contextmanager def _master_port_env(port: Optional[int]): """Pin ``MASTER_PORT`` to the reserved engine port for the duration of ``Omni()``.""" @@ -162,7 +146,9 @@ def boot(cls, intent: Dict[str, Any]) -> "VLLMOmniBackend": pass deploy_config_path = _resolve_deploy_config(str(intent["deploy_config"])) - omni_kwargs = _assemble_omni_kwargs(intent) + omni_kwargs = dict(intent.get("omni_kwargs") or {}) + if intent.get("enable_sleep_mode"): + omni_kwargs["enable_sleep_mode"] = True ports = intent.get("ports") boot_master_port = int(ports.master_port) if ports is not None else None logger.info( @@ -369,9 +355,6 @@ def _require_ack_success(action: str, stage_id: int, task_id: str, acks: object) def ack_field(ack: object, name: str, default: object = None) -> object: return ack.get(name, default) if isinstance(ack, Mapping) else getattr(ack, name, default) - def ack_error(ack: object) -> object: - return ack_field(ack, "error_msg", ack_field(ack, "error")) - def validate_result(result: object) -> None: nonlocal success_count if result is None: @@ -387,10 +370,11 @@ def validate_result(result: object) -> None: if status is None: raise RuntimeError(f"vllm-omni {action} returned no ACK status for stage {stage_id}: result={result!r}") if status != "SUCCESS": + error = ack_field(result, "error_msg", ack_field(result, "error")) raise RuntimeError( f"vllm-omni {action} failed on stage {stage_id}: worker rank " f"{ack_field(result, 'rank', '?')} answered status={status!r} " - f"error={ack_error(result)!r}" + f"error={error!r}" ) ack_stage_id = ack_field(result, "stage_id") ack_task_id = ack_field(result, "task_id") From a4e46fff0039b227c1e82ecb905d39d188d27233 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Mon, 21 Sep 2026 14:29:16 +0800 Subject: [PATCH 21/23] refactor(rollout): tighten vllm-omni backend boundaries --- .../rollout/engine/vllm_omni/backends/base.py | 2 - .../engine/vllm_omni/backends/native.py | 83 +++++++++---------- unirl/utils/peft_merge.py | 7 +- 3 files changed, 43 insertions(+), 49 deletions(-) diff --git a/unirl/rollout/engine/vllm_omni/backends/base.py b/unirl/rollout/engine/vllm_omni/backends/base.py index 4397c9dff..0b0270829 100644 --- a/unirl/rollout/engine/vllm_omni/backends/base.py +++ b/unirl/rollout/engine/vllm_omni/backends/base.py @@ -11,7 +11,6 @@ Optional, Protocol, Sequence, - runtime_checkable, ) if TYPE_CHECKING: @@ -70,7 +69,6 @@ class OmniRawResult(Protocol): multimodal_output: Optional[dict] -@runtime_checkable class Backend(Protocol): """The seam every ``vllm_omni`` collaborator reaches the runtime through.""" diff --git a/unirl/rollout/engine/vllm_omni/backends/native.py b/unirl/rollout/engine/vllm_omni/backends/native.py index b79f94b67..e085de514 100644 --- a/unirl/rollout/engine/vllm_omni/backends/native.py +++ b/unirl/rollout/engine/vllm_omni/backends/native.py @@ -63,20 +63,6 @@ def _resolve_deploy_config(name: str) -> str: return path -def _tp_from_stage_configs(stage_configs: Sequence[Any]) -> Dict[int, int]: - """Extract ``{stage_id: tensor_parallel_size}`` from the runtime's configs.""" - tp_map: Dict[int, int] = {} - for entry in stage_configs: - engine_args = entry.engine_args - tp = engine_args.get("tensor_parallel_size") - if tp is None: - parallel_config = engine_args.get("parallel_config") - if parallel_config is not None: - tp = parallel_config.get("tensor_parallel_size") - tp_map[int(entry.stage_id)] = int(tp) if tp is not None else 1 - return tp_map - - @contextmanager def _master_port_env(port: Optional[int]): """Pin ``MASTER_PORT`` to the reserved engine port for the duration of ``Omni()``.""" @@ -104,12 +90,10 @@ def __init__( runtime: Dict[str, Any], *, tokenizer: Optional[Any], - tp_per_stage: Dict[int, int], ) -> None: self._omni: Optional[Any] = omni self._rt = runtime self._tokenizer = tokenizer - self._tp_per_stage = dict(tp_per_stage) @classmethod def boot(cls, intent: Dict[str, Any]) -> "VLLMOmniBackend": @@ -204,15 +188,12 @@ def boot(cls, intent: Dict[str, Any]) -> "VLLMOmniBackend": omni, rt, tokenizer=tokenizer, - tp_per_stage=_tp_from_stage_configs(omni.engine.stage_configs), ) except BaseException: logger.exception("VLLM-Omni boot failed; tearing down any engine processes") if omni is not None: try: - close = getattr(omni, "close", None) - if callable(close): - close() + omni.close() except Exception: logger.exception("Failed to close the half-booted vLLM-Omni engine") terminate_descendants(os.getpid(), name_prefix=_ENGINE_PROC_PREFIX) @@ -285,7 +266,16 @@ def num_stages(self) -> int: return int(self._require_omni().engine.num_stages) def tp_per_stage(self) -> Dict[int, int]: - return dict(self._tp_per_stage) + tp_map: Dict[int, int] = {} + for entry in self._require_omni().engine.stage_configs: + engine_args = entry.engine_args + tp = engine_args.get("tensor_parallel_size") + if tp is None: + parallel_config = engine_args.get("parallel_config") + if parallel_config is not None: + tp = parallel_config.get("tensor_parallel_size") + tp_map[int(entry.stage_id)] = int(tp) if tp is not None else 1 + return tp_map def _stage_ids(self) -> List[int]: return list(range(self.num_stages())) @@ -431,14 +421,19 @@ def wake_task(self) -> None: torch.cuda.synchronize() def ping(self) -> bool: - return self._omni is not None + omni = self._omni + if omni is None: + return False + try: + return bool(omni.is_running) + except Exception: + logger.exception("vLLM-Omni health check failed") + return False def shutdown(self) -> None: if self._omni is not None: try: - close = getattr(self._omni, "close", None) - if callable(close): - close() + self._omni.close() finally: self._omni = None @@ -565,8 +560,9 @@ def set_lora_handle( DIFFRL_LORA_NAME, DIFFRL_LORA_PATH, ) + from unirl.utils.peft_merge import adapt_lora_for_vllm - lora_tensors = self._wrap_peft_envelope(lora_tensors) + lora_tensors = adapt_lora_for_vllm(lora_tensors) self._remove_existing_lora(int(DIFFRL_LORA_INT_ID)) from unirl.distributed.weight_sync.transfer.sgl_compat import ( @@ -608,8 +604,9 @@ def set_lora_copy( DIFFRL_LORA_NAME, DIFFRL_LORA_PATH, ) + from unirl.utils.peft_merge import adapt_lora_for_vllm - lora_tensors = self._wrap_peft_envelope(lora_tensors) + lora_tensors = adapt_lora_for_vllm(lora_tensors) self._remove_existing_lora(int(DIFFRL_LORA_INT_ID)) cpu_tensors = { @@ -632,16 +629,6 @@ def set_lora_copy( ), ) - @staticmethod - def _wrap_peft_envelope(lora_tensors: Dict[str, Any]) -> Dict[str, Any]: - """Wrap canonical wire keys in the PEFT envelope vllm-omni expects.""" - from unirl.utils.peft_merge import adapt_lora_for_vllm - - first_key = next(iter(lora_tensors), "") - if lora_tensors and not first_key.startswith("base_model.model."): - return adapt_lora_for_vllm(lora_tensors) - return lora_tensors - def _remove_existing_lora(self, adapter_id: int) -> None: """Drop the existing adapter on every stage before re-adding.""" for sid in self._stage_ids(): @@ -681,15 +668,19 @@ def _group_by_request(flat_outputs: Sequence[Any], n: int) -> List[List[Any]]: """Group ``Omni.generate``'s flat output list into per-request lists.""" grouped: List[List[Any]] = [[] for _ in range(n)] for out in flat_outputs: - rid = getattr(out, "request_id", "") or "" - if "_" in rid: - idx_part = rid.split("_", 1)[0] - try: - idx = int(idx_part) - except ValueError: - continue - if 0 <= idx < n: - grouped[idx].append(out) + try: + rid = out.request_id + idx_part, suffix = rid.split("_", 1) + if not suffix: + raise ValueError("empty request-id suffix") + idx = int(idx_part) + except (AttributeError, TypeError, ValueError) as exc: + raise RuntimeError( + f"vllm-omni returned an invalid request_id: {getattr(out, 'request_id', None)!r}" + ) from exc + if not 0 <= idx < n: + raise RuntimeError(f"vllm-omni request_id index {idx} is outside the prompt batch of size {n}") + grouped[idx].append(out) return grouped diff --git a/unirl/utils/peft_merge.py b/unirl/utils/peft_merge.py index 6aef50419..c7e273ec4 100644 --- a/unirl/utils/peft_merge.py +++ b/unirl/utils/peft_merge.py @@ -346,7 +346,12 @@ def extract_lora_tensors( def adapt_lora_for_vllm(tensors: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Wrap canonical-format LoRA keys in the vllm-omni PEFT envelope.""" + """Idempotently wrap canonical LoRA keys in the vllm-omni PEFT envelope.""" + enveloped = [name.startswith(_PEFT_PREFIX) for name in tensors] + if all(enveloped): + return tensors + if any(enveloped): + raise ValueError("LoRA tensors mix canonical and PEFT-enveloped keys") return {f"{_PEFT_PREFIX}{k}": v for k, v in tensors.items()} From 775de2741300bdd65f39723cbb16729f0f13ccd5 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Wed, 23 Sep 2026 13:43:16 +0800 Subject: [PATCH 22/23] docs(install): align extras notes with the 0.28 pins --- INSTALL.md | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/INSTALL.md b/INSTALL.md index 8a27c831d..b1bb21f9f 100644 --- a/INSTALL.md +++ b/INSTALL.md @@ -58,7 +58,7 @@ uv pip install -e ".[sglang,train,infer]" | `sglang` | `sglang[diffusion]`, `checkpoint-engine`, `flash-attn-4`, `flash-linear-attention[conv1d]`, torch +cu130 stack, PyAV | SGLang-based AR/VLM and diffusion recipes | | `fastvideo` | FastVideo pinned to an upstream Git commit | WAN 2.1 / 2.2 rollout; [the extra does not currently resolve](#fastvideo-installation-blocker) | | `train` | `wandb`, `aiohttp`, `math-verify` | Training runs and local math-answer scoring | -| `cosmos3` | `diffusers>=0.39` | [Cosmos3 SFT](unirl/models/cosmos3/README.md); apply the [version constraint](#cosmos3-version-prerequisite) | +| `cosmos3` | `diffusers>=0.39` | [Cosmos3 SFT](unirl/models/cosmos3/README.md); uv's `diffusers==0.40.0` override already satisfies this extra | | `infer` | `accelerate`, `timm` | HunyuanImage3, Janus-Pro, and similar models | | `eval` | `torchvision`, `paddlepaddle`, `paddleocr`, `python-Levenshtein` | OCR-based reward components | | `veomni` | `veomni` | Recipes using the [VeOmni training backend](unirl/train/backend/veomni/) | @@ -89,24 +89,19 @@ Prefer these extras over the legacy [`requirements.txt`](requirements.txt) and The `fastvideo` extra pins [hao-ai-lab/FastVideo@2095477](https://github.com/hao-ai-lab/FastVideo/blob/2095477eac7e289c7a7ab13acb367ca60687c304/pyproject.toml), which requires `transformers==4.57.3` and `wandb>=0.21.0`. The transformers pin -conflicts with UniRL's `transformers>=5.6,<5.7`, so `.[fastvideo]` does not +conflicts with UniRL's `transformers>=5.6,<5.13`, so `.[fastvideo]` does not resolve — a separate venv does not help, because UniRL's base deps still apply. Adding `train` also conflicts on `wandb`. Use `$FASTVIDEO_PATH` as in the [FastVideo engine README](unirl/rollout/engine/fastvideo/README.md) until the extra is solvable. -### Cosmos3 version prerequisite +### Cosmos3 -`cosmos3` asks for `diffusers>=0.39`, but uv's override `diffusers>=0.38.0` -[replaces](https://docs.astral.sh/uv/concepts/resolution/#dependency-overrides) -that floor instead of intersecting with it. Include `cosmos3` in the extras and -pass `--constraint` when installing: +`cosmos3` asks for `diffusers>=0.39`. The uv override pins `diffusers==0.40.0`, +which already satisfies that floor, so install it as a normal extra: ```bash -COSMOS3_CONSTRAINTS="$(mktemp)" -printf '%s\n' 'diffusers>=0.39' > "$COSMOS3_CONSTRAINTS" -uv pip install -e ".[vllm,train,infer,cosmos3]" \ - --constraint "$COSMOS3_CONSTRAINTS" +uv pip install -e ".[vllm,train,infer,cosmos3]" ``` ## Environment From c23fd66ce978b811de2e6b40225f2bed0bbdc870 Mon Sep 17 00:00:00 2001 From: CjhHa1 Date: Thu, 24 Sep 2026 09:47:14 +0800 Subject: [PATCH 23/23] refactor(rollout): tidy vllm-omni 0.28 docs and boundaries - Describe the split AR/diffusion sleep-wake path in engine docstrings. - Drop the stale sglang-unification rationale from the transformers pin. - Access vLLM 0.28 workspace-lane state directly; negative lanes are rejected by use_workspace_lane itself. - Drop a redundant str() on stage metadata and read stage_configs once. - Rename stage YAML/config references in examples to deploy config. - Hoist distributed_executor_backend to the top level in HI3 deploy configs; resolved stage configs are unchanged. --- .../bagel/bagel_it2i_managed_editscore.yaml | 2 +- examples/diffusion/bagel/bagel_it2i_vllmomni.yaml | 2 +- .../hunyuan_video15_t2v_vllmomni_colocate.yaml | 2 +- .../diffusion/sd3/sd3_vllmomni_lora_separate.yaml | 4 ++-- examples/unified_model/hi3_vllmomni.yaml | 2 +- examples/unified_model/hi3_vllmomni_veomni_ep.yaml | 2 +- pyproject.toml | 3 +-- unirl/rollout/engine/vllm_omni/backends/native.py | 7 +++---- .../deploy_configs/hunyuan_image3_ar_rl.yaml | 2 +- .../deploy_configs/hunyuan_image3_i2t_rl.yaml | 2 +- .../deploy_configs/hunyuan_image3_it2i_rl.yaml | 3 +-- .../deploy_configs/hunyuan_image3_t2i_rl.yaml | 3 +-- unirl/rollout/engine/vllm_omni/engine.py | 4 ++-- .../vllm_omni/patches/compat_moe_workspace.py | 13 ++++--------- 14 files changed, 21 insertions(+), 30 deletions(-) diff --git a/examples/diffusion/bagel/bagel_it2i_managed_editscore.yaml b/examples/diffusion/bagel/bagel_it2i_managed_editscore.yaml index e06c2f087..c7d56e45f 100644 --- a/examples/diffusion/bagel/bagel_it2i_managed_editscore.yaml +++ b/examples/diffusion/bagel/bagel_it2i_managed_editscore.yaml @@ -137,7 +137,7 @@ backend: total_steps: 100000 lora_cfg: _target_: unirl.train.configs.LoraConfig - rank: 64 # must be <= the stage YAML's max_lora_rank (64) + rank: 64 # must be <= the deploy config's max_lora_rank (64) alpha: 128 dropout: 0.0 bias: none diff --git a/examples/diffusion/bagel/bagel_it2i_vllmomni.yaml b/examples/diffusion/bagel/bagel_it2i_vllmomni.yaml index 723f85f1f..bce22115e 100644 --- a/examples/diffusion/bagel/bagel_it2i_vllmomni.yaml +++ b/examples/diffusion/bagel/bagel_it2i_vllmomni.yaml @@ -134,7 +134,7 @@ backend: total_steps: 100000 lora_cfg: _target_: unirl.train.configs.LoraConfig - rank: 64 # must be <= the stage YAML's max_lora_rank (64) + rank: 64 # must be <= the deploy config's max_lora_rank (64) alpha: 128 dropout: 0.0 bias: none diff --git a/examples/diffusion/hunyuan_video15/hunyuan_video15_t2v_vllmomni_colocate.yaml b/examples/diffusion/hunyuan_video15/hunyuan_video15_t2v_vllmomni_colocate.yaml index 8a839f6a5..5c7a16aac 100644 --- a/examples/diffusion/hunyuan_video15/hunyuan_video15_t2v_vllmomni_colocate.yaml +++ b/examples/diffusion/hunyuan_video15/hunyuan_video15_t2v_vllmomni_colocate.yaml @@ -20,7 +20,7 @@ # sync (sd3/qwen v2 use it). Pushes the trained LoRA adapter into the # co-located sibling engine in-process; the engine runs base + adapter, # which is mathematically the merged model the separate recipe pushed. -# hv15's stage config sets enable_lora/max_lora_rank=64 so the adapter +# hv15's deploy config sets enable_lora/max_lora_rank=64 so the adapter # actually applies. This avoids the CUDA-IPC path's SGLang dependency, # which the vllm-omni-only venv (two-venv image) does not provide. # diff --git a/examples/diffusion/sd3/sd3_vllmomni_lora_separate.yaml b/examples/diffusion/sd3/sd3_vllmomni_lora_separate.yaml index 782b140ad..36a5e5521 100644 --- a/examples/diffusion/sd3/sd3_vllmomni_lora_separate.yaml +++ b/examples/diffusion/sd3/sd3_vllmomni_lora_separate.yaml @@ -6,7 +6,7 @@ # - LoRA training (lora_cfg kept), pushed as the bare ADAPTER each sync # (RemoteLoraWeightSync ships lora_A/lora_B to the engine's set_lora_from_tensors). # - No NCCL group / merged-model broadcast: rank 0 ships the few-MB adapter over a -# plain Ray RPC to each rollout Worker. The vLLM-Omni sd35_t2i stage config +# plain Ray RPC to each rollout Worker. The vLLM-Omni sd35_t2i deploy config # enables LoRA (enable_lora: true, max_lora_rank: 32), so the engine accepts it. num_devices: 8 @@ -87,7 +87,7 @@ rollout: config: _target_: unirl.rollout.engine.vllm_omni.config.VLLMOmniEngineConfig model_path: ${oc.env:PRETRAINED_MODEL,stabilityai/stable-diffusion-3.5-medium} - # sd35_t2i stage config enables LoRA (enable_lora: true, max_lora_rank: 32), + # sd35_t2i deploy config enables LoRA (enable_lora: true, max_lora_rank: 32), # so the engine accepts the adapter pushed by RemoteLoraWeightSync. modality: sd3_t2i # Separate slabs do not time-share GPUs, so sleep/wake is unnecessary. diff --git a/examples/unified_model/hi3_vllmomni.yaml b/examples/unified_model/hi3_vllmomni.yaml index bc3398763..be6666cf4 100644 --- a/examples/unified_model/hi3_vllmomni.yaml +++ b/examples/unified_model/hi3_vllmomni.yaml @@ -118,7 +118,7 @@ backend: # texts per prompt; the DiT engine (modality dit_recaption, single diffusion # stage, GPUs 4-7) renders M distinct-noise images per recaption. UnifiedModelTrainer # wires both with remote() in the shared placement; each engine clears -# CUDA_VISIBLE_DEVICES for its multi-GPU HI3 modality and its stage YAML's +# CUDA_VISIBLE_DEVICES for its multi-GPU HI3 modality and its deploy config's # runtime.devices pins the physical cards (a real partition, NOT anchor+pop). # Both run enable_sleep_mode so the trainer can sleep/wake them around the # colocate train phase. The two engines share ONE backbone/LoRA via sync. diff --git a/examples/unified_model/hi3_vllmomni_veomni_ep.yaml b/examples/unified_model/hi3_vllmomni_veomni_ep.yaml index fcd710c79..f88a9c001 100644 --- a/examples/unified_model/hi3_vllmomni_veomni_ep.yaml +++ b/examples/unified_model/hi3_vllmomni_veomni_ep.yaml @@ -124,7 +124,7 @@ backend: # texts per prompt; the DiT engine (modality dit_recaption, single diffusion # stage, GPUs 4-7) renders M distinct-noise images per recaption. UnifiedModelTrainer # wires both with remote() in the shared placement; each engine clears -# CUDA_VISIBLE_DEVICES for its multi-GPU HI3 modality and its stage YAML's +# CUDA_VISIBLE_DEVICES for its multi-GPU HI3 modality and its deploy config's # runtime.devices pins the physical cards (a real partition, NOT anchor+pop). # Both run enable_sleep_mode so the trainer can sleep/wake them around the # colocate train phase. The two engines share ONE backbone/LoRA via sync. diff --git a/pyproject.toml b/pyproject.toml index 1896ce081..f94eae7cc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -86,8 +86,7 @@ vllm = [ "vllm==0.28.0 ; sys_platform == 'linux'", "vllm-omni==0.28.0 ; sys_platform == 'linux'", # vllm-omni 0.28 needs >=5.10.1,<5.15 (5.10.0 is yanked, 5.15 has a known - # construction regression); exact because there is no uv.lock to hold it, and - # 5.12.1 is what sglang 0.5.15+ pins, so unifying later is a bump not a port. + # construction regression); exact because there is no uv.lock to hold it. "transformers==5.12.1 ; sys_platform == 'linux'", "torch==2.13.0+cu130 ; sys_platform == 'linux'", "torchvision==0.28.0+cu130 ; sys_platform == 'linux'", diff --git a/unirl/rollout/engine/vllm_omni/backends/native.py b/unirl/rollout/engine/vllm_omni/backends/native.py index e085de514..3f179e2f5 100644 --- a/unirl/rollout/engine/vllm_omni/backends/native.py +++ b/unirl/rollout/engine/vllm_omni/backends/native.py @@ -165,16 +165,16 @@ def boot(cls, intent: Dict[str, Any]) -> "VLLMOmniBackend": fcntl.flock(lock_file, fcntl.LOCK_UN) lock_file.close() + stage_configs = omni.engine.stage_configs try: from omegaconf import OmegaConf - stage_configs = omni.engine.stage_configs resolved_stage_configs = OmegaConf.to_container( OmegaConf.create(stage_configs), resolve=True, ) except Exception: # noqa: BLE001 - config logging must never block boot - resolved_stage_configs = omni.engine.stage_configs + resolved_stage_configs = stage_configs logger.info( "VLLM-Omni resolved runtime stage configs (after all overrides):\n%s", pformat(resolved_stage_configs, sort_dicts=True), @@ -281,8 +281,7 @@ def _stage_ids(self) -> List[int]: return list(range(self.num_stages())) def _stage_type(self, stage_id: int) -> str: - metadata = self._require_omni().engine.get_stage_metadata(stage_id) - return str(metadata.stage_type) + return self._require_omni().engine.get_stage_metadata(stage_id).stage_type @staticmethod def _require_rpc_success( diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml index 9fbd4a294..4954d65e4 100644 --- a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_ar_rl.yaml @@ -1,6 +1,7 @@ # HunyuanImage-3 text-only AR deployment for the t2t adapter. pipeline: unirl_hi3_ar_text async_chunk: false +distributed_executor_backend: mp stages: - stage_id: 0 @@ -15,7 +16,6 @@ stages: max_num_batched_tokens: 32768 tensor_parallel_size: 4 pipeline_parallel_size: 1 - distributed_executor_backend: mp hf_overrides: rope_parameters: mrope_section: [0, 32, 32] diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml index 8b0098343..dd99ba799 100644 --- a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_i2t_rl.yaml @@ -1,6 +1,7 @@ # HunyuanImage-3 multimodal AR deployment for the i2t adapter. pipeline: unirl_hi3_ar_multimodal_text async_chunk: false +distributed_executor_backend: mp stages: - stage_id: 0 @@ -15,7 +16,6 @@ stages: max_num_batched_tokens: 32768 tensor_parallel_size: 4 pipeline_parallel_size: 1 - distributed_executor_backend: mp hf_overrides: rope_parameters: mrope_section: [0, 32, 32] diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_it2i_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_it2i_rl.yaml index 634e08989..17fdb9d9c 100644 --- a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_it2i_rl.yaml +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_it2i_rl.yaml @@ -1,6 +1,7 @@ # HunyuanImage-3 image-to-image deployment with UniRL worker extensions. pipeline: hunyuan_image_3_moe async_chunk: false +distributed_executor_backend: mp stages: - stage_id: 0 @@ -15,7 +16,6 @@ stages: max_num_batched_tokens: 32768 tensor_parallel_size: 4 pipeline_parallel_size: 1 - distributed_executor_backend: mp hf_overrides: rope_parameters: mrope_section: [0, 32, 32] @@ -33,7 +33,6 @@ stages: max_num_seqs: 1 enforce_eager: true trust_remote_code: true - distributed_executor_backend: mp parallel_config: tensor_parallel_size: 4 enable_expert_parallel: true diff --git a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_t2i_rl.yaml b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_t2i_rl.yaml index 7f7bf433c..20839b06f 100644 --- a/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_t2i_rl.yaml +++ b/unirl/rollout/engine/vllm_omni/deploy_configs/hunyuan_image3_t2i_rl.yaml @@ -1,6 +1,7 @@ # HunyuanImage-3 text-to-image deployment with UniRL worker extensions. pipeline: hunyuan_image_3_moe async_chunk: false +distributed_executor_backend: mp stages: - stage_id: 0 @@ -15,7 +16,6 @@ stages: max_num_batched_tokens: 32768 tensor_parallel_size: 4 pipeline_parallel_size: 1 - distributed_executor_backend: mp hf_overrides: rope_parameters: mrope_section: [0, 32, 32] @@ -35,7 +35,6 @@ stages: devices: "4,5,6,7" max_num_seqs: 1 enforce_eager: true - distributed_executor_backend: mp vae_use_slicing: false vae_use_tiling: false cache_backend: null diff --git a/unirl/rollout/engine/vllm_omni/engine.py b/unirl/rollout/engine/vllm_omni/engine.py index 194b38139..0033499b4 100644 --- a/unirl/rollout/engine/vllm_omni/engine.py +++ b/unirl/rollout/engine/vllm_omni/engine.py @@ -129,7 +129,7 @@ def _mark_consistently_offloaded(self) -> None: @distributed(dispatch_mode=Dispatch.BROADCAST) def sleep(self) -> None: - """Fan ``handle_sleep_task`` to every stage's workers (level 1).""" + """Sleep every stage at level 1 (AR via EngineCore, diffusion via worker task).""" if self._is_offloaded and not self._transition_failed: return try: @@ -145,7 +145,7 @@ def sleep(self) -> None: @distributed(dispatch_mode=Dispatch.BROADCAST) def wake_up(self) -> None: - """Fan ``handle_wake_task`` to every stage's workers + restore LoRA.""" + """Wake every stage (AR via EngineCore, diffusion via worker task) + restore LoRA.""" if self._transition_failed: # Recover an unknown partial stage state to one known boundary # before attempting another wake. If this retry fails, retain the diff --git a/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py b/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py index 4c5c8f6bc..2a58912d2 100644 --- a/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py +++ b/unirl/rollout/engine/vllm_omni/patches/compat_moe_workspace.py @@ -105,23 +105,18 @@ def patch_moe_workspace_pool() -> None: try: from vllm.device_allocator.cumem import CuMemAllocator from vllm.v1.worker.ubatching import dbo_current_ubatch_id - from vllm.v1.worker.workspace import WorkspaceManager + from vllm.v1.worker.workspace import WorkspaceManager, _workspace_lane WorkspaceManager._ensure_workspace_size except (ImportError, AttributeError): return - try: - from vllm.v1.worker.workspace import _workspace_lane - except ImportError: - _workspace_lane = None - def workspace_id_provider(manager: Any) -> int | None: # Delegate invalid lanes before indexing. A lane can otherwise flatten # onto the next ubatch's valid slot and release an unrelated workspace. - lane = 0 if _workspace_lane is None else _workspace_lane.get() - num_lanes = getattr(manager, "_num_lanes", 1) - if lane < 0 or lane >= num_lanes: + lane = _workspace_lane.get() + num_lanes = manager._num_lanes + if lane >= num_lanes: return None return dbo_current_ubatch_id() * num_lanes + lane