Add NPU support for `torchmonarch`

[xiaopyyy]

Summary

torchmonarch currently targets GPU-based execution (e.g., CUDA) and does not support NPU accelerators Issue #1653 and Issue #2296 . Similar to the ongoing discussion in Issue #1649 and Issue #1683 , extending Monarch to additional accelerator backends would broaden its applicability across diverse hardware environments.

Current status

We’ve implemented experimental Ascend NPU support and have successfully run Monarch workloads on Ascend machines:

• Single-node tensor execution
• Multi-node tensor execution

These results indicate that Monarch can be extended to support NPUs without requiring changes to user-facing APIs or application code.

Evidence

Attached screenshots:

1. Monarch tensor execution on a single Ascend machine

2. Monarch multi-machine tensor execution on two Ascend machines

Proposed next steps

Before submitting code, we’d like maintainer guidance on:

• Be upstreamed incrementally via Issues → PRs in torchmonarch, or
• Whether NPU support should live as a separate sub-project / experimental backend

If helpful, we can start with a small, gated PR (e.g., minimal NPU device discovery / backend abstraction) to align on design and review expectations.

Goal

Contribute NPU support in a way that is aligned with Monarch’s architecture, avoids long-lived forks, and expands Monarch’s hardware ecosystem.

[tianyi-ge]

thanks @xiaopyyy. so far we did not find a good abstraction for hardware vendors, so @KaisennHu designed and developed a parallel codebase structure (e.g. torch-sys-ascend and hccl-sys) and control it with compilation and made it work first.

If monarch team is planing to make hardware libraries as plugins, we are happy to get involved in discussing and contributing

[KaisennHu]

By replacing torch-sys-cuda and nccl-sys with torch-sys-ascend and hccl-sys, we now support seamless distributed tensor reduction and multi-node scaling on Ascend. Usage is show as below.

• Distributed tensor reduction example:

import torch
import monarch
from monarch.actor import this_host
 
mesh = this_host().spawn_procs(per_host={"npu": 2})
with mesh.activate():
    a = torch.ones(3, 4, device="npu")
    r = a.reduce("npu", "sum")
 
for i in range(2):
    print(f"result_a on npu-{i}: ", monarch.inspect(a, npu=i))
    print(f"result_r on npu-{i}: ", monarch.inspect(r, npu=i))

• Multi-node scaling example:

import logging
import asyncio
import torch
import os
 
from monarch.actor import Actor, endpoint
from monarch.job import SlurmJob
 
 
logging.basicConfig(
    level=logging.INFO,
    format="%(name)s %(asctime)s %(levelname)s %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    force=True,
)

logger: logging.Logger = logging.getLogger(__name__)

class NPUCompute(Actor):
    def __init__(self, multiplier: float):
        super().__init__()
        self.multiplier = multiplier
 
    @endpoint
    def compute(self, x: torch.Tensor) -> torch.Tensor:
        x = x.npu()
        return x * self.multiplier
 
async def main():
    num_nodes = 2
    npus_per_node = 4
    mesh_name = "mesh0"
    slurm_job = SlurmJob(
        meshes={mesh_name: num_nodes},
        job_name="monarch_example",
        partition="monarch_ascend",
        time_limit="06:00:00",
    )
 
    try:
        job_state = slurm_job.state()
        proc_mesh = job_state.mesh0.spawn_procs({"npus": npus_per_node})
        npu_actors = proc_mesh.spawn("npu_compute", NPUCompute, 2.0)
        input_tensor = torch.ones(3, 4)
        result = await npu_actors.compute.call(input_tensor)
        logger.info(f"succeed to get results: {result}")
        
    finally:
        slurm_job.kill()
        logger.info("Job terminated successfully")
 
 
if __name__ == "__main__":
    asyncio.run(main())