diff --git a/csrc/adam/cpu_adam.cpp b/csrc/adam/cpu_adam.cpp index 72b7b2a97473..1b899bf14f4c 100644 --- a/csrc/adam/cpu_adam.cpp +++ b/csrc/adam/cpu_adam.cpp @@ -8,8 +8,34 @@ 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++)"); + // The update, the rescale and the rollback release the GIL, so that a CPU Adam step in another + // thread lets the training thread run. + m.def("adam_update", + &ds_adam_step, + "DeepSpeed CPU Adam update (C++); grad_scale multiplies each gradient as it is read", + "optimizer_id"_a, + "step"_a, + "lr"_a, + "beta1"_a, + "beta2"_a, + "epsilon"_a, + "weight_decay"_a, + "bias_correction"_a, + "params"_a, + "grads"_a, + "exp_avg"_a, + "exp_avg_sq"_a, + "grad_scale"_a = 1.0f, + pybind11::call_guard()); + m.def("adam_rescale_step", + &ds_adam_rescale_step, + "DeepSpeed CPU Adam: turn the last step, taken on grads * old_scale, into the step on " + "grads * new_scale (C++)", + pybind11::call_guard()); + m.def("adam_rollback", + &ds_adam_rollback, + "DeepSpeed CPU Adam rollback (C++)", + pybind11::call_guard()); m.def("create_adam", &create_adam_optimizer, "DeepSpeed CPU Adam (C++)"); m.def("destroy_adam", &destroy_adam_optimizer, "DeepSpeed CPU Adam destroy (C++)"); diff --git a/csrc/adam/cpu_adam_impl.cpp b/csrc/adam/cpu_adam_impl.cpp index d5912272376d..ad13629cffbc 100644 --- a/csrc/adam/cpu_adam_impl.cpp +++ b/csrc/adam/cpu_adam_impl.cpp @@ -32,9 +32,11 @@ static std::unordered_map> s_optimizers; // C++ interface -template +template void Adam_Optimizer::Step_1(ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t _param_size, @@ -62,6 +64,7 @@ void Adam_Optimizer::Step_1(ds_params_precision_t* _params, #pragma omp parallel for if (parallel) for (size_t k = t; k < offset; k++) { float grad = (float)grads[k]; + if (_grad_scale != 1.0f) { grad = grad * _grad_scale; } float param = (float)_params[k]; float momentum = _exp_avg[k]; float variance = _exp_avg_sq[k]; @@ -86,9 +89,11 @@ void Adam_Optimizer::Step_1(ds_params_precision_t* _params, } } -template +template void Adam_Optimizer::Step_4(ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t _param_size, @@ -155,9 +160,11 @@ int create_adam_optimizer(int optimizer_id, return 0; } -template +template void Adam_Optimizer::Step_8(ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t _param_size, @@ -180,7 +187,9 @@ void Adam_Optimizer::Step_8(ds_params_precision_t* _params, parallel); } -template +template void step_invoker(std::shared_ptr opt, void* _params, void* grads, @@ -190,7 +199,7 @@ void step_invoker(std::shared_ptr opt, bool parallel) { opt->Step_8((ds_params_precision_t*)(_params), - (ds_params_precision_t*)(grads), + (ds_grad_precision_t*)(grads), (ds_state_precision_t*)(_exp_avg), (ds_state_precision_t*)(_exp_avg_sq), _param_size, @@ -198,17 +207,20 @@ void step_invoker(std::shared_ptr opt, } std::map< - std::tuple, + std::tuple, std::function, void*, void*, void*, void*, size_t, bool)>> invokers; // Fill map with template functions for each type -template +template void create_invoker() { invokers[std::tuple(c10::CppTypeToScalarType(), - c10::CppTypeToScalarType())] = - step_invoker; + c10::CppTypeToScalarType(), + c10::CppTypeToScalarType())] = + step_invoker; } struct InvokerInitializer { InvokerInitializer() @@ -218,6 +230,8 @@ struct InvokerInitializer { create_invoker(); create_invoker(); create_invoker(); + // fp32 weights on bf16 gradients: SuperOffload copies them to the CPU in bf16. + create_invoker(); } } _invoker_initializer; @@ -231,11 +245,13 @@ void invoke(std::shared_ptr opt, { c10::ScalarType params_type = at::typeMetaToScalarType(params.options().dtype()); c10::ScalarType state_type = at::typeMetaToScalarType(exp_avg.options().dtype()); + c10::ScalarType grad_type = at::typeMetaToScalarType(grads.options().dtype()); - auto it = invokers.find(std::tuple(params_type, state_type)); + auto it = invokers.find(std::tuple(params_type, state_type, grad_type)); if (it == invokers.end()) { throw std::runtime_error("Adam optimizer with param type "s + c10::toString(params_type) + - " and state type "s + c10::toString(state_type) + + ", state type "s + c10::toString(state_type) + " and grad type "s + + c10::toString(grad_type) + " is not supported on current hardware"s); } @@ -259,7 +275,8 @@ int ds_adam_step(int optimizer_id, torch::Tensor& params, torch::Tensor& grads, torch::Tensor& exp_avg, - torch::Tensor& exp_avg_sq) + torch::Tensor& exp_avg_sq, + float grad_scale) { auto params_c = params.contiguous(); auto grads_c = grads.contiguous(); @@ -269,13 +286,110 @@ int ds_adam_step(int optimizer_id, std::shared_ptr opt = std::static_pointer_cast(s_optimizers[optimizer_id]); opt->IncrementStep(step, beta1, beta2); - opt->update_state(lr, epsilon, weight_decay, bias_correction); + opt->update_state(lr, epsilon, weight_decay, bias_correction, grad_scale); invoke(opt, params_c, grads_c, exp_avg_c, exp_avg_sq_c, params_c.numel()); return 0; } +template +void Adam_Optimizer::Rescale_Step(float* _params, + ds_grad_precision_t* grads, + float* _exp_avg, + float* _exp_avg_sq, + size_t _param_size, + float old_scale, + float new_scale) +{ + if (_weight_decay > 0 && !_adamw_mode) { + // Adam's L2 decay adds the params of before the step to the gradient, which are gone. + throw std::runtime_error( + "Adam rescale step supports AdamW only, or Adam without weight decay"); + } + size_t rounded_size = 0; +#if defined(__AVX512__) or defined(__AVX256__) or defined(__NEON__) + Rescale_Step_AVX<8>( + &rounded_size, _params, grads, _exp_avg, _exp_avg_sq, _param_size, old_scale, new_scale); +#endif + if (_param_size <= rounded_size) return; + float momentum_shift = (1 - _betta1) * (new_scale - old_scale); + float variance_shift = (1 - _betta2) * (new_scale * new_scale - old_scale * old_scale); + float step_size = -1 * _alpha / _bias_correction1; +#pragma omp parallel for + for (size_t k = rounded_size; k < _param_size; k++) { + float grad = (float)grads[k]; + float momentum = _exp_avg[k]; + float variance = _exp_avg_sq[k]; + float param = _params[k]; + param = momentum / (sqrt(variance) * _bias_correction2 + _eps) * -step_size + param; + momentum = grad * momentum_shift + momentum; + // std::max(x, 0) keeps a NaN x. + variance = std::max(grad * grad * variance_shift + variance, 0.0f); + param = momentum / (sqrt(variance) * _bias_correction2 + _eps) * step_size + param; + _params[k] = param; + _exp_avg[k] = momentum; + _exp_avg_sq[k] = variance; + } +} + +int ds_adam_rescale_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, + float old_scale, + float new_scale) +{ + TORCH_CHECK(params.scalar_type() == torch::kFloat32 && + exp_avg.scalar_type() == torch::kFloat32 && + exp_avg_sq.scalar_type() == torch::kFloat32, + "Adam rescale step supports float32 params and states only"); + TORCH_CHECK(grads.scalar_type() == torch::kFloat32 || grads.scalar_type() == torch::kBFloat16, + "Adam rescale step supports float32 and bfloat16 gradients only"); + TORCH_CHECK(params.is_contiguous() && grads.is_contiguous() && exp_avg.is_contiguous() && + exp_avg_sq.is_contiguous(), + "Adam rescale step updates its tensors in place, so they must be contiguous"); + TORCH_CHECK(grads.numel() == params.numel() && exp_avg.numel() == params.numel() && + exp_avg_sq.numel() == params.numel(), + "Adam rescale step needs as many gradients and states as params"); + + std::shared_ptr opt = + std::static_pointer_cast(s_optimizers[optimizer_id]); + // As the step being rescaled had this step and these betas, beta1^step and beta2^step stay. + opt->IncrementStep(step, beta1, beta2); + opt->update_state(lr, epsilon, weight_decay, bias_correction); + + float* params_ptr = params.data_ptr(); + float* exp_avg_ptr = exp_avg.data_ptr(); + float* exp_avg_sq_ptr = exp_avg_sq.data_ptr(); + if (grads.scalar_type() == torch::kBFloat16) { + opt->Rescale_Step(params_ptr, + (c10::BFloat16*)grads.data_ptr(), + exp_avg_ptr, + exp_avg_sq_ptr, + params.numel(), + old_scale, + new_scale); + } else { + opt->Rescale_Step(params_ptr, + grads.data_ptr(), + exp_avg_ptr, + exp_avg_sq_ptr, + params.numel(), + old_scale, + new_scale); + } + return 0; +} + void adamw_rollback_inplace(float* params, const float* grads, float* momentum, @@ -318,7 +432,10 @@ void adamw_rollback_inplace(float* params, const float grad = grads[i]; momentum[i] = (momentum[i] - one_minus_beta1 * grad) / beta1; - variance[i] = (variance[i] - one_minus_beta2 * grad * grad) / beta2; + // exp_avg_sq is a sum of squares, so it is never negative. The inverse can leave a tiny + // negative rounding residual (in about half of the elements that were 0 before the step), + // and sqrt() of it in the next step makes the parameter NaN. std::max(x, 0) keeps a NaN x. + variance[i] = std::max((variance[i] - one_minus_beta2 * grad * grad) / beta2, 0.0f); } --step_count; @@ -601,9 +718,9 @@ class ZenFlowAdam { torch::Tensor& M = grp.exp_avg[now_state]; torch::Tensor& V = grp.exp_avg_sq[now_state]; - auto it = invokers.find(std::tuple(P.scalar_type(), M.scalar_type())); + auto it = invokers.find(std::tuple(P.scalar_type(), M.scalar_type(), G.scalar_type())); TORCH_CHECK(it != invokers.end(), - "ZenFlowAdam: unsupported param/state dtype combination"); + "ZenFlowAdam: unsupported param/state/grad dtype combination"); auto fn = it->second; char* pp = static_cast(P.data_ptr()); @@ -612,13 +729,14 @@ class ZenFlowAdam { char* vp = static_cast(V.data_ptr()); char* sp = grp.stale.defined() ? static_cast(grp.stale.data_ptr()) : nullptr; const size_t pe = P.element_size(); + const size_t ge = G.element_size(); const size_t se = M.element_size(); const size_t numel = P.numel(); pool_->parallel_for(numel, kZenAdamAlign, [=](size_t b, size_t e) { const size_t len = e - b; // parallel=false: each pinned thread runs its slice serially. - fn(opt, pp + b * pe, gp + b * pe, mp + b * se, vp + b * se, len, false); + fn(opt, pp + b * pe, gp + b * ge, mp + b * se, vp + b * se, len, false); if (sp) std::memcpy(sp + b * pe, pp + b * pe, len * pe); }); } diff --git a/csrc/includes/cpu_adam.h b/csrc/includes/cpu_adam.h index abab3ecda999..c961152a4e22 100644 --- a/csrc/includes/cpu_adam.h +++ b/csrc/includes/cpu_adam.h @@ -16,13 +16,15 @@ #include "simd.h" #include "sve.h" -#define STEP(SPAN) \ - template \ - 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, \ +#define STEP(SPAN) \ + template \ + void Step_##SPAN(ds_params_precision_t* _params, \ + ds_grad_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 { @@ -41,36 +43,46 @@ class Adam_Optimizer { _betta1_t(1.0), _betta2_t(1.0), _step(0), + _grad_scale(1.0f), _adamw_mode(adamw_mode) { } ~Adam_Optimizer() {} #if defined(__AVX512__) or defined(__AVX256__) or defined(__NEON__) - template + template void Step_AVX(size_t* rounded_size, ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t param_size, bool parallel = true); #endif #if defined(__SVE__) && defined(__ARM_FEATURE_SVE) - template + template void Step_SVE(size_t* rounded_size, ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t param_size, bool parallel = true); #endif #if defined(__RVV__) && defined(__riscv_v_intrinsic) && __riscv_v_intrinsic >= 12000 - template + template void Step_RVV(size_t* rounded_size, ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t param_size, @@ -79,6 +91,25 @@ class Adam_Optimizer { STEP(1) STEP(4) STEP(8) +#if defined(__AVX512__) or defined(__AVX256__) or defined(__NEON__) + template + void Rescale_Step_AVX(size_t* rounded_size, + float* _params, + ds_grad_precision_t* grads, + float* _exp_avg, + float* _exp_avg_sq, + size_t param_size, + float old_scale, + float new_scale); +#endif + template + void Rescale_Step(float* _params, + ds_grad_precision_t* grads, + float* _exp_avg, + float* _exp_avg_sq, + size_t _param_size, + float old_scale, + float new_scale); inline void IncrementStep(size_t step, float beta1, float beta2) { if (beta1 != _betta1 || beta2 != _betta2) { @@ -102,11 +133,17 @@ class Adam_Optimizer { } } } - inline void update_state(float lr, float epsilon, float weight_decay, bool bias_correction) + // grad_scale multiplies the gradients as the step reads them, as unscaling and clipping does. + inline void update_state(float lr, + float epsilon, + float weight_decay, + bool bias_correction, + float grad_scale = 1.0f) { _alpha = lr; _eps = epsilon; _weight_decay = weight_decay; + _grad_scale = grad_scale; _bias_correction1 = 1.0f; _bias_correction2 = 1.0f; @@ -130,14 +167,19 @@ class Adam_Optimizer { float _bias_correction1; float _bias_correction2; + float _grad_scale; + bool _adamw_mode; }; #if defined(__AVX512__) or defined(__AVX256__) or defined(__NEON__) -template +template void Adam_Optimizer::Step_AVX(size_t* rounded_size, ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t _param_size, @@ -145,7 +187,8 @@ void Adam_Optimizer::Step_AVX(size_t* rounded_size, { #if !defined(__AVX512__) && !defined(__NEON__) if (std::is_same_v || - std::is_same_v) { + std::is_same_v || + std::is_same_v) { return; } #endif @@ -177,6 +220,8 @@ void Adam_Optimizer::Step_AVX(size_t* rounded_size, AVX_Data weight_decay4; if (_weight_decay > 0) weight_decay4.data = (_adamw_mode ? SIMD_SET(w_decay) : SIMD_SET(_weight_decay)); + AVX_Data grad_scale_4; + grad_scale_4.data = SIMD_SET(_grad_scale); 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; @@ -186,6 +231,7 @@ void Adam_Optimizer::Step_AVX(size_t* rounded_size, for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { AVX_Data grad_4[span]; simd_load(grad_4, grads + i); + if (_grad_scale != 1.0f) { simd_mul(grad_4, grad_4, grad_scale_4); } AVX_Data momentum_4[span]; simd_load(momentum_4, _exp_avg + i); @@ -222,20 +268,110 @@ void Adam_Optimizer::Step_AVX(size_t* rounded_size, } *rounded_size = new_rounded_size; } + +// Turns the last step of fp32 params and states, taken on grads * old_scale, into the step on +// grads * new_scale, in one pass where undoing it and stepping again take two. The moments move by +// the difference of the two gradients, and the params by the difference of the two Adam updates: +// AdamW's weight decay does not depend on the gradient. +template +void Adam_Optimizer::Rescale_Step_AVX(size_t* rounded_size, + float* _params, + ds_grad_precision_t* grads, + float* _exp_avg, + float* _exp_avg_sq, + size_t _param_size, + float old_scale, + float new_scale) +{ +#if !defined(__AVX512__) && !defined(__NEON__) + if (std::is_same_v) { return; } +#endif + AVX_Data momentum_shift_4; + momentum_shift_4.data = SIMD_SET((1 - _betta1) * (new_scale - old_scale)); + AVX_Data variance_shift_4; + variance_shift_4.data = + SIMD_SET((1 - _betta2) * (new_scale * new_scale - old_scale * old_scale)); + + 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); + AVX_Data minus_step_size_4; + minus_step_size_4.data = SIMD_SET(-step_size); + + AVX_Data zero_4; + zero_4.data = SIMD_SET(0.0f); + + size_t 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 + for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { + AVX_Data grad_4[span]; + simd_load(grad_4, grads + i); + + AVX_Data momentum_4[span]; + simd_load(momentum_4, _exp_avg + i); + + AVX_Data variance_4[span]; + simd_load(variance_4, _exp_avg_sq + i); + + AVX_Data param_4[span]; + simd_load(param_4, _params + i); + + // The update of the last step, as Step_AVX computed it from these moments. + AVX_Data update_4[span]; + simd_sqrt(update_4, variance_4); + simd_fma(update_4, update_4, bias2_sqrt, eps_4); + simd_div(update_4, momentum_4, update_4); + simd_fma(param_4, update_4, minus_step_size_4, param_4); + + simd_fma(momentum_4, grad_4, momentum_shift_4, momentum_4); + simd_mul(grad_4, grad_4, grad_4); + simd_fma(variance_4, grad_4, variance_shift_4, variance_4); + // The shift can leave a rounding residual below 0 where the old variance was 0, + // whose sqrt is NaN. + simd_max(variance_4, zero_4, variance_4); + + simd_sqrt(update_4, variance_4); + simd_fma(update_4, update_4, bias2_sqrt, eps_4); + simd_div(update_4, momentum_4, update_4); + simd_fma(param_4, update_4, step_size_4, param_4); + + simd_store(_params + i, param_4); + simd_store(_exp_avg + i, momentum_4); + simd_store(_exp_avg_sq + i, variance_4); + } + } + *rounded_size = new_rounded_size; +} #endif #if defined(__SVE__) && defined(__ARM_FEATURE_SVE) -template +template void Adam_Optimizer::Step_SVE(size_t* rounded_size, ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t _param_size, bool parallel) { if constexpr (std::is_same_v && - std::is_same_v) { + std::is_same_v && + std::is_same_v) { + // The SVE kernel does not scale the gradient; the scalar loop steps a scaled one. + if (_grad_scale != 1.0f) { return; } sve_adam_update(_params, grads, _exp_avg, @@ -256,17 +392,23 @@ void Adam_Optimizer::Step_SVE(size_t* rounded_size, #endif #if defined(__RVV__) && defined(__riscv_v_intrinsic) && __riscv_v_intrinsic >= 12000 -template +template void Adam_Optimizer::Step_RVV(size_t* rounded_size, ds_params_precision_t* _params, - ds_params_precision_t* grads, + ds_grad_precision_t* grads, ds_state_precision_t* _exp_avg, ds_state_precision_t* _exp_avg_sq, size_t _param_size, bool parallel) { if constexpr (std::is_same_v && - std::is_same_v) { + std::is_same_v && + std::is_same_v) { + // The RVV kernel does not scale the gradient; the scalar loop steps a scaled one. + if (_grad_scale != 1.0f) { return; } rvv_adam_update(_params, grads, _exp_avg, @@ -306,7 +448,23 @@ int ds_adam_step(int optimizer_id, torch::Tensor& params, torch::Tensor& grads, torch::Tensor& exp_avg, - torch::Tensor& exp_avg_sq); + torch::Tensor& exp_avg_sq, + float grad_scale); + +int ds_adam_rescale_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, + float old_scale, + float new_scale); int ds_adam_rollback(int optimizer_id, size_t step, diff --git a/csrc/includes/simd.h b/csrc/includes/simd.h index c9bbd63bcb88..b725dff72ac1 100644 --- a/csrc/includes/simd.h +++ b/csrc/includes/simd.h @@ -42,6 +42,8 @@ inline void writeAs(void* dst, const T& val) #define SIMD_FMA(x, y, c) _mm512_fmadd_ps(x, y, c) #define SIMD_SQRT(x) _mm512_sqrt_ps(x) #define SIMD_DIV(x, y) _mm512_div_ps(x, y) +// Returns y when either operand is NaN. +#define SIMD_MAX(x, y) _mm512_max_ps(x, y) #define SIMD_AND(x, y) _mm512_and_ps(x, y) #define SIMD_ANDNOT(x, y) _mm512_andnot_ps(x, y) #define SIMD_OR(x, y) _mm512_or_ps(x, y) @@ -105,6 +107,8 @@ static void store_16_f32_as_bf16_nearest(__m512 v, void* data) #define SIMD_FMA(x, y, c) _mm256_fmadd_ps(x, y, c) #define SIMD_SQRT(x) _mm256_sqrt_ps(x) #define SIMD_DIV(x, y) _mm256_div_ps(x, y) +// Returns y when either operand is NaN. +#define SIMD_MAX(x, y) _mm256_max_ps(x, y) #define SIMD_AND(x, y) _mm256_and_ps(x, y) #define SIMD_ANDNOT(x, y) _mm256_andnot_ps(x, y) #define SIMD_OR(x, y) _mm256_or_ps(x, y) @@ -128,6 +132,8 @@ static void store_16_f32_as_bf16_nearest(__m512 v, void* data) #define SIMD_FMA(x, y, c) vfmaq_f32(c, x, y) #define SIMD_SQRT(x) vsqrtq_f32(x) #define SIMD_DIV(x, y) vdivq_f32(x, y) +// Returns NaN when either operand is NaN. +#define SIMD_MAX(x, y) vmaxq_f32(x, y) #define SIMD_AND(x, y) \ vreinterpretq_f32_u32(vandq_u32(vreinterpretq_u32_f32(x), vreinterpretq_u32_f32(y))) // x86 andnot(x, y) computes ~x & y; NEON bic(a, b) computes a & ~b, so the operands swap. @@ -305,6 +311,13 @@ inline void simd_div(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data* src_a_r) #pragma unroll for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_DIV(src_a_l[i].data, src_a_r[i].data); } } +// max(src_a_l, src_a_r) keeps a NaN of src_a_r on every ISA. +template +inline void simd_max(AVX_Data* dst, AVX_Data src_a_l, AVX_Data* src_a_r) +{ +#pragma unroll + for (size_t i = 0; i < span; ++i) { dst[i].data = SIMD_MAX(src_a_l.data, src_a_r[i].data); } +} template inline void simd_and(AVX_Data* dst, AVX_Data* src_a_l, AVX_Data src_a_r) { diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 284a9c7b5e38..72de35133fbb 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -4999,6 +4999,10 @@ def save_checkpoint(self, save_dir, tag=None, client_state={}, save_latest=True, if save_latest and self.global_rank == 0: with open(os.path.join(save_dir, 'latest'), 'w') as fd: fd.write(tag) + elif self.super_offload(): + # The decoupled writer reads the fp32 weights and Adam state from shared memory, which SuperOffload + # updates in the next backward, before step() would commit this save. + self._commit_decoupled_checkpoint() dist.barrier() diff --git a/deepspeed/runtime/superoffload/superoffload_stage3.py b/deepspeed/runtime/superoffload/superoffload_stage3.py index d6a4270f5aaa..adce271f234c 100644 --- a/deepspeed/runtime/superoffload/superoffload_stage3.py +++ b/deepspeed/runtime/superoffload/superoffload_stage3.py @@ -3,10 +3,14 @@ # DeepSpeed Team -import time -from typing import List +from collections import defaultdict, deque +from concurrent.futures import Future +from typing import List, NamedTuple -from deepspeed.runtime.superoffload.superoffload_utils import SuperOffloadCPUOptimizer, TaskKeys, ResultKeys, EventTypes +import torch + +from deepspeed.ops.adam import DeepSpeedCPUAdam +from deepspeed.runtime.superoffload.superoffload_utils import SuperOffloadCPUAdam from deepspeed.runtime.zero.partition_parameters import Parameter, Tensor from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3 from deepspeed.runtime.utils import combine_grad_norm_groups, is_invalid_grad_norm @@ -15,6 +19,9 @@ from deepspeed.accelerator import get_accelerator OPTIMIZER_STEP_TIMER = 'optimizer_step' +# Larger gradients go to the CPU in pieces of this size (see _copy_grad_to_cpu). On B200 over PCIe Gen5, 16 MiB pieces +# copy at 98% of the rate of whole gradients; smaller ones lose more to the switches between pieces. +_GRAD_COPY_PIECE_BYTES = 16 << 20 def _validate_superoffload_accelerator(): @@ -24,6 +31,29 @@ def _validate_superoffload_accelerator(): f"SuperOffload only supports NVIDIA CUDA GPUs, but found accelerator '{accelerator.device_name()}'.") +def _cpu_adam_settings(optimizer): + """Return (adamw_mode, fp32_optimizer_states) for the CPU Adam that steps the CPU sub-groups of ``optimizer``.""" + if any(group.get("amsgrad", False) for group in optimizer.param_groups): + raise ValueError( + "SuperOffload steps the CPU sub-groups with DeepSpeedCPUAdam, which does not support AMSGrad.") + if isinstance(optimizer, DeepSpeedCPUAdam): + return optimizer.adam_w_mode, optimizer.fp32_optimizer_states + if isinstance(optimizer, torch.optim.AdamW): + return True, True + if isinstance(optimizer, torch.optim.Adam): + return optimizer.defaults.get("decoupled_weight_decay", False), True + raise ValueError(f"SuperOffload steps the CPU sub-groups with DeepSpeedCPUAdam, so it supports only Adam and " + f"AdamW, but the optimizer is {type(optimizer).__name__}.") + + +class _CPUStep(NamedTuple): + """A CPU step submitted during backward, and what changing or undoing it needs.""" + # Its result tells whether the sub-group stepped: one whose gradients were not finite is left to step(). + stepped: Future + hyperparams: dict + grad_scale: float + + class SuperOffloadOptimizer_Stage3(DeepSpeedZeroOptimizer_Stage3): def __init__( @@ -36,59 +66,110 @@ def __init__( **kwargs, ): _validate_superoffload_accelerator() - - self.sub_group_to_param_num = {} - self.sub_group_grad_partition_counts = {} - self._submitted_cpu_sub_groups = set() - self.async_cpuadam_num = 0 - self.max_grad_numel = 0 + adamw_mode, fp32_optimizer_states = _cpu_adam_settings(init_optimizer) + + # Per sub-group, the ds_ids of the params whose gradient for this step is in the fp32 gradient buffer. + self.sub_group_grads_ready = defaultdict(set) + # Sub-groups whose CPU step was submitted during backward, until step() is done or the step is rolled back. + self._submitted_cpu_sub_groups = {} + # The submitted sub-groups whose step has not been collected yet, in the order the CPU Adam thread runs them. + self._uncollected_cpu_steps = deque() + # Host addresses of the CPU sub-groups' fp32 weights and gradients registered with CUDA. + self._registered_host_addresses = [] + # Sub-groups stepped during backward whose new weights are converted into a GPU staging buffer, mapped to the + # event of that conversion. + self._staged_lp_params = {} + self._lp_staging_buffers = {} + # Per staging buffer, the event after which step() has copied it into the params. + self._lp_staging_free_events = {} + self._lp_staging_stream = None + # Stream of the gradient copies to the CPU, so that they overlap the backward compute. + self._grad_copy_stream = None super().__init__(module, init_optimizer, param_names, timers, ds_config, **kwargs) - optimizer_configs = [] - for pg in self.optimizer.param_groups: - optimizer_configs.append({ - "lr": pg["lr"], - "betas": pg["betas"], - "eps": pg["eps"], - "weight_decay": pg["weight_decay"], - "amsgrad": pg["amsgrad"], - }) - cpuadam_cores_perc = kwargs.get("cpuadam_cores_perc", 0.8) - self.superoffload_cpu_optimizer = SuperOffloadCPUOptimizer(optimizer_config=optimizer_configs, - cpuadam_cores_perc=cpuadam_cores_perc, - max_grad_numel=self.max_grad_numel, - pin_memory=self.offload_optimizer_pin_memory) - - def _create_fp16_sub_groups(self, params_group): - - params_group_numel = sum([param.partition_numel() for param in params_group]) - sub_group_size = self.sub_group_size - - if sub_group_size is None or sub_group_size >= params_group_numel: - global_idx = len(self.sub_group_to_param_num) - self.sub_group_to_param_num[global_idx] = len(params_group) - self.max_grad_numel = max(self.max_grad_numel, params_group_numel) - return [params_group] - - sub_groups = [] - sub_group = [] - local_sub_group_size = 0 - - for param in params_group: - sub_group.append(param) - local_sub_group_size += param.partition_numel() - - if local_sub_group_size >= sub_group_size or id(param) == id(params_group[-1]): - self.max_grad_numel = max(self.max_grad_numel, local_sub_group_size) - sub_groups.append(sub_group) - global_idx = len(self.sub_group_to_param_num) - self.sub_group_to_param_num[global_idx] = len(sub_group) - - sub_group = [] - local_sub_group_size = 0 - - return sub_groups + self.sub_group_to_param_num = [len(group) for group in self.fp16_groups] + # step() changes a step taken during backward with adam_rescale_step when the gradients clip or the loss scale + # changed, and undoes it with adam_rollback when it skips the step. Both need fp32 AdamW, and adam_rollback + # bias correction. With other settings, step() runs the CPU sub-groups after the overflow and clipping checks. + bias_correction = all(group.get("bias_correction", True) for group in self.optimizer.param_groups) + self.speculative_cpu_step = (adamw_mode and fp32_optimizer_states + and self.master_weights_and_grads_dtype == torch.float32 and bias_correction) + self.cpu_adam = SuperOffloadCPUAdam(self.optimizer.param_groups[0], adamw_mode, fp32_optimizer_states, + kwargs.get("cpuadam_cores_perc", 0.8)) + # With bf16 gradients accumulated into fp32 master weights in one micro-step, the gradients of the CPU + # sub-groups go to the CPU in bf16, half the bytes of fp32, and the worker promotes them in its kernel. + self._half_cpu_grads = {} + if (self.dtype == torch.bfloat16 and self.gradient_accumulation_dtype == torch.bfloat16 + and self.master_weights_and_grads_dtype == torch.float32 and fp32_optimizer_states + and self.gradient_accumulation_steps == 1): + for sub_group_id, device in self.subgroup_to_device.items(): + if device == 'cpu' and not self._swappable_optimizer_subgroup(sub_group_id): + self._half_cpu_grads[sub_group_id] = torch.zeros( + self.fp32_partitioned_groups_flat[sub_group_id].numel(), dtype=torch.bfloat16) + self._pin_cpu_buffers() + self._grad_copy_stream = get_accelerator().Stream() + # The pieces of a large gradient take turns on the copy stream and a second one, each with an event + # (_copy_grad_to_cpu). + self._grad_copy_streams = (self._grad_copy_stream, get_accelerator().Stream()) + self._grad_copy_events = (get_accelerator().Event(), get_accelerator().Event()) + if self.speculative_cpu_step and self._cpu_sub_group_weights_pinned(): + # Default priority is the lowest, so the conversions never delay the backward kernels. + self._lp_staging_stream = get_accelerator().Stream() + + def _pin_cpu_buffers(self): + """Keep the CPU sub-groups' fp32 weights and gradients (and bf16 gradients) page-locked. + + offload_optimizer.pin_memory pins the gradients but not the fp32 weights. The gradient copies of backward and + the weight copies of step() would then run from pageable memory, at a fraction of the bandwidth and blocking + the thread that issues them. Registering their memory with the device runtime keeps the copies asynchronous. + """ + if not self.offload_optimizer_pin_memory: + return + accelerator = get_accelerator() + for sub_group_id, device in self.subgroup_to_device.items(): + # Swapped sub-groups share one gradient buffer and do not keep their weights in memory. + if device != 'cpu' or self._swappable_optimizer_subgroup(sub_group_id): + continue + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + for tensor in (fp32_param.data, fp32_param.grad, self._half_cpu_grads.get(sub_group_id)): + if tensor is None or tensor.is_pinned(): + continue + storage = tensor.untyped_storage() + if storage.data_ptr() in self._registered_host_addresses: + continue + try: + registered = accelerator.register_host_memory(storage.data_ptr(), storage.nbytes()) + except Exception as e: + registered = False + logger.warning(f"SuperOffload could not page-lock its CPU buffers: {e}") + if not registered: + return + self._registered_host_addresses.append(storage.data_ptr()) + + def _cpu_sub_group_weights_pinned(self): + """Whether the fp32 weights of the CPU sub-groups are page-locked, which the GPU staging of step() needs.""" + weights = [ + self.fp32_partitioned_groups_flat[sub_group_id].data + for sub_group_id, device in self.subgroup_to_device.items() + if device == 'cpu' and not self._swappable_optimizer_subgroup(sub_group_id) + ] + return bool(weights) and all(weight.is_pinned() for weight in weights) + + def _worker_grad(self, sub_group_id): + """The gradient of a sub-group that the CPU worker steps on: the bf16 one if it has one, else the fp32 one.""" + half_grad = self._half_cpu_grads.get(sub_group_id) + return half_grad if half_grad is not None else self.fp32_partitioned_groups_flat[sub_group_id].grad.data + + def destroy(self): + self.cpu_adam.close() + super().destroy() + + def _unpin_offload_buffers(self): + for address in self._registered_host_addresses: + get_accelerator().unregister_host_memory(address) + self._registered_host_addresses = [] + super()._unpin_offload_buffers() def _optimizer_step(self, sub_group_id): param_group_id = self.sub_group_to_group_id[sub_group_id] @@ -103,20 +184,21 @@ def step_with_gradscaler(optimizer): cur_device = self.subgroup_to_device[sub_group_id] if cur_device != 'cpu': + self._sync_backup_optimizer_group(param_group_id) self.backup_optimizer.param_groups[param_group_id]['params'] = [fp32_param] step_with_gradscaler(self.backup_optimizer) self.backup_optimizer.param_groups[param_group_id]['params'] = [] - @instrument_w_nvtx - def independent_gradient_partition_epilogue(self): - super().independent_gradient_partition_epilogue() - self.sub_group_grad_partition_counts.clear() - @instrument_w_nvtx def _reassign_or_swap_out_partitioned_parameters(self, sub_group_id): if self.subgroup_to_device[sub_group_id] == 'cpu': - self.fp16_partitioned_groups_flat[sub_group_id].data.copy_( - self.fp32_partitioned_groups_flat[sub_group_id].data) + lp_flat = self.fp16_partitioned_groups_flat[sub_group_id] + fp32_flat = self.fp32_partitioned_groups_flat[sub_group_id].data + if fp32_flat.is_pinned() and get_accelerator().on_accelerator(lp_flat): + # copy_ from a CPU tensor of another dtype converts on the CPU into a pageable temporary and copies + # that. Copying the pinned fp32 weights and converting on the GPU keeps the CPU free and the copy fast. + fp32_flat = fp32_flat.to(lp_flat.device, non_blocking=True) + lp_flat.data.copy_(fp32_flat) self._unflatten_partitioned_parameters(sub_group_id) return @@ -128,9 +210,79 @@ def _reassign_or_swap_out_partitioned_parameters(self, sub_group_id): self._partitioned_params_swap_out(sub_group_id) @instrument_w_nvtx - def _reassign_or_swap_out_partitioned_parameters_async(self, sub_group_id, updated_param): - """Asynchronously update partitioned parameters with optimized values.""" - self.fp32_partitioned_groups_flat[sub_group_id].data.copy_(updated_param, non_blocking=True) + def _stage_lp_params(self, sub_group_id): + """Convert the weights of a CPU step taken during backward to the low-precision dtype on the GPU. + + Backward may still gather the current params, and step() changes them only after the overflow and clipping + checks, so the conversion goes to a staging buffer that step() copies in. + """ + stream = self._lp_staging_stream + lp_flat = self.fp16_partitioned_groups_flat[sub_group_id] + fp32_flat = self.fp32_partitioned_groups_flat[sub_group_id].data + if stream is None or lp_flat is None or not get_accelerator().on_accelerator(lp_flat): + return + if not fp32_flat.is_pinned(): + return + staging = self._lp_staging_buffers.get(sub_group_id) + if staging is None: + # Allocate it on the stream that writes it. The caching allocator hands the current stream blocks freed + # while kernels queued on that stream may still use them, which another stream does not wait for. + with get_accelerator().stream(stream): + staging = torch.empty_like(lp_flat) + self._lp_staging_buffers[sub_group_id] = staging + free_event = self._lp_staging_free_events.pop(sub_group_id, None) + if free_event is not None: + # The previous step() may still be copying this buffer into the params. + stream.wait_event(free_event) + with get_accelerator().stream(stream): + staging.copy_(fp32_flat.to(staging.device, non_blocking=True)) + event = get_accelerator().Event() + event.record(stream) + self._staged_lp_params[sub_group_id] = event + + def _drop_staged_lp_params(self, sub_group_id): + """Discard a staged conversion before the worker changes the sub-group's weights, once it has read them.""" + event = self._staged_lp_params.pop(sub_group_id, None) + if event is not None: + event.synchronize() + + @instrument_w_nvtx + def _apply_staged_lp_params(self, sub_group_id): + event = self._staged_lp_params.pop(sub_group_id) + stream = get_accelerator().current_stream() + stream.wait_event(event) + self.fp16_partitioned_groups_flat[sub_group_id].data.copy_(self._lp_staging_buffers[sub_group_id]) + free_event = get_accelerator().Event() + free_event.record(stream) + self._lp_staging_free_events[sub_group_id] = free_event + self._unflatten_partitioned_parameters(sub_group_id) + + def _collect_cpu_steps(self, wait=False): + """Stage the new weights of the CPU steps submitted during backward that have finished. With wait, wait for + all of them, and raise what a step raised. + + The CPU Adam thread runs the steps in the order they were submitted, so they finish in that order. + """ + while self._uncollected_cpu_steps: + sub_group_id = self._uncollected_cpu_steps[0] + stepped = self._submitted_cpu_sub_groups[sub_group_id].stepped + if not wait and not stepped.done(): + return + self._uncollected_cpu_steps.popleft() + if stepped.result(): + self._stage_lp_params(sub_group_id) + else: + # Its gradients were not finite, so step() runs it after the overflow check. + del self._submitted_cpu_sub_groups[sub_group_id] + + def _copies_grad_directly(self, sub_group_id): + """Whether a gradient of the sub-group goes straight into the master-dtype gradient the worker reads. + + That needs a CPU sub-group and the only micro-step of the step, so that no earlier micro-step left a gradient + in the reduce buffer. fp16 keeps the reduce buffer, which its overflow check scans. + """ + return (self.subgroup_to_device[sub_group_id] == 'cpu' and self.dtype != torch.float16 + and self.is_gradient_accumulation_boundary() and self.micro_step_id == 0) @instrument_w_nvtx def partition_grads(self, params_to_release: List[Parameter], grad_partitions: List[Tensor]) -> None: @@ -142,7 +294,29 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L # Accumulate gradient into the grad_buffer, mirroring base class logic grad_buffer = self._DeepSpeedZeroOptimizer_Stage3__param_id_to_grad_partition[param.ds_id].narrow( 0, 0, grad_partition.numel()) - if self.micro_step_id == 0: + # The buffer holds an older step's gradient until the param's first gradient of this step. Testing the + # param rather than micro_step_id also keeps both gradients of a param that two reentrant checkpoint + # segments use. + first_grad = param.ds_id not in self._offload_boundary_param_ids + self._offload_boundary_param_ids.add(param.ds_id) + if self._copies_grad_directly(i): + if first_grad: + # The gradient goes to the CPU once, into the gradient the worker reads, instead of into the + # reduce buffer, back to the GPU for the norm, and out again. It is rounded to the accumulation + # dtype as the reduce buffer would round it, and goes in that dtype if the worker reads bf16. + grad_buffer = grad_partition.to(self.gradient_accumulation_dtype) + if i not in self._half_cpu_grads: + grad_buffer = grad_buffer.to(self.master_weights_and_grads_dtype) + else: + # A second gradient in this step: the first one is only in the gradient the worker reads, which + # the copy stream may still be writing. The sum is taken in the accumulation dtype, as in the + # reduce buffer; the first gradient was rounded to it, so converting it back is exact. + cpu_grad_tensor = self._worker_grad(i).narrow(0, dest_offset, grad_partition.numel()) + get_accelerator().current_stream().wait_stream(self._grad_copy_stream) + grad_buffer = cpu_grad_tensor.to(grad_partition.device, non_blocking=True) + grad_buffer = grad_buffer.to(self.gradient_accumulation_dtype) + grad_buffer.add_(grad_partition.to(grad_buffer.dtype).view(grad_buffer.shape)) + elif first_grad: grad_buffer.copy_(grad_partition, non_blocking=True) grad_buffer = grad_buffer.to(grad_partition.device, non_blocking=True) elif get_accelerator().on_accelerator(grad_buffer): @@ -154,37 +328,58 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L grad_buffer = cuda_grad_buffer if self.is_gradient_accumulation_boundary(): + if i in self._submitted_cpu_sub_groups: + # Another gradient for a sub-group that has already stepped: undo that step while the fp32 + # buffer still holds the gradient it used. One whose gradients were not finite did not step. + self._collect_cpu_steps(wait=True) + if i in self._submitted_cpu_sub_groups: + self._rollback_cpu_sub_group(i) self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer) - fp32_grad_tensor = self.fp32_partitioned_groups_flat[i].grad.narrow( - 0, dest_offset, grad_buffer.numel()) - fp32_grad_tensor.copy_(grad_buffer.to(dtype=self.master_weights_and_grads_dtype), non_blocking=True) + cpu_grad_tensor = self._worker_grad(i).narrow(0, dest_offset, grad_buffer.numel()) + grad_to_copy = grad_buffer.to(dtype=cpu_grad_tensor.dtype) + self._copy_grad_to_cpu(cpu_grad_tensor, grad_to_copy) - self.sub_group_grad_partition_counts[i] = self.sub_group_grad_partition_counts.get(i, 0) + 1 - if self.sub_group_grad_partition_counts[i] == self.sub_group_to_param_num[i]: - completed_sub_groups.append(i) + self.sub_group_grads_ready[i].add(param.ds_id) + if len(self.sub_group_grads_ready[i]) == self.sub_group_to_param_num[i]: + completed_sub_groups.append(i) if self.is_gradient_accumulation_boundary() and completed_sub_groups: - get_accelerator().current_stream().synchronize() - for i in completed_sub_groups: - if self.subgroup_to_device[i] == 'cpu' and not self.clip_grad: - param_group_id = self.sub_group_to_group_id[i] - fp32_param = self.fp32_partitioned_groups_flat[i] - current_lr = self.optimizer.param_groups[param_group_id]['lr'] - - self.superoffload_cpu_optimizer.async_step(param_group_id, - i, - fp32_param.data, - fp32_param.grad.data, - lr=current_lr) - self._submitted_cpu_sub_groups.add(i) - self.async_cpuadam_num += 1 - - result = self.superoffload_cpu_optimizer.get_result() - if result is not None: - self._reassign_or_swap_out_partitioned_parameters_async(result[TaskKeys.SUB_GROUP_ID], - result[ResultKeys.UPDATED_PARAM]) - self.async_cpuadam_num -= 1 + steps = [] + finite_on_device = [] + # A param that gets a second gradient adds its sub-group again. + for i in dict.fromkeys(completed_sub_groups): + param_group_id = self.sub_group_to_group_id[i] + if (self.subgroup_to_device[i] == 'cpu' and self.speculative_cpu_step + and self._betas_nonzero(param_group_id)): + hyperparams = self._hyperparams(param_group_id) + # The gradients still carry the loss scale, which step() would remove. + grad_scale = 1.0 / self.loss_scale + + self._drop_staged_lp_params(i) + steps.append((i, hyperparams, grad_scale)) + finite_on_device.append(self._grads_finite(i)) + if steps: + # The CPU Adam thread reads the gradients that the copy stream writes, and waits for them itself, so + # that this thread goes on queueing the backward kernels. + copy_stream = self._grad_copy_stream + copy_stream.wait_stream(get_accelerator().current_stream()) + with get_accelerator().stream(copy_stream): + # Behind the gradient copies: a copy to the CPU on the compute stream would wait for them in the + # copy engine, and hold up the kernels queued after it. + finite = [flag.to("cpu", non_blocking=True) for flag in finite_on_device] + copies_done = get_accelerator().Event() + copies_done.record(copy_stream) + for flag in finite_on_device: + flag.record_stream(copy_stream) + for (i, hyperparams, grad_scale), grads_finite in zip(steps, finite): + stepped = self.cpu_adam.submit(self._step_when_copied, i, copies_done, grads_finite, hyperparams, + grad_scale) + self._submitted_cpu_sub_groups[i] = _CPUStep(stepped, hyperparams, grad_scale) + self._uncollected_cpu_steps.append(i) + + if self.is_gradient_accumulation_boundary(): + self._collect_cpu_steps() for param in params_to_release: if not get_accelerator().is_synchronized_device(): @@ -192,20 +387,55 @@ def partition_grads(self, params_to_release: List[Parameter], grad_partitions: L param.grad.record_stream(get_accelerator().current_stream()) param.grad = None + def _copy_grad_to_cpu(self, cpu_grad_tensor, grad_to_copy): + """Queue the copy of a gradient to the CPU on the gradient copy stream, behind the kernels that computed it. + + The copy engine keeps serving a stream while it has a copy ready, so a copy to the CPU on another stream, such + as the one .item() makes in backward, waits behind all the gradient copies queued before it: up to 28 ms with + OPT-30B. A large gradient goes in pieces that take turns on two streams, each waiting for the one before, so + that neither stream has a copy ready while a piece runs and the copies of other streams go in between. + """ + copy_stream = self._grad_copy_stream + copy_stream.wait_stream(get_accelerator().current_stream()) + piece = _GRAD_COPY_PIECE_BYTES // grad_to_copy.element_size() + if grad_to_copy.numel() <= piece: + with get_accelerator().stream(copy_stream): + cpu_grad_tensor.copy_(grad_to_copy, non_blocking=True) + # The copy reads the gradient later, so keep the allocator from handing its memory out before then. + grad_to_copy.record_stream(copy_stream) + return + streams = self._grad_copy_streams + events = self._grad_copy_events + previous = None + for turn, start in enumerate(range(0, grad_to_copy.numel(), piece)): + stream = streams[turn % 2] + if previous is not None: + stream.wait_event(previous) + with get_accelerator().stream(stream): + cpu_grad_tensor[start:start + piece].copy_(grad_to_copy[start:start + piece], non_blocking=True) + previous = events[turn % 2] + previous.record(stream) + # What waits for the copy stream, as step() and the CPU steps do, then waits for the last piece too. + copy_stream.wait_event(previous) + for stream in streams: + grad_to_copy.record_stream(stream) + @instrument_w_nvtx def step(self, closure=None): """ Not supporting closure. """ - self._wait_for_async_operations() + # Sub-groups that did not complete in backward may still have gradient copies in flight. + self._grad_copy_stream.synchronize() + self._collect_cpu_steps(wait=True) + self._complete_grads() self._pre_step() self._partition_all_parameters() self._overflow_check_and_loss_scale_update(update_scale=False) if self.overflow: - if not self.clip_grad: - self._handle_overflow_rollback() + self._rollback_submitted_cpu_sub_groups() self._loss_scale_update_and_overflow_cleanup() return @@ -213,8 +443,7 @@ def step(self, closure=None): scaled_global_grad_norm = combine_grad_norm_groups(norm_groups) if is_invalid_grad_norm(scaled_global_grad_norm): self.overflow = True - if not self.clip_grad: - self._handle_overflow_rollback() + self._rollback_submitted_cpu_sub_groups() self._loss_scale_update_and_overflow_cleanup() return self._loss_scale_update_and_overflow_cleanup() @@ -224,166 +453,157 @@ def step(self, closure=None): timer_names.add(OPTIMIZER_STEP_TIMER) self.timers(OPTIMIZER_STEP_TIMER).start() - if self.clip_grad: - self._step_with_clipping(scaled_global_grad_norm, timer_names) - else: - self._step_without_clipping(scaled_global_grad_norm, timer_names) + # CPU Adam removes the loss scale and clips as it reads the gradients, by the factor unscale_and_clip_grads + # multiplies them by. + grad_scale = float(self._unscale_and_clip_factor(scaled_global_grad_norm)) + for sub_group_id, group in enumerate(self.fp16_groups): + self._prepare_sub_group(sub_group_id, timer_names) + if self.subgroup_to_device[sub_group_id] != 'cpu': + self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) + self._optimizer_step(sub_group_id) + elif sub_group_id in self._submitted_cpu_sub_groups: + self._rescale_cpu_step(sub_group_id, grad_scale) + else: + param_group_id = self.sub_group_to_group_id[sub_group_id] + self._cpu_step_and_wait(sub_group_id, self._worker_grad(sub_group_id), + self._hyperparams(param_group_id), grad_scale) + + if sub_group_id in self._staged_lp_params: + self._apply_staged_lp_params(sub_group_id) + else: + self._reassign_or_swap_out_partitioned_parameters(sub_group_id) + self._release_sub_group(sub_group_id, timer_names) self.timers(OPTIMIZER_STEP_TIMER).stop() self._post_step(timer_names) def _post_step(self, timer_names): super()._post_step(timer_names) + for sub_group_id in list(self._staged_lp_params): + self._drop_staged_lp_params(sub_group_id) self._submitted_cpu_sub_groups.clear() - def _step_without_clipping(self, scaled_global_grad_norm, timer_names): - """Fast path: async CPU steps already completed during backward.""" - for sub_group_id, group in enumerate(self.fp16_groups): - self._prepare_sub_group(sub_group_id, timer_names) - self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) - self._optimizer_step(sub_group_id) - self._reassign_or_swap_out_partitioned_parameters(sub_group_id) - self._release_sub_group(sub_group_id, timer_names) - - def _step_with_clipping(self, scaled_global_grad_norm, timer_names): - """Clipping path: no async steps were done during backward, - so we unscale+clip first, then step all sub-groups.""" - for sub_group_id, group in enumerate(self.fp16_groups): - self._prepare_sub_group(sub_group_id, timer_names) - self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) + def _complete_grads(self): + """Fill the gradient the worker reads for the params that got no gradient in the boundary micro-step. - if self.subgroup_to_device[sub_group_id] == 'cpu': - param_group_id = self.sub_group_to_group_id[sub_group_id] - fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] - current_lr = self.optimizer.param_groups[param_group_id]['lr'] - self._sync_cpu_optimizer_step(param_group_id, - sub_group_id, - fp32_param.data, - fp32_param.grad.data, - lr=current_lr) - else: - self._optimizer_step(sub_group_id) - - self._reassign_or_swap_out_partitioned_parameters(sub_group_id) - self._release_sub_group(sub_group_id, timer_names) - - def _wait_for_async_operations(self, timeout_seconds=60): - """Wait for all pending asynchronous CPU optimizer operations to complete with timeout error. - - Args: - timeout_seconds (int): Maximum time to wait before throwing an error. Default is 60 seconds. + A param used only in earlier micro-steps has its gradient in the reduce buffer. A param not used in this step + gets a zero gradient, as set_none_gradients_to_zero gives it in ZeRO-1/2; zeroing its reduce buffer also keeps + an older inf out of the fp16 overflow check. Their sub-groups did not step during backward, so step() runs + them. """ - if self.async_cpuadam_num > 0: - logger.info(f"[INFO] {self.async_cpuadam_num} asynchronous CPU optimizer operations pending...") - if self.async_cpuadam_num == 0: - return - - start_time = time.time() - initial_pending_ops = self.async_cpuadam_num - - while self.async_cpuadam_num > 0: - result = self.superoffload_cpu_optimizer.get_result() - if result is None: - current_time = time.time() - elapsed_time = current_time - start_time - - # Throw error if we've been waiting longer than the timeout - if elapsed_time >= timeout_seconds: - raise RuntimeError( - f"SuperOffload CPU optimizer timeout after {elapsed_time:.1f} seconds. " - f"Still waiting for {self.async_cpuadam_num}/{initial_pending_ops} async operations to complete. " - f"This indicates a deadlock or critical performance issue in the CPU optimizer.") - - time.sleep(0.001) # 1ms sleep + filled = False + for i, group in enumerate(self.fp16_groups): + ready = self.sub_group_grads_ready[i] + if len(ready) == self.sub_group_to_param_num[i]: continue - - self._reassign_or_swap_out_partitioned_parameters_async(result[TaskKeys.SUB_GROUP_ID], - result[ResultKeys.UPDATED_PARAM]) - self.async_cpuadam_num -= 1 - - def _wait_for_single_async_result(self, event_type: str, timeout_seconds=60): - """Wait for a single asynchronous CPU-Adam optimizer operation with timeout. - - Args: - event_type (str): Type of operation expected ('adam_step' or 'rollback'). - timeout_seconds (int): Maximum time to wait before throwing an error. Default is 60 seconds. + for param in group: + if param.ds_id in ready: + continue + _, dest_offset, num_elements = self.grad_position[self.get_param_id(param)] + grad_buffer = self._DeepSpeedZeroOptimizer_Stage3__param_id_to_grad_partition[param.ds_id].narrow( + 0, 0, num_elements) + if param.ds_id not in self._offload_boundary_param_ids: + grad_buffer.zero_() + if not get_accelerator().on_accelerator(grad_buffer): + grad_buffer = grad_buffer.to(get_accelerator().current_device_name(), non_blocking=True) + self.norm_for_param_grads[self.get_param_id(param)] = self._constant_buffered_norm2(grad_buffer) + cpu_grad_tensor = self._worker_grad(i).narrow(0, dest_offset, num_elements) + cpu_grad_tensor.copy_(grad_buffer.to(dtype=cpu_grad_tensor.dtype), non_blocking=True) + filled = True + if filled: + get_accelerator().current_stream().synchronize() + self.sub_group_grads_ready.clear() + + def _grads_finite(self, sub_group_id): + """Whether the sub-group's gradients are finite, as a device tensor, so that this thread does not wait for the + norms. A step on inf or nan cannot be rolled back.""" + norms = [self.norm_for_param_grads[self.get_param_id(param)] for param in self.fp16_groups[sub_group_id]] + return torch.isfinite(torch.stack(norms)).all() + + def _betas_nonzero(self, param_group_id): + """Whether a step of the param group can be rolled back: the rollback divides the moments by beta1 and beta2. + Checked at each step, as momentum schedules change the betas.""" + return all(beta != 0 for beta in self.optimizer.param_groups[param_group_id]["betas"]) + + def _hyperparams(self, param_group_id): + """The hyperparameters of a param group, which LR, momentum or weight decay schedules change.""" + group = self.optimizer.param_groups[param_group_id] + hyperparams = {key: group[key] for key in ("lr", "betas", "eps", "weight_decay")} + hyperparams["bias_correction"] = group.get("bias_correction", True) + return hyperparams + + def _step_when_copied(self, sub_group_id, copies_done, grads_finite, hyperparams, grad_scale): + """In the CPU Adam thread: step a sub-group once its gradients are on the CPU, and tell whether it did. + + A step on inf or nan cannot be rolled back, so step() runs those sub-groups. """ - start_time = time.time() - - while True: - result = self.superoffload_cpu_optimizer.get_result(expected_event_type=event_type) - if result is not None: - self._reassign_or_swap_out_partitioned_parameters_async(result[TaskKeys.SUB_GROUP_ID], - result[ResultKeys.UPDATED_PARAM]) - break - - current_time = time.time() - elapsed_time = current_time - start_time - - # Throw error if we've been waiting longer than the timeout - if elapsed_time >= timeout_seconds: - raise RuntimeError(f"SuperOffload CPU optimizer timeout after {elapsed_time:.1f} seconds. " - f"This indicates a deadlock or critical performance issue in the CPU optimizer.") - - time.sleep(0.001) # 1ms sleep - - def _sync_cpu_optimizer_step(self, - param_group_id: int, - sub_group_id: int, - fp32_param_data, - fp32_grad_data, - rollback: bool = False, - lr: float = None, - timeout_seconds: int = 60): - event_type = EventTypes.ROLLBACK if rollback else EventTypes.ADAM_STEP - self.superoffload_cpu_optimizer.async_step(param_group_id, - sub_group_id, - fp32_param_data, - fp32_grad_data, - rollback=rollback, - lr=lr) - # Wait for completion - self._wait_for_single_async_result(event_type, timeout_seconds) - - def _handle_overflow_rollback(self): - """Handle gradient overflow by rolling back CPU optimizer states.""" - for sub_group_id in self._submitted_cpu_sub_groups: - param_group_id = self.sub_group_to_group_id[sub_group_id] - fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] - self._sync_cpu_optimizer_step(param_group_id, - sub_group_id, - fp32_param.data, - fp32_param.grad.data, - rollback=True) - self._submitted_cpu_sub_groups.clear() + copies_done.synchronize() + if not grads_finite.item(): + return False + self._cpu_adam_step(sub_group_id, self._worker_grad(sub_group_id), hyperparams, grad_scale) + return True + + def _cpu_adam_step(self, sub_group_id, grad, hyperparams, grad_scale=1.0, rollback=False): + """In the CPU Adam thread: step a CPU sub-group, or undo its last step.""" + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + update = self.cpu_adam.rollback if rollback else self.cpu_adam.step + # The Adam state stays where ZeRO-3 CPU offload keeps it, so that checkpoints save and load it as they are. + update(self.optimizer.state[fp32_param], fp32_param.data, grad, hyperparams, grad_scale) - def _handle_gradient_clipping(self, scaled_global_grad_norm): - """Handle gradient clipping with CPU optimizer rollback and re-optimization.""" - for sub_group_id, _ in enumerate(self.fp16_groups): - if self.subgroup_to_device[sub_group_id] == 'cpu': - param_group_id = self.sub_group_to_group_id[sub_group_id] - fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + def _cpu_step_and_wait(self, sub_group_id, grad, hyperparams, grad_scale=1.0, rollback=False): + self._drop_staged_lp_params(sub_group_id) + self.cpu_adam.submit(self._cpu_adam_step, sub_group_id, grad, hyperparams, grad_scale, rollback).result() - # Rollback CPU optimizer states - self._sync_cpu_optimizer_step(param_group_id, - sub_group_id, - fp32_param.data, - fp32_param.grad.data, - rollback=True) + def _rescale_cpu_step(self, sub_group_id, grad_scale): + """Turn the CPU step a sub-group took during backward into the step on its gradients at grad_scale. - # Clip gradients and re-optimize - self.unscale_and_clip_grads(sub_group_id, scaled_global_grad_norm) - - current_lr = self.optimizer.param_groups[param_group_id]['lr'] - self._sync_cpu_optimizer_step(param_group_id, - sub_group_id, - fp32_param.data, - fp32_param.grad.data, - rollback=False, - lr=current_lr) + That step read the gradients unclipped and with the loss scale of the backward, which + _loss_scale_update_and_overflow_cleanup may have raised since. One pass over the weights and moments changes + the step, where undoing it and stepping again would take two, and the undo is slower than a step. + """ + step = self._submitted_cpu_sub_groups[sub_group_id] + if step.grad_scale == grad_scale: + return + hyperparams = self._hyperparams(self.sub_group_to_group_id[sub_group_id]) + if hyperparams != step.hyperparams: + # Changed since backward, so the step taken with the old ones is undone, and taken again. + self._rollback_cpu_sub_group(sub_group_id) + self._cpu_step_and_wait(sub_group_id, self._worker_grad(sub_group_id), hyperparams, grad_scale) + return + self._drop_staged_lp_params(sub_group_id) + self.cpu_adam.submit(self._cpu_adam_rescale, sub_group_id, hyperparams, step.grad_scale, grad_scale).result() - @instrument_w_nvtx - def check_clip_grads(self, total_norm): - """Check if gradients need to be clipped based on the global norm.""" - unscaled_norm = total_norm / self.loss_scale - return self.clip_grad and unscaled_norm > self.clip_grad + def _cpu_adam_rescale(self, sub_group_id, hyperparams, old_scale, new_scale): + """In the CPU Adam thread: turn the last step of a CPU sub-group into the step on its gradient at new_scale.""" + fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] + self.cpu_adam.rescale_step(self.optimizer.state[fp32_param], fp32_param.data, self._worker_grad(sub_group_id), + hyperparams, old_scale, new_scale) + + def _rollback_submitted_cpu_sub_groups(self): + """Undo the CPU steps taken during backward, for a skipped step.""" + for sub_group_id in list(self._submitted_cpu_sub_groups): + self._rollback_cpu_sub_group(sub_group_id) + + def _rollback_cpu_sub_group(self, sub_group_id): + step = self._submitted_cpu_sub_groups.pop(sub_group_id) + # The rollback reads the gradient that the step read. + self._cpu_step_and_wait(sub_group_id, + self._worker_grad(sub_group_id), + step.hyperparams, + step.grad_scale, + rollback=True) + + def load_state_dict(self, + state_dict_list, + load_optimizer_states=True, + load_from_fp32_weights=False, + checkpoint_folder=None, + load_serial=None, + param_shapes=None): + get_accelerator().current_stream().synchronize() + super().load_state_dict(state_dict_list, + load_optimizer_states=load_optimizer_states, + load_from_fp32_weights=load_from_fp32_weights, + checkpoint_folder=checkpoint_folder, + load_serial=load_serial, + param_shapes=param_shapes) diff --git a/deepspeed/runtime/superoffload/superoffload_utils.py b/deepspeed/runtime/superoffload/superoffload_utils.py index 266f8e922a2c..143363654860 100644 --- a/deepspeed/runtime/superoffload/superoffload_utils.py +++ b/deepspeed/runtime/superoffload/superoffload_utils.py @@ -3,313 +3,125 @@ # DeepSpeed Team """ -SuperOffload utilities for 1) running CPU optimizers in separate processes. - +SuperOffload's CPU Adam, which steps the CPU sub-groups in a thread of the training process. """ -from typing import Dict, Optional, Any +import os +from concurrent.futures import Future, ThreadPoolExecutor +from typing import Any, Callable, Dict, List + import torch -import torch.multiprocessing as mp -import psutil from deepspeed.ops.adam import DeepSpeedCPUAdam -from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import CPUAdamBuilder from deepspeed.utils import logger +from deepspeed.utils.logging import should_log_le -class TaskKeys: - PARAM_DATA = "param_data" - PARAM_GRAD = "param_grad" - PARAM_GROUP_ID = "param_group_id" - SUB_GROUP_ID = "sub_group_id" - ROLLBACK = "rollback" - LR = "lr" - - -class ResultKeys: - UPDATED_PARAM = "updated_param" - EVENT_TYPE = "event_type" - - -class EventTypes: - ADAM_STEP = "adam_step" - ROLLBACK = "rollback" - - -def _allocate_worker_grad_buffer(max_grad_numel: int, pin_memory: bool) -> torch.Tensor: - # Scratch destination for the GPU->CPU grad copy. Pinning is optional: it - # speeds DMA but is not required for correctness, so honor - # offload_optimizer.pin_memory rather than always page-locking. - if pin_memory: - return get_accelerator().pin_empty(max_grad_numel, dtype=torch.float32, device='cpu') - return torch.empty(max_grad_numel, dtype=torch.float32, device='cpu') - - -def superoffload_optimizer_worker(param_queue: mp.SimpleQueue, - result_queue: mp.SimpleQueue, - optimizer_config: Dict[str, Any], - max_grad_numel: int, - pin_memory: bool = True) -> None: - """ - This function runs in a separate process and continuously processes optimization - tasks from the parameter queue. It creates a DeepSpeedCPUAdam optimizer and - applies optimization steps to parameters received from the main process. - - Args: - param_queue: Queue for receiving optimization tasks - result_queue: Queue for sending back optimization results - optimizer_config: Configuration dictionary for the optimizer containing - lr, betas, eps, weight_decay, and amsgrad parameters - max_grad_numel: Maximum number of elements expected in gradient tensors - pin_memory: Whether to page-lock the reusable grad scratch buffer - """ - cpu_tensor = torch.randn(1, device="cpu") - cpu_param = torch.nn.Parameter(cpu_tensor) - - try: - if isinstance(optimizer_config, list): - pg_configs = optimizer_config - else: - pg_configs = [optimizer_config] - - first_cfg = pg_configs[0] - optimizer = DeepSpeedCPUAdam([cpu_param], - lr=first_cfg["lr"], - betas=first_cfg["betas"], - eps=first_cfg["eps"], - weight_decay=first_cfg["weight_decay"], - amsgrad=first_cfg["amsgrad"]) - for cfg in pg_configs[1:]: - dummy = torch.nn.Parameter(torch.randn(1, device="cpu")) - optimizer.add_param_group({ - "params": [dummy], - "lr": cfg["lr"], - "betas": cfg["betas"], - "eps": cfg["eps"], - "weight_decay": cfg["weight_decay"], - "amsgrad": cfg["amsgrad"], - }) - except KeyError as e: - error_msg = f"Missing required optimizer config key: {e}" - logger.error(error_msg) - result_queue.put({"error": error_msg}) +def _bind_thread(cores: List[int]) -> None: + """Run the calling thread on these cores; the threads it starts later, such as its OpenMP threads, inherit them.""" + if not cores: return + try: + os.sched_setaffinity(0, cores) + except (AttributeError, OSError) as e: + logger.debug(f"Could not set the CPU affinity of a SuperOffload thread: {e}") - # Pre-allocate reusable host buffer for gradients - grad_buffer = _allocate_worker_grad_buffer(max_grad_numel, pin_memory) - - while True: - try: - task = param_queue.get() - - if task is None: - logger.debug("Received termination signal, shutting down worker") - break - - param_data = task[TaskKeys.PARAM_DATA] - param_grad = task[TaskKeys.PARAM_GRAD] - param_group_id = task[TaskKeys.PARAM_GROUP_ID] - sub_group_id = task[TaskKeys.SUB_GROUP_ID] - rollback = task.get(TaskKeys.ROLLBACK, False) - task_lr = task.get(TaskKeys.LR, None) - - logger.debug(f"Processing param_group_id: {param_group_id}, sub_group_id: {sub_group_id}") - - del task[TaskKeys.PARAM_DATA] - del task[TaskKeys.PARAM_GRAD] - task.clear() - - if task_lr is not None: - optimizer.param_groups[param_group_id]['lr'] = task_lr - - grad_numel = param_grad.numel() - if grad_numel > max_grad_numel: - error_msg = ( - f"Gradient size {grad_numel} exceeds pre-allocated buffer size {max_grad_numel}. " - f"This indicates insufficient buffer allocation. Please increase max_grad_numel parameter.") - result_queue.put({"error": error_msg}) - break - - param_grad_cpu = grad_buffer[:grad_numel].view_as(param_grad) - param_grad_cpu.copy_(param_grad, non_blocking=False) - - fp32_param = torch.nn.Parameter(param_data) - fp32_param.grad = param_grad_cpu - - optimizer.param_groups[param_group_id]['params'] = [fp32_param] - - if rollback: - logger.debug(f"Rolling back optimizer state for sub_group_id: {sub_group_id}") - optimizer.rollback_subgroup(sub_group_id) - else: - optimizer.step_subgroup(sub_group_id) - - # Send result back to main process - event_type = EventTypes.ROLLBACK if rollback else EventTypes.ADAM_STEP - result_queue.put({ - TaskKeys.PARAM_GROUP_ID: param_group_id, - TaskKeys.SUB_GROUP_ID: sub_group_id, - ResultKeys.UPDATED_PARAM: fp32_param.data, - ResultKeys.EVENT_TYPE: event_type, - }) - - # Clean up references to free memory - optimizer.param_groups[param_group_id]['params'] = [] - del param_grad_cpu, fp32_param.grad, fp32_param, param_grad, param_data - - except KeyError as e: - error_msg = f"Missing required task key: {e}" - logger.error(error_msg) - result_queue.put({"error": error_msg}) - break - except Exception as e: - error_msg = f"Unexpected error in worker process: {e}" - logger.error(error_msg) - result_queue.put({"error": error_msg}) - break - # Clean up the reusable host buffer. unpin_memory is a no-op on torch and - # frees immediately on the native backend. - if 'grad_buffer' in locals(): - if pin_memory: - get_accelerator().unpin_memory(grad_buffer) - del grad_buffer - logger.debug("Cleaned up worker grad buffer") +def _kernel_args(hyperparams: Dict[str, Any]): + beta1, beta2 = hyperparams["betas"] + return (hyperparams["lr"], beta1, beta2, hyperparams["eps"], hyperparams["weight_decay"], + hyperparams["bias_correction"]) - logger.debug("Worker process terminated") +class SuperOffloadCPUAdam: + """DeepSpeed's CPU Adam for the CPU sub-groups, run by one thread of the training process on the cores given to + CPU Adam. -class SuperOffloadCPUOptimizer: + The native kernels release the GIL, so the training thread goes on queueing backward kernels while a sub-group + steps. The thread runs what is submitted in order. The caller keeps the Adam state of each sub-group and passes it + to each step. + """ - def __init__(self, - optimizer_config: Dict[str, Any], - cpuadam_cores_perc: float = 0.8, - max_grad_numel: int = 1000000, - pin_memory: bool = True) -> None: + def __init__(self, defaults: Dict[str, Any], adamw_mode: bool, fp32_optimizer_states: bool, + cpuadam_cores_perc: float) -> None: + """defaults: lr, betas, eps and weight_decay to create the native optimizer with; each step passes its own.""" if not 0 < cpuadam_cores_perc <= 1: raise ValueError("cpuadam_cores_perc must be between 0 and 1") - - self.max_grad_numel = max_grad_numel - self.mp_context = mp.get_context('spawn') - self.param_queue = self.mp_context.SimpleQueue() - self.result_queue = self.mp_context.SimpleQueue() - - self.cpuadam_process = self.mp_context.Process( - target=superoffload_optimizer_worker, - args=(self.param_queue, self.result_queue, optimizer_config, max_grad_numel, pin_memory), - daemon=True, - ) - self.cpuadam_process.start() - - # Set CPU affinity for better performance isolation - self._set_cpu_affinity(cpuadam_cores_perc) - - def _set_cpu_affinity(self, cpuadam_cores_perc: float) -> None: - """ - Set CPU affinity for the main (Pytorch) process and worker (CPU Adam) process. - - Args: - cpuadam_cores_perc: Percentage of cores to allocate to the worker (CPU Adam) process - """ try: - current_process = psutil.Process() - all_cores = current_process.cpu_affinity() - num_cores = len(all_cores) - - split_idx = int((1 - cpuadam_cores_perc) * num_cores) - pt_cores = all_cores[:split_idx] - cpuadam_cores = all_cores[split_idx:] - - # Set affinity for main process (PyTorch) - current_process.cpu_affinity(pt_cores) - - # Set affinity for optimizer process (CPU Adam) - optimizer_process = psutil.Process(self.cpuadam_process.pid) - optimizer_process.cpu_affinity(cpuadam_cores) - - logger.debug(f"Set CPU affinity - PyTorch cores: {pt_cores}, " - f"Optimizer cores: {cpuadam_cores}") - - except (psutil.AccessDenied, psutil.NoSuchProcess, AttributeError) as e: - logger.debug(f"Could not set CPU affinities for superoffload optimizer process: {e}") - except Exception as e: - logger.warning(f"Unexpected error setting CPU affinity: {e}") - - def async_step(self, - param_group_id: int, - sub_group_id: int, - fp32_param: torch.Tensor, - fp32_grad: torch.Tensor, - rollback: bool = False, - lr: float = None) -> None: - """ - Queue parameter for optimization in the worker process. - """ - if not self.cpuadam_process.is_alive(): - raise RuntimeError("Worker process is not alive") - - task = { - TaskKeys.PARAM_DATA: fp32_param, - TaskKeys.PARAM_GRAD: fp32_grad, - TaskKeys.PARAM_GROUP_ID: param_group_id, - TaskKeys.SUB_GROUP_ID: sub_group_id, - TaskKeys.ROLLBACK: rollback, - } - if lr is not None: - task[TaskKeys.LR] = lr - self.param_queue.put(task) - - def get_result(self, expected_event_type: str = None) -> Optional[Dict[str, Any]]: - """ - Get result from worker process with optional event type validation. - - Args: - expected_event_type (str, optional): Expected event type ('adam_step' or 'rollback'). - If provided, validates that the result matches. - """ - if self.result_queue.empty(): - return None - - result = self.result_queue.get() - - if "error" in result: - raise RuntimeError(f"Error in worker process: {result['error']}") - - # Validate event type if expected_event_type is provided - if expected_event_type is not None: - result_event_type = result.get(ResultKeys.EVENT_TYPE) - if result_event_type != expected_event_type: - raise RuntimeError(f"Event type mismatch: expected '{expected_event_type}', got '{result_event_type}'") - - return result + cores = sorted(os.sched_getaffinity(0)) + except AttributeError: # not Linux + cores = [] + split = int((1 - cpuadam_cores_perc) * len(cores)) + # The training thread keeps the first cores and CPU Adam gets the others, so they do not compete. + _bind_thread(cores[:split]) + self._thread = ThreadPoolExecutor(max_workers=1, + thread_name_prefix="SuperOffloadCPUAdam", + initializer=_bind_thread, + initargs=(cores[split:], )) + + self.fp32_optimizer_states = fp32_optimizer_states + self._ops = CPUAdamBuilder().load() + self._opt_id = DeepSpeedCPUAdam.optimizer_id + DeepSpeedCPUAdam.optimizer_id += 1 + self._ops.create_adam(self._opt_id, defaults["lr"], defaults["betas"][0], defaults["betas"][1], + defaults["eps"], defaults["weight_decay"], adamw_mode, should_log_le("info")) + # Holds the gradient adam_rollback reads: in fp32, with the loss scale removed. + self._grad_buffer = None + + def submit(self, fn: Callable, *args) -> Future: + """Run fn(*args) in the CPU Adam thread, after what was submitted before it.""" + return self._thread.submit(fn, *args) + + def step(self, + state: Dict[str, Any], + weights: torch.Tensor, + grad: torch.Tensor, + hyperparams: Dict[str, Any], + grad_scale: float = 1.0) -> None: + """Step the weights on grad * grad_scale, as DeepSpeedCPUAdam does; state starts empty.""" + if not state: + state["step"] = 0 + state_dtype = torch.float if self.fp32_optimizer_states else weights.dtype + state["exp_avg"] = torch.zeros_like(weights, dtype=state_dtype) + state["exp_avg_sq"] = torch.zeros_like(weights, dtype=state_dtype) + state["step"] += 1 + # adam_update reads bf16 gradients of fp32 weights as they are, so they go to the CPU at half the bytes, and + # multiplies them by grad_scale as it reads them. + self._ops.adam_update(self._opt_id, state["step"], *_kernel_args(hyperparams), weights, grad, state["exp_avg"], + state["exp_avg_sq"], grad_scale) + + def rescale_step(self, state: Dict[str, Any], weights: torch.Tensor, grad: torch.Tensor, + hyperparams: Dict[str, Any], old_scale: float, new_scale: float) -> None: + """Turn the last step of the weights, taken on grad * old_scale with these hyperparameters, into the step on + grad * new_scale, in one pass over the weights and moments.""" + self._ops.adam_rescale_step(self._opt_id, state["step"], *_kernel_args(hyperparams), weights, grad, + state["exp_avg"], state["exp_avg_sq"], old_scale, new_scale) + + def rollback(self, + state: Dict[str, Any], + weights: torch.Tensor, + grad: torch.Tensor, + hyperparams: Dict[str, Any], + grad_scale: float = 1.0) -> None: + """Undo the last step of the weights, which used this gradient, grad_scale and these hyperparameters.""" + # adam_rollback reads gradients in the weights' dtype only. + self._ops.adam_rollback(self._opt_id, state["step"], *_kernel_args(hyperparams), weights, + self._unscaled(grad, weights.dtype, grad_scale), state["exp_avg"], state["exp_avg_sq"]) + state["step"] -= 1 + + def _unscaled(self, grad: torch.Tensor, dtype: torch.dtype, grad_scale: float) -> torch.Tensor: + if grad_scale == 1.0 and grad.dtype == dtype: + return grad + buffer = self._grad_buffer + if buffer is None or buffer.dtype != dtype or buffer.numel() < grad.numel(): + buffer = self._grad_buffer = torch.empty(grad.numel(), dtype=dtype) + unscaled = buffer[:grad.numel()].view_as(grad) + # The step of a sub-group submitted during backward read the gradient before the loss scale is removed. + torch.mul(grad, grad_scale, out=unscaled) + return unscaled def close(self) -> None: - """ - Shutdown the worker process gracefully. - - Sends termination signal to worker and waits for clean shutdown. - If the process doesn't terminate within the timeout, it will be forcefully killed. - """ - if not self.cpuadam_process.is_alive(): - logger.debug("Worker process already terminated") - return - - # Send termination signal - self.param_queue.put(None) - - # Wait for graceful shutdown - self.cpuadam_process.join(timeout=5) - - if self.cpuadam_process.is_alive(): - logger.warning("Optimizer process did not terminate cleanly within timeout, " - "forcefully terminating") - self.cpuadam_process.terminate() - self.cpuadam_process.join(timeout=2) - - # Last resort: kill the process - if self.cpuadam_process.is_alive(): - logger.error("Failed to terminate optimizer process, killing it") - self.cpuadam_process.kill() - self.cpuadam_process.join() - - logger.debug("SuperOffload CPU optimizer closed successfully") + """Finish what was submitted, stop the thread and free the native optimizer.""" + self._thread.shutdown() + self._ops.destroy_adam(self._opt_id) diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index c1393d1569a5..e3246d4fc618 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -314,16 +314,25 @@ def _enforce_optimizer_offload(): backup_gpu_tensor = torch.randn(1, device=get_accelerator().device_name()).to(self.dtype) backup_gpu_param = torch.nn.Parameter(backup_gpu_tensor) assert type(init_optimizer) == DeepSpeedCPUAdam, 'Hybrid Optimizer Only Supports DeepSpeedCPUAdam' - self.backup_optimizer = torch.optim.AdamW([backup_gpu_param], - lr=self.optimizer.param_groups[0]["lr"], - betas=self.optimizer.param_groups[0]["betas"], - eps=self.optimizer.param_groups[0]["eps"], - weight_decay=self.optimizer.param_groups[0]["weight_decay"], - amsgrad=self.optimizer.param_groups[0]["amsgrad"]) - # Multiple param_groups configs for back-up optimizer - if len(self.optimizer.param_groups) > 1: - for i in range(1, len(self.optimizer.param_groups)): - self.backup_optimizer.add_param_group(self.optimizer.param_groups[i]) + if not all(group.get("bias_correction", True) for group in self.optimizer.param_groups): + raise ValueError( + "offload_optimizer.ratio < 1 steps the sub-groups kept on the device with torch Adam, " + "which always applies bias correction; bias_correction=False is not supported.") + # The same update as DeepSpeedCPUAdam: decoupled weight decay, or L2 when adam_w_mode is off. + backup_optimizer_class = torch.optim.AdamW if init_optimizer.adam_w_mode else torch.optim.Adam + self.backup_optimizer = backup_optimizer_class([backup_gpu_param], + lr=self.optimizer.param_groups[0]["lr"], + betas=self.optimizer.param_groups[0]["betas"], + eps=self.optimizer.param_groups[0]["eps"], + weight_decay=self.optimizer.param_groups[0]["weight_decay"], + amsgrad=self.optimizer.param_groups[0]["amsgrad"]) + # Each param group gets its own dict; _sync_backup_optimizer_group copies the hyperparameters before a + # step. A dict shared with self.optimizer would also get the CPU sub-groups that + # _set_fp32_optimizer_param_groups adds for a checkpoint, and keep them after a load replaces + # self.optimizer's dicts. + for group in self.optimizer.param_groups[1:]: + backup_group = {key: value for key, value in group.items() if key != "params"} + self.backup_optimizer.add_param_group(dict(backup_group, params=[])) self._initialize_zenflow_stage3_epilogue(zenflow_config, overlap_comm) @@ -1239,6 +1248,23 @@ def _create_fp16_sub_groups(self, params_group): return sub_groups + def _sync_backup_optimizer_group(self, param_group_id): + """Copy the hyperparameters of a param group to the optimizer of the sub-groups kept on the device. + + LR, momentum and weight decay schedules, set_lr() and checkpoint loads change self.optimizer's groups only. + """ + backup_group = self.backup_optimizer.param_groups[param_group_id] + main_group = self.optimizer.param_groups[param_group_id] + for key in ("lr", "betas", "eps", "weight_decay"): + backup_group[key] = main_group[key] + + def _sub_group_optimizer(self, sub_group_id): + """The optimizer that steps a sub-group and holds its state.""" + if (self.offload_optimizer and self.subgroup_to_device[sub_group_id] != 'cpu' + and not (self.use_muon and self.sub_groups_using_muon[sub_group_id])): + return self.backup_optimizer + return self.optimizer + def _optimizer_step(self, sub_group_id): param_group_id = self.sub_group_to_group_id[sub_group_id] fp32_param = self.fp32_partitioned_groups_flat[sub_group_id] @@ -1260,6 +1286,7 @@ def step_with_gradscaler(optimizer): step_with_gradscaler(self.optimizer) self.optimizer.param_groups[param_group_id]['params'] = [] else: + self._sync_backup_optimizer_group(param_group_id) self.backup_optimizer.param_groups[param_group_id]['params'] = [fp32_param] step_with_gradscaler(self.backup_optimizer) self.backup_optimizer.param_groups[param_group_id]['params'] = [] @@ -3009,6 +3036,10 @@ def dump_post_step_gradients(self): @instrument_w_nvtx def unscale_and_clip_grads(self, sub_group_id, total_norm): + self.fp32_partitioned_groups_flat[sub_group_id].grad.mul_(self._unscale_and_clip_factor(total_norm)) + + def _unscale_and_clip_factor(self, total_norm): + """The factor unscale_and_clip_grads multiplies the gradients by: it removes the loss scale and clips.""" # compute combined scale factor for this group combined_scale = self.loss_scale if self.clip_grad > 0.: @@ -3016,8 +3047,7 @@ def unscale_and_clip_grads(self, sub_group_id, total_norm): clip = ((total_norm / self.loss_scale) + 1e-6) / self.clip_grad clip = torch.clamp(clip, min=1.0) combined_scale = clip * self.loss_scale - - self.fp32_partitioned_groups_flat[sub_group_id].grad.mul_(1. / combined_scale) + return 1. / combined_scale def _check_overflow(self, partition_gradients=True): self.overflow = self.has_overflow(partition_gradients) @@ -3192,7 +3222,8 @@ def _get_fp32_opt_state_partition(self, param, release_swap_buffers, optim_state if optim_state_key is None: fp32_opt_state = fp32_param.narrow(0, dest_offset, num_elements) else: - fp32_opt_state = self.optimizer.state[fp32_param][optim_state_key].narrow(0, dest_offset, num_elements) + optimizer = self._sub_group_optimizer(group_idx) + fp32_opt_state = optimizer.state[fp32_param][optim_state_key].narrow(0, dest_offset, num_elements) if self._swappable_optimizer_subgroup(group_idx) and release_swap_buffers: self._release_swap_buffers(group_idx) @@ -3472,12 +3503,53 @@ def _rigid_state_dict(self): state_dict['ds_zero_partition_groups'] = zero3_partition_group_metadata self._set_fp32_optimizer_param_groups() - state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict() + # The sub-groups kept on the device have their state in the backup optimizer. Save it with the other + # sub-groups', as DeepSpeedCPUAdam keeps it, so that the checkpoint does not depend on offload_optimizer.ratio. + backup_state_params = self._add_backup_optimizer_state() + try: + state_dict[OPTIMIZER_STATE_DICT] = self.optimizer.state_dict() + finally: + for fp32_param in backup_state_params: + del self.optimizer.state[fp32_param] state_dict[FP32_FLAT_GROUPS] = self.fp32_partitioned_groups_flat self._clear_fp32_optimizer_param_groups() return state_dict + def _add_backup_optimizer_state(self): + """Put the Adam state of the sub-groups kept on the device into self.optimizer.state; return their params.""" + added = [] + if getattr(self, "backup_optimizer", None) is None: + return added + for sub_group_id, fp32_param in enumerate(self.fp32_partitioned_groups_flat): + state = self.backup_optimizer.state.get(fp32_param) + if state and self._sub_group_optimizer(sub_group_id) is self.backup_optimizer: + self.optimizer.state[fp32_param] = { + "step": int(state["step"]), + "exp_avg": state["exp_avg"], + "exp_avg_sq": state["exp_avg_sq"], + } + added.append(fp32_param) + return added + + def _move_state_to_backup_optimizer(self): + """Move the loaded Adam state of the sub-groups kept on the device to the backup optimizer.""" + if getattr(self, "backup_optimizer", None) is None: + return + for sub_group_id, fp32_param in enumerate(self.fp32_partitioned_groups_flat): + if self._sub_group_optimizer(sub_group_id) is not self.backup_optimizer: + continue + state = self.optimizer.state.pop(fp32_param, None) + self.backup_optimizer.state.pop(fp32_param, None) + if state: + # torch Adam keeps the moments in the param dtype and the step count as a float32 tensor; + # DeepSpeedCPUAdam keeps an int step and, with fp32_optimizer_states, fp32 moments for bf16 params. + self.backup_optimizer.state[fp32_param] = { + "step": torch.tensor(float(state["step"])), + "exp_avg": state["exp_avg"].to(fp32_param.dtype), + "exp_avg_sq": state["exp_avg_sq"].to(fp32_param.dtype), + } + def state_dict(self): """ Returns a dict containing the current state of this :class:`FP16_Optimizer` instance. @@ -3584,6 +3656,7 @@ def _rigid_load_state_dict(self, state_dict, load_optimizer_states=True): self._set_fp32_optimizer_param_groups() self.optimizer.load_state_dict(state_dict[OPTIMIZER_STATE_DICT]) self._clear_fp32_optimizer_param_groups() + self._move_state_to_backup_optimizer() if self.swap_optimizer: # Purge the swapped optimizer state, it was initialized to the freshly created model and not the checkpoint @@ -3714,6 +3787,7 @@ def load_hp_checkpoint_state_from_checkpoint_dir_stage3(self, checkpoint_dir): self.optimizer.state[fp32_param]['step'] = state_step else: self.optimizer.state[fp32_param][key] = key_tensor + self._move_state_to_backup_optimizer() for param_group in self.optimizer.param_groups: # Generally, the hyperparameters of each parameter should be the same, we can obtain from any parameter. diff --git a/tests/unit/ops/adam/test_cpu_adam.py b/tests/unit/ops/adam/test_cpu_adam.py index 762049e0abfd..422b7e1e6c3b 100644 --- a/tests/unit/ops/adam/test_cpu_adam.py +++ b/tests/unit/ops/adam/test_cpu_adam.py @@ -6,6 +6,8 @@ import math import os import sys +import threading +import time import torch import numpy as np import pytest @@ -111,6 +113,38 @@ def test_cpu_adam_strict_state_updates(model_size, adamw_mode, weight_decay, bia torch.testing.assert_close(state['exp_avg_sq'], ref_exp_avg_sq, rtol=3e-5, atol=2e-6) +def test_cpu_adam_step_releases_gil(): + """A CPU Adam step in another thread lets this thread run, as SuperOffload's worker thread needs.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + + # Large enough that a step takes several ms, longer than the wakeups of this thread are apart. + param = torch.nn.Parameter(torch.zeros(2**26)) + param.grad = torch.full_like(param, 0.5) + optimizer = DeepSpeedCPUAdam([param]) + optimizer.step() + steps = [] + done = threading.Event() + + def step(): + for _ in range(5): + start = time.perf_counter() + optimizer.step() + steps.append((start, time.perf_counter())) + done.set() + + thread = threading.Thread(target=step) + thread.start() + wakeups = [] + while not done.is_set(): + time.sleep(0.001) + wakeups.append(time.perf_counter()) + thread.join() + # Holding the GIL during a step would keep this thread from waking up until the step is done, so none of its + # wakeups would fall inside a step. Being descheduled now and then only loses a few. + inside = sum(any(a + 1e-3 < t < b - 1e-3 for a, b in steps) for t in wakeups) + assert inside >= 10, f"{inside} wakeups during {len(steps)} steps of {sum(b - a for a, b in steps) * 1e3:.0f} ms" + + def test_cpu_adam_rejects_amsgrad(): from deepspeed.ops.adam import DeepSpeedCPUAdam @@ -123,6 +157,130 @@ def test_cpu_adam_rejects_amsgrad(): assert len(optimizer.param_groups) == 1 +def test_cpu_adam_steps_fp32_params_on_bf16_grads(): + """fp32 params step on bf16 gradients as on the same gradients in fp32, and a gradient dtype that has no kernel + is rejected instead of being read as the params' dtype.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + + op = CPUAdamBuilder().load() + opt_id = DeepSpeedCPUAdam.optimizer_id + DeepSpeedCPUAdam.optimizer_id += 1 + lr, beta1, beta2, eps, wd = 1e-3, 0.9, 0.999, 1e-8, 0.01 + n = 100003 # not a multiple of the SIMD block, so the scalar tail runs too + torch.manual_seed(0) + params = torch.randn(n) + grads = [torch.randn(n).to(torch.bfloat16) for _ in range(3)] + results = [] + for grad_dtype in (torch.bfloat16, torch.float): + # A new native optimizer for each, as it carries beta**step over from the steps before. + op.create_adam(opt_id, lr, beta1, beta2, eps, wd, True, False) + p, m, v = params.clone(), torch.zeros(n), torch.zeros(n) + for step, grad in enumerate(grads, 1): + op.adam_update(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad.to(grad_dtype), m, v) + results.append((p, m, v)) + op.destroy_adam(opt_id) + for name, a, b in zip(("params", "exp_avg", "exp_avg_sq"), *results): + assert torch.equal(a, b), name + + op.create_adam(opt_id, lr, beta1, beta2, eps, wd, True, False) + try: + with pytest.raises(RuntimeError, match="grad type"): + op.adam_update(opt_id, 1, lr, beta1, beta2, eps, wd, True, params.clone(), grads[0].half(), torch.zeros(n), + torch.zeros(n)) + finally: + op.destroy_adam(opt_id) + + +def test_cpu_adam_scales_grads_as_it_reads_them(): + """adam_update with grad_scale steps bit for bit as on the gradient scaled before, as unscale_and_clip_grads + scales it, for fp32 and bf16 gradients of fp32 params.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + + op = CPUAdamBuilder().load() + opt_id = DeepSpeedCPUAdam.optimizer_id + DeepSpeedCPUAdam.optimizer_id += 1 + lr, beta1, beta2, eps, wd, scale = 1e-3, 0.9, 0.999, 1e-8, 0.01, 0.37 + n = 100003 # not a multiple of the SIMD block, so the scalar tail runs too + torch.manual_seed(0) + params = torch.randn(n) + for grad_dtype in (torch.bfloat16, torch.float): + grads = [torch.randn(n).to(grad_dtype) for _ in range(3)] + results = [] + for in_kernel in (True, False): + # A new native optimizer for each, as it carries beta**step over from the steps before. + op.create_adam(opt_id, lr, beta1, beta2, eps, wd, True, False) + p, m, v = params.clone(), torch.zeros(n), torch.zeros(n) + for step, grad in enumerate(grads, 1): + if in_kernel: + op.adam_update(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad, m, v, scale) + else: + op.adam_update(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad.float() * scale, m, v) + results.append((p, m, v)) + op.destroy_adam(opt_id) + for name, a, b in zip(("params", "exp_avg", "exp_avg_sq"), *results): + assert torch.equal(a, b), f"{grad_dtype} {name}" + + +@pytest.mark.parametrize("earlier_steps", [0, 3], ids=["first_step", "later_step"]) +@pytest.mark.parametrize("grad_dtype", [torch.bfloat16, torch.float], ids=["bf16", "fp32"]) +def test_cpu_adam_rescale_step(earlier_steps, grad_dtype): + """adam_rescale_step turns a step on grad * old_scale into the step on grad * new_scale from the same state, up to + rounding: 7e-6 of the update for the params and 1e-6 of the largest moment (2.5e-7 and 1e-6 on the first step, + where the moments of before were 0). Rescaling to the same scale leaves the moments as they are, and takes the + update out of the params and adds it back.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + + op = CPUAdamBuilder().load() + opt_id = DeepSpeedCPUAdam.optimizer_id + DeepSpeedCPUAdam.optimizer_id += 1 + lr, beta1, beta2, eps, wd, old_scale, new_scale = 1e-3, 0.9, 0.999, 1e-8, 0.01, 1.0, 0.2 + n = 100003 # not a multiple of the SIMD block, so the scalar tail runs too + torch.manual_seed(0) + params = torch.randn(n) * 0.02 + earlier = [torch.randn(n) * 1e-3 for _ in range(earlier_steps)] + grad = (torch.randn(n) * 1e-3).to(grad_dtype) + step = earlier_steps + 1 + results = [] + for rescale in (True, False): + op.create_adam(opt_id, lr, beta1, beta2, eps, wd, True, False) + p, m, v = params.clone(), torch.zeros(n), torch.zeros(n) + for t, g in enumerate(earlier, 1): + op.adam_update(opt_id, t, lr, beta1, beta2, eps, wd, True, p, g, m, v) + before = p.clone() + if rescale: + op.adam_update(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad, m, v, old_scale) + stepped = (p.clone(), m.clone(), v.clone()) + op.adam_rescale_step(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad, m, v, old_scale, old_scale) + assert torch.equal(m, stepped[1]) and torch.equal(v, stepped[2]) + torch.testing.assert_close(p, stepped[0], rtol=0, atol=1e-5 * (stepped[0] - before).abs().max().item()) + op.adam_rescale_step(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad, m, v, old_scale, new_scale) + else: + op.adam_update(opt_id, step, lr, beta1, beta2, eps, wd, True, p, grad, m, v, new_scale) + results.append((p, m, v)) + op.destroy_adam(opt_id) + (p, m, v), (p_ref, m_ref, v_ref) = results + torch.testing.assert_close(p, p_ref, rtol=0, atol=1e-4 * (p_ref - before).abs().max().item()) + torch.testing.assert_close(m, m_ref, rtol=0, atol=2e-6 * m_ref.abs().max().item()) + torch.testing.assert_close(v, v_ref, rtol=0, atol=1e-5 * v_ref.abs().max().item()) + + +def test_cpu_adam_rescale_step_rejects_l2_weight_decay(): + """Adam's L2 decay adds the params of before the step to the gradient, which a rescale cannot recover.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + + op = CPUAdamBuilder().load() + opt_id = DeepSpeedCPUAdam.optimizer_id + DeepSpeedCPUAdam.optimizer_id += 1 + op.create_adam(opt_id, 1e-3, 0.9, 0.999, 1e-8, 0.01, False, False) + try: + p, g, m, v = torch.randn(64), torch.randn(64), torch.zeros(64), torch.zeros(64) + op.adam_update(opt_id, 1, 1e-3, 0.9, 0.999, 1e-8, 0.01, True, p, g, m, v) + with pytest.raises(RuntimeError, match="AdamW only"): + op.adam_rescale_step(opt_id, 1, 1e-3, 0.9, 0.999, 1e-8, 0.01, True, p, g, m, v, 1.0, 0.5) + finally: + op.destroy_adam(opt_id) + + @pytest.mark.parametrize('dtype', [torch.half, torch.bfloat16, torch.float], ids=["fp16", "bf16", "fp32"]) @pytest.mark.parametrize('model_size', [ @@ -443,6 +601,23 @@ def test_rollback_subgroup_basic(self, dtype): # Verify step count decremented assert optimizer.state[subgroup_id]['step'] == 0, "Step count should be decremented after rollback" + def test_rollback_of_first_step_keeps_exp_avg_sq_non_negative(self): + """exp_avg_sq is 0 before the first step. Rolling that step back must not leave it negative, because sqrt() + of a negative value in the next step turns the parameter into NaN.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + + model_size = 4096 + param = torch.nn.Parameter(torch.randn(model_size, device='cpu')) + optimizer = DeepSpeedCPUAdam([param]) + param.grad = torch.randn(model_size, device='cpu') + + optimizer.step_subgroup(0) + optimizer.rollback_subgroup(0) + assert (optimizer.state[0]['exp_avg_sq'] >= 0).all() + + optimizer.step_subgroup(0) + assert torch.isfinite(param.data).all() + def test_rollback_subgroup_uninitialized_error(self): """Test that rollback_subgroup raises error for uninitialized subgroup.""" from deepspeed.ops.adam import DeepSpeedCPUAdam diff --git a/tests/unit/v1/pin_memory/test_offload_route.py b/tests/unit/v1/pin_memory/test_offload_route.py index ce50bff43a97..64925b17e5bc 100644 --- a/tests/unit/v1/pin_memory/test_offload_route.py +++ b/tests/unit/v1/pin_memory/test_offload_route.py @@ -132,20 +132,3 @@ def _recording_synchronize(): gc.collect() # Released once the copies are known to have completed. assert begin not in manager._ranges - - -def test_superoffload_grad_buffer_unpinned_when_disabled(): - """offload_optimizer.pin_memory=False must allocate a regular CPU buffer.""" - from deepspeed.runtime.superoffload.superoffload_utils import _allocate_worker_grad_buffer - buffer = _allocate_worker_grad_buffer(32, pin_memory=False) - assert get_accelerator().is_pinned(buffer) is False - - -def test_superoffload_grad_buffer_pinned_when_enabled(monkeypatch): - """offload_optimizer.pin_memory=True keeps the current pinned allocation.""" - _require_native(monkeypatch) - from deepspeed.runtime.superoffload.superoffload_utils import _allocate_worker_grad_buffer - accel = get_accelerator() - buffer = _allocate_worker_grad_buffer(32, pin_memory=True) - assert accel.is_pinned(buffer) is True - assert accel.unpin_memory(buffer) is True diff --git a/tests/unit/v1/zero/test_invalid_grad_norm.py b/tests/unit/v1/zero/test_invalid_grad_norm.py index 2039e549d306..4315a45ce8f8 100644 --- a/tests/unit/v1/zero/test_invalid_grad_norm.py +++ b/tests/unit/v1/zero/test_invalid_grad_norm.py @@ -154,18 +154,19 @@ def test_bf16_optimizer_reports_invalid_step_and_recovers(monkeypatch): def test_superoffload_rolls_back_only_subgroups_submitted_this_step(): optimizer = object.__new__(SuperOffloadOptimizer_Stage3) - optimizer._submitted_cpu_sub_groups = {1} + optimizer._submitted_cpu_sub_groups = {1: SimpleNamespace(hyperparams={"lr": 0.1}, grad_scale=1.0)} + optimizer._half_cpu_grads = {} optimizer.sub_group_to_group_id = {0: 10, 1: 11} parameter = SimpleNamespace(data=torch.tensor([1.0]), grad=SimpleNamespace(data=torch.tensor([2.0]))) optimizer.fp32_partitioned_groups_flat = [parameter, parameter] rollbacks = [] - optimizer._sync_cpu_optimizer_step = lambda *args, **kwargs: rollbacks.append((args, kwargs)) + optimizer._cpu_step_and_wait = lambda *args, **kwargs: rollbacks.append((args, kwargs)) - optimizer._handle_overflow_rollback() + optimizer._rollback_submitted_cpu_sub_groups() - assert [args[1] for args, _ in rollbacks] == [1] + assert [args[0] for args, _ in rollbacks] == [1] assert rollbacks[0][1]["rollback"] - assert optimizer._submitted_cpu_sub_groups == set() + assert optimizer._submitted_cpu_sub_groups == {} class TestZeROBFloat16Stage3InvalidNormRecovery(DistributedTest): diff --git a/tests/unit/v1/zero/test_superoffload.py b/tests/unit/v1/zero/test_superoffload.py new file mode 100644 index 000000000000..c88c3f50412e --- /dev/null +++ b/tests/unit/v1/zero/test_superoffload.py @@ -0,0 +1,447 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""SuperOffload against ZeRO-3 CPU offload, which applies the same DeepSpeedCPUAdam update in the main process.""" + +import time + +import pytest +import torch +import torch.utils.checkpoint + +import deepspeed +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import CPUAdamBuilder +from unit.common import DistributedTest + +if not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]: + pytest.skip("cpu-adam is not compatible", allow_module_level=True) +if get_accelerator().device_name() != "cuda": + pytest.skip("SuperOffload supports only CUDA", allow_module_level=True) + +HIDDEN = 64 + + +class PartlyUsedModel(torch.nn.Module): + """`a` is always used. `b` and `c` are used, skipped, multiplied by 0, which gives them a zero gradient, or used + twice, which gives them two gradients in one step.""" + + def __init__(self, hidden=HIDDEN): + super().__init__() + self.hidden = hidden + self.a = torch.nn.Linear(hidden, hidden) + self.b = torch.nn.Linear(hidden, hidden) + self.c = torch.nn.Linear(hidden, hidden) + + def forward(self, x, uses): + h = torch.nn.functional.gelu(self.a(x)) + for name in ("b", "c"): + use = uses.get(name, "use") + if use == "use": + h = h + getattr(self, name)(x) + elif use == "zero": + h = h + 0 * getattr(self, name)(x) + elif use in ("twice", "twice_in_one_backward"): + layer = getattr(self, name) + if use == "twice": + # The backward of a reentrant checkpoint runs on its own, so the param's gradient hook fires twice. + again = torch.utils.checkpoint.checkpoint(layer, h, use_reentrant=True) + else: + again = layer(h) + h = h + layer(x) + again + return h.float().pow(2).mean() + + +def _config(super_offload, + clip=0.0, + ga=1, + precision="bf16", + ratio=1.0, + sub_group_size=1e9, + reduce_bucket_size=None, + **optimizer_params): + optimizer_params.setdefault("lr", 3e-3) + optimizer_params.setdefault("weight_decay", 0.01) + optimizer_type = optimizer_params.pop("type", "AdamW") + config = { + "train_micro_batch_size_per_gpu": 4, + "gradient_accumulation_steps": ga, + "gradient_clipping": clip, + "steps_per_print": 1000, + "optimizer": { + "type": optimizer_type, + "params": optimizer_params + }, + "zero_optimization": { + "stage": 3, + "sub_group_size": sub_group_size, + "offload_optimizer": { + "device": "cpu", + "ratio": ratio, + "super_offload": super_offload, + }, + }, + } + if reduce_bucket_size is not None: + config["zero_optimization"]["reduce_bucket_size"] = reduce_bucket_size + if precision == "fp16": + config["fp16"] = {"enabled": True, "initial_scale_power": 24} + elif precision != "fp32": + config["bf16"] = { + "enabled": True, + "bf16_master_weights_and_grads": precision in ("bf16_master", "bf16_states"), + "bf16_optimizer_states": precision == "bf16_states", + } + return config + + +def _engine(config, two_groups=False, hidden=HIDDEN): + torch.manual_seed(0) + model = PartlyUsedModel(hidden) + params = model.parameters() + if two_groups: + # Weights with weight decay and biases without, as most training scripts set them up. + named = list(model.named_parameters()) + params = [ + dict(params=[p for n, p in named if n.endswith("weight")]), + dict(params=[p for n, p in named if n.endswith("bias")], weight_decay=0.0) + ] + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=params, config=config) + return engine + + +def _masters(engine): + return torch.cat([t.detach().float().flatten().cpu() for t in engine.optimizer.fp32_partitioned_groups_flat]) + + +def _close(engine): + if hasattr(engine.optimizer, "cpu_adam"): + engine.optimizer.cpu_adam.close() + + +def _train(engine, steps, uses_for, first_step=0, before_step=None): + """Train `steps` optimizer steps and return the master weights after each one. + + uses_for(step, micro_step) gives the `uses` of PartlyUsedModel.forward. before_step(engine, step), if given, runs + before each step, as an LR or momentum schedule would. + """ + dtype = torch.float32 + if engine.fp16_enabled(): + dtype = torch.float16 + elif engine.bfloat16_enabled(): + dtype = torch.bfloat16 + generator = torch.Generator().manual_seed(1 + first_step + 100 * dist.get_rank()) + masters = [] + for step in range(first_step, first_step + steps): + if before_step is not None: + before_step(engine, step) + for micro_step in range(engine.gradient_accumulation_steps()): + x = torch.randn(4, engine.module.hidden, generator=generator).to(engine.device, dtype) + loss = engine(x, uses_for(step, micro_step)) + engine.backward(loss) + engine.step() + masters.append(_masters(engine)) + return masters + + +def _all_used(step, micro_step): + return {} + + +def _b_unused_on_even_steps(unused): + return lambda step, micro_step: {"b": "use" if step % 2 else unused} + + +def _b_used_twice(unused): + # ZeRO-3 keeps only the last gradient when the hook of a param fires twice in a step, so the reference uses the + # param twice in one backward, where autograd sums the two gradients. + use = "twice_in_one_backward" if unused == "zero" else "twice" + return lambda step, micro_step: {"b": use} + + +def _b_and_c_in_one_micro_step(unused): + return lambda step, micro_step: {"b": unused, "c": "use"} if micro_step else {"b": "use", "c": unused} + + +# name: (config overrides, uses_for(unused), compare exactly) +CASES = { + "all_used": (dict(), lambda unused: _all_used, True), + "unused_param": (dict(), _b_unused_on_even_steps, True), + # Clipping is on but never triggers, so the CPU steps still run during backward and must stay exact. + "unused_param_clip": (dict(clip=100.0), _b_unused_on_even_steps, True), + "grad_accumulation": (dict(ga=2), _b_and_c_in_one_micro_step, True), + # The second gradient adds to the first in the gradient accumulation dtype, bf16 or fp32. ZeRO-3 takes a param + # once per reduce bucket, so a small bucket keeps the two gradients apart. + "param_used_twice": (dict(reduce_bucket_size=1000), _b_used_twice, True), + "param_used_twice_fp32": (dict(reduce_bucket_size=1000, precision="fp32"), _b_used_twice, True), + "adam": (dict(type="Adam", adam_w_mode=False), lambda unused: _all_used, True), + "no_bias_correction": (dict(bias_correction=False), lambda unused: _all_used, True), + "bf16_master_weights": (dict(precision="bf16_master"), lambda unused: _all_used, True), + "bf16_optimizer_states": (dict(precision="bf16_states"), lambda unused: _all_used, True), + "two_param_groups": (dict(two_groups=True, sub_group_size=4500), _b_unused_on_even_steps, True), + # Clipping turns the steps taken during backward into steps on the clipped gradients, exact up to rounding. In + # fp32 the rounding stays small; bf16 weights would amplify it. + "clipping": (dict(clip=0.05, precision="fp32"), lambda unused: _all_used, False), + # A step cannot be undone with a beta of 0 (the rollback divides by the betas), so those CPU steps wait for + # step() and stay exact. + "clipping_beta1_zero": (dict(clip=0.05, precision="fp32", betas=[0.0, 0.999]), lambda unused: _all_used, True), + "clipping_beta2_zero": (dict(clip=0.05, precision="fp32", betas=[0.9, 0.0]), lambda unused: _all_used, True), + # Overflow steps are skipped; the steps taken during backward are undone, exact up to rounding. + "fp16_overflow": (dict(precision="fp16", sub_group_size=4500), lambda unused: _all_used, False), +} + + +class TestSuperOffloadMatchesZeRO3(DistributedTest): + world_size = 1 + + @pytest.mark.parametrize("case", list(CASES)) + def test(self, case): + overrides, uses_for, exact = CASES[case] + overrides = dict(overrides) + two_groups = overrides.pop("two_groups", False) + # The reference gives an unused param an exactly zero gradient, which SuperOffload must give a param that + # gets no gradient (as ZeRO-1/2 do). + engine = _engine(_config(super_offload=False, **overrides), two_groups) + expected = _train(engine, 8, uses_for("zero")) + _close(engine) + engine = _engine(_config(super_offload=True, **overrides), two_groups) + actual = _train(engine, 8, uses_for("skip")) + _close(engine) + _assert_same(actual, expected, exact) + + +def _assert_same(actual, expected, exact, atol=1e-5): + for step, (a, e) in enumerate(zip(actual, expected)): + assert torch.isfinite(a).all(), f"step {step}" + if exact: + assert torch.equal(a, e), f"step {step}: max diff {(a - e).abs().max().item()}" + else: + torch.testing.assert_close(a, e, rtol=0, atol=atol, msg=f"step {step}") + + +def _change_hyperparams(engine, step): + for group in engine.optimizer.param_groups: + group["betas"] = (0.9 - 0.05 * step, 0.999 - 0.01 * step) + group["eps"] = 1e-8 * 10**step + group["weight_decay"] = 0.01 + 0.1 * step + + +class TestSuperOffloadFollowsZeRO3(DistributedTest): + world_size = 1 + + def test_hyperparameter_changes(self): + """Momentum and weight decay schedules change betas and weight_decay of the param groups between steps.""" + results = [] + for super_offload in (False, True): + engine = _engine(_config(super_offload=super_offload), two_groups=True) + results.append(_train(engine, 6, _all_used, before_step=_change_hyperparams)) + _close(engine) + _assert_same(results[1], results[0], exact=True) + + def test_large_grads(self): + """The bf16 gradient of a 4090 x 4090 weight, 31.9 MiB, goes to the CPU in two pieces of at most 16 MiB. A CPU + step that did not wait for the second piece would read part of the previous gradient. Only to pinned memory + are the copies asynchronous.""" + results = [] + for super_offload in (False, True): + config = _config(super_offload=super_offload) + config["zero_optimization"]["offload_optimizer"]["pin_memory"] = True + engine = _engine(config, hidden=4090) + results.append(_train(engine, 3, _all_used)) + _close(engine) + _assert_same(results[1], results[0], exact=True) + + def test_fp16_loss_scale_increase(self): + """Dynamic loss scaling raises the scale in step(), after the CPU steps taken during backward used the old + one; those steps are changed to the new one, exact up to rounding, which fp16 weights amplify (1.4e-5 after 6 + steps; 3.9e-3 when the steps are kept).""" + results = [] + for super_offload in (False, True): + config = _config(super_offload=super_offload, precision="fp16") + config["fp16"] = {"enabled": True, "initial_scale_power": 8, "loss_scale_window": 2, "hysteresis": 1} + engine = _engine(config) + results.append(_train(engine, 6, _all_used)) + assert engine.optimizer.loss_scale > 2**8 + _close(engine) + _assert_same(results[1], results[0], exact=False, atol=1e-4) + + def test_cpu_steps_run_during_backward(self): + engine = _engine(_config(super_offload=True)) + cpu_adam = engine.optimizer.cpu_adam + calls = [] + + def recording(name): + update = getattr(cpu_adam, name) + + def record(state, weights, grad, *args): + calls.append((name, grad.dtype)) + update(state, weights, grad, *args) + + return record + + cpu_adam.step, cpu_adam.rollback = recording("step"), recording("rollback") + x = torch.randn(4, HIDDEN).to(engine.device, torch.bfloat16) + engine.backward(engine(x, {})) + cpu_sub_groups = list(engine.optimizer.subgroup_to_device.values()).count("cpu") + # The sub-groups step once their gradients reach the CPU, without waiting for step(), on the bf16 gradients + # as they went to the CPU. + deadline = time.monotonic() + 30 + while len(calls) < cpu_sub_groups and time.monotonic() < deadline: + time.sleep(0.01) + assert calls == [("step", torch.bfloat16)] * cpu_sub_groups + engine.step() + assert calls == [("step", torch.bfloat16)] * cpu_sub_groups + # CPU Adam steps the weights where they are, without moving them to shared memory. + assert not any(t.is_shared() for t in engine.optimizer.fp32_partitioned_groups_flat) + _close(engine) + + def test_clipping_bf16(self): + """Clipping with bf16 gradients, which the worker scales as it reads them. Where a master differs by an ulp, + its bf16 weight can round the other way, which the next steps amplify: 3.5e-7 after 2 steps, 6e-5 after 8 + (3e-5 when the steps of backward are undone and taken again), so only 2 steps are compared.""" + results = [] + for super_offload in (False, True): + engine = _engine(_config(super_offload=super_offload, clip=0.05)) + results.append(_train(engine, 2, _all_used)) + _close(engine) + _assert_same(results[1], results[0], exact=False) + + def test_clipped_step_rescales_cpu_steps(self): + """A step that clips changes each CPU step taken during backward to the clipped gradients in one pass, + instead of undoing it and stepping again.""" + engine = _engine(_config(super_offload=True, clip=0.05)) + cpu_adam = engine.optimizer.cpu_adam + calls = [] + + def recording(name): + update = getattr(cpu_adam, name) + + def record(*args): + calls.append(name) + update(*args) + + return record + + for name in ("step", "rescale_step", "rollback"): + setattr(cpu_adam, name, recording(name)) + x = torch.randn(4, HIDDEN).to(engine.device, torch.bfloat16) + engine.backward(engine(x, {})) + engine.step() + assert engine.get_global_grad_norm() > 0.05 + cpu_sub_groups = list(engine.optimizer.subgroup_to_device.values()).count("cpu") + assert calls == ["step"] * cpu_sub_groups + ["rescale_step"] * cpu_sub_groups + _close(engine) + + def test_cpu_buffers_pinned(self): + config = _config(super_offload=True) + config["zero_optimization"]["offload_optimizer"]["pin_memory"] = True + engine = _engine(config) + # The gradient copies of backward and the weight copies of step() stay asynchronous only from page-locked + # memory; ZeRO-3 pins the gradients but not the fp32 weights. + assert all(t.is_pinned() and t.grad.is_pinned() for t in engine.optimizer.fp32_partitioned_groups_flat) + assert engine.optimizer._lp_staging_stream is not None + _close(engine) + + +class InfInjectionModel(torch.nn.Module): + """With sub_group_size 1000 and ratio 0.5, `lin` is a sub-group kept on the GPU and `emb` a CPU sub-group.""" + + def __init__(self): + super().__init__() + self.lin = torch.nn.Linear(32, 32) + self.emb = torch.nn.Embedding(512, 32) + + def forward(self, tokens, inject): + loss = self.lin(self.emb(tokens)).float().pow(2).mean() + # inf * weight puts an inf gradient into `lin` only, so `emb` gets a finite gradient in the skipped step. + return loss + self.lin.weight[0, 0].float() * inject + + +class TestSuperOffloadSkippedStep(DistributedTest): + world_size = 1 + + def test_sub_group_with_finite_grads(self): + """A step skipped for an inf gradient elsewhere undoes the CPU sub-group's step; its Adam state must stay + valid. Each step looks up new embedding rows, which have zero Adam state before their first gradient.""" + masters = {} + for super_offload in (False, True): + config = _config(super_offload=super_offload, ratio=0.5, sub_group_size=1000, lr=1e-3, weight_decay=0.0) + torch.manual_seed(0) + model = InfInjectionModel() + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + masters[super_offload] = [] + for step in range(6): + tokens = (torch.arange(step * 64, step * 64 + 64).view(4, 16) % 512).to(engine.device) + loss = engine(tokens, float("inf") if step == 1 else 0.0) + engine.backward(loss) + engine.step() + masters[super_offload].append(_masters(engine)) + assert engine.skipped_steps == 1 + _close(engine) + for step, (a, e) in enumerate(zip(masters[True], masters[False])): + assert torch.isfinite(a).all(), f"step {step}" + # Undoing a step is exact up to rounding, which Adam amplifies for the rows whose state was zero. + torch.testing.assert_close(a, e, rtol=0, atol=1e-4, msg=f"step {step}") + + +class TestSuperOffloadTwoRanks(DistributedTest): + world_size = 2 + + def test_unused_param(self): + engine = _engine(_config(super_offload=False)) + expected = _train(engine, 6, _b_unused_on_even_steps("zero")) + _close(engine) + engine = _engine(_config(super_offload=True)) + actual = _train(engine, 6, _b_unused_on_even_steps("skip")) + _close(engine) + for step, (a, e) in enumerate(zip(actual, expected)): + assert torch.equal(a, e), f"step {step}: max diff {(a - e).abs().max().item()}" + + +def _resume_matches_uninterrupted(config, checkpoint_dir, two_groups): + """Train 4 steps, save, train 4 more; a new engine loads the checkpoint and trains the same 4 steps.""" + engine = _engine(config, two_groups) + _train(engine, 4, _all_used) + engine.save_checkpoint(checkpoint_dir, tag="t") + uninterrupted = _train(engine, 4, _all_used, first_step=4) + _close(engine) + engine = _engine(config, two_groups) + engine.load_checkpoint(checkpoint_dir, tag="t") + resumed = _train(engine, 4, _all_used, first_step=4) + _close(engine) + for step, (r, u) in enumerate(zip(resumed, uninterrupted)): + assert torch.equal(r, u), f"step {step} after the resume: max diff {(r - u).abs().max().item()}" + + +class TestSuperOffloadCheckpoint(DistributedTest): + world_size = 2 + + # Betas of 0.5 make beta**t exact, so the resume must be bit-exact (CPU Adam recomputes beta**t on resume). + @pytest.mark.parametrize("two_groups", [False, True]) + def test_resume(self, tmpdir, two_groups): + config = _config(super_offload=True, sub_group_size=4500, betas=[0.5, 0.5]) + _resume_matches_uninterrupted(config, str(tmpdir), two_groups) + + +def test_cpu_adam_steps_and_rolls_back_on_bf16_grads(): + """A step on a bf16 gradient matches DeepSpeedCPUAdam on the same gradient in fp32, and the rollback undoes it.""" + from deepspeed.ops.adam import DeepSpeedCPUAdam + from deepspeed.runtime.superoffload.superoffload_utils import SuperOffloadCPUAdam + + hyperparams = dict(lr=1e-2, betas=(0.9, 0.999), eps=1e-8, weight_decay=0.01, bias_correction=True) + torch.manual_seed(0) + weights, grad = torch.randn(4099), torch.randn(4099).to(torch.bfloat16) + expected = torch.nn.Parameter(weights.clone()) + expected.grad = grad.float() + DeepSpeedCPUAdam([expected], **hyperparams).step() + original = weights.clone() + # All the cores go to CPU Adam, which leaves the affinity of this thread as it is. + cpu_adam = SuperOffloadCPUAdam(hyperparams, adamw_mode=True, fp32_optimizer_states=True, cpuadam_cores_perc=1.0) + try: + state = {} + cpu_adam.submit(cpu_adam.step, state, weights, grad, hyperparams).result() + assert torch.equal(weights, expected.data) + cpu_adam.submit(cpu_adam.rollback, state, weights, grad, hyperparams).result() + torch.testing.assert_close(weights, original, rtol=0, atol=1e-6) + finally: + cpu_adam.close() diff --git a/tests/unit/v1/zero/test_zero_partial_offload.py b/tests/unit/v1/zero/test_zero_partial_offload.py new file mode 100644 index 000000000000..dd0c8064d0d1 --- /dev/null +++ b/tests/unit/v1/zero/test_zero_partial_offload.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""ZeRO-3 with offload_optimizer.ratio < 1, where the sub-groups kept on the device are stepped by a torch Adam.""" + +import glob +import os + +import pytest +import torch + +import deepspeed +import deepspeed.comm as dist +from deepspeed.accelerator import get_accelerator +from deepspeed.ops.op_builder import CPUAdamBuilder +from deepspeed.utils import safe_get_full_optimizer_state +from unit.common import DistributedTest + +if not deepspeed.ops.__compatible_ops__[CPUAdamBuilder.NAME]: + pytest.skip("cpu-adam is not compatible", allow_module_level=True) + +HIDDEN = 64 + + +class Net(torch.nn.Module): + + def __init__(self): + super().__init__() + self.layers = torch.nn.ModuleList(torch.nn.Linear(HIDDEN, HIDDEN) for _ in range(3)) + + def forward(self, x): + for layer in self.layers: + x = torch.nn.functional.gelu(layer(x)) + return x.float().pow(2).mean() + + +def _config(ratio, super_offload=False, **optimizer_params): + optimizer_params.setdefault("lr", 3e-3) + optimizer_params.setdefault("weight_decay", 0.01) + # 3 sub-groups: the first two weights, the last weight, the biases. Ratio 0.5 keeps the first on the device. + return { + "train_micro_batch_size_per_gpu": 4, + "steps_per_print": 1000, + "optimizer": { + "type": "AdamW", + "params": optimizer_params + }, + "zero_optimization": { + "stage": 3, + "sub_group_size": 4500, + "offload_optimizer": { + "device": "cpu", + "ratio": ratio, + "super_offload": super_offload + }, + }, + "bf16": { + "enabled": True + }, + } + + +def _engine(config): + torch.manual_seed(0) + model = Net() + named = list(model.named_parameters()) + # Weights with weight decay and biases without, as most training scripts set them up. + params = [ + dict(params=[p for n, p in named if n.endswith("weight")]), + dict(params=[p for n, p in named if n.endswith("bias")], weight_decay=0.0) + ] + engine, _, _, _ = deepspeed.initialize(model=model, model_parameters=params, config=config) + return engine + + +def _train(engine, steps, first_step=0): + generator = torch.Generator().manual_seed(1 + first_step + 100 * dist.get_rank()) + dtype = torch.bfloat16 if engine.bfloat16_enabled() else torch.float32 + masters = [] + for _ in range(steps): + x = torch.randn(4, HIDDEN, generator=generator).to(engine.device, dtype) + loss = engine(x) + engine.backward(loss) + engine.step() + masters.append( + torch.cat([t.detach().float().flatten().cpu() for t in engine.optimizer.fp32_partitioned_groups_flat])) + return masters + + +def _close(engine): + if hasattr(engine.optimizer, "cpu_adam"): + engine.optimizer.cpu_adam.close() + + +def _device_params(engine): + optimizer = engine.optimizer + return [optimizer.fp32_partitioned_groups_flat[i] for i, d in optimizer.subgroup_to_device.items() if d != "cpu"] + + +def _adam_state(engine): + return { + (name, key): safe_get_full_optimizer_state(p, key).cpu() + for name, p in engine.module.named_parameters() + for key in ("exp_avg", "exp_avg_sq") + } + + +class TestPartialOffloadHyperparameters(DistributedTest): + world_size = 1 + + @pytest.mark.parametrize("super_offload", [False, True]) + def test_first_step_at_zero_lr(self, super_offload): + if super_offload and get_accelerator().device_name() != "cuda": + pytest.skip("SuperOffload supports only CUDA") + config = _config(ratio=0.5, super_offload=super_offload) + config["scheduler"] = { + "type": "WarmupLR", + "params": { + "warmup_min_lr": 0.0, + "warmup_max_lr": 3e-3, + "warmup_num_steps": 4, + "warmup_type": "linear" + } + } + engine = _engine(config) + assert _device_params(engine) and "cpu" in engine.optimizer.subgroup_to_device.values() + before = torch.cat([t.detach().float().flatten().cpu() for t in engine.optimizer.fp32_partitioned_groups_flat]) + after = _train(engine, 1)[0] + assert torch.equal(after, before) + # A momentum or weight decay schedule changes the other hyperparameters of the param groups. + for group in engine.optimizer.param_groups: + group["betas"], group["eps"], group["weight_decay"] = (0.8, 0.99), 1e-6, 0.2 + _train(engine, 1) + optimizer = engine.optimizer + # (The scheduler has already set the lr of the next step.) + for i in {optimizer.sub_group_to_group_id[i] for i, d in optimizer.subgroup_to_device.items() if d != "cpu"}: + for key in ("betas", "eps", "weight_decay"): + assert optimizer.backup_optimizer.param_groups[i][key] == optimizer.param_groups[i][key], key + _close(engine) + + +class TestPartialOffloadCheckpoint(DistributedTest): + world_size = 2 + + def test_resume(self, tmpdir): + """Betas of 0.5 make beta**t exact, so the resume is bit-exact (CPU Adam recomputes beta**t on resume).""" + config = _config(ratio=0.5, betas=[0.5, 0.5]) + engine = _engine(config) + _train(engine, 4) + engine.save_checkpoint(str(tmpdir), tag="t") + uninterrupted = _train(engine, 4, first_step=4) + engine = _engine(config) + engine.load_checkpoint(str(tmpdir), tag="t") + resumed = _train(engine, 4, first_step=4) + for step, (r, u) in enumerate(zip(resumed, uninterrupted)): + assert torch.equal(r, u), f"step {step} after the resume: max diff {(r - u).abs().max().item()}" + + +class TestPartialOffloadCheckpointPortability(DistributedTest): + world_size = 1 + + @pytest.mark.parametrize("load_ratio", [0.25, 0.75, 1.0]) + def test_load_with_another_ratio(self, tmpdir, load_ratio): + engine = _engine(_config(ratio=0.5)) + _train(engine, 3) + saved = _adam_state(engine) + engine.save_checkpoint(str(tmpdir), tag="t") + engine = _engine(_config(ratio=load_ratio)) + engine.load_checkpoint(str(tmpdir), tag="t") + loaded = _adam_state(engine) + for key, value in saved.items(): + assert torch.equal(loaded[key], value), key + _train(engine, 1) + engine.save_checkpoint(str(tmpdir), tag="again") + + def test_resume_with_bf16_master_weights(self, tmpdir): + """torch Adam keeps the moments of the device sub-groups in the param dtype, also after a load that returns + fp32 moments for bf16 params (DeepSpeedCPUAdam with fp32_optimizer_states).""" + config = _config(ratio=0.5) + config["bf16"]["bf16_master_weights_and_grads"] = True + engine = _engine(config) + _train(engine, 2) + engine.save_checkpoint(str(tmpdir), tag="t") + engine = _engine(config) + engine.load_checkpoint(str(tmpdir), tag="t") + assert torch.isfinite(_train(engine, 1)[0]).all() + for fp32_param in _device_params(engine): + state = engine.optimizer.backup_optimizer.state[fp32_param] + assert state["exp_avg"].dtype == state["exp_avg_sq"].dtype == fp32_param.dtype + + def test_load_checkpoint_without_device_state(self, tmpdir): + """Checkpoints saved before the state of the sub-groups kept on the device was saved have none for them.""" + engine = _engine(_config(ratio=0.5)) + _train(engine, 3) + engine.save_checkpoint(str(tmpdir), tag="t") + device_ids = [i for i, d in engine.optimizer.subgroup_to_device.items() if d != "cpu"] + path, = glob.glob(os.path.join(str(tmpdir), "t", "*zero_pp_rank_0_mp_rank_00_optim_states.pt")) + checkpoint = torch.load(path, weights_only=False) + for i in device_ids: + checkpoint["optimizer_state_dict"]["optimizer_state_dict"]["state"].pop(i, None) + torch.save(checkpoint, path) + + engine = _engine(_config(ratio=0.5)) + engine.load_checkpoint(str(tmpdir), tag="t") + _train(engine, 1) + # Only the sub-groups kept on the device are stepped by the backup optimizer. + assert set(engine.optimizer.backup_optimizer.state) == set(_device_params(engine)) + engine.save_checkpoint(str(tmpdir), tag="again") + + +class TestPartialOffloadAdamSettings(DistributedTest): + world_size = 1 + + def test_l2_weight_decay(self): + """With adam_w_mode off, the sub-groups kept on the device apply L2 weight decay, as DeepSpeedCPUAdam does.""" + results = [] + for ratio in (1.0, 0.5): + config = _config(ratio=ratio, weight_decay=0.1) + config["optimizer"] = {"type": "Adam", "params": {**config["optimizer"]["params"], "adam_w_mode": False}} + del config["bf16"] + engine = _engine(config) + results.append(_train(engine, 4)) + for step, (partial, cpu) in enumerate(zip(results[1], results[0])): + # torch Adam and DeepSpeedCPUAdam round differently. + torch.testing.assert_close(partial, cpu, rtol=0, atol=1e-6, msg=f"step {step}") + + def test_bias_correction_off_is_rejected(self): + config = _config(ratio=0.5) + config["optimizer"]["params"]["bias_correction"] = False + with pytest.raises(ValueError, match="bias_correction"): + _engine(config)