Description

ActivationsStore.get_activations assumes CUDA in two places. The multi-hook sibling in the same class, get_multi_hook_activations, already avoids both.

The first is the autocast device, which is hardcoded:

with torch.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=self.autocast_lm):

On a non-CUDA device PyTorch does not raise here. It warns and turns autocast off:

UserWarning: CUDA is not available or torch_xla is imported. Disabling autocast.

So on MPS, CPU and XPU, autocast_lm=True silently runs the model in full precision. On tiny-stories-1M the activations come back bit-identical to autocast_lm=False. I've changed this to _get_input_token_device(self.model).type, which is the same expression get_multi_hook_activations already uses, and the activations then differ by 1.3e-3, which is bfloat16 rounding.

The second is the staging buffer, which is allocated with no device:

stacked_activations = torch.zeros((n_batches, n_context, self.d_in))

It lands on the default device no matter where the model ran, so the slice-assignments below it copy the activations off the accelerator, and get_raw_llm_batch moves them straight back with .to(self.device). On an A100 with act_store_device="cuda", get_activations returns its result on cpu today and on cuda:0 with this change. Allocating on self.device removes that round trip and makes the .to(self.device) in get_raw_llm_batch a no-op. It is also what the tests already expect: test_activations_store.py:174 and :206 assert activations.device == store.device, which passes on main only because they run on CPU. Timing just the buffer step on MPS, a 537 MB batch goes from 68.8 ms to 21.4 ms.

The values are unchanged there. With the store on CUDA and again on CPU, autocast off and on, the activations are bit-identical before and after this change, and autocast_lm still engages on CUDA either way. I also ran tests/training/test_activations_store.py on that machine unpatched and then patched and compared the results test by test. The same seven were already failing in both runs and nothing new breaks.

This is a real change for anyone already setting autocast_lm=True on MPS or CPU. Those runs now get the bfloat16 forward pass the flag promises instead of a silent float32 one, so their activations shift by bfloat16 rounding and their memory use drops. That's the intent, but it isn't a no-op for existing runs. The flag defaults to False, so nothing changes for anyone who hasn't opted in.

Alternatives I ruled out

• Allocating on layerwise_activations.device rather than self.device looks equivalent but is worse. With act_store_device="cpu" and the model on the GPU it builds the buffer in VRAM and then ships all of it to the host anyway. At 537 MB that measured 141 ms against 55 ms on main, and it holds an extra copy in VRAM while it does so, which hurts exactly the people who picked that setting to save memory.

• Dropping the staging buffer altogether, so this path matches the multi-hook one, would change behavior. The buffer pins the output to float32, and its fixed shape turns a hook whose width disagrees with d_in into a RuntimeError instead of a silently wrong result.

Type of change

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

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• 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

You have tested formatting, typing and tests

• I have run make check-ci to check format and linting. (you can run make format to format code if needed.)