Add evaluate_guards option to DynamicShapesConfig #27432
Conversation
laithsakka
requested review from
NickLucche,
ProExpertProg,
WoosukKwon,
hmellor,
houseroad,
mgoin,
robertgshaw2-redhat,
sighingnow,
simon-mo,
tlrmchlsmth,
yewentao256,
youkaichao and
zou3519
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a valuable evaluate_guards option for debugging dynamic shapes in torch.compile, which will be very helpful for understanding and debugging recompilations. The implementation includes a significant and well-executed refactoring of the compilation wrapper and decorator. The new tests are comprehensive and effectively validate the new functionality. I've identified a critical bug in vllm/compilation/backends.py due to an incorrect comparison, and a minor redundancy in vllm/config/compilation.py. Once these issues are addressed, this PR will be a solid contribution.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Which PR is this built on top of? |
|
This pull request has merge conflicts that must be resolved before it can be |
|
Ok, should we wait for those to merge before reviewing this one? |
| backed/unbacked. | ||
| """ | ||
|
|
||
| evaluate_guards: bool = False |
There was a problem hiding this comment.
I think I want to get into a state where this is always on during vLLM warmup, and then optionally during runtime. This involves making sure the guards are actually correct (maybe this requires unbacked). Not for this PR, but we should do this in a follow-up.
| options["guard_filter_fn"] = lambda x: [False for _ in x] | ||
| if vllm_config.compilation_config.dynamic_shapes_config.evaluate_guards: | ||
| options["guard_filter_fn"] = lambda x: [ | ||
| entry.guard_type == "SHAPE_ENV" for entry in x |
There was a problem hiding this comment.
this seems brittle..
can we add a test in pytorch for this or something? "if you change how the dynamic shape guards look then you need to change vLLM"
There was a problem hiding this comment.
yeh we shall if we dont have any. I will file issue and assign to me.
| f"when shape changes from 10 to 5, but the " | ||
| f"model ran successfully without error." | ||
| ) | ||
| except Exception as e: |
There was a problem hiding this comment.
is there a more specific exception we can add?
| else: | ||
| print(f"❌ {config_desc}: Unexpected failure") | ||
| print(f" Error: {e}") | ||
| raise |
There was a problem hiding this comment.
can you add a test that we dropped 0/1 specialization in backed mode?
| if "guard" in error_msg.lower() or "recompile" in error_msg.lower(): | ||
| print(f"✅ {config_desc}: Expected failure due to guard violation") | ||
| print(f" Error (truncated): {error_msg[:150]}") | ||
| else: | ||
| # Unexpected error type | ||
| print(f"❌ {config_desc}: Unexpected error type") | ||
| print(f" Error: {e}") | ||
| raise |
There was a problem hiding this comment.
tests aren't really supposed to print. why not pytest.fail for the failure case?
zou3519
left a comment
zou3519
left a comment
There was a problem hiding this comment.
seems reasonable, just had testing qs
|
|
||
| model(input1) | ||
|
|
||
| if evaluate_guards and not is_01_specialization: |
There was a problem hiding this comment.
Can we use with pytest.raises()?
| # 2 for any, we assume it's 0. | ||
| for s, r in fake_mode.shape_env.var_to_range.items(): | ||
| if r.lower == 2: | ||
| fake_mode.shape_env.var_to_range[s] = ValueRanges(0, r.upper) |
There was a problem hiding this comment.
Can you import ValueRanges locally as it's an internal?
|
|
||
| fake_mode = detect_fake_mode() | ||
|
|
||
| if ( |
There was a problem hiding this comment.
Why do we only do this with evaluate_guards on?
There was a problem hiding this comment.
Oh is it because otherwise we drop all guards anyway?
|
This pull request has merge conflicts that must be resolved before it can be |
Signed-off-by: Laith Sakka <lsakka@meta.com>
4 participants
evaluate_guards is A debug mode to detect and fail if Dynamo ever specializes a dynamic shape by
guarding on it. When True, dynamic shape guards are not dropped from Dynamo.
And a failure will be triggered if recompilation ever happens due to that.
Enabling this allow observing the dynamic shapes guards in the tl-parse artifact.
One exception that evaluate_guards wont verify is counter 0/1 specializations.
Backed dynamic shapes by default will assume sizes >=2. we will not check nor fail
if that guard fails for 0/1 inputs sizes.