feat(activation):enable aligned permutation of crossmodel gen

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

Author: crabshellman

src/lm_saes/__init__.py

@@ -6,6 +6,7 @@
     ActivationFactoryDatasetSource,
     ActivationFactoryTarget,
     ActivationWriterConfig,
+    BufferShuffleConfig,
     CrossCoderConfig,
     DatasetConfig,
     FeatureAnalyzerConfig,
@@ -38,6 +39,7 @@
     "ActivationFactoryDatasetSource",
     "ActivationFactoryConfig",
     "ActivationWriterConfig",
+    "BufferShuffleConfig",
     "ActivationFactoryTarget",
     "load_dataset",
     "load_model",

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
+                hook_points=cfg.hook_points, batch_size=cfg.batch_size, buffer_size=cfg.buffer_size, buffer_shuffle_config=cfg.buffer_shuffle_config
             )
 
         processors = [build_batchler()] if cfg.target >= ActivationFactoryTarget.BATCHED_ACTIVATIONS_1D else []

src/lm_saes/activation/processors/activation.py

@@ -7,6 +7,7 @@
 from transformer_lens import HookedTransformer
 
 from lm_saes.activation.processors.core import BaseActivationProcessor
+from lm_saes.config import BufferShuffleConfig
 
 
 @dataclass
@@ -23,6 +24,7 @@ class ActivationBuffer:
     """
 
     buffer: list[dict[str, Any]] = field(default_factory=list)
+    generator: torch.Generator = torch.Generator()  # Generator passed from ActivationBatchler
 
     def __len__(self) -> int:
         """Get the number of samples in the buffer.
@@ -41,7 +43,7 @@ def cat(self, activations: dict[str, Any]) -> "ActivationBuffer":
         Returns:
             ActivationBuffer: New buffer containing concatenated activations
         """
-        return ActivationBuffer(buffer=self.buffer + [activations])
+        return ActivationBuffer(buffer=self.buffer + [activations], generator=self.generator)
 
     def consume(self) -> dict[str, torch.Tensor | list[Any]]:
         """Consume the buffer and return the activations as a dictionary."""
@@ -68,7 +70,7 @@ def yield_batch(self, batch_size: int) -> tuple[dict[str, torch.Tensor | list[An
         data = self.consume()
         batch = {k: v[:batch_size] for k, v in data.items()}
         buffer = {k: v[batch_size:] for k, v in data.items()}
-        return batch, ActivationBuffer(buffer=[buffer])
+        return batch, ActivationBuffer(buffer=[buffer], generator=self.generator)
 
     def shuffle(self) -> "ActivationBuffer":
         """Randomly shuffle all samples in the buffer.
@@ -81,9 +83,11 @@ def shuffle(self) -> "ActivationBuffer":
             isinstance(data[k], torch.Tensor) for k in data.keys()
         ), "All data must be tensors to perform shuffling"
         data = cast(dict[str, torch.Tensor], data)
-        perm = torch.randperm(data[list(data.keys())[0]].shape[0])
+        
+        # Use the passed generator for shuffling
+        perm = torch.randperm(data[list(data.keys())[0]].shape[0], generator=self.generator, device=self.generator.device)
         buffer = {k: v[perm] for k, v in data.items()}
-        return ActivationBuffer(buffer=[buffer])
+        return ActivationBuffer(buffer=[buffer], generator=self.generator)
 
 
 class ActivationGenerator(BaseActivationProcessor[Iterable[dict[str, Any]], Iterable[dict[str, Any]]]):
@@ -254,7 +258,7 @@ class ActivationBatchler(BaseActivationProcessor[Iterable[dict[str, Any]], Itera
             data will be refilled into the buffer whenever the buffer is less than half full, and then re-shuffled.
     """
 
-    def __init__(self, hook_points: list[str], batch_size: int, buffer_size: Optional[int] = None):
+    def __init__(self, hook_points: list[str], batch_size: int, buffer_size: Optional[int] = None, buffer_shuffle_config: Optional[BufferShuffleConfig] = None):
         """Initialize the ActivationBatchler.
 
         Args:
@@ -265,6 +269,10 @@ def __init__(self, hook_points: list[str], batch_size: int, buffer_size: Optiona
         self.hook_points = hook_points
         self.batch_size = batch_size
         self.buffer_size = buffer_size
+        self.perm_generator = torch.Generator()
+        if buffer_shuffle_config is not None:
+            self.perm_generator = torch.Generator(buffer_shuffle_config.generator_device)
+            self.perm_generator.manual_seed(buffer_shuffle_config.perm_seed)  # Set seed if provided
 
     def process(self, data: Iterable[dict[str, Any]], **kwargs) -> Iterable[dict[str, Any]]:
         """Process input data by batching activations.
@@ -283,7 +291,7 @@ def process(self, data: Iterable[dict[str, Any]], **kwargs) -> Iterable[dict[str
         Raises:
             AssertionError: If hook points are missing or tensors have invalid shapes
         """
-        buffer = ActivationBuffer()
+        buffer = ActivationBuffer(generator=self.perm_generator)
         pbar = tqdm(total=self.buffer_size, desc="Buffer monitor", miniters=1)
 
         for d in data:

src/lm_saes/config.py

@@ -231,6 +231,13 @@ def __le__(self, other: "ActivationFactoryTarget") -> bool:
         return self.stage <= other.stage
 
 
+class BufferShuffleConfig(BaseConfig):
+    perm_seed: int = 42
+    """ Perm seed for aligned permutation for generating activations. If `None`, will not use manual seed for Generator. """
+    generator_device: Optional[str]= None
+    """ The device to be assigned for the torch.Generator. If 'None', generator will be initialized on cpu as pytorch default. """
+    
+    
 class ActivationFactoryConfig(BaseConfig):
     sources: list[ActivationFactoryDatasetSource | ActivationFactoryActivationsSource]
     """ List of sources to use for activations. Can be a dataset or a path to activations. """
@@ -254,6 +261,8 @@ class ActivationFactoryConfig(BaseConfig):
         else None
     )
     """ Buffer size for online shuffling. If `None`, no shuffling will be performed. """
+    buffer_shuffle_config: 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

@@ -17,6 +17,7 @@
     ActivationFactoryTarget,
     ActivationWriterConfig,
     BaseSAEConfig,
+    BufferShuffleConfig,
     DatasetConfig,
     FeatureAnalyzerConfig,
     InitializerConfig,
@@ -123,6 +124,9 @@ class GenerateActivationsSettings(BaseSettings):
 
     buffer_size: Optional[int] = None
     """Size of the buffer for activation generation"""
+    
+    buffer_shuffle_config: Optional[BufferShuffleConfig] = None
+    """"Manual seed and device of generator for generating randomperm in buffer"""
 
     total_tokens: Optional[int] = None
     """Optional total number of tokens to generate"""
@@ -198,6 +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,
     )
 
     # Configure activation writer