プログラミング
PyTorch全体を1ページで
The whole of PyTorch on one page (tensor.khalilli.ai)
要約
この記事は、PyTorchの内部構造を包括的に解説するシリーズの導入部です。Pythonからコンパイルされたコード、ディスパッチャー、カーネルに至るまで、8つの階層を高速に traverser します。特に、PythonとC++間の境界を越える際のオーバーヘッドや、ディスパッチャーが操作をどのようにルーティングするかについて説明し、PyTorchの複雑なアーキテクチャへの理解を深めるためのロードマップを提供します。
全文翻訳
目次
秋
領域
十二のアイデア
このシリーズの描き方
読み方
今言えること
試してみる
ドアを選ぶ
参考文献
図1。このシリーズ全体が対象とするプログラム。
あなたはこのようなものを何千回も入力したことがあるでしょう。このシリーズは、その終わりまでに、これらの行のすべてが何をするかを知るために存在します。それらのPython、それらが到達するC++、それらが記録するグラフ、それらが選択するカーネル、それらが使用するメモリ、そしてそれらが実行される2つのクロック。それらの各単語は、下のフロアで平易な意味を持ちます。
これはパート0、マップです。まず、すべてのレイヤーを一度、速く下に進みます。次に、領域を描きます。次に、コードベースの残りの部分を予測可能にする12のアイデア。次に、このシリーズの仕組みと、それを読む方法。ここにあるものは何も完全な物語を得ません。ここにあるすべては場所を得て、すべての完全な物語はそれを待つ番号付きの部分を持っています。
始める前に一つの約束があります。このシリーズの測定されたすべての数値は、数値が表示されている場所にリンクされている、自分で実行できる小さなスクリプトから来ています。私はこれらをtorch 2.11.0 [1]を搭載したApple M3 Maxラップトップで測定しました。あなたの数値は異なるでしょう。それらが作るパターンはそうではありません。
秋
PyTorchは深いです。あなたのキーボードとチップの間には8つのレベルがあります。私はそれらをフロアと呼び、このメーターはそれらすべてを示します。それはシリーズ全体を通して戻ってくるので、あなたは常にあなたがどれだけ深いかを知ることができます。
図2。深さメーター。オレンジ色の点はあなたがどこにいるかを示します。
建物を学ぶ最速の方法は、一度立ち止まらずに通り抜けることです。それがこのセクションです。
フロア1:Python
torch.randnはPython関数のように見えます。Pythonにそれが実際に何であるかを尋ねてください。
>>> type(torch.randn)
<class 'builtin_function_or_method'>
Pythonはそのタイプをコンパイル済みコードで書かれた関数にのみ与えます。コンパイル済みコードとは、あなたがそれをインストールする前にマシン命令に翻訳されたコードのことです。そのため、中に読むべきPython本体はなく、デバッガーが停止する行もありません。
では、それらのマシン命令はどこにありますか?共有ライブラリの中にあります。共有ライブラリは、プログラムが実行中にロードするコンパイル済みコードのファイルです。それらはディスク上のtorchパッケージの中にあり、それらを見ることができます(証明):
torch._C -> _C.cpython-312-darwin.so (49 KB、ローダー)
libtorch_cpu.dylib 206.5 MB (テンソルとカーネル)
libtorch_python.dylib 28.5 MB (境界のPython側)
p0_the_library.py 証明、読み取りまたは実行可能
"""証明:PyTorchのコンパイル済み部分が実際にどこにあるか。 torch._Cは薄いコンパイル済みのスタブです。フレームワークの重みは、その隣にある共有ライブラリにあります。ファイルとそのサイズを出力します。"""
import glob
import os
import torch
stub = torch._C.__file__
print(f"torch {torch.__version__}")
print(f"torch._C -> {os.path.basename(stub)}"
f" ({os.path.getsize(stub)/1024:.0f} KB stub)")
libdir = os.path.join(os.path.dirname(stub), "lib")
for lib in ["libtorch_cpu.dylib", "libtorch_python.dylib"]:
p = os.path.join(libdir, lib)
if os.path.exists(p):
print(f"{lib:24s} {os.path.getsize(p)/1024/1024:6.1f} MB")
ダウンロードして実行する
サイズを読み取り、それらを見てください。
図3。ファイルサイズで等倍で描画。Pythonが見ることができるPyTorchの部分はオレンジ色の点です。
Pythonから見えるPyTorchの部分は49 KBのファイルで、その唯一の仕事は他の2つをロードすることです。実際の本体は235 MBのコンパイル済みコードです。import torchはそれをあなたのプロセスにロードし、その後、torch.randnを呼び出すことはその本体にジャンプすることを意味します。今日はこれらのファイルが存在することだけを知っていれば十分です。
これはコードベースの最初の正直な驚きです:あなたが一日中書くPythonは、そのレイヤーの中で最も小さいものです。
境界
呼び出しはすぐにPythonを離れます。どこに着陸しますか?
前のフロアの28.5 MBのライブラリ内にあるTHPVariable_randnという名前のC++関数に着陸します。そして、ここに奇妙な事実があります。この関数はPyTorchリポジトリには存在しません。リポジトリをクローンして、名前を検索しても何も見つかりません。プログラムはビルド中にこの関数と、その何千もの兄弟を書き込みます。アイデア4は理由を説明し、パート5はそのプログラムを示します。
図4。二つの言語の間の境界。すべてのテンソル操作がそれを横切ります。
この境界を越えるにはコストがかかります。コストだけを見るために、最小限の操作を時間計測します。ほとんど算術演算がそれを隠さないようにします(証明):
add、1要素:0.538マイクロ秒/呼び出し
add、4M要素:337.264マイクロ秒/呼び出し
p2_dispatch_cost.py 証明、読み取りまたは実行可能
"""証明:1つのEager Opの固定コスト、およびサイズがそれを隠す理由。2つのサイズで同じ`a + b`を時間計測します。1要素のaddはほぼ純粋な機械(ディスパッチ、ラッパー、割り当て)であり、4M要素のaddはほぼ純粋な算術です。CPU、単一プロセス。"""
import time
import torch
def per_op_us(a, b, iters):
# ウォームアップ
for _ in range(2000):
a + b
t0 = time.perf_counter()
for _ in range(iters):
a + b
return (time.perf_counter() - t0) / iters * 1e6
tiny = per_op_us(torch.ones(1), torch.ones(1), 200_000)
big_n = 4_000_000
big = per_op_us(torch.ones(big_n), torch.ones(big_n), 2_000)
print(f"torch {torch.__version__}, cpu")
print(f"add, 1 element : {tiny:8.3f} us/op")
print(f"add, 4M elements : {big:8.3f} us/op")
print(f"machinery share of the tiny op: ~all of it")
print(f"ops/sec you can issue from python: {1e6/tiny:,.0f}")
# Toll meterの背後にあるスイープ:12のサイズでの同じadd、
# 4のべき乗で1から4M要素まで。ウィジェットの軸の各ドットは、この出力の1行です。
import json
sweep = []
for k in range(12):
n = 4 ** k
iters = max(1_000, min(200_000, 40_000_000 // max(n, 1)))
us = per_op_us(torch.ones(n), torch.ones(n), iters)
sweep.append({"n": n, "us": round(us, 3)})
print(f"add, {n:>9,} elements : {us:9.3f} us/op")
print("JSON_SWEEP=" + json.dumps(sweep))
ダウンロードして実行する
1要素のaddはほとんど計算を行いません。そのため、0.54マイクロ秒はほぼ純粋なクロスオーバーコストです:Pythonを離れ、引数をチェックし、結果オブジェクトを構築し、返します。半マイクロ秒は取るに足らないように聞こえます。それはPythonが毎秒約190万操作しか発行できないことを意味し、単一のトレーニングステップには数千の操作が含まれます。この数値を覚えておいてください。アイデア6で戻ってきます。
ディスパッチャー
境界の下で、呼び出しはPyTorchで最も奇妙なマシンに到達します:ディスパッチャー。ディスパッチャーは、各操作に対して、どのコード片がどの順序で実行されるかを決定するルーターです。
それが決定しなければならないことを見てください。あなたの3行は「勾配を記録する」とは決して言いませんでした。あなたのコードにそれをオンにするif文はありません。それでもどこかで、この行列乗算が後方()のために記憶されるべきだと決定されました。その何かがディスパッチャーです。すべての操作は、固定されたレイヤーのスタックを通過します。各レイヤーは呼び出しに作用したり、変更したり、変更せずに通過させたりできます。Autograd、勾配を計算するPyTorchの部分は、そのようなレイヤーの一つです。混合精度も別のレイヤーです。この実行では、Autogradのみがアクティブです。
図5。4つのレイヤーが、計算が始まる前に呼び出しに触れます。今日アクティブなのはハイライトされたものだけです。
カーネル
スタックの底で、具体的な関数が選択されます。選択された、が正しい言葉です。このtorchビルドには3,677の登録済み操作名があります(証明はカウントを出力します)、そして名前は関数本体ではありません。操作addmm、モデル(x)の背後にある行列乗算は、CPU用と各種類のGPU用、各データ型用、密なテンソル用と疎なテンソル用に別々の本体を持っています。このような本体、つまり1つのデバイスと1つのデータ型のために書かれたものは、カーネルと呼ばれます。ディスパッチャーの最後の仕事は、1つを選択することです。
図6。1つの名前、本体のグリッド。ディスパッチャーは呼び出しごとに正確に1つのセルを選択します。
p4_micro_proofs.py 証明、読み取りまたは実行可能
"""パート0で引用されたマイクロ証明:ストレージ共有、ビューエラー、ミューテーション履歴の書き換え、no_gradレイヤー、float32吸収。"""
import torch
print(f"torch {torch.__version__}\n")
# 1. テンソルはストレージのウィンドウです
x = torch.arange(6, dtype=torch.float32)
v = x.view(2, 3)
print("same bytes under both:", x.data_ptr() == v.data_ptr())
print("v.stride():", v.stride(), " v.t().stride():", v.t().stride())
try:
v.t().view(