ROCm HW support

QUICKSTART.md

@@ -181,7 +181,7 @@ python -m pip install -e packages/tracker -e packages/node -e packages/p2p -e pa
 python -m pip install "transformers>=5.12" accelerate safetensors
 python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.3
 ```
-
+ <!-- .venv-rocm/bin/pip install --index-url https://rocm.nightlies.amd.com/v2/gfx1151/ --upgrade  -->
 Keep this separate from a known-good CPU `.venv` until ROCm is verified on that
 machine. ROCm wheels are large and host-runtime-sensitive; a failed ROCm install
 should not break the CPU fallback environment.

packages/node/meshnet_node/hardware.py

@@ -123,6 +123,24 @@ def _detect_nvidia_smi_gpu_memory() -> dict | None:
     return None
 
 
+def _detect_torch_cuda_inventory(torch_module) -> dict | None:
+    """Return torch-visible CUDA/HIP GPU metadata without running kernels."""
+    try:
+        if not torch_module.cuda.is_available() or torch_module.cuda.device_count() < 1:
+            return None
+        idx = torch_module.cuda.current_device()
+        name = torch_module.cuda.get_device_name(idx)
+        props = torch_module.cuda.get_device_properties(idx)
+        vram_mb = int(props.total_memory // (1024 * 1024))
+        gpu = {"gpu_name": name, "vram_mb": max(0, vram_mb)}
+        gcn_arch = getattr(props, "gcnArchName", None)
+        if gcn_arch:
+            gpu["gcn_arch"] = str(gcn_arch)
+        return gpu
+    except Exception:
+        return None
+
+
 def _torch_cuda_is_executable(torch_module) -> bool:
     """Return True only if this Python process can execute a CUDA tensor op."""
     try:
@@ -139,7 +157,7 @@ def _torch_cuda_is_executable(torch_module) -> bool:
 def _gpu_inventory_profile(gpu: dict | None, ram_mb: int) -> dict | None:
     if gpu is None:
         return None
-    return {
+    profile = {
         "device": "cpu",
         "gpu_name": gpu["gpu_name"],
         "vram_mb": gpu["vram_mb"],
@@ -148,20 +166,26 @@ def _gpu_inventory_profile(gpu: dict | None, ram_mb: int) -> dict | None:
         "ram_mb": ram_mb,
         "cuda_available": False,
     }
+    if gpu.get("gcn_arch"):
+        profile["gcn_arch"] = gpu["gcn_arch"]
+    return profile
 
 
 def detect_hardware() -> dict:
     """Detect GPU model and available VRAM. Returns hardware profile dict."""
     ram_mb = _detect_ram_mb()
+    torch_gpu: dict | None = None
     try:
         import torch  # type: ignore[import]
+
+        torch_gpu = _detect_torch_cuda_inventory(torch)
         if _torch_cuda_is_executable(torch):
-            idx = torch.cuda.current_device()
+            if torch_gpu is None:
-            name = torch.cuda.get_device_name(idx)
+                torch_gpu = _detect_torch_cuda_inventory(torch)
-            props = torch.cuda.get_device_properties(idx)
+            name = torch_gpu["gpu_name"] if torch_gpu is not None else "CUDA GPU"
-            vram_mb = props.total_memory // (1024 * 1024)
+            vram_mb = torch_gpu["vram_mb"] if torch_gpu is not None else 0
             shared_vram_mb = max(0, ram_mb // 2)
-            return {
+            profile = {
                 "device": "cuda",
                 "gpu_name": name,
                 "vram_mb": vram_mb,
@@ -170,9 +194,16 @@ def detect_hardware() -> dict:
                 "ram_mb": ram_mb,
                 "cuda_available": True,
             }
+            if torch_gpu is not None and torch_gpu.get("gcn_arch"):
+                profile["gcn_arch"] = torch_gpu["gcn_arch"]
+            return profile
     except ImportError:
         pass
 
+    torch_inventory = _gpu_inventory_profile(torch_gpu, ram_mb)
+    if torch_inventory is not None:
+        return torch_inventory
+
     nvidia_gpu = _gpu_inventory_profile(_detect_nvidia_smi_gpu_memory(), ram_mb)
     if nvidia_gpu is not None:
         return nvidia_gpu

packages/node/meshnet_node/model_backend.py

@@ -39,6 +39,26 @@ class KVCacheMiss(ModelBackendError):
     """
 
 
+def _torch_cuda_is_executable(torch_module: Any) -> bool:
+    """Return True only when this process can actually execute a CUDA/HIP op.
+
+    On ROCm, ``torch.cuda.is_available()`` can be true for an AMD GPU even when
+    the installed PyTorch wheel has no runnable kernels for that GPU target.
+    Loading weights onto such a device can segfault in native code, so the model
+    backend must use the same executable-device check as startup hardware
+    detection rather than trusting inventory alone.
+    """
+    try:
+        if not torch_module.cuda.is_available():
+            return False
+        probe = torch_module.empty((1,), device="cuda")
+        probe += 1
+        torch_module.cuda.synchronize()
+        return True
+    except Exception:
+        return False
+
+
 @dataclass(frozen=True)
 class TensorPayload:
     body: bytes
@@ -209,7 +229,7 @@ class TorchModelShard:
             ) from exc
 
         self.torch = torch
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device("cuda" if _torch_cuda_is_executable(torch) else "cpu")
         load_source = str(cache_dir) if cache_dir is not None and (cache_dir / "config.json").exists() else model_id
         quant_config, dtype, uses_quantized_weights = _model_load_plan(
             AutoConfig,

tests/test_node_startup.py

@@ -128,6 +128,59 @@ def test_nvidia_smi_without_torch_cuda_keeps_cpu_execution(monkeypatch):
     assert hw["ram_mb"] == 80 * 1024
 
 
+def test_torch_rocm_inventory_is_reported_when_kernels_are_not_executable(monkeypatch):
+    """ROCm can expose GPU metadata even when this torch wheel cannot run kernels."""
+    import meshnet_node.hardware as hardware_mod
+
+    class FakeProps:
+        total_memory = 64 * 1024 * 1024 * 1024
+        gcnArchName = "gfx1151"
+
+    class FakeCuda:
+        @staticmethod
+        def is_available():
+            return True
+
+        @staticmethod
+        def device_count():
+            return 1
+
+        @staticmethod
+        def current_device():
+            return 0
+
+        @staticmethod
+        def get_device_name(_idx):
+            return "AMD Radeon 8060S"
+
+        @staticmethod
+        def get_device_properties(_idx):
+            return FakeProps()
+
+        @staticmethod
+        def synchronize():
+            raise AssertionError("synchronize should not run after empty() fails")
+
+    fake_torch = types.SimpleNamespace(
+        cuda=FakeCuda(),
+        empty=lambda *args, **kwargs: (_ for _ in ()).throw(
+            RuntimeError("HIP error: invalid device function")
+        ),
+    )
+
+    monkeypatch.setattr(hardware_mod, "_detect_ram_mb", lambda: 125 * 1024)
+    monkeypatch.setitem(sys.modules, "torch", fake_torch)
+
+    hw = hardware_mod.detect_hardware()
+
+    assert hw["device"] == "cpu"
+    assert hw["cuda_available"] is False
+    assert hw["gpu_name"] == "AMD Radeon 8060S"
+    assert hw["vram_mb"] == 64 * 1024
+    assert hw["shared_vram_mb"] == 64_000
+    assert hw["gcn_arch"] == "gfx1151"
+
+
 def test_memory_budget_uses_ram_for_cpu_and_shared_memory_for_cuda():
     assert _memory_budget("cpu", vram_mb=8192, ram_mb=80 * 1024, shared_vram_mb=40 * 1024) == (
         80 * 1024,

tests/test_real_model_backend.py

@@ -24,6 +24,7 @@ from meshnet_node.model_backend import (
     _should_partial_materialize_shard,
     _decoder_attention_mask,
     _int_tensor_header,
+    _torch_cuda_is_executable,
     build_quantization_config,
     validate_quantization,
 )
@@ -209,6 +210,26 @@ def test_bitsandbytes_configs_are_created_lazily(monkeypatch):
     ]
 
 
+def test_rocm_inventory_without_executable_kernels_is_not_used_as_cuda():
+    class FakeCuda:
+        @staticmethod
+        def is_available():
+            return True
+
+        @staticmethod
+        def synchronize():
+            raise AssertionError("synchronize should not run after empty() fails")
+
+    fake_torch = types.SimpleNamespace(
+        cuda=FakeCuda(),
+        empty=lambda *args, **kwargs: (_ for _ in ()).throw(
+            RuntimeError("HIP error: invalid device function")
+        ),
+    )
+
+    assert _torch_cuda_is_executable(fake_torch) is False
+
+
 def test_head_forward_accepts_text_prompt_and_returns_bfloat16_activations():
     node = TorchNodeServer(backend=_FakeBackend())
     port = node.start()