[Possible Bug] Running multiple initialization seeds

[Abzinger]

Hi,

I wanted to train a network with multiple initial seeds such as this

cfg = TrainConfig(sae_config, optimizer=optimizer, init_seeds=[0,42], 
                batch_size=b_size, layers=['6'], 
                run_name=run_name, save_dir=save_dir)
trainer = Trainer(cfg, tokenized, gpt)

However I get the error KeyError: 'h.6' from line 375 raw = self.saes[name], see log below

File "script_training_sae.py", line 97, in <module>
    trainer.fit()
  File "sparsify/sparsify/trainer.py", line 481, in fit
    self.model(x)
  File "lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.11/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 1062, in forward
    transformer_outputs = self.transformer(
                          ^^^^^^^^^^^^^^^^^
  File "lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.11/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 922, in forward
    outputs = block(
              ^^^^^^
  File "lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "lib/python3.11/site-packages/torch/nn/modules/module.py", line 1845, in _call_impl
    return inner()
           ^^^^^^^
  File "/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1806, in inner
    hook_result = hook(self, args, result)
                  ^^^^^^^^^^^^^^^^^^^^^^^^
  File "sparsify/sparsify/trainer.py", line 375, in hook
    raw = self.saes[name]
          ~~~~~~~~~^^^^^^
KeyError: 'h.6'

Same error when training using hookpoints. (I was able to successfully train on layer 6 with a single initial_seed)

Possible Bug

I think the problem might be the following:
On line 77
name = f"{hook}/seed{seed}" if len(cfg.init_seeds) > 1 else hook
self.saes[name] = SparseCoder(
input_widths[hook], cfg.sae, device, dtype=torch.float32
)
the dictionary sae is defined with the keys name = f"{hook}/seed{seed}"
and apparently in fit() (line 375) raw=self.saes[name] takes only f"{hook}". The problem could be that the variable name is being reset from f"{hook}/seed{seed}" to f"{hook} in lines 328:

name_to_module = {
            name: self.model.base_model.get_submodule(name)
            for name in self.cfg.hookpoints
        }

I hope the error is reproducible.

[luciaquirke]

Yep I've seen this error before, I think your diagnosis is right too

[JordanMcCann]

Traced this. There are two layers to it.

The self.saes keys are f"{hook}/seed{seed}" when len(init_seeds) > 1 (trainer.py ~L78), but the hook looks them up by bare hookpoint:

name, _, _ = module_to_name[module].partition("/")   # ~L351
...
raw = self.saes[name]                                 # ~L395  -> KeyError 'h.6'

The partition("/") is a no-op here. module_to_name comes from name_to_module, which is keyed by self.cfg.hookpoints - no seed suffix (~L318) - so module_to_name[module] is already just "h.6". There's never a slash to strip, and the bare name can't match the seeded self.saes keys.
The deeper issue: the hook assumes one SAE per hookpoint, but multi-seed makes N. So it's not only the lookup - self.saes[name], maybe_wrapped[name], and num_tokens_since_fired[name] all need to resolve to the right seed(s). The minimal fix that matches "train independent SAEs per seed in parallel" loops the per-hookpoint block over the seed variants:

base, _, _ = module_to_name[module].partition("/")
names = [base] if base in self.saes else [k for k in self.saes if k.split("/")[0] == base]
for name in names:
    raw = self.saes[name]
    wrapped = maybe_wrapped[name]
    out = wrapped(x=inputs, y=outputs, dead_mask=...)
    # ...existing per-SAE loss / backward bookkeeping...

One question before I PR it: is multi-seed meant to be N independent SAEs trained simultaneously on the same hookpoint? If so I'll send the loop version plus the matching change in the step/aggregation code. If it's meant to pick a single seed, the fix is smaller.

[luciaquirke]

(ignoring this due to overuse of AI)

[AnshKetchum]

Still open? Would love to try tackling this.

[luciaquirke]

Yes, still open!

[AnshKetchum]

Sounds good, working on replicating this issue, should have something by end of week. If possible, would love to have the issue assigned to me.