triton

Build triton under virtualenv explode memory

Hi the following is what I did:

  1. clone LLVM and build with 
    cmake -G Ninja -S llvm -B build -DCMAKE_INSTALL_PREFIX=../bin -DCMAKE_BUILD_TYPE=Debug -DLLVM_ENABLE_ASSERTIONS=ON -DLLVM_ENABLE_PROJECTS="mlir;llvm" -DLLVM_TARGETS_TO_BUILD="host;NVPTX;AMDGPU" -DLLVM_PARALLEL_COMPILE_JOBS=32 -DLLVM_PARALLEL_LINK_JOBS=4
ninja
  2. clone triton repository
  3. make python virtualenv python -m venv .venv --prompt triton
  4. activate virtualenv source .venv/bin/activate
  5. install dependencies pip install ninja cmake wheel
  6. build triton with LLVM_INCLUDE_DIRS=../llvm-project/build/include LLVM_LIBRARY_DIR=../llvm-project/build/lib LLVM_SYSPATH=../llvm-project/build CMAKE_BUILD_TYPE=Debug pip install -e python

And I cannot build it since it cost extremely large memory, even my computer has 48G is not enough.

Thus I would like to ask:

  1. why?? I've built LLVM and does triton really need such a big memory to build??
  2. how can I finish the building??
Solution

  • I hit this issue as well. Wow does memory usage go out of control with 24 threads (Ryzen 5900X). Looking in the setup.py file for Triton, I saw that it uses MAX_JOBS on Linux, so, in the Triton source directory:

    MAX_JOBS="6" pip3 install .

    That keeps the maximum number of threads to 6. It seems each thread will use 1-2.5GB of RAM, so for 32 threads which I see in your question, you'd need at least 80Gb of RAM.

    I was following instructions here when I hit this issue: https://docs.vllm.ai/en/latest/getting_started/amd-installation.html

    Modify your step #6 to include MAX_JOBS="[maximum number of parallel jobs]", which I think with 48GB of RAM, not to go past 16. Maybe try 8 to be conservative.