ziplab/mesa

This is the official PyTorch implementation for "Mesa: A Memory-saving Training Framework for Transformers".

A Memory-saving Training Framework for Transformers

This is the official PyTorch implementation for Mesa: A Memory-saving Training Framework for Transformers.

By Zizheng Pan, Peng Chen, Haoyu He, Jing Liu, Jianfei Cai and Bohan Zhuang.

Installation

1. Create a virtual environment with anaconda.

   conda create -n mesa python=3.7 -y
   conda activate mesa
   
   # Install PyTorch, we use PyTorch 1.7.1 with CUDA 10.1 
   pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html
   
   # Install ninja
   pip install ninja

2. Build and install Mesa.

   # clone this repo
   git clone https://github.com/zhuang-group/Mesa
   # build
   cd Mesa/
   # You need to have an NVIDIA GPU
   python setup.py develop

Usage

1. Prepare your policy and save as a text file, e.g. policy.txt.

   on gelu: # layer tag, choices: fc, conv, gelu, bn, relu, softmax, matmul, layernorm
       by_index: all # layer index
       enable: True # enable for compressing
       level: 256 # we adopt 8-bit quantization by default
       ema_decay: 0.9 # the decay rate for running estimates
       
       by_index: 1 2 # e.g. exluding GELU layers that indexed by 1 and 2.
       enable: False

2. Next, you can wrap your model with Mesa by:

   import mesa as ms
   ms.policy.convert_by_num_groups(model, 3)
   # or convert by group size with ms.policy.convert_by_group_size(model, 64)
   
   # setup compression policy
   ms.policy.deploy_on_init(model, '[path to policy.txt]', verbose=print, override_verbose=False)

   That's all you need to use Mesa for memory saving.

   Note that convert_by_num_groups and convert_by_group_size only recognize nn.XXX, if your code has functional operations, such as Q@K and F.Softmax, you may need to manually setup these layers. For example:

   import mesa as ms
   # matrix multipcation (before)
   out = [e-mail](-2, -1)
   # with Mesa
   self.mm = ms.MatMul(quant_groups=3)
   out = self.mm(q, k.transpose(-2, -1))
   
   # sofmtax (before)
   attn = attn.softmax(dim=-1)
   # with Mesa
   self.softmax = ms.Softmax(dim=-1, quant_groups=3)
   attn = self.softmax(attn)

3. You can also target one layer by:

   import mesa as ms
   # previous 
   self.act = nn.GELU()
   # with Mesa
   self.act = ms.GELU(quant_groups=[num of quantization groups])

Demo projects for DeiT and Swin

We provide demo projects to replicate our results of training DeiT and Swin with Mesa, please refer to DeiT-Mesa and Swin-Mesa.

Results on ImageNet

Model	Param (M)	FLOPs (G)	Train Memory (MB)	Top-1 (%)
DeiT-Ti	5	1.3	4,171	71.9
DeiT-Ti w/ Mesa	5	1.3	1,858	72.1
DeiT-S	22	4.6	8,459	79.8
DeiT-S w/ Mesa	22	4.6	3,840	80.0
DeiT-B	86	17.5	17,691	81.8
DeiT-B w/ Mesa	86	17.5	8,616	81.8
Swin-Ti	29	4.5	11,812	81.3
Swin-Ti w/ Mesa	29	4.5	5,371	81.3
PVT-Ti	13	1.9	7,800	75.1
PVT-Ti w/ Mesa	13	1.9	3,782	74.9

Memory footprint at training time is measured with a batch size of 128 and an image resolution of 224x224 on a single GPU.

Citation

If you find our work interesting or helpful to your research, please consider citing Mesa.

@article{pan2021mesa,
      title={Mesa: A Memory-saving Training Framework for Transformers}, 
      author={Zizheng Pan and Peng Chen and Haoyu He and Jing Liu and Jianfei Cai and Bohan Zhuang},
      journal={arXiv preprint arXiv:2111.11124}
      year={2021}
}

License

This repository is released under the Apache 2.0 license as found in the LICENSE file.

Acknowledgments

This repository has adopted part of the quantization codes from ActNN, we thank the authors for their open-sourced code.

mesa's Issues

How do you measure train memory?

Thanks for your great work! I have a question about the memory usage.

How do you get the "train memory" in this table https://github.com/zhuang-group/Mesa#results-on-imagenet?
Is it the total memory required for training? Do you get it analytically or empirically using torch.cuda.memory_allocated?

If I understand correctly, Mesa compresses fp16 activations to int8, so it can at most reduce the activation memory by 2x.
There are also other components in total memory, so Mesa cannot reduce the total memory by 2x.
However, in your results, Mesa reduces memory by more than 2x for some models. How is this possible?

