From 77567b51f502fdebb53c47fbe37c2ca958c300b6 Mon Sep 17 00:00:00 2001
From: Du Li <duli1@microsoft.com>
Date: Fri, 19 Apr 2024 16:57:09 -0700
Subject: [PATCH] 1. change the from capability to ds_feature 2. used
 __compatible_ops__ 3. adding more op name to constants.py

---
 accelerator/abstract_accelerator.py | 13 ++++++------
 accelerator/constants.py            | 31 ++++++++++++++++++++++++-----
 accelerator/cuda_accelerator.py     |  7 +++----
 3 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/accelerator/abstract_accelerator.py b/accelerator/abstract_accelerator.py
index 39675fa3305d..f266474369ec 100644
--- a/accelerator/abstract_accelerator.py
+++ b/accelerator/abstract_accelerator.py
@@ -6,14 +6,15 @@
 import abc
 from abc import ABC
 from .constants import *
-
+from ..deepspeed.git_version_info import compatible_ops as __compatible_ops__
 
 class DeepSpeedAccelerator(ABC):
 
     def __init__(self):
         self._name = None
         self._communication_backend_name = None
-        self._capabilities: dict[str, bool] = {ZERO_1: False, ZERO_2: False, ZERO_3: False, SPARSE_ATTN: False}
+        self._ds_features: dict[str, bool] = {ZERO_1: False, ZERO_2: False, ZERO_3: False}
+        self._ds_features.update({op: compatibility for op, compatibility in __compatible_ops__})
 
     @abc.abstractmethod
     def is_synchronized_device(self):
@@ -290,8 +291,8 @@ def build_extension(self):
     def export_envs(self):
         ...
 
-    def get_capability(self, key):
-        return self._capabilities[key]
+    def get_ds_feature(self, key):
+        return self._ds_features[key]
 
-    def set_capability(self, key, value):
-        self._capabilities[key] = value
+    def set_ds_feature(self, key, value):
+        self._ds_features[key] = value
diff --git a/accelerator/constants.py b/accelerator/constants.py
index 5470248af7d5..3f404b708d73 100644
--- a/accelerator/constants.py
+++ b/accelerator/constants.py
@@ -3,8 +3,29 @@
 
 # DeepSpeed Team
 
-#A list of constants used in the DeepSpeed capabilities dictionary
-ZERO_1= "zero1"
-ZERO_2= "zero2"
-ZERO_3= "zero3"
-SPARSE_ATTN= "sparse_attn"
+#A list of constants used in the DeepSpeed feature dictionary
+
+OP_ASYNC_IO = "async_io"
+OP_CCL_COMM = "deepspeed_ccl_comm"
+OP_CPU_ADAGRAD = "cpu_adagrad"
+OP_CPU_ADAM = "cpu_adam"
+OP_CPU_LION = "cpu_lion"
+OP_EVOFORMER_ATTN = "evoformer_attn"
+OP_FP_QUANTIZER = "fp_quantizer"
+OP_FUSED_ADAM = "fused_adam"
+OP_FUSED_LAMB = "fused_lamb"
+OP_FUSED_LION = "fused_lion"
+OP_INFERENCE_CORE_OPS = "inference_core_ops"
+OP_CUTLASS_OPS = "cutlass_ops"
+OP_QUANTIZER = "quantizer"
+OP_RAGGED_DEVICE_OPS = "ragged_device_ops"
+OP_RAGGED_OPS = "ragged_ops"
+OP_RANDOM_LTD = "random_ltd"
+OP_SPARSE_ATTN = "sparse_attn"
+OP_SPATIAL_INFERENCE = "spatial_inference"
+OP_STOCHASTIC_TRANSFORMER = "stochastic_transformer"
+OP_TRANSFORMER = "transformer"
+OP_TRANSFORMER_INFERENCE = "transformer_inference"
+ZERO_1 = "zero1"
+ZERO_2 = "zero2"
+ZERO_3 = "zero3"
\ No newline at end of file
diff --git a/accelerator/cuda_accelerator.py b/accelerator/cuda_accelerator.py
index 2a40580e3bad..159e26c60cfc 100644
--- a/accelerator/cuda_accelerator.py
+++ b/accelerator/cuda_accelerator.py
@@ -30,10 +30,9 @@ def __init__(self):
         self._communication_backend_name = 'nccl'
         if pynvml is None:
             self._init_pynvml()
-        self.set_capability(ZERO_1, True)
-        self.set_capability(ZERO_2, True)
-        self.set_capability(ZERO_3, True)
-        self.set_capability(SPARSE_ATTN, True)
+        self.set_ds_feature(ZERO_1, True)
+        self.set_ds_feature(ZERO_2, True)
+        self.set_ds_feature(ZERO_3, True)
 
     def _init_pynvml(self):
         global pynvml