[BugFix] Fix DeepSeek-R1 hang with DP and MTP #30119
Conversation
![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
The pull request addresses a hang issue related to CUDA graphs when using DeepSeek-R1 with Data Parallelism (DP) and Model-Tensor Parallelism (MTP). The fix aims to ensure that the drafter model correctly utilizes piecewise CUDA graphs when the main model is operating with any CUDA graph mode. While the overall intent of the fix appears to be achieved, there is a logical redundancy in the condition used to determine use_cudagraphs for the drafter, which could be simplified for better readability and maintainability.
LucasWilkinson
added
the
ready
| ( | ||
| is_graph_capturing | ||
| and cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE | ||
| ) | ||
| or (cudagraph_runtime_mode != CUDAGraphMode.NONE) |
There was a problem hiding this comment.
the basic idea here is during a DP dummy run we may end up using FULL CGs for the target model; this was set to only using drafter PIECEWISE CGs when he target model was using PIECEWISE CGs (mostly for CG capture) meaning if the target model used FULL CGs this would revert the drafter to eager. This didn't match the behavior of non-dummy runs where if the main model used FULL CGs the drafter would still use PIECEWISE CGs.
Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
2 participants
There's and edge case where when the dummy_run runs FULL cudagraphs and as a result we mistakenly disable cudagraphs for the drafter; this means that the dummy_run doesn't call
pad_for_cudagraphafter_pad_batch_across_dpresulting in a hang.This PR is a temporary hack but we should rework the drafter similar to #28579 where we pad for cudagraphs before padding for DP.