Remove PyTorch 2.5.0 checks

[RdoubleA]

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:

  torchtune/torchtune/utils/_import_guard.py

  Line 13 in 7d29c21

  torch_version_ge("2.5.0")

• Switch to per layer compile by default, which brings down compile time from 10 mins to under 1 minute:

  torchtune/torchtune/training/_compile.py

  Line 45 in 7d29c21

  if torch_version_ge("2.5.0"):

• Update this flex attention related docstring to say "If flex attention is supported by the current hardware":

  torchtune/torchtune/modules/attention_utils.py

  Line 119 in 7d29c21

  torch version >= 2.5.0, this is done by creating a mask_mod function with the

• Remove this check in reparametrize state dict hook:

  torchtune/torchtune/modules/common_utils.py

  Line 150 in 7d29c21

  if torch.__version__ < "2.5.0.dev20240906":

• Default use_streams to True in activation offloading:

  torchtune/torchtune/training/_activation_offloading.py

  Line 36 in 7d29c21

  after torch-2.5.0.dev20240907. Default: True if a later torch build is found, else False.

  and update related docs:

  torchtune/docs/source/tutorials/memory_optimizations.rst

  Line 111 in 7d29c21

  usage of streams, make sure you are on a torch version later than PyTorch 2.5.0.dev20240907.

• Remove these decorators checking for 2.5 on flex attention tests:

  torchtune/tests/torchtune/modules/test_attention_utils.py

  Line 87 in 7d29c21

  reason="Please install a nightly build of torch (>=2.5.0) to run this test.",

• Update lora distributed and lora single device recipe docstrings:

  torchtune/recipes/lora_finetune_distributed.py

  Line 77 in 7d29c21

  the computation. This option is currently only available on PyTorch nightly 2.5.0.dev20240907

[JP-sDEV]

Hi @RdoubleA, can I be assigned this issue?

[RdoubleA]

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]

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:

• Flex attention requiring torch >= 2.5.0, can just remove this line:

  torchtune/torchtune/utils/_import_guard.py

  Line 13 in 7d29c21

  torch_version_ge("2.5.0")

• Switch to per layer compile by default, which brings down compile time from 10 mins to under 1 minute:

  torchtune/torchtune/training/_compile.py

  Line 45 in 7d29c21

  if torch_version_ge("2.5.0"):

• Update this flex attention related docstring to say "If flex attention is supported by the current hardware":

  torchtune/torchtune/modules/attention_utils.py

  Line 119 in 7d29c21

  torch version >= 2.5.0, this is done by creating a mask_mod function with the

• Remove this check in reparametrize state dict hook:

  torchtune/torchtune/modules/common_utils.py

  Line 150 in 7d29c21

  if torch.__version__ < "2.5.0.dev20240906":

• Default use_streams to True in activation offloading:

  torchtune/torchtune/training/_activation_offloading.py

  Line 36 in 7d29c21

  after torch-2.5.0.dev20240907. Default: True if a later torch build is found, else False.

  
  and update related docs:

  torchtune/docs/source/tutorials/memory_optimizations.rst

  Line 111 in 7d29c21

  usage of streams, make sure you are on a torch version later than PyTorch 2.5.0.dev20240907.

• Remove these decorators checking for 2.5 on flex attention tests:

  torchtune/tests/torchtune/modules/test_attention_utils.py

  Line 87 in 7d29c21

  reason="Please install a nightly build of torch (>=2.5.0) to run this test.",

• Update lora distributed and lora single device recipe docstrings:

  torchtune/recipes/lora_finetune_distributed.py

  Line 77 in 7d29c21

  the computation. This option is currently only available on PyTorch nightly 2.5.0.dev20240907