在实际深度学习中训练或部署大模型的成本变化往往不是算法本身决定的而是底层芯片架构决定的。过去十年AI 算力从通用 CPU 逐步迁移到 GPU又从 GPU 派生出以 TPU 为代表的训练专用芯片和以 LPU 为代表的推理专用芯片。理解这些芯片的区别不能只看“谁算得快”还要看数据流方式、计算单元组织方式、片上存储带宽以及软件栈如何配合。很多开发者只在 PyTorch 里写一句model.to(cuda)却不太清楚这行代码背后为什么在 GPU 上执行效率高换到 TPU 或 LPU 上时又为什么需要重写部分算子。这篇博客会用工程视角梳理 GPU、TPU、LPU 的架构演进逻辑对比各自的适用场景并结合 Ollama、PyTorch、WSL 等常见软件栈说明 AI 芯片到底是怎么被“用起来”的。1. 先从 AI 计算的瓶颈说起为什么通用 CPU 不够用1.1 神经网络计算的真实负载特征神经网络的计算负载和传统 Web 服务完全不同。一个请求事务只涉及少量内存读取和基本运算而一次大模型前向推理需要完成海量矩阵乘法和激活函数计算。以 Transformer 结构为例核心操作是QK^T、Softmax、AV、MLP线性变换整个计算过程都集中在高维矩阵乘法和规约操作上。从计算机体系结构的角度看这类负载有三个特征计算密度高每个中间数据往往会被多个输出复用。数据并行性强同一个权重矩阵会对不同的输入批次同时运算。访存模式规律可以按固定形状分块预取和流水线相对容易。CPU 虽然逻辑控制能力强、单核延迟低但片上核心数量有限SIMD 宽度也不足以支撑大矩阵的并行乘加。更关键的是CPU 从内存取数据的路径很长而神经网络计算要求“数据喂得上计算才停不下来”。所以 AI 计算转向了“大并行 高带宽”的架构路线。1.2 “算得快”并不等于“跑得快”很多人只盯着峰值算力比如一个 GPU 声称有 100 TFLOPS但实际模型推理帧率不高。问题往往出在瓶颈转移到了访存带宽和通信同步上。一条训练链路上至少包含三个瓶颈点计算吞吐单位时间能做多少次浮点运算。内存带宽单位时间能搬多少数据到计算单元。通信开销多卡训练时梯度同步和参数交换成本。领域专用架构的核心思路就是针对特定计算模式把这三者重新平衡。有的芯片牺牲通用性把大量晶体管放在脉动阵列上让矩阵乘法的数据流始终在相邻单元之间流动有的芯片则针对推理阶段把注意力机制和自回归解码的访存模式硬编码成专用指令。这样做的代价是灵活度下降但换来的是单位功耗和单位面积上的有效算力大幅提升。注意业内常说的“AI 芯片”并不是一回事。云端训练芯片、云端推理芯片、终端 NPU、自动驾驶芯片虽然都叫 AI 芯片但侧重点完全不同。后面讨论的 GPU、TPU、LPU属于面向数据中心场景的典型代表。2. GPU从图形渲染到通用并行计算的架构基础2.1 GPU 为什么天然适合矩阵运算GPU 最早为图形渲染设计图形渲染要处理大量顶点和像素每个像素的颜色计算又相互独立因此 GPU 天然采用 SIMT 架构大量线程并行执行同一套指令但作用于不同数据。这种架构和神经网络训练高度匹配。神经网络里最常见的操作是Y X * W b其中 X 是一个[batch, seq_len, hidden_size]的张量W 是[hidden_size, output_size]的权重矩阵。整个计算可以被拆成成千上万个独立乘加单元恰好落在 GPU 的并行模型内。从硬件角度看GPU 和 CPU 的关键区别在三块线程数多GPU 有数千个 CUDA 核心或流处理器。带宽高HBM 显存提供的带宽远高于普通系统内存。上下文切换快通过 warp 调度掩盖访存延迟。2.2 CUDA 生态是 GPU 的护城河GPU 硬件本身并不是唯一价值真正的壁垒是 CUDA 生态。开发者写一个算子编译后能在不同版本的 GPU 上运行PyTorch、TensorFlow 的默认后端也会自动调用 cuDNN、cuBLAS 这些库。相比之下其他芯片即使峰值算力接近也很难在短期内兼容如此庞大的软件生态。CUDA 的执行模型可以简单理解成三层Grid一次内核启动包含多个线程块。Block多个线程组成线程块可共享显存。Thread最小的执行单元。从软件代码来看一个最简单的 CUDA 向量加法如下__global__ void vector_add(float* a, float* b, float* c, int n) { int i blockIdx.x * blockDim.x threadIdx.x; if (i n) { c[i] a[i] b[i]; } }这段代码的重点不是怎么计算而是blockIdx.x、blockDim.x、threadIdx.x如何协作把线程铺满整个数组。实际 AI 框架里很少直接写这种底层代码但理解这个模型有助于排查“GPU 利用率上不去”的问题。2.3 从训练到推理GPU 哪个阶段更擅长训练阶段GPU 的优势非常明显。反向传播需要保存中间激活值计算梯度时需要大量矩阵乘法和卷积GPU 大显存和并行计算能力都能充分发挥。推理阶段则开始出现分歧。大模型推理常见的自回归解码是串行生成 token 的过程生成第 N 个 token 时KVCache 会持续增加计算量相对固定但访存需求随时间累积。如果只是用 GPU 做单纯的前向推理显存会浪费在 KVCache 的预留空间上计算单元的利用率也可能因为批处理太小而不够理想。这也是后来出现推理专用芯片的原因推理阶段不需要像训练那样频繁更新权重、不需要保存大量中间梯度它的瓶颈更偏向“权重搬运”和“访存带宽”。3. TPU用脉动阵列把矩阵乘法做成数据流3.1 TPU 的定位面向训练和推理的专用矩阵加速器TPU 是 Google 为深度神经网络设计的专用 ASIC。它和 GPU 的路线不同GPU 是通用并行计算芯片TPU 则把大量芯片面积用在矩阵乘法单元上尽量减少控制逻辑和缓存管理开销。TPU 的核心计算单元是“脉动阵列”。所谓脉动阵列就是多个乘加单元按二维网格排布数据像心跳一样在相邻单元之间流动。权重提前固定在每个单元内部输入数据从左侧流入部分结果从下方流出形成一种流水线式的高密度计算。3.2 为什么脉动阵列效率高用一个简单的矩阵乘法来理解。假设输出矩阵是 CC 的每个元素都是 A 的某行和 B 的某列的内积。在传统 GPU 上计算单元需要从寄存器或共享内存里反复取数据在脉动阵列上数据从上到下、从左到右移动每个计算单元不需要反复访问内存而是使用相邻单元传递过来的数据。这样做带来两个好处减少访存次数数据在阵列内部流动不用每次从片上缓存取。提高计算密度同样的芯片面积下可以放更多的乘加单元。从架构图上看TPU 的脉动阵列是一个规则的二维排布。输入激活从左端进入权重预加载到每个单元部分和从底部流出再通过累加器完成最终结果。3.3 使用 TPU 时的软件约束TPU 不能直接运行任意 CUDA 代码它依赖 TensorFlow 或 JAX 的 XLA 编译器把计算图转换为 TPU 可执行指令。开发者写模型时通常不需要手写底层算子但必须注意动态 shape 会限制 XLA 优化效果。自定义算子需要封装成 XLA 兼容操作。某些模型结构如果包含大量逐元素操作脉动阵列的利用率可能不如 GPU。一位开发者的经验是标准 Transformer 模型迁移到 TPU 的收益明显但图像前处理里的一些随机操作和自定义损失函数迁移成本会上升。4. LPU面向大语言模型推理的专用处理单元4.1 LPU 要解决什么问题大语言模型推理有两个显著瓶颈。一个是权重矩阵太大HBM 带宽不足导致 token 生成速度受限另一个是自回归解码的串行性质使得并行度难以无限提高。GPU 在训练阶段表现优秀但在长上下文、大并发推理场景下会同时受到显存容量和带宽限制。LPU 的思路是把资源集中在语言模型推理最关键的路径上。它不追求所有算子都通用而是针对 Transformer 解码过程中的矩阵乘法和注意力机制做专门优化同时简化调度逻辑让计算单元尽量保持满载。4.2 LPU 的架构重点权重驻留与数据流优化LPU 在设计上有几个经典取舍提高片上 SRAM 容量降低对 HBM 的依赖。针对 KVCache 访问模式做专用缓冲。采用大规模并行但更简单的计算核心减少复杂调度带来的空转。强化批量推理能力让同一个模型同时服务更多用户。和 GPU 相比LPU 的调度模型更接近“推理专用处理器”它不需要处理图形渲染、不需要支撑训练反向传播甚至不需要像 GPU 那样运行一个完整的可编程线程模型而是用指令流驱动固定的计算管线。4.3 一个容易混淆的概念LPU 不等于“通用加速器”LPU 必须和具体模型结构匹配。针对 Transformer 的 LPU如果去跑一个 CNN 模型效果可能远不如 GPU。原因很简单CNN 的卷积计算和 Transformer 的矩阵乘法不同注意力机制中的 Softmax、LayerNorm、旋转位置编码等操作也不是所有芯片都会用硬件电路直接实现。所以在做技术选型时不能只看“峰值性能多高”而是要看“目标模型的主要算子是否落在芯片优化范围内”。5. GPU、TPU、LPU 的关键架构对比与选型逻辑5.1 从规格到适用场景的横向对比下面这个表格可以帮开发者快速建立整体印象维度GPUTPULPU架构路线SIMT 大规模并行脉动阵列 ASIC语言模型推理专用核心优化目标通用并行、生态丰富矩阵乘法密度、训练吞吐解码吞吐、低延迟、批量推理典型负载训练、推理、图形、科学计算TensorFlow/JAX 大模型训练Transformer 类模型推理软件栈CUDA、ROCm、PyTorchXLA、TensorFlow、JAX专用推理引擎/SDK编程灵活度高中低单位功耗算力中高高最大优势生态成熟、场景全面大规模训练效率高推理场景能效比强主要限制显存带宽和功耗算子覆盖有限、动态图支持弱模型结构适配要求高5.2 选型时的四种典型场景追求通用训练平台、需要同时跑多种模型优先考虑 GPU。固定用 TensorFlow/JAX 训练超大模型能接受编译器约束TPU 值得评估。大模型线上推理服务在线时长和并发量是核心 KPILPU 会体现能效优势。快速原型验证、经常调整模型结构先别锁定专用芯片等模型稳定后再做推理硬件迁移。需要强调架构选型从来不是只看芯片本身。云厂商的部署方式、数据中心的供电和散热、模型框架版本兼容性都决定了同一块芯片在真实业务中的表现。6. 软件栈与硬件协同从 PyTorch 到 CUDA、Ollama 的 GPU 调用链路6.1 PyTorch 如何调用 GPUPyTorch 调 GPU 最少需要三件事安装 CUDA 版 PyTorch、机器上有可用的 NVIDIA 驱动、显存足够放模型。一个典型的安装命令如下pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121这里cu121表示 PyTorch 针对 CUDA 12.1 编译运行时需要驱动版本兼容。如果显卡驱动太旧PyTorch 会报 CUDA 驱动不够新的错误。验证 GPU 是否可用的标准代码是import torch print(torch.__version__) print(torch.cuda.is_available()) if torch.cuda.is_available(): print(torch.cuda.get_device_name(0)) print(torch.cuda.get_device_capability(0))很多问题都出在第一步torch.cuda.is_available()返回False。常见原因是 PyTorch 安装成了 CPU 版本或者驱动和 PyTorch 的 CUDA 版本不匹配。6.2 Ollama 如何调用 GPU 运行大模型Ollama 是当前非常流行的大模型本地运行工具。它是否使用 GPU取决于安装版本、驱动、显卡品牌以及模型量化格式。常见做法是安装后先查看服务日志和状态ollama serve ollama list运行模型时可以用OLLAMA_DEBUG1查看详细日志OLLAMA_DEBUG1 ollama run qwen2.5:7b日志中会显示是否加载了 CUDA 或 Metal 后端。如果 Ollama 0.32.6 没有识别到 GPU一般从这几个方向排查驱动是否安装成功nvidia-smi是否正常输出。WSL 环境是否更新了内核驱动。容器环境是否把/dev/dxg设备映射进容器。Ollama 的OLLAMA_HOST、OLLAMA_MODELS配置是否导致它运行在错误的执行环境中。6.3 WSL 里的 GPU 调用问题WSL 2 支持通过 WSLg 和 GPU-PV 调用 Windows 侧 NVIDIA 驱动但开发者经常遇到一个经典报错failed to initialize nvml: gpu access blocked by the operating system这个报错说明 NVIDIA Management Library 初始化失败通常原因是 WSL 里装了一套 Linux 版本的 NVIDIA 驱动而 WSL 2 环境下不需要也不应该在 Linux 内部再安装驱动应该使用 Windows 侧安装的 NVIDIA 驱动然后在 WSL 内部安装 CUDA Toolkit。正确的检查顺序nvidia-smi如果这个命令无法输出 GPU 信息说明驱动链路有问题。如果能在 WSL 里看到 GPU但仍然无法使用 GPU 跑深度学习再检查 CUDA Toolkit 版本和 PyTorch 版本。注意WSL 2 里不要重复安装 NVIDIA Linux 驱动否则会出现权限冲突和 NVML 初始化失败。正确做法是只装 Windows 驱动再在 WSL 内安装 CUDA toolkit 和 PyTorch。7. 常见问题排查从 GPU 无法识别到显存不足7.1 GPU 识别不到问题现象常见原因检查方式解决思路nvidia-smi 无输出驱动未安装或未加载检查设备管理器、运行 nvidia-smi重装匹配版本驱动PyTorch 报 CUDA not available安装了 CPU 版 PyTorchpip show torch查看版本按 CUDA 版本重装WSL 中 NVML 初始化失败在 WSL 内装 Linux 驱动查看/usr/lib/wsl/lib链接改用 Windows 驱动WSL 内不装驱动Ollama 未识别 GPU后端加载失败OLLAMA_DEBUG1查看日志检查驱动和容器设备映射显卡被系统占用GPU 分片、驱动异常nvidia-smi -L查看进程重启服务或重置 GPU7.2 显存不足与优化方向大模型部署中最常见的错误是CUDA out of memory。这个报错不等于模型太大无法运行可能是 KVCache 预留过多、批次过大、内存碎片或显存泄漏。处理优先级建议先减批次或减少并发。再检查是否缓存了过多中间张量。然后考虑模型量化比如从 FP16 转为 INT8。最后考虑模型并行或换更大显存的设备。在训练场景还可以考虑梯度累积、混合精度、梯度检查点这些方法。实际项目中一个推理服务占用多少 GPU 显存需要结合模型参数量、量化位宽、上下文长度和并发数一起估算。7.3 如何观察 GPU 占用情况用nvidia-smi可以看到显存占用、利用率、温度、功耗和进程信息。需要持续监控时可以刷新nvidia-smi -l 2只看某个进程使用情况nvidia-smi --query-compute-appspid,used_memory,process_name --formatcsv如果发现 GPU 利用率很低但任务很慢通常在向“访存瓶颈”方向排查。nvidia-smi里的利用率只是 SM 利用率并不反映内存带宽是否打满。8. 实际部署中的 GPU 直通、多卡与容器方案8.1 Docker 是否需要 GPU 直通Docker 本身不会自动使用宿主机的 GPU需要借助 NVIDIA Container Toolkit。如果容器里运行 PyTorch 或 Ollama 时看不到 GPU往往会报could not select device driver with capabilities: [[gpu]]这是因为没有给 Docker 配置 NVIDIA runtime。常见安装方式distribution$(. /etc/os-release;echo $ID$VERSION_ID) curl -s -L https://nvidia.github.io/libnvidia-container/gpgkey | sudo apt-key add - curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list sudo apt-get update sudo apt-get install -y nvidia-container-toolkit sudo nvidia-ctk runtime configure --runtimedocker sudo systemctl restart docker然后运行容器时加参数docker run --gpus all --shm-size8g -it pytorch/pytorch:latest bash8.2 多卡并行与 GPU 设备编号在多卡服务器上PyTorch 默认从 GPU 0 开始分配。开发者可以通过环境变量控制对哪些卡可见CUDA_VISIBLE_DEVICES2,3 python train.py这样做的好处是进程只看到逻辑设备 0 和 1代码里就不用改设备编号。CUDA 代码里的device_id也是按这个逻辑编号来索引的。对于多卡训练最简单的方式是使用torch.nn.DataParallel但它有负载不均衡的问题。生产环境更推荐DistributedDataParalleltorchrun --nproc_per_node4 train.py8.3 如何实现跨机器的 GPU 推理本地显存不够时除了换卡还能通过推理框架把模型分配到多块 GPU 或多台机器上。常见方案包括张量并行、流水线并行和数据并行。对于 Ollama 这类工具社区版本倾向于单机部署跨机器调用通常要通过负载均衡层把请求分发到不同节点而不会在一台机器上自动拼接多个 GPU 的显存。如果需要大显存跑大模型要么选择大显存 GPU要么使用支持跨机的推理框架比如 vLLM 分布式推理、TensorRT-LLM 多卡配置。9. 不同硬件平台的实际适配差异9.1 NVIDIA GPU 之外的选择AMD 的 ROCm 和 Intel GPU 是常见的替代品。PyTorch 对 ROCm 的支持已经比较完善但在 WSL 环境里AMD GPU 的兼容性往往不如 NVIDIA。热搜词里提到的“wsl 中 ollama 如何调用 amd gpu”说明 AMD 用户在 WSL 下的配置路径并不统一。Ollama 在 Linux 原生环境下对 AMD GPU 的支持主要通过 ROCm 实现如果运行在 WSL 里需要确认 Windows 侧的 AMD 驱动支持 WSL 的 GPU-PV 机制。不同品牌的 GPU 会有完全不同的软件栈迁移前最好先在目标环境跑一个最小推理测试。9.2 昇腾等国内 AI 芯片的适配思路昇腾系列芯片是国内云服务商常见 AI 算力选项。它对应的软件栈是 CANN而不是 CUDA。使用昇腾前需要安装 NPU 驱动、固件和 CANN Toolkit然后把 PyTorch 替换为昇腾适配版本或者通过 MindSpore 运行模型。常见的用户困惑是“昇腾有哪些 GPU”这个问法不够准确。昇腾硬件通常叫 NPU不是 GPU它有自己的算子体系。迁移项目时需要重点检查模型中是否使用了 CUDA 特有算子比如某些深度可分离卷积的自定义实现这类算子必须重写为 CANN 算子或社区等价实现。9.3 没有独立 GPU 时怎么办在只有 CPU 的环境里可以选择更小的模型、量化位数更低的版本或者使用在线 API。Ollama 在 CPU 上运行小模型是可行的但并发和延迟表现会差很多。如果开发机没有独显可以租用云 GPU 进行训练设备本地只做代码调试。这样既能控制成本又能让 GPU 资源集中在真正需要的任务上。10. 面向未来的领域专用架构趋势10.1 从“通用硬件”到“软硬协同设计”芯片性能增长已经不能完全靠制程工艺提升架构设计越来越依赖和软件深度绑定。GPU 靠 CUDA 生态取胜TPU 靠 XLA 编译器做计算图优化LPU 则直接和推理引擎深度集成。未来的 AI 芯片一定不是“硬件造完软件自适配”而是从第一步就考虑算子分布、显存布局和编译器优化。10.2 推理专用芯片会越来越重要随着大模型从训练走向部署推理开销在总成本中占比会持续上升。训练阶段需要高精度矩阵乘法和梯度更新推理阶段更需要低延迟、高吞吐和低功耗。LPU 这类推理专用架构的价值会体现在长上下文和超高并发场景比如在线客服、代码补全、智能助手。10.3 NPU 进入 PC 和终端设备不只是数据中心PC 和手机上也出现了 NPU。像 AMD Ryzen AI 9 HX 370 这类处理器已经把 AI 加速单元集成进 CPU 芯片中。这种架构适合端侧推理比如语音识别、图像分类、本地小模型运行。它解决的问题不是替代云端 GPU而是让一部分推理任务不需要联网、不需要上传隐私数据就能完成。10.4 异构计算会是长期主线未来不会出现一种芯片统一所有场景。训练用 GPU 或 TPU云端推理用 GPU 或 LPU终端设备用 NPU多种芯片通过任务调度协同工作。开发者的核心能力不再是只写 CUDA 代码而是能判断某个计算任务应该放到哪个处理器上以及如何设计模型结构让软硬件协同更高效。11. 最佳实践与可复用清单11.1 芯片选型前的评估清单训练还是推理为主模型结构是 Transformer 还是 CNN是动态 shape 还是固定 shape团队熟悉 CUDA 还是 XLA还是愿意接受新工具链部署环境是在云服务器、容器、WSL还是裸机是否有多卡训练或跨机推理需求允许的功耗和成本范围是多少11.2 GPU 环境检查清单驱动版本是否满足 CUDA 版本要求。CUDA Toolkit 和 PyTorch 版本是否匹配。在 WSL 中是否避免重复安装 GPU 驱动。Docker 环境是否安装了 NVIDIA Container Toolkit。nvidia-smi是否能正常输出 GPU 信息。首次运行前用最小模型测试 GPU 调用链路是否走通。11.3 推理性能排查清单GPU 利用率低的可能原因是算子太小、数据没搬到 GPU或者模型绑定了 CPU。显存不足时优先检查 KVCache 和系统显存碎片。延迟高时关注权重加载时间、模型量化策略和并发调度。功耗异常时检查是否有多个推理进程同时竞争 GPU 资源。多卡并行时确保批次大小能被卡数整除。11.4 长期工程建议模型跑通之后把 GPU 资源占用记录成基线方便后续替换硬件时做对比。使用自动混精度时先跑一个验证集确认精度下降在可接受范围。大模型推理服务要预留显存余量不要让 KVCache 把显存打到 100%否则会出现抖动。不要为了通用性牺牲模型质量但也不要因为是“AI 芯片”就忽略实际负载形态。这篇博客的核心判断是GPU、TPU、LPU 不是同一个问题的不同答案它们分别针对 AI 工作负载的不同阶段。GPU 用成熟生态赢下训练环节TPU 用脉动阵列强化矩阵计算密度LPU 则在推理吞吐上做深度优化。在实际项目中先确认瓶颈是计算、访存、通信还是生态兼容再选硬件才不会走弯路。对于正在入门 AI 芯片的开发者最有价值的练习不是背参数而是亲自把 PyTorch、Ollama 或自定义 CUDA 算子在不同硬件环境中跑通观察计算、显存和延迟的变化这时候你对架构差异的理解会远比读规格书更深刻。