在使用PyTorch或TensorFlow等框架时,开发体验感觉即时且具有交互性。您可以定义一个张量,执行一个加法,并打印结果。这种设计选择,通常称为立即执行模式(eager execution),优先考虑开发人员的效率和调试的方便性。然而,这种便利性掩盖了Python程序描述操作的方式与硬件加速器执行这些操作的方式之间存在的显著隔阂。这种隔阂被称为框架与硬件的差距。这种差距存在的原因是,现代深度学习加速器(如GPU和TPU)依赖大规模并行和高吞吐量来实现高性能。相比之下,运行在CPU上的Python解释器本质上是串行和动态的。连接这两种环境会引入额外开销,从而严重限制机器学习模型的性能,无论底层硬件可能有多强大。动态调度的开销在解释型环境中,框架逐个执行操作。当解释器遇到 c = a + b 这样的代码行时,它会在任何计算发生之前执行一系列检查。它必须验证数据类型,检查张量形状的兼容性,为结果分配内存,最后为硬件后端选择适合的实现(内核)。这个过程称为动态调度。虽然调度单个操作的开销以微秒计,但一个深度神经网络可能涉及数十万个此类操作。如果GPU上的实际数学计算所需时间少于CPU准备和调度指令所花费的时间,那么加速器将空闲等待任务。这会创建一个程序“CPU受限”的场景,这意味着昂贵的GPU资源未得到充分利用。下图展示了立即执行模式的流程,突出了主机CPU在GPU内核执行之间引入的延迟。digraph G { rankdir=TB; node [fontname="Helvetica,Arial,sans-serif", shape=box, style=filled]; edge [fontname="Helvetica,Arial,sans-serif"]; subgraph cluster_host { label = "主机 (CPU)"; style = filled; color = "#f8f9fa"; node [color="#adb5bd", fillcolor="#ffffff"]; Python [label="Python 解释器\n(读取操作 1)"]; Dispatch [label="框架调度器\n(检查类型, 形状)"]; Launch [label="驱动启动\n(cudaLaunchKernel)"]; Python2 [label="Python 解释器\n(读取操作 2)"]; Dispatch2 [label="框架调度器"]; Launch2 [label="驱动启动"]; } subgraph cluster_device { label = "设备 (GPU)"; style = filled; color = "#e7f5ff"; node [color="#4dabf7", fillcolor="#d0ebff"]; Kernel1 [label="执行内核 1\n(加法)"]; Idle [label="空闲 / 等待状态", style=dashed, color="#fa5252", fontcolor="#fa5252"]; Kernel2 [label="执行内核 2\n(Relu)"]; } Python -> Dispatch; Dispatch -> Launch; Launch -> Kernel1 [label="异步调用"]; Kernel1 -> Idle; Idle -> Kernel2; Launch -> Python2 [label="返回控制权", style=dashed]; Python2 -> Dispatch2; Dispatch2 -> Launch2; Launch2 -> Kernel2; }立即执行模式下的操作序列,显示了CPU调度开销如何造成GPU利用率的空闲间隙。内存带宽和算子粒度差距的第二个主要方面涉及数据在内存中的流动方式。在标准框架实现中,每个操作都被视为一个独立的函数调用。考虑神经网络中一个典型的层,它执行卷积、添加偏置并应用ReLU激活函数:$$y = \text{ReLU}(\text{Conv}(x, W) + b)$$以立即执行模式运行时,框架将其作为三个独立步骤处理:卷积: 从全局内存读取输入 $x$ 和权重 $W$,计算,并将中间结果写入全局内存。添加偏置: 从全局内存读取中间结果和偏置 $b$,进行相加,并将新结果写回全局内存。ReLU: 从内存读取结果,应用激活函数,并将最终输出写入内存。这种模式效率低下,因为访问全局内存 (VRAM) 远慢于进行算术逻辑运算。数据在芯片的计算单元及其主内存之间反复来回移动。硬件加速器在算术运算与内存访问之比(算术强度)高时表现最佳。通过将算子视为细粒度、独立的单元,框架人为地降低了这种强度。内存带宽成为瓶颈,在计算核心饱和之前,有效吞吐量就已经达到上限。下图比较了朴素执行与优化方法的执行时间分解,其中开销和内存访问被最小化。{ "layout": { "title": "执行时间分解:立即执行与优化执行", "barmode": "stack", "xaxis": { "title": "执行模式", "showgrid": false }, "yaxis": { "title": "归一化时间", "showgrid": true, "gridcolor": "#e9ecef" }, "showlegend": true, "plot_bgcolor": "#ffffff", "width": 600, "height": 400, "margin": {"l": 50, "r": 50, "t": 50, "b": 50} }, "data": [ { "type": "bar", "name": "内核计算", "x": ["立即执行", "优化执行"], "y": [40, 40], "marker": {"color": "#4dabf7"} }, { "type": "bar", "name": "内存读写等待", "x": ["立即执行", "优化执行"], "y": [50, 10], "marker": {"color": "#ff6b6b"} }, { "type": "bar", "name": "CPU调度开销", "x": ["立即执行", "优化执行"], "y": [30, 2], "marker": {"color": "#adb5bd"} } ] }计算时间与内存访问及调度开销的比较。优化执行显著减少了非计算任务。内核启动瓶颈每当框架指示GPU运行一个函数时,它都会启动一次“内核启动”。这次启动会告知GPU调度器要运行的代码、参数以及如何将线程映射到数据。对于大规模矩阵乘法(如大型语言模型中的乘法),计算时间足够长,启动开销可以忽略不计。然而,现代网络架构常包含许多小型、逐元素的运算(激活、归一化、形状变换)。对于这些小型算子,配置和启动内核所需时间可能超过执行时间。如果一个模型需要启动50个不同的内核来处理单个输入图像,并且每次启动都产生固定开销,那么延迟就会累积。这对于需要低延迟的推理工作负载来说尤其是个问题。弥合差距为了解决这些低效问题,我们不能仅仅依赖更快的Python解释器或更快的硬件。解决方案在于重构执行策略本身。算子融合: 不将中间结果写入内存,编译器可以生成一个单独的内核,一次性完成卷积、偏置添加和ReLU运算。数据保留在GPU的快速寄存器或L1缓存中,减少内存带宽压力。静态调度: 通过提前分析计算图,编译器可以消除运行时动态调度检查的需要。图捕获: 编译器将程序意图捕获为中间表示 (IR),这使得它们能够将模型视为一个整体,而不是一系列独立的 Python 命令。ML编译器栈主要存在是为了自动化这些优化,将框架的高级数学意图转换为硬件所需的严格、高效的指令流。