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