kent/consciousness
1
0
Fork 1
Code Issues 1 Pull requests 5 Projects Releases Packages Wiki Activity Actions

amygdala lib: move_to_cpu=True to avoid bf16 SVD on CUDA

Browse source 
torch.svd doesn't support bf16 on CUDA; moving activations to CPU
first makes pca_aggregator work.

Co-Authored-By: Proof of Concept <poc@bcachefs.org>
This commit is contained in:
Kent Overstreet 2026-04-18 22:19:23 -04:00
parent 2ea89b1cb0
commit 7f6d94417e
1 changed files with 1 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

1
training/amygdala_training/train_with_library.py
View file

@ -185,6 +185,7 @@ def main() -> None:
aggregator=aggregator, aggregator=aggregator,
batch_size=args.batch_size, batch_size=args.batch_size,
show_progress=False, show_progress=False,
move_to_cpu=True,
) )
# sv.layer_activations is a dict {layer_idx: tensor[hidden]} # sv.layer_activations is a dict {layer_idx: tensor[hidden]}
for l_idx, layer in enumerate(target_layers): for l_idx, layer in enumerate(target_layers):