Compilation failed with torch=1.12.1

Hello! Thank you for the great work.

I am currently trying to use the Mesa with PyTorch 1.12.1 version, but got errors during the system compilation.

I also checked that the internal operations have changed at PyTorch 1.10 as mentioned at #1 .

I think there are major changes at PyTorch 1.12, too.
I used this command to construct the environment.

conda create -n env python=3.8.8
conda activate env
conda install -c anaconda cudatoolkit==11.3.1
conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch

And I got this error while compilation.

running install
/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/cmd.py:66: SetuptoolsDeprecationWarning: setup.py install is deprecated.
!!

        ********************************************************************************
        Please avoid running ``setup.py`` directly.
        Instead, use pypa/build, pypa/installer or other
        standards-based tools.

        See https://blog.ganssle.io/articles/2021/10/setup-py-deprecated.html for details.
        ********************************************************************************

!!
  self.initialize_options()
/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/cmd.py:66: EasyInstallDeprecationWarning: easy_install command is deprecated.
!!

        ********************************************************************************
        Please avoid running ``setup.py`` and ``easy_install``.
        Instead, use pypa/build, pypa/installer or other
        standards-based tools.

        See https://github.com/pypa/setuptools/issues/917 for details.
        ********************************************************************************

!!
  self.initialize_options()
