WebGLでGPT-2の実行：GPUシェーダープログラミングの失われたアートの再発見

5月、24 2025•16分読む•2461語

序文：数週間前、WebGLとシェーダーを使用してGPT-2を実装しました（Github Repo）ハッカーニュースのフロントページを作成しました（議論）。一般的な需要によって、GPUシェーダープログラミング（汎用コンピューティング用）の背後にある主なアイデアに関する短い記事を以下に示します。

一般的な目的GPUプログラミングの起源

2000年代初頭、NvidiaはGeForce 3（2001）およびGeforce FX（2003）でプログラム可能なシェーダーを導入しました。以前のGPUの所定の変換と効果に限定される代わりに、開発者は現在、レンダリングパイプラインを前例のない制御を与えられ、はるかに洗練された視覚効果を可能にしました。これらのプログラム可能なシェーダーは、最新のGPUコンピューティングの基礎を築きました。

研究者はすぐに、特定の計算（マトリックスやベクトルを含む線形代数など）がGPUのグラフィック操作としてそれらをキャストすることで加速できることを発見しました。ただし、No-GraphicsタスクにOpenGLのGLSLなどのシェーダー言語を使用するのは面倒でした。 2000年代半ばまでに、GPUへのより簡単な非グラフィックインターフェースの必要性が明らかになり、Nvidiaは新しい機会を見ました。

の需要に触発されました 汎用GPU（GPGPU） 2006年11月、Nvidiaがリリースしました cuda、 統一されたデバイスアーキテクチャを計算します。 CUDAは、グラフィックスAPIの仲介者なしで、開発者がGPUの計算能力に直接アクセスできる並列コンピューティングプラットフォームおよびプログラミングモデルです。 CUDAを使用すると、C/C ++コードを記述して、並列カーネルの簡単な拡張機能とGPUメモリの管理を明示的に管理するCPUで実行することができます。これは、開発者がグラフィック固有の概念を無視し、汎用GPUコンピューティングの障壁を劇的に低下させることができることを意味しました。 CUDAに続いて、NVIDIAエコシステムを超えて汎用コンピューティングを拡大したOpenCLが登場しました。

グラフィックアピvsコンピューティングAPI

OpenGLのような従来のグラフィックスAPIは、画像のレンダリングに合わせた固定機能パイプラインを中心にしています。パイプラインは、頂点処理、ラスター化、フラグメント処理などの段階で構成されています。各段階はシェーダーでプログラム可能である可能性がありますが、全体的なフローは固定されています。計算にOpenGLを使用するには、多くのボイラープレートが必要でした。データをテクスチャ形式に詰め、画面外のフレームバッファを使用して結果をキャプチャする必要があり、多くの場合、複数のレンダリングパスを実行してマルチステージアルゴリズムを達成しました。

対照的に、OpenCLとCUDAは、GPUを1つの巨大なSIMDプロセッサとして扱うための直接計算モデルを公開します。

• シェーダーではなくカーネル：関数を書き、数千コピーを起動して並行して実行します（頂点やフラグメントの概念はありません）。
• 生のバッファー：フロートまたは整数の配列を割り当て、それらを直接読み取り/書き込み、明示的なコピーコールでホストとデバイスの間で行き来します。
• ユーザー駆動型パイプライン：レンダリング段階の定義された固定シーケンスを使用する代わりに、何が起こるかを正確に定義します。

その結果、線形代数、シミュレーション、物理学、ML、および独立した計算をバルクで計算するだけのアルゴリズムにははるかに自然に適合します。

OpenGLでは、計算の出力は最終的にフレームバッファーのピクセルまたはテクスチャの値になります。 OPENCLでは、出力は任意の形式（フロートアレイ、計算された数字など）のデータにすることができます。これは、CPUに戻るか、さらなる計算で使用します。これにより、OpenCLは、数値結果が必要な一般的なアルゴリズムにより適しています。

シェーダーを使用してGPT-2を実装します

以下には、テクスチャ、フレームバッファ、頂点、フラグメントシェーダー、およびシェーダーを使用してGPUでGPT-2を実行するためにハイジャックした他のグラフィックス固有の概念をカバーしています。

テクスチャとフレームバッファ：データバス

従来のグラフィックレンダリングでは、a テクスチャ GPUメモリに保存されているピクセルデータの2D（または3D）配列です。テクスチャを三角形にマッピングすると、フラグメントシェーダーが「サンプル」して色の値を調べます。グラフィックとしてのコンピューティックパラダイムでは、色の代わりに数値データを保存および操作するために、この同じメカニズムをハイジャックします。

• テンソルとしてのテクスチャ：各テクスチャは、フローティングポイント値の配列（シングルチャネルR32F形式を使用します）で、各ピクセルはマトリックスまたはベクトルの1つの要素をエンコードします。 H×Wテクスチャを画像のH×W RGBピクセルを保持していると考えるのと同じように、ここでは、重量マトリックスまたはアクティベーションマップのH×Wスカラー値を保持します。

