Support for ARM wheel

[csahithi]

Hello @fzyzcjy , currently I see a wheel available only for x86 - https://pypi.org/project/torch-memory-saver/0.0.9rc3/#files. Can we also get a wheel for aarch64?
When I try to install torch-memory-saver for ARM, it is unable to find the wheel (like it does in x86) and starts building from source instead and fails

 × Failed to build `torch-memory-saver @ file:///home/torch_memory_saver`
  ├─▶ The build backend returned an error
  ╰─▶ Call to `setuptools.build_meta:__legacy__.build_wheel` failed (exit status: 1)

      [stdout]
      Detected platform: cuda
      running bdist_wheel
      running build
      running build_py
      creating build/lib.linux-aarch64-cpython-312/torch_memory_saver
      copying torch_memory_saver/__init__.py -> build/lib.linux-aarch64-cpython-312/torch_memory_saver
      copying torch_memory_saver/binary_wrapper.py -> build/lib.linux-aarch64-cpython-312/torch_memory_saver
      copying torch_memory_saver/entrypoint.py -> build/lib.linux-aarch64-cpython-312/torch_memory_saver
      copying torch_memory_saver/testing_utils.py -> build/lib.linux-aarch64-cpython-312/torch_memory_saver
      copying torch_memory_saver/utils.py -> build/lib.linux-aarch64-cpython-312/torch_memory_saver
      creating build/lib.linux-aarch64-cpython-312/torch_memory_saver/hooks
      copying torch_memory_saver/hooks/__init__.py -> build/lib.linux-aarch64-cpython-312/torch_memory_saver/hooks
      copying torch_memory_saver/hooks/base.py -> build/lib.linux-aarch64-cpython-312/torch_memory_saver/hooks
      copying torch_memory_saver/hooks/mode_preload.py -> build/lib.linux-aarch64-cpython-312/torch_memory_saver/hooks
      copying torch_memory_saver/hooks/mode_torch.py -> build/lib.linux-aarch64-cpython-312/torch_memory_saver/hooks
      running build_ext
      building 'torch_memory_saver_hook_mode_preload' extension
      creating build/temp.linux-aarch64-cpython-312/csrc
      aarch64-linux-gnu-g++ -fno-strict-overflow -Wsign-compare -DNDEBUG -g -O2 -Wall -fPIC -DPy_LIMITED_API=0x03090000
      -DUSE_CUDA=1 -DTMS_HOOK_MODE_PRELOAD=1 -I/usr/local/cuda/include -I/home/.cache/uv/builds-v0/.tmpwYZ4kV/include
      -I/usr/include/python3.12 -c csrc/api_forwarder.cpp -o build/temp.linux-aarch64-cpython-312/csrc/api_forwarder.o -std=c++17
      -O3

      [stderr]
      In file included from csrc/api_forwarder.h:3,
                       from csrc/api_forwarder.cpp:2:
      csrc/macro.h:46:10: fatal error: cuda_runtime_api.h: No such file or directory
         46 | #include <cuda_runtime_api.h>
            |          ^~~~~~~~~~~~~~~~~~~~
      compilation terminated.
      error: command '/usr/bin/aarch64-linux-gnu-g++' failed with exit code 1

      hint: This error likely indicates that you need to install a library that provides "cuda_runtime_api.h" for
      `file:///home/torch_memory_saver`

[faradawn]

Same error with CUDA 13 on ARM. Upvote ^

[fzyzcjy]

hmm I am a bit confused

      csrc/macro.h:46:10: fatal error: cuda_runtime_api.h: No such file or directory
         46 | #include <cuda_runtime_api.h>

installing from source should not have such issue... could you please try whether install from source in x86 has this issue?

[csahithi]

Installing from source on x86 doesn't show this issue. I've observed this only on ARM

[fzyzcjy]

Interesting... could you please examine a bit the difference b/t x86 and arm? Also I have successfully used mem saver in grace blackwell machines, thus seems it will not error.

[csahithi]

My observation has been that I have been able to run it on both x86 and ARM (grace blackwell) if I run it inside a container. But when I try to run it on baremetal, it works on x86 but seeing the above error on ARM (grace blackwell). Please let me know if you want to me to check for something in specific and I can provide that info.

[fzyzcjy]

Interesting... then maybe it is b/c your baremetal arm machine has a different system configuration?

my first guess is whether

torch_memory_saver/setup.py

Line 13 in 0f3d6ed

def _find_platform_home(platform):

or

torch_memory_saver/setup.py

Line 106 in 0f3d6ed

include_dirs = [str((platform_home / 'include').resolve())]

does not suite your baremetal case