Inconsistent behavior with module renaming in NNsight

[Butanium]

When using the new rename feature in NNsight, I've discovered some inconsistent behaviors while working on unit tests:

from nnsight import LanguageModel
import torch as th

gpt2 = LanguageModel("gpt2", rename=dict(ln_f="norm"))

# Case 1: Accessing module outputs
with gpt2.trace("hello"):
    # This fails with AttributeError: 'LayerNorm' object has no attribute 'output'
    norm_out = gpt2.transformer.ln_f.output.save()  
    # Consider: Provide a more helpful error when accessing a renamed module
    
    # This works as expected
    norm_out = gpt2.transformer.norm.output.save()
    h = gpt2.transformer.h[0].output[0].save()

# Case 2: Using module forward
with gpt2.trace("hello"):
    # This fails with NNsightError: 'GPT2LMHeadModel' object has no attribute 'norm' while it's supposed to be renamed
    ln_out = gpt2.transformer.norm(h)
    
    # This works, but is inconsistent with the renaming
    ln_out = gpt2.transformer.ln_f(h)

1. Original module names remain accessible after renaming

   • Attempting to access model.old_name.output fails with a generic AttributeError rather than something more specific like "Tried to access renamed module"
   • Note: Original names can still be used in traces (not necessarily problematic)

2. Forward method doesn't work with renamed modules

   • Users cannot use the renamed module's forward method (model.new_name(x))
   • This is counterintuitive as the renamed module should behave the same as the original

I feel like better handling the renaming in get_attribute is the way to go

[JadenFiotto-Kaufman]

@Butanium This PR https://github.com/ndif-team/nnsight/pull/325 fixes the calling problem. Good catch.

[Butanium]

@JadenFiotto-Kaufman thanks first bug resolved! However this other code also fails. This means that accessing a children of a renamed module to apply it during a trace doesn't work

gpt2 = LanguageModel("gpt2", rename=dict(transformer="model"))

with gpt2.trace("hello"):
    ln_out = gpt2.model.ln_f.output.save()
    h = gpt2.model.h[0].output[0].save()
print("successfully accessed renamed module")
with gpt2.trace("hello"):
    ln_applied_deprecated = gpt2.transformer.ln_f(h)
print("successfully applied deprecated ln_f")
with gpt2.trace("hello"):
    # This fails with NNsightError: 'GPT2LMHeadModel' object has no attribute 'model'
    ln_applied = gpt2.model.ln_f(h)

[JadenFiotto-Kaufman]

@Butanium Fixed in this branch! https://github.com/ndif-team/nnsight/tree/issues/321 Would you mind pulling it and testing it out? Also if you have the time, would you mind adding some pytests around renaming under /tests ? That would be very helpful.

[Butanium]

fyi, my nnterp renaming tests are passing, see: https://github.com/Butanium/nnterp/blob/0.4renamingdev/tests/test_model_renaming.py. I have on my todo to add some tests to nnsight directly but this might come only after the neurips deadline.

[AdamBelfki3]

Should this issue be closed @JadenFiotto-Kaufman ?

[Butanium]

Maybe let's wait until the tests are implemented?