[zero] faster flatten/unflatten (cpp version)

[stas00]

(this PR has evolved and has been edited from its original version which proposed using apex's functions)

This PR is trying to switch to the fastest version of flatten/unflatten and consistently use it through all zero stages.

as discovered in later comments of this PR #910 (comment) UtilsBuilder().load().flatten is on par with apex_C.flatten speed-wise and 2-3x faster than torch._utils._flatten_dense_tensors
(about the same for unflatten)

So this PR was modified to switch to UtilsBuilder().load().flatten and UtilsBuilder().load().unflatten

Let's try first if it may work loading it at import-level since it will trigger build the first time for JIT. If it doesn't work we can fold it into each class's __init__ like zero2.py had it prior to this PR.

Fixes: #877

@jeffra

[jeffra]

Good point @stas00, curious on @samyam's thoughts here too but we do have apex's (un)flatten ops ported into our JIT-compatible ops. We use them in zero-2, take a look here:

DeepSpeed/deepspeed/runtime/zero/stage2.py

Lines 124 to 127 in 7fcc891

	# Load pre-installed or JIT compile (un)flatten ops
	util_ops = UtilsBuilder().load()
	self.flatten = util_ops.flatten
	self.unflatten = util_ops.unflatten

We might be able to use these in z3 instead so we don't have to depend on apex.

[stas00]

Oh! then why doesn't zero3 use those? not fully optimized yet?

yes, apex is being phased out, so it's definitely a good idea to either ask pytorch to port this to the core or to copy away.

I haven't done the benchmarks so I suppose the first question is if someone did pytorch vs. apex vs your version.

[jeffra]

We should be able to use it as a drop in replacement I suspect. The (un)flatten code in our op is taken directly from apex, which was taken from pytorch haha take a look, it's not doing anything fancy:

https://github.com/microsoft/DeepSpeed/blob/master/csrc/utils/flatten_unflatten.cpp

[stas00]

This is totally weird then because they don't use the c++ implementation in the core:
https://github.com/pytorch/pytorch/blob/33b95c6bac5cf9c52fa646bf7335664b0556263d/torch/_utils.py#L248

[stas00]

Also there is a related issue: #877 if this apex code is removed it'll close this issue as well.

[jeffra]

This is totally weird then because they don't use the c++ implementation in the core:
https://github.com/pytorch/pytorch/blob/33b95c6bac5cf9c52fa646bf7335664b0556263d/torch/_utils.py#L248

haha yep, I've also dug into this and don't understand why torch doesn't just use these ops directly. Seems odd...

I believe @RezaYazdaniAminabadi did some small benchmarking at one point between the torch (un)flatten functions vs the cpp op and the cpp op performed significantly better. But I think we sometimes forget to use it in practice.

[stas00]

It looks like this is an identical code, c++

https://github.com/pytorch/pytorch/blob/547346d66350b4ac325941752f53fa92f756f6a9/torch/csrc/utils/tensor_flatten.h#L10

inline at::Tensor flatten_dense_tensors(at::TensorList tensors) {
  static auto flatten = [](const at::Tensor &t) { return t.contiguous().view({-1}); };
  if (tensors.size() == 1)
    return flatten(tensors[0]);
  return at::cat(fmap(tensors, flatten));
}

python:

https://github.com/pytorch/pytorch/blob/33b95c6bac5cf9c52fa646bf7335664b0556263d/torch/_utils.py#L248

def _flatten_dense_tensors(tensors):
    if len(tensors) == 1:
        return tensors[0].contiguous().view(-1)
    flat = torch.cat([t.contiguous().view(-1) for t in tensors], dim=0)
    return flat

I wonder why it'd be faster other than bypassing the python machinery for intermediary results.

[stas00]

I believe @RezaYazdaniAminabadi did some small benchmarking at one point between the torch (un)flatten functions vs the cpp op and the cpp op performed significantly better. But I think we sometimes forget to use it in practice.

@RezaYazdaniAminabadi, do you by chance have the benchmark still? If it's faster we should notify the pytorch devs to rectify this omission - perhaps this function pair is not widely used.

[stas00]

I did some flatten benchmarking and the results are in:

I compared:

1. torch._utils._flatten_dense_tensors
2. UtilsBuilder().load().flatten
3. apex_C.flatten

Results:

1. UtilsBuilder().load().flatten and apex_C.flatten are on par
2. torch._utils._flatten_dense_tensors is 2 times slower than the first 2

For unflatten it's about the same with the gap for pytorch version at about 2-3x slower than the first 2 versions.

This was on RTX-3090. W/ prebuilt deepspeed.

I verified with cProfile, line_profiler and timeit - all getting the same results.

Setup:

pip install line_profiler

Benchmark for flatten

#!/usr/bin/env python

import argparse

import gc

import torch
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
from deepspeed.ops.op_builder import UtilsBuilder

from apex_C import flatten as flatten_apex

util_ops = UtilsBuilder().load()
flatten = util_ops.flatten
unflatten = util_ops.unflatten

torch.manual_seed(0)
# emulate a small typical model weights
x = [torch.rand((512,512)).cuda(), torch.rand((512,1024)).cuda(), torch.rand((512,30000)).cuda()]
t = x * 30

# warm up and check that the same output is produced
flat_py = _flatten_dense_tensors(t)
flat_cpp = flatten(t)
flat_apex = flatten_apex(t)
#numel = flat_cpp.numel()
assert torch.eq(flat_py, flat_cpp).all(), "both produce the same tensor"
assert torch.eq(flat_py, flat_apex).all(), "both produce the same tensor"

TIMES = 1000

# the programs being tested
def py():
    for i in range(TIMES):
        flat = _flatten_dense_tensors(t)

def cpp():
    for i in range(TIMES):
        flat = flatten(t)

def apex():
    for i in range(TIMES):
        flat = flatten_apex(t)

#### cProfile ####

import cProfile

def cprofileme():
    print("--------------- cProfile -----------------")
    print("py")
    cProfile.run("py()", sort=-1)
    gc.collect(); torch.cuda.empty_cache()
    print("cpp")
    cProfile.run("cpp()", sort=-1)
    gc.collect(); torch.cuda.empty_cache()
    print("apex")
    cProfile.run("apex()", sort=-1)
    gc.collect(); torch.cuda.empty_cache()

#### timeit ####

import timeit

def timeme():
    print("--------------- timeit -----------------")
    print(f'py  ={timeit.Timer("py()", globals=globals()).timeit(number=1)}')
    gc.collect(); torch.cuda.empty_cache()
    print(f'cpp ={timeit.Timer("cpp()", globals=globals()).timeit(number=1)}')
    gc.collect(); torch.cuda.empty_cache()
    print(f'apex={timeit.Timer("apex()", globals=globals()).timeit(number=1)}')
    gc.collect(); torch.cuda.empty_cache()

#### line_profiler ####
# this one requires a special way to be called
# pip install line_profiler
# kernprof -l flatten_bench.py -l; python -m line_profiler  flatten_bench.py.lprof

def line_profileme():
    print("--------------- line_profier -----------------")
    print("py")
    profile(py)()
    gc.collect(); torch.cuda.empty_cache()
    print("cpp")
    profile(cpp)()
    gc.collect(); torch.cuda.empty_cache()
    print("apex")
    profile(apex)()
    gc.collect(); torch.cuda.empty_cache()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-l", action='store_true')
    parser.add_argument("-c", action='store_true')
    parser.add_argument("-t", action='store_true')
    args = parser.parse_args()
    if args.l:
        line_profileme()
    elif args.c:
        cprofileme()
    elif args.t:
        timeme()

It looks like that if I mix cProfile with timeit, the latter gets invalid results. So have to run each profiler separately.

Also before I added empty_cache the results were quite invalid - the first to run was getting 32x faster outcome!

Output:

(main-38) ✘ /hf/deepspeed> ./flatten_bench.py -t
--------------- timeit -----------------
py  =0.1116456389427185
cpp =0.06466162856668234
apex=0.06333659961819649
(main-38) /hf/deepspeed> ./flatten_bench.py -c
--------------- cProfile -----------------
py
         184004 function calls in 0.129 seconds

   Ordered by: standard name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.129    0.129 <string>:1(<module>)
     1000    0.009    0.000    0.128    0.000 _utils.py:248(_flatten_dense_tensors)
     1000    0.017    0.000    0.094    0.000 _utils.py:264(<listcomp>)
        1    0.001    0.001    0.129    0.129 flatten_bench.py:33(py)
        1    0.000    0.000    0.129    0.129 {built-in method builtins.exec}
     1000    0.000    0.000    0.000    0.000 {built-in method builtins.len}
     1000    0.024    0.000    0.024    0.000 {built-in method cat}
    90000    0.012    0.000    0.012    0.000 {method 'contiguous' of 'torch._C._TensorBase' objects}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}
    90000    0.065    0.000    0.065    0.000 {method 'view' of 'torch._C._TensorBase' objects}


cpp
         1004 function calls in 0.063 seconds

   Ordered by: standard name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.063    0.063 <string>:1(<module>)
        1    0.001    0.001    0.063    0.063 flatten_bench.py:37(cpp)
        1    0.000    0.000    0.063    0.063 {built-in method builtins.exec}
     1000    0.062    0.000    0.062    0.000 {built-in method deepspeed.ops.utils_op.flatten}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}


apex
         1004 function calls in 0.072 seconds

   Ordered by: standard name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.072    0.072 <string>:1(<module>)
        1    0.001    0.001    0.072    0.072 flatten_bench.py:41(apex)
     1000    0.071    0.000    0.071    0.000 {built-in method apex_C.flatten}
        1    0.000    0.000    0.072    0.072 {built-in method builtins.exec}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}

Benchmark for unflatten

#!/usr/bin/env python

import argparse
import gc
import torch
from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
from deepspeed.ops.op_builder import UtilsBuilder

from apex_C import flatten as flatten_apex
from apex_C import unflatten as unflatten_apex

util_ops = UtilsBuilder().load()
flatten = util_ops.flatten
unflatten = util_ops.unflatten

torch.manual_seed(0)
# emulate a small typical model weights
x = [torch.rand((512,512)).cuda(), torch.rand((512,1024)).cuda(), torch.rand((512,30000)).cuda()]
unflat_t = x * 30

