Shaun ZhouandClaude Opus 5
141cb907f2
移除"训练必须有 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>
113 lines
3.5 KiB
Python
113 lines
3.5 KiB
Python
"""设备解析：cuda / mps / cpu 三后端的选择、AMP 默认值与显存统计。"""
|
||
|
||
import pytest
|
||
|
||
torch = pytest.importorskip("torch")
|
||
|
||
from utils.device import ( # noqa: E402
|
||
BACKEND_PRIORITY,
|
||
DeviceNotAvailableError,
|
||
PeakMemory,
|
||
amp_supported,
|
||
available_backends,
|
||
default_amp,
|
||
describe_device,
|
||
memory_stats,
|
||
pin_memory_supported,
|
||
resolve_amp,
|
||
resolve_device,
|
||
synchronize,
|
||
)
|
||
|
||
|
||
def test_available_backends_is_ordered_and_always_has_cpu():
|
||
backends = available_backends()
|
||
assert backends[-1] == "cpu"
|
||
assert backends == [b for b in BACKEND_PRIORITY if b in backends]
|
||
assert len(set(backends)) == len(backends)
|
||
|
||
|
||
def test_auto_picks_the_highest_priority_available_backend():
|
||
device = resolve_device("auto")
|
||
assert device.type == available_backends()[0]
|
||
assert resolve_device("") == device # 空串等同 auto
|
||
|
||
|
||
def test_cpu_is_always_resolvable():
|
||
assert resolve_device("cpu").type == "cpu"
|
||
|
||
|
||
def test_explicit_backend_matches_hardware_or_raises():
|
||
"""显式请求不可用的后端必须报错，不能悄悄退回 CPU——那会让训练慢几十倍还不吭声。"""
|
||
for backend in ("cuda", "mps"):
|
||
if backend in available_backends():
|
||
assert resolve_device(backend).type == backend
|
||
else:
|
||
with pytest.raises(DeviceNotAvailableError, match=backend):
|
||
resolve_device(backend)
|
||
|
||
|
||
def test_unknown_device_raises():
|
||
with pytest.raises(DeviceNotAvailableError, match="未知设备"):
|
||
resolve_device("tpu")
|
||
|
||
|
||
def test_amp_defaults_follow_the_backend():
|
||
# CUDA 有 tensor core，默认开；MPS 上 autocast 实测更慢，默认关；CPU 不支持
|
||
assert default_amp("cuda") is True
|
||
assert default_amp("mps") is False
|
||
assert default_amp("cpu") is False
|
||
assert amp_supported("cuda") and amp_supported("mps")
|
||
assert not amp_supported("cpu")
|
||
|
||
|
||
def test_resolve_amp_handles_auto_and_explicit():
|
||
assert resolve_amp("auto", "cuda") is True
|
||
assert resolve_amp("auto", "mps") is False
|
||
assert resolve_amp(None, "mps") is False
|
||
assert resolve_amp(True, "mps") is True # 显式打开，MPS 支持 autocast
|
||
assert resolve_amp(True, "cpu") is False # CPU 不支持，显式打开也无效
|
||
assert resolve_amp(False, "cuda") is False
|
||
|
||
|
||
def test_pin_memory_only_for_cuda():
|
||
assert pin_memory_supported("cuda")
|
||
assert not pin_memory_supported("mps")
|
||
assert not pin_memory_supported("cpu")
|
||
|
||
|
||
def test_describe_device_is_readable():
|
||
assert describe_device("cpu") == "CPU"
|
||
assert "MPS" in describe_device("mps")
|
||
|
||
|
||
def test_memory_stats_shape_per_backend():
|
||
assert memory_stats("cpu") == {}
|
||
stats = memory_stats(resolve_device("auto"))
|
||
device_type = resolve_device("auto").type
|
||
if device_type == "cuda":
|
||
assert {"allocated_mb", "reserved_mb", "peak_allocated_mb"} <= set(stats)
|
||
elif device_type == "mps":
|
||
assert {"allocated_mb", "driver_allocated_mb"} <= set(stats)
|
||
|
||
|
||
def test_peak_memory_tracks_growth_on_accelerators():
|
||
device = resolve_device("auto")
|
||
peak = PeakMemory(device)
|
||
|
||
if device.type == "cpu":
|
||
assert peak.peak_mb == 0.0 and peak.reserved_mb == 0.0
|
||
return
|
||
|
||
# MPS 没有峰值 API，靠采样，所以峰值是下界（exact=False）
|
||
assert peak.exact == (device.type != "mps")
|
||
buffer = torch.empty(8 * 1024 * 1024 // 4, device=device) # 约 8 MB
|
||
peak.update()
|
||
synchronize(device)
|
||
assert peak.peak_mb >= 4.0
|
||
del buffer
|
||
|
||
|
||
def test_synchronize_is_a_noop_on_cpu():
|
||
synchronize("cpu") # 不应抛异常
|