pytorch/torch/optim
haozhe.zhu 1c3fe84033 [optim] add fused_adagrad support for CPU device (#124905)
Support fused_sgd_kernel support for CPU.

## Bench result:
32 core/sockets ICX
Test Scripts:
https://gist.github.com/zhuhaozhe/79e842e0a6e25d6d7fa1e4598807272c
https://gist.github.com/zhuhaozhe/b4c6998a509dcea1796dd05b3005c969
```
Tensor Size: 262144, Num Tensor 4, Num Threads: 1
_single_tensor_adagrad time: 0.2500 seconds
_fused_adagrad time: 0.0933 seconds
Tensor Size: 4194304, Num Tensor 32, Num Threads: 32
_single_tensor_adagrad time: 2.8819 seconds
_fused_adagrad time: 1.7591 seconds
```
## Test Plan:
```
python test_optim.py -k test_fused_matches_forloop
python test_optim.py -k test_fused_large_tensor
python test_optim.py -k test_can_load_older_state_dict
python test_optim.py -k test_grad_scaling_autocast_fused_optimizers
python test_torch.py -k test_grad_scaling_autocast_fused
python test_torch.py -k test_params_invalidated_with_grads_invalidated_between_unscale_and_step
```

Co-authored-by: Jane (Yuan) Xu <31798555+janeyx99@users.noreply.github.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/124905
Approved by: https://github.com/jgong5, https://github.com/janeyx99
2024-05-13 01:16:20 +00:00
..
_multi_tensor Part 1: UFMT partial files in torch/optim due to the pr-sanity-checks (#124053) 2024-04-16 03:17:18 +00:00
__init__.py Part 1: UFMT partial files in torch/optim due to the pr-sanity-checks (#124053) 2024-04-16 03:17:18 +00:00
__init__.pyi
_functional.py Part 1: UFMT partial files in torch/optim due to the pr-sanity-checks (#124053) 2024-04-16 03:17:18 +00:00
adadelta.py Disable dynamo on functional optims if capturable=False (#123619) 2024-05-07 22:17:01 +00:00
adagrad.py [optim] add fused_adagrad support for CPU device (#124905) 2024-05-13 01:16:20 +00:00
adam.py Disable dynamo on functional optims if capturable=False (#123619) 2024-05-07 22:17:01 +00:00
adam.pyi
adamax.py Disable dynamo on functional optims if capturable=False (#123619) 2024-05-07 22:17:01 +00:00
adamax.pyi
adamw.py Disable dynamo on functional optims if capturable=False (#123619) 2024-05-07 22:17:01 +00:00
adamw.pyi
asgd.py Fix user warning for tensor LR (#123752) 2024-05-09 00:52:43 +00:00
asgd.pyi
lbfgs.py Part 1: UFMT partial files in torch/optim due to the pr-sanity-checks (#124053) 2024-04-16 03:17:18 +00:00
lbfgs.pyi
lr_scheduler.py add typing in torch.optim.lr_scheduler (#125556) 2024-05-10 19:28:00 +00:00
nadam.py Disable dynamo on functional optims if capturable=False (#123619) 2024-05-07 22:17:01 +00:00
nadam.pyi
optimizer.py [optim] add fused_adagrad support for CPU device (#124905) 2024-05-13 01:16:20 +00:00
radam.py Disable dynamo on functional optims if capturable=False (#123619) 2024-05-07 22:17:01 +00:00
radam.pyi
rmsprop.py Disable dynamo on functional optims if capturable=False (#123619) 2024-05-07 22:17:01 +00:00
rmsprop.pyi
rprop.py Disable dynamo on functional optims if capturable=False (#123619) 2024-05-07 22:17:01 +00:00
rprop.pyi
sgd.py Add LR as tensor tests (#123750) 2024-05-01 04:46:49 +00:00
sgd.pyi
sparse_adam.py Part 3: UFMT fix the rest files in torch/optim due to the pr-sanity-checks (#124055) 2024-04-16 03:22:39 +00:00
sparse_adam.pyi
swa_utils.py Part 3: UFMT fix the rest files in torch/optim due to the pr-sanity-checks (#124055) 2024-04-16 03:22:39 +00:00