Fix remote execution for encoder models

src/nnsight/intervention/envoy.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import importlib
 import inspect
 import os
 import warnings
@@ -75,6 +76,11 @@ class Envoy(Batchable):
         _alias (Aliaser): Aliaser object for managing aliases
     """
 
+    #: Subclasses synthesized by :meth:`_preserved_subclass`, keyed on
+    #: ``(base class, mount point)``. Shared across Envoy instances so the
+    #: classes stay resolvable by qualified name and therefore picklable.
+    _PRESERVED_CLASSES: Dict[Tuple[Type["Envoy"], str], Type["Envoy"]] = {}
+
     def __init__(
         self,
         module: torch.nn.Module,
@@ -752,6 +758,54 @@ def _add_envoy(self, module: torch.nn.Module, name: str) -> Envoy:
 
         return envoy
 
+    @staticmethod
+    def _preserved_subclass(base: Type["Envoy"], mount_point: str) -> Tuple[Type["Envoy"], bool]:
+        """Return a subclass of ``base`` whose ``mount_point`` is shadowed.
+
+        The class is memoized on ``(base, mount_point)`` and bound into the
+        module that defines ``base``, so that ``pickle`` and ``importlib`` can
+        resolve it by qualified name.
+
+        Previously one anonymous class was synthesized per Envoy instance, named
+        ``f"{base.__name__}.Preserved"``. That name contains a dot and was never
+        bound in any module, so the class could not be serialized:
+
+            PicklingError: Can't pickle <class 'nnsight.intervention.envoy.Envoy.Preserved'>:
+            attribute lookup Envoy.Preserved on nnsight.intervention.envoy failed
+
+        Remote execution therefore failed for every model with a submodule named
+        ``input`` or ``output`` -- which is every BERT- and ESM-family model,
+        since ``BertLayer.output`` is a ``BertOutput`` submodule -- while
+        decoder-only models, which never take this path, worked.
+
+        Since the attributes installed on the class are fully determined by
+        ``(base, mount_point)``, one shared class per pair is equivalent to one
+        per instance. Per-instance state stays in ``self.__dict__``.
+
+        Returns:
+            (cls, created): ``created`` is False when the class came from cache
+            and its attributes are already installed.
+        """
+
+        key = (base, mount_point)
+
+        cached = Envoy._PRESERVED_CLASSES.get(key)
+        if cached is not None:
+            return cached, False
+
+        name = f"{base.__name__}__nns_{mount_point}"
+
+        cls = type(
+            name,
+            (base,),
+            {"__module__": base.__module__, "__qualname__": name},
+        )
+
+        Envoy._PRESERVED_CLASSES[key] = cls
+        setattr(importlib.import_module(base.__module__), name, cls)
+
+        return cls, True
+
     def _handle_overloaded_mount(self, envoy: Envoy, mount_point: str) -> None:
         """If a given module already has an attribute of the same name as something nnsight wants to add, we need to rename it.
 
@@ -766,20 +820,15 @@ def _handle_overloaded_mount(self, envoy: Envoy, mount_point: str) -> None:
             f"Module `{self.path}` of type `{type(self._module)}` has pre-defined a `{mount_point}` attribute. nnsight access for `{mount_point}` will be mounted at `.nns_{mount_point}` instead of `.{mount_point}` for this module only."
         )
 
-        # If we already shifted a mount point dont create another new class.
-        if "Preserved" in self.__class__.__name__:
+        new_cls, created = Envoy._preserved_subclass(self.__class__, mount_point)
 
-            new_cls = self.__class__
+        object.__setattr__(self, "__class__", new_cls)
 
-        else:
-
-            new_cls = type(
-                f"{self.__class__.__name__}.Preserved",
-                (self.__class__,),
-                {},
-            )
-
-            object.__setattr__(self, "__class__", new_cls)
+        if not created:
+            # The shared class already carries the shadowing attributes; only
+            # the per-instance child envoy below still has to be installed.
+            self.__dict__[mount_point] = envoy
+            return
 
         # Get the normal proxy mount point
         mount = getattr(Envoy, mount_point)

src/nnsight/modeling/language.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-
+import json
 import warnings
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
@@ -543,7 +543,32 @@ def _batch(
         return tuple(), kwargs
 
     def _remoteable_model_key(self) -> str:
-        return super()._remoteable_model_key()
+        """Include ``automodel`` so non-causal heads survive remote execution.
+
+        The base implementation serializes only ``repo_id`` and ``revision``, so
+        the server reconstructs every model through the ``LanguageModel``
+        default of ``AutoModelForCausalLM``. Any encoder or masked-LM checkpoint
+        then fails to provision:
+
+            Failed to provision model: Unrecognized configuration class
+            EsmConfig for this kind of AutoModel: AutoModelForCausalLM
+
+        which covers every protein and RNA language model, along with BERT-style
+        encoders generally. ``_remoteable_from_model_key`` already merges the key
+        JSON into the constructor kwargs, and ``TransformersMixin`` accepts a
+        string ``automodel`` and resolves it via ``getattr(modeling_auto, ...)``,
+        so carrying the class name is enough.
+
+        Omitted when it is the default, to keep keys for existing causal-LM
+        deployments byte-identical.
+        """
+
+        key = json.loads(super()._remoteable_model_key())
+
+        if self.automodel is not AutoModelForCausalLM:
+            key["automodel"] = self.automodel.__name__
+
+        return json.dumps(key)
 
     def _remoteable_persistent_objects(self) -> dict:
         persistent_objects = super()._remoteable_persistent_objects()

tests/test_envoy_overloaded_mount.py

@@ -0,0 +1,115 @@
+"""Encoder models cannot be executed remotely: the Envoy class they get is unpicklable.
+
+`Envoy._handle_overloaded_mount` fires whenever a wrapped module already defines
+an attribute nnsight wants to mount -- in practice `.output`, which every
+HuggingFace BERT- and ESM-family layer defines as a submodule
+(`BertLayer.output = BertOutput(...)`). It resolves the collision by synthesizing
+a subclass at runtime:
+
+    new_cls = type(f"{self.__class__.__name__}.Preserved", (self.__class__,), {})
+    object.__setattr__(self, "__class__", new_cls)
+
+The synthesized class has a dot in `__name__`, is never bound in its defining
+module, and so cannot be looked up by qualified name. Pickling it fails, and
+remote execution of any encoder model fails on NDIF with
+`RemoteException: name 'hooked_output' is not defined`.
+
+Decoder models never enter this path, which is why GPT-2, Pythia and Llama work
+remotely and BERT and ESM do not.
+
+The local tests need no NDIF access and run on tiny models in seconds. The remote
+tests are skipped unless NDIF_KEY is set.
+
+    pytest test_envoy_overloaded_mount.py -v
+"""
+
+import os
+import pickle
+
+import pytest
+
+from nnsight import Envoy, LanguageModel
+
+# Tiny models so the local tests are CI-cheap. Both are public.
+TINY_DECODER = "sshleifer/tiny-gpt2"
+TINY_ENCODER = "hf-internal-testing/tiny-random-BertModel"
+
+needs_ndif = pytest.mark.skipif(
+    not os.environ.get("NDIF_KEY"), reason="NDIF_KEY not set"
+)
+
+
+@pytest.fixture(scope="module")
+def decoder():
+    return LanguageModel(TINY_DECODER)
+
+
+@pytest.fixture(scope="module")
+def encoder():
+    from transformers import AutoModelForMaskedLM
+
+    return LanguageModel(TINY_ENCODER, automodel=AutoModelForMaskedLM)
+
+
+# ── local: no cluster required ───────────────────────────────────────────────
+
+
+def test_decoder_layer_envoy_keeps_the_base_class(decoder):
+    assert type(decoder.transformer.h[0]) is Envoy
+
+
+def test_decoder_layer_envoy_class_is_picklable(decoder):
+    pickle.dumps(type(decoder.transformer.h[0]))
+
+
+def test_encoder_layer_envoy_class_name_has_no_dot(encoder):
+    """A dot in __name__ makes the class unresolvable by qualified name."""
+    cls = type(encoder.bert.encoder.layer[0])
+    assert "." not in cls.__name__, (
+        f"synthesized class is named {cls.__name__!r}; a dot in __name__ means "
+        "pickle and importlib cannot resolve it"
+    )
+
+
+def test_encoder_layer_envoy_class_is_resolvable_in_its_module(encoder):
+    import importlib
+
+    cls = type(encoder.bert.encoder.layer[0])
+    module = importlib.import_module(cls.__module__)
+    assert getattr(module, cls.__name__, None) is cls, (
+        f"{cls.__module__}.{cls.__name__} does not resolve back to the class"
+    )
+
+
+def test_encoder_layer_envoy_class_is_picklable(encoder):
+    """The failing test. Passes for decoders, fails for every encoder."""
+    pickle.dumps(type(encoder.bert.encoder.layer[0]))
+
+
+def test_remapped_accessor_is_present(encoder):
+    """Guards the intended behavior of the collision handling itself.
+
+    Whatever fix lands must keep `.nns_output` mounted, so this should pass
+    before and after.
+    """
+    assert hasattr(type(encoder.bert.encoder.layer[0]), "nns_output")
+
+
+# ── remote: needs NDIF_KEY ───────────────────────────────────────────────────
+
+
+@needs_ndif
+def test_remote_decoder_returns_activations():
+    lm = LanguageModel("EleutherAI/pythia-160m")
+    with lm.trace("The capital of France is", remote=True):
+        h = lm.gpt_neox.layers[6].output[0].save()
+    assert h.shape[-1] == 768
+
+
+@needs_ndif
+def test_remote_encoder_returns_activations():
+    """Fails with: RemoteException: name 'hooked_output' is not defined."""
+    bm = LanguageModel("google-bert/bert-base-uncased")
+    with bm.trace("The capital of France is [MASK].", remote=True):
+        h = bm.bert.encoder.layer[6].nns_output[0].save()
+    assert h.shape[-1] == 768