Remove PyTorch 2.5.0 checks

RdoubleA commented 1 week ago

PyTorch 2.5.0 is officially released which includes features such as FlexAttention that is now public API and other compile features. We can now remove the following checks:

[ ] Flex attention requiring torch >= 2.5.0, can just remove this line: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/torchtune/utils/_import_guard.py#L13
[ ] Switch to per layer compile by default, which brings down compile time from 10 mins to under 1 minute: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/torchtune/training/_compile.py#L45
[ ] Update this flex attention related docstring to say "If flex attention is supported by the current hardware": https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/torchtune/modules/attention_utils.py#L119
[ ] Remove this check in reparametrize state dict hook: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/torchtune/modules/common_utils.py#L150
[ ] Default use_streams to True in activation offloading: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/torchtune/training/_activation_offloading.py#L36 and update related docs: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/docs/source/tutorials/memory_optimizations.rst#L111
[ ] Remove these decorators checking for 2.5 on flex attention tests: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/tests/torchtune/modules/test_attention_utils.py#L87
[ ] Update lora distributed and lora single device recipe docstrings: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/recipes/lora_finetune_distributed.py#L77

JP-sDEV commented 1 week ago

Hi @RdoubleA, can I be assigned this issue?

RdoubleA commented 1 week ago

Hi @JP-sDEV, sure! feel free to take on as many or as few of the items and link this issue in your PR. Feel free to comment if you run into any trouble and please include me or @joecummings as reviewers. Appreciate it!

JP-sDEV commented 6 days ago

Thanks! I have opened a pull request that completes all the tasks in this issue.

PyTorch 2.5.0 is officially released which includes features such as FlexAttention that is now public API and other compile features. We can now remove the following checks:

[x] Flex attention requiring torch >= 2.5.0, can just remove this line: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/torchtune/utils/_import_guard.py#L13

[x] Switch to per layer compile by default, which brings down compile time from 10 mins to under 1 minute: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/torchtune/training/_compile.py#L45

[x] Update this flex attention related docstring to say "If flex attention is supported by the current hardware": https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/torchtune/modules/attention_utils.py#L119

[x] Remove this check in reparametrize state dict hook: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/torchtune/modules/common_utils.py#L150

[x] Default use_streams to True in activation offloading: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/torchtune/training/_activation_offloading.py#L36 and update related docs: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/docs/source/tutorials/memory_optimizations.rst#L111

[x] Remove these decorators checking for 2.5 on flex attention tests: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/tests/torchtune/modules/test_attention_utils.py#L87

[x] Update lora distributed and lora single device recipe docstrings: https://github.com/pytorch/torchtune/blob/7d29c21149ace4be9fdc2709a03c0a21ba191010/recipes/lora_finetune_distributed.py#L77

pytorch / torchtune

Remove PyTorch 2.5.0 checks #1861