deepspeed-fork content for 1.16.1

Signed-off-by: SW publisher <sw_publisher@habana-labs.com>

Author: SW publisher

.github/workflows/nv-nightly.yml

@@ -38,10 +38,6 @@ jobs:
           git rev-parse --short HEAD
           pip install .
 
-      - name: Install datasets
-        run: |
-          pip install datasets
-
       - name: Install deepspeed
         run: |
           pip install .[dev,1bit,autotuning,inf]

.pre-commit-config.yaml

@@ -39,7 +39,7 @@ repos:
         name: check-torchdist
         entry: ./scripts/check-torchdist.py
         language: python
-        exclude: ^(deepspeed/comm/|docs/|benchmarks/|scripts/check-torchdist.py|deepspeed/moe/sharded_moe.py|deepspeed/runtime/comm/coalesced_collectives.py|deepspeed/elasticity/elastic_agent.py|deepspeed/launcher/launch.py|tests/unit/comm/test_dist.py)
+        exclude: ^(deepspeed/comm/|docs/|benchmarks/|scripts/check-torchdist.py|deepspeed/moe/sharded_moe.py|deepspeed/runtime/comm/coalesced_collectives.py|deepspeed/elasticity/elastic_agent.py|deepspeed/launcher/launch.py|tests/unit/comm/test_dist.py|deepspeed/runtime/zero/utils.py|deepspeed/tools/pg_sim/ut/base.py|deepspeed/tools/pg_sim/pg.py)
         # Specific deepspeed/ files are excluded for now until we wrap ProcessGroup in deepspeed.comm
 
 -   repo: local

accelerator/abstract_accelerator.py

@@ -280,6 +280,10 @@ def create_op_builder(self, class_name):
     def get_op_builder(self, class_name):
         ...
 
+    @abc.abstractmethod
+    def get_compile_backend(self):
+        ...
+
     @abc.abstractmethod
     def build_extension(self):
         ...

accelerator/cpu_accelerator.py

@@ -302,3 +302,6 @@ def build_extension(self):
 
     def export_envs(self):
         return []
+
+    def get_compile_backend(self):
+        return "inductor"

accelerator/cuda_accelerator.py

@@ -360,3 +360,6 @@ def build_extension(self):
 
     def export_envs(self):
         return ['NCCL']
+
+    def get_compile_backend(self):
+        return "inductor"

accelerator/hpu_accelerator.py

@@ -288,6 +288,17 @@ def get_op_builder(self, class_name):
         else:
             return self.class_dict['NotImplementedBuilder'] if 'NotImplementedBuilder' in self.class_dict else None
 
+    def get_compile_backend(self):
+        return "hpu_backend"
+
+    #shall be removed once moving to torch.compile
+    def wrap_in_hpu_graph(self, module):
+        if self.hpu.is_lazy():
+            module = self.hpu.wrap_in_hpu_graph(module)
+        else:
+            print("Warning: hpu graphs in eager mode is not supported, ignoring")
+        return module
+
     def build_extension(self):
         from torch.utils.cpp_extension import BuildExtension
         return BuildExtension

accelerator/mps_accelerator.py

@@ -258,3 +258,6 @@ def build_extension(self):
 
     def export_envs(self):
         return []
+
+    def get_compile_backend(self):
+        return "inductor"

accelerator/npu_accelerator.py

@@ -278,3 +278,6 @@ def build_extension(self):
 
     def export_envs(self):
         return ['ASCEND', 'HCCL', 'LD_LIBRARY', 'PATH']
+
+    def get_compile_backend(self):
+        return "inductor"

build.txt

@@ -0,0 +1 @@
++hpu.synapse.v1.16.1

csrc/adam/cpu_adam_impl.cpp

@@ -244,13 +244,17 @@ int ds_adam_step(int optimizer_id,
     opt->IncrementStep(step, beta1, beta2);
     opt->update_state(lr, epsilon, weight_decay, bias_correction);
 
+    bool bit16_precision = false;
+    if ((params.options().dtype() == at::kHalf) || (params.options().dtype() == at::kBFloat16))
+        bit16_precision = true;
+
     opt->Step_8(params_ptr,
                 grads_ptr,
                 exp_avg_ptr,
                 exp_avg_sq_ptr,
                 params_c.numel(),
                 nullptr,
-                (params.options().dtype() == at::kHalf));
+                bit16_precision);
 
 #if defined(__ENABLE_CUDA__) or defined(__ENABLE_CANN__)
     opt->SynchronizeStreams();