wehub-resource-sync
279 lines
10 KiB
C++
279 lines
10 KiB
C++
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
#pragma once
|
|
|
|
#define NOMINMAX // Windows idiosyncrasy
|
|
// https://stackoverflow.com/questions/4913922/possible-problems-with-nominmax-on-visual-c
|
|
|
|
#include <stdio.h>
|
|
#include <torch/extension.h>
|
|
#include <cassert>
|
|
#include <cstdint>
|
|
#include "simd.h"
|
|
|
|
#define STEP(SPAN) \
|
|
template <typename ds_params_precision_t, typename ds_state_precision_t> \
|
|
void Step_##SPAN(ds_params_precision_t* _params, \
|
|
ds_params_precision_t* grads, \
|
|
ds_state_precision_t* _exp_avg, \
|
|
ds_state_precision_t* _exp_avg_sq, \
|
|
size_t _param_size, \
|
|
bool parallel = true);
|
|
|
|
class Adam_Optimizer {
|
|
public:
|
|
Adam_Optimizer(float alpha = 1e-3,
|
|
float betta1 = 0.9,
|
|
float betta2 = 0.999,
|
|
float eps = 1e-8,
|
|
float weight_decay = 0,
|
|
bool adamw_mode = true)
|
|
: _alpha(alpha),
|
|
_betta1(betta1),
|
|
_betta2(betta2),
|
|
_eps(eps),
|
|
_weight_decay(weight_decay),
|
|
_betta1_t(1.0),
|
|
_betta2_t(1.0),
|
|
_step(0),
|
|
_adamw_mode(adamw_mode)
|
|
{
|
|
}
|
|
~Adam_Optimizer() {}
|
|
|
|
#if defined(__AVX512__) or defined(__AVX256__)
|
|
template <int span, typename ds_params_precision_t, typename ds_state_precision_t>
|
|
void Step_AVX(size_t* rounded_size,
|
|
ds_params_precision_t* _params,
|
|
ds_params_precision_t* grads,
|
|
ds_state_precision_t* _exp_avg,
|
|
ds_state_precision_t* _exp_avg_sq,
|
|
size_t param_size,
|
|
bool parallel = true);
|
|
#endif
|
|
STEP(1)
|
|
STEP(4)
|
|
STEP(8)
|
|
inline void IncrementStep(size_t step, float beta1, float beta2)
|
|
{
|
|
if (beta1 != _betta1 || beta2 != _betta2) {
|
|
_step = step;
|
|
_betta1 = beta1;
|
|
_betta2 = beta2;
|
|
_betta1_t = std::pow(_betta1, step);
|
|
_betta2_t = std::pow(_betta2, step);
|
|
} else {
|
|
if (step == _step + 1) { // first optimizer step increase
|
|
_step++;
|
|
_betta1_t *= _betta1;
|
|
_betta2_t *= _betta2;
|
|
} else if (step ==
|
|
_step) { // no need to update step; beta1_t and beta2_t already updated
|
|
return;
|
|
} else { // support step increase not equal to 1
|
|
_betta1_t = std::pow(_betta1, step);
|
|
_betta2_t = std::pow(_betta2, step);
|
|
_step = step;
|
|
}
|
|
}
|
|
}
|
|
inline void update_state(float lr, float epsilon, float weight_decay, bool bias_correction)
|
|
{
|
|
_alpha = lr;
|
|
_eps = epsilon;
|
|
_weight_decay = weight_decay;
|
|
|
|
_bias_correction1 = 1.0f;
|
|
_bias_correction2 = 1.0f;
|
|
if (bias_correction == 1) {
|
|
_bias_correction1 = 1 - _betta1_t;
|
|
_bias_correction2 = 1 / sqrt(1 - _betta2_t);
|
|
}
|
|
}
|
|
|
|
private:
|
|
float _alpha;
|
|
float _betta1;
|
|
float _betta2;
|
|
float _eps;
|
|
float _weight_decay;
|
|
|
|
float _betta1_t;
|
|
float _betta2_t;
|
|
size_t _step;
|
|
|
|
float _bias_correction1;
|
|
float _bias_correction2;
|
|
|
|
bool _adamw_mode;
|
|
};
|
|
|
|
#if defined(__AVX512__) or defined(__AVX256__)
|
|
template <int span, typename ds_params_precision_t, typename ds_state_precision_t>
|
|
void Adam_Optimizer::Step_AVX(size_t* rounded_size,
|
|
ds_params_precision_t* _params,
|
|
ds_params_precision_t* grads,
|
|
ds_state_precision_t* _exp_avg,
|
|
ds_state_precision_t* _exp_avg_sq,
|
|
size_t _param_size,
|
|
bool parallel)
|
|
{
|
|
#if !defined(__AVX512__)
|
|
if (std::is_same_v<ds_params_precision_t, c10::BFloat16> ||
|
|
std::is_same_v<ds_state_precision_t, c10::BFloat16>) {
|
|
return;
|
|
}
|
|
#endif
|
|
size_t new_rounded_size = 0;
|
|
|
|
AVX_Data betta1_4;
|
|
betta1_4.data = SIMD_SET(_betta1);
|
|
AVX_Data betta2_4;
|
|
betta2_4.data = SIMD_SET(_betta2);
|
|
|
|
float betta1_minus1 = 1 - _betta1;
|
|
float betta2_minus1 = 1 - _betta2;
|
|
AVX_Data betta1_minus1_4;
|
|
betta1_minus1_4.data = SIMD_SET(betta1_minus1);
|
|
AVX_Data betta2_minus1_4;
|
|
betta2_minus1_4.data = SIMD_SET(betta2_minus1);
|
|
|
|
AVX_Data bias2_sqrt;
|
|
bias2_sqrt.data = SIMD_SET(_bias_correction2);
|
|
|
|
AVX_Data eps_4;
|
|
eps_4.data = SIMD_SET(_eps);
|
|
|
|
float step_size = -1 * _alpha / _bias_correction1;
|
|
AVX_Data step_size_4;
|
|
step_size_4.data = SIMD_SET(step_size);
|
|
|
|
float w_decay = -1 * _alpha * _weight_decay;
|
|
AVX_Data weight_decay4;
|
|
if (_weight_decay > 0)
|
|
weight_decay4.data = (_adamw_mode ? SIMD_SET(w_decay) : SIMD_SET(_weight_decay));
|
|
new_rounded_size = ROUND_DOWN(_param_size, SIMD_WIDTH * span);
|
|
for (size_t t = 0; t < new_rounded_size; t += TILE) {
|
|
size_t copy_size = TILE;
|
|
if ((t + TILE) > new_rounded_size) copy_size = new_rounded_size - t;
|
|
size_t offset = copy_size + t;
|
|
#pragma omp parallel for if (parallel)
|
|
for (size_t i = t; i < offset; i += SIMD_WIDTH * span) {
|
|
AVX_Data grad_4[span];
|
|
simd_load<span>(grad_4, grads + i);
|
|
|
|
AVX_Data momentum_4[span];
|
|
simd_load<span>(momentum_4, _exp_avg + i);
|
|
|
|
AVX_Data variance_4[span];
|
|
simd_load<span>(variance_4, _exp_avg_sq + i);
|
|
|
|
AVX_Data param_4[span];
|
|
simd_load<span>(param_4, _params + i);
|
|
|
|
if (_weight_decay > 0 && !_adamw_mode) {
|
|
simd_fma<span>(grad_4, param_4, weight_decay4, grad_4);
|
|
}
|
|
|
|
simd_mul<span>(momentum_4, momentum_4, betta1_4);
|
|
simd_fma<span>(momentum_4, grad_4, betta1_minus1_4, momentum_4);
|
|
simd_mul<span>(variance_4, variance_4, betta2_4);
|
|
simd_mul<span>(grad_4, grad_4, grad_4);
|
|
simd_fma<span>(variance_4, grad_4, betta2_minus1_4, variance_4);
|
|
simd_sqrt<span>(grad_4, variance_4);
|
|
simd_fma<span>(grad_4, grad_4, bias2_sqrt, eps_4);
|
|
simd_div<span>(grad_4, momentum_4, grad_4);
|
|
|
|
if (_weight_decay > 0 && _adamw_mode) {
|
|
simd_fma<span>(param_4, param_4, weight_decay4, param_4);
|
|
}
|
|
|
|
simd_fma<span>(param_4, grad_4, step_size_4, param_4);
|
|
|
|
simd_store<span>(_params + i, param_4);
|
|
simd_store<span>(_exp_avg + i, momentum_4);
|
|
simd_store<span>(_exp_avg_sq + i, variance_4);
|
|
}
|
|
}
|
|
*rounded_size = new_rounded_size;
|
|
}
|
|
#endif
|
|
|
|
int create_adam_optimizer(int optimizer_id,
|
|
float alpha = 1e-3,
|
|
float betta1 = 0.9,
|
|
float betta2 = 0.999,
|
|
float eps = 1e-8,
|
|
float weight_decay = 0,
|
|
bool adamw_mode = true,
|
|
bool should_log = false);
|
|
|
|
int ds_adam_step(int optimizer_id,
|
|
size_t step,
|
|
float lr,
|
|
float beta1,
|
|
float beta2,
|
|
float epsilon,
|
|
float weight_decay,
|
|
bool bias_correction,
|
|
torch::Tensor& params,
|
|
torch::Tensor& grads,
|
|
torch::Tensor& exp_avg,
|
|
torch::Tensor& exp_avg_sq);
|
|
|
|
int ds_adam_rollback(int optimizer_id,
|
|
size_t step,
|
|
float lr,
|
|
float beta1,
|
|
float beta2,
|
|
float epsilon,
|
|
float weight_decay,
|
|
bool bias_correction,
|
|
torch::Tensor& params,
|
|
torch::Tensor& grads,
|
|
torch::Tensor& exp_avg,
|
|
torch::Tensor& exp_avg_sq);
|
|
|
|
int destroy_adam_optimizer(int optimizer_id);
|
|
|
|
// ZenFlowAdam: the native CPU Adam backing ZenFlow's overlapped optimizer step. The handle
|
|
// indexes a pinned thread pool; the optimizer runs in a dedicated process (run_worker) and
|
|
// is driven from the main process through the shared-memory control block below.
|
|
int zenflow_adam_create(int optimizer_id, std::vector<int> zf_affinity);
|
|
|
|
void zenflow_adam_register_group(int handle,
|
|
torch::Tensor param,
|
|
torch::Tensor grad0,
|
|
torch::Tensor grad1,
|
|
torch::Tensor exp_avg0,
|
|
torch::Tensor exp_avg1,
|
|
torch::Tensor exp_avg_sq0,
|
|
torch::Tensor exp_avg_sq1,
|
|
torch::Tensor stale);
|
|
|
|
void zenflow_adam_destroy(int handle);
|
|
|
|
#if defined(__linux__)
|
|
// The optimizer runs in a separate process and coordinates with the main process through two
|
|
// process-shared semaphores in a shared-memory control block. ctrl_size/ctrl_init/ctrl_exit
|
|
// set it up and tear it down; the worker process loops in run_worker; the main process drives
|
|
// each step with submit (non-blocking) / wait.
|
|
int64_t zenflow_adam_ctrl_size();
|
|
void zenflow_adam_ctrl_init(uintptr_t control_ptr, int num_groups);
|
|
void zenflow_adam_run_worker(int handle, uintptr_t control_ptr);
|
|
void zenflow_adam_submit(uintptr_t control_ptr,
|
|
int now_state,
|
|
int64_t step,
|
|
std::vector<float> lr,
|
|
std::vector<float> beta1,
|
|
std::vector<float> beta2,
|
|
std::vector<float> eps,
|
|
std::vector<float> weight_decay,
|
|
std::vector<uint8_t> bias_correction);
|
|
bool zenflow_adam_wait(uintptr_t control_ptr, double timeout_s);
|
|
void zenflow_adam_ctrl_exit(uintptr_t control_ptr);
|
|
#endif
|