ここ2日間、GPUメモリプロファイルのグラフをいくつも眺めていました。GPUメモリプロファイルから何を抽出すればよいのか、あまり情報がなかったので、1行で処理できるスクリプトを書いてブログ記事を書きました。 それで、それらの奇妙な形は何を意味するのでしょうか? クイックチュートリアル🧵
まず、任意のコード ブロックをプロファイリングするには、この Python ファイルをプロジェクトにドロップし、以下のようなコードを使用してプロファイリング データを取得します。 https://t.co/iYTzEFksVR
以下のような可視化が得られます。3つのトレーニングステップで3つの繰り返しパターンがすぐに確認できます。良いスタートですね。これらの縞模様は、レイヤーごとにメモリ操作が実行されていることを示しています。その下の定数帯は、モデルパラメータとオプティマイザーの状態を表しています。 しかし、それらすべては何ですか?
最初の学習ステップを分解してみましょう。半三角形、つまり順方向パスから始まります。そして、巨大な赤い平行四辺形であるロジットの作成で終わります。最初の急激なスパイクは、ソフトマックスのための一時メモリです。ロジットが巨大なのは、語彙が151kあるためです。
次に、古い.gradの割り当てをクリアするzero_grad()を呼び出します。zero_grad()の最後には小さなスパイクがあり、その後大きな崖があります。スパイクは.backward()がソフトマックスから開始されたときに発生し、崖はロジットテンソルが不要になったために発生します。これにより、LMヘッドグラディエント（緑の帯）が生成されます。
上の図の右側に波のようなパターンがあることに注目してください。これは活性化チェックポイントを有効にしているためです。これにより、各層で活性化が再計算され、小さなスパイクが発生します。また、.backward() が層ごとに進むにつれて、入力テンソルが解放されます。
最後に、私のお気に入りの部分、optimizer.step() です。これは、旗竿のある寺院のように、実に美しく、二つの基礎の上に立っています。オレンジ色のストライプは、実際には .backward() の最後で生成された埋め込み層のグラデーションです。
上の図では、緑のストライプはAdamのgrad^2の平方根を保持するために必要なストレージです。Fused Adamはレイヤーごとに更新を計算し、ストライプを生成します。上部では、これらの更新を一括適用しています。スパイクは更新を適用する前の除算に必要な一時バッファです。
これらの美しいビジュアライゼーションについては、まだまだ語りたいことがたくさんあります。もし誤りや、他に知っておくべき興味深い点などがあれば、ぜひ教えてください。 詳細は私のブログをご覧ください: GPUメモリプロファイリングの取得と解釈方法 https://t.co/3uPt0S6RIp