Add support for non-CUDA `torchmonarch`

[allenwang28]

torchmonarch currently builds only with USE_TENSOR_ENGINE=1 and is only built against CUDA devices, but providing torchmonarch wheels for non-CUDA devices would provide access to "Monarch Core" (actor<>actor communications) on any hardware running linux. See requests:

• AMD (AMD Support #1649)
• Mac OS (Mac build is not working as monarch assumes the presence of cuda #145, Publish artifacts to PyPi during releases to enable easier testing on MAC, etc. #1173)
• TPUs (not created yet but I know there is interest)

It'd also provide a build/publishing path for advanced features, like hardware-specific tensor engine, RDMA etc.

To do this, we could follow the lead of other PyTorch libraries and build and push to the Pytorch index

[colin2328]

Extending this issue to include the plan

Build ‘tensor_engine’ and ‘rdma’ as python ‘extras’ (optional dependencies); e.g enable ‘pip install torchmonarch[tensor_engine, rdma, grad_generator]’. Continue to leverage rust ‘features’ (cfg).

[johnwhumphreys]

Update: some diffs are out to test how things work; but it's still WIP.

[johnwhumphreys]

@allenwang28 to sync with @zdevito on the static linking wrap-up and follow-ups.

[allenwang28]

A few updates:

First - @zdevito's changes for static linking have landed. Theoretically, pip install torchmonarch should work on all x86 machines. For the remaining pieces, I was looking into enabling ARM and on Mac. The goal is that users can just pip install torchmonarch-nightly or pip install torchmonarch (for next stable release) and PyPI can link to the correct architecture.

In practice, this means adding in new wheel building workflows - for aarch64 and Mac - which publish to PyPI, just like we do for x86. This is doable, but the challenge I discovered in my investigations with #2029 that the docker images typically used for these architectures don't come with Conda. We heavily rely on Conda for linking to Python, so we end up running into a several build issues.

IIUC our options are:

1. To add conda into these docker environments, or
2. Change the Rust build such that it no longer has a static Python building. One option is to use the extension-module option of PyO3 - this coincidentally would be the same fix for enabling uv as well.

I plan to explore the feasibility of the extension-module path first.

[allenwang28]

for extension-module: #2148 seems to work locally

[johnwhumphreys]

Zach had landed the static linking for everything. So, for x86 workflows; you should be able to pip install everything fine. For non-cuda torch-monarch, Allen was aiming to add ARM support (GB200 machines, and maybe for MAC). basically need to add workflows per component. Unfortunately, pytorch test infra does not have the same docker images you use for x86; related to extension module things (similar problems to UV). Allen has something that seems to work; focusing on CI before landing.

Zach's PR for the static linking should land before we do the release. A lot of the python installs do not work well without this fix. Zach will try to land it today. @zdevito and @allenwang28 can sync on this today as they are both going at it from different directions.

[allenwang28]

libpython dependency has been removed, so what remains to be done is to is to add ARM and Mac specific workflows. This is needed to enable the pip install torchmonarch route for non x86 users.

I plan to do two things:

1. Modify workflows to depend on uv
2. Add the ARM and Mac workflows to nightly wheel publishing, and extend this to the official wheel routes

I believe the main pain point of this issue to be solved, e.g. requirement to install a specific version of torch, but the issue isn't fully completed. I propose deprioritizing the full resolution of this issue to the next stable release, i.e. that this should not block Monarch v0.2 cc @colin2328

[colin2328]

@allenwang28 lets close this task as we have a notorch whl for x86 now.

let's add new task for 0.3 that includes your plans