Add a rename_modules_dict arg to Envoy that allows renaming modules

[Butanium]

After discussing with @cadentj some alternatives to transformer_lens (which sometimes has a significant implementation gap with HuggingFace), he suggested adding a renaming feature to nnsight itself, as it might be useful not only for transformers.

The basic usage would be to convert module names from, e.g., GPT-2 to the standard LLaMA / Gemma format using:

The basic usage would be to convert module names from e.g. gpt2 to the standard llama / gemma format using

dict_change = {
    "transformer": "model",
    "self_attention": "self_attn",
    "attn": "self_attn",
    "h": "layers",
}
patched_model = LanguageModel("gpt2", rename_modules_dict=dict_change)
with patched_model.trace("hello"):
    out = patched_model.model.layers[0].output[0].save()

[Butanium]

@JadenFiotto-Kaufman I can try to add tests if needed but I wanted to first check if you were ok with adding this feature to the nnsight core.

[JadenFiotto-Kaufman]

@JadenFiotto-Kaufman I can try to add tests if needed but I wanted to first check if you were ok with adding this feature to the nnsight core.

Hey this is an awesome idea and something I've wanted to add to nnsight. Essentially we want aliasing right?

I think it should be a dictionary where the keys are valid regex strings that map to an alias. This dictionary would live at Envoy._aliases and would be checked on the Envoy.getattr method. If a key matches the reverse mapping, it returns the real underlying envoy for that alias.

Does that sound like a good solution?

[Butanium]

Hmm, would you still pass the alias dict through the NNSight class init?

I guess one advantage of renaming over aliasing is that when you print the model, you actually see modules' new names.

What advantage do you see of using aliasing instead of renaming?

[Butanium]

@JadenFiotto-Kaufman bump