[P2] Add nnsight/NDIF backend for interventions - #241
Conversation
There was a problem hiding this comment.
💡 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".
44aadc0 to
86e1dd3
Compare
|
@codex review. |
There was a problem hiding this comment.
💡 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".
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.
86e1dd3 to
b28797c
Compare
|
@codex review. |
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.IntervenableNdifModelsupporting 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.pip install "pyvene[ndif]". Only the extra needs Python ≥ 3.10 (nnsight 0.7); core pyvene still supports 3.9.ndif_backend_101.ipynb) and per-intervention examples.Testing Done
Full suite on Python 3.9 / transformers 4.57 (no regressions):
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):
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:
[Your Priority] Your Title