Fix: Skip JIT warmup when fusion is disabled via arguments#2827
Fix: Skip JIT warmup when fusion is disabled via arguments#2827kisseternity wants to merge 11 commits intoNVIDIA:mainfrom
Conversation
asolergi-nv
left a comment
asolergi-nv
left a comment
There was a problem hiding this comment.
Looks nice, but we should check exhaustively the conditions that trigger the torch.compiled functions. Could you also add geglu? Thanks!
Thanks for the feedback! I've updated the code to add exhaustive condition checks based on the MLP & shared experts implementation. Please let me know if any additional conditions should be considered. |
chtruong814
added
needs-follow-up
asolergi-nv
left a comment
asolergi-nv
left a comment
There was a problem hiding this comment.
Just remove some comments & we are ready!
Addressed, thanks! |
|
/ok to test 201e816 |
|
@mcore-oncall Hello, could you please approve the pending workflow run so the CI checks can proceed? The branch has been updated and is ready for testing. Thanks! |
|
/ok to test 11eb966 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/22682094495 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/22684909191 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/22685078517 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/22685362301 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/22685853817 |
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/22694249359 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/22695853372 |
4 participants
What does this PR do ?
This PR adds conditional checks to _warmup_jit_function in initialize.py to ensure JIT fusions (Bias+GeLU/SwiGLU and Bias+Dropout+Add) are only warmed up when they are enabled via arguments (e.g., checks args.bias_gelu_fusion, args.bias_dropout_fusion). This prevents unnecessary warmup execution when users explicitly disable these fusions using flags like --no-bias-gelu-fusion.
Pre-checks
Core 0.8)