nnsight: Acesso Transparente aos Internals de Redes Neurais

nnsight (/ɛn.saɪt/) permite que pesquisadores interpretem e manipulem os internals de qualquer modelo PyTorch, com a capacidade única de executar o mesmo código localmente em modelos pequenos ou remotamente em modelos massivos (70B+) via NDIF.

GitHub: ndif-team/nnsight (730+ stars) Paper: NNsight and NDIF: Democratizing Access to Foundation Model Internals (ICLR 2025)

Proposta de Valor Principal

Escreva uma vez, execute em qualquer lugar: O mesmo código de interpretabilidade funciona em GPT-2 localmente ou em Llama-3.1-405B remotamente. Basta alternar remote=True.

# Execução local (modelo pequeno)
with model.trace("Hello world"):
    hidden = model.transformer.h[5].output[0].save()

# Execução remota (modelo massivo) - mesmo código!
with model.trace("Hello world", remote=True):
    hidden = model.model.layers[40].output[0].save()

Quando Usar nnsight

Use nnsight quando você precisa:

• Executar experimentos de interpretabilidade em modelos grandes demais para GPUs locais (70B, 405B)
• Trabalhar com qualquer arquitetura PyTorch (transformers, Mamba, modelos customizados)
• Realizar intervenções de geração multi-token
• Compartilhar ativações entre diferentes prompts
• Acessar os internals completos do modelo sem reimplementação

Considere alternativas quando:

• Você quer API consistente entre modelos → Use TransformerLens
• Você precisa de intervenções declarativas e compartilháveis → Use pyvene
• Você está treinando SAEs → Use SAELens
• Você trabalha apenas com modelos pequenos localmente → TransformerLens pode ser mais simples

Instalação

# Instalação básica
pip install nnsight

# Para suporte vLLM
pip install "nnsight[vllm]"

Para execução remota NDIF, inscreva-se em login.ndif.us para obter uma chave de API.

Conceitos Principais

Wrapper LanguageModel

from nnsight import LanguageModel

# Carrega modelo (usa HuggingFace internamente)
model = LanguageModel("openai-community/gpt2", device_map="auto")

# Para modelos maiores
model = LanguageModel("meta-llama/Llama-3.1-8B", device_map="auto")

Contexto de Tracing

O context manager trace permite execução diferida - operações são coletadas em um grafo computacional:

from nnsight import LanguageModel

model = LanguageModel("gpt2", device_map="auto")

with model.trace("The Eiffel Tower is in") as tracer:
    # Acessa a saída de qualquer módulo
    hidden_states = model.transformer.h[5].output[0].save()

    # Acessa padrões de atenção
    attn = model.transformer.h[5].attn.attn_dropout.input[0][0].save()

    # Modifica ativações
    model.transformer.h[8].output[0][:] = 0  # Zera a camada 8

    # Obtém saída final
    logits = model.output.save()

# Após sair do contexto, acessa valores salvos
print(hidden_states.shape)  # [batch, seq, hidden]

Objetos Proxy

Dentro de trace, acessos a módulos retornam objetos Proxy que registram operações:

with model.trace("Hello"):
    # Estes são todos objetos Proxy - operações são diferidas
    h5_out = model.transformer.h[5].output[0]  # Proxy
    h5_mean = h5_out.mean(dim=-1)              # Proxy
    h5_saved = h5_mean.save()                   # Salva para acesso posterior

Workflow 1: Análise de Ativações

Passo a Passo

from nnsight import LanguageModel
import torch

model = LanguageModel("gpt2", device_map="auto")

prompt = "The capital of France is"

with model.trace(prompt) as tracer:
    # 1. Coleta ativações de múltiplas camadas
    layer_outputs = []
    for i in range(12):  # GPT-2 tem 12 camadas
        layer_out = model.transformer.h[i].output[0].save()
        layer_outputs.append(layer_out)

    # 2. Obtém padrões de atenção
    attn_patterns = []
    for i in range(12):
        # Acessa pesos de atenção (após softmax)
        attn = model.transformer.h[i].attn.attn_dropout.input[0][0].save()
        attn_patterns.append(attn)

    # 3. Obtém logits finais
    logits = model.output.save()

# 4. Analisa fora do contexto
for i, layer_out in enumerate(layer_outputs):
    print(f"Saída camada {i}: {layer_out.shape}")
    print(f"Norma camada {i}: {layer_out.norm().item():.3f}")

# 5. Encontra principais predições
probs = torch.softmax(logits[0, -1], dim=-1)
top_tokens = probs.topk(5)
for token, prob in zip(top_tokens.indices, top_tokens.values):
    print(f"{model.tokenizer.decode(token)}: {prob.item():.3f}")

Checklist

• Carrega modelo com wrapper LanguageModel
• Usa contexto trace para operações
• Chama .save() em valores que você precisa após o contexto
• Acessa valores salvos fora do contexto
• Usa .shape, .norm(), etc. para análise

Workflow 2: Patching de Ativações

Passo a Passo

from nnsight import LanguageModel
import torch

model = LanguageModel("gpt2", device_map="auto")

prompt_limpo = "The Eiffel Tower is in"
prompt_corrompido = "The Colosseum is in"

# 1. Obtém ativações limpas
with model.trace(prompt_limpo) as tracer:
    hidden_limpo = model.transformer.h[8].output[0].save()

# 2. Faz patch limpo em execução corrompida
with model.trace(prompt_corrompido) as tracer:
    # Substitui saída da camada 8 com ativações limpas
    model.transformer.h[8].output[0][:] = hidden_limpo

    logits_patch = model.output.save()

# 3. Compara predições
paris_token = model.tokenizer.encode(" Paris")[0]
rome_token = model.tokenizer.encode(" Rome")[0]

probs_patch = torch.softmax(logits_patch[0, -1], dim=-1)
print(f"Prob Paris: {probs_patch[paris_token].item():.3f}")
print(f"Prob Roma: {probs_patch[rome_token].item():.3f}")

Varredura Sistemática de Patching

def patch_layer_position(layer, position, clean_cache, prompt_corrompido):
    """Faz patch de camada/posição única de limpo para corrompido."""
    with model.trace(prompt_corrompido) as tracer:
        # Obtém ativação atual
        current = model.transformer.h[layer].output[0]

        # Faz patch apenas de posição específica
        current[:, position, :] = clean_cache[layer][:, position, :]

        logits = model.output.save()

    return logits

# Varre todas as camadas e posições
results = torch.zeros(12, seq_len)
for layer in range(12):
    for pos in range(seq_len):
        logits = patch_layer_position(layer, pos, hidden_limpo, prompt_corrompido)
        results[layer, pos] = compute_metric(logits)

Workflow 3: Execução Remota com NDIF

Execute os mesmos experimentos em modelos massivos sem GPUs locais.

Passo a Passo

from nnsight import LanguageModel

# 1. Carrega modelo grande (será executado remotamente)
model = LanguageModel("meta-llama/Llama-3.1-70B")

# 2. Mesmo código, basta adicionar remote=True
with model.trace("The meaning of life is", remote=True) as tracer:
    # Acessa internals do modelo 70B!
    layer_40_out = model.model.layers[40].output[0].save()
    logits = model.output.save()

# 3. Resultados retornados de NDIF
print(f"Forma camada 40: {layer_40_out.shape}")

# 4. Geração com intervenções
with model.trace(remote=True) as tracer:
    with tracer.invoke("What is 2+2?"):
        # Intervém durante geração
        model.model.layers[20].output[0][:, -1, :] *= 1.5

    output = model.generate(max_new_tokens=50)

Setup NDIF

1. Inscreva-se em login.ndif.us
2. Obtenha chave de API
3. Defina variável de ambiente ou passe para nnsight:

import os
os.environ["NDIF_API_KEY"] = "your_key"

# Ou configure diretamente
from nnsight import CONFIG
CONFIG.API_KEY = "your_key"

Modelos Disponíveis em NDIF

• Llama-3.1-8B, 70B, 405B
• Modelos DeepSeek-R1
• Vários modelos open-weight (verifique ndif.us para lista atual)

Workflow 4: Compartilhamento de Ativações Entre Prompts

Compartilha ativações entre diferentes inputs em um único trace.

from nnsight import LanguageModel

model = LanguageModel("gpt2", device_map="auto")

with model.trace() as tracer:
    # Primeiro prompt
    with tracer.invoke("The cat sat on the"):
        hidden_gato = model.transformer.h[6].output[0].save()

    # Segundo prompt - injeta ativações do gato
    with tracer.invoke("The dog ran through the"):
        # Substitui com ativações do gato na camada 6
        model.transformer.h[6].output[0][:] = hidden_gato
        cachorro_com_gato = model.output.save()

