nnsight – ディープラーニングモデルの内部を覗いて編集する

それは何か
nnsight は、あらゆる PyTorch モデル（例：Hugging Face Transformers、Diffusers、vLLM）のフォワードパスをトレースし、モデル実行中に中間テンソルを読み取ったり変更したりできる Python ライブラリです。
ローカルで動作するか、NDIF サービスを使用して、自分でホストできない非常に大きなモデルでも動作します。

仕組み
nnsight でモデルをラップし（例：TransformersModel）、with model.trace(...): ブロックを開きます。ブロック内では、レイヤーの .output にアクセスし、.save() で値を読み取り、または上書きする通常の Python コードを記述します。ライブラリはこれらの操作を実際のフォワードパスに注入するため、手動でフックを追加することなく、編集したテンソルを使用してモデルが計算します。

主な機能

• 内部の読み書き – 隠れ状態、アテンションマップ、勾配を検査し、その場で置き換えます。
• 生成フック – 各生成ステップで介入し、ロジットを収集するか、ロジットレンズを実行します。
• バッチ単位のトレース – tracer.invoke を使用して、1 回のパスで複数のプロンプトを処理します。
• 勾配キャプチャ – 任意の内部テンソルに関する勾配を計算します。
• リモート実行 – remote=True を設定し、NDIF ホストモデルで同じトレースを実行します。
• 多くのランタイムをサポート – Hugging Face、Diffusers、vLLM 用のラッパーに加え、任意の torch.nn.Module 用の汎用 NNsight ラッパー。

典型的な使用法

from nnsight import TransformersModel
model = TransformersModel("openai-community/gpt2", dispatch=True)

with model.trace("The Eiffel Tower is in the city of"):
    # 最初のトランスフォーマーブロックの出力をゼロにする
    model.transformer.h[0].output[:] = 0
    # ネットワークのより深い隠れ状態を保存する
    hidden = model.transformer.h[6].output.save()
    # 後で使用するために最終ロジットを保持する
    logits = model.output.logits.save()

print(hidden.shape)          # torch.Size([1, 10, 768])
print(logits.argmax(-1))     # トークン予測

同じパターンは、生成、勾配抽出、またはリモート実行にも機能します。

誰が使うか

• モデル内部を調査し、メカニスティック解釈可能性を行い、または「ロジットレンズ」可視化を構築する研究者。
• LLM 用のデバッグや編集ツールを構築する開発者。
• 基盤モデルにカスタム介入を実行する必要があるエージェント（README には、nnsight を Claude、OpenAI Codex などに公開する方法もリストされています）。

詳細情報

• 完全なドキュメント：https://www.nnsight.net
• 設計論文（arXiv 2407.14561）とリポジトリ内の NNsight.md 実装ガイド。
• README にリンクされているインタラクティブな Colab チュートリアル。

インストール

pip install nnsight

---

すべての詳細はリポジトリの README から直接取得されています。