{"id":"nnsight","name":"nnsight-remote-interpretability","summary":"nnsightを用いてニューラルネットワーク内部の解釈と操作のガイダンスを提供し、オプションでNDIFリモート実行も可能です。","body":"# nnsight: Transparent Access to Neural Network Internals\n\nnnsight (/ɛn.saɪt/) enables researchers to interpret and manipulate the internals of any PyTorch model, with the unique capability of running the same code locally on small models or remotely on massive models (70B+) via NDIF.\n\n**GitHub**: [ndif-team/nnsight](https://github.com/ndif-team/nnsight) (730+ stars)\n**Paper**: [NNsight and NDIF: Democratizing Access to Foundation Model Internals](https://arxiv.org/abs/2407.14561) (ICLR 2025)\n\n## Key Value Proposition\n\n**Write once, run anywhere**: The same interpretability code works on GPT-2 locally or Llama-3.1-405B remotely. Just toggle `remote=True`.\n\n```python\n# Local execution (small model)\nwith model.trace(\"Hello world\"):\n    hidden = model.transformer.h[5].output[0].save()\n\n# Remote execution (massive model) - same code!\nwith model.trace(\"Hello world\", remote=True):\n    hidden = model.model.layers[40].output[0].save()\n```\n\n## When to Use nnsight\n\n**Use nnsight when you need to:**\n- Run interpretability experiments on models too large for local GPUs (70B, 405B)\n- Work with any PyTorch architecture (transformers, Mamba, custom models)\n- Perform multi-token generation interventions\n- Share activations between different prompts\n- Access full model internals without reimplementation\n\n**Consider alternatives when:**\n- You want consistent API across models → Use **TransformerLens**\n- You need declarative, shareable interventions → Use **pyvene**\n- You're training SAEs → Use **SAELens**\n- You only work with small models locally → **TransformerLens** may be simpler\n\n## Installation\n\n```bash\n# Basic installation\npip install nnsight\n\n# For vLLM support\npip install \"nnsight[vllm]\"\n```\n\nFor remote NDIF execution, sign up at [login.ndif.us](https://login.ndif.us) for an API key.\n\n## Core Concepts\n\n### LanguageModel Wrapper\n\n```python\nfrom nnsight import LanguageModel\n\n# Load model (uses HuggingFace under the hood)\nmodel = LanguageModel(\"openai-community/gpt2\", device_map=\"auto\")\n\n# For larger models\nmodel = LanguageModel(\"meta-llama/Llama-3.1-8B\", device_map=\"auto\")\n```\n\n### Tracing Context\n\nThe `trace` context manager enables deferred execution - operations are collected into a computation graph:\n\n```python\nfrom nnsight import LanguageModel\n\nmodel = LanguageModel(\"gpt2\", device_map=\"auto\")\n\nwith model.trace(\"The Eiffel Tower is in\") as tracer:\n    # Access any module's output\n    hidden_states = model.transformer.h[5].output[0].save()\n\n    # Access attention patterns\n    attn = model.transformer.h[5].attn.attn_dropout.input[0][0].save()\n\n    # Modify activations\n    model.transformer.h[8].output[0][:] = 0  # Zero out layer 8\n\n    # Get final output\n    logits = model.output.save()\n\n# After context exits, access saved values\nprint(hidden_states.shape)  # [batch, seq, hidden]\n```\n\n### Proxy Objects\n\nInside `trace`, module accesses return Proxy objects that record operations:\n\n```python\nwith model.trace(\"Hello\"):\n    # These are all Proxy objects - operations are deferred\n    h5_out = model.transformer.h[5].output[0]  # Proxy\n    h5_mean = h5_out.mean(dim=-1)              # Proxy\n    h5_saved = h5_mean.save()                   # Save for later access\n```\n\n## Workflow 1: Activation Analysis\n\n### Step-by-Step\n\n```python\nfrom nnsight import LanguageModel\nimport torch\n\nmodel = LanguageModel(\"gpt2\", device_map=\"auto\")\n\nprompt = \"The capital of France is\"\n\nwith model.trace(prompt) as tracer:\n    # 1. Collect activations from multiple layers\n    layer_outputs = []\n    for i in range(12):  # GPT-2 has 12 layers\n        layer_out = model.transformer.h[i].output[0].save()\n        layer_outputs.append(layer_out)\n\n    # 2. Get attention patterns\n    attn_patterns = []\n    for i in range(12):\n        # Access attention weights (after softmax)\n        attn = model.transformer.h[i].attn.attn_dropout.input[0][0].save()\n        attn_patterns.append(attn)\n\n    # 3. Get final logits\n    logits = model.output.save()\n\n# 4. Analyze outside context\nfor i, layer_out in enumerate(layer_outputs):\n    print(f\"Layer {i} output shape: {layer_out.shape}\")\n    print(f\"Layer {i} norm: {layer_out.norm().item():.3f}\")\n\n# 5. Find top predictions\nprobs = torch.softmax(logits[0, -1], dim=-1)\ntop_tokens = probs.topk(5)\nfor token, prob in zip(top_tokens.indices, top_tokens.values):\n    print(f\"{model.tokenizer.decode(token)}: {prob.item():.3f}\")\n```\n\n### Checklist\n- [ ] Load model with LanguageModel wrapper\n- [ ] Use trace context for operations\n- [ ] Call `.save()` on values you need after context\n- [ ] Access saved values outside context\n- [ ] Use `.shape`, `.norm()`, etc. for analysis\n\n## Workflow 2: Activation Patching\n\n### Step-by-Step\n\n```python\nfrom nnsight import LanguageModel\nimport torch\n\nmodel = LanguageModel(\"gpt2\", device_map=\"auto\")\n\nclean_prompt = \"The Eiffel Tower is in\"\ncorrupted_prompt = \"The Colosseum is in\"\n\n# 1. Get clean activations\nwith model.trace(clean_prompt) as tracer:\n    clean_hidden = model.transformer.h[8].output[0].save()\n\n# 2. Patch clean into corrupted run\nwith model.trace(corrupted_prompt) as tracer:\n    # Replace layer 8 output with clean activations\n    model.transformer.h[8].output[0][:] = clean_hidden\n\n    patched_logits = model.output.save()\n\n# 3. Compare predictions\nparis_token = model.tokenizer.encode(\" Paris\")[0]\nrome_token = model.tokenizer.encode(\" Rome\")[0]\n\npatched_probs = torch.softmax(patched_logits[0, -1], dim=-1)\nprint(f\"Paris prob: {patched_probs[paris_token].item():.3f}\")\nprint(f\"Rome prob: {patched_probs[rome_token].item():.3f}\")\n```\n\n### Systematic Patching Sweep\n\n```python\ndef patch_layer_position(layer, position, clean_cache, corrupted_prompt):\n    \"\"\"Patch single layer/position from clean to corrupted.\"\"\"\n    with model.trace(corrupted_prompt) as tracer:\n        # Get current activation\n        current = model.transformer.h[layer].output[0]\n\n        # Patch only specific position\n        current[:, position, :] = clean_cache[layer][:, position, :]\n\n        logits = model.output.save()\n\n    return logits\n\n# Sweep over all layers and positions\nresults = torch.zeros(12, seq_len)\nfor layer in range(12):\n    for pos in range(seq_len):\n        logits = patch_layer_position(layer, pos, clean_hidden, corrupted)\n        results[layer, pos] = compute_metric(logits)\n```\n\n## Workflow 3: Remote Execution with NDIF\n\nRun the same experiments on massive models without local GPUs.\n\n### Step-by-Step\n\n```python\nfrom nnsight import LanguageModel\n\n# 1. Load large model (will run remotely)\nmodel = LanguageModel(\"meta-llama/Llama-3.1-70B\")\n\n# 2. Same code, just add remote=True\nwith model.trace(\"The meaning of life is\", remote=True) as tracer:\n    # Access internals of 70B model!\n    layer_40_out = model.model.layers[40].output[0].save()\n    logits = model.output.save()\n\n# 3. Results returned from NDIF\nprint(f\"Layer 40 shape: {layer_40_out.shape}\")\n\n# 4. Generation with interventions\nwith model.trace(remote=True) as tracer:\n    with tracer.invoke(\"What is 2+2?\"):\n        # Intervene during generation\n        model.model.layers[20].output[0][:, -1, :] *= 1.5\n\n    output = model.generate(max_new_tokens=50)\n```\n\n### NDIF Setup\n\n1. Sign up at [login.ndif.us](https://login.ndif.us)\n2. Get API key\n3. Set environment variable or pass to nnsight:\n\n```python\nimport os\nos.environ[\"NDIF_API_KEY\"] = \"your_key\"\n\n# Or configure directly\nfrom nnsight import CONFIG\nCONFIG.API_KEY = \"your_key\"\n```\n\n### Available Models on NDIF\n\n- Llama-3.1-8B, 70B, 405B\n- DeepSeek-R1 models\n- Various open-weight models (check [ndif.us](https://ndif.us) for current list)\n\n## Workflow 4: Cross-Prompt Activation Sharing\n\nShare activations between different inputs in a single trace.\n\n```python\nfrom nnsight import LanguageModel\n\nmodel = LanguageModel(\"gpt2\", device_map=\"auto\")\n\nwith model.trace() as tracer:\n    # First prompt\n    with tracer.invoke(\"The cat sat on the\"):\n        cat_hidden = model.transformer.h[6].output[0].save()\n\n    # Second prompt - inject cat's activations\n    with tracer.invoke(\"The dog ran through the\"):\n        # Replace with cat's activations at layer 6\n        model.transformer.h[6].output[0][:] = cat_hidden\n        dog_with_cat = model.output.save()\n\n# The dog prompt now has cat's internal representations\n```\n\n## Workflow 5: Gradient-Based Analysis\n\nAccess gradients during backward pass.\n\n```python\nfrom nnsight import LanguageModel\nimport torch\n\nmodel = LanguageModel(\"gpt2\", device_map=\"auto\")\n\nwith model.trace(\"The quick brown fox\") as tracer:\n    # Save activations and enable gradient\n    hidden = model.transformer.h[5].output[0].save()\n    hidden.retain_grad()\n\n    logits = model.output\n\n    # Compute loss on specific token\n    target_token = model.tokenizer.encode(\" jumps\")[0]\n    loss = -logits[0, -1, target_token]\n\n    # Backward pass\n    loss.backward()\n\n# Access gradients\ngrad = hidden.grad\nprint(f\"Gradient shape: {grad.shape}\")\nprint(f\"Gradient norm: {grad.norm().item():.3f}\")\n```\n\n**Note**: Gradient access not supported for vLLM or remote execution.\n\n## Common Issues & Solutions\n\n### Issue: Module path differs between models\n```python\n# GPT-2 structure\nmodel.transformer.h[5].output[0]\n\n# LLaMA structure\nmodel.model.layers[5].output[0]\n\n# Solution: Check model structure\nprint(model._model)  # See actual module names\n```\n\n### Issue: Forgetting to save\n```python\n# WRONG: Value not accessible outside trace\nwith model.trace(\"Hello\"):\n    hidden = model.transformer.h[5].output[0]  # Not saved!\n\nprint(hidden)  # Error or wrong value\n\n# RIGHT: Call .save()\nwith model.trace(\"Hello\"):\n    hidden = model.transformer.h[5].output[0].save()\n\nprint(hidden)  # Works!\n```\n\n### Issue: Remote timeout\n```python\n# For long operations, increase timeout\nwith model.trace(\"prompt\", remote=True, timeout=300) as tracer:\n    # Long operation...\n```\n\n### Issue: Memory with many saved activations\n```python\n# Only save what you need\nwith model.trace(\"prompt\"):\n    # Don't save everything\n    for i in range(100):\n        model.transformer.h[i].output[0].save()  # Memory heavy!\n\n    # Better: save specific layers\n    key_layers = [0, 5, 11]\n    for i in key_layers:\n        model.transformer.h[i].output[0].save()\n```\n\n### Issue: vLLM gradient limitation\n```python\n# vLLM doesn't support gradients\n# Use standard execution for gradient analysis\nmodel = LanguageModel(\"gpt2\", device_map=\"auto\")  # Not vLLM\n```\n\n## Key API Reference\n\n| Method/Property | Purpose |\n|-----------------|---------|\n| `model.trace(prompt, remote=False)` | Start tracing context |\n| `proxy.save()` | Save value for access after trace |\n| `proxy[:]` | Slice/index proxy (assignment patches) |\n| `tracer.invoke(prompt)` | Add prompt within trace |\n| `model.generate(...)` | Generate with interventions |\n| `model.output` | Final model output logits |\n| `model._model` | Underlying HuggingFace model |\n\n## Comparison with Other Tools\n\n| Feature | nnsight | TransformerLens | pyvene |\n|---------|---------|-----------------|--------|\n| Any architecture | Yes | Transformers only | Yes |\n| Remote execution | Yes (NDIF) | No | No |\n| Consistent API | No | Yes | Yes |\n| Deferred execution | Yes | No | No |\n| HuggingFace native | Yes | Reimplemented | Yes |\n| Shareable configs | No | No | Yes |\n\n## Reference Documentation\n\nFor detailed API documentation, tutorials, and advanced usage, see the `references/` folder:\n\n| File | Contents |\n|------|----------|\n| [references/README.md](references/README.md) | Overview and quick start guide |\n| [references/api.md](references/api.md) | Complete API reference for LanguageModel, tracing, proxy objects |\n| [references/tutorials.md](references/tutorials.md) | Step-by-step tutorials for local and remote interpretability |\n\n## External Resources\n\n### Tutorials\n- [Getting Started](https://nnsight.net/start/)\n- [Features Overview](https://nnsight.net/features/)\n- [Remote Execution](https://nnsight.net/notebooks/features/remote_execution/)\n- [Applied Tutorials](https://nnsight.net/applied_tutorials/)\n\n### Official Documentation\n- [Official Docs](https://nnsight.net/documentation/)\n- [NDIF Info](https://ndif.us/)\n- [Community Forum](https://discuss.ndif.us/)\n\n### Papers\n- [NNsight and NDIF Paper](https://arxiv.org/abs/2407.14561) - Fiotto-Kaufman et al. (ICLR 2025)\n\n## Architecture Support\n\nnnsight works with any PyTorch model:\n- **Transformers**: GPT-2, LLaMA, Mistral, etc.\n- **State Space Models**: Mamba\n- **Vision Models**: ViT, CLIP\n- **Custom architectures**: Any nn.Module\n\nThe key is knowing the module structure to access the right components.","author":"@Orchestra-Research","ownerProfile":null,"authorContacts":null,"sourceUrl":"https://github.com/Orchestra-Research/AI-Research-SKILLs/tree/main/04-mechanistic-interpretability/nnsight","license":"MIT","category":"writing","lang":"en","tokens":3224,"stars":0,"calls30d":1,"claimed":false,"visibility":"public","origin":"crawler","version":"0.1.0","createdAt":"2026-08-22","updatedAt":"2026-08-22","files":[{"path":"references/api.md","size":7160,"sha256":"750fe5a6b7092c839a0602d2d00c13cc9168fb0bdb73eb3279441cda7aef5b39"},{"path":"references/README.md","size":1913,"sha256":"8cc09397004cb4855704dd490ea1993dee5935aea52f0e02eb5b960c4da733b2"},{"path":"references/tutorials.md","size":7970,"sha256":"a46505bde75ab0c4703fdf6fd0248b4d02820d519f763a3324e44d29fca0333e"}],"requires":{"mcp":[],"tools":[]},"safety":{"flags":[],"scannedAt":"2026-08-22","hasScripts":false,"networkEndpoints":["arxiv.org","discuss.ndif.us","login.ndif.us","ndif.us","nnsight.net"]}}