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): …`	여러 프롬프트를 병렬 스레드로 실행; 각 호출은 순차적으로 실행되지만 `.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'):
    # 첫 번째 레이어의 활성화를 0으로 설정
    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 블록이 토큰 예측에 어떻게 기여하는지 분석.
인과 탐색 – 활성화를 0으로, 노이즈를 추가하거나 교체하여 정보 흐름에 대한 가설 검증.
모델 편집 연구 – 역으로 되돌릴 수 있는 편집(예: "에펠 탑" 이후 항상 "파리"를 출력하게 하기).
사용자 정의 아키텍처 디버깅 – 임의의 torch.nn.Module에서 동일한 API를 사용해 포워드 패스 동작 검증.
효율적인 배치 실험 – invoke/session 기계를 사용해 한 번의 포워드 패스로 여러 프롬프트 실행.

제한 사항 및 주의사항 (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스러운 방법을 연구자에게 제공합니다. 현대 기초 모델에 대한 기계적 작업을 위한 진정한 연구용 도구입니다.