Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Cherry-pick fused_adam_cuda related patches and add its unit tests to the ROCm extension test script #93

Open
wants to merge 5 commits into
base: master
Choose a base branch
from
Open

Cherry-pick fused_adam_cuda related patches and add its unit tests to the ROCm extension test script #93

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a9983db
Add dist_adam test to the extension unit test script
Sep 9, 2022
3126ca1
Add distributed_fused_adam_v2 and _v3 back
Sep 9, 2022
63789df
Store optimizer state_dict in uint8 tensor instead of bytearray (#1446)
timmoon10 Aug 10, 2022
cb0bfa2
Optimize distributed Adam kernels and implementation (#1463)
timmoon10 Aug 25, 2022
a6e45fb
Revert "Add distributed_fused_adam_v2 and _v3 back"
Sep 9, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 29 additions & 9 deletions apex/contrib/csrc/optimizers/multi_tensor_distopt_adam.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,37 @@ void multi_tensor_fused_adam_cuda(
int chunk_size,
at::Tensor noop_flag,
std::vector<std::vector<at::Tensor>> tensor_lists,
at::Tensor per_tensor_beta1,
at::Tensor per_tensor_beta2,
at::Tensor per_tensor_bias_correction,
at::Tensor per_tensor_eps,
at::Tensor per_tensor_weight_decay,
at::Tensor grad_scale,
float lr,
float grad_scale,
float beta1,
float beta2,
float eps,
int step,
int mode);
int mode,
int bias_correction,
float weight_decay);

void multi_tensor_fused_adam_with_param_remainders_cuda(
int chunk_size,
at::Tensor noop_flag,
std::vector<std::vector<at::Tensor>> tensor_lists,
at::Tensor grad_scale,
float lr,
float beta1,
float beta2,
float eps,
int step,
int mode,
int bias_correction,
float weight_decay);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("multi_tensor_fused_adam", &multi_tensor_fused_adam_cuda,
"Multi tensor Adam optimized CUDA implementation.");
m.def("multi_tensor_fused_adam",
&multi_tensor_fused_adam_cuda,
"CUDA kernels for multi-tensor Adam, "
"with param copy");
m.def("multi_tensor_fused_adam_with_param_remainders",
&multi_tensor_fused_adam_with_param_remainders_cuda,
"CUDA kernel for multi-tensor Adam, "
"with stored param remainders and param copy");
}
Loading