exo + tinygrad：Linux 节点设备能力自动探测（NVIDIA / AMD / CPU 安全兜底）

在使用exo构建本地或分布式推理集群时，每个节点的设备能力探测是非常关键的一步。
exo 需要知道：

• 当前节点使用CPU 还是 GPU
• GPU 型号、显存大小
• 节点的大致算力能力（FLOPS）

本文结合tinygrad，介绍一套稳定、可落地、工程级安全的 Linux 设备能力探测方案，适用于：

• exo worker / node
• 本地推理
• Docker / 裸机 / 无 GPU 环境

一、为什么 exo 需要设备能力探测？

在 exo 中，每个节点都会向调度层上报自身能力，例如：

• 是否支持 CUDA
• 显存大小是否满足模型加载
• 是否能承担 FP16 / INT8 推理

如果设备探测逻辑不健壮，可能会导致：

• ❌ 节点启动即崩溃
• ❌ GPU 探测失败导致服务不可用
• ❌ Docker / 云环境下无法运行

设计原则只有一个：

GPU 是加速项，CPU 是兜底项
任何情况下，exo 节点都必须能启动

二、整体设计思路（exo 场景）

本方案基于tinygrad 的真实执行后端：

fromtinygradimportDevice Device.DEFAULT

而不是单纯扫描系统硬件。

探测优先级

1. NVIDIA GPU（CUDA / NV / GPU）
2. AMD GPU
3. CPU（最终兜底）

任何 GPU 探测失败，立即回退 CPU，不抛异常。

三、完整代码（exo 节点可直接使用）

git clone https://github.com/exo-explore/exo.git cd exo pip install -e .

启动不起来 搜索linux_device_capabilities进行修改

#vim /home/michah/exo/exo/topology/device_capabilities.py async def linux_device_capabilities() -> DeviceCapabilities: import psutil from tinygrad import Device if DEBUG >= 2: print(f"tinygrad {Device.DEFAULT=}") # ----------------------------- # NVIDIA / CUDA 路径（安全版） # ----------------------------- if Device.DEFAULT in ("CUDA", "NV", "GPU"): try: import pynvml pynvml.nvmlInit() handle = pynvml.nvmlDeviceGetHandleByIndex(0) gpu_raw_name = pynvml.nvmlDeviceGetName(handle).upper() gpu_name = ( gpu_raw_name.rsplit(" ", 1)[0] if gpu_raw_name.endswith("GB") else gpu_raw_name ) gpu_memory_info = pynvml.nvmlDeviceGetMemoryInfo(handle) if DEBUG >= 2: print(f"NVIDIA device {gpu_name=} {gpu_memory_info=}") pynvml.nvmlShutdown() return DeviceCapabilities( model=f"Linux Box ({gpu_name})", chip=gpu_name, memory=gpu_memory_info.total // 2**20, flops=CHIP_FLOPS.get( gpu_name, DeviceFlops(fp32=0, fp16=0, int8=0), ), ) except Exception as e: # 关键：任何 NVML / 驱动 / import 错误都回退 CPU if DEBUG >= 1: print(f"[WARN] NVIDIA GPU detection failed, fallback to CPU: {e}") # ----------------------------- # AMD 路径（原样保留） # ----------------------------- if Device.DEFAULT == "AMD": try: import pyamdgpuinfo gpu_raw_info = pyamdgpuinfo.get_gpu(0) gpu_name = gpu_raw_info.name gpu_memory_info = gpu_raw_info.memory_info["vram_size"] if DEBUG >= 2: print(f"AMD device {gpu_name=} {gpu_memory_info=}") return DeviceCapabilities( model=f"Linux Box ({gpu_name})", chip=gpu_name, memory=gpu_memory_info // 2**20, flops=CHIP_FLOPS.get( gpu_name, DeviceFlops(fp32=0, fp16=0, int8=0), ), ) except Exception as e: if DEBUG >= 1: print(f"[WARN] AMD GPU detection failed, fallback to CPU: {e}") # ----------------------------- # CPU / 兜底路径 # ----------------------------- return DeviceCapabilities( model=f"Linux Box (Device: {Device.DEFAULT})", chip=f"Unknown Chip (Device: {Device.DEFAULT})", memory=psutil.virtual_memory().total // 2**20, flops=DeviceFlops(fp32=0, fp16=0, int8=0), )

四、关键实现解析（exo 视角）

1️⃣ 为什么使用Device.DEFAULT

Device.DEFAULT

代表tinygrad 实际使用的后端，而不是系统“可能存在”的设备。

这对 exo 非常重要：

exo 调度依据的是「真实可执行能力」，不是硬件清单

2️⃣ NVIDIA GPU 探测为什么必须 try / except

常见失败场景：

• Docker 容器无/dev/nvidia*
• 没装 NVIDIA Driver
• pynvml 安装了但 NVML 初始化失败
• MIG / 权限问题

因此：

exceptException:fallback to CPU

这是 exo 节点稳定运行的关键。

3️⃣ GPU 名称清洗的意义

"NVIDIA RTX 3090 24GB"→"NVIDIA RTX 3090"

目的：

• 避免CHIP_FLOPS查表失败
• 统一 exo 节点能力上报口径

4️⃣ CPU 兜底是 exo 的生命线

psutil.virtual_memory()

无论：

• 裸机
• 云主机
• CI
• Docker

exo 节点都必须可启动、可注册、可上报。

五、exo 中的典型使用场景

• worker 启动时自动上报能力
• 调度前判断模型是否可加载
• 异构节点（CPU / GPU 混合）统一管理
• GPU 节点异常时自动降级