Repository navigation
[Core] Composable parallel: Phase 1 strategy overlay (TP/DP/stage-replica) (RFC #4084) - #4281
Conversation
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
This PR is substantial (>22 files). For such a large change affecting core parallelism infrastructure, please run the L3 tests locally and paste the results here.
Phase 1 strategy overlay implementation looks architecturally sound. The declarative front-end approach for composable parallelism is well-designed.
|
@hsliuustc0106, thanks again. I ran the L3 suite for the omni paths on an 8×H100 node and pulled the results together below. All 8 errors and a few of the failures turned out to be missing test-harness dependencies / gated-model access in the runner. Once those were installed/granted they go green: 5 opencc-dependent tests -> passed on rerun after pip install opencc test_glm_tts voice-clone (×2): model-side dtype bug - float input vs BFloat16 bias in a conv1d layer. On overhead: I benchmarked the adapter path (StrategySpec -> translation -> engine kwargs) against driving the engine directly (main), on the same model/config (Qwen, TP=1, H100, warmup discarded, 3 timed iters). The translated EngineCore config is byte-identical, and decode throughput matches within noise: baseline (typed engine): 8086.1 tok/s (std 30.4) |
557df8d to
00f6978
Compare
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
Reviewing from a UX and default-YAML perspective. The core design (separate strategy intent from deploy layout) is solid — a few things that would make the feature more adoptable out of the box.
| thinker: | ||
| - axis: stage_replica | ||
| size: 2 | ||
| routing: round_robin |
There was a problem hiding this comment.
This is the only strategy example and it only demonstrates stage_replica. The PR supports tp, dp, pp, ep axes but provides no example YAMLs for them.
Consider adding a strategy_tp2.yaml (e.g. thinker: [{axis: tp, size: 2}]) alongside a matching deploy_qwen2_5_omni_tp2.yaml. TP is the most common parallelism axis and an example would help users understand the mental model.
Alternatively, a single strategy_examples.yaml with commented-out sections for each axis would serve as a living reference.
There was a problem hiding this comment.
TP example request Added a TP=2 pair: strategy_tp2.yaml + deploy_qwen2_5_omni_tp2.yaml. Verified e2e on Qwen2.5-Omni (the strategy alone sharded the thinker TP=2). Thanks!
| # python3 end2end.py --query-type text --num-prompts 6 \ | ||
| # --deploy-config deploy_qwen2_5_omni_replica.yaml \ | ||
| # --strategy-config strategy_stage_replica.yaml --omni-lb-policy round-robin | ||
| async_chunk: false |
There was a problem hiding this comment.
The header comment shows --omni-lb-policy round-robin passed on the CLI, but the PR's async_omni_engine._apply_strategy_lb_policy auto-derives the LB policy from the strategy's routing field. These two paths contradict each other.
The example should drop --omni-lb-policy and let the strategy drive it — that demonstrates the auto-derive feature. Add a comment like:
# omni-lb-policy is auto-derived from the strategy's routing field;
# only pass --omni-lb-policy if you need to override it.
There was a problem hiding this comment.
Good catch! removed it and added a note that the LB policy auto-derives from the strategy's routing field
| @@ -0,0 +1,63 @@ | |||
| # Qwen2.5-Omni deploy laid out for stage replication (4-GPU box). | |||
There was a problem hiding this comment.
This deploy YAML assumes exactly 4 GPUs (thinker on 0,1; talker on 2; code2wav on 2). Users on 2-GPU or 8-GPU boxes will have to figure this out themselves.
Consider either:
- Making the GPU count explicit in the filename (e.g.
deploy_qwen2_5_omni_replica_4gpu.yaml), or - Adding a prominent comment:
# Requires >= 3 GPUs (thinker uses 2, talker+code2wav share 1)
A common pattern is to provide variants per GPU count, or at least document how to adapt the devices fields.
There was a problem hiding this comment.
added # Requires >= 3 GPUs: thinker uses 2 (stage_replica=2), talker + code2wav share 1
|
|
||
| MESH_AXIS_KINDS: tuple[MeshAxisKind, ...] = ( | ||
| "tp", | ||
| "dp", |
There was a problem hiding this comment.
11 axis kinds are declared but only tp, dp, stage_replica are translatable in Phase 1. The rest (pp, ep, sp_ulysses, sp_ring, cfg, vae_pp, hsdp, stage_pp, cp) are reserved.
Consider adding a docstring or comment listing which axes are currently wired up vs reserved, so users reading the spec don't try to use an unsupported axis and get a confusing error. Something like:
# Phase 1: tp, dp, stage_replica
# Reserved for future phases: pp, ep, sp_ulysses, sp_ring, cfg, vae_pp, hsdp, stage_pp, cpThere was a problem hiding this comment.
Added a docstring listing wired axes (tp/dp/pp/ep/stage_replica) vs reserved (sp_ulysses/sp_ring/cfg/vae_pp/hsdp/stage_pp/cp, which raise on translation)
| omni_config_group.add_argument( | ||
| "--strategy-config", | ||
| type=str, | ||
| default=None, |
There was a problem hiding this comment.
The --strategy-config help text says "opt-in; ignored on the legacy YAML path." This is important behavior that users might miss.
Consider either:
- Emitting a warning when
--strategy-configis passed but the model uses the legacy YAML path (so the flag is silently ignored), or - Updating the help text to be more explicit: "Only works with registry-based models (e.g. qwen2_5_omni). Silently ignored for models using legacy pipeline YAML."
There was a problem hiding this comment.
logs a warning when --strategy-config is passed but the model resolves via the legacy yaml path. Thanks!
|
Thanks @hsliuustc0106 for your review. I’ve addressed the comments. Please take a look when you have a chance. Thank you. |
| @@ -0,0 +1,69 @@ | |||
| # Qwen2.5-Omni deploy laid out for stage replication. | |||
There was a problem hiding this comment.
currently, the deploy yamls are placed under src config folder and there does not exist strategy yaml for users, what do you suggest to avoid change user behavior if possible? currently we have added a lot of recipes, this needs more consideration:)
There was a problem hiding this comment.
Good point @hsliuustc0106, and agreed it needs care. The intent is to keep this fully additive:
- So that nothing changes unless you opt in. All built-in deploy recipes work exactly as today; this PR doesn't touch any of them. Strategy only kicks in when you pass
--strategy-config, otherwise the path is byte-identical to main. - We don't ship strategy YAMLs as recipes. They stay user authored, the only ones in the tree are under examples/ as references, not in vllm_omni/deploy/. So the recipe surface users rely on is unchanged.
- The common case needs no new file. Plain
--tensor-parallel-sizeetc. still work and get wrapped into the equivalent spec internally, so simple users never see a strategy file.
For the bigger question of where strategy files should live long term (and whether built-in recipes eventually pair with one), I'd defer that to #4021 so we don't migrate the recipe layout twice. For Phase 1: zero strategy recipes shipped, examples only. WDYT?
There was a problem hiding this comment.
LGTM, cc @alex-jw-brooks , we need to move it faster with this PR
There was a problem hiding this comment.
Hey, sorry for the delay @kushanam @hsliuustc0106 - I will review this PR first thing in the morning 🙏
hsliuustc0106
added
high priority
There was a problem hiding this comment.
Thanks for your patience and work @kushanam! Some suggestions and discussion points.
Also please fix DCO when you get the chance!
| raise DuplicateAxisKindError(f"axis kind {kind!r} appears more than once in the spec stack") | ||
|
|
||
| owner = _resolve_l1_owner(spec) | ||
| if kind == "dp": |
There was a problem hiding this comment.
I know the spec types the kind to MeshAxisKind, but I think it would be better to use a str enum instead for some of these types that are widely used. Otherwise refactors etc may break these types of changes in the future
There was a problem hiding this comment.
Done. class MeshAxisKind(str, Enum) string comparisons still work
There was a problem hiding this comment.
I think it would be better to also compare against the enum members instead of the raw strings to prevent typos. I can also open a follow-up for this though - IMO it would be useful to add this before we get too far along, and also clean up load_and_resolve_stage_configs to not return a tuple, but we don't need to block this PR for it
| descriptor = f"stage_id={key}" | ||
| else: | ||
| matches = [s for s in stages if getattr(s, "model_stage", None) == key] | ||
| descriptor = f"model_stage={key!r}" |
There was a problem hiding this comment.
can you clarify why we match against both stage and stage id? This seems like it could cause confusing behaviors imo
There was a problem hiding this comment.
Since key type picks the field (str -> model_stage, int -> stage_id, bool rejected) I think it's unambiguous. Added a docstring spelling this out. Does it satisfy your concern @alex-jw-brooks?
There was a problem hiding this comment.
Thanks @kushanam! Is there a reason why we can't just use the model_stage given that it's not optional on stage configs? This still feels unnecessarily flexible to me. I think it's more practical to refer to stages by name rather than indices as keys, because in addition to being more readable, its more understandable if we have situations like either of the following:
- An architecture supports multiple pipeline configs, and the user either picks one or its selected based on the HF config for what is actually supported. In this case, the index for a given stage may not be the same across pipeline config variations, but the name should be
- A model's stage separations change (e.g., splitting one stage to two). This would mess up both names and indices, but with names we could do potential validation that may not be possible based purely on indices
…ranslationError (taxonomy align)
Thank you @alex-jw-brooks for your thorough review. I addressed the comments. Please have a look when you got a chance! |
…ject#8 docstring) Signed-off-by: Kushan A. <kahmadian@nvidia.com>
…ranslationError (taxonomy align) Signed-off-by: Kushan A. <kahmadian@nvidia.com>
b550c95 to
f48a24c
Compare
Signed-off-by: Kushan A. <kahmadian@nvidia.com>
There was a problem hiding this comment.
Thanks @kushanam! I think this is close to being ready to merge, but added some additional discussion points. The most important things are the public docs, the test skips, and the dual key stuff though (since it affects API a bit), I think once that is resolved we can merge + iterate on this design in the next phase.
Also please resolve conflicts when you get the chance!
| raise DuplicateAxisKindError(f"axis kind {kind!r} appears more than once in the spec stack") | ||
|
|
||
| owner = _resolve_l1_owner(spec) | ||
| if kind == "dp": |
There was a problem hiding this comment.
I think it would be better to also compare against the enum members instead of the raw strings to prevent typos. I can also open a follow-up for this though - IMO it would be useful to add this before we get too far along, and also clean up load_and_resolve_stage_configs to not return a tuple, but we don't need to block this PR for it
| if not derived: | ||
| return | ||
| explicit = kwargs.get("omni_lb_policy") | ||
| user_set = explicit is not None and str(explicit) != "random" |
There was a problem hiding this comment.
As a follow-up, I think we should make sure defaults/fallbacks for all of these are initialized in a way that is easy to follow (i.e., things like lora stuff, async chunk, and this are currently a bit scattered).
It should be a separate PR as this effort continues in case there are conflicts
There was a problem hiding this comment.
Agreed, will do this as a separate follow-up after this PR lands. I'll open a tracking issue covering the lora, async-chunk, and omni_lb_policy defaults so they all get centralized in one pass
| type=str, | ||
| default=None, | ||
| choices=["random", "round-robin", "least-queue-length"], | ||
| help="StagePool load-balancer policy for replicated stages (engine-level knob).", |
There was a problem hiding this comment.
I think engine-level knob should actually be pipeline level or orchestrator level. Engine level is confusing because every stage has its own engine, so it sounds like this is either stage level or ambiguously configuring all stages of a specific engine type
There was a problem hiding this comment.
Renamed to "orchestrator-level knob" everywhere
|
|
||
| def test_device_check_survives_cli_override(): | ||
| if not _DEPLOY.exists(): | ||
| pytest.skip("qwen2_5_omni deploy config not found") |
There was a problem hiding this comment.
Can you remove these conditional skips? If the deploy config is moved, these tests should fail so that we know to update them rather than skip
There was a problem hiding this comment.
Removed. All three _DEPLOY.exists() skip guards in this file are gone so the tests now fail loudly if the deploy config moves
| descriptor = f"stage_id={key}" | ||
| else: | ||
| matches = [s for s in stages if getattr(s, "model_stage", None) == key] | ||
| descriptor = f"model_stage={key!r}" |
There was a problem hiding this comment.
Thanks @kushanam! Is there a reason why we can't just use the model_stage given that it's not optional on stage configs? This still feels unnecessarily flexible to me. I think it's more practical to refer to stages by name rather than indices as keys, because in addition to being more readable, its more understandable if we have situations like either of the following:
- An architecture supports multiple pipeline configs, and the user either picks one or its selected based on the HF config for what is actually supported. In this case, the index for a given stage may not be the same across pipeline config variations, but the name should be
- A model's stage separations change (e.g., splitting one stage to two). This would mess up both names and indices, but with names we could do potential validation that may not be possible based purely on indices
|
|
||
| The output is a plain sizing struct keyed by the real ``OmniEngineArgs`` / | ||
| ``EngineArgs`` field names, so this module runs on CPU and the deploy layer can | ||
| splat the result into a stage's engine args. |
There was a problem hiding this comment.
…ject#8 docstring) Signed-off-by: Kushan A. <kahmadian@nvidia.com>
…ranslationError (taxonomy align) Signed-off-by: Kushan A. <kahmadian@nvidia.com>
Signed-off-by: Kushan A. <kahmadian@nvidia.com>
- Narrow strategy role key to str (model_stage); drop int/bool branch - Remove conditional pytest.skip on deploy-YAML path - Rename "engine-level knob" -> "orchestrator-level knob" (omni_lb_policy) - Add docs/configuration/composable_parallel.md Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com>
…ntegration stage_config._PIPELINE_REGISTRY -> pipeline_registry.OMNI_PIPELINES; StageConfigFactory now in config_factory. Caught by CPU test collection. Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com>
…m_registry _create_from_registry on main now takes pipeline_cfg as positional arg 2 and returns (stages, omni_lb_policy). Update 2 tests in test_factory_integration.py to pass pipeline_cfg and 1 test in test_config_factory.py to unpack the tuple. Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com>
…y move The _reconcile_strategy_with_cli warning now emits from config_factory.py; update test_cli_overrides_strategy_with_warning to attach its handler to vllm_omni.config.config_factory. Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com>
…n B) Opt-in strategy.yaml that overlays per-role parallel sizing (TP, stage replica, ...) onto merged stage configs after pipeline+deploy merge and before CLI overrides. No-op when no strategy is supplied. Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: kushanam <kushanam@users.noreply.github.com>
Apply strategy-derived omni_lb_policy at engine construction, expose --strategy-config on vllm serve (and headless), validate dp routing policy, enforce strategy sizing against the post-CLI layout, and harden strategy.yaml parsing. Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: kushanam <kushanam@users.noreply.github.com>
Modernize typing imports/annotations (UP006/UP035/UP045) and add the Error suffix to composable-parallel exception classes (N818). Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: kushanam <kushanam@users.noreply.github.com>
Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com> Signed-off-by: kushanam <kushanam@users.noreply.github.com>
TP=2 example pair (strategy_tp2 + deploy_qwen2_5_omni_tp2), spec wired/reserved-axis docs, replica example lb-policy/GPU notes, and a legacy-path warning + sharper help text for --strategy-config. Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com>
…ject#8 docstring) Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: kahmadian <kahmadian@users.noreply.github.com>
…ranslationError (taxonomy align) Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com> Signed-off-by: kahmadian <kahmadian@users.noreply.github.com>
Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com>
- Narrow strategy role key to str (model_stage); drop int/bool branch - Remove conditional pytest.skip on deploy-YAML path - Rename "engine-level knob" -> "orchestrator-level knob" (omni_lb_policy) - Add docs/configuration/composable_parallel.md Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com>
…ntegration stage_config._PIPELINE_REGISTRY -> pipeline_registry.OMNI_PIPELINES; StageConfigFactory now in config_factory. Caught by CPU test collection. Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com>
…m_registry _create_from_registry on main now takes pipeline_cfg as positional arg 2 and returns (stages, omni_lb_policy). Update 2 tests in test_factory_integration.py to pass pipeline_cfg and 1 test in test_config_factory.py to unpack the tuple. Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com>
…y move The _reconcile_strategy_with_cli warning now emits from config_factory.py; update test_cli_overrides_strategy_with_warning to attach its handler to vllm_omni.config.config_factory. Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com>
Parse and pass stage_overrides into load_and_resolve_stage_configs from run_headless via a shared parse_stage_overrides helper, restoring parity with the standard engine path. Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com>
Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: Kushan A <kahmadian@nvidia.com>
aa62af9 to
6869fd1
Compare
Thank you @fhfuih, I will address them in the following PR if you don't mind :)
@alex-jw-brooks Sorry for the delay and thanks for you attention :) Rebased. |
hsliuustc0106
added
ready
…lica) (RFC vllm-project#4084) (vllm-project#4281) Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: kushanam <kushanam@users.noreply.github.com> Signed-off-by: Kushan A <kahmadian@nvidia.com> Signed-off-by: kahmadian <kahmadian@users.noreply.github.com> Co-authored-by: kushanam <kushanam@users.noreply.github.com> Co-authored-by: kahmadian <kahmadian@users.noreply.github.com> Signed-off-by: Wojciech Kutak <wkutak@nvidia.com>
…lica) (RFC vllm-project#4084) (vllm-project#4281) Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: kushanam <kushanam@users.noreply.github.com> Signed-off-by: Kushan A <kahmadian@nvidia.com> Signed-off-by: kahmadian <kahmadian@users.noreply.github.com> Co-authored-by: kushanam <kushanam@users.noreply.github.com> Co-authored-by: kahmadian <kahmadian@users.noreply.github.com>
…lica) (RFC vllm-project#4084) (vllm-project#4281) Signed-off-by: Kushan A. <kahmadian@nvidia.com> Signed-off-by: kushanam <kushanam@users.noreply.github.com> Signed-off-by: Kushan A <kahmadian@nvidia.com> Signed-off-by: kahmadian <kahmadian@users.noreply.github.com> Co-authored-by: kushanam <kushanam@users.noreply.github.com> Co-authored-by: kahmadian <kahmadian@users.noreply.github.com>
4 participants
Purpose
Implements Phase 1 of RFC #4084 ([RFC]: Composable Parallel Strategies for
vLLM-Omni), adding an opt-in, additive declarative front-end for vLLM-Omni
parallelism. Users describe per-role parallelism in a small
strategy.yaml; theorchestrator overlays the derived sizing onto the already-merged stage configs
(after
merge_pipeline_deploy, before CLI overrides) and then leaves runtimerequest routing to vLLM-Omni's existing load balancer. When no strategy is
supplied, behavior is unchanged.
This follows the design consensus reached in the RFC discussion:
the strategy owns each axis's degree, the deploy YAML owns physical
devices,and a validator enforces
count(devices) == product(degrees).axes it declares; it overrides defaults and errors only on a conflicting
explicit value, so the additive rollout stays painless.
dpis realizedintra-engine by vLLM's DP load balancer;
stage_replicaspins up independentstage backends balanced by omni's
StagePool— with a pre-flight check that astage_replicaplan can never imply cross-replica vLLM DP.can't honor) is rejected with an actionable error rather than routed at random.
What landed:
vllm_omni/config/composable_parallel/— pure, native types:spec.py,routing.py,aggregation.py,translator.py(
StrategySpec → OmniParallelConfig),apply.py(override-after-merge +pre-spawn device check),
strategy_yaml.py(loads astrategy.yaml).OrchestratorArgs.strategy_config,--strategy-configon
vllm serve(head + headless) and the Qwen2.5-Omni example, threadedthrough
load_and_resolve_stage_configsintoStageConfigFactory, with apost-CLI reconciliation pass that warns when a CLI arg overrides a
strategy-declared axis and re-runs the device check against the effective
layout. A strategy-derived
omni_lb_policyis applied at engine construction.Axis coverage:
tp,dp,pp,ep, andstage_replicaare translated.tpandstage_replicaare validated by the live e2e below;dp/pp/epsizing is covered by CPU unit tests (Phase 1 per the RFC is DP/TP — the
remaining schemes' full semantics arrive in follow-up RFCs).
Test Plan
CPU unit tests (no GPU) — translation, override/conflict, device guard, YAML
parsing, closed-hierarchy guards, factory integration against the real bundled
deploy, post-CLI reconciliation, lb-policy precedence, and the serve CLI flag:
Live e2e (H100x8), Qwen2.5-Omni, from
examples/offline_inference/qwen2_5_omni/:Test Result
CPU unit tests: all green.
Live e2e — control (no strategy):
RC=0, produced text + audio, no errors.Live e2e — strategy (
stage_replica=2on the thinker, 4 prompts):RC=0,all 4 prompts produced text + audio, clean StagePool shutdown, no
errors. The new paths fired exactly as intended (no
--omni-lb-policypassed)Essential Elements of an Effective PR Description Checklist
strategy_stage_replica.yaml,deploy_qwen2_5_omni_replica.yaml, and--strategy-config/--deploy-config/--omni-lb-policyflags in the Qwen2.5-Omni example).