前言

截至目前（2026 年），海光（Hygon）GPU 主要基于 DCU（Deep Computing Unit）架构，其硬件兼容 AMD 的 CDNA 架构，并通过 ROCm（Radeon Open Compute）平台 提供 AI 加速支持。因此，在海光 GPU 上运行 AI 模型的核心路径是：使用 ROCm 生态 + 兼容框架。

ROCm 是 AMD 为自家 GPU（包括 Instinct 加速卡、Radeon 专业卡和部分消费级显卡）打造的 GPU 计算生态系统，对标 NVIDIA 的 CUDA。

PyTorch、TensorFlow、ONNX Runtime 等均有官方或社区 ROCm 适配版

前提条件

确认你的海光 GPU 型号和驱动

海光 DCU（如 Z100、K100 等）需运行在 国产操作系统（如麒麟、统信 UOS、中科方德）或 特定 Linux 发行版（如 CentOS/RHEL 7/8）。
必须安装 海光官方提供的 ROCm 兼容驱动栈（通常由服务器厂商或海光提供，非公开下载）。

验证命令（需 root 或 sudo）：

1 2	rocminfo # 查看 ROCm 设备信息 clinfo # 查看 OpenCL 支持（部分场景）

确保 PyTorch 支持 ROCm

官方 PyTorch 提供 ROCm 版本（支持 AMD CDNA 架构，海光 DCU 与之兼容）。
需使用 特定版本的 PyTorch + ROCm（如 ROCm 5.4.2 / 5.6 / 5.7 对应 PyTorch 2.0+）。

推荐方案：使用 PyTorch + ROCm

步骤 1

安装 ROCm（由海光或系统厂商提供）

由于海光的 ROCm 是定制版本，通常不能直接使用 AMD 官方 ROCm。

你需要：

联系服务器供应商获取 海光认证的 ROCm 驱动和运行时
或使用预装 ROCm 的国产 OS 镜像（如“海光 AI 开发环境”）

注意：

海光 ROCm 的版本可能滞后于 AMD 官方，需确认支持的 PyTorch 版本。

步骤 2

安装 ROCm 版 PyTorch

使用官方 ROCm wheel（如果兼容）

1 2	# 示例：PyTorch 2.1 + ROCm 5.6（仅当海光 ROCm 兼容时） pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.6

步骤 3

验证 PyTorch 是否识别海光 GPU

import torch
print(torch.__version__)
print(torch.cuda.is_available())      # 应返回 True（ROCm 中 CUDA API 兼容）
print(torch.cuda.device_count())
print(torch.cuda.get_device_name(0))  # 可能显示 "gfx908" 或 "Hygon DCU"

💡 在 ROCm 中，torch.cuda 接口仍然可用（AMD/海光做了 CUDA API 兼容层），所以代码无需修改。

运行 AI 模型

from transformers import AutoTokenizer, AutoModelForCausalLM

model_name = "bert-base-chinese"  # 或任何支持的模型

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")  # 自动使用海光 DCU

inputs = tokenizer("你好，海光！", return_tensors="pt").to("cuda")
outputs = model(**inputs)

✅ 只要模型不依赖 CUDA 特有算子（如 vLLM 的 PagedAttention、FlashAttention 的 CUDA 内核），就能正常运行。

注意事项

项目	说明
不支持 vLLM / TensorRT-LLM	这些框架强依赖 NVIDIA CUDA，无法在海光 DCU 上运行
FlashAttention 可能不可用	官方 FlashAttention 无 ROCm 后端，但可尝试 FlashAttention-ROCm（实验性）
性能调优需专用工具	使用 `rocprof`、`Omniperf` 等海光/ROCm 性能分析工具
分布式训练	支持 `torch.distributed` + ROCm-aware NCCL（需海光提供通信库）

替代方案

如果 PyTorch ROCm 不可用

使用 ONNX Runtime + DirectML（仅限推理）

将模型导出为 ONNX
使用 ONNX Runtime 的 ROCm 执行提供者（需编译支持）

使用 MindSpore（华为）或 PaddlePaddle（百度）

PaddlePaddle 已宣布支持海光 DCU（需使用其定制版本）
MindSpore 对国产芯片支持较好，但生态较封闭

总结

海光 GPU 调用 AI 模型路线图

建议操作

联系海光或服务器厂商，获取：
- 兼容的 ROCm 版本
- PyTorch / TensorFlow 安装包
- 性能优化指南
优先使用标准 PyTorch + Hugging Face，避免依赖 CUDA 专属扩展
测试小模型（如 BERT、ResNet）验证环境是否正常

我是码客，我是全栈工程师，我为自己代言。

海光GPU上怎样调用AI模型

前言