fix: handle LayerNorm folding correctly in load_and_process_state_dict

[VedantMadane]

Fixes #219

Problem

load_and_process_state_dict(state_dict, fold_ln=True) has two failure modes when called directly (outside the from_pretrained path):

1. Unfolded state dict: fold_layer_norm correctly removes LN keys from the state dict, but the model's LayerNorm modules are never replaced with LayerNormPre. This leaves mismatched architecture (modules expect w/b params that no longer exist) and broken hooks after loading.

2. Already-folded state dict: fold_layer_norm crashes with KeyError because it tries to access blocks.{l}.ln1.w keys that were already removed when the model was saved.

The workaround (from the issue) was a 3-step dance:

model.load_and_process_state_dict(state_dict, fold_ln=False)
model.process_weights_(fold_ln=True)
model.setup()

Fix

1. Detect already-folded state dicts: Check if LN weight keys (.ln1.w, .ln2.w, ln_final.w) exist before folding. If missing, skip with a warning instead of crashing.

2. Replace LN modules when folding: Move the LayerNorm -> LayerNormPre (and RMSNorm -> RMSNormPre) module replacement from process_weights_ into load_and_process_state_dict, so direct callers get the same correct behavior.

3. Re-attach hooks: Call self.setup() after loading when fold_ln=True to ensure hooks are properly connected to the new modules.

4. Simplify process_weights_: Since the module replacement now lives in load_and_process_state_dict, process_weights_ can simply delegate without duplicating the logic.

What this enables

# Both of these now work correctly:
model = HookedTransformer(cfg)
model.load_and_process_state_dict(unfolded_state_dict, fold_ln=True)  # was broken

model = HookedTransformer(cfg)
model.load_and_process_state_dict(already_folded_dict, fold_ln=True)  # was KeyError

[jlarson4]

In the future, please run make check-format & resolve any formatting errors before submitting code. There is a ton of noise in the File Changes diff that make this contribution difficult to review.

For a change like this that is meant to address specific failures, it is important to add new Unit Tests covering these previous failed states (to prevent regression) and any new features/functionalities you add (to make sure they don't break in the future).

[VedantMadane]

Apologies for the formatting noise — the diff is now cleaned up to contain only the logic changes. No auto-formatter reformatting of existing lines.

[jlarson4]

Apologies for the formatting noise — the diff is now cleaned up to contain only the logic changes. No auto-formatter reformatting of existing lines.

Thank you! I will take a look at reviewing this soon