amp_C_frontend.cpp 3.87 KB
Newer Older
1
2
3
4
5
6
7
8
#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);

Simon Layton's avatar
Simon Layton committed
9
10
11
12
13
14
15
16
17
void multi_tensor_sgd_cuda(
  int chunk_size,
  at::Tensor noop_flag,
  std::vector<std::vector<at::Tensor>> tensor_lists,
  float wd,
  float momentum,
  float dampening,
  float lr,
  bool nesterov,
18
  bool first_run,
19
20
  bool wd_after_momentum,
  float scale);
Simon Layton's avatar
Simon Layton committed
21

22
23
24
25
26
void multi_tensor_axpby_cuda(
  int chunk_size,
  at::Tensor noop_flag,
  std::vector<std::vector<at::Tensor>> tensor_lists,
  float a,
27
28
  float b,
  int arg_to_check);
29

30
std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda(
31
32
  int chunk_size,
  at::Tensor noop_flag,
33
34
  std::vector<std::vector<at::Tensor>> tensor_lists,
  at::optional<bool> per_tensor_python);
35

36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
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);

56
void multi_tensor_adam_cuda(
57
58
59
  int chunk_size,
  at::Tensor noop_flag,
  std::vector<std::vector<at::Tensor>> tensor_lists,
60
  const float lr,
61
62
63
  const float beta1,
  const float beta2,
  const float epsilon,
64
  const int step,
65
  const int mode,
66
67
  const int bias_correction,
  const float weight_decay);
68

Andrew Tulloch's avatar
Andrew Tulloch committed
69
70
71
72
73
74
75
76
77
78
79

void multi_tensor_adagrad_cuda(
  int chunk_size,
  at::Tensor noop_flag,
  std::vector<std::vector<at::Tensor>> tensor_lists,
  const float lr,
  const float epsilon,
  const int mode,
  const float weight_decay);


80
void multi_tensor_novograd_cuda(
81
82
83
  int chunk_size,
  at::Tensor noop_flag,
  std::vector<std::vector<at::Tensor>> tensor_lists,
84
  at::Tensor grad_norms,
85
86
87
88
89
90
  const float lr,
  const float beta1,
  const float beta2,
  const float epsilon,
  const int step,
  const int bias_correction,
91
92
  const float weight_decay,
  const int grad_averaging,
93
  const int mode,
94
  const int norm_type);
95

96
void multi_tensor_lamb_cuda(
97
98
99
100
101
102
103
104
105
106
107
  int chunk_size,
  at::Tensor noop_flag,
  std::vector<std::vector<at::Tensor>> tensor_lists,
  const float lr,
  const float beta1,
  const float beta2,
  const float epsilon,
  const int step,
  const int bias_correction,
  const float weight_decay,
  const int grad_averaging,
108
  const int mode,
109
  const float max_grad_norm);
110

111
112
113
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");
Simon Layton's avatar
Simon Layton committed
114
115
  m.def("multi_tensor_sgd", &multi_tensor_sgd_cuda,
        "Fused SGD optimizer for list of contiguous tensors");
116
117
  m.def("multi_tensor_axpby", &multi_tensor_axpby_cuda,
        "out = a*x + b*y for a list of contiguous tensors");
118
119
  m.def("multi_tensor_l2norm", &multi_tensor_l2norm_cuda,
        "Computes L2 norm for a list of contiguous tensors");
120
121
122
123
  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");
124
125
  m.def("multi_tensor_adam", &multi_tensor_adam_cuda,
        "Compute and apply gradient update to parameters for Adam optimizer");
Andrew Tulloch's avatar
Andrew Tulloch committed
126
127
  m.def("multi_tensor_adagrad", &multi_tensor_adagrad_cuda,
        "Compute and apply gradient update to parameters for Adam optimizer");
128
129
  m.def("multi_tensor_novograd", &multi_tensor_novograd_cuda,
        "Compute and apply gradient update to parameters for Adam optimizer");
130
131
  m.def("multi_tensor_lamb", &multi_tensor_lamb_cuda,
        "Computes and apply update for LAMB optimizer");
132
}