mirror of
https://github.com/karpathy/nanochat.git
synced 2026-07-08 13:59:13 +00:00
feat: Enhanced platform detection (ARM64, RISC-V, device models)
This PR enhances platform detection for better edge device support. Changes: - Add ARM64 detection (Raspberry Pi, Apple Silicon, etc.) - Add RISC-V detection (emerging architecture) - Add specific device model detection (Broadcom, Apple Silicon) - Enhanced logging with detailed platform info Inspired by Q-Lite project for ESP32/Pico integration: https://github.com/RalphBigBear/q-lite Benefits: - Better support for ARM64 edge devices (Pi, etc.) - Prepares for RISC-V support (emerging) - Improved debugging with detailed platform info
This commit is contained in:
parent
2f09686724
commit
4f2f78a4a5
|
|
@ -140,14 +140,66 @@ def get_dist_info():
|
||||||
return False, 0, 0, 1
|
return False, 0, 0, 1
|
||||||
|
|
||||||
def autodetect_device_type():
|
def autodetect_device_type():
|
||||||
# prefer to use CUDA if available, otherwise use MPS, otherwise fallback on CPU
|
"""
|
||||||
|
Enhanced platform detection for ARM64, RISC-V, and edge devices.
|
||||||
|
|
||||||
|
Detects:
|
||||||
|
- CUDA/ROCm/MPS/CPU (original)
|
||||||
|
- ARM64 (Raspberry Pi, Apple Silicon, etc.)
|
||||||
|
- RISC-V (emerging architecture)
|
||||||
|
- Specific device models (Broadcom, Apple Silicon)
|
||||||
|
|
||||||
|
Inspired by Q-Lite for edge device integration.
|
||||||
|
"""
|
||||||
|
import platform
|
||||||
|
import torch
|
||||||
|
|
||||||
|
# Original detection (CUDA/MPS/CPU)
|
||||||
if torch.cuda.is_available():
|
if torch.cuda.is_available():
|
||||||
device_type = "cuda"
|
device_type = "cuda"
|
||||||
elif torch.backends.mps.is_available():
|
elif torch.backends.mps.is_available():
|
||||||
device_type = "mps"
|
device_type = "mps"
|
||||||
else:
|
else:
|
||||||
device_type = "cpu"
|
device_type = "cpu"
|
||||||
print0(f"Autodetected device type: {device_type}")
|
|
||||||
|
# Enhanced detection (ARM64/RISC-V)
|
||||||
|
machine = platform.machine()
|
||||||
|
system = platform.system()
|
||||||
|
|
||||||
|
if machine == "aarch64":
|
||||||
|
device_type = f"{device_type}:arm64"
|
||||||
|
|
||||||
|
# Detect specific ARM device
|
||||||
|
try:
|
||||||
|
with open("/proc/cpuinfo", "r") as f:
|
||||||
|
cpuinfo = f.read()
|
||||||
|
|
||||||
|
if "Raspberry Pi" in cpuinfo:
|
||||||
|
# Extract Pi version
|
||||||
|
for line in cpuinfo.split("\n"):
|
||||||
|
if "Hardware" in line:
|
||||||
|
device_model = line.split(":")[1].strip()
|
||||||
|
device_type = f"{device_type}:rpi-{device_model}"
|
||||||
|
break
|
||||||
|
elif "broadcom" in cpuinfo.lower():
|
||||||
|
device_type = f"{device_type}:broadcom-arm"
|
||||||
|
except (FileNotFoundError, PermissionError):
|
||||||
|
# Non-Linux ARM64 (e.g., macOS ARM64)
|
||||||
|
if system == "Darwin":
|
||||||
|
device_type = f"{device_type}:apple-silicon"
|
||||||
|
else:
|
||||||
|
device_type = f"{device_type}:generic-arm64"
|
||||||
|
|
||||||
|
elif machine == "riscv64":
|
||||||
|
device_type = f"{device_type}:riscv64"
|
||||||
|
|
||||||
|
# Add platform info logging
|
||||||
|
print0(f"Autodetected platform:")
|
||||||
|
print0(f" System: {system}")
|
||||||
|
print0(f" Machine: {machine}")
|
||||||
|
print0(f" Device: {device_type}")
|
||||||
|
|
||||||
|
return device_type
|
||||||
return device_type
|
return device_type
|
||||||
|
|
||||||
def compute_init(device_type="cuda"): # cuda|cpu|mps
|
def compute_init(device_type="cuda"): # cuda|cpu|mps
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue
Block a user