From 9ca60160179aacc1efe87121e1ec259ce8d0a476 Mon Sep 17 00:00:00 2001
From: Logan Adams <114770087+loadams@users.noreply.github.com>
Date: Fri, 6 Dec 2024 14:29:00 -0800
Subject: [PATCH 1/3] Pin HPU tests (#6831)

HPU tests are impacted by the same issue as other tests that use
transformers latest. This PR pins to a version of transformers before
the fix.
---
 .github/workflows/hpu-gaudi2.yml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/.github/workflows/hpu-gaudi2.yml b/.github/workflows/hpu-gaudi2.yml
index 9f1a9d973ca2..a06f871b7c56 100644
--- a/.github/workflows/hpu-gaudi2.yml
+++ b/.github/workflows/hpu-gaudi2.yml
@@ -111,6 +111,8 @@ jobs:
         run: |
           git clone https://github.com/huggingface/transformers
           cd transformers
+          # if needed switch to the last known good SHA until transformers@master is fixed
+          git checkout 6c3f168b3
           git rev-parse --short HEAD
           pip install .
 

From 9a41ccaf446cc2026a30040bc55a21df1256cb64 Mon Sep 17 00:00:00 2001
From: Hoa La <87745938+lvhoaa@users.noreply.github.com>
Date: Mon, 9 Dec 2024 12:56:54 -0500
Subject: [PATCH 2/3] Flops profiler support einops.einsum (#6755)

- Added support for FlopsProfiler to include einops.einsum operation
- Added _patch_miscellaneous_operations() and
_reload_miscellaneous_operations() to include this operation and
potentially include other miscellaneous operations in the future
- Added _einops_einsum_flops_compute() that mimic already-existed
_einsum_flops_compute()

---------

Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
---
 .../profiling/flops_profiler/profiler.py      | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/deepspeed/profiling/flops_profiler/profiler.py b/deepspeed/profiling/flops_profiler/profiler.py
index f87f1beb7e4e..c6b607aad813 100644
--- a/deepspeed/profiling/flops_profiler/profiler.py
+++ b/deepspeed/profiling/flops_profiler/profiler.py
@@ -16,6 +16,7 @@
 from deepspeed.moe.layer import MoE
 from deepspeed.utils.timer import FORWARD_GLOBAL_TIMER, BACKWARD_GLOBAL_TIMER, STEP_GLOBAL_TIMER
 from deepspeed.utils.torch import required_torch_version
+import einops
 
 Tensor = torch.Tensor
 
@@ -82,6 +83,7 @@ def start_profile(self, ignore_list=None):
         self.reset_profile()
         _patch_functionals()
         _patch_tensor_methods()
+        _patch_miscellaneous_operations()
 
         def register_module_hooks(module, ignore_list):
             if ignore_list and type(module) in ignore_list:
@@ -137,6 +139,7 @@ def stop_profile(self):
         if self.started and self.func_patched:
             _reload_functionals()
             _reload_tensor_methods()
+            _reload_miscellaneous_operations()
             self.func_patched = False
 
         def remove_profile_attrs(module):
@@ -787,6 +790,29 @@ def _einsum_flops_compute(equation, *operands):
     raise NotImplementedError("Unsupported einsum operation.")
 
 
+def _einops_einsum_flops_compute(*args):
+    """
+    Count flops for the einops.einsum operation.
+    """
+    *operands, equation = args
+    input_shapes = [o.shape for o in operands]
+
+    # Re-map equation so that same equation with different alphabet
+    # representations will look the same.
+    letter_order = OrderedDict((k, 0) for k in equation if k.isalpha()).keys()
+    mapping = {ord(x): 97 + i for i, x in enumerate(letter_order)}
+    equation = equation.translate(mapping)
+
+    np_arrs = [np.zeros(s) for s in input_shapes]
+    optim = np.einsum_path(equation, *np_arrs, optimize="optimal")[1]
+    for line in optim.split("\n"):
+        if "optimized flop" in line.lower():
+            flop = int(float(line.split(":")[-1]))
+            return flop, 0
+
+    raise NotImplementedError("Unsupported einops.einsum operation.")
+
+
 def _tensor_addmm_flops_compute(self, mat1, mat2, *, beta=1, alpha=1, out=None):
     """
     Count flops for the tensor addmm operation.
@@ -937,6 +963,10 @@ def _patch_tensor_methods():
     torch.baddbmm = wrapFunc(torch.baddbmm, _tensor_addmm_flops_compute)
 
 
+def _patch_miscellaneous_operations():
+    einops.einsum = wrapFunc(einops.einsum, _einops_einsum_flops_compute)
+
+
 def _reload_functionals():
     # torch.nn.functional does not support importlib.reload()
     F.linear = old_functions[F.linear.__str__]
@@ -995,6 +1025,10 @@ def _reload_tensor_methods():
     torch.baddbmm = old_functions[torch.baddbmm.__str__]
 
 
+def _reload_miscellaneous_operations():
+    einops.einsum = old_functions[einops.einsum.__str__]
+
+
 def _rnn_flops(flops, rnn_module, w_ih, w_hh, input_size):
     gates_size = w_ih.shape[0]
     # matrix matrix mult ih state and internal state

From 08b907a226c64fbadace5dba71ed45bbe22bc0ef Mon Sep 17 00:00:00 2001
From: Logan Adams <114770087+loadams@users.noreply.github.com>
Date: Mon, 9 Dec 2024 12:24:33 -0800
Subject: [PATCH 3/3] Pin pytest-subtests version for accelerate tests (#6842)

---
 .github/workflows/nv-accelerate-v100.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/nv-accelerate-v100.yml b/.github/workflows/nv-accelerate-v100.yml
index 346055e2685f..56b1b1ab12d2 100644
--- a/.github/workflows/nv-accelerate-v100.yml
+++ b/.github/workflows/nv-accelerate-v100.yml
@@ -52,7 +52,7 @@ jobs:
           pip install .[testing]
           # force protobuf version due to issues
           pip install "protobuf<4.21.0"
-          # tmp fix: force newer datasets version
-          #pip install "datasets>=2.0.0"
+          # force pytest-subtests version due to issues
+          pip install pytest-subtests==0.13.1
           pip list
           pytest $PYTEST_OPTS --color=yes --durations=0 --verbose tests/deepspeed