Hotdry.
ai-systems

Tinygrad五年架构演进:从最小化深度学习框架到多后端编译优化

分析Tinygrad从2020年诞生至今的架构演进路径,涵盖其RISC-like设计哲学、多后端编译系统、延迟评估优化策略,以及向零依赖目标迈进的技术决策。

引言:最小化哲学的五年实践

2020 年 10 月 17 日,George Hotz(geohot)在 GitHub 上提交了 tinygrad 的第一个 commit。五年后的今天,这个代码库仅包含 18,935 行代码(不含测试),却能够运行 LLaMA 和 Stable Diffusion 等大型模型,在 RTX 4090 上实现 7B LLaMA 模型超过 30 tokens/s 的推理速度。tinygrad 的演进历程不仅是一个深度学习框架的发展史,更是对 "最小化" 软件设计哲学的长期实践。

正如 Hotz 在《Five years of tinygrad》中所言:"我花了五年时间在 18,935 行代码上,现在很多人也投入了数年时间。可能还需要五年时间。但这是与 NVIDIA 竞争的正确过程。" 这种对代码简洁性的执着追求,反映了 tinygrad 团队对深度学习框架本质的深刻理解:软件栈的完全主权是硬件竞争力的前提

架构演进:从 12 个原始操作到多后端编译系统

RISC-like 设计哲学

tinygrad 采用了一种 RISC-like 的设计方法,整个框架仅依赖 12 个原始操作:addmulreducemovloadstorefencebarrierwaitphiloopendloop。这种极简主义与传统的 "CISC" 框架(如 XLA)形成鲜明对比,后者通常包含数百个高级操作。

这种设计选择带来了几个关键优势:

  1. 编译优化可见性最大化:由于操作集极小,编译器能够更清晰地看到整个计算图的结构,为内核融合等优化提供了理想条件。
  2. 后端实现简化:每个后端只需实现这 12 个操作,大大降低了移植到新硬件平台的成本。
  3. 代码可维护性:小规模的核心代码库使得深度理解和修改成为可能,这与大多数深度学习框架的 "黑盒" 特性形成对比。

多后端编译架构

经过五年的演进,tinygrad 已经支持九个不同的运行时后端,所有这些后端都通过相同的单元测试:

  • CPU 后端:NumPy
  • GPU 后端:OpenCL、Metal、CUDA
  • 编译器后端:Triton、LLVM、C(Clang)
  • 其他:WebGPU、RISC-V 向量扩展(社区贡献)

每个后端的实现代码通常少于 1,000 行,这种 "小表面积" 策略鼓励社区贡献新的硬件支持。后端架构的核心是 UOps(Universal Operations)系统,它将计算图转换为设备无关的中间表示,然后通过特定于设备的编译器生成最终的可执行代码。

性能优化:延迟评估与内核融合策略

延迟评估系统

tinygrad 采用延迟评估(lazy evaluation)策略,这意味着操作不会立即执行,而是在调用.realize()方法时才进行实际计算。这种设计允许编译器在运行时之前看到完整的计算图,为优化提供了前所未有的机会。

延迟评估的核心优势在于跨操作边界的内核融合。传统的深度学习框架通常在每个操作边界处产生内存传输和同步开销,而 tinygrad 能够将整个前向传播和反向传播过程融合到单个 GPU 内核中。例如,一个典型的矩阵乘法操作可能涉及 reshape、broadcast-multiply 和 reduce-sum 等多个步骤,在 tinygrad 中这些步骤可以被编译成一个单一的 matmul 内核。

内存管理优化

tinygrad 的内存管理策略同样体现了最小化哲学:

  1. 零拷贝数据移动:通过 BufferCopy 操作在主机和设备之间高效传输数据
  2. 设备内存池:重用已分配的内存块,减少分配 / 释放开销
  3. 常量折叠优化:在编译时计算常量表达式,减少运行时计算

这些优化使得 tinygrad 在资源受限的环境中表现出色。例如,团队正在开发 int4/int3 量化内核,目标是在 8GB 显存卡上运行 30B 参数的模型。

自动微分系统重构与分布式训练支持

自动微分系统演进

tinygrad 的自动微分系统经历了多次重构,从最初的简单实现演变为现在的高效系统。关键改进包括:

  1. 符号微分与自动微分的结合:对于简单操作使用符号微分,复杂操作使用自动微分
  2. 梯度检查点优化:在内存和计算之间智能权衡
  3. 二阶导数支持:通过自动微分的自动微分实现

系统采用基于 tape 的反向模式自动微分,每个操作都记录其前向传播信息,在反向传播时按相反顺序应用梯度计算。

分布式训练架构

虽然 tinygrad 最初专注于单设备训练,但近年来已经扩展到分布式训练支持:

  1. 数据并行:通过 AllReduce 操作同步梯度
  2. 模型并行:支持跨多个设备的模型分割
  3. 流水线并行:将模型层分配到不同设备,重叠计算和通信

