Repository navigation
feat(offloader): RFC #7204 canonical offload policy + CLI shortcuts - #8964
Asthenia0412 wants to merge 7 commits into
Conversation
Implements the core configuration unification from RFC #7204: - New --diffusion-offload {none,module,layer} CLI shortcut - New --diffusion-offload-config JSON with expanded schema (mode, components, pin_memory) - New --host-weight-runtime-config JSON (mode, root, registration_limit_gib) - CanonicalOffloadPolicy dataclass with single-source validation - Backward-compatible: all 10 legacy CLI flags continue to work unchanged - 45+ unit tests covering every RFC acceptance criterion Files changed: +vllm_omni/diffusion/offloader/config_parser.py (538 new lines) +tests/diffusion/offloader/test_offload_config_parser.py (484 new lines) ~vllm_omni/diffusion/offloader/base.py (+185/-53 lines refactor) ~vllm_omni/diffusion/data.py (+56 lines) ~vllm_omni/config/omni_config.py (+34 lines) ~vllm_omni/entrypoints/cli/serve.py (+44 lines) Resolves #7204 (Phase 1 + core canonical parser). Signed-off-by: Asthenia0412 <Asthenia0412@users.noreply.github.com>
- Convert isinstance(x, (A, B)) to isinstance(x, A | B) for Python 3.10+ union syntax - Run ruff format on new files - Pass ruff-check and ruff-format pre-commit hooks Signed-off-by: Asthenia0412 <Asthenia0412@users.noreply.github.com>
Local ruff 0.8.4 did not reformats assert statements that ruff 0.14.10 (CI-pinned version) rewrites into inline form. Run with correct version to pre-empt the pre-commit failure. Signed-off-by: Asthenia0412 <Asthenia0412@users.noreply.github.com>
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
|
This PR appears to belong to: docs/design/module/diffusion/index.md, docs/design/module/vllm_omni_config.md, docs/design/module/entrypoints.md, docs/design/module/host_weight_runtime.md, docs/design/module/diffusion/offloader.md. Module owners: @Isotr0py @david6666666 @princepride Routing: @Isotr0py via module of the changed files, module named in the PR description, CODEOWNERS; @david6666666 via module of the changed files, module named in the PR description; @princepride via module of the changed files, module named in the PR description @Asthenia0412, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer. Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment. |
Signed-off-by: Asthenia0412 <Asthenia0412@users.noreply.github.com>
Omni ReviewBot: supersededThe CI failure noted on |
Signed-off-by: Asthenia0412 <Asthenia0412@users.noreply.github.com>
Omni ReviewBot: supersededThe CI failure noted on |
PR introduced a duplicate registration of --diffusion-offload-config in serve.py (lines ~920 and ~1023). Argparse rejects conflicting option strings, causing mkdocs build to fail at generate_argparse.py startup. Remove the first (legacy) registration and keep the canonical RFC #7204 version added by this PR. Signed-off-by: Asthenia0412 <Asthenia0412@users.noreply.github.com>
Signed-off-by: Asthenia0412 <Asthenia0412@users.noreply.github.com>
3b12bae to
45f4179
Compare
Dong1017
left a comment
Dong1017
left a comment
There was a problem hiding this comment.
I reproduced this through the production config-to-backend path, not only the standalone parser. _DiffusionConfigProjection.from_kwargs followed by OffloadConfig.from_od_config resolves the old layerwise flag and old list-shaped config to LAYER_WISE, but diffusion_offload_shortcut='layer' still resolves to NONE. The new component-map schema raises TypeError in the old validator. Please wire the canonical parser and its resolved policy into that path while preserving the legacy inputs. This is a CPU config integration check; no model or device execution was required to reproduce the missing wiring.
| ) | ||
| omni_config_group.add_argument( | ||
| "--diffusion-offload-config", | ||
| dest="diffusion_offload_config", |
There was a problem hiding this comment.
These new inputs do not reach backend selection: _DiffusionConfigProjection.from_kwargs still calls the old list-shaped parser, and OffloadConfig.from_od_config still calls resolve_offload, which does not read diffusion_offload_shortcut. In a direct production-path check the new shortcut resolves to NONE, and the new component-map schema is rejected, while both legacy controls select LAYER_WISE. Please connect the canonical parser to this path and preserve the old input behavior.
4 participants
Background
vLLM-Omni offloader currently exposes 10 independent CLI flags with no unified parser, no JSON input, and no shortcut for common cases. Tracked by RFC #7204 — Simplify Diffusion Offloader Configuration + Canonical CLI.
Problem
pin_cpu_memory/pin_memorynaming split across two entry points.OffloadConfig.from_od_confighardcodes strategy mapping in one ~100-line method.Solution
Canonical parser that consumes both new RFC #7204 inputs and all legacy boolean flags:
New CLI inputs
--diffusion-offload {none,module,layer}--diffusion-offload-config <JSON>mode,components.{dit,text_encoder}.{weight_transfer,resident_layers},pin_memory--host-weight-runtime-config <JSON>mode/root/registration_limit_gibNew module
vllm_omni/diffusion/offloader/config_parser.py(538 lines):CanonicalOffloadPolicydataclass — single source of truthparse_offload_policy()— unified entry point (shortcut > JSON config > legacy booleans)parse_hwr_policy()+validate_joint()— cross-validationModified
OffloadConfig.from_od_configrefactored into canonical path + legacy path. Three backend classes unchanged.Backward compatibility: all 10 legacy CLI flags continue to work unchanged.
Testing
45+ unit tests covering shortcut expansion, JSON config parsing, legacy boolean mapping, every RFC §3/§4 validation rejection case, HWR joint validation, policy roundtrip.
Diff Summary
+ vllm_omni/diffusion/offloader/config_parser.py+ tests/diffusion/offloader/test_offload_config_parser.py~ vllm_omni/diffusion/offloader/base.py~ vllm_omni/diffusion/data.py~ vllm_omni/config/omni_config.py~ vllm_omni/entrypoints/cli/serve.pyOpen follow-ups (per RFC §8, not in this PR)
Closes part of #7204.