From 5a905be8f9a722e6604ddd979484e2da0fca7abf Mon Sep 17 00:00:00 2001 From: Hoa La Date: Sun, 17 Nov 2024 11:35:02 -0500 Subject: [PATCH 1/5] Fix: add support for einops.einsum in FlopsProfiler --- .../profiling/flops_profiler/profiler.py | 30 +++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/deepspeed/profiling/flops_profiler/profiler.py b/deepspeed/profiling/flops_profiler/profiler.py index f87f1beb7e4e..afdd81ba0d5a 100644 --- a/deepspeed/profiling/flops_profiler/profiler.py +++ b/deepspeed/profiling/flops_profiler/profiler.py @@ -16,6 +16,8 @@ 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 @@ -786,6 +788,28 @@ def _einsum_flops_compute(equation, *operands): return flop, 0 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): """ @@ -932,7 +956,8 @@ def _patch_tensor_methods(): torch.add = wrapFunc(torch.add, _add_flops_compute) torch.Tensor.add = wrapFunc(torch.Tensor.add, _add_flops_compute) - torch.einsum = wrapFunc(torch.einsum, _einsum_flops_compute) + torch.einsum = wrapFunc(torch.einsum, _einsum_flops_compute) + einops.einsum = wrapFunc(einops.einsum, _einops_einsum_flops_compute) torch.baddbmm = wrapFunc(torch.baddbmm, _tensor_addmm_flops_compute) @@ -991,6 +1016,7 @@ def _reload_tensor_methods(): torch.Tensor.add = old_functions[torch.Tensor.add.__str__] torch.einsum = old_functions[torch.einsum.__str__] + einops.einsum = old_functions[einops.einsum.__str__] torch.baddbmm = old_functions[torch.baddbmm.__str__] @@ -1256,4 +1282,4 @@ def get_model_profile(model, if as_string: return number_to_string(flops), macs_to_string(macs), params_to_string(params) - return flops, macs, params + return flops, macs, params \ No newline at end of file From 3fdec4dd28e728e1f765fd214d344aa1bb0885fc Mon Sep 17 00:00:00 2001 From: Hoa La Date: Sun, 17 Nov 2024 11:52:00 -0500 Subject: [PATCH 2/5] Fix: add support for einops.einsum in FlopsProfiler --- deepspeed/profiling/flops_profiler/profiler.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/deepspeed/profiling/flops_profiler/profiler.py b/deepspeed/profiling/flops_profiler/profiler.py index afdd81ba0d5a..113cee5042fd 100644 --- a/deepspeed/profiling/flops_profiler/profiler.py +++ b/deepspeed/profiling/flops_profiler/profiler.py @@ -84,6 +84,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: @@ -139,6 +140,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): @@ -957,11 +959,14 @@ def _patch_tensor_methods(): torch.Tensor.add = wrapFunc(torch.Tensor.add, _add_flops_compute) torch.einsum = wrapFunc(torch.einsum, _einsum_flops_compute) - einops.einsum = wrapFunc(einops.einsum, _einops_einsum_flops_compute) 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__] @@ -1016,9 +1021,12 @@ def _reload_tensor_methods(): torch.Tensor.add = old_functions[torch.Tensor.add.__str__] torch.einsum = old_functions[torch.einsum.__str__] - einops.einsum = old_functions[einops.einsum.__str__] 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): From 6c51b3ca201eaa877002cc7a036795df697fcb4a Mon Sep 17 00:00:00 2001 From: Hoa La Date: Mon, 18 Nov 2024 12:21:27 -0500 Subject: [PATCH 3/5] format precommit and edit requirements --- bin/deepspeed | 2 +- bin/deepspeed.pt | 2 +- bin/dsr | 2 +- deepspeed/accelerator | 2 +- deepspeed/ops/csrc | 2 +- deepspeed/ops/op_builder | 2 +- deepspeed/profiling/flops_profiler/profiler.py | 12 ++++++------ requirements/requirements.txt | 1 + 8 files changed, 13 insertions(+), 12 deletions(-) diff --git a/bin/deepspeed b/bin/deepspeed index 6b7685641019..7f8898257a27 120000 --- a/bin/deepspeed +++ b/bin/deepspeed @@ -1 +1 @@ -ds \ No newline at end of file +ds diff --git a/bin/deepspeed.pt b/bin/deepspeed.pt index 6b7685641019..7f8898257a27 120000 --- a/bin/deepspeed.pt +++ b/bin/deepspeed.pt @@ -1 +1 @@ -ds \ No newline at end of file +ds diff --git a/bin/dsr b/bin/dsr index 747bf4722c42..a3a3b689ba18 120000 --- a/bin/dsr +++ b/bin/dsr @@ -1 +1 @@ -ds_report \ No newline at end of file +ds_report diff --git a/deepspeed/accelerator b/deepspeed/accelerator index b61ffface2d0..15fddb37dd6d 120000 --- a/deepspeed/accelerator +++ b/deepspeed/accelerator @@ -1 +1 @@ -../accelerator/ \ No newline at end of file +../accelerator/ diff --git a/deepspeed/ops/csrc b/deepspeed/ops/csrc index 338652302505..fd4e1dc24e66 120000 --- a/deepspeed/ops/csrc +++ b/deepspeed/ops/csrc @@ -1 +1 @@ -../../csrc \ No newline at end of file +../../csrc diff --git a/deepspeed/ops/op_builder b/deepspeed/ops/op_builder index db4f9c335065..05f303450d1c 120000 --- a/deepspeed/ops/op_builder +++ b/deepspeed/ops/op_builder @@ -1 +1 @@ -../../op_builder \ No newline at end of file +../../op_builder diff --git a/deepspeed/profiling/flops_profiler/profiler.py b/deepspeed/profiling/flops_profiler/profiler.py index 113cee5042fd..c6b607aad813 100644 --- a/deepspeed/profiling/flops_profiler/profiler.py +++ b/deepspeed/profiling/flops_profiler/profiler.py @@ -18,7 +18,6 @@ from deepspeed.utils.torch import required_torch_version import einops - Tensor = torch.Tensor module_flop_count = [] @@ -790,6 +789,7 @@ def _einsum_flops_compute(equation, *operands): return flop, 0 raise NotImplementedError("Unsupported einsum operation.") + def _einops_einsum_flops_compute(*args): """ Count flops for the einops.einsum operation. @@ -958,14 +958,14 @@ def _patch_tensor_methods(): torch.add = wrapFunc(torch.add, _add_flops_compute) torch.Tensor.add = wrapFunc(torch.Tensor.add, _add_flops_compute) - torch.einsum = wrapFunc(torch.einsum, _einsum_flops_compute) + torch.einsum = wrapFunc(torch.einsum, _einsum_flops_compute) 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() @@ -1023,8 +1023,8 @@ def _reload_tensor_methods(): torch.einsum = old_functions[torch.einsum.__str__] torch.baddbmm = old_functions[torch.baddbmm.__str__] - - + + def _reload_miscellaneous_operations(): einops.einsum = old_functions[einops.einsum.__str__] @@ -1290,4 +1290,4 @@ def get_model_profile(model, if as_string: return number_to_string(flops), macs_to_string(macs), params_to_string(params) - return flops, macs, params \ No newline at end of file + return flops, macs, params diff --git a/requirements/requirements.txt b/requirements/requirements.txt index 296398f680cc..1af4c69c5807 100755 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -1,3 +1,4 @@ +einops hjson msgpack ninja From a196890653b8c8d519a82ae32f46ea6b4b67b1f6 Mon Sep 17 00:00:00 2001 From: Hoa La Date: Tue, 19 Nov 2024 15:29:06 -0500 Subject: [PATCH 4/5] add einops to requirements-dev --- requirements/requirements-dev.txt | 1 + requirements/requirements.txt | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 2e2c880fbeb0..c6447d1198ae 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -3,6 +3,7 @@ clang-format==18.1.3 comet_ml>=3.41.0 deepspeed-kernels ; sys_platform == 'linux' docutils<0.18 +einops future importlib-metadata>=4 mup diff --git a/requirements/requirements.txt b/requirements/requirements.txt index 1af4c69c5807..296398f680cc 100755 --- a/requirements/requirements.txt +++ b/requirements/requirements.txt @@ -1,4 +1,3 @@ -einops hjson msgpack ninja From 211b5c1c57085f08209396c594f678041d0dc44f Mon Sep 17 00:00:00 2001 From: Hoa La Date: Tue, 19 Nov 2024 15:40:16 -0500 Subject: [PATCH 5/5] remove unwanted changes --- bin/deepspeed | 2 +- bin/deepspeed.pt | 2 +- bin/dsr | 2 +- deepspeed/accelerator | 2 +- deepspeed/ops/csrc | 2 +- deepspeed/ops/op_builder | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/bin/deepspeed b/bin/deepspeed index 7f8898257a27..6b7685641019 120000 --- a/bin/deepspeed +++ b/bin/deepspeed @@ -1 +1 @@ -ds +ds \ No newline at end of file diff --git a/bin/deepspeed.pt b/bin/deepspeed.pt index 7f8898257a27..6b7685641019 120000 --- a/bin/deepspeed.pt +++ b/bin/deepspeed.pt @@ -1 +1 @@ -ds +ds \ No newline at end of file diff --git a/bin/dsr b/bin/dsr index a3a3b689ba18..747bf4722c42 120000 --- a/bin/dsr +++ b/bin/dsr @@ -1 +1 @@ -ds_report +ds_report \ No newline at end of file diff --git a/deepspeed/accelerator b/deepspeed/accelerator index 15fddb37dd6d..b61ffface2d0 120000 --- a/deepspeed/accelerator +++ b/deepspeed/accelerator @@ -1 +1 @@ -../accelerator/ +../accelerator/ \ No newline at end of file diff --git a/deepspeed/ops/csrc b/deepspeed/ops/csrc index fd4e1dc24e66..338652302505 120000 --- a/deepspeed/ops/csrc +++ b/deepspeed/ops/csrc @@ -1 +1 @@ -../../csrc +../../csrc \ No newline at end of file diff --git a/deepspeed/ops/op_builder b/deepspeed/ops/op_builder index 05f303450d1c..db4f9c335065 120000 --- a/deepspeed/ops/op_builder +++ b/deepspeed/ops/op_builder @@ -1 +1 @@ -../../op_builder +../../op_builder \ No newline at end of file