# warm up and check that the same output is produced
flat_py = _flatten_dense_tensors(unflat_t)
flat_cpp = flatten(unflat_t)
flat_apex = flatten_apex(unflat_t)
#numel = flat_cpp.numel()
assert torch.eq(flat_py, flat_cpp).all(), "both produce the same tensor"
assert torch.eq(flat_py, flat_apex).all(), "both produce the same tensor"

flat_t = flat_py
unflat_py = _unflatten_dense_tensors(flat_py, unflat_t)
for i in range(len(unflat_t)): assert torch.eq(unflat_t[i], unflat_py[i]).all()
unflat_cpp = _unflatten_dense_tensors(flat_cpp, unflat_t)
for i in range(len(unflat_t)): assert torch.eq(unflat_t[i], unflat_cpp[i]).all()
unflat_apex = _unflatten_dense_tensors(flat_apex, unflat_t)
for i in range(len(unflat_t)): assert torch.eq(unflat_t[i], unflat_apex[i]).all()

# the programs being tested
def py():
    for i in range(1000):
        unflat = _unflatten_dense_tensors(flat_t, unflat_t)

def cpp():
    for i in range(1000):
        unflat = unflatten(flat_t, unflat_t)

def apex():
    for i in range(1000):
        unflat = unflatten_apex(flat_t, unflat_t)


#### cProfile ####

import cProfile

def cprofileme():
    print("--------------- cProfile -----------------")
    print("py")
    cProfile.run("py()", sort=-1)
    gc.collect(); torch.cuda.empty_cache()
    print("cpp")
    cProfile.run("cpp()", sort=-1)
    gc.collect(); torch.cuda.empty_cache()
    print("apex")
    cProfile.run("apex()", sort=-1)
    gc.collect(); torch.cuda.empty_cache()

#### timeit ####

import timeit

def timeme():
    print("--------------- timeit -----------------")
    print(f'py  ={timeit.Timer("py()", globals=globals()).timeit(number=1)}')
    gc.collect(); torch.cuda.empty_cache()
    print(f'cpp ={timeit.Timer("cpp()", globals=globals()).timeit(number=1)}')
    gc.collect(); torch.cuda.empty_cache()
    print(f'apex={timeit.Timer("apex()", globals=globals()).timeit(number=1)}')
    gc.collect(); torch.cuda.empty_cache()

