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>
57 lines
1.6 KiB
Python
Executable File
57 lines
1.6 KiB
Python
Executable File
#!/usr/bin/env python
|
|
"""环境自检：确认依赖与可用的计算后端。
|
|
|
|
python scripts/check_env.py
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from utils.device import ( # noqa: E402
|
|
available_backends,
|
|
default_amp,
|
|
describe_device,
|
|
memory_stats,
|
|
resolve_device,
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
print(f"Python: {sys.version.split()[0]}")
|
|
try:
|
|
import torch
|
|
except ImportError:
|
|
print("PyTorch: 未安装 —— 先执行 `uv sync`")
|
|
return 1
|
|
|
|
print(f"PyTorch: {torch.__version__} (CUDA 构建: {torch.version.cuda})")
|
|
print(f"可用后端（按优先级）: {available_backends()}")
|
|
|
|
print("\n后端明细:")
|
|
print(f" cuda: available={torch.cuda.is_available()} device_count={torch.cuda.device_count()}")
|
|
print(
|
|
f" mps : available={torch.backends.mps.is_available()} "
|
|
f"built={torch.backends.mps.is_built()}"
|
|
)
|
|
|
|
device = resolve_device("auto")
|
|
print(f"\ndevice=auto 会选中: {describe_device(device)}")
|
|
print(f" AMP 默认: {'开' if default_amp(device) else '关'}")
|
|
stats = memory_stats(device)
|
|
if stats:
|
|
print(" " + " ".join(f"{k}={v:.0f}" for k, v in stats.items()))
|
|
|
|
if device.type == "cpu":
|
|
print("\n⚠️ 只有 CPU 可用，训练会非常慢（实测比 MPS 慢约 30 倍），建议换有 GPU 的机器。")
|
|
return 1
|
|
print("\n训练环境就绪。")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|