|
#include <torch/extension.h> |
|
|
|
void multi_tensor_scale_cuda( |
|
int chunk_size, |
|
at::Tensor noop_flag, |
|
std::vector<std::vector<at::Tensor>> tensor_lists, |
|
float scale); |
|
|
|
void multi_tensor_axpby_cuda( |
|
int chunk_size, |
|
at::Tensor noop_flag, |
|
std::vector<std::vector<at::Tensor>> tensor_lists, |
|
float a, |
|
float b, |
|
int arg_to_check); |
|
|
|
std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda( |
|
int chunk_size, |
|
at::Tensor noop_flag, |
|
std::vector<std::vector<at::Tensor>> tensor_lists, |
|
at::optional<bool> per_tensor_python); |
|
|
|
void multi_tensor_lamb_stage1_cuda( |
|
int chunk_size, |
|
at::Tensor noop_flag, |
|
std::vector<std::vector<at::Tensor>> tensor_lists, |
|
at::Tensor per_tensor_decay, |
|
const int step, |
|
const float beta1, |
|
const float beta2, |
|
const float epsilon, |
|
const float global_grad_norm, |
|
const float max_global_grad_norm); |
|
|
|
void multi_tensor_lamb_stage2_cuda( |
|
int chunk_size, |
|
at::Tensor noop_flag, |
|
std::vector<std::vector<at::Tensor>> tensor_lists, |
|
at::Tensor per_tensor_param_norm, |
|
at::Tensor per_tensor_update_norm, |
|
const float step_size); |
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { |
|
m.def("multi_tensor_scale", &multi_tensor_scale_cuda, |
|
"Fused overflow check + scale for a list of contiguous tensors"); |
|
m.def("multi_tensor_axpby", &multi_tensor_axpby_cuda, |
|
"out = a*x + b*y for a list of contiguous tensors"); |
|
m.def("multi_tensor_l2norm", &multi_tensor_l2norm_cuda, |
|
"Computes L2 norm for a list of contiguous tensors"); |
|
m.def("multi_tensor_lamb_stage1_cuda", &multi_tensor_lamb_stage1_cuda, |
|
"Computes update part of LAMB optimizer"); |
|
m.def("multi_tensor_lamb_stage2_cuda", &multi_tensor_lamb_stage2_cuda, |
|
"Completes application of gradient to parameters for LAMB optimizer"); |
|
} |
|
|