Transformer installation fails

[anderson1008]

I am following the instruction published on Nvidia web to install Nemo.
See this link
One of step is to install the Transformer engine. It failed with the following signature.

[29/35] Building CUDA object common/CMakeFiles/transformer_engine.dir/fused_attn/fused_attn_fp8.cu.o
      FAILED: common/CMakeFiles/transformer_engine.dir/fused_attn/fused_attn_fp8.cu.o
      /usr/local/cuda/bin/nvcc -forward-unknown-to-host-compiler -Dtransformer_engine_EXPORTS -I/workspace/TransformerEngine/transformer_engine -I/workspace/TransformerEngine/transformer_engine/common/include -I/workspace/TransformerEngine/transformer_engine/common/../../3rdparty/cudnn-frontend/include -I/workspace/TransformerEngine/build/cmake/common/string_headers -isystem /usr/local/cuda/targets/x86_64-linux/include --threads 4 --expt-relaxed-constexpr -O3 -O3 -DNDEBUG -std=c++17 "--generate-code=arch=compute_70,code=[compute_70,sm_70]" "--generate-code=arch=compute_80,code=[compute_80,sm_80]" "--generate-code=arch=compute_89,code=[compute_89,sm_89]" "--generate-code=arch=compute_90,code=[compute_90,sm_90]" -Xcompiler=-fPIC -MD -MT common/CMakeFiles/transformer_engine.dir/fused_attn/fused_attn_fp8.cu.o -MF common/CMakeFiles/transformer_engine.dir/fused_attn/fused_attn_fp8.cu.o.d -x cu -c /workspace/TransformerEngine/transformer_engine/common/fused_attn/fused_attn_fp8.cu -o common/CMakeFiles/transformer_engine.dir/fused_attn/fused_attn_fp8.cu.o
      Killed
      Killed
      Killed
      Killed
      [30/35] Building CUDA object common/CMakeFiles/transformer_engine.dir/fused_attn/utils.cu.o
      FAILED: common/CMakeFiles/transformer_engine.dir/fused_attn/utils.cu.o
      /usr/local/cuda/bin/nvcc -forward-unknown-to-host-compiler -Dtransformer_engine_EXPORTS -I/workspace/TransformerEngine/transformer_engine -I/workspace/TransformerEngine/transformer_engine/common/include -I/workspace/TransformerEngine/transformer_engine/common/../../3rdparty/cudnn-frontend/include -I/workspace/TransformerEngine/build/cmake/common/string_headers -isystem /usr/local/cuda/targets/x86_64-linux/include --threads 4 --expt-relaxed-constexpr -O3 -O3 -DNDEBUG -std=c++17 "--generate-code=arch=compute_70,code=[compute_70,sm_70]" "--generate-code=arch=compute_80,code=[compute_80,sm_80]" "--generate-code=arch=compute_89,code=[compute_89,sm_89]" "--generate-code=arch=compute_90,code=[compute_90,sm_90]" -Xcompiler=-fPIC -MD -MT common/CMakeFiles/transformer_engine.dir/fused_attn/utils.cu.o -MF common/CMakeFiles/transformer_engine.dir/fused_attn/utils.cu.o.d -x cu -c /workspace/TransformerEngine/transformer_engine/common/fused_attn/utils.cu -o common/CMakeFiles/transformer_engine.dir/fused_attn/utils.cu.o
      Killed
      [31/35] Building CUDA object common/CMakeFiles/transformer_engine.dir/transpose/cast_transpose_fusion.cu.o
      [32/35] Building CUDA object common/CMakeFiles/transformer_engine.dir/fused_attn/fused_attn_f16_max512_seqlen.cu.o
      [33/35] Building CUDA object common/CMakeFiles/transformer_engine.dir/fused_attn/fused_attn_f16_arbitrary_seqlen.cu.o
      [34/35] Building CUDA object common/CMakeFiles/transformer_engine.dir/layer_norm/ln_fwd_cuda_kernel.cu.o
      ninja: build stopped: subcommand failed.
      Traceback (most recent call last):
        File "/workspace/TransformerEngine/setup.py", line 336, in _build_cmake
          subprocess.run(command, cwd=build_dir, check=True)
        File "/usr/lib/python3.10/subprocess.py", line 526, in run
          raise CalledProcessError(retcode, process.args,
      subprocess.CalledProcessError: Command '['/usr/local/lib/python3.10/dist-packages/cmake/data/bin/cmake', '--build', '/workspace/TransformerEngine/build/cmake']' returned non-zero exit status 1.
      
      During handling of the above exception, another exception occurred:
      
      Traceback (most recent call last):
        File "<string>", line 2, in <module>
        File "<pip-setuptools-caller>", line 34, in <module>
        File "/workspace/TransformerEngine/setup.py", line 611, in <module>
          main()
        File "/workspace/TransformerEngine/setup.py", line 596, in main
          setuptools.setup(
        File "/usr/local/lib/python3.10/dist-packages/setuptools/__init__.py", line 103, in setup
          return distutils.core.setup(**attrs)
        File "/usr/local/lib/python3.10/dist-packages/setuptools/_distutils/core.py", line 185, in setup
          return run_commands(dist)
        File "/usr/local/lib/python3.10/dist-packages/setuptools/_distutils/core.py", line 201, in run_commands
          dist.run_commands()
        File "/usr/local/lib/python3.10/dist-packages/setuptools/_distutils/dist.py", line 969, in run_commands
          self.run_command(cmd)
        File "/usr/local/lib/python3.10/dist-packages/setuptools/dist.py", line 989, in run_command
          super().run_command(command)
        File "/usr/local/lib/python3.10/dist-packages/setuptools/_distutils/dist.py", line 988, in run_command
          cmd_obj.run()
        File "/usr/local/lib/python3.10/dist-packages/wheel/bdist_wheel.py", line 368, in run
          self.run_command("build")
        File "/usr/local/lib/python3.10/dist-packages/setuptools/_distutils/cmd.py", line 318, in run_command
          self.distribution.run_command(command)
        File "/usr/local/lib/python3.10/dist-packages/setuptools/dist.py", line 989, in run_command
          super().run_command(command)
        File "/usr/local/lib/python3.10/dist-packages/setuptools/_distutils/dist.py", line 988, in run_command
          cmd_obj.run()
        File "/usr/local/lib/python3.10/dist-packages/setuptools/_distutils/command/build.py", line 131, in run
          self.run_command(cmd_name)
        File "/usr/local/lib/python3.10/dist-packages/setuptools/_distutils/cmd.py", line 318, in run_command
          self.distribution.run_command(command)
        File "/usr/local/lib/python3.10/dist-packages/setuptools/dist.py", line 989, in run_command
          super().run_command(command)
        File "/usr/local/lib/python3.10/dist-packages/setuptools/_distutils/dist.py", line 988, in run_command
          cmd_obj.run()
        File "/workspace/TransformerEngine/setup.py", line 368, in run
          ext._build_cmake(
        File "/workspace/TransformerEngine/setup.py", line 338, in _build_cmake
          raise RuntimeError(f"Error when running CMake: {e}")
      RuntimeError: Error when running CMake: Command '['/usr/local/lib/python3.10/dist-packages/cmake/data/bin/cmake', '--build', '/workspace/TransformerEngine/build/cmake']' returned non-zero exit status 1.
      [end of output]
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
  ERROR: Failed building wheel for transformer-engine
  Running setup.py clean for transformer-engine
Failed to build transformer-engine
ERROR: Could not build wheels for transformer-engine, which is required to install pyproject.toml-based projects

[anderson1008]

It seems that process is killed due to OOM. I tried to start the docker with the following cmd, but it still does not help.

docker run -m 16g --memory-swap=17g --shm-size=16g --oom-kill-disable --rm --gpus all --ulimit memlock=-1 --ulimit stack=67108864 -it nvcr.io/nvidia/pytorch:24.02-py3

In demsg

[ 4518.680704] Memory cgroup out of memory: Killed process 22194 (cicc) total-vm:2187596kB, anon-rss:1907520kB, file-rss:15168kB, shmem-rss:0kB, UID:0 pgtables:4064kB oom_score_adj:0

And it says my kernel does not support OOM kill disable.

WARNING: Your kernel does not support OomKillDisable. OomKillDisable discarded.

Can anyone please clarify how much memory is needed to install transformer? My system has 32G.
Most importantly, how can I get around this?

[anderson1008]

I might be hitting a known bug here:
see here

Apparently, building flash attention module will take up to 60GB of memory. So, for machine does not have enough RAM, we can get around this by creating a large swap partition.

Here is my workaround:

sudo swapoff -a
sudo fallocate -l 128G /swapfile
sudo chmod 600 /swapfile 
sudo mkswap /swapfile
sudo swapon --show
sudo swapon /swapfile
docker run -m 16g --memory-swap=128g --shm-size=16g --rm --gpus all --ulimit memlock=-1 --ulimit stack=67108864 -it nvcr.io/nvidia/pytorch:24.02-py3 

Yay! It works for me!

[anderson1008]

Found workaround.