プロジェクト：tinygrad/tinygrad · 作者：tiny corp / George Hotz ほか · License：MIT
これは読書ノートです。私は開発に参加していません。

概要

tinygrad は PyTorch と micrograd の間に位置する深層学習フレームワークで、「読み切れるサイズ」を目標にしています。少数の Python ファイルで Tensor、自動微分、IR コンパイラ、JIT、複数のハードウェアバックエンドを実装し、訓練と LLM 推論を動かせます。

中核コンポーネント

  • Tensor ライブラリ：PyTorch 風の Eager API と自動微分
  • IR コンパイラ：カーネル融合と lowering、JIT + Graph 実行
  • 学習ツール：nn、optim、datasets
  • LLM 推論：GGUF 対応

ハードウェアバックエンド

CUDA / PTX、AMD ROCm、Apple Metal、OpenCL、Qualcomm QCOM、WebGPU、LLVM IR（CPU）、ONNX インポート。

興味深いのは、新しいアクセラレータを追加するのにフレームワーク全体を書き直す必要がないことです。少数の低レベル演算子を実装すれば、融合とスケジューリングは IR コンパイラが担当します。

読む価値のある部分

Lazy Execution と Kernel Fusion

DEBUG=3 または 4 で、融合後のカーネルと生成コードを確認できます。JIT キャプチャとカーネル再生で、小さい演算子のオーバーヘッドを減らします。

デバイスごとに最適なカーネル設定を探索します。同じ計算グラフでも、GPU によって最適な実装が大きく異なるという現実を反映しています。

Process Replay

変更前後で生成カーネルを比較し、性能劣化を防ぎます。「動けばいい」プロジェクトに欠けがちな工程です。

私の評価

tinygrad は、フレームワーク内部の仕組みを学ぶことと、新しいハードウェアバックエンドの実現可能性を検証することに向いています。PyTorch の代替や本番訓練基盤ではありません。しかし「演算子がどう GPU コードになるか」を知りたいなら、多くのフレームワークより透明です。

元プロジェクトを見る