feat(activation):enable aligned permutation of crossmodel gen

This is achieved via setting same seeds for randperm generators and making sure they all lie in cuda device

Author: crabshellman

src/lm_saes/activation/factory.py

@@ -207,7 +207,7 @@ def build_batchler():
             """Create batchler for batched-activations-1d target."""
             assert cfg.batch_size is not None, "Batch size must be provided for outputting batched-activations-1d"
             return ActivationBatchler(
-                hook_points=cfg.hook_points, batch_size=cfg.batch_size, buffer_size=cfg.buffer_size, buffer_shuffle_config=cfg.buffer_shuffle_config
+                hook_points=cfg.hook_points, batch_size=cfg.batch_size, buffer_size=cfg.buffer_size, buffer_shuffle_config=cfg.buffer_shuffle
             )
 
         processors = [build_batchler()] if cfg.target >= ActivationFactoryTarget.BATCHED_ACTIVATIONS_1D else []

src/lm_saes/config.py

@@ -261,7 +261,7 @@ class ActivationFactoryConfig(BaseConfig):
         else None
     )
     """ Buffer size for online shuffling. If `None`, no shuffling will be performed. """
-    buffer_shuffle_config: Optional[BufferShuffleConfig] = None
+    buffer_shuffle: Optional[BufferShuffleConfig] = None
     """" Manual seed and device of generator for generating randomperm in buffer. """
     ignore_token_ids: Optional[list[int]] = None
     """ Tokens to ignore in the activations. """

src/lm_saes/runner.py

@@ -125,7 +125,7 @@ class GenerateActivationsSettings(BaseSettings):
     buffer_size: Optional[int] = None
     """Size of the buffer for activation generation"""
 
-    buffer_shuffle_config: Optional[BufferShuffleConfig] = None
+    buffer_shuffle: Optional[BufferShuffleConfig] = None
     """"Manual seed and device of generator for generating randomperm in buffer"""
 
     total_tokens: Optional[int] = None
@@ -202,7 +202,7 @@ def generate_activations(settings: GenerateActivationsSettings) -> None:
         model_batch_size=settings.model_batch_size,
         batch_size=settings.batch_size,
         buffer_size=settings.buffer_size,
-        buffer_shuffle_config=settings.buffer_shuffle_config,
+        buffer_shuffle=settings.buffer_shuffle,
     )
 
     # Configure activation writer