Skip to content

[P2] Add nnsight/NDIF backend for interventions - #241

Open
melatg0 wants to merge 1 commit into
stanfordnlp:mainfrom
melatg0:nnsight_integration
Open

[P2] Add nnsight/NDIF backend for interventions#241
melatg0 wants to merge 1 commit into
stanfordnlp:mainfrom
melatg0:nnsight_integration

Conversation

@melatg0

@melatg0 melatg0 commented Jul 2, 2026

Copy link
Copy Markdown

Description

Adds an nnsight/NDIF backend so pyvene interventions can run remotely on NDIF (or locally through nnsight). The main point is being able to intervene on big hosted models like Llama-3.1-405B without a local GPU using the same pyvene API, just remote=True.

  • New IntervenableNdifModel supporting forward, generate, serial interventions, and save/load, plus a torch-only helper so the code passes NDIF's server allowlist.
  • get_remote_weights() on trainable interventions to send their weights for remote runs.
  • Optional install: pip install "pyvene[ndif]". Only the extra needs Python ≥ 3.10 (nnsight 0.7); core pyvene still supports 3.9.
  • Adds a tutorial (ndif_backend_101.ipynb) and per-intervention examples.

Testing Done

Full suite on Python 3.9 / transformers 4.57 (no regressions):

$ python -m unittest discover -p '*TestCase.py'
   Ran 100 tests in 108s
   OK (skipped=9)

The 9 skipped are the opt-in remote NDIF tests. Running them against live NDIF
(Python 3.12, NDIF_REMOTE_TESTS=1, Llama-3.1-8B):


$ NDIF_REMOTE_TESTS=1 python -m unittest tests.integration_tests.NdifBackendTestCase
    Ran 28 tests in 52s
    OK

Also verified NdifBackendCorrectnessTestCase (backend matches HuggingFace and native pyvene exactly, torch.allclose 0.0), and ran ndif_backend_101.ipynb end to end (GPT-2 local + Llama-3.1-405B remote) with no errors.

Checklist:

  • My PR title strictly follows the format: [Your Priority] Your Title
  • I have attached the testing log above
  • I provide enough comments to my code
  • I have changed documentations
  • I have added tests for my changes

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 44aadc063e

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread pyvene/models/ndif_remote_helper.py Outdated
Comment thread pyvene/models/ndif_remote_helper.py Outdated
@melatg0
melatg0 force-pushed the nnsight_integration branch from 44aadc0 to 86e1dd3 Compare July 6, 2026 03:02
@melatg0

melatg0 commented Jul 6, 2026

Copy link
Copy Markdown
Author

@codex review.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 86e1dd383d

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread pyvene/models/intervenable_base.py Outdated
Comment thread pyvene/models/intervenable_base.py
Comment thread pyvene/models/ndif_remote_helper.py Outdated
Add an ndif backend so pyvene interventions can run through nnsight,
locally or remotely on NDIF, letting users intervene on large models
(e.g. Llama-3.1-405B) without a local GPU.

- IntervenableNdifModel: local/remote forward, generate, serial
  interventions, and save/load.
- ndif_remote_helper: import-free (torch-only) implementations of every
  intervention type so NDIF's server-side allowlist accepts them.
- get_remote_weights() on trainable interventions to serialize rotation
  matrices, masks, and autoencoder weights for remote execution.
- nnsight as an optional [ndif] dependency.
- ndif_backend_101 tutorial (GPT-2 local + Llama-3.1-405B remote).
- Integration tests including NdifBackendCorrectnessTestCase, which
  asserts exact allclose against HuggingFace and native pyvene.
- examples/nnsight_examples.py covering each intervention type.
@melatg0
melatg0 force-pushed the nnsight_integration branch from 86e1dd3 to b28797c Compare July 6, 2026 23:45
@melatg0

melatg0 commented Jul 6, 2026

Copy link
Copy Markdown
Author

@codex review.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant