fix for intel devices torch compile configs - #3952
Conversation
|
Pls note, this pr rely on unslothai/unsloth-zoo#451 |
![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 fixes for torch.compile configurations on Intel devices (XPU) by making the options device-specific. The changes correctly abstract away device-specific calls like torch.cuda.synchronize() to a generic device_synchronize() function. My main feedback is to refactor a small piece of duplicated code where compile options for 'xpu' and the default case are identical. Consolidating this will improve code clarity and maintainability. Otherwise, the changes look good.
| elif DEVICE_TYPE == "xpu": | ||
| # XPU-specific torch_compile_options (disable CUDA-specific options) | ||
| new_options = """torch_compile_options = { | ||
| "epilogue_fusion" : True, | ||
| "max_autotune" : False, | ||
| "shape_padding" : True, | ||
| "trace.enabled" : False, | ||
| }""" | ||
| else: | ||
| # Default options for other device types (hip, etc.) | ||
| new_options = """torch_compile_options = { | ||
| "epilogue_fusion" : True, | ||
| "max_autotune" : False, | ||
| "shape_padding" : True, | ||
| "trace.enabled" : False, | ||
| }""" |
There was a problem hiding this comment.
The torch_compile_options for xpu and the default else case are identical. You can combine these blocks to reduce code duplication and improve maintainability.
else:
# Default options for other device types (xpu, hip, etc.)
new_options = """torch_compile_options = {
"epilogue_fusion" : True,
"max_autotune" : False,
"shape_padding" : True,
"trace.enabled" : False,
}"""
|
|
||
| if RLTrainer_name == "GRPOTrainer": | ||
| new_options = """torch_compile_options = { | ||
| # Generate torch_compile_options based on device type |
There was a problem hiding this comment.
NIT: Isn't it simpler if we have base compile_options and then add on top of them?
Because I see all the 3 backends have same configs.
…ic extensions - Extract common options into base_options shared by all device types - CUDA devices get additional CUDA-specific options - XPU, HIP, and other devices use base options only - Reduces code duplication and improves maintainability
for more information, see https://pre-commit.ci
* fix for intel devices * Refactor torch_compile_options to use base options with device-specific extensions - Extract common options into base_options shared by all device types - CUDA devices get additional CUDA-specific options - XPU, HIP, and other devices use base options only - Reduces code duplication and improves maintainability * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: danielhanchen <danielhanchen@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
3 participants
fix for intel devices torch compile configs