Description

SAETrainer.load_trainer_state calls torch.load without map_location, so the restored tensors come back on whatever device the checkpoint was written from instead of the device the trainer is configured for.

act_freq_scores and n_forward_passes_since_fired are created on cfg.device in __init__ and reset on cfg.device in _reset_running_sparsity_stats, so only the restore path disagrees with the rest of the class. Every other torch.load in the codebase already passes map_location, which is the five call sites in pretrained_sae_loaders.py, and this is the only one that does not.

It is reachable through the ordinary resume path, since LanguageModelSAETrainingRunner calls load_trainer_state for resume_from_checkpoint and MultiSAETrainer calls it per SAE. Resuming a run on a different device from the one it was trained on then fails on the first sparsity accumulation in _train_step:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, mps:0 and cpu!

Resuming a CUDA checkpoint on a machine without CUDA is worse, since it raises inside torch.load before any state is restored.

The fix passes the trainer's own configured device:

state_dict = torch.load(
    checkpoint_path / TRAINER_STATE_FILENAME, map_location=self.cfg.device
)

Nothing changes when a checkpoint is saved and resumed on the same device, which is why the existing round-trip test did not catch it, since it builds both trainers on the default device.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

Two tests, both failing on current main:

• test_load_trainer_state_reads_the_checkpoint_onto_the_configured_device asserts torch.load receives map_location, and runs everywhere including CPU-only CI. On main it fails as assert None == 'cpu'.
• test_load_trainer_state_moves_an_accelerator_checkpoint_onto_cpu writes the checkpoint from an accelerator and resumes on CPU. It is skipped without a second device, following the GPU-gated tests in test_activations_store.py. On main it fails as assert 'mps' == 'cpu'.

Reverting the one-line change while keeping both tests fails both. The full tests/training/test_sae_trainer.py passes at 24 passed and 1 skipped, where the skip is the pre-existing fp16 case.

You have tested formatting, typing and tests

• I have run make check-ci to check format and linting.

ruff check and ruff format --check are clean. pyright reports three more instances of Type of parameter "ts_model" is unknown, the same pre-existing report the other ts_model tests in that file already produce, since the fixture is untyped.

Performance Check.

Not applicable. This only changes which device the running-sparsity tensors are restored onto when resuming, so L0, CE loss, MSE loss and feature dashboards are unaffected.