From c1ba6a104f56327d6cf0fb605e32cbdc53e5c174 Mon Sep 17 00:00:00 2001 From: hipudding Date: Tue, 14 Nov 2023 21:37:16 +0800 Subject: [PATCH] [CANN] Support cpu offload optimizer for Ascend NPU (#4568) Support cpu_adam, cpu_adagrad and cpu_lion optimizer for Ascend NPU. All these optimizer are running on host, the difference between each backend is the way to copy params back to device. This commit add a new symbol called "__ENABLE_CANN__". This symbol can compile code adapted to NPU. The NPU builder adds the required header files and libraries for compiling, according to CANN's compilation manual. Note that there's no FusedLion implementation for NPU, test_cpu_lion test case should disabled until FusedLion optimizer implemented. Besides, when NPU is selected as the accelerator, ds_report will show torch_npu and CANN informations. With this PR, deepspeed test cases in [huggingface/accelerate](https://github.com/huggingface/accelerate/tree/main/tests/deepspeed) are all passed. It's a part of feature list for Ascend NPU support, @see #4567 --------- Co-authored-by: Olatunji Ruwase --- csrc/adagrad/cpu_adagrad.cpp | 19 +++++++-- csrc/adam/cpu_adam_impl.cpp | 27 +++++++++---- csrc/includes/cpu_adagrad.h | 37 ++++++++++++++++- csrc/includes/cpu_adam.h | 38 +++++++++++++++++- csrc/includes/cpu_lion.h | 37 ++++++++++++++++- csrc/lion/cpu_lion_impl.cpp | 19 +++++++-- deepspeed/env_report.py | 31 ++++++++++++++ op_builder/npu/__init__.py | 4 +- op_builder/npu/builder.py | 76 +++++++++++++++++++++++++++++------ op_builder/npu/cpu_adagrad.py | 25 ++++++++++++ op_builder/npu/cpu_adam.py | 25 ++++++++++++ op_builder/npu/cpu_lion.py | 25 ++++++++++++ op_builder/npu/no_impl.py | 9 +++++ 13 files changed, 343 insertions(+), 29 deletions(-) create mode 100644 op_builder/npu/cpu_adagrad.py create mode 100644 op_builder/npu/cpu_adam.py create mode 100644 op_builder/npu/cpu_lion.py diff --git a/csrc/adagrad/cpu_adagrad.cpp b/csrc/adagrad/cpu_adagrad.cpp index 8eebe00349be..563255176500 100644 --- a/csrc/adagrad/cpu_adagrad.cpp +++ b/csrc/adagrad/cpu_adagrad.cpp @@ -47,6 +47,8 @@ void Adagrad_Optimizer::Step_1(float* _params, size_t offset = copy_size + t; #if defined(__ENABLE_CUDA__) if ((t / TILE) >= 2) { cudaStreamSynchronize(_streams[_buf_index]); } +#elif defined(__ENABLE_CANN__) + if ((t / TILE) >= 2) { aclrtSynchronizeStream(_streams[_buf_index].stream()); } #endif #pragma omp parallel for for (size_t k = t; k < offset; k++) { @@ -62,7 +64,7 @@ void Adagrad_Optimizer::Step_1(float* _params, grad += _eps; grad = momentum / grad; param = grad * step_size + param; -#if defined(__ENABLE_CUDA__) +#if defined(__ENABLE_CUDA__) or defined(__ENABLE_CANN__) if (dev_params) _doubled_buffer[_buf_index][k - t] = param; #endif if (half_precision) @@ -79,6 +81,17 @@ void Adagrad_Optimizer::Step_1(float* _params, _doubled_buffer[_buf_index], dev_params + t, (copy_size), _streams[_buf_index]); _buf_index = !_buf_index; } +#elif defined(__ENABLE_CANN__) + if (dev_params) { + size_t memcpy_size = copy_size * sizeof(_doubled_buffer[_buf_index][0]); + aclrtMemcpy(dev_params + t, + memcpy_size, + _doubled_buffer[_buf_index], + memcpy_size, + aclrtMemcpyKind::ACL_MEMCPY_HOST_TO_DEVICE); + + _buf_index = !_buf_index; + } #endif } } @@ -180,7 +193,7 @@ int ds_adagrad_step(int optimizer_id, opt->update_state(lr, epsilon, weight_decay); opt->Step_8(params_ptr, grads_ptr, exp_avg_sq_ptr, params_c.numel()); -#if defined(__ENABLE_CUDA__) +#if defined(__ENABLE_CUDA__) or defined(__ENABLE_CANN__) opt->SynchronizeStreams(); #endif return 0; @@ -196,7 +209,7 @@ int ds_adagrad_step_plus_copy(int optimizer_id, torch::Tensor& exp_avg_sq, torch::Tensor& gpu_params) { -#if defined(__ENABLE_CUDA__) +#if defined(__ENABLE_CUDA__) or defined(__ENABLE_CANN__) auto params_c = params.contiguous(); auto gpu_params_c = gpu_params.contiguous(); auto exp_avg_sq_c = exp_avg_sq.contiguous(); diff --git a/csrc/adam/cpu_adam_impl.cpp b/csrc/adam/cpu_adam_impl.cpp index 742cb4292777..9a4a8d956519 100644 --- a/csrc/adam/cpu_adam_impl.cpp +++ b/csrc/adam/cpu_adam_impl.cpp @@ -61,6 +61,8 @@ void Adam_Optimizer::Step_1(float* _params, size_t offset = copy_size + t; #if defined(__ENABLE_CUDA__) if ((t / TILE) >= 2) { cudaStreamSynchronize(_streams[_buf_index]); } +#elif defined(__ENABLE_CANN__) + if ((t / TILE) >= 2) { aclrtSynchronizeStream(_streams[_buf_index].stream()); } #endif #pragma omp parallel for for (size_t k = t; k < offset; k++) { @@ -81,7 +83,7 @@ void Adam_Optimizer::Step_1(float* _params, grad = momentum / grad; if (_weight_decay > 0 && _adamw_mode) { param += w_decay * param; } param = grad * step_size + param; -#if defined(__ENABLE_CUDA__) +#if defined(__ENABLE_CUDA__) or defined(__ENABLE_CANN__) if (dev_params) _doubled_buffer[_buf_index][k - t] = param; #endif if (half_precision) @@ -96,6 +98,17 @@ void Adam_Optimizer::Step_1(float* _params, launch_param_update( _doubled_buffer[_buf_index], dev_params + t, (copy_size), _streams[_buf_index]); + _buf_index = !_buf_index; + } +#elif defined(__ENABLE_CANN__) + if (dev_params) { + size_t memcpy_size = copy_size * sizeof(_doubled_buffer[_buf_index][0]); + aclrtMemcpy(dev_params + t, + memcpy_size, + _doubled_buffer[_buf_index], + memcpy_size, + aclrtMemcpyKind::ACL_MEMCPY_HOST_TO_DEVICE); + _buf_index = !_buf_index; } #endif @@ -239,7 +252,7 @@ int ds_adam_step(int optimizer_id, nullptr, (params.options().dtype() == at::kHalf)); -#if defined(__ENABLE_CUDA__) +#if defined(__ENABLE_CUDA__) or defined(__ENABLE_CANN__) opt->SynchronizeStreams(); #endif return 0; @@ -257,18 +270,18 @@ int ds_adam_step_plus_copy(int optimizer_id, torch::Tensor& grads, torch::Tensor& exp_avg, torch::Tensor& exp_avg_sq, - torch::Tensor& gpu_params) + torch::Tensor& device_params) { -#if defined(__ENABLE_CUDA__) +#if defined(__ENABLE_CUDA__) or defined(__ENABLE_CANN__) auto params_c = params.contiguous(); - auto gpu_params_c = gpu_params.contiguous(); + auto device_params_c = device_params.contiguous(); auto exp_avg_c = exp_avg.contiguous(); auto exp_avg_sq_c = exp_avg_sq.contiguous(); auto grads_c = grads.contiguous(); float* params_ptr = (float*)params_c.data_ptr(); float* grads_ptr = (float*)grads_c.data_ptr(); - ds_half_precision_t* gpu_params_ptr = (ds_half_precision_t*)gpu_params_c.data_ptr(); + ds_half_precision_t* device_params_ptr = (ds_half_precision_t*)device_params_c.data_ptr(); float* exp_avg_ptr = (float*)exp_avg_c.data_ptr(); float* exp_avg_sq_ptr = (float*)exp_avg_sq_c.data_ptr(); @@ -281,7 +294,7 @@ int ds_adam_step_plus_copy(int optimizer_id, exp_avg_ptr, exp_avg_sq_ptr, params_c.numel(), - gpu_params_ptr, + device_params_ptr, (params.options().dtype() == at::kHalf)); opt->SynchronizeStreams(); diff --git a/csrc/includes/cpu_adagrad.h b/csrc/includes/cpu_adagrad.h index ba40fcf7b62a..59888adf17c3 100644 --- a/csrc/includes/cpu_adagrad.h +++ b/csrc/includes/cpu_adagrad.h @@ -18,6 +18,10 @@ #include "cuda.h" #include "custom_cuda_layers.h" typedef __half ds_half_precision_t; +#elif defined(__ENABLE_CANN__) +#include "acl/acl.h" +#include "torch_npu/csrc/core/npu/NPUStream.h" +typedef c10::Half ds_half_precision_t; #else typedef unsigned short ds_half_precision_t; #endif @@ -41,6 +45,11 @@ class Adagrad_Optimizer { _streams[0] = TrainingContext::Instance().GetCurrentStream(); _streams[1] = TrainingContext::Instance().GetNewStream(); + _buf_index = false; +#elif defined(__ENABLE_CANN__) + aclrtMallocHost((void**)_doubled_buffer, TILE * sizeof(float)); + aclrtMallocHost((void**)(_doubled_buffer + 1), TILE * sizeof(float)); + _buf_index = false; #endif } @@ -49,6 +58,9 @@ class Adagrad_Optimizer { #if defined(__ENABLE_CUDA__) cudaFreeHost(_doubled_buffer[0]); cudaFreeHost(_doubled_buffer[1]); +#elif defined(__ENABLE_CANN__) + aclrtFreeHost(_doubled_buffer[0]); + aclrtFreeHost(_doubled_buffer[1]); #endif } #if defined(__AVX512__) or defined(__AVX256__) @@ -69,6 +81,11 @@ class Adagrad_Optimizer { { for (int i = 0; i < 2; i++) cudaStreamSynchronize(_streams[i]); } +#elif defined(__ENABLE_CANN__) + inline void SynchronizeStreams() + { + for (int i = 0; i < 2; i++) aclrtSynchronizeStream(_streams[i].stream()); + } #endif inline void IncrementStep(size_t step) { @@ -95,6 +112,11 @@ class Adagrad_Optimizer { bool _buf_index; float* _doubled_buffer[2]; cudaStream_t _streams[2]; +#elif defined(__ENABLE_CANN__) + float* _doubled_buffer[2]; + c10_npu::NPUStream _streams[2] = {c10_npu::getCurrentNPUStream(), + c10_npu::getNPUStreamFromPool()}; + bool _buf_index; #endif }; @@ -125,6 +147,8 @@ void Adagrad_Optimizer::Step_AVX(size_t* rounded_size, size_t offset = copy_size + t; #if defined(__ENABLE_CUDA__) if ((t / TILE) >= 2) { cudaStreamSynchronize(_streams[_buf_index]); } +#elif defined(__ENABLE_CANN__) + if ((t / TILE) >= 2) { aclrtSynchronizeStream(_streams[_buf_index].stream()); } #endif #pragma omp parallel for for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { @@ -149,7 +173,7 @@ void Adagrad_Optimizer::Step_AVX(size_t* rounded_size, simd_fma(param_4, grad_4, step_size_4, param_4); simd_store(_params + i, param_4, half_precision); -#if defined(__ENABLE_CUDA__) +#if defined(__ENABLE_CUDA__) or defined(__ENABLE_CANN__) if (dev_params) { simd_store(_doubled_buffer[_buf_index] + (i - t), param_4, half_precision); } @@ -167,6 +191,17 @@ void Adagrad_Optimizer::Step_AVX(size_t* rounded_size, _buf_index = !_buf_index; } +#elif defined(__ENABLE_CANN__) + if (dev_params) { + size_t memcpy_size = copy_size * sizeof(_doubled_buffer[_buf_index][0]); + if (half_precision) memoryCopySize /= 2; + aclrtMemcpy(dev_params + t, + memcpy_size, + _doubled_buffer[_buf_index], + memcpy_size, + aclrtMemcpyKind::ACL_MEMCPY_HOST_TO_DEVICE); + + _buf_index = !_buf_index; #endif } *rounded_size = new_rounded_size; diff --git a/csrc/includes/cpu_adam.h b/csrc/includes/cpu_adam.h index c4f7edcd7410..44d3ed3cac61 100644 --- a/csrc/includes/cpu_adam.h +++ b/csrc/includes/cpu_adam.h @@ -19,6 +19,10 @@ #include "cuda.h" #include "custom_cuda_layers.h" typedef __half ds_half_precision_t; +#elif defined(__ENABLE_CANN__) +#include "acl/acl.h" +#include "torch_npu/csrc/core/npu/NPUStream.h" +typedef c10::Half ds_half_precision_t; #else #include typedef unsigned short ds_half_precision_t; @@ -57,6 +61,11 @@ class Adam_Optimizer { _streams[0] = TrainingContext::Instance().GetCurrentStream(); _streams[1] = TrainingContext::Instance().GetNewStream(); + _buf_index = false; +#elif defined(__ENABLE_CANN__) + aclrtMallocHost((void**)_doubled_buffer, TILE * sizeof(float)); + aclrtMallocHost((void**)(_doubled_buffer + 1), TILE * sizeof(float)); + _buf_index = false; #endif } @@ -65,6 +74,9 @@ class Adam_Optimizer { #if defined(__ENABLE_CUDA__) cudaFreeHost(_doubled_buffer[0]); cudaFreeHost(_doubled_buffer[1]); +#elif defined(__ENABLE_CANN__) + aclrtFreeHost(_doubled_buffer[0]); + aclrtFreeHost(_doubled_buffer[1]); #endif } @@ -87,6 +99,11 @@ class Adam_Optimizer { { for (int i = 0; i < 2; i++) cudaStreamSynchronize(_streams[i]); } +#elif defined(__ENABLE_CANN__) + inline void SynchronizeStreams() + { + for (int i = 0; i < 2; i++) aclrtSynchronizeStream(_streams[i].stream()); + } #endif inline void IncrementStep(size_t step, float beta1, float beta2) { @@ -142,6 +159,11 @@ class Adam_Optimizer { float* _doubled_buffer[2]; cudaStream_t _streams[2]; bool _buf_index; +#elif defined(__ENABLE_CANN__) + float* _doubled_buffer[2]; + c10_npu::NPUStream _streams[2] = {c10_npu::getCurrentNPUStream(), + c10_npu::getNPUStreamFromPool()}; + bool _buf_index; #endif }; @@ -192,6 +214,9 @@ void Adam_Optimizer::Step_AVX(size_t* rounded_size, size_t offset = copy_size + t; #if defined(__ENABLE_CUDA__) if ((t / TILE) >= 2) { cudaStreamSynchronize(_streams[_buf_index]); } +#elif defined(__ENABLE_CANN__) + if ((t / TILE) >= 2) { aclrtSynchronizeStream((_streams[_buf_index].stream()); + } #endif #pragma omp parallel for for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { @@ -227,7 +252,7 @@ void Adam_Optimizer::Step_AVX(size_t* rounded_size, simd_fma(param_4, grad_4, step_size_4, param_4); simd_store(_params + (i >> rshft), param_4, half_precision); -#if defined(__ENABLE_CUDA__) +#if defined(__ENABLE_CUDA__) or defined(__ENABLE_CANN__) if (dev_params) { simd_store(_doubled_buffer[_buf_index] + (i - t), param_4, half_precision); } @@ -246,6 +271,17 @@ void Adam_Optimizer::Step_AVX(size_t* rounded_size, _buf_index = !_buf_index; } +#elif defined(__ENABLE_CANN__) + if (dev_params) { + size_t memcpy_size = copy_size * sizeof(_doubled_buffer[_buf_index][0]); + if (half_precision) memoryCopySize /= 2; + aclrtMemcpy(dev_params + t, + memcpy_size, + _doubled_buffer[_buf_index], + memcpy_size, + aclrtMemcpyKind::ACL_MEMCPY_HOST_TO_DEVICE); + + _buf_index = !_buf_index; #endif } *rounded_size = new_rounded_size; diff --git a/csrc/includes/cpu_lion.h b/csrc/includes/cpu_lion.h index 76034ceb3459..d83fe9473332 100644 --- a/csrc/includes/cpu_lion.h +++ b/csrc/includes/cpu_lion.h @@ -19,6 +19,10 @@ #include "cuda.h" #include "custom_cuda_layers.h" typedef __half ds_half_precision_t; +#elif defined(__ENABLE_CANN__) +#include "acl/acl.h" +#include "torch_npu/csrc/core/npu/NPUStream.h" +typedef c10::Half ds_half_precision_t; #else #include typedef unsigned short ds_half_precision_t; @@ -46,6 +50,11 @@ class Lion_Optimizer { _streams[0] = TrainingContext::Instance().GetCurrentStream(); _streams[1] = TrainingContext::Instance().GetNewStream(); + _buf_index = false; +#elif defined(__ENABLE_CANN__) + aclrtMallocHost((void**)_doubled_buffer, TILE * sizeof(float)); + aclrtMallocHost((void**)(_doubled_buffer + 1), TILE * sizeof(float)); + _buf_index = false; #endif } @@ -54,6 +63,9 @@ class Lion_Optimizer { #if defined(__ENABLE_CUDA__) cudaFreeHost(_doubled_buffer[0]); cudaFreeHost(_doubled_buffer[1]); +#elif defined(__ENABLE_CANN__) + aclrtFreeHost(_doubled_buffer[0]); + aclrtFreeHost(_doubled_buffer[1]); #endif } @@ -75,6 +87,11 @@ class Lion_Optimizer { { for (int i = 0; i < 2; i++) cudaStreamSynchronize(_streams[i]); } +#elif defined(__ENABLE_CANN__) + inline void SynchronizeStreams() + { + for (int i = 0; i < 2; i++) aclrtSynchronizeStream(_streams[i].stream()); + } #endif inline void IncrementStep(size_t step, float beta1, float beta2) { @@ -102,6 +119,11 @@ class Lion_Optimizer { float* _doubled_buffer[2]; cudaStream_t _streams[2]; bool _buf_index; +#elif defined(__ENABLE_CANN__) + float* _doubled_buffer[2]; + c10_npu::NPUStream _streams[2] = {c10_npu::getCurrentNPUStream(), + c10_npu::getNPUStreamFromPool()}; + bool _buf_index; #endif }; @@ -149,6 +171,8 @@ void Lion_Optimizer::Step_AVX(size_t* rounded_size, size_t offset = copy_size + t; #if defined(__ENABLE_CUDA__) if ((t / TILE) >= 2) { cudaStreamSynchronize(_streams[_buf_index]); } +#elif defined(__ENABLE_CANN__) + if ((t / TILE) >= 2) { aclrtSynchronizeStream(_streams[_buf_index].stream()); } #endif #pragma omp parallel for for (size_t i = t; i < offset; i += SIMD_WIDTH * span) { @@ -178,7 +202,7 @@ void Lion_Optimizer::Step_AVX(size_t* rounded_size, simd_fma(momentum_4, grad_4, betta2_minus1_4, momentum_4); simd_store(_params + (i >> rshft), param_4, half_precision); -#if defined(__ENABLE_CUDA__) +#if defined(__ENABLE_CUDA__) or defined(__ENABLE_CANN__) if (dev_params) { simd_store(_doubled_buffer[_buf_index] + (i - t), param_4, half_precision); } @@ -196,6 +220,17 @@ void Lion_Optimizer::Step_AVX(size_t* rounded_size, _buf_index = !_buf_index; } +#elif defined(__ENABLE_CANN__) + if (dev_params) { + size_t memcpy_size = copy_size * sizeof(_doubled_buffer[_buf_index][0]); + if (half_precision) memoryCopySize /= 2; + aclrtMemcpy(dev_params + t, + memcpy_size, + _doubled_buffer[_buf_index], + memcpy_size, + aclrtMemcpyKind::ACL_MEMCPY_HOST_TO_DEVICE); + + _buf_index = !_buf_index; #endif } *rounded_size = new_rounded_size; diff --git a/csrc/lion/cpu_lion_impl.cpp b/csrc/lion/cpu_lion_impl.cpp index 5c24e23b4b20..28314cf5b6e1 100644 --- a/csrc/lion/cpu_lion_impl.cpp +++ b/csrc/lion/cpu_lion_impl.cpp @@ -54,6 +54,8 @@ void Lion_Optimizer::Step_1(float* _params, size_t offset = copy_size + t; #if defined(__ENABLE_CUDA__) if ((t / TILE) >= 2) { cudaStreamSynchronize(_streams[_buf_index]); } +#elif defined(__ENABLE_CANN__) + if ((t / TILE) >= 2) { aclrtSynchronizeStream(_streams[_buf_index].stream()); } #endif #pragma omp parallel for for (size_t k = t; k < offset; k++) { @@ -72,7 +74,7 @@ void Lion_Optimizer::Step_1(float* _params, } momentum = momentum * _betta2; momentum = grad * betta2_minus1 + momentum; -#if defined(__ENABLE_CUDA__) +#if defined(__ENABLE_CUDA__) or defined(__ENABLE_CANN__) if (dev_params) _doubled_buffer[_buf_index][k - t] = param; #endif if (half_precision) @@ -86,6 +88,17 @@ void Lion_Optimizer::Step_1(float* _params, launch_param_update( _doubled_buffer[_buf_index], dev_params + t, (copy_size), _streams[_buf_index]); + _buf_index = !_buf_index; + } +#elif defined(__ENABLE_CANN__) + if (dev_params) { + size_t memcpy_size = copy_size * sizeof(_doubled_buffer[_buf_index][0]); + aclrtMemcpy(dev_params + t, + memcpy_size, + _doubled_buffer[_buf_index], + memcpy_size, + aclrtMemcpyKind::ACL_MEMCPY_HOST_TO_DEVICE); + _buf_index = !_buf_index; } #endif @@ -201,7 +214,7 @@ int ds_lion_step(int optimizer_id, nullptr, (params.options().dtype() == at::kHalf)); -#if defined(__ENABLE_CUDA__) +#if defined(__ENABLE_CUDA__) or defined(__ENABLE_CANN__) opt->SynchronizeStreams(); #endif return 0; @@ -218,7 +231,7 @@ int ds_lion_step_plus_copy(int optimizer_id, torch::Tensor& exp_avg, torch::Tensor& gpu_params) { -#if defined(__ENABLE_CUDA__) +#if defined(__ENABLE_CUDA__) or defined(__ENABLE_CANN__) auto params_c = params.contiguous(); auto gpu_params_c = gpu_params.contiguous(); auto exp_avg_c = exp_avg.contiguous(); diff --git a/deepspeed/env_report.py b/deepspeed/env_report.py index 2c3a9e701d4d..912262453748 100644 --- a/deepspeed/env_report.py +++ b/deepspeed/env_report.py @@ -80,6 +80,32 @@ def nvcc_version(): return ".".join(release) +def installed_cann_path(): + if "ASCEND_HOME_PATH" in os.environ or os.path.exists(os.environ["ASCEND_HOME_PATH"]): + return os.environ["ASCEND_HOME_PATH"] + return None + + +def installed_cann_version(): + import re + ascend_path = installed_cann_path() + if ascend_path is None: + return f"CANN_HOME does not exist, unable to compile NPU op(s)" + cann_version = "" + for dirpath, _, filenames in os.walk(os.path.realpath(ascend_path)): + if cann_version: + break + install_files = [file for file in filenames if re.match(r"ascend_.*_install\.info", file)] + if install_files: + filepath = os.path.join(dirpath, install_files[0]) + with open(filepath, "r") as f: + for line in f: + if line.find("version") != -1: + cann_version = line.strip().split("=")[-1] + break + return cann_version + + def get_shm_size(): try: shm_stats = os.statvfs('/dev/shm') @@ -122,6 +148,11 @@ def debug_report(): ("deepspeed wheel compiled w.", f"torch {torch_info['version']}, " + (f"hip {torch_info['hip_version']}" if hip_version else f"cuda {torch_info['cuda_version']}")) ]) + elif get_accelerator().device_name() == 'npu': + import torch_npu + report.extend([("deepspeed wheel compiled w.", f"torch {torch_info['version']}"), + ("torch_npu install path", torch_npu.__path__), ("torch_npu version", torch_npu.__version__), + ("ascend_cann version", installed_cann_version())]) else: report.extend([("deepspeed wheel compiled w.", f"torch {torch_info['version']} ")]) diff --git a/op_builder/npu/__init__.py b/op_builder/npu/__init__.py index 0d9e76b3903f..6ad9124d36f0 100644 --- a/op_builder/npu/__init__.py +++ b/op_builder/npu/__init__.py @@ -4,6 +4,8 @@ # DeepSpeed Team '''Copyright The Microsoft DeepSpeed Team''' -# NPU related operators will be added in the future. from .fused_adam import FusedAdamBuilder from .no_impl import NotImplementedBuilder +from .cpu_adam import CPUAdamBuilder +from .cpu_adagrad import CPUAdagradBuilder +from .cpu_lion import CPULionBuilder diff --git a/op_builder/npu/builder.py b/op_builder/npu/builder.py index 7773388737a2..0dea2e78915e 100644 --- a/op_builder/npu/builder.py +++ b/op_builder/npu/builder.py @@ -3,6 +3,13 @@ # DeepSpeed Team +import re +import os +try: + import torch_npu +except ImportError as e: + pass + try: # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed # if successful this also means we're doing a local install and not JIT compile path @@ -13,22 +20,67 @@ class NPUOpBuilder(OpBuilder): + _ascend_path = None + _torch_npu_path = None + _cann_version = None + + def __init__(self, name): + super().__init__(name) + self._ascend_path = self.installed_cann_path() + self._torch_npu_path = os.path.join(os.path.dirname(os.path.abspath(torch_npu.__file__))) + try: + self._cann_version = self.installed_cann_version(self.name) + except BaseException: + print(f"{self.name} ascend_cann is missing, npu ops cannot be compiled!") - def builder(self): - from torch.utils.cpp_extension import CppExtension as ExtensionBuilder + def cann_defs(self): + if self._cann_version: + return '-D__ENABLE_CANN__' + return '-D__DISABLE_CANN__' - compile_args = {'cxx': self.strip_empty_entries(self.cxx_args())} + def installed_cann_path(self): + if "ASCEND_HOME_PATH" in os.environ or os.path.exists(os.environ["ASCEND_HOME_PATH"]): + return os.environ["ASCEND_HOME_PATH"] + return None - cpp_ext = ExtensionBuilder(name=self.absolute_name(), - sources=self.strip_empty_entries(self.sources()), - include_dirs=self.strip_empty_entries(self.include_paths()), - libraries=self.strip_empty_entries(self.libraries_args()), - extra_compile_args=compile_args) + def installed_cann_version(self, name=""): + ascend_path = self.installed_cann_path() + assert ascend_path is not None, "CANN_HOME does not exist, unable to compile NPU op(s)" + cann_version = "" + for dirpath, _, filenames in os.walk(os.path.realpath(ascend_path)): + if cann_version: + break + install_files = [file for file in filenames if re.match(r"ascend_.*_install\.info", file)] + if install_files: + filepath = os.path.join(dirpath, install_files[0]) + with open(filepath, "r") as f: + for line in f: + if line.find("version") != -1: + cann_version = line.strip().split("=")[-1] + break + return cann_version - return cpp_ext + def include_paths(self): + paths = super().include_paths() + paths += [os.path.join(self._ascend_path, 'include'), os.path.join(self._torch_npu_path, 'include')] + return paths def cxx_args(self): - return [] + args = super().cxx_args() + args += ['-O3', '-std=c++17', '-g', '-Wno-reorder', '-fopenmp'] + args += ['-fstack-protector-all', '-Wl,-z,relro,-z,now,-z,noexecstack', '-Wl,--disable-new-dtags,--rpath'] + args += [ + self.cann_defs(), + self.cpu_arch(), + self.simd_width(), '-L' + os.path.join(self._ascend_path, 'lib64'), + '-L' + os.path.join(self._torch_npu_path, 'lib') + ] + return args - def libraries_args(self): - return [] + def extra_ldflags(self): + flags = super().extra_ldflags() + flags += [ + '-L' + os.path.join(self._ascend_path, 'lib64'), '-lascendcl', + '-L' + os.path.join(self._torch_npu_path, 'lib'), '-ltorch_npu' + ] + return flags diff --git a/op_builder/npu/cpu_adagrad.py b/op_builder/npu/cpu_adagrad.py new file mode 100644 index 000000000000..161bc82efe1c --- /dev/null +++ b/op_builder/npu/cpu_adagrad.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import NPUOpBuilder + + +class CPUAdagradBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAGRAD" + NAME = "cpu_adagrad" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adagrad.{self.NAME}_op' + + def sources(self): + return ['csrc/adagrad/cpu_adagrad.cpp'] + + def include_paths(self): + args = super().include_paths() + args += ['csrc/includes'] + return args diff --git a/op_builder/npu/cpu_adam.py b/op_builder/npu/cpu_adam.py new file mode 100644 index 000000000000..a4e9569c0f33 --- /dev/null +++ b/op_builder/npu/cpu_adam.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import NPUOpBuilder + + +class CPUAdamBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_ADAM" + NAME = "cpu_adam" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.adam.{self.NAME}_op' + + def sources(self): + return ['csrc/adam/cpu_adam.cpp', 'csrc/adam/cpu_adam_impl.cpp'] + + def include_paths(self): + args = super().include_paths() + args += ['csrc/includes'] + return args diff --git a/op_builder/npu/cpu_lion.py b/op_builder/npu/cpu_lion.py new file mode 100644 index 000000000000..6917e0fd03d0 --- /dev/null +++ b/op_builder/npu/cpu_lion.py @@ -0,0 +1,25 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from .builder import NPUOpBuilder + + +class CPULionBuilder(NPUOpBuilder): + BUILD_VAR = "DS_BUILD_CPU_LION" + NAME = "cpu_lion" + + def __init__(self): + super().__init__(name=self.NAME) + + def absolute_name(self): + return f'deepspeed.ops.lion.{self.NAME}_op' + + def sources(self): + return ['csrc/lion/cpu_lion.cpp', 'csrc/lion/cpu_lion_impl.cpp'] + + def include_paths(self): + args = super().include_paths() + args += ['csrc/includes'] + return args diff --git a/op_builder/npu/no_impl.py b/op_builder/npu/no_impl.py index f17973fda401..5b1771fabc22 100644 --- a/op_builder/npu/no_impl.py +++ b/op_builder/npu/no_impl.py @@ -22,3 +22,12 @@ def load(self, verbose=True): def sources(self): return [] + + def cxx_args(self): + return [] + + def extra_ldflags(self): + return [] + + def include_paths(self): + return []