真正驱动神经网络的是什么:Python背后的AI编译器
究竟是什么驱动着神经网络:从 Python 表达式出发,追踪大型语言模型如何通过计算图、操作符融合、MLIR 风格的中间表示,最终实现向硬件的转换。
任何拥有编译器背景并了解一些机器学习知识的人可能都从未将这两个领域联系起来。词法分析器、语法分析器、抽象语法树以及LLVM IR属于同一个领域,而神经网络则属于另一个领域。它们在AI/ML编译器中相遇,而第一个令人惊讶的事实是:深度学习模型实际上并不是在Python上运行的。
这听起来很奇怪,因为像下面这样的代码看起来似乎全是Python在处理一切:
output = model(input)
不过,Python 并不会亲自执行数十亿次矩阵乘法运算,其背后其实还存在着更为复杂的处理过程。
从公式到硬件指令
想象一个从输入开始,经过矩阵乘法运算,然后再继续前行的小型数据流网络:
Input
↓
Matrix Multiplication
↓
ReLU
↓
Matrix Multiplication
↓
Output
在 Python 中,其核心代码可能如下所示:
y = relu(x @ W + b)
看起来很简单,但要想想机器需要什么。CPU或GPU根本不具备这种功能的概念:
relu()
或此运算符:
x @ W
正如代码中所显示的那样。硬件负责执行指令。模型的高级描述必须被转化为芯片能够高效运行的形式,而这一转化工作正是人工智能/机器学习编译器的职责。
从模型到处理器的整个流程的简化图示:
AI Model
↓
PyTorch / TensorFlow
↓
Computational Graph
↓
Intermediate Representation (IR)
↓
Optimisations
↓
Lowering
↓
Hardware-specific Code
↓
Machine Instructions
↓
CPU / GPU / NPU
阶段太多了,所以最好一个一个来。
该框架的功能
在 PyTorch 中,同样的思路可以这样实现:
y = torch.relu(x @ W + b)
开发者不想处理GPU寄存器、内存布局、指令调度以及成千上万的硬件线程。他们只需说明想要进行的计算是什么。Python和PyTorch提供了便捷的方式来实现这一表达,而其他工具则需要将这种描述转化为硬件可理解的格式。
计算图
神经网络可以被视作一个计算图。其表达形式为:
y = relu(x @ W + b)
可以将其表示为一系列相连操作的图形:
现在编译器有了具体的工作对象。一旦这些操作变得可见,它就可以判断是否能让计算速度更快。这就是优化变得至关重要的地方。
运算符融合与内存访问
以包含如下链式结构的模型为例:
MatMul
↓
Add
↓
ReLU
编译器可能会发现这些步骤可以合并,或针对目标硬件进行优化。其中一项关键技术是运算符融合:不必将每个操作视为独立的步骤,而是将某些操作合并起来,从而减少不必要的工作量,尤其是数据在内存与计算单元之间的传输。
这一点非常重要,因为在现代人工智能任务中,数据传输的成本很高。瓶颈不仅在于结果计算的速度,还在于输入数据传递到计算单元的效率。这也是该领域如此有趣的原因之一。
中间表示
接下来是来自经典编译器课程的一个熟悉概念——中间表示(IR)。传统编译器很少直接从这一步开始:
Source Code
翻译为:
Machine Code
它会经过中间表示阶段。人工智能编译器遵循相同原理,但其表示方式必须能够捕捉对机器学习而言重要的概念。而普通编译器则会考虑:
variables
functions
loops
机器学习编译器还需要表示诸如以下内容:
tensors
shapes
matrix operations
convolutions
attention
memory layouts
这正是像 MLIR 这样的项目如此重要的原因。MLIR 提供了用于在多个抽象层次上表示和转换计算的基础设施。
逐步降低抽象层次
编译器会逐渐从高级形式转向更接近目标硬件的形式。可以将其想象成走下楼梯的过程:
High-level ML operation
↓
Tensor operation
↓
Lower-level operation
↓
Hardware-oriented operation
↓
Machine instructions
在每个层级上,抽象概念都会变得更加具体,直到编译器生成目标硬件能够使用的代码。
最底层的硬件
现在我们来到了编写代码时没人会考虑的部分:
model(input)
最终的计算可能在CPU、GPU、NPU、TPU或其他人工智能加速器上完成。它们根本不在乎原始代码中是怎么写的:
y = relu(x @ W + b)
它们需要为该设备量身定制的、结构合理的可执行指令与数据。GPU可能需要特定于GPU的代码和执行模型,CPU则使用CPU指令和向量化技术,而加速器或许需要完全不同的后端支持。因此,同一个高级模型会根据其目标平台的不同而被以不同的方式编译。
那么Python的作用是什么呢?
Python并未消失,而且依然非常有用。它是定义模型、控制执行流程、加载数据以及构建应用程序的接口。只不过它并不一定需要亲自执行所有的数值运算。一个简化的理解方式是:
这改变了人们对人工智能的认知。仅从模型层面来看,情况其实很简单:
Data → Model → Prediction
加入编译器后,中间又多了一层:
Model
↓
Compiler
↓
Hardware
模型决定了需要执行哪些计算。编译器则负责确定如何表示这些计算并对其进行优化。而硬件则负责实际执行这些计算。
为何需要另一层
既然已有框架存在,为何还要添加编译器层?因为人工智能硬件正变得越来越专业化。有CPU、GPU、TPU、NPU、推理加速器以及定制芯片,它们各自具有不同的特性。在某种架构上运行良好的模型,在另一种架构上可能需要完全不同的调整。编译器起到了连接两者的作用。其一就是明确模型想要执行哪些计算:
WHAT
The ML model wants to compute
另一个方面是特定设备应如何高效地完成该任务:
HOW
A specific piece of hardware should execute it efficiently
那座桥梁的重要性与日俱增。
新领域中的经典理念
任何了解传统编译器的人都能认出其基本的处理流程:
Source Code
↓
Lexer
↓
Parser
↓
AST
↓
IR
↓
Optimisation
↓
Machine Code
机器学习在模型中也运用了类似的理念:
ML Model
↓
ML Representation
↓
ML IR
↓
Optimisation
↓
Lowering
↓
Hardware Code
领域不同,问题也更为复杂,但思维方式依然一致:表示、转换、优化、降低复杂度、执行。
一个自然的下一步想法是构建一个小型人工智能编译器来了解其工作原理。它无法与LLVM、XLA或MLIR相媲美,目标只是打造一个足够简单以便理解的工具,或许可以从一种简单的机器学习语言开始:
Simple ML Language
↓
Parser
↓
ML AST
↓
ML IR
↓
Simple Optimiser
↓
Lowering
↓
LLVM / Backend
一行内的堆栈显示
人工智能通常通过非常简单的界面呈现:
model(input)
在那次调用背后,实际上存在着一套完整的软硬件体系:框架、图结构、中间表示形式、编译器、优化技术、后端系统、硬件,以及最终的机器指令。运行一个模型意味着要启动整个编译器和硬件生态系统,而Python仅作为上层接口存在。
相关阅读
- 无需大语言模型,利用pgvector和Node.js在PostgreSQL中实现语义搜索 —— 通过pgvector、MiniLM嵌入模型以及Node.js在PostgreSQL中构建基于语义的搜索功能:包括标记、分块、向量、IVFFlat索引、距离运算符以及如何解读搜索得分。