• フィルタリングなしでサンプリング：次のような関数を使用します texelFetch 正確な整数座標でテクスチャデータを読み取るには、補間をバイパスします。これにより、「ランダムアクセス」は、列と列ごとにCPUアレイにインデックスを作成することに似た、重量とアクティベーションアレイに読み書きされます。

a フレームバッファオブジェクト（FBO） 画面からテクスチャの1つにレンダリング出力をリダイレクトできる軽量コンテナです。

1. テクスチャをレンダリングターゲットとして取り付けます：テクスチャをFBOにバインドすることにより、通常はモニター用に運命づけられているドローコールが、代わりにそのテクスチャに書き込みます。フラグメントシェーダー out 変数は、GPUメモリへの書き込みポートになります。

2. オフスクリーンとピンポンレンダリング：さまざまなテクスチャを連続して取り付けることができるため、それらの間に「ピンポン」があります。 テクスチャa、次のパスはからです テクスチャa 書き込み中 テクスチャb、 等々。これにより、最終的にデータをCPUにコピーすることができなくなります。

3. ハイスループットデータバス：これらはすべて、GPUのVRAMバスで完全に発生します。バインディングテクスチャとフレームバッファーは、GPUでのポインター交換だけです。セットアップすると、フラグメントシェーダーは、システムメモリに触れることなく、並列、読み取り、コンピューティング、および書き込みで何百万ものコアを通過します。

一緒に、テクスチャとFBOが形成されます データバス シェーダーベースのコンピューティングエンジン：テクスチャは、ニューラルネットワーク（重量、中間アクティベーション、および出力）の生のビットを保持し、フレームバッファーによりシェーダーをシームレスに通過させ、最終的なロジットをCPUに明示的に引き戻すまで、すべてを高速GPUパイプラインに保持します。

計算カーネルとしてのフラグメントシェーダー

フラグメントシェーダーは魔法が起こる場所です。フラグメントシェーダーを使用してディスプレイ用のピクセルをシェーディングする代わりに、コンピューティートカーネルとしてハイジャックします。各フラグメントの呼び出しは、単一の出力値を計算する「スレッド」になります。 GPUはこれらの数千を並行して発射し、ニューラルネットワーク操作のための大規模なスループットを提供します。

以下は、マトリックス増殖のためのフラグメントシェーダーの例です。

// Matrix Multiply (C = A × B)
#version 300 es
precision highp float;

uniform sampler2D u_A;    // Texture holding matrix A (M x K)
uniform sampler2D u_B;    // Texture holding matrix B (K x N)
uniform int        u_K;   // Shared inner dimension
out vec4           outColor;

void main() {
  // Determine the output coordinate (i, j) from the fragment’s pixel position
  ivec2 coord = ivec2(gl_FragCoord.xy);
  float sum = 0.0;

  // Perform the dot–product along the K dimension
  for (int k = 0; k  u_K; ++k) {
    float a = texelFetch(u_A, ivec2(k, coord.y), 0).r;
    float b = texelFetch(u_B, ivec2(coord.x, k), 0).r;
    sum += a * b;
  }

  // Write the result into the single‐channel R component of the output texture
  outColor = vec4(sum, 0.0, 0.0, 1.0);
}

ここに：

• ピクセルあたりの作業アイテム：各フラグメントは、1つのマトリックス要素（i、j）に対応します。 GPUは、シェーダーコア全体ですべて（i、j）に対してこのループを実行します。
• 正確なインデックス：Texelfetchは、整数座標によって単一のフロートを読み取ります。
• 書き込みバック：autcolor.rへの割り当ては、その計算値を（i、j）のバウンドFBOのテクスチャに直接書き込みます。

これが別のフラグメントシェーダーですが、GELUアクティベーション関数の場合：

// GELU Activation
#version 300 es
precision highp float;
uniform sampler2D u_X;
out vec4 o;

void main() {
  ivec2 c = ivec2(gl_FragCoord.xy);
  float x = texelFetch(u_X, c, 0).r;
  float t = 0.5 * (1.0 + tanh(0.79788456 * (x + 0.044715 * x*x*x)));
  o = vec4(x * t, 0, 0, 1);
}

共有頂点シェーダー

すべての操作は、操作の数学が発生する場所であるため、独自のフラグメントシェーダーを取得します。一方、頂点シェーダーはシンプルで同じです。それがするのは、ビューポート全体を覆う2つの三角形を描画することです。

// Shared Vertex Shader
#version 300 es
in vec2 a_position;
out vec2 v_uv;
void main() {
  v_uv = a_position * 0.5 + 0.5;  // map [-1,+1] to [0,1]
  gl_Position = vec4(a_position, 0, 1);
}

• フルスクリーンクワッド：2つの三角形がビューポートをカバーしています。フラグメント段階のすべてのピクセルは、1つのテンソル要素にマップします。
• 再利用可能：頂点の作業はすべての操作で同一であるため、一度コンパイルし、すべてのマトリックスマルチプリ、アクティベーション、およびバイアスADDパスで再利用します。

