Open RdoubleA opened 1 week ago
Hi @RdoubleA, can I be assigned this issue?
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!
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 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: