Skip to content

Fix remote execution for encoder models - #700

Open
elliottower wants to merge 2 commits into
ndif-team:mainfrom
elliottower:fix/remote-encoder-models
Open

Fix remote execution for encoder models#700
elliottower wants to merge 2 commits into
ndif-team:mainfrom
elliottower:fix/remote-encoder-models

Conversation

@elliottower

Copy link
Copy Markdown

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant