DeepSpeed content for 1.23.0

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

Author: SW publisher

.pre-commit-config.yaml

@@ -5,55 +5,6 @@
 # Learn more about CODEOWNERS syntax here:
 # https://docs.github.com/en/repositories/managing-your-repositorys-settings-and-features/customizing-your-repository/about-code-owners
 
+*          oelayan@habana.ai lbachar@habana.ai mkovalenko@habana.ai nsonnenschein@habana.ai snahir@habana.ai mmalekan@habana.ai
 
-# top-level repo folders
-/.github/ @loadams
-/azure/ @loadams
-/benchmarks/ @guanhuawang @tjruwase
-/bin/ @loadams
-/csrc/ @tjruwase
-/deepspeed/ @loadams @tjruwase
-/docker/ @loadams @guanhuawang
-/docs/ @loadams @tjruwase
-/examples/ @jomayeri @tohtana
-/op_builder/ @loadams @tjruwase @jomayeri
-/release/ @loadams @jomayeri
-/requirements/ @loadams
-/scripts/ @loadams @tjruwase
-/tests/ @tjruwase @loadams @tohtana
-
-# deepspeed
-/deepspeed/autotuning/ @loadams
-/deepspeed/checkpoint/ @tjruwase
-/deepspeed/comm/ @guanhuawang
-/deepspeed/compression/ @tjruwase
-/deepspeed/elasticity/ @tjruwase
-/deepspeed/launcher/ @loadams
-/deepspeed/module_inject/ @hwchen2017 @loadams
-/deepspeed/moe/ @tohtana
-/deepspeed/monitor/ @tjruwase
-/deepspeed/nebula/ @tjruwase
-/deepspeed/nvme/ @tjruwase @jomayeri
-/deepspeed/ops/ @tohtana
-/deepspeed/pipe/ @tohtana @loadams
-/deepspeed/profiling/ @loadams
-/deepspeed/sequence/ @tohtana
-/deepspeed/utils/ @tjruwase @tohtana
-
-# inference
-/deepspeed/inference/ @hwchen2017 @tohtana
-/deepspeed/model_implementations/@tohtana @loadams
-
-# training
-/deepspeed/runtime/ @tjruwase @tohtana
-/deepspeed/runtime/activation_checkpointing/ @tjruwase
-/deepspeed/runtime/checkpoint_engine/ @tjruwase
-/deepspeed/runtime/comm/ @guanhuawang
-/deepspeed/runtime/compression/ @tjruwase
-/deepspeed/runtime/data_pipeline/ @tjruwase
-/deepspeed/runtime/domino/ @guanhuawang @hwchen2017
-/deepspeed/runtime/fp16/ @tjruwase @tohtana
-/deepspeed/runtime/fp16/onebit/ @tjruwase
-/deepspeed/runtime/pipe/ @loadams @tohtana
-/deepspeed/runtime/swap_tensor/ @tjruwase @jomayeri
-/deepspeed/runtime/zero/ @tjruwase @tohtana
+CODEOWNERS oelayan@habana.ai oabramovich@habana.ai

CODEOWNERS

@@ -301,6 +301,14 @@ def get_op_builder(self, class_name):
         else:
             return self.class_dict['NotImplementedBuilder'] if 'NotImplementedBuilder' in self.class_dict else None
 
+    #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
@@ -309,6 +317,7 @@ def export_envs(self):
         return []
 
     def visible_devices_envs(self):
+        # TODO SW-195658: remove WA to not return HABANA_VISIBLE_MODULES once SW-195657 is resolved
         # Current way deepspeed set this env var is not applicable with all HPU instances
         # User has to follow instructions in:
         # https://docs.habana.ai/en/latest/PyTorch/Reference/PT_Multiple_Tenants_on_HPU/Multiple_Workloads_Single_Docker.html

accelerator/hpu_accelerator.py

@@ -67,7 +67,7 @@ def get_accelerator():
                     f"XPU_Accelerator requires intel_extension_for_pytorch, which is not installed on this system.")
         elif accelerator_name == "xpu.external":
             try:
-                import intel_extension_for_deepspeed  # noqa: F401 # type: ignore
+                from intel_extension_for_deepspeed import XPU_Accelerator  # noqa: F401 # type: ignore
             except ImportError as e:
                 raise ValueError(
                     f"XPU_Accelerator external requires intel_extension_for_deepspeed, which is not installed on this system."
@@ -224,6 +224,12 @@ def get_accelerator():
         ds_accelerator = CPU_Accelerator()
     elif accelerator_name == "xpu.external":
         # XPU_Accelerator is already imported in detection stage
+        try:
+            from intel_extension_for_deepspeed import XPU_Accelerator  # noqa: F811
+        except ImportError as e:
+            raise ValueError(
+                f"XPU_Accelerator external requires intel_extension_for_deepspeed, which is not installed on this system."
+            )
         ds_accelerator = XPU_Accelerator()
     elif accelerator_name == "xpu":
         from .xpu_accelerator import XPU_Accelerator
@@ -258,7 +264,7 @@ def get_accelerator():
 def set_accelerator(accel_obj):
     global ds_accelerator
     _validate_accelerator(accel_obj)
-    if accel_logger is not None:
+    if accel_logger is not None and accel_obj is not None:
         accel_logger.info(f"Setting ds_accelerator to {accel_obj._name} (model specified)")
     ds_accelerator = accel_obj
 

accelerator/real_accelerator.py

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

build.txt

@@ -24,6 +24,7 @@
 
 at::Tensor quantize(torch::Tensor& out,
                     torch::Tensor& val,
+                    torch::Tensor& scale,
                     int group_size,
                     int stochastic_rounding,
                     int q_bits,
@@ -59,6 +60,7 @@ at::Tensor quantize(torch::Tensor& out,
 
 void dequantize(torch::Tensor& val,
                 torch::Tensor& val_q,
+                torch::Tensor& scale,
                 int group_size,
                 int q_mantisa_bits,
                 int q_exponent_bits)

csrc/fp_quantizer/fp_quantize.cpp

@@ -81,7 +81,7 @@ def __init__(self, args, active_resources):
         if not os.path.exists(self.results_dir):
             try:
                 os.makedirs(self.results_dir, exist_ok=True)
-                logger.info(f"Created autotuning results directory: {self.exps_dir}")
+                logger.info(f"Created autotuning results directory: {self.results_dir}")
             except:
                 logger.error(
                     f"Failed to create {self.results_dir}, please check results_dir in the autotuning config file is accessible by all the nodes in the job."

deepspeed/autotuning/autotuner.py

@@ -144,7 +144,7 @@
     "zero_optimization": {
         "stage": 3
     },
-    "memory_break_down": False
+    "memory_breakdown": False
 }
 
 DEFAULT_TUNING_SPACE_ZERO_0 = {"zero_optimization": {"stage": 0}}

deepspeed/autotuning/constants.py

@@ -77,27 +77,12 @@ def run_collective(self, name, **kwargs):
             return CCLHandler(self.ccl_comm_op)
 
     def all_reduce(self, tensor, op=ReduceOp.SUM, group=None, async_op=False):
-        use_caching = False
-        if use_caching:
-            match_id = f"{tensor.size()}-{op}"
-            name = "all_reduce_caching"
-            if name in self.available_coll:
-                group = self.get_all_ranks_from_group(group)
-                return self.ccl_comm_op.all_reduce_caching(tensor, op, match_id, group, async_op)
-            else:
-                return self.run_collective(name=name,
-                                           tensor=tensor,
-                                           op=op,
-                                           match_id=match_id,
-                                           group=group,
-                                           async_op=async_op)
+        name = "all_reduce"
+        if name in self.available_coll:
+            group = self.get_all_ranks_from_group(group)
+            return self.ccl_comm_op.all_reduce(tensor, op, group, async_op)
         else:
-            name = "all_reduce"
-            if name in self.available_coll:
-                group = self.get_all_ranks_from_group(group)
-                return self.ccl_comm_op.all_reduce(tensor, op, group, async_op)
-            else:
-                return self.run_collective(name=name, tensor=tensor, op=op, group=group, async_op=async_op)
+            return self.run_collective(name=name, tensor=tensor, op=op, group=group, async_op=async_op)
 
     def inference_all_reduce(self, tensor, op=ReduceOp.SUM, group=None):
         name = "inference_all_reduce"

deepspeed/comm/ccl.py

@@ -4,16 +4,18 @@
 # DeepSpeed Team
 
 import torch
+from deepspeed.utils.torch import required_torch_version
 
 try:
-    import torch.utils._pytree as pytree
-    from torch._functorch.aot_autograd import create_aot_dispatcher_function
-    from torch._inductor.lowering import register_lowering, fallbacks, add_needs_realized_inputs
-    from torch._inductor.ir import TensorBox, FallbackKernel, Layout, IRNode
-    from torch._inductor.virtualized import V
-    from torch._inductor.scheduler import Scheduler
-
-    original_create_aot_dispatcher_function = create_aot_dispatcher_function
+    if required_torch_version(min_version=2.6):
+        import torch.utils._pytree as pytree
+        from torch._functorch.aot_autograd import create_aot_dispatcher_function
+        from torch._inductor.lowering import register_lowering, fallbacks, add_needs_realized_inputs
+        from torch._inductor.ir import TensorBox, FallbackKernel, Layout, IRNode
+        from torch._inductor.virtualized import V
+        from torch._inductor.scheduler import Scheduler
+
+        original_create_aot_dispatcher_function = create_aot_dispatcher_function
 except ImportError:
     pass