From 19d1be71c22e4b2d348717b1f76a37de38e26ab9 Mon Sep 17 00:00:00 2001 From: Elliot Tower Date: Tue, 25 Aug 2026 23:43:23 -0400 Subject: [PATCH 1/2] Fix unpicklable Envoy classes for modules that define .output --- src/nnsight/intervention/envoy.py | 73 ++++++++++++++--- tests/test_envoy_overloaded_mount.py | 115 +++++++++++++++++++++++++++ 2 files changed, 176 insertions(+), 12 deletions(-) create mode 100644 tests/test_envoy_overloaded_mount.py diff --git a/src/nnsight/intervention/envoy.py b/src/nnsight/intervention/envoy.py index 2f1a253d..661be0c0 100755 --- a/src/nnsight/intervention/envoy.py +++ b/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 : + 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) diff --git a/tests/test_envoy_overloaded_mount.py b/tests/test_envoy_overloaded_mount.py new file mode 100644 index 00000000..b7a32995 --- /dev/null +++ b/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 From 79b52a02b5421c84f01c436f474e1d77ec791221 Mon Sep 17 00:00:00 2001 From: Elliot Tower Date: Tue, 25 Aug 2026 23:43:22 -0400 Subject: [PATCH 2/2] Send automodel in the remote model key --- src/nnsight/modeling/language.py | 29 +++++++++++++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/src/nnsight/modeling/language.py b/src/nnsight/modeling/language.py index fc77f4e6..ad582613 100755 --- a/src/nnsight/modeling/language.py +++ b/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()