ndif-team/nnsight

The nnsight package enables interpreting and manipulating the internals of deep learned models.

nnsight – PyTorchモデルの内部を解釈・編集する

何ができるか – nnsight は pip install nnsight でインストール可能なPythonライブラリで、任意のPyTorchモデル（GPT‑2、LLaMAなど）の内部状態にアクセスし、各レイヤーの隠れ状態テンソルを読み取ったり、実行中に変更したり、中間値の勾配を計算したり、モデルの永続的編集を可能にします。ローカルのGPU/CPUで動作し、非常に大きなモデルの場合はNDインスティテュートのリモートインフラ上で実行することもできます。

なぜ重要か – 現代の基礎モデル（GPT‑2、LLaMAなど）はブラックボックスです。特定の予測の理由を理解したい、因果仮説を検証したい、あるいはモデル編集技術をプロトタイピングしたい研究者にとって、モデルのコードを書き換えることなく内部状態にアクセス・操作できるクリーンな方法が必要です。nnsight は、フック挿入やトレースのボイラープレートを抽象化した、高レベルでPythonらしいAPIを提供します。

コア機能（READMEに記載）

機能	使用方法	結果
活性化のアクセス	`with model.trace(prompt): hidden = model.transformer.h[5].output[0].save()`	指定されたプロンプトに対するレイヤー5の隠れ状態を含む実際のテンソル。
インプレース介入	`model.transformer.h[0].output[0][:] = 0` をトレースブロック内で	変更された活性化でフォワードパスが継続され、因果効果をテスト可能。
中間テンソルの勾配	`with loss.backward(): grad = hs.grad.save()`	定義した損失関数に対して任意のテンソル（例：隠れ状態）の勾配を取得。
バッチレベルの呼び出し	`with tracer.invoke(prompt): …`	複数のプロンプトを並列スレッドで実行；各invokeは順次実行され、`.save()`で値を共有。
ステップ単位制御付き生成	`with model.generate(prompt, max_new_tokens=5) as tracer: for step in tracer.iter[:]: …`	あなたが特定のステップで介入できる自己回帰生成。
モデル編集	`with model.edit() as edited: edited.transformer.h[0].output[0][:] = 0`	編集を永続的に反映した新しい `LanguageModel` インスタンスを生成し、元のモデルは変更されない。
形状のみのスキャン	`with model.scan(prompt): dim = nnsight.save(layer.output.shape[-1])`	完全なフォワードパスなしでテンソルの形状を取得。
キャッシュとセッション	`cache = tracer.cache()` または `with model.session() as s: …`	複数のトレース間で以前にキャプチャされた活性化を再利用し、効率を向上。
リモート実行	`CONFIG.set_default_api_key(..); model = LanguageModel('meta-llama/Meta-Llama-3.1-8B'); with model.trace(..., remote=True): …`	トレースコードをNDIFのクラウドサービス上で実行可能。ローカルでは収まらないモデルに有用。
vLLM統合	`from nnsight.modeling.vllm import VLLM; model = VLLM('gpt2', ...)`	同じトレースAPIを維持しつつ、高スループットの推論を実現。
任意のPyTorchモデル	`NNsight(net)` ここで `net` は任意の `torch.nn.Module`	Transformer以外のモデル（例：単純なフィードフォワードネット）にも同じトレース・介入ツールが利用可能。

クイックスタート例（READMEより）

from nnsight import LanguageModel

model = LanguageModel('openai-community/gpt2', device_map='auto', dispatch=True)

with model.trace('The Eiffel Tower is in the city of'):
    # 最初のレイヤーの活性化をゼロに
    model.transformer.h[0].output[0][:] = 0
    # 最終隠れ状態とロジットを保存
    hidden = model.transformer.h[-1].output[0].save()
    logits = model.output.save()

print(model.tokenizer.decode(logits.logits.argmax(dim=-1)[0]))

このスニペットはモデルの読み込み、トレースの開始、レイヤーへの介入、最終予測の取得を示しています。

一般的な利用例

機械的解釈 – アテンションヘッドやMLPブロックがトークン予測にどのように寄与しているかを検証。
因果プロービング – 活性化をゼロに、ノイズを加え、または置き換えることで情報フローに関する仮説を検証。
モデル編集研究 – 可逆的な編集（例：「エッフェル塔」の後に「パリ」を常に出力させる）を実装。
カスタムアーキテクチャのデバッグ – 任意の torch.nn.Module で同じAPIを使用し、フォワードパスの挙動を検証。
効率的なバッチ実験 – invoke/session機構を使って、1回のフォワードパスで多数のプロンプトを実行。

制限事項と注意点（READMEに記載）

実行順序が重要 – トレース内では、モジュールを実行順にアクセスする必要があり、そうでないと OutOfOrderError のデッドロックに陥る。
スレッドベースの同期 – ライブラリはトレースコードを別ワーカースレッドで実行；.save()を呼び出すまで値は利用不可。
無限イテレータ – tracer.iter[:] は決して返らないため、その後のコードは invoke ブロック内に置く必要がある。
リモート実行にはAPIキーが必要 かつ、モデルがNDIFプラットフォーム上で利用可能でなければならない。
vLLM統合はvLLMがサポートするモデルに限定（現時点ではTransformerスタイルの言語モデルのみ）。

詳細情報の入手先

ドキュメントサイト – https://www.nnsight.net
論文 – NNsight and NDIF: Democratizing Access to Foundation Model Internals (arXiv 2407.14561)
Discord & フォーラム – READMEにリンクされたコミュニティサポート・ディスカッションチャンネル。
Colabウォークスルー – インタラクティブなノートブックで実際に試せる。

まとめ（TL;DR）

nnsight は、バッチ処理、勾配抽出、リモート実行、vLLMのような高性能バックエンドをサポートしながら、任意のPyTorchモデルの隠れ状態をトレース、読み取り、変更、永続的に編集するための簡潔でPythonらしい方法を研究者に提供します。現代の基礎モデルに対する機械的作業に特化した本格的なツールです。