Olga Gerasimova 2efa5eaa65 swa avoid stream sync (#157705) ...

Summary:
When AveragedModel updates_parameters it calls self.n_averaged == 0 for each parameter, where n_averated is a buffer on GPU. Moving check before the cycle to call sync once

It improves update_parameter from 74ms to 57ms ~22% improvement
{F1980011097}
{F1980011111}

Test Plan:
CI

Rollback Plan:

Differential Revision: D77723025

Pull Request resolved: https://github.com/pytorch/pytorch/pull/157705
Approved by: https://github.com/albanD, https://github.com/Skylion007, https://github.com/janeyx99

2025-07-07 20:47:35 +00:00

..

_multi_tensor

[BE][PYFMT] migrate PYFMT for {torch,test}/{nn,optim}/** to ruff format (#144548)

2025-06-14 11:27:04 +00:00

__init__.py

[BE][optim] Make pyright recognize exported symbols (#135043)

2024-09-04 21:53:46 +00:00

_adafactor.py

updated adafactor doc #154862 (#155248)

2025-06-27 23:23:19 +00:00

_functional.py

[BE][PYFMT] migrate PYFMT for {torch,test}/{nn,optim}/** to ruff format (#144548)

2025-06-14 11:27:04 +00:00

adadelta.py

[BE][12/16] fix typos in torch/ (#156602)

2025-07-02 22:55:29 +00:00

adagrad.py

[BE][12/16] fix typos in torch/ (#156602)

2025-07-02 22:55:29 +00:00

adam.py

[BE][12/16] fix typos in torch/ (#156602)

2025-07-02 22:55:29 +00:00

adamax.py

[BE][12/16] fix typos in torch/ (#156602)

2025-07-02 22:55:29 +00:00

adamw.py

PEP585 update - torch/nn torch/optim torch/package torch/profiler torch/serialization torch/sparse torch/xpu (#145175)

2025-01-21 16:57:27 +00:00

asgd.py

[BE][PYFMT] migrate PYFMT for {torch,test}/{nn,optim}/** to ruff format (#144548)

2025-06-14 11:27:04 +00:00

lbfgs.py

mypy 1.16.0 (#155821)

2025-06-14 18:18:43 +00:00

lr_scheduler.py

[BE][PYFMT] migrate PYFMT for {torch,test}/{nn,optim}/** to ruff format (#144548)

2025-06-14 11:27:04 +00:00

nadam.py

[BE][12/16] fix typos in torch/ (#156602)

2025-07-02 22:55:29 +00:00

optimizer.py

[BE][PYFMT] migrate PYFMT for {torch,test}/{nn,optim}/** to ruff format (#144548)

2025-06-14 11:27:04 +00:00

radam.py

[BE][12/16] fix typos in torch/ (#156602)

2025-07-02 22:55:29 +00:00

rmsprop.py

[BE][12/16] fix typos in torch/ (#156602)

2025-07-02 22:55:29 +00:00

rprop.py

[BE][12/16] fix typos in torch/ (#156602)

2025-07-02 22:55:29 +00:00

sgd.py

[BE][12/16] fix typos in torch/ (#156602)

2025-07-02 22:55:29 +00:00

sparse_adam.py

[BE][PYFMT] migrate PYFMT for {torch,test}/{nn,optim}/** to ruff format (#144548)

2025-06-14 11:27:04 +00:00

swa_utils.py

swa avoid stream sync (#157705)

2025-07-07 20:47:35 +00:00