Commit Graph

8 Commits

Author SHA1 Message Date
Isalia20
49f6cce736 [MPS] grad scaler (#150255)
Fixes #142397

Basic implementation is done. What's left:
- [x] Different dtype/device tensors in the TensorList
- [x] fast path for grouping the foreach kernel
- [x] Tests

Regarding tests, I found some tests in `test/test_torch.py` for GradScaler but I couldn't figure out what is the best way to enable the test for MPS device.

By removing `@onlyNativeDeviceTypes`, one enables the tests for MPS but also enables tests for all other devices which are not included in the native device types. If I put:
`instantiate_device_type_tests(TestTorchDeviceType, globals(), allow_mps=True)`

This enables lots of tests in that class for MPS which were not(?) being tested before? This part needs some clarification

Pull Request resolved: https://github.com/pytorch/pytorch/pull/150255
Approved by: https://github.com/malfet

Co-authored-by: Nikita Shulga <2453524+malfet@users.noreply.github.com>
2025-04-06 17:06:55 +00:00
Aaron Orenstein
f2cfe8b59f PEP585 update - mostly toplevels (#145178)
See #145101 for details.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/145178
Approved by: https://github.com/bobrenjc93
2025-01-22 02:21:14 +00:00
yousufmo
305ba62906 Add support to GradScaler for respecting an already set grad_scale value (#123429)
Fixes #123428

Co-authored-by: Yousuf Mohamed-Ahmed <youmed.tech@gmail.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/123429
Approved by: https://github.com/ezyang
2024-06-27 22:40:54 +00:00
Aaron Orenstein
afe15d2d2f Flip default value for mypy disallow_untyped_defs [3/11] (#127840)
See #127836 for details.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/127840
Approved by: https://github.com/oulgen
2024-06-08 18:28:01 +00:00
Aaron Gokaslan
9c4fc5fa34 [BE][Ez]: Fix minor potential perf regression from #123960 (#124013)
The `non_blocking` arg here is useless if the values are all eagerly consumed, so revert the change.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/124013
Approved by: https://github.com/ezyang
2024-04-15 16:51:45 +00:00
Aaron Gokaslan
1d6c5972c1 [BE]: Optimize min/max/sum comprehensions C419 (#123960)
Automatic fixes that replaces certain list comprehensions with generator ones where appropriate so that they are immediately consumed. This is preview functionality in ruff for rule C419 and it was automatically applied.

Co-authored-by: Nikita Shulga <2453524+malfet@users.noreply.github.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/123960
Approved by: https://github.com/malfet
2024-04-12 23:54:15 +00:00
wangjiangben-hw
354a436d96 Remove device assert in Gradscaler (#119362)
Fixes #119358

Co-authored-by: Edward Z. Yang <ezyang@mit.edu>
Co-authored-by: ydwu4 <ydwu2014@gmail.com>
Co-authored-by: PyTorch UpdateBot <pytorchupdatebot@users.noreply.github.com>
Co-authored-by: Bin Bao <binbao@meta.com>
Co-authored-by: Shuqiang Zhang <sqzhang@meta.com>
Co-authored-by: Adnan Akhundov <aakhundov@meta.com>
Co-authored-by: Ting Lu <tingl@nvidia.com>
Co-authored-by: Yang Chen <yangche@fb.com>
Co-authored-by: cyy <cyyever@outlook.com>
Co-authored-by: Animesh Jain <anijain@umich.edu>
Co-authored-by: Jason Ansel <jansel@meta.com>
Co-authored-by: Eddie Yan <eddiey@nvidia.com>
Co-authored-by: wz337 <wz337@cornell.edu>
Co-authored-by: Xuehai Pan <XuehaiPan@pku.edu.cn>
Co-authored-by: Anthony Alayo <anthony.alayo@applovin.com>
Co-authored-by: leslie-fang-intel <leslie.fang@intel.com>
Co-authored-by: Yifu Wang <yifu@fb.com>
Co-authored-by: Yukio Siraichi <yukio.siraichi@gmail.com>
Co-authored-by: atalman <atalman@fb.com>
Co-authored-by: PyTorch MergeBot <pytorchmergebot@users.noreply.github.com>
Co-authored-by: Jeff Daily <jeff.daily@amd.com>
Co-authored-by: Nikita Shulga <2453524+malfet@users.noreply.github.com>
Co-authored-by: haozhe.zhu <haozhe.zhu@intel.com>
Co-authored-by: lezcano <lezcano-93@hotmail.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/119362
Approved by: https://github.com/ezyang
2024-02-22 08:02:18 +00:00
CaoE
bacbad5bc9 add GradScaler on CPU (#109993)
Step 2 of https://github.com/pytorch/pytorch/issues/111559.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/109993
Approved by: https://github.com/jgong5, https://github.com/ezyang
2024-01-29 23:42:35 +00:00