Remote Execution on NDIF¶
NDIF (the National Deep Inference Fabric) is the hosted service that runs nnsight
intervention code on shared GPU pods. You write the same trace you would run
locally, add remote=True, and nnsight serializes the traced block, ships it to a
server that holds the real weights, runs the forward pass with your interventions,
and streams the .save()d results back — no local GPU required.
This lets you probe models far larger than your hardware could hold (Llama-3.1-70B,
405B, DeepSeek, ...): the same code you write against GPT-2 locally scales up by
swapping the model id and adding remote=True.
About this notebook. Every runnable cell below is executed for real against a live NDIF instance with
remote=True— the outputs you see came back over the wire from a server that holds the weights. Set your API key in the setup cell to run them yourself (see Getting an NDIF account below). A handful of cells that need a specific large model or a separate importable module are shown as illustrative code blocks and marked as such.
Setup¶
In 0.8 the primary HuggingFace wrapper is TransformersModel (the older
LanguageModel still works but warns on construction — it is now a thin deprecated
alias for TransformersModel(task="text-generation")).
CONFIG.API.HOST already points at the public service (https://api.ndif.us), so
the only thing to set is your API key. Point it elsewhere — a lab deployment, a local
dev server — with CONFIG.API.HOST, the NDIF_HOST environment variable, or by
passing the URL as remote= on a single call.
import nnsight
from nnsight import TransformersModel, CONFIG
CONFIG.API.APIKEY = "YOUR_API_KEY" # from login.ndif.us -- see below
# The live status display uses in-place terminal updates that render as noise in a
# saved notebook, so we silence it here; leave it on (the default) for interactive use.
CONFIG.APP.REMOTE_LOGGING = False
/home/localjadenfk/miniconda3/envs/ndif2/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
Getting an NDIF Account¶
Create a free account at login.ndif.us to get your API key. Every remote request is keyed against it. There are a few ways to set it.
Option 1 — Config (persistent, recommended once per machine).
set_default_api_key sets CONFIG.API.APIKEY and writes it to
~/.config/nnsight/config.yaml, so you only do this once:
from nnsight import CONFIG
CONFIG.set_default_api_key("YOUR_API_KEY")
Option 2 — the login helper (HuggingFace-style).
login() prompts for the key with getpass (never echoed), verifies it against the
service, and persists it. whoami() reports the identity behind the stored key:
from nnsight import login, whoami
login() # prompts; or login("YOUR_KEY") to skip the prompt
whoami() # {"email": ..., "tags": [...]}
The same commands work from a terminal: nnsight login and nnsight whoami.
Option 3 — environment variable.
export NDIF_API_KEY="YOUR_API_KEY"
Option 4 — Google Colab. Store the key as a Colab Userdata secret named
NDIF_API_KEY; nnsight reads it automatically.
Checking Available Models¶
A remote request only runs against a model that is deployed on NDIF. Check what is
live before submitting with nnsight.status() (the older nnsight.ndif_status()
still works but is deprecated), or visit the status page.
nnsight.is_model_running(...) answers the yes/no question for one repo id.
print(nnsight.is_model_running("openai-community/gpt2"))
print(nnsight.status())
True NDIF Service: Up 🟢 Model Class Repo ID Revision Level State ----------------- --------------------- -------- ----- ------- TransformersModel openai-community/gpt2 main HOT RUNNING
nnsight.status() returns an NdifStatus — a dict-like view over the deployments, so
you can also inspect it programmatically (for repo_id in s:, s[repo_id],
s.status).
Comparing Environments¶
The most common cause of "works locally, fails remotely" is a package version
mismatch — your block is serialized against your local interpreter, and the server
unpickles and runs it against its own. nnsight.compare() diffs your local
environment against NDIF's, flagging the critical packages (nnsight, torch,
transformers):
import nnsight
print(nnsight.compare())
Package Local Version Remote Version Status
------------ ------------- -------------- ----------
nnsight 0.8.0 0.8.0 ✓
torch 2.6.0 2.6.0 ✓
transformers 4.45.0 4.45.0 ✓
Because nnsight serializes intervention code by source (not by pickle reference),
you don't need the same Python version as the server — the block is rebuilt from
source on the other side. Diverging torch/transformers (or a helper module whose
local version differs from the server's) are the risk to watch. The request also
carries persistent ids that name objects by where they sit in your local module
tree; when that tree doesn't line up with the server's, one fails to resolve and the
run errors server-side (a persistent-id lookup failure). compare() is the quick
check before you submit.
Loading a Model for Remote Use¶
When you instantiate a model for remote use you do not pass dispatch=True (and
you don't need device_map — there are no local weights to place): nnsight builds a
lightweight skeleton on the meta device — the architecture is constructed (so
model.transformer.h[0].output is a real path you can hook) but no weights are
allocated and nothing is downloaded. That is what lets a machine with no GPU write
intervention code against a 405B model in seconds.
Here we build GPT-2 this way and use it for every remote call below. To target a
model too big to hold locally, you would swap in its id exactly the same way —
TransformersModel("meta-llama/Llama-3.1-70B").
model = TransformersModel("openai-community/gpt2") # built on the meta device
Basic Remote Tracing¶
Add remote=True to any .trace() call to execute on NDIF. Your interventions are
packaged, sent to the server, and only the .save()d results come back:
with model.trace("The Eiffel Tower is in the city of", remote=True):
logit = model.lm_head.output[0][-1].argmax(dim=-1).save()
print(f"Prediction: {model.tokenizer.decode(logit)!r}")
Prediction: ' Paris'
While it runs (with the status display left on), a single in-place status line reports the job's lifecycle:
| Status | Meaning |
|---|---|
RECEIVED |
Request validated with your API key |
QUEUED |
Waiting in the model's queue |
DISPATCHED |
Forwarded to a model deployment; about to run |
RUNNING |
Your interventions are executing on the GPU pod |
LOG |
A print(...) from inside your block |
COMPLETED |
Results ready for download |
ERROR |
Server-side exception; raised locally as RemoteError |
We silenced the display in the setup cell (CONFIG.APP.REMOTE_LOGGING = False) so the
saved outputs stay clean; leave it on for interactive use. Turn on
CONFIG.APP.DEBUG = True for payload/result byte sizes and a per-status timeline.
Remote Generation¶
model.generate(..., remote=True) runs autoregressive generation and returns
token ids — read them off tracer.result:
with model.generate("The Eiffel Tower is in the city of",
max_new_tokens=3, do_sample=False, remote=True) as tracer:
out = tracer.result.save()
print(f"Generated: {model.tokenizer.decode(out[0])!r}")
Generated: 'The Eiffel Tower is in the city of Paris, and'
For per-step interventions inside generation, iterate the steps with
for _ in tracer.iter[:N]: and act on model.<module>.output each step (see the
Multiple Token Generation tutorial). model.generate returns ids; model.pipe(...)
runs the whole task pipeline and returns its decoded records.
Saving Results¶
.save() is the transmission mechanism. Only values touched by .save() (or
nnsight.save(...)) are shipped back; everything else lives only on the server for
the duration of the run and is discarded when the job finishes. Your local variable
name is matched to the saved value when the result is pushed back into your frame.
with model.trace("Hello", remote=True):
hidden = model.transformer.h[5].output # not saved -> not returned
answer = model.lm_head.output[0][-1].argmax(-1).save() # returned
# 'hidden' is undefined here; 'answer' is a real tensor.
In 0.8, .save() raises if called outside a trace (it used to be a silent
no-op) — it only makes sense inside a with model.trace(...): block.
Minimize what you save
Every .save() transfers data over the internet. Save argmax indices instead of full
logit tensors, specific positions instead of whole sequences. For large activations,
call .detach().cpu() inside the block before saving — the deserializer already maps
tensors to CPU on your end, but detaching drops the autograd graph and does the
conversion server-side, so you ship a smaller payload:
with model.trace("Hello", remote=True):
hidden = model.transformer.h[0].output.detach().cpu().save()
Sessions¶
A session bundles multiple traces into a single NDIF job: one queue wait, one
transport round-trip, and values captured in an earlier trace are available in a
later one without a .save() round-trip (only cross-process results need
.save()). This is the way to run a multi-step experiment — such as activation
patching — without paying the queue cost several times.
remote=True goes on model.session(...), not on the inner model.trace(...)
calls. Here is activation patching across three traces in one job:
- Trace 1 captures the final-layer hidden state at the last position. No
.save()is needed — the value is reused in a later trace within the same session. - Trace 2 records the clean baseline prediction for a different prompt.
- Trace 3 patches the captured hidden state into that prompt and re-reads the prediction.
with model.session(remote=True):
with model.trace("The Eiffel Tower is in the city of"):
hs = model.transformer.h[-1].output[:, -1, :]
with model.trace("The Colosseum is in the city of"):
clean = model.lm_head.output[0, -1].argmax(dim=-1).save()
with model.trace("The Colosseum is in the city of"):
model.transformer.h[-1].output[:, -1, :] = hs
patched = model.lm_head.output[0, -1].argmax(dim=-1).save()
print(f"Clean: {model.tokenizer.decode(clean)!r}")
print(f"Patched: {model.tokenizer.decode(patched)!r}")
Clean: ' T' Patched: ' Paris'
Patching the last layer's representation flips the Colosseum prediction toward the Eiffel Tower's answer — captured, patched, and compared in a single request.
Session gotchas
.save()is still required for any value returned to your process; cross-trace sharing inside the session is free, cross-process (server -> you) is not.- Put
remote=Trueonmodel.session(...), never on the inner traces. - Build everything inside the session block — variables defined outside it can't be referenced inside.
- Sessions cut queue and transport overhead, not GPU time.
- One trace fails -> the whole session aborts. Structure fault-tolerant pipelines as separate jobs.
Ways to Run a Job¶
Every example so far used the default blocking submission. There are three ways to submit, trading off how your process waits for the result.
Blocking (default)¶
The client holds one websocket open until COMPLETED and pushes the saved values
straight back into your variables — this is what every cell above did. Simplest, and
right for interactive work:
with model.trace("...", remote=True): # blocking=True is the default
answer = model.lm_head.output[0][-1].argmax(-1).save()
# 'answer' is populated here.
Non-blocking¶
Pass blocking=False to submit the job, get its id immediately, and poll for the
result later — useful for long jobs, or to manage many jobs at once. The job is
submitted when the with block exits; tracer.backend then holds the
RemoteBackend with the assigned job_id. Calling the backend polls once, returning
None until the job is COMPLETED and then the saves dict (keyed by your saved
variable names):
import time
with model.trace("The Eiffel Tower is in the city of",
remote=True, blocking=False) as tracer:
output = model.lm_head.output[0][-1].argmax(dim=-1).save()
backend = tracer.backend # the RemoteBackend, now holding the job id
print(f"job id: {backend.job_id}")
print(f"status: {backend.status.name}")
while True: # backend() returns None until COMPLETED
result = backend()
if result is not None:
break
time.sleep(1)
print(f"keys: {list(result.keys())}")
print(f"prediction: {model.tokenizer.decode(result['output'])!r}")
job id: cb6c6dd6673d4742b376ad451f9121a9 status: RECEIVED
keys: ['output'] prediction: ' Paris'
blocking=False works on model.session(...) too. You can also reattach to a job
later by constructing a RemoteBackend(model.to_model_key(), blocking=False, job_id="...") and calling it — because job_id is already set, the first call polls
instead of resubmitting.
Async¶
For asyncio code, AsyncRemoteBackend runs the same streaming request on the event
loop: build one, pass it as the backend= for your trace, then await it for the
saves dict — keeping the loop free while the job runs. This is what you want to fire
several jobs concurrently with asyncio.gather:
from nnsight.intervention.backends.remote import AsyncRemoteBackend
backend = AsyncRemoteBackend(model.to_model_key())
with model.trace("The Eiffel Tower is in the city of", backend=backend):
answer = model.lm_head.output[0][-1].argmax(dim=-1).save()
result = await backend # in a plain script: asyncio.run(...) a coroutine
print(f"keys: {list(result.keys())}")
print(f"prediction: {model.tokenizer.decode(result['answer'])!r}")
keys: ['answer'] prediction: ' Paris'
async for update in backend instead streams each raw status update as it arrives, if
you want to react to the lifecycle rather than only the final result.
Print Statements Appear as LOG¶
Anything you print(...) inside the block runs on the server and is forwarded as a
LOG status message, rendered inline by the status display — handy for debugging
without saving intermediate values. (With REMOTE_LOGGING off, as in this notebook,
the LOG line isn't shown, but the server still executes the print.) Print
summaries, not raw multi-megabyte tensors — a huge print floods the websocket.
with model.trace("The Eiffel Tower is in the city of", remote=True):
logits = model.lm_head.output
pred = logits[0, -1].argmax(dim=-1)
print(f"Predicted token id: {pred}") # appears as a LOG status server-side
decoded = model.tokenizer.decode(pred).save()
print(f"Result: {decoded!r}")
Result: ' Paris'
Remote Gradients¶
Backward passes work through the same interface: wrap the loss in
with loss.backward(): and read any tensor's .grad inside, .save()-ing the
gradient you want back.
One thing differs from a local run: NDIF loads models with gradients off — every
parameter is served with requires_grad=False, so a plain forward builds no autograd
graph and a backward has nothing to differentiate. To track gradients, turn grad on
for a tensor the graph should flow from — the input embeddings are the natural choice
— with requires_grad_(True). The backward then runs, and any downstream tensor's
.grad is populated:
with model.trace("The Eiffel Tower is in the city of", remote=True):
model.transformer.wte.output.requires_grad_(True) # grad is off by default; enable it here
hs = model.transformer.h[5].output
with model.lm_head.output.sum().backward():
grad = hs.grad.save()
print(f"Gradient shape: {tuple(grad.shape)}")
print(f"Gradient norm: {grad.norm():.4f}")
Gradient shape: (1, 10, 768) Gradient norm: 954368.0000
Running Your Own Code on the Server¶
Your intervention block is serialized by source and rebuilt on the server, so
ordinary Python and PyTorch work as written. nnsight serializes with
cloudpickle, which ships functions and
classes defined in your session (__main__) by value — their source travels with
the request automatically. So a helper you define right in the notebook just works
remotely, no registration needed:
def steer(hidden, strength=5.0):
"""A helper defined locally — cloudpickle ships it by value with the request."""
return hidden + strength * hidden.mean(dim=1, keepdim=True)
with model.trace("The Eiffel Tower is in the city of", remote=True):
model.transformer.h[5].output[0][:] = steer(model.transformer.h[5].output[0])
steered = model.lm_head.output[0][-1].argmax(dim=-1).save()
print(f"Steered prediction: {model.tokenizer.decode(steered)!r}")
Steered prediction: ' London'
The gap is code that lives in a separate importable module you wrote (not
installed on the server). By default cloudpickle would pickle it by reference — the
server would import your_module and fail with ModuleNotFoundError. Register it to
force by-value serialization so its source ships too:
import nnsight
import my_utils # a local file, not installed on the server
nnsight.register(my_utils) # ship my_utils by value (wraps
# cloudpickle.register_pickle_by_value)
with model.trace("Hello world", remote=True):
vec = my_utils.steer(model.transformer.h[5].output[0]).save()
nnsight.register accepts a module object or its name as a string
(nnsight.register("my_utils")). In practice you rarely call it by hand: before its
first request the remote backend auto-registers every module in your working tree
(via pull_env), so local project files are usually shipped already. Reach for
register for editable installs or a locally-upgraded package that also exists on the
server. Define helpers at module (or notebook) scope — dynamically created functions
(for example functions built with exec, or closures generated inside a loop) can
confuse the source-based serializer.
To verify offline, remote="local" deserializes the block with your local modules
hidden, mimicking a server that lacks your source, so an unshipped helper raises
ModuleNotFoundError exactly as the real server would.
Custom Host¶
NDIF has staging and self-hosted deployments (the setup cell at the top pointed at a
local one). Point at a host persistently
(CONFIG.API.HOST = "https://staging.api.ndif.us"; CONFIG.save()), via the
NDIF_HOST env var, or per-call by passing the URL as remote:
with model.trace("...", remote="https://self-hosted.example.com"):
out = model.lm_head.output.save()
Usage Limits¶
Each request has a one hour time limit; a job that exceeds it is terminated. NDIF is shared infrastructure and the fairness policy is still evolving — if you run into limits or have feedback, reach out on the NDIF Discord.