この構造を念頭に置いて、すべての「シェーダーパス」は本当に次のとおりです。

1. 頂点シェーダー：ビューポートに2つの三角形をマップします
2. フラグメントシェーダー：ピクセルごとにトランスの数学の小さなピースを1つ実行します

チェーンパス

ボンネットの下では、すべてのニューラルネットワーク操作は、マトリックスの増加、活性化関数、またはバイアスの追加であろうと、4つの単純なGPUステップに沸騰します。

1. 入力をテクスチャ（重み、アクティベーション、または中間結果）としてバインドします。
2. 新鮮な出力テクスチャをオフスクリーンフレームバッファ（FBO）に取り付けます。
3. 共有頂点シェーダーを使用して、フルスクリーンのクワッドを描きます。
4. ピクセルあたりの実際の計算を実行するフラグメントシェーダーを実行します。

すべてのWebglボイラープレートは私たちに住んでいます _runPass ヘルパー、したがって、GPT-2フォワードループの各パスは、単一の宣言的な呼び出しのように感じます。

private _runPass(
  name: string,
  inputs: { [u: string]: WebGLTexture },
  ints: { [u: string]: number },
  outTex: WebGLTexture,
  W: number,
  H: number
) {
  // Grab the WebGL2 context and compiled shader program for this pass
  const gl = this.gl;
  const prog = this.programs[name]; // This has our vertex and frag shaders
  gl.useProgram(prog);

  // BOILERPLATE: Bind all input textures as uniforms

  // BOILERPLATE: Bind an FBO and attach our empty texture to it.

  // BOILERPLATE: Set up the full-screen quad geometry

  // Draw into our texture
  gl.viewport(0, 0, W, H);            // Ensure viewport matches texture size
  gl.drawArrays(gl.TRIANGLES, 0, 6);  // Runs our shaders

  // Clean up: Unbind FBO
  gl.bindFramebuffer(gl.FRAMEBUFFER, null);
}

フォワードパス：レイヤーごと

各パスは結果をVRAMに残すため、最後までCPUにラウンドトリップのコストを支払うことはありません。以下は、フォワードパス全体の高レベルの説明です。

1. 埋め込みをアップロードします：トークン+の位置埋め込みをCPUに計算し、1つのテクスチャとしてGPUに送信します。
2. 変圧器層（合計12）：
   • 正規化とプロジェクト：レイヤーの正規化を適用し、GPUで完全に注意とフィードフォワードサブレーヤーを実行します。
   • 注意：クエリ、キー、値を計算します。注意の重みを計算します。値を結合します。
   • フィードフォワード：2つのマトリックスは、その間にGELUの活性化を掛けます。
   • 残差：各Substepでレイヤーの入力を追加します。
3. 最終的な正規化と出力：最後のレイヤー正規化を行い、出力重量マトリックスを掛け、結果のロジットをCPUに戻します。

ロジットがCPUに戻ったら、ソフトマックスとサンプル（TOP-KまたはTOP-P）を適用して、次のトークンを選択します。その後、新しいトークンがコンテキストに追加されてから、プロセスが再びやり直されます。

これらの操作を接続することにより、最終ロジットまでGPU上にGPT-2パイプライン全体を保持します。これは、プログラム可能なシェーダーがグラフィックパイプラインを汎用の並列エンジンとして扱う方法です。

制限

WebGLをハイジャックすると、GPUで機械学習モデルを実行できますが、いくつかの重要な制限があります。

• 共有/ローカルメモリはありません：フラグメントシェーダーは、グローバルテクスチャのみを読み取り/書き込むことができます。ブロッキングやデータの再利用用のオンチップスクラッチパッドはないため、要素ごとのパスに限定されます。
• テクスチャサイズの制限：GPUは、最大2Dテクスチャディメンション（16 k×16 Kなど）を強制します。より大きなものはすべて、タイルに手動で分割する必要があり、簿記と追加の抽選コールを追加する必要があります。
• 同期や原子はありません：パス内のフラグメント間をバリアまたは調整することはできません。削減、散布/収集、およびその他のデータ依存パターンが困難または不可能です。
• ドローコールと精度のオーバーヘッド：すべてのNeural-net操作には、FBOをバインドし、テクスチャを交換し、CPUオーバーヘッドを負担するドローコール（レイヤーあたり数十）を発行する必要があります。さらに、16または32ビットのフロートに縛られています（経由 EXT_color_buffer_float）、二重精度または整数のテクスチャなし。

まとめると、これらの制約により、シェーダーベースのコンピューティートは興味深い教育プロジェクトになりますが、現実世界での使用のための歴史的な斬新さにすぎません。 CUDAやOpenCLなどのAPIを計算して、同じことを達成するためのより簡単でより良いツールを提供します。

読んでくれてありがとう！コードを表示して、デモをローカルにリポジトリで実行できます ここ。私に連絡してください x 提案がある場合。