wehub-resource-sync
52 lines
2.3 KiB
C++
52 lines
2.3 KiB
C++
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
#include "cpu_adam.h"
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
|
|
{
|
|
using namespace pybind11::literals;
|
|
m.def("adam_update", &ds_adam_step, "DeepSpeed CPU Adam update (C++)");
|
|
m.def("adam_rollback", &ds_adam_rollback, "DeepSpeed CPU Adam rollback (C++)");
|
|
m.def("create_adam", &create_adam_optimizer, "DeepSpeed CPU Adam (C++)");
|
|
m.def("destroy_adam", &destroy_adam_optimizer, "DeepSpeed CPU Adam destroy (C++)");
|
|
|
|
// ZenFlowAdam: the native CPU Adam backing ZenFlow's overlapped optimizer step. create /
|
|
// register_group / destroy set up the handle-indexed pinned pool (used by the worker process).
|
|
m.def("zenflow_adam_create", &zenflow_adam_create, "ZenFlowAdam create (C++)");
|
|
m.def("zenflow_adam_register_group",
|
|
&zenflow_adam_register_group,
|
|
"ZenFlowAdam register a parameter group (C++)");
|
|
m.def("zenflow_adam_destroy",
|
|
&zenflow_adam_destroy,
|
|
"ZenFlowAdam destroy (C++)",
|
|
pybind11::call_guard<pybind11::gil_scoped_release>());
|
|
|
|
#if defined(__linux__)
|
|
// The optimizer runs in a separate process, coordinated through a shared-memory semaphore
|
|
// control block. submit/wait/run_worker release the GIL so the optimizer process overlaps
|
|
// the Python training thread.
|
|
m.def(
|
|
"zenflow_adam_ctrl_size", &zenflow_adam_ctrl_size, "ZenFlowAdam control block size (C++)");
|
|
m.def("zenflow_adam_ctrl_init", &zenflow_adam_ctrl_init, "ZenFlowAdam control init (C++)");
|
|
m.def("zenflow_adam_run_worker",
|
|
&zenflow_adam_run_worker,
|
|
"ZenFlowAdam optimizer-process worker loop (C++)",
|
|
pybind11::call_guard<pybind11::gil_scoped_release>());
|
|
m.def("zenflow_adam_submit",
|
|
&zenflow_adam_submit,
|
|
"ZenFlowAdam submit an overlapped step (C++)",
|
|
pybind11::call_guard<pybind11::gil_scoped_release>());
|
|
m.def("zenflow_adam_wait",
|
|
&zenflow_adam_wait,
|
|
"ZenFlowAdam wait for a submitted step (C++)",
|
|
pybind11::call_guard<pybind11::gil_scoped_release>());
|
|
m.def("zenflow_adam_ctrl_exit",
|
|
&zenflow_adam_ctrl_exit,
|
|
"ZenFlowAdam cross-process exit (C++)",
|
|
pybind11::call_guard<pybind11::gil_scoped_release>());
|
|
#endif
|
|
}
|