running bdist_egg
running egg_info
writing Mesa.egg-info/PKG-INFO
writing dependency_links to Mesa.egg-info/dependency_links.txt
writing top-level names to Mesa.egg-info/top_level.txt
reading manifest file 'Mesa.egg-info/SOURCES.txt'
adding license file 'LICENSE'
writing manifest file 'Mesa.egg-info/SOURCES.txt'
installing library code to build/bdist.linux-x86_64/egg
running install_lib
running build_py
copying mesa/custom_quant.py -> build/lib.linux-x86_64-cpython-38/mesa
copying mesa/custom_layer_norm.py -> build/lib.linux-x86_64-cpython-38/mesa
copying mesa/custom_relu.py -> build/lib.linux-x86_64-cpython-38/mesa
copying mesa/custom_softmax.py -> build/lib.linux-x86_64-cpython-38/mesa
copying mesa/custom_matmul.py -> build/lib.linux-x86_64-cpython-38/mesa
copying mesa/custom_conv.py -> build/lib.linux-x86_64-cpython-38/mesa
copying mesa/packbit.py -> build/lib.linux-x86_64-cpython-38/mesa
copying mesa/__init__.py -> build/lib.linux-x86_64-cpython-38/mesa
copying mesa/custom_bn.py -> build/lib.linux-x86_64-cpython-38/mesa
copying mesa/policy.py -> build/lib.linux-x86_64-cpython-38/mesa
copying mesa/custom_fc.py -> build/lib.linux-x86_64-cpython-38/mesa
copying mesa/custom_gelu.py -> build/lib.linux-x86_64-cpython-38/mesa
running build_ext
building 'mesa.native' extension
Emitting ninja build file /home/sungho/BackRazor_Neurips22/Mesa/build/temp.linux-x86_64-cpython-38/build.ninja...
Compiling objects...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
[1/1] c++ -MMD -MF /home/sungho/BackRazor_Neurips22/Mesa/build/temp.linux-x86_64-cpython-38/native.o.d -pthread -B /home/sungho/anaconda3/envs/test/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/torch/include -I/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -I/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/torch/include/TH -I/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/torch/include/THC -I/home/sungho/anaconda3/envs/test/include/python3.8 -c -c /home/sungho/BackRazor_Neurips22/Mesa/native.cpp -o /home/sungho/BackRazor_Neurips22/Mesa/build/temp.linux-x86_64-cpython-38/native.o -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=native -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++14
FAILED: /home/sungho/BackRazor_Neurips22/Mesa/build/temp.linux-x86_64-cpython-38/native.o 
c++ -MMD -MF /home/sungho/BackRazor_Neurips22/Mesa/build/temp.linux-x86_64-cpython-38/native.o.d -pthread -B /home/sungho/anaconda3/envs/test/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/torch/include -I/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -I/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/torch/include/TH -I/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/torch/include/THC -I/home/sungho/anaconda3/envs/test/include/python3.8 -c -c /home/sungho/BackRazor_Neurips22/Mesa/native.cpp -o /home/sungho/BackRazor_Neurips22/Mesa/build/temp.linux-x86_64-cpython-38/native.o -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=native -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++14
cc1plus: warning: command line option ‘-Wstrict-prototypes’ is valid for C/ObjC but not for C++
/home/sungho/BackRazor_Neurips22/Mesa/native.cpp: In function ‘std::tuple<at::Tensor, at::Tensor> conv2d_backward(const at::Tensor&, const at::Tensor&, const at::Tensor&, c10::ArrayRef<long int>, c10::ArrayRef<long int>, c10::ArrayRef<long int>, int64_t, bool, bool, bool, std::array<bool, 2>)’:
/home/sungho/BackRazor_Neurips22/Mesa/native.cpp:26:16: error: ‘cudnn_convolution_backward’ is not a member of ‘at’
     return at::cudnn_convolution_backward(
                ^~~~~~~~~~~~~~~~~~~~~~~~~~
/home/sungho/BackRazor_Neurips22/Mesa/native.cpp:26:16: note: suggested alternative: ‘mps_convolution_backward’
     return at::cudnn_convolution_backward(
                ^~~~~~~~~~~~~~~~~~~~~~~~~~
                mps_convolution_backward
/home/sungho/BackRazor_Neurips22/Mesa/native.cpp: In function ‘at::Tensor softmax_backward_cpu(const at::Tensor&, const at::Tensor&, int64_t, const at::Tensor&)’:
/home/sungho/BackRazor_Neurips22/Mesa/native.cpp:200:69: error: cannot convert ‘const at::Tensor’ to ‘c10::ScalarType’ for argument ‘4’ to ‘at::Tensor at::_softmax_backward_data(const at::Tensor&, const at::Tensor&, int64_t, c10::ScalarType)’
     return at::_softmax_backward_data(grad_output, output, dim, self);
                                                                     ^
/home/sungho/BackRazor_Neurips22/Mesa/native.cpp: In function ‘at::Tensor softmax_backward_cuda(const at::Tensor&, const at::Tensor&, int64_t, const at::Tensor&)’:
/home/sungho/BackRazor_Neurips22/Mesa/native.cpp:212:69: error: cannot convert ‘const at::Tensor’ to ‘c10::ScalarType’ for argument ‘4’ to ‘at::Tensor at::_softmax_backward_data(const at::Tensor&, const at::Tensor&, int64_t, c10::ScalarType)’
     return at::_softmax_backward_data(grad_output, output, dim, self);
                                                                     ^
ninja: build stopped: subcommand failed.
Traceback (most recent call last):
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1808, in _run_ninja_build
    subprocess.run(
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/subprocess.py", line 516, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['ninja', '-v']' returned non-zero exit status 1.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "setup.py", line 15, in <module>
    setup(
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/__init__.py", line 107, in setup
    return distutils.core.setup(**attrs)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/core.py", line 185, in setup
    return run_commands(dist)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/core.py", line 201, in run_commands
    dist.run_commands()
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/dist.py", line 969, in run_commands
    self.run_command(cmd)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/dist.py", line 1234, in run_command
    super().run_command(command)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/dist.py", line 988, in run_command
    cmd_obj.run()
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/command/install.py", line 80, in run
    self.do_egg_install()
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/command/install.py", line 129, in do_egg_install
    self.run_command('bdist_egg')
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/cmd.py", line 318, in run_command
    self.distribution.run_command(command)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/dist.py", line 1234, in run_command
    super().run_command(command)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/dist.py", line 988, in run_command
    cmd_obj.run()
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/command/bdist_egg.py", line 164, in run
    cmd = self.call_command('install_lib', warn_dir=0)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/command/bdist_egg.py", line 150, in call_command
    self.run_command(cmdname)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/cmd.py", line 318, in run_command
    self.distribution.run_command(command)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/dist.py", line 1234, in run_command
    super().run_command(command)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/dist.py", line 988, in run_command
    cmd_obj.run()
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/command/install_lib.py", line 11, in run
    self.build()
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/command/install_lib.py", line 111, in build
    self.run_command('build_ext')
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/cmd.py", line 318, in run_command
    self.distribution.run_command(command)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/dist.py", line 1234, in run_command
    super().run_command(command)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/dist.py", line 988, in run_command
    cmd_obj.run()
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/command/build_ext.py", line 84, in run
    _build_ext.run(self)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/command/build_ext.py", line 345, in run
    self.build_extensions()
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 765, in build_extensions
    build_ext.build_extensions(self)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/command/build_ext.py", line 467, in build_extensions
    self._build_extensions_serial()
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/command/build_ext.py", line 493, in _build_extensions_serial
    self.build_extension(ext)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/command/build_ext.py", line 246, in build_extension
    _build_ext.build_extension(self, ext)
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/setuptools/_distutils/command/build_ext.py", line 548, in build_extension
    objects = self.compiler.compile(
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 586, in unix_wrap_ninja_compile
    _write_ninja_file_and_compile_objects(
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1487, in _write_ninja_file_and_compile_objects
    _run_ninja_build(
  File "/home/sungho/anaconda3/envs/test/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1824, in _run_ninja_build
    raise RuntimeError(message) from e
RuntimeError: Error compiling objects for extension

Can I get some help to solve this compilation error?

Does not compile with new pytorch versions

pip3 install torch==1.10.0+cu113 torchvision==0.11.1+cu113 torchaudio==0.10.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
conda install pytorch=1.10 torchvision torchaudio cudatoolkit=11.3 -c pytorch
conda install pytorch==1.9 torchvision torchaudio cudatoolkit=11.1 -c pytorch -c conda-forge -y

FAILED: /home/Desktop/Mesa/build/temp.linux-x86_64-3.7/native.o 
c++ -MMD -MF /home/Desktop/Mesa/build/temp.linux-x86_64-3.7/native.o.d -pthread -B /home/miniconda3/envs/mesa2/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/home/miniconda3/envs/mesa2/lib/python3.7/site-packages/torch/include -I/home/miniconda3/envs/mesa2/lib/python3.7/site-packages/torch/include/torch/csrc/api/include -I/home/miniconda3/envs/mesa2/lib/python3.7/site-packages/torch/include/TH -I/home/miniconda3/envs/mesa2/lib/python3.7/site-packages/torch/include/THC -I/home/miniconda3/envs/mesa2/include/python3.7m -c -c /home/Desktop/Mesa/native.cpp -o /home/Desktop/Mesa/build/temp.linux-x86_64-3.7/native.o -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=native -D_GLIBCXX_USE_CXX11_ABI=0 -std=c++14
cc1plus: warning: command line option ‘-Wstrict-prototypes’ is valid for C/ObjC but not for C++
/home/Desktop/Mesa/native.cpp: In function ‘at::Tensor gelu_backward_cpu(const at::Tensor&, const at::Tensor&)’:
/home/Desktop/Mesa/native.cpp:92:24: error: ‘gelu_backward_cpu’ is not a member of ‘at::native’; did you mean ‘glu_backward_cpu’?
   92 |     return at::native::gelu_backward_cpu(grad_output, input);
      |                        ^~~~~~~~~~~~~~~~~
      |                        glu_backward_cpu
/home/Desktop/Mesa/native.cpp: In function ‘at::Tensor gelu_backward_cuda(const at::Tensor&, const at::Tensor&)’:
/home/Desktop/Mesa/native.cpp:98:24: error: ‘gelu_backward_cuda’ is not a member of ‘at::native’; did you mean ‘glu_backward_cuda’?
   98 |     return at::native::gelu_backward_cuda(grad_output, input);
      |                        ^~~~~~~~~~~~~~~~~~
      |                        glu_backward_cuda
/home/Desktop/Mesa/native.cpp: In function ‘at::Tensor softmax_backward_cpu(const at::Tensor&, const at::Tensor&, int64_t, const at::Tensor&)’:
/home/Desktop/Mesa/native.cpp:190:24: error: ‘softmax_backward_cpu’ is not a member of ‘at::native’; did you mean ‘col2im_backward_cpu’?
  190 |     return at::native::softmax_backward_cpu(grad_output, output, dim, self);
      |                        ^~~~~~~~~~~~~~~~~~~~
      |                        col2im_backward_cpu
/home/Desktop/Mesa/native.cpp: In function ‘at::Tensor softmax_backward_cuda(const at::Tensor&, const at::Tensor&, int64_t, const at::Tensor&)’:
/home/Desktop/Mesa/native.cpp:198:24: error: ‘softmax_backward_cuda’ is not a member of ‘at::native’; did you mean ‘col2im_backward_cuda’?
  198 |     return at::native::softmax_backward_cuda(grad_output, output, dim, self);
      |                        ^~~~~~~~~~~~~~~~~~~~~
      |                        col2im_backward_cuda
ninja: build stopped: subcommand failed.
Traceback (most recent call last):
  File "/home/miniconda3/envs/mesa2/lib/python3.7/site-packages/torch/utils/cpp_extension.py", line 1723, in _run_ninja_build
    env=env)
  File "/home/miniconda3/envs/mesa2/lib/python3.7/subprocess.py", line 512, in run
    output=stdout, stderr=stderr)
subprocess.CalledProcessError: Command '['ninja', '-v']' returned non-zero exit status 1.

These versions seem to work:

conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=10.2 -c pytorch
conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge
pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html