12345678910111213141516171819202122232425262728293031323334353637383940414243 |
- // Copyright (c) Microsoft Corporation.
- // SPDX-License-Identifier: Apache-2.0
- // DeepSpeed Team
- #include "cpu_lion.h"
- // C++ interface
- void multi_tensor_lion(int chunk_size,
- at::Tensor noop_flag,
- std::vector<std::vector<at::Tensor>> tensor_lists, /*gpmv*/
- const float lr,
- const float beta1,
- const float beta2,
- const int step,
- const int mode,
- const float weight_decay)
- {
- static bool initialized = false;
- if (!initialized) {
- create_lion_optimizer(0);
- initialized = true;
- }
- for (int i = 0; i < tensor_lists[0].size(); i++) {
- ds_lion_step(0,
- step,
- lr,
- beta1,
- beta2,
- weight_decay,
- tensor_lists[1][i],
- tensor_lists[0][i],
- tensor_lists[2][i]);
- }
- }
- PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
- {
- m.def("multi_tensor_lion",
- &multi_tensor_lion,
- "Compute and apply gradient update to parameters for Lion optimizer");
- }
|