Fix remote execution for encoder models

[elliottower]

Hi all, greatly appreciate all the hard work that goes into NDIF. I am working on a mechinterp project on RNA and DNA foundation models and was looking into whether it was possible to run them here, and it looks like with these two bug fixes it will be. It should also help anyone else who wants to run encoder models in the future. Thanks.

---

1. Envoy classes for modules that define .output can't be pickled

_handle_overloaded_mount builds a class per Envoy instance named <cls>.Preserved. The dot makes the name unresolvable and the class is never bound in a module, so pickle can't serialize it:

PicklingError: Can't pickle <class 'nnsight.intervention.envoy.Envoy.Preserved'>:
attribute lookup Envoy.Preserved on nnsight.intervention.envoy failed

Remotely this shows up as RemoteException: name 'hooked_output' is not defined. It hits any model with an input or output submodule, so all of BERT and ESM (BertLayer.output is a BertOutput). Decoder-only models never take this path, which is why they work today and encoders don't.

• classes are memoized on (base class, mount point) and bound into the module that defines the base
• the attributes installed are fully determined by that pair, so one shared class is equivalent to one per instance
• per-instance state stays in __dict__
• side effect: a 12-layer BERT goes from 12 synthesized classes to 1

2. automodel isn't sent to the server

LanguageModel._remoteable_model_key serializes only repo_id and revision, so the server rebuilds every model as AutoModelForCausalLM:

Failed to provision model: Unrecognized configuration class EsmConfig
for this kind of AutoModel: AutoModelForCausalLM

• automodel is now in the key
• _remoteable_from_model_key already merges the key JSON into constructor kwargs, and TransformersMixin resolves a string automodel, so no server change is needed
• omitted when it's the default, so keys for existing causal-LM deployments are byte identical

Tests

tests/test_envoy_overloaded_mount.py, tiny models, no cluster needed. 3 fail on main, all pass here. The two remote tests skip without NDIF_KEY.

Reproduced and fixed on two independent stacks:

• macOS x86_64, Python 3.12, torch 2.2.2, transformers 4.57.6, nnsight 0.7.0 from PyPI
• Linux, Python 3.11, torch 2.6.0, transformers 5.13.1, nnsight from source at 87f4dad

Same failure and same fix on both, spanning transformers 4.x and 5.x and torch either side of the >=2.4.0 floor. On Linux I ran this branch against main in the same container:

• tests/test_tiny.py passes on both
• RNABERT and SpliceBERT envoy classes go from unpicklable to picklable
• per-position embedding shift under a single-nucleotide substitution is identical to 4 decimals on both branches, so no numerical behavior changes

[JadenFiotto-Kaufman]

Confirmed and fixed, by a different route than your patch, so I would like your read on it before this closes.

Your diagnosis holds on 0.8. The synthesized class survived the rewrite, renamed from .Preserved to __Overloaded, still bound to no module. With a decoder as a control:

gpt2   remote -> OK
bert   remote -> TypeError: cannot pickle 'eproperty' object

Your fix makes that class picklable: cache it per (base, names) and bind it into sys.modules[base.__module__] so cloudpickle finds it by reference. That works, and I had it applied and passing before I changed my mind about it.

What I went after instead is the class itself. It exists only to override output on one envoy so the submodule can keep the name. Give the name to nnsight's .output and move the submodule, and there is nothing to override, so the envoy stays a plain Envoy and pickles like any other.

On a BERT attention block:

attention.output      # the block's forward output, as on every other module
attention.E_output    # the BertSelfOutput submodule

This is a breaking change and the one thing I am least sure of. .nns_output is gone. The submodule keeps its own path, so .path and named_modules() still report ...attention.output, and print(model) labels it E_output/output. Attribute access is the only spelling that moved, so get() and rename are written against E_output too.

Landed as 518d03f on 0.8, with you as co-author. Sixteen tests cover it, including a real BERT through remote="local", which is the path your issue was actually blocked on.

Two things I would value your view on, since you hit this in real use rather than in a test:

1. Does E_ read right at a call site? It is short and it sorts next to the module name in the repr, but I do not love it.
2. Is there anywhere you count on that submodule answering to the plain name, where this would break you?

One thing that was our fault: this was opened against main, and 0.8 is where development happens. Our CLAUDE.md says to base PRs on dev, which has not been true for a while. We are fixing that.

Leaving this open until you have had a look.