# O prompt do cachorro agora tem representações internas do gato

Workflow 5: Análise Baseada em Gradientes

Acessa gradientes durante backward pass.

from nnsight import LanguageModel
import torch

model = LanguageModel("gpt2", device_map="auto")

with model.trace("The quick brown fox") as tracer:
    # Salva ativações e habilita gradiente
    hidden = model.transformer.h[5].output[0].save()
    hidden.retain_grad()

    logits = model.output

    # Computa loss em token específico
    target_token = model.tokenizer.encode(" jumps")[0]
    loss = -logits[0, -1, target_token]

    # Backward pass
    loss.backward()

# Acessa gradientes
grad = hidden.grad
print(f"Forma gradiente: {grad.shape}")
print(f"Norma gradiente: {grad.norm().item():.3f}")

Nota: Acesso a gradientes não é suportado para vLLM ou execução remota.

Problemas Comuns e Soluções

Problema: Caminho do módulo diferente entre modelos

# Estrutura GPT-2
model.transformer.h[5].output[0]

# Estrutura LLaMA
model.model.layers[5].output[0]

# Solução: Verifica estrutura do modelo
print(model._model)  # Vê nomes reais dos módulos

Problema: Esqueceu de salvar

# ERRADO: Valor não acessível fora do trace
with model.trace("Hello"):
    hidden = model.transformer.h[5].output[0]  # Não salvo!

print(hidden)  # Erro ou valor errado

# CORRETO: Chama .save()
with model.trace("Hello"):
    hidden = model.transformer.h[5].output[0].save()

print(hidden)  # Funciona!

Problema: Timeout remoto

# Para operações longas, aumenta timeout
with model.trace("prompt", remote=True, timeout=300) as tracer:
    # Operação longa...

Problema: Memória com muitas ativações salvas

# Salva apenas o que você precisa
with model.trace("prompt"):
    # Não salva tudo
    for i in range(100):
        model.transformer.h[i].output[0].save()  # Pesado em memória!

    # Melhor: salva camadas específicas
    camadas_chave = [0, 5, 11]
    for i in camadas_chave:
        model.transformer.h[i].output[0].save()

Problema: Limitação de gradiente vLLM

# vLLM não suporta gradientes
# Use execução padrão para análise de gradientes
model = LanguageModel("gpt2", device_map="auto")  # Não vLLM

Referência de API Chave

Método/Propriedade	Propósito
model.trace(prompt, remote=False)	Inicia contexto de tracing
proxy.save()	Salva valor para acesso após trace
proxy[:]	Slice/índice de proxy (assignment faz patch)
tracer.invoke(prompt)	Adiciona prompt dentro de trace
model.generate(...)	Gera com intervenções
model.output	Logits de saída final do modelo
model._model	Modelo HuggingFace subjacente

Comparação com Outras Ferramentas

Funcionalidade	nnsight	TransformerLens	pyvene
Qualquer arquitetura	Sim	Apenas Transformers	Sim
Execução remota	Sim (NDIF)	Não	Não
API consistente	Não	Sim	Sim
Execução diferida	Sim	Não	Não
HuggingFace nativo	Sim	Reimplementado	Sim
Configs compartilháveis	Não	Não	Sim

Documentação de Referência

Para documentação detalhada de API, tutoriais e uso avançado, veja a pasta references/:

Arquivo	Conteúdo
references/README.md	Visão geral e guia de início rápido
references/api.md	Referência completa de API para LanguageModel, tracing, objetos proxy
references/tutorials.md	Tutoriais passo a passo para interpretabilidade local e remota

Recursos Externos

Tutoriais

• Getting Started
• Features Overview
• Remote Execution
• Applied Tutorials

Documentação Oficial

• Official Docs
• NDIF Info
• Community Forum

Artigos

• NNsight and NDIF Paper - Fiotto-Kaufman et al. (ICLR 2025)

Suporte a Arquiteturas

nnsight funciona com qualquer modelo PyTorch:

• Transformers: GPT-2, LLaMA, Mistral, etc.
• Modelos State Space: Mamba
• Modelos Vision: ViT, CLIP
• Arquiteturas customizadas: Qualquer nn.Module

A chave é conhecer a estrutura do módulo para acessar os componentes certos.