chore: remove PyTorch 2.5.0 checks #1877

JP-sDEV · 2024-10-21T16:21:32Z

Context

What is the purpose of this PR? Is it to

add a new feature
fix a bug
update tests and/or documentation

Changelog

What are the changes made in this PR?

modify documentation and checks for PyTorch 2.5.0

Note: Line 69 in torchtune/torchtune/training/_activation_offloading.py, should this if block be removed? or modified to check for pytorch 2.5.0? - was not included in the issue to modify this block.

 if use_streams is None:
            # Default to True if an acceptable torch is installed (later nightly/version or from source)
            **self.use_streams = torch.__version__ >= "2.5.0.dev20240907"**

Test plan

Please make sure to do each of the following if applicable to your PR. If you're unsure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.

run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
add unit tests for any new functionality
update docstrings for any new or updated methods or classes
run unit tests via pytest tests
run recipe tests via pytest tests -m integration_test
manually run any new or modified recipes with sufficient proof of correctness
include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

UX

If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Here is a docstring example
and a tutorial example

I did not change any public API
I have added an example to docs or docstrings

pytorch-bot · 2024-10-21T16:21:36Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1877

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 2 Active SEVs

There are 2 currently active SEVs. If your PR is affected, please view them below:

✅ No Failures

As of commit 37d5a01 with merge base 4b6877a ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

RdoubleA

Thanks again for doing these - looks great. A few minor nits and one comment we should decide on.

RdoubleA · 2024-10-22T16:34:00Z

docs/source/tutorials/memory_optimizations.rst

@@ -108,7 +108,7 @@ checkpointing, where all activations will either be recomputed later in the back

 To enable activation offloading, use the ``enable_activation_offloading`` config entry or flag
 in our lora finetuning single device recipe, e.g. ``enable_activation_offloading=True``. To allow
-usage of streams, make sure you are on a torch version later than PyTorch 2.5.0.dev20240907.
+usage of streams, make sure you are on a torch version equal to or later than PyTorch.


Suggested change

usage of streams, make sure you are on a torch version equal to or later than PyTorch.

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

bumping this

I think we need a merge here? The docs in main read:

To enable activation offloading, use enable_activation_offloading=True. If you are on torch version later than PyTorch 2.5.0, it will allow the usage of multiple CUDA streams automatically.

I believe we can just revert the changes here to leave the file as-as since it's been updated in another PR @JP-sDEV

RdoubleA · 2024-10-22T16:34:46Z

torchtune/training/_activation_offloading.py

@@ -33,7 +33,7 @@ class OffloadActivations(saved_tensors_hooks):

        use_streams (Optional[bool]): Whether or not to use streams for performance optimization where
            the communications get overlapped with the computation. Requires a torch build
-            after torch-2.5.0.dev20240907. Default: True if a later torch build is found, else False.
+            after torch-2.5.0.]. Default: True.


Suggested change

after torch-2.5.0.]. Default: True.

after torch-2.5.0. Default: True.

bumping this

RdoubleA · 2024-10-22T16:37:18Z

torchtune/training/_compile.py

-                Compiling full model with torch.compile...
-                For faster compile times via per-layer compile, please run on PyTorch nightlies.
-                """
+            log.warning(


I'm not sure if we want to retain the fallback logic for older pytorch versions. If so, then the if-else should remain the same and only the warning message should be updated. any thoughts? @ebsmothers @felipemello1

Sorry just seeing this now. I think if we claim to not support PyTorch < 2.5 then we shouldn't leave in the full-model compile option at all. For the same reason I'm ambivalent about leaving in the log warning.. really if we want to check someone is at least on the latest stable PyTorch we should just do it in a single consolidated place. So not the end of the world to keep the warning in, but personally I'd just take it out.

Where did we land on this?

I vote remove it altogether. We claim to only support latest stable PyTorch so no need to keep this around. We can treat adding a consolidated PyTorch version check as a separate task

I can't wait to simplify our attention dispatching

JP-sDEV · 2024-10-22T17:02:17Z

Thanks again for doing these - looks great. A few minor nits and one comment we should decide on.

Thanks, will keep up with the updates! Once finalized, let me know what updates are needed. Cheers! :)

RdoubleA · 2024-11-12T18:50:19Z

Hey @JP-sDEV, pinging this thread again to see if you're able to resolve the merge conflicts so we can get this landed. I believe we agreed that we can remove the fallback logic in _compile.py, so would be great to make that change as well.

JP-sDEV · 2024-11-12T20:09:20Z

Hi @RdoubleA, apologies for the delay - will resolve the conflicts and remove the fallback logic first thing tomorrow morning (EST)

RdoubleA · 2024-11-14T14:41:18Z

torchtune/training/_activation_offloading.py

@@ -67,7 +67,7 @@ class OffloadActivations(saved_tensors_hooks):
    def __init__(
        self,
        use_pin_memory: bool = True,
-        use_streams: Optional[bool] = None,
+        use_streams: Optional[bool] = True,
        max_fwd_stash_size: int = 5,
        min_offload_size: int = 1024,
    ) -> None:


let's remove the check below and make use_streams: bool = True

are you referring to this use_streams check found in __init__?

if use_streams is None: # Default to True if an acceptable torch is installed (later nightly/version or from source) self.use_streams = torch.__version__ >= "2.5.0.dev20240907" else: self.use_streams = use_streams

or should it be changed like this, where use_steams is set to False?

if use_streams is False: # Default to True if an acceptable torch is installed (later nightly/version or from source) self.use_streams = torch.__version__ >= "2.5.0.dev20240907" else: self.use_streams = use_streams

yep! it would just be:

self.use_streams = use_streams

I have noticed that OffloadingActivations in torchtune/training/_activation_offloading.py still checks for torch version: 2.5.0.dev20240907. Should this check also be removed?

# for streaming if self.use_streams: if torch.__version__ < "2.5.0.dev20240907": raise RuntimeError( "OffloadActivations with use_streams=True requires PyTorch 2.5.0.dev20240907 or later." )

Yes, good catch. Let's remove that as well. I believe these may have been added after I put the issue up, or maybe I just missed it

sounds good, will update all 2.5.0.dev20240907 checks in the file

RdoubleA

Thanks for the follow up, a few more comments and it's good to go

torchtune/modules/attention_utils.py

RdoubleA · 2024-11-14T15:33:53Z

The gpu test failure is interesting, let me take a closer look

SalmanMohammadi · 2024-11-14T15:47:16Z

Can we update these checks here? Are we able to remove them?

torchtune/recipes/qat_distributed.py

Line 139 in 4b6877a

and not utils.torch_version_ge("2.4.0")

torchtune/recipes/full_finetune_distributed.py

Line 128 in 4b6877a

and not utils.torch_version_ge("2.4.0")

RdoubleA · 2024-11-15T18:48:17Z

tests/torchtune/modules/test_attention_utils.py

@@ -82,10 +82,6 @@ def test_packed_block_causal_mask_sdpa(self, seq_lens):
        )
        torch.testing.assert_close(actual, expected)

-    @pytest.mark.skipif(


Ok looks like we need to keep this check, in case the hardware that runs the gpu tests on GitHub CI does not support flex attention

RdoubleA · 2024-11-15T18:49:42Z

@JP-sDEV thanks for pushing this through, just one more comment which should fix the failing gpu test

JP-sDEV · 2024-11-18T13:36:19Z

@JP-sDEV thanks for pushing this through, just one more comment which should fix the failing gpu test

I am reading the comments and reviewing the commits, and if I understand correctly - the changes needed for checking hardware compatibility is something outside this issue?

If so, I'll merge with the main branch once resolved :)

RdoubleA · 2024-11-18T15:28:26Z

The update is that we just need to keep the decorator that was removed on test_packed_block_causalmask_flex, then it should be good to go

codecov-commenter · 2024-11-18T20:12:06Z

Codecov Report

Attention: Patch coverage is 10.00000% with 9 lines in your changes missing coverage. Please review.

Project coverage is 64.66%. Comparing base (4b6877a) to head (37d5a01).
Report is 10 commits behind head on main.

Files with missing lines	Patch %	Lines
torchtune/training/_compile.py	0.00%	7 Missing ⚠️
torchtune/modules/common_utils.py	0.00%	1 Missing ⚠️
torchtune/training/_activation_offloading.py	50.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1877      +/-   ##
==========================================
- Coverage   67.51%   64.66%   -2.86%     
==========================================
  Files         318      318              
  Lines       17684    17641      -43     
==========================================
- Hits        11940    11408     -532     
- Misses       5744     6233     +489

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚨 Try these New Features:

Flaky Tests Detection - Detect and resolve failed and flaky tests

RdoubleA · 2024-11-18T20:37:10Z

@JP-sDEV I went ahead and quickly made the change - thank you again for your help!

JP-sDEV · 2024-11-19T15:41:01Z

No problem! Thank you for guiding me through the issue and for being patient! :)

chore: remove PyTorch 2.5.0 checks

b91b1f1

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 21, 2024

RdoubleA reviewed Oct 22, 2024

View reviewed changes

JP-sDEV added 3 commits November 14, 2024 09:25

chore: remove whitespace

da5178f

chore: remove fallback warning for Pytorch < 2.5

e2cfa3d

chore: rebase with main

4408138

RdoubleA reviewed Nov 14, 2024

View reviewed changes

SalmanMohammadi reviewed Nov 14, 2024

View reviewed changes

torchtune/modules/attention_utils.py Show resolved Hide resolved

docs: update to remove optional use_streams

9970128

refactor: remove pytorch 2.5.0.dev check

4027c2e

RdoubleA reviewed Nov 15, 2024

View reviewed changes

re-add flex attention check to test

37d5a01

RdoubleA approved these changes Nov 18, 2024

View reviewed changes

RdoubleA merged commit 65d07c9 into pytorch:main Nov 18, 2024
17 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

chore: remove PyTorch 2.5.0 checks #1877

chore: remove PyTorch 2.5.0 checks #1877

JP-sDEV commented Oct 21, 2024

pytorch-bot bot commented Oct 21, 2024 •

edited

Loading

RdoubleA left a comment

RdoubleA Oct 22, 2024

RdoubleA Nov 14, 2024

SalmanMohammadi Nov 14, 2024

SalmanMohammadi Nov 14, 2024

RdoubleA Oct 22, 2024

RdoubleA Nov 14, 2024

RdoubleA Oct 22, 2024

ebsmothers Oct 25, 2024

SalmanMohammadi Nov 5, 2024

ebsmothers Nov 5, 2024

SalmanMohammadi Nov 5, 2024

JP-sDEV commented Oct 22, 2024

RdoubleA commented Nov 12, 2024

JP-sDEV commented Nov 12, 2024

RdoubleA Nov 14, 2024

JP-sDEV Nov 14, 2024 •

edited

Loading

RdoubleA Nov 14, 2024 •

edited

Loading

JP-sDEV Nov 14, 2024

RdoubleA Nov 14, 2024

JP-sDEV Nov 14, 2024

RdoubleA left a comment

RdoubleA commented Nov 14, 2024

SalmanMohammadi commented Nov 14, 2024

RdoubleA Nov 15, 2024

RdoubleA commented Nov 15, 2024

JP-sDEV commented Nov 18, 2024

RdoubleA commented Nov 18, 2024

codecov-commenter commented Nov 18, 2024 •

edited

Loading

RdoubleA commented Nov 18, 2024

JP-sDEV commented Nov 19, 2024

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

	after torch-2.5.0.]. Default: True.
	after torch-2.5.0. Default: True.

chore: remove PyTorch 2.5.0 checks #1877

chore: remove PyTorch 2.5.0 checks #1877

Conversation

JP-sDEV commented Oct 21, 2024

Context

Changelog

Test plan

UX

pytorch-bot bot commented Oct 21, 2024 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1877

❗ 2 Active SEVs

✅ No Failures

RdoubleA left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JP-sDEV commented Oct 22, 2024

RdoubleA commented Nov 12, 2024

JP-sDEV commented Nov 12, 2024

Choose a reason for hiding this comment

JP-sDEV Nov 14, 2024 • edited Loading

Choose a reason for hiding this comment

RdoubleA Nov 14, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

RdoubleA left a comment

Choose a reason for hiding this comment

RdoubleA commented Nov 14, 2024

SalmanMohammadi commented Nov 14, 2024

Choose a reason for hiding this comment

RdoubleA commented Nov 15, 2024

JP-sDEV commented Nov 18, 2024

RdoubleA commented Nov 18, 2024

codecov-commenter commented Nov 18, 2024 • edited Loading

Codecov Report

RdoubleA commented Nov 18, 2024

JP-sDEV commented Nov 19, 2024

pytorch-bot bot commented Oct 21, 2024 •

edited

Loading

JP-sDEV Nov 14, 2024 •

edited

Loading

RdoubleA Nov 14, 2024 •

edited

Loading

codecov-commenter commented Nov 18, 2024 •

edited

Loading