#### line_profiler ####
# this one requires a special way to be called
# pip install line_profiler
# kernprof -l unflatten_bench.py -l; python -m line_profiler unflatten_bench.py.lprof

def line_profileme():
    print("--------------- line_profier -----------------")
    print("py")
    profile(py)()
    gc.collect(); torch.cuda.empty_cache()
    print("cpp")
    profile(cpp)()
    gc.collect(); torch.cuda.empty_cache()
    print("apex")
    profile(apex)()
    gc.collect(); torch.cuda.empty_cache()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-l", action='store_true')
    parser.add_argument("-c", action='store_true')
    parser.add_argument("-t", action='store_true')
    args = parser.parse_args()
    if args.l:
        line_profileme()
    elif args.c:
        cprofileme()
    elif args.t:
        timeme()

Sample output:

(main-38) /hf/deepspeed> ./unflatten_bench.py -c
--------------- cProfile -----------------
py
         361004 function calls in 0.255 seconds

   Ordered by: standard name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.255    0.255 <string>:1(<module>)
     1000    0.063    0.000    0.245    0.000 _utils.py:284(_unflatten_dense_tensors)
        1    0.009    0.009    0.254    0.254 unflatten_bench.py:38(py)
        1    0.000    0.000    0.255    0.255 {built-in method builtins.exec}
    90000    0.005    0.000    0.005    0.000 {method 'append' of 'list' objects}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}
    90000    0.094    0.000    0.094    0.000 {method 'narrow' of 'torch._C._TensorBase' objects}
    90000    0.011    0.000    0.011    0.000 {method 'numel' of 'torch._C._TensorBase' objects}
    90000    0.072    0.000    0.072    0.000 {method 'view_as' of 'torch._C._TensorBase' objects}


cpp
         1004 function calls in 0.082 seconds

   Ordered by: standard name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.082    0.082 <string>:1(<module>)
        1    0.009    0.009    0.082    0.082 unflatten_bench.py:42(cpp)
        1    0.000    0.000    0.082    0.082 {built-in method builtins.exec}
     1000    0.073    0.000    0.073    0.000 {built-in method deepspeed.ops.utils_op.unflatten}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}


apex
         1004 function calls in 0.081 seconds

   Ordered by: standard name

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.000    0.000    0.081    0.081 <string>:1(<module>)
        1    0.009    0.009    0.081    0.081 unflatten_bench.py:46(apex)
     1000    0.073    0.000    0.073    0.000 {built-in method apex_C.unflatten}
        1    0.000    0.000    0.081    0.081 {built-in method builtins.exec}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Profiler' objects}

kernprof -l unflatten_bench.py -l; python -m line_profiler  unflatten_bench.py.lprof
--------------- line_profier -----------------
py
cpp
apex
Wrote profile results to unflatten_bench.py.lprof
Timer unit: 1e-06 s

Total time: 0.254009 s
File: unflatten_bench.py
Function: py at line 38

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    38                                           def py():
    39      1001        356.0      0.4      0.1      for i in range(1000):
    40      1000     253653.0    253.7     99.9          unflat = _unflatten_dense_tensors(flat_t, unflat_t)

Total time: 0.088684 s
File: unflatten_bench.py
Function: cpp at line 42

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    42                                           def cpp():
    43      1001        304.0      0.3      0.3      for i in range(1000):
    44      1000      88380.0     88.4     99.7          unflat = unflatten(flat_t, unflat_t)

Total time: 0.087492 s
File: unflatten_bench.py
Function: apex at line 46

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    46                                           def apex():
    47      1001        334.0      0.3      0.4      for i in range(1000):
    48      1000      87158.0     87.2     99.6          unflat = unflatten_apex(flat_t, unflat_t)

[stas00]

Proposed to pytorch to switch to the cpp version pytorch/pytorch#55240 so down the road we won't need the workaround.

[stas00]

It looks good now. As I had to turn some functions into methods to make this work, I'm not sure if I placed them in the most intuitive place between other existing methods - please feel free to move them or tell me where to move them. Either way works.

[stas00]

FYI, pytorch has now replaced the slower python version with the cpp version pytorch/pytorch#58006