Replace cuda_graph assertion with explicit ValueError; fix custom_op type annotations - #8336
Conversation
PyTorch's torch.library.infer_schema does not accept builtin list[int] annotations; use typing.List[int] instead so deepspeed imports cleanly on recent PyTorch releases. Signed-off-by: Chakshu Dhannawat <chakshu.dhannawat1@gmail.com>
When replace_with_kernel_inject and enable_cuda_graph are both requested for a model whose policy does not support CUDA graphs, raise a clear ValueError instead of an internal assertion. Fixes deepspeedai#8330 Signed-off-by: Chakshu Dhannawat <chakshu.dhannawat1@gmail.com>
![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.
Reviewed commit: 5e4cde7fc1
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| intermediate_size=64, | ||
| torch_dtype=torch.bfloat16, | ||
| ) | ||
| model = LlamaForCausalLM(model_config).to("cuda") |
There was a problem hiding this comment.
When this inference test is selected on a non-CUDA accelerator where get_accelerator().is_available() is true (for example XPU/HPU), the class does not skip, but moving the model with .to("cuda") fails before deepspeed.init_inference() can raise the ValueError being tested. Please either skip unless the selected accelerator is CUDA, or move the model to get_accelerator().device_name() so accelerator-specific test runs do not fail for an unrelated reason.
Useful? React with 👍 / 👎.
1 participant
This PR addresses #8330 by replacing the internal assertion in
replace_with_policywith a clearValueErrorwhen bothreplace_with_kernel_inject=Trueandenable_cuda_graph=Trueare requested for a model whose policy does not support CUDA graphs.Changes:
deepspeed/module_inject/replace_module.py: raise an actionableValueErrorinstead of asserting.tests/unit/inference/test_inference_config.py: add a regression test using a tiny Llama model that verifies the new error message.deepspeed/compile/custom_ops/tp_collectives.py: replacelist[int]annotations withtyping.List[int]soimport deepspeedworks on PyTorch 2.6+ (the new regression test needs this to import cleanly).The regression test was verified on a single H100 with: