编写 Python 脚本快速诊断 AMD GPU 健康状态

为什么需要程序化的 GPU 健康检查

在 AMD Instinct GPU 上部署大模型推理服务时，很多开发者习惯依赖rocm-smi或rocminfo等命令行工具来确认环境状态。这些工具虽然直观，但在自动化运维流程或容器化部署场景中显得力不从心。当我们需要在 CI/CD 流水线中自动判断节点可用性，或者在大规模集群启动前快速筛选健康实例时，手动执行命令并肉眼核对输出不仅效率低下，还容易因人为疏忽漏掉关键隐患。

更棘手的是，命令行工具往往只能反映驱动层面的状态，难以直接揭示深度学习框架（如 PyTorch）是否能真正调用到硬件。有时候rocm-smi显示正常，但 Python 脚本运行时却报“未找到 CUDA 设备”，这种割裂感通常源于环境变量配置错误或用户权限问题。为了解决这一痛点，我们需要一种更贴近应用层的诊断方式——直接利用 PyTorch 接口编写脚本，在代码层面完成从设备可见性、显存容量到核心计算特性（如 BF16 支持）的全方位体检。

构建 health_check.py 核心逻辑

编写诊断脚本的核心思路是绕过繁琐的系统命令，直接通过 PyTorch 的后端接口与 GPU 对话。在 ROCm 环境下，PyTorch 依然沿用cuda作为后端别名，这使得代码具有较好的通用性。我们需要重点关注三个维度：设备是否可见、显存资源是否充足、以及是否支持关键的 BF16 数据类型。

首先，脚本必须能够捕获设备不可见的异常情况。这通常是因为HIP_VISIBLE_DEVICES环境变量被错误限制，或者当前用户缺乏访问/dev/kfd和/dev/dri的权限。其次，显存信息的获取不能仅看总量，可用显存才是决定能否加载大模型的关键。最后，针对 AMD Instinct MI300 等新架构显卡，BF16（Brain Floating Point 16）的支持情况直接影响推理精度与性能，必须在启动前予以确认。

以下是一个经过生产环境验证的health_check.py脚本示例，它封装了上述所有检查逻辑，并提供了清晰的错误指引：

importtorchimportsysimportosdefcheck_rocm_health():# 第一步：检查后端可用性# 在 ROCm 环境中，torch.cuda.is_available() 同样用于检测 HIP 后端ifnottorch.cuda.is_available():print("❌ 严重错误：未检测到可用的 ROCm 设备")print(" 可能原因：")print(" 1. 驱动未正确安装或版本不匹配")print(" 2. 当前用户不在 video/render 用户组")print(" 3. HIP_VISIBLE_DEVICES 环境变量限制了设备可见性")print(" 建议操作：运行 'rocm-smi' 确认驱动状态，并检查用户组权限。")returnFalse# 第二步：获取设备数量与详细信息device_count=torch.cuda.device_count()print(f"✅ 成功检测到{device_count}个加速卡")all_checks_passed=Trueforiinrange(device_count):print(f"\n--- 正在诊断卡{i}---")try:props=torch.cuda.get_device_properties(i)free_mem,total_mem_bytes=torch.cuda.mem_get_info(i)# 单位换算为 GBtotal_mem_gb=total_mem_bytes/1024**3free_mem_gb=free_mem/1024**3print(f" 设备名称：{props.name}")print(f" 显存总量：{total_mem_gb:.2f}GB")print(f" 可用显存：{free_mem_gb:.2f}GB ({(free_mem_gb/total_mem_gb)*100:.1f}%)")# 第三步：关键特性检查 (BF16 支持)# AMD Instinct MI200/MI300 系列通常 major >= 9ifprops.major>=9:print(" ✅ 支持 BF16 加速 (推荐用于大模型推理)")else:print(" ⚠️ 未检测到原生 BF16 支持，可能需要降级使用 FP16")# 显存预警iffree_mem_gb<2.0:print(" ⚠️ 警告：可用显存过低，可能无法加载大型模型")exceptExceptionase:print(f" ❌ 读取卡{i}信息失败：{str(e)}")all_checks_passed=Falsereturnall_checks_passedif__name__=="__main__":print("🚀 开始 AMD GPU 健康自检...")ifcheck_rocm_health():print("\n🎉 环境健康检查通过，准备启动服务")sys.exit(0)else:print("\n💥 检查失败，请根据上述提示修复环境")sys.exit(1)

深度解析 BF16 支持与显存诊断

在上述脚本中，对 BF16 支持的判断逻辑尤为关键。AMD 的新一代 Instinct GPU（如 MI250X, MI300X）引入了强大的 BF16 计算单元，这对于运行 Llama 3、Qwen 等大语言模型至关重要。BF16 相比 FP16 拥有更大的动态范围，能有效避免梯度溢出，同时在推理阶段能显著降低显存占用并提升吞吐量。

脚本通过props.major属性来间接判断架构代际。通常情况下，计算能力主版本号大于等于 9 的设备（对应 GFX90A 及更新架构）都原生支持 BF16。如果脚本输出警告信息，意味着当前硬件可能较旧，或者驱动未能正确识别新特性。在这种情况下，强行启用 BF16 可能导致推理结果异常甚至崩溃，开发者需在 vLLM 启动参数中显式指定--dtype float16进行兼容。

显存诊断部分则采用了torch.cuda.mem_get_info()接口。与查看静态总量不同，实时获取“可用显存”能帮助我们预判 OOM（内存溢出）风险。例如，若一张 80GB 的显卡仅剩 2GB 可用，即便设备可见，也无法加载 7B 以上的模型。脚本中的阈值判断（如小于 2GB 报警）是一个可配置的启发式规则，实际生产中可根据目标模型的大小动态调整这一阈值。

集成到自动化运维流程

将health_check.py纳入自动化运维体系，能极大提升部署效率。在 Kubernetes 或 Slurm 集群中，可以将此脚本作为 Pod 的initContainer或作业的前置钩子（Pre-job Hook）。只有当脚本退出码为 0 时，主容器才会启动或作业才会被调度。这种“故障左移”的策略，避免了因个别节点驱动异常导致整个推理服务启动失败或运行中突然崩溃。

此外，该脚本也可嵌入 CI/CD 流水线。每当基础镜像更新或驱动升级时，自动运行健康检查，确保新的软件栈与底层硬件完美契合。对于多卡环境，脚本会遍历所有设备，任何一张卡的异常都会导致整体检查失败，从而防止出现“部分可用”的亚健康状态，保障生产环境的稳定性。

通过这种程序化的诊断方式，我们将原本依赖经验的“玄学”排查转化为标准化的代码逻辑。这不仅减少了人工干预成本，更为大规模 AMD GPU 集群的稳健运行筑起了一道坚实防线。在正式拉起 vLLM 服务之前，花几秒钟运行一次这个脚本，往往能节省数小时的故障定位时间。

200小时GPU算力已就位，快来领取：https://marketing.csdn.net/questions/Q2604140858304426315?utm_source=AIpaper