nnsight – 딥러닝 모델 내부를 들여다보고 편집하기

그것이 무엇인가
nnsight는 모든 PyTorch 모델(예: Hugging Face Transformers, Diffusers, vLLM)의 순방향 패스를 추적하고 모델이 실행되는 동안 중간 텐서를 읽거나 수정할 수 있는 Python 라이브러리입니다.
로컬에서 작동하거나 NDIF 서비스를 통해 직접 호스팅할 수 없는 매우 큰 모델에서도 작동합니다.

작동 방식
nnsight로 모델을 래핑하고(예: TransformersModel) with model.trace(...): 블록을 엽니다. 블록 내에서 레이어의 .output에 접근하고 .save()로 값을 읽거나 덮어쓰는 일반 Python 코드를 작성합니다. 라이브러리는 이러한 작업을 실제 순방향 패스에 주입하므로 수동으로 훅을 추가하지 않고도 편집된 텐서를 사용하여 모델이 계산합니다.

주요 기능

• 내부 읽기/쓰기 – 숨겨진 상태, 어텐션 맵, 그래디언트를 검사하고 즉시 교체합니다.
• 생성 훅 – 각 생성 단계에서 개입하고, 로짓을 수집하거나 로짓 렌즈를 실행합니다.
• 배치 단위 추적 – tracer.invoke를 사용하여 한 번의 패스에서 여러 프롬프트를 처리합니다.
• 그래디언트 캡처 – 내부 텐서에 대한 그래디언트를 계산합니다.
• 원격 실행 – 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"):
    # 첫 번째 트랜스포머 블록의 출력을 0으로 설정
    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에서 직접 가져왔습니다.