Skip to content

Conversation

justinvyu
Copy link
Contributor

@justinvyu justinvyu commented Mar 4, 2025

Summary

The Ray Train torch AMP wrapper does not provide any value and has a worse user experience and is using deprecated APIs (torch.cuda.amp.GradScaler). It's easier to just have the user manage the mixed precision training themselves, as the native torch API is already pretty lightweight. Therefore, deprecate this API and its related tests.

As a side-node, the unit test here is testing whether or not the torch AMP API works, which is a pattern that we should avoid in the future since we don't need to duplicate tests that should already exist within torch.

Here's the migration message printed to users:

The `accelerate`, `backward`, and `prepare_optimizer` utility methods in the `ray.train.torch` module are deprecated 
and will be removed in a future release. Please use the native PyTorch mixed precision API directly, or a library such as 
Lightning or HuggingFace Transformers/Accelerate.
See this issue for more context: https://github.com/ray-project/ray/issues/49454

@justinvyu justinvyu changed the title [train][ci] Remove torch amp wrapper utility [train] Remove torch amp wrapper utility Mar 4, 2025
@justinvyu justinvyu changed the title [train] Remove torch amp wrapper utility [train] Deprecate torch amp wrapper utility Mar 4, 2025
@justinvyu justinvyu changed the title [train] Deprecate torch amp wrapper utility [train] Deprecate torch amp wrapper utilities Mar 4, 2025
@justinvyu justinvyu enabled auto-merge (squash) March 4, 2025 21:43
@github-actions github-actions bot added the go add ONLY when ready to merge, run all tests label Mar 4, 2025
Signed-off-by: Justin Yu <[email protected]>
@github-actions github-actions bot disabled auto-merge March 4, 2025 23:23
@justinvyu justinvyu enabled auto-merge (squash) March 4, 2025 23:40
@justinvyu justinvyu merged commit 528caf4 into ray-project:master Mar 5, 2025
6 checks passed
@justinvyu justinvyu deleted the remove_torch_amp_test branch March 5, 2025 01:36
abrarsheikh pushed a commit that referenced this pull request Mar 8, 2025
The Ray Train torch AMP wrapper does not provide any value and has a
worse user experience and is using deprecated APIs
(`torch.cuda.amp.GradScaler`). It's easier to just have the user manage
the mixed precision training themselves, as the native torch API is
already pretty lightweight. Therefore, deprecate this API and its
related tests.

---------

Signed-off-by: Justin Yu <[email protected]>
Signed-off-by: Abrar Sheikh <[email protected]>
park12sj pushed a commit to park12sj/ray that referenced this pull request Mar 18, 2025
The Ray Train torch AMP wrapper does not provide any value and has a
worse user experience and is using deprecated APIs
(`torch.cuda.amp.GradScaler`). It's easier to just have the user manage
the mixed precision training themselves, as the native torch API is
already pretty lightweight. Therefore, deprecate this API and its
related tests.

---------

Signed-off-by: Justin Yu <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
community-backlog go add ONLY when ready to merge, run all tests
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants