Enable torch version dependent compilation of record_module and iter_params #7362
Conversation
tohtana
left a comment
tohtana
left a comment
There was a problem hiding this comment.
Hi @deepcharm,
Thank you for the improvement. I confirmed that this works with the latest PyTorch (v2.7).
I think it is good to merge this PR now, but I remember that the compiler of the older version threw an error. I'm not sure which version fixed it. We might need to turn on/off compiling this function depending on the PyTorch version.
hi @tohtana, thanks for pointing this out - let me try to find the failing compiler version and then turn on/off that respectively. |
|
@deepcharm - is this ready to be merged? |
Hi @loadams, @tohtana Kindly review, thanks. |
Signed-off-by: Max Kovalenko <[email protected]>
This reverts commit 11612773b3d68aa5b8d72bad1de4b1714ea1193a. Signed-off-by: Max Kovalenko <[email protected]>
Dynamo breaks the graph because currently compile is disabled for record_module. Re-enabling the compile should avoid the graph break and improve the performance. Signed-off-by: Max Kovalenko <[email protected]>
Signed-off-by: Max Kovalenko <[email protected]>
deepcharm
changed the title
…params (deepspeedai#7362) Dynamo breaks graphs because currently compile is disabled for a number of functions such as `iter_params` and `record_module`. The above functions compile successfully for at least PyTorch version 2.7.0. We enable the compilation based on the user PyTorch version using a new `compiler.enable(min_version=None)` decorator. This should avoid the corresponding graph breaks and improve the performance. --------- Signed-off-by: Max Kovalenko <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]> Co-authored-by: Logan Adams <[email protected]>
…params (deepspeedai#7362) Dynamo breaks graphs because currently compile is disabled for a number of functions such as `iter_params` and `record_module`. The above functions compile successfully for at least PyTorch version 2.7.0. We enable the compilation based on the user PyTorch version using a new `compiler.enable(min_version=None)` decorator. This should avoid the corresponding graph breaks and improve the performance. --------- Signed-off-by: Max Kovalenko <[email protected]> Co-authored-by: Masahiro Tanaka <[email protected]> Co-authored-by: Logan Adams <[email protected]>
3 participants
Dynamo breaks graphs because currently compile is disabled for a number of functions such as
iter_paramsandrecord_module.The above functions compile successfully for at least PyTorch version 2.7.0.
We enable the compilation based on the user PyTorch version using a new
compiler.enable(min_version=None)decorator.This should avoid the corresponding graph breaks and improve the performance.