From 1d3d3018cda4f828fc47774d8baeb82080eecc21 Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Thu, 9 Jul 2026 01:07:53 +0300 Subject: [PATCH] ROCm HW support --- QUICKSTART.md | 2 +- packages/node/meshnet_node/hardware.py | 43 ++++++++++++++--- packages/node/meshnet_node/model_backend.py | 22 ++++++++- tests/test_node_startup.py | 53 +++++++++++++++++++++ tests/test_real_model_backend.py | 21 ++++++++ 5 files changed, 133 insertions(+), 8 deletions(-) diff --git a/QUICKSTART.md b/QUICKSTART.md index f7cb70a..e686bb8 100644 --- a/QUICKSTART.md +++ b/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 ``` - + 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. diff --git a/packages/node/meshnet_node/hardware.py b/packages/node/meshnet_node/hardware.py index d929b4b..6fdf9c6 100644 --- a/packages/node/meshnet_node/hardware.py +++ b/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() - name = torch.cuda.get_device_name(idx) - props = torch.cuda.get_device_properties(idx) - vram_mb = props.total_memory // (1024 * 1024) + if torch_gpu is None: + torch_gpu = _detect_torch_cuda_inventory(torch) + name = torch_gpu["gpu_name"] if torch_gpu is not None else "CUDA GPU" + 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 diff --git a/packages/node/meshnet_node/model_backend.py b/packages/node/meshnet_node/model_backend.py index ec0249f..063c309 100644 --- a/packages/node/meshnet_node/model_backend.py +++ b/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, diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index 0385f46..d574391 100644 --- a/tests/test_node_startup.py +++ b/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, diff --git a/tests/test_real_model_backend.py b/tests/test_real_model_backend.py index 64640c0..57a6db1 100644 --- a/tests/test_real_model_backend.py +++ b/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()