分布式训练的实现保持了 tinygrad 的最小化哲学,通过少量精心设计的原语支持复杂的并行策略。

工程实践:零依赖目标与 LLVM 移除计划

向零依赖迈进

tinygrad 团队目前正在积极移除 LLVM 依赖,目标是实现 "零依赖(除纯 Python 外)驱动 AMD GPU"。这一决策基于几个关键考虑:

  1. 构建复杂性降低:移除 LLVM 可以简化构建过程,特别是在嵌入式环境中
  2. 启动时间优化:避免 LLVM 的初始化开销
  3. 代码透明度:完全控制整个软件栈,便于调试和优化

正如 Hotz 指出的:"只有傻瓜才会从流片芯片开始;这很昂贵而且不是困难的部分。一旦你拥有能够训练 SOTA 模型的完全主权软件栈,芯片就变得如此简单。"

解构公司模式

tinygrad 的开发组织模式同样具有创新性。团队采用 "解构公司" 模式:

  • 几乎完全公开:主要通过 GitHub 和 Discord 运作
  • 贡献者招聘:通过代码贡献获得工作机会
  • 极简会议:每周仅一次会议,专注于使 tinygrad 变得更好
  • 多元化收入:通过计算机销售部门(年收入约 200 万美元)和 AMD 合同资助开发

这种组织模式与代码库的最小化哲学相一致,减少了传统公司结构带来的开销和复杂性。

技术挑战与限制

尽管 tinygrad 取得了显著进展,但仍面临一些挑战:

  1. 生态系统成熟度:相比 PyTorch 和 TensorFlow,tinygrad 的生态系统较小,缺少成熟的工具链和预训练模型库
  2. 多后端维护成本:随着支持的硬件平台增加,维护成本呈线性增长
  3. 性能调优复杂性:虽然基础性能优秀,但针对特定硬件的深度优化需要专业知识
  4. 文档和教程:相对于主流框架,学习资源有限

对深度学习框架设计的启示

tinygrad 五年的演进为深度学习框架设计提供了几个重要启示:

1. 最小化作为竞争优势

在大多数框架追求功能丰富性的时代,tinygrad 证明了最小化可以成为竞争优势。小代码库意味着:

  • 更快的理解曲线
  • 更低的维护成本
  • 更高的优化可见性
  • 更容易的硬件移植

2. 软件栈主权的重要性

tinygrad 的经验验证了 Hotz 的观点:"AMD、Amazon、Tesla 和 Groq 都流片了不错的芯片,但只有 Google 和 NVIDIA 的芯片曾被认真用于训练。因为它们有软件。" 软件栈的完全主权是硬件竞争力的前提

3. 延迟评估的威力

延迟评估不仅仅是编程语言特性,在深度学习框架中,它成为了性能优化的核心机制。通过延迟执行,编译器获得了优化整个计算图的能力,这是即时执行框架难以实现的。

4. 社区驱动的创新

tinygrad 的 "解构公司" 模式展示了开源社区如何驱动技术创新。通过降低贡献门槛和保持完全透明,项目吸引了高质量的贡献者,形成了良性的创新循环。

未来展望

展望未来,tinygrad 的发展方向包括:

  1. 完全移除 LLVM 依赖:实现真正的零依赖软件栈
  2. 更多硬件支持:扩展到新兴的 AI 加速器
  3. 量化支持增强:支持更低精度的推理和训练
  4. 编译器优化深化:进一步挖掘延迟评估的优化潜力
  5. 生态系统建设:发展更丰富的工具链和模型库

结论

tinygrad 五年的演进历程展示了深度学习框架设计的另一种可能性:通过极简主义实现高性能,通过完全透明建立信任,通过社区驱动加速创新。在 AI 硬件竞争日益激烈的今天,tinygrad 的经验提醒我们:软件栈的质量和主权可能比硬件规格更重要

正如 Hotz 总结的:" 我们的使命是使 petaflop 商品化。" 这一使命不仅关乎技术,更关乎 AI 计算的民主化 —— 通过最小化、透明化的软件栈,让更多人能够访问和利用强大的计算能力。

在深度学习框架日益复杂化的趋势中,tinygrad 坚持的最小化哲学提供了一种反直觉但有效的替代方案。它的成功证明了:有时候,最好的部分就是没有部分


资料来源

  1. George Hotz, "Five years of tinygrad", geohot.github.io, 2025-12-29
  2. "tinygrad: The Ultra-Minimal Deep-Learning Library That Runs LLaMA and Stable Diffusion", blog.brightcoding.dev, 2025-09-08
  3. tinygrad GitHub 仓库:https://github.com/tinygrad/tinygrad
查看归档