feat: 支持 cuda / mps / cpu 三后端，Mac 上可用 MPS 训练
移除"训练必须有 Nvidia GPU"的硬性守卫，改为按后端能力分派。
- utils/device: resolve_device("auto") 按 cuda > mps > cpu 挑最快的可用后端；
显式指定不可用的后端时直接报错，不会悄悄退回 CPU——训练慢几十倍却不吭声是最坏的情况
- AMP 随后端取默认值：CUDA 开、MPS 关。MPS 上 autocast 数值正常但实测更慢
（tiny_conv @128²：fp32 76 ms/步、fp16 90 ms、bf16 108 ms），需要时 amp: true 可强开
- PeakMemory 统一显存峰值统计：CUDA 读内置计数器，MPS 无峰值 API 改用逐 step 采样
（是下界，日志与产出里已标注）
- profile_inference 支持 MPS，并提示 torch.profiler 没有 MPS 的 ProfilerActivity，
算子明细只反映 CPU 侧下发耗时
- train.py / eval.py 新增 --device，配置默认改为 auto
- check_env 改为列出全部可用后端与 auto 的选择结果
实测：tiny_conv 训练在 MPS 上比 CPU 快约 32 倍。
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
4fc0fa5353
commit
141cb907f2
+174
-39
@@ -1,54 +1,189 @@
|
||||
"""设备检查。
|
||||
"""设备解析与能力查询：统一处理 cuda / mps / cpu 三种后端。
|
||||
|
||||
项目硬性约定（见 CLAUDE.MD）：禁止在没有 Nvidia GPU 的设备上训练。
|
||||
所有训练入口必须先调用 `require_cuda()`，检查不通过直接抛错退出。
|
||||
设计要点：
|
||||
- `resolve_device("auto")` 按 cuda > mps > cpu 的优先级自动挑最快的后端；
|
||||
- 显式指定但不可用时直接报错，而不是悄悄退回 CPU——训练慢 30 倍却不吭声是最坏的情况；
|
||||
- AMP 的默认值随后端而变：CUDA 上开（有 tensor core），MPS 上关
|
||||
（实测 tiny_conv @128² 每步 fp32 76 ms、fp16 90 ms，autocast 只剩转换开销反而更慢）；
|
||||
- 显存统计跨后端统一：CUDA 有内置峰值计数器，MPS 没有，靠采样取最大值。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
class NoNvidiaGpuError(RuntimeError):
|
||||
"""当前设备不具备可用的 Nvidia GPU。"""
|
||||
#: 自动选择时的优先级
|
||||
BACKEND_PRIORITY = ("cuda", "mps", "cpu")
|
||||
|
||||
|
||||
def cuda_devices() -> list[str]:
|
||||
"""返回可用的 Nvidia GPU 名称列表；没有 torch 或没有 GPU 时返回空列表。"""
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
return []
|
||||
if not torch.cuda.is_available() or torch.version.cuda is None:
|
||||
return []
|
||||
return [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]
|
||||
class DeviceNotAvailableError(RuntimeError):
|
||||
"""请求的设备在当前机器上不可用。"""
|
||||
|
||||
|
||||
def require_cuda() -> list[str]:
|
||||
"""训练前的守卫；无 Nvidia GPU 时抛出 `NoNvidiaGpuError`。"""
|
||||
try:
|
||||
import torch
|
||||
except ImportError as exc:
|
||||
raise NoNvidiaGpuError("未安装 PyTorch，无法确认 Nvidia GPU，训练已中止。") from exc
|
||||
|
||||
if torch.version.cuda is None:
|
||||
raise NoNvidiaGpuError(
|
||||
f"当前 PyTorch ({torch.__version__}) 不是 CUDA 构建（可能是 CPU/ROCm/MPS 版本），"
|
||||
"本项目禁止在非 Nvidia GPU 设备上训练。"
|
||||
)
|
||||
if not torch.cuda.is_available() or torch.cuda.device_count() == 0:
|
||||
raise NoNvidiaGpuError(
|
||||
"未检测到可用的 Nvidia GPU，本项目禁止在无 GPU 的设备上训练。"
|
||||
"评估/推理请使用 eval.py 并显式指定 device=cpu。"
|
||||
)
|
||||
return cuda_devices()
|
||||
def _cuda_ready() -> bool:
|
||||
return torch.cuda.is_available() and torch.cuda.device_count() > 0
|
||||
|
||||
|
||||
def resolve_device(device: str = "cuda"):
|
||||
"""把配置中的 device 字符串解析为 `torch.device`。
|
||||
def _mps_ready() -> bool:
|
||||
return torch.backends.mps.is_available() and torch.backends.mps.is_built()
|
||||
|
||||
仅用于评估/推理等允许 CPU 的场景；训练路径请先走 `require_cuda()`。
|
||||
|
||||
def available_backends() -> list[str]:
|
||||
"""按优先级返回当前可用的后端，末尾必定有 cpu。"""
|
||||
backends = []
|
||||
if _cuda_ready():
|
||||
backends.append("cuda")
|
||||
if _mps_ready():
|
||||
backends.append("mps")
|
||||
backends.append("cpu")
|
||||
return backends
|
||||
|
||||
|
||||
def resolve_device(spec: str = "auto") -> torch.device:
|
||||
"""把配置里的 device 字符串解析成 `torch.device`。
|
||||
|
||||
Args:
|
||||
spec: "auto" 自动挑最优后端；也可写 "cuda" / "cuda:1" / "mps" / "cpu"。
|
||||
|
||||
Raises:
|
||||
DeviceNotAvailableError: 显式请求的后端不可用。
|
||||
"""
|
||||
import torch
|
||||
spec = str(spec).strip().lower()
|
||||
if spec in ("auto", ""):
|
||||
return torch.device(available_backends()[0])
|
||||
|
||||
if device.startswith("cuda") and not torch.cuda.is_available():
|
||||
raise NoNvidiaGpuError(f"请求了 {device}，但当前设备没有可用的 Nvidia GPU。")
|
||||
return torch.device(device)
|
||||
backend = spec.split(":", 1)[0]
|
||||
if backend == "cuda" and not _cuda_ready():
|
||||
raise DeviceNotAvailableError(
|
||||
f"请求了 {spec}，但当前没有可用的 Nvidia GPU"
|
||||
f"（PyTorch {torch.__version__}，CUDA 构建：{torch.version.cuda}）。"
|
||||
f"可用后端：{available_backends()}，或用 device=auto 自动选择。"
|
||||
)
|
||||
if backend == "mps" and not _mps_ready():
|
||||
raise DeviceNotAvailableError(
|
||||
f"请求了 {spec}，但当前 PyTorch 不支持 MPS"
|
||||
f"（available={torch.backends.mps.is_available()}, built={torch.backends.mps.is_built()}）。"
|
||||
f"可用后端：{available_backends()}。"
|
||||
)
|
||||
if backend not in ("cuda", "mps", "cpu"):
|
||||
raise DeviceNotAvailableError(f"未知设备 {spec}，支持 cuda / mps / cpu / auto。")
|
||||
return torch.device(spec)
|
||||
|
||||
|
||||
def describe_device(device: torch.device | str) -> str:
|
||||
"""给日志用的人类可读描述。"""
|
||||
device = torch.device(device)
|
||||
if device.type == "cuda":
|
||||
index = device.index or 0
|
||||
return f"{torch.cuda.get_device_name(index)} (cuda:{index})"
|
||||
if device.type == "mps":
|
||||
return "Apple Silicon GPU (Metal / MPS)"
|
||||
return "CPU"
|
||||
|
||||
|
||||
def amp_supported(device: torch.device | str) -> bool:
|
||||
"""该后端是否支持 autocast 混合精度。"""
|
||||
return torch.device(device).type in ("cuda", "mps")
|
||||
|
||||
|
||||
def default_amp(device: torch.device | str) -> bool:
|
||||
"""AMP 的默认开关：只有 CUDA 默认开。
|
||||
|
||||
MPS 支持 autocast 且数值正常，但实测比 fp32 更慢，所以默认关闭；
|
||||
需要省显存时可以在配置里显式 `amp: true` 打开。
|
||||
"""
|
||||
return torch.device(device).type == "cuda"
|
||||
|
||||
|
||||
def resolve_amp(setting: bool | str | None, device: torch.device | str) -> bool:
|
||||
"""解析配置里的 `amp` 字段：`auto`/`None` 随后端，`true`/`false` 显式指定。"""
|
||||
device = torch.device(device)
|
||||
if setting is None or (isinstance(setting, str) and setting.lower() == "auto"):
|
||||
return default_amp(device)
|
||||
if not bool(setting):
|
||||
return False
|
||||
return amp_supported(device)
|
||||
|
||||
|
||||
def pin_memory_supported(device: torch.device | str) -> bool:
|
||||
"""只有 CUDA 的锁页内存有意义；MPS 上开启会告警且无收益。"""
|
||||
return torch.device(device).type == "cuda"
|
||||
|
||||
|
||||
def synchronize(device: torch.device | str) -> None:
|
||||
"""等待该设备上的异步算子完成——计时前后必须调用，否则测到的是下发时间。"""
|
||||
device = torch.device(device)
|
||||
if device.type == "cuda":
|
||||
torch.cuda.synchronize(device)
|
||||
elif device.type == "mps":
|
||||
torch.mps.synchronize()
|
||||
|
||||
|
||||
def memory_stats(device: torch.device | str) -> dict[str, float]:
|
||||
"""统一的显存查询，单位 MB。
|
||||
|
||||
CUDA 返回 allocated / reserved / peak；MPS 只有当前分配量与驱动占用
|
||||
（Metal 没有暴露峰值计数器）；CPU 返回空字典。
|
||||
"""
|
||||
device = torch.device(device)
|
||||
mb = 1024.0 * 1024.0
|
||||
if device.type == "cuda":
|
||||
return {
|
||||
"allocated_mb": torch.cuda.memory_allocated(device) / mb,
|
||||
"reserved_mb": torch.cuda.memory_reserved(device) / mb,
|
||||
"peak_allocated_mb": torch.cuda.max_memory_allocated(device) / mb,
|
||||
"peak_reserved_mb": torch.cuda.max_memory_reserved(device) / mb,
|
||||
}
|
||||
if device.type == "mps":
|
||||
return {
|
||||
"allocated_mb": torch.mps.current_allocated_memory() / mb,
|
||||
"driver_allocated_mb": torch.mps.driver_allocated_memory() / mb,
|
||||
}
|
||||
return {}
|
||||
|
||||
|
||||
class PeakMemory:
|
||||
"""跨后端的显存峰值统计。
|
||||
|
||||
CUDA 有内置峰值计数器，直接读；MPS 没有，只能在 `update()` 时采样取最大值——
|
||||
因此 MPS 上的峰值是**下界**（采样点之间的瞬时尖峰会被漏掉）。
|
||||
"""
|
||||
|
||||
def __init__(self, device: torch.device | str) -> None:
|
||||
self.device = torch.device(device)
|
||||
self._peak = 0.0
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
self._peak = 0.0
|
||||
if self.device.type == "cuda":
|
||||
torch.cuda.reset_peak_memory_stats(self.device)
|
||||
self.update()
|
||||
|
||||
def update(self) -> None:
|
||||
if self.device.type != "mps":
|
||||
return
|
||||
self._peak = max(self._peak, torch.mps.current_allocated_memory() / (1024.0 * 1024.0))
|
||||
|
||||
@property
|
||||
def peak_mb(self) -> float:
|
||||
if self.device.type == "cuda":
|
||||
return torch.cuda.max_memory_allocated(self.device) / (1024.0 * 1024.0)
|
||||
if self.device.type == "mps":
|
||||
self.update()
|
||||
return self._peak
|
||||
return 0.0
|
||||
|
||||
@property
|
||||
def reserved_mb(self) -> float:
|
||||
"""CUDA 是缓存分配器保留量；MPS 是 Metal 驱动占用；CPU 为 0。"""
|
||||
if self.device.type == "cuda":
|
||||
return torch.cuda.max_memory_reserved(self.device) / (1024.0 * 1024.0)
|
||||
if self.device.type == "mps":
|
||||
return torch.mps.driver_allocated_memory() / (1024.0 * 1024.0)
|
||||
return 0.0
|
||||
|
||||
@property
|
||||
def exact(self) -> bool:
|
||||
"""峰值是否精确（MPS 靠采样，只能给下界）。"""
|
||||
return self.device.type != "mps"
|
||||
|
||||
Reference in New Issue
Block a user