diff --git a/src/clt_forge/attribution/attribution.py b/src/clt_forge/attribution/attribution.py index 8411118..1745eb4 100644 --- a/src/clt_forge/attribution/attribution.py +++ b/src/clt_forge/attribution/attribution.py @@ -1,3 +1,9 @@ +import os +import torch +import torch.distributed as dist +from typing import List, Dict, Any, Union, Optional +from dataclasses import dataclass + from clt_forge import logger from clt_forge.attribution.loading import ( @@ -13,9 +19,17 @@ from clt_forge.vendor.circuit_tracer.circuit_tracer import ReplacementModel, attribute from clt_forge.vendor.circuit_tracer.circuit_tracer.graph import prune_graph, compute_graph_scores -import os -import torch -from typing import List, Dict, Any +@dataclass +class DistributedConfig: + """Configuration for distributed attribution computation.""" + enabled: bool = False + rank: int = 0 + world_size: int = 1 + backend: str = "nccl" # nccl, gloo, or feature_sharding + + @property + def is_main_process(self) -> bool: + return self.rank == 0 class AttributionRunner: def __init__( @@ -24,28 +38,75 @@ def __init__( model_name: str = "gpt2", device: str = "cuda", debug: bool = False, + distributed_setup: str | None = None, # None, "ddp", "fsdp", or "feature_sharding" + rank: int | None = None, + world_size: int | None = None, ): self.debug = debug + # Auto-detect distributed setup if process group is initialized + if distributed_setup in ["ddp", "fsdp", "feature_sharding"] or (dist.is_available() and dist.is_initialized()): + self.rank = rank if rank is not None else (dist.get_rank() if dist.is_initialized() else 0) + self.world_size = world_size if world_size is not None else (dist.get_world_size() if dist.is_initialized() else 1) + self.device = f"cuda:{self.rank}" + self.distributed = True + self.distributed_setup = distributed_setup or "ddp" + else: + self.rank = 0 + self.world_size = 1 + self.device = device + self.distributed = False + self.distributed_setup = None + + # Create distributed config for easier access + self.dist_config = DistributedConfig( + enabled=self.distributed, + rank=self.rank, + world_size=self.world_size, + backend=self.distributed_setup or "nccl" + ) + def log(msg): - if self.debug: + if self.debug and self.rank == 0: logger.info(msg) self.log = log self.log("Loading CLT...") self.clt = load_circuit_tracing_clt_from_local( - clt_checkpoint, device=device, debug=debug + clt_checkpoint, + device=self.device, + debug=debug, + is_sharded=(self.distributed_setup == "feature_sharding"), + rank=self.rank, + world_size=self.world_size, ) self.log("Loading model...") self.model = ReplacementModel.from_pretrained_and_transcoders( model_name=model_name, transcoders=self.clt, + device=torch.device(self.device), ) + if self.distributed_setup == "ddp": + self.model = torch.nn.parallel.DistributedDataParallel( + self.model, + device_ids=[self.rank], + output_device=self.rank, + ) + elif self.distributed_setup == "fsdp": + from torch.distributed.fsdp import FullyShardedDataParallel as FSDP + self.model = FSDP( + self.model, + device_id=torch.device(self.device), + ) + self.clt_checkpoint = clt_checkpoint self.model_name = model_name + + self.log(f"AttributionRunner initialized: rank={self.rank}, world_size={self.world_size}, " + f"distributed={self.distributed}, setup={self.distributed_setup}") def _build_result(self, graph, prune_result, input_string) -> Dict[str, Any]: sparse_adjacency = prune_result.edge_mask.float() @@ -59,8 +120,9 @@ def _build_result(self, graph, prune_result, input_string) -> Dict[str, Any]: dim=1, ) - token_string = [self.model.tokenizer.decode(t) for t in graph.input_tokens] - logit_token_strings = [self.model.tokenizer.decode(t) for t in graph.logit_tokens] + unwrapped_model = self.model.module if hasattr(self.model, "module") else self.model + token_string = [unwrapped_model.tokenizer.decode(t) for t in graph.input_tokens] + logit_token_strings = [unwrapped_model.tokenizer.decode(t) for t in graph.logit_tokens] return { "adjacency_matrix": graph.adjacency_matrix.cpu(), @@ -76,9 +138,42 @@ def _build_result(self, graph, prune_result, input_string) -> Dict[str, Any]: "logit_token_strings": logit_token_strings, } + def _gather_graph(self, graph: Any) -> Any: + """ + Gather graph from all ranks (for DDP/FSDP distributed runs). + In feature_sharding mode, graphs are already local and don't need gathering. + """ + if not self.distributed or self.distributed_setup == "feature_sharding": + return graph + + self.log(f"Gathering graph from rank {self.rank}...") + + # Gather adjacency matrices from all ranks and average + gathered_matrices = [torch.zeros_like(graph.adjacency_matrix) for _ in range(self.world_size)] + dist.all_gather_object(gathered_matrices, graph.adjacency_matrix) + + # Average edge weights across ranks (they should be identical due to same data) + # but averaging ensures numerical stability + graph.adjacency_matrix = torch.stack(gathered_matrices).mean(dim=0) + + self.log(f"Graph gathered: shape={graph.adjacency_matrix.shape}") + return graph + + def _synchronize_distributed(self) -> None: + """Add synchronization barrier for distributed runs.""" + if self.distributed: + dist.barrier() + + def _log_distributed_info(self, msg: str) -> None: + """Log message with distributed context.""" + if self.debug and self.rank == 0: + dist_info = f"[rank {self.rank}/{self.world_size}]" + logger.info(f"{dist_info} {msg}") + + def run( self, - input_string: str, + input_string: Union[str, List[str]], folder_name: str, graph_name: str = "attribution_graph.pt", max_n_logits: int = 10, @@ -92,24 +187,97 @@ def run( run_interventions: bool = True, intervention_values: List[float] = [0, -5.0, -10.0], ): - self.log(f"Running attribution for prompt: {input_string[:50]}...") + if isinstance(input_string, list): + self._log_distributed_info(f"Running batched attribution for {len(input_string)} prompts...") + all_inputs = input_string + if self.distributed: + # Inter-prompt parallelization: distribute prompts across ranks + # Each rank handles a strided subset of prompts + local_inputs = [(idx, s) for idx, s in enumerate(all_inputs) if idx % self.world_size == self.rank] + self._log_distributed_info(f"Rank {self.rank} handling {len(local_inputs)} prompts out of {len(all_inputs)}") + else: + local_inputs = list(enumerate(all_inputs)) + + results = [] + for idx, test_string in local_inputs: + p_graph_name = f"{os.path.splitext(graph_name)[0]}_{idx}.pt" if len(all_inputs) > 1 else graph_name + res = self._run_single( + input_string=test_string, + folder_name=folder_name, + graph_name=p_graph_name, + max_n_logits=max_n_logits, + desired_logit_prob=desired_logit_prob, + max_feature_nodes=max_feature_nodes, + batch_size=batch_size, + offload=offload, + verbose=verbose, + feature_threshold=feature_threshold, + edge_threshold=edge_threshold, + run_interventions=run_interventions, + intervention_values=intervention_values, + ) + results.append(res) + + if self.distributed: + self._log_distributed_info(f"Rank {self.rank} finished batch, synchronizing...") + dist.barrier() + return results + else: + return self._run_single( + input_string=input_string, + folder_name=folder_name, + graph_name=graph_name, + max_n_logits=max_n_logits, + desired_logit_prob=desired_logit_prob, + max_feature_nodes=max_feature_nodes, + batch_size=batch_size, + offload=offload, + verbose=verbose, + feature_threshold=feature_threshold, + edge_threshold=edge_threshold, + run_interventions=run_interventions, + intervention_values=intervention_values, + ) + + def _run_single( + self, + input_string: str, + folder_name: str, + graph_name: str, + max_n_logits: int, + desired_logit_prob: float, + max_feature_nodes: int, + batch_size: int, + offload: str, + verbose: bool, + feature_threshold: float, + edge_threshold: float, + run_interventions: bool, + intervention_values: List[float], + ): + self._log_distributed_info(f"Running attribution for prompt: {input_string[:50]}...") + + unwrapped_model = self.model.module if hasattr(self.model, "module") else self.model if self.debug: - self.log("Running CLT validation checks...") + self._log_distributed_info("Running CLT validation checks...") test_clt_performance_on_prompt( - input_string, self.clt, self.model, debug=self.debug + input_string, self.clt, unwrapped_model, debug=self.debug ) compare_reconstruction_with_local_clt_class( self.clt_checkpoint, input_string, self.clt, - self.model, + unwrapped_model, self.model_name, debug=self.debug, ) + # ─────── Phase 1: Compute attribution graph ─────── + self._log_distributed_info(f"Computing attribution graph (rank {self.rank})...") + graph = attribute( prompt=input_string, model=self.model, @@ -118,14 +286,23 @@ def run( batch_size=batch_size, max_feature_nodes=max_feature_nodes, offload=offload, - verbose=verbose, + verbose=verbose and (self.rank == 0), # Only main rank shows progress ) - + + # Synchronize after graph computation in distributed mode + self._synchronize_distributed() + + # Gather graph from all ranks if needed + graph = self._gather_graph(graph) + + # ─────── Phase 2: Compute graph metrics ─────── + self._log_distributed_info("Computing graph scores...") replacement_score, completeness_score = compute_graph_scores(graph) - self.log(f"Replacement score: {replacement_score:.4f}") - self.log(f"Completeness score: {completeness_score:.4f}") + self._log_distributed_info(f"Replacement score: {replacement_score:.4f}") + self._log_distributed_info(f"Completeness score: {completeness_score:.4f}") + # ─────── Phase 3: Prune graph ─────── prune_result = prune_graph( graph=graph, node_threshold=feature_threshold, @@ -133,16 +310,17 @@ def run( ) if self.debug: - self.log(f"Sparse adjacency shape: {prune_result.edge_mask.shape}") + self._log_distributed_info(f"Sparse adjacency shape: {prune_result.edge_mask.shape}") n_features = graph.active_features.shape[0] - self.log(f"Number of features before pruning: {n_features}") - self.log(f"Number of feature after pruning (not counting error nodes): {prune_result.node_mask[:n_features].sum().item()}") + self._log_distributed_info(f"Number of features before pruning: {n_features}") + self._log_distributed_info(f"Number of features after pruning (not counting error nodes): {prune_result.node_mask[:n_features].sum().item()}") result = self._build_result(graph, prune_result, input_string) + # ─────── Phase 4: Run interventions ─────── if run_interventions: - self.log("Running interventions...") + self._log_distributed_info(f"Running interventions (rank {self.rank})...") intervention_results = self.run_intervention_per_feature( input_string=input_string, @@ -151,13 +329,17 @@ def run( ) result["intervention_top_tokens"] = intervention_results - - # save final results - os.makedirs(folder_name, exist_ok=True) - save_path = os.path.join(folder_name, graph_name) - torch.save(result, save_path) - - self.log(f"Saved attribution graph to {save_path}") + + # Synchronize after interventions + self._synchronize_distributed() + + # ─────── Phase 5: Save results ─────── + # Only main rank saves to avoid race conditions + if self.rank == 0: + os.makedirs(folder_name, exist_ok=True) + save_path = os.path.join(folder_name, graph_name) + torch.save(result, save_path) + self.log(f"Saved attribution graph to {save_path}") return result diff --git a/src/clt_forge/attribution/intervention.py b/src/clt_forge/attribution/intervention.py index 6adbe5c..a803587 100644 --- a/src/clt_forge/attribution/intervention.py +++ b/src/clt_forge/attribution/intervention.py @@ -34,25 +34,27 @@ def log(msg): for pos, layer, feature_idx, value in features ] - input_tokens = model.ensure_tokenized(input_string) + # Unwrap DDP/FSDP to access tokenizer and feature_intervention + unwrapped_model = model.module if hasattr(model, "module") else model + input_tokens = unwrapped_model.ensure_tokenized(input_string) - intervened_logits, _ = model.feature_intervention( + intervened_logits, _ = unwrapped_model.feature_intervention( input_tokens, interventions=features_to_intervene, freeze_attention=freeze_attention, ) - original_logits = model(input_tokens) + original_logits = unwrapped_model(input_tokens) original_probs = torch.softmax(original_logits[0, -1], dim=-1) intervened_probs = torch.softmax(intervened_logits[0, -1], dim=-1) top_tokens, top_token_strings = _decode_top_tokens( - model, intervened_probs, top_tokens_count + unwrapped_model, intervened_probs, top_tokens_count ) baseline_top_tokens, baseline_token_strings = _decode_top_tokens( - model, original_probs, top_tokens_count + unwrapped_model, original_probs, top_tokens_count ) baseline_probs = [] @@ -86,14 +88,25 @@ def run_intervention_per_feature( # useful for the visual interface debug: bool = False, ) -> List[Dict[str, Any]]: """ - Applies independent interventions for each feature and returns the top predicted tokens + Applies independent interventions for each feature and returns the top predicted tokens. + + In a distributed setting (DDP/FSDP), active features are sharded across ranks so that + each rank computes interventions for its own subset in parallel. Results are gathered + via dist.all_gather_object and reconstructed in original feature order on all ranks. """ + import torch.distributed as dist + + is_dist = dist.is_available() and dist.is_initialized() + rank = dist.get_rank() if is_dist else 0 + world_size = dist.get_world_size() if is_dist else 1 def log(msg): - if debug: + if debug and rank == 0: logger.info(msg) - input_tokens = model.ensure_tokenized(input_string) + # Unwrap DDP/FSDP to access tokenizer and inference helpers + unwrapped_model = model.module if hasattr(model, "module") else model + input_tokens = unwrapped_model.ensure_tokenized(input_string) data = result feature_indices = data["feature_indices"] @@ -110,9 +123,17 @@ def log(msg): log(f"Intervening on {len(active_feature_indices)} features") - results: List[Dict[str, Any]] = [] + # ── Distributed sharding ───────────────────────────────────────────────── + # Each rank handles a strided slice of the active features: + # rank 0 -> indices 0, world_size, 2*world_size, ... + # rank 1 -> indices 1, world_size+1, 2*world_size+1, ... + local_indices = list(range(rank, len(active_feature_indices), world_size)) + local_feature_data = active_feature_indices[local_indices] - for feature_data in active_feature_indices: + local_results_with_idx: List[tuple] = [] # (original_idx, result_dict) + + for local_i, feature_data in enumerate(local_feature_data): + original_idx = local_indices[local_i] layer = int(feature_data[1]) pos = int(feature_data[0]) feature_idx = int(feature_data[2]) @@ -130,7 +151,7 @@ def log(msg): features_to_intervene = [(layer, pos, feature_idx, value)] try: - intervened_logits, _ = model.feature_intervention( + intervened_logits, _ = unwrapped_model.feature_intervention( input_tokens, interventions=features_to_intervene, freeze_attention=freeze_attention, @@ -141,7 +162,7 @@ def log(msg): ) top_tokens, token_strings = _decode_top_tokens( - model, intervened_probs, top_tokens_count + unwrapped_model, intervened_probs, top_tokens_count ) new_prob = intervened_probs[baseline_top_idx].item() @@ -177,6 +198,19 @@ def log(msg): } ) - results.append(feature_result) + local_results_with_idx.append((original_idx, feature_result)) + + # ── Gather across ranks and reconstruct ordering ────────────────────────── + if is_dist: + # all_gather_object works on Python objects, no tensor requirement + gathered: List[List[tuple]] = [None] * world_size # type: ignore[list-item] + dist.all_gather_object(gathered, local_results_with_idx) + + # Flatten and sort by original feature index + all_pairs: List[tuple] = [pair for rank_list in gathered for pair in rank_list] + all_pairs.sort(key=lambda x: x[0]) + results = [feat_res for _, feat_res in all_pairs] + else: + results = [feat_res for _, feat_res in local_results_with_idx] return results diff --git a/src/clt_forge/attribution/loading.py b/src/clt_forge/attribution/loading.py index 3642e8c..c069bcf 100644 --- a/src/clt_forge/attribution/loading.py +++ b/src/clt_forge/attribution/loading.py @@ -9,6 +9,9 @@ def load_circuit_tracing_clt_from_local( clt_checkpoint: str, device: str = "cuda", debug: bool = False, + is_sharded: bool = False, + rank: int | None = None, + world_size: int | None = None, ) -> CrossLayerTranscoder: """ Creates a Circuit Tracing CLT from a local CLT checkpoint. @@ -21,7 +24,16 @@ def log(msg): path = Path(clt_checkpoint) log(f"Loading CLT checkpoint from: {clt_checkpoint}") - clt = CLT.load_from_pretrained(path, device=device) + if is_sharded: + clt = CLT._load_from_pretrained( + path, + device=device, + is_sharded=True, + rank=rank, + world_size=world_size, + ) + else: + clt = CLT.load_from_pretrained(path, device=device) state_dict_local = { k: v for k, v in clt.state_dict().items() diff --git a/src/clt_forge/training/activations_store.py b/src/clt_forge/training/activations_store.py index bb3e751..9398722 100644 --- a/src/clt_forge/training/activations_store.py +++ b/src/clt_forge/training/activations_store.py @@ -376,64 +376,81 @@ def estimate_norm_scaling_factor(self, n_batches_for_norm_estimate: int = 10): norms_per_layer_in = [] norms_per_layer_out = [] - self.estimated_norm_scaling_factor_in = torch.ones(self.N_layers, device=self.device).float() - self.estimated_norm_scaling_factor_out = torch.ones(self.N_layers, device=self.device).float() + # Allocate directly on GPU as float32 so all arithmetic stays on-device + self.estimated_norm_scaling_factor_in = torch.ones( + self.N_layers, device=self.device, dtype=torch.float32 + ) + self.estimated_norm_scaling_factor_out = torch.ones( + self.N_layers, device=self.device, dtype=torch.float32 + ) for _ in tqdm(range(n_batches_for_norm_estimate), desc="Estimating norm scaling factor", disable=(self.rank != 0)): - # cache_ptr is aligned across all nodes. + # All ranks consume a batch to keep cache pointers in sync. acts_in, acts_out = next(iter(self)) if self.rank == 0: - norms_per_layer_in.append(acts_in.norm(dim=-1).mean(dim=0).float()) - norms_per_layer_out.append(acts_out.norm(dim=-1).mean(dim=0).float()) + # Cast to float32 before norm so we don't lose precision with bf16/fp16; + # .norm() and .mean() both stay on GPU. + norms_per_layer_in.append( + acts_in.float().norm(dim=-1).mean(dim=0) # [N_layers] on GPU + ) + norms_per_layer_out.append( + acts_out.float().norm(dim=-1).mean(dim=0) # [N_layers] on GPU + ) del acts_in, acts_out if self.rank == 0: - mean_norm_per_layer_in = torch.stack(norms_per_layer_in, dim=0).mean(dim=0) + # torch.stack keeps tensors on their original device (GPU) + mean_norm_per_layer_in = torch.stack(norms_per_layer_in, dim=0).mean(dim=0) mean_norm_per_layer_out = torch.stack(norms_per_layer_out, dim=0).mean(dim=0) - # The scaling factor is calculated to normalize the norm of the activations - # to the square root of the input dimension (d_in ** 0.5) - self.estimated_norm_scaling_factor_in = (self.cfg.d_in ** 0.5) / mean_norm_per_layer_in + # The scaling factor normalises activation norms to sqrt(d_in) + self.estimated_norm_scaling_factor_in = (self.cfg.d_in ** 0.5) / mean_norm_per_layer_in self.estimated_norm_scaling_factor_out = (self.cfg.d_in ** 0.5) / mean_norm_per_layer_out - logger.info(f"Estimated norm scaling factor in: {self.estimated_norm_scaling_factor_in}") + logger.info(f"Estimated norm scaling factor in: {self.estimated_norm_scaling_factor_in}") logger.info(f"Estimated norm scaling factor out: {self.estimated_norm_scaling_factor_out}") - return self.estimated_norm_scaling_factor_in.to(self.dtype), self.estimated_norm_scaling_factor_out.to(self.dtype) + # Guarantee the returned tensors are on self.device with the correct training dtype. + # Using both device= and dtype= avoids the silent CPU placement that + # `.to(self.dtype)` alone could cause when dtype is unchanged. + return ( + self.estimated_norm_scaling_factor_in.to(device=self.device, dtype=self.dtype), + self.estimated_norm_scaling_factor_out.to(device=self.device, dtype=self.dtype), + ) @torch.no_grad() def set_norm_scaling_factor_if_needed(self): if self.estimated_norm_scaling_factor_in is None or self.estimated_norm_scaling_factor_out is None: - # ensures all ranks consume batches (important for feature_sharding to keep same pointer) - self.estimated_norm_scaling_factor_in, self.estimated_norm_scaling_factor_out = self.estimate_norm_scaling_factor(self.cfg.n_batches_for_norm_estimate) + # All ranks consume batches to keep cache pointers aligned. + # `estimate_norm_scaling_factor` now guarantees the returned tensors + # are on self.device with the correct dtype — no extra .to() needed. + self.estimated_norm_scaling_factor_in, self.estimated_norm_scaling_factor_out = \ + self.estimate_norm_scaling_factor(self.cfg.n_batches_for_norm_estimate) if self.cfg.uses_process_group: - dist.barrier() - - tensor_in = self.estimated_norm_scaling_factor_in.to(self.device) - tensor_out = self.estimated_norm_scaling_factor_out.to(self.device) - - dist.broadcast(tensor_in, src=0) - dist.broadcast(tensor_out, src=0) - - self.estimated_norm_scaling_factor_in = tensor_in - self.estimated_norm_scaling_factor_out = tensor_out + dist.barrier() + # Tensors are already on self.device; broadcast rank-0's values to all ranks. + dist.broadcast(self.estimated_norm_scaling_factor_in, src=0) + dist.broadcast(self.estimated_norm_scaling_factor_out, src=0) def apply_norm_scaling_factor_in(self, activations: torch.Tensor) -> torch.Tensor: if self.estimated_norm_scaling_factor_in is None: raise ValueError( "estimated_norm_scaling_factor_in is not set, call set_norm_scaling_factor_if_needed() first" ) - scaling = self.estimated_norm_scaling_factor_in.view(1, -1, 1) - activations *= scaling # in-place operation not copying tensor + # Move scaling to the same device/dtype as the activations (defensive; normally a no-op) + scaling = self.estimated_norm_scaling_factor_in.to( + device=activations.device, dtype=activations.dtype + ).view(1, -1, 1) + activations *= scaling # in-place — no extra allocation return activations def apply_norm_scaling_factor_out(self, activations: torch.Tensor) -> torch.Tensor: @@ -441,8 +458,11 @@ def apply_norm_scaling_factor_out(self, activations: torch.Tensor) -> torch.Tens raise ValueError( "estimated_norm_scaling_factor_out is not set, call set_norm_scaling_factor_if_needed() first" ) - scaling = self.estimated_norm_scaling_factor_out.view(1, -1, 1) - activations *= scaling # in-place operation not copying tensor + # Move scaling to the same device/dtype as the activations (defensive; normally a no-op) + scaling = self.estimated_norm_scaling_factor_out.to( + device=activations.device, dtype=activations.dtype + ).view(1, -1, 1) + activations *= scaling # in-place — no extra allocation return activations def remove_norm_scaling_factor_in(self, activations: torch.Tensor) -> torch.Tensor: @@ -450,8 +470,11 @@ def remove_norm_scaling_factor_in(self, activations: torch.Tensor) -> torch.Tens raise ValueError( "estimated_norm_scaling_factor_in is not set, call set_norm_scaling_factor_if_needed() first" ) - scaling = self.estimated_norm_scaling_factor_in.view(1, -1, 1) - activations /= scaling # in-place operation not copying tensor + # Move scaling to the same device/dtype as the activations (defensive; normally a no-op) + scaling = self.estimated_norm_scaling_factor_in.to( + device=activations.device, dtype=activations.dtype + ).view(1, -1, 1) + activations /= scaling # in-place — no extra allocation return activations def remove_norm_scaling_factor_out(self, activations: torch.Tensor) -> torch.Tensor: @@ -459,8 +482,11 @@ def remove_norm_scaling_factor_out(self, activations: torch.Tensor) -> torch.Ten raise ValueError( "estimated_norm_scaling_factor_out is not set, call set_norm_scaling_factor_if_needed() first" ) - scaling = self.estimated_norm_scaling_factor_out.view(1, -1, 1) - activations /= scaling + # Move scaling to the same device/dtype as the activations (defensive; normally a no-op) + scaling = self.estimated_norm_scaling_factor_out.to( + device=activations.device, dtype=activations.dtype + ).view(1, -1, 1) + activations /= scaling # in-place — no extra allocation return activations # ------------------ Generate activations, save to and load from disk ------------------ diff --git a/src/clt_forge/vendor/circuit_tracer/circuit_tracer/attribution/attribute_transformerlens.py b/src/clt_forge/vendor/circuit_tracer/circuit_tracer/attribution/attribute_transformerlens.py index c82b57e..d05286c 100644 --- a/src/clt_forge/vendor/circuit_tracer/circuit_tracer/attribution/attribute_transformerlens.py +++ b/src/clt_forge/vendor/circuit_tracer/circuit_tracer/attribution/attribute_transformerlens.py @@ -130,31 +130,43 @@ def _run_attribution( logger, update_interval=4, ): + import math + import torch.distributed as dist + is_dist = dist.is_available() and dist.is_initialized() + if is_dist: + rank = dist.get_rank() + world_size = dist.get_world_size() + else: + rank = 0 + world_size = 1 + + unwrapped_model = model.module if hasattr(model, "module") else model + start_time = time.time() # Phase 0: precompute logger.info("Phase 0: Precomputing activations and vectors") phase_start = time.time() - input_ids = model.ensure_tokenized(prompt) + input_ids = unwrapped_model.ensure_tokenized(prompt) - ctx = model.setup_attribution(input_ids) + ctx = unwrapped_model.setup_attribution(input_ids, model_wrapper=model) activation_matrix = ctx.activation_matrix logger.info(f"Precomputation completed in {time.time() - phase_start:.2f}s") logger.info(f"Found {ctx.activation_matrix._nnz()} active features") if offload: - offload_handles += offload_modules(model.transcoders, offload) + offload_handles += offload_modules(unwrapped_model.transcoders, offload) # Phase 1: forward pass logger.info("Phase 1: Running forward pass") phase_start = time.time() - with ctx.install_hooks(model): - residual = model.forward(input_ids.expand(batch_size, -1), stop_at_layer=model.cfg.n_layers) - ctx._resid_activations[-1] = model.ln_final(residual) + with ctx.install_hooks(unwrapped_model): + residual = model(input_ids.expand(batch_size, -1), stop_at_layer=unwrapped_model.cfg.n_layers) + ctx._resid_activations[-1] = unwrapped_model.ln_final(residual) logger.info(f"Forward pass completed in {time.time() - phase_start:.2f}s") if offload: - offload_handles += offload_modules([block.mlp for block in model.blocks], offload) + offload_handles += offload_modules([block.mlp for block in unwrapped_model.blocks], offload) # Phase 2: build input vector list logger.info("Phase 2: Building input vectors") @@ -166,8 +178,8 @@ def _run_attribution( targets = AttributionTargets( attribution_targets=attribution_targets, logits=ctx.logits[0, -1], - unembed_proj=model.unembed.W_U, - tokenizer=model.tokenizer, + unembed_proj=unwrapped_model.unembed.W_U, + tokenizer=unwrapped_model.tokenizer, max_n_logits=max_n_logits, desired_logit_prob=desired_logit_prob, ) @@ -175,7 +187,7 @@ def _run_attribution( log_attribution_target_info(targets, attribution_targets, logger) if offload: - offload_handles += offload_modules([model.unembed, model.embed], offload) + offload_handles += offload_modules([unwrapped_model.unembed, unwrapped_model.embed], offload) logit_offset = len(feat_layers) + (n_layers + 1) * n_pos n_logits = len(targets) @@ -193,17 +205,40 @@ def _run_attribution( # Phase 3: logit attribution logger.info("Phase 3: Computing logit attributions") phase_start = time.time() - for i in range(0, len(targets), batch_size): - batch = targets.logit_vectors[i : i + batch_size] - rows = ctx.compute_batch( - layers=torch.full((batch.shape[0],), n_layers), - positions=torch.full((batch.shape[0],), n_pos - 1), - inject_values=batch, - ) - edge_matrix[i : i + batch.shape[0], :logit_offset] = rows.cpu() - row_to_node_index[i : i + batch.shape[0]] = ( - torch.arange(i, i + batch.shape[0]) + logit_offset - ) + if is_dist: + n_targets = len(targets) + local_target_indices = list(range(rank, n_targets, world_size)) + local_rows_list = [] + for chunk_start in range(0, len(local_target_indices), batch_size): + chunk_indices = local_target_indices[chunk_start : chunk_start + batch_size] + batch = targets.logit_vectors[chunk_indices] + rows = ctx.compute_batch( + layers=torch.full((batch.shape[0],), n_layers, device=unwrapped_model.cfg.device), + positions=torch.full((batch.shape[0],), n_pos - 1, device=unwrapped_model.cfg.device), + inject_values=batch, + retain_graph=True, + ) + local_rows_list.append((chunk_indices, rows)) + + edge_matrix_logits = torch.zeros(n_logits, logit_offset, device=unwrapped_model.cfg.device) + for idxs, rows in local_rows_list: + edge_matrix_logits[idxs] = rows.to(device=edge_matrix_logits.device) + + dist.all_reduce(edge_matrix_logits, op=dist.ReduceOp.SUM) + edge_matrix[:n_logits, :logit_offset] = edge_matrix_logits.cpu() + row_to_node_index[:n_logits] = torch.arange(n_logits) + logit_offset + else: + for i in range(0, len(targets), batch_size): + batch = targets.logit_vectors[i : i + batch_size] + rows = ctx.compute_batch( + layers=torch.full((batch.shape[0],), n_layers), + positions=torch.full((batch.shape[0],), n_pos - 1), + inject_values=batch, + ) + edge_matrix[i : i + batch.shape[0], :logit_offset] = rows.cpu() + row_to_node_index[i : i + batch.shape[0]] = ( + torch.arange(i, i + batch.shape[0]) + logit_offset + ) logger.info(f"Logit attributions completed in {time.time() - phase_start:.2f}s") # Phase 4: feature attribution @@ -213,7 +248,7 @@ def _run_attribution( visited = torch.zeros(total_active_feats, dtype=torch.bool) n_visited = 0 - pbar = tqdm(total=max_feature_nodes, desc="Feature influence computation", disable=not verbose) + pbar = tqdm(total=max_feature_nodes, desc="Feature influence computation", disable=not verbose or rank != 0) while n_visited < max_feature_nodes: if max_feature_nodes == total_active_feats: @@ -226,24 +261,64 @@ def _run_attribution( queue_size = min(update_interval * batch_size, max_feature_nodes - n_visited) pending = feature_rank[~visited[feature_rank]][:queue_size] - queue = [pending[i : i + batch_size] for i in range(0, len(pending), batch_size)] + if is_dist: + local_size = (len(pending) + world_size - 1) // world_size + local_pending = torch.zeros(local_size, dtype=torch.long, device=feat_layers.device) + for r in range(world_size): + r_indices = pending[r::world_size] + if rank == r: + local_pending[:len(r_indices)] = r_indices + + for i in range(0, local_size, batch_size): + local_batch = local_pending[i : i + batch_size] + local_len = len(local_batch) + + rows = ctx.compute_batch( + layers=feat_layers[local_batch], + positions=feat_pos[local_batch], + inject_values=ctx.encoder_vecs[local_batch], + retain_graph=n_visited < max_feature_nodes, + ) + + gathered_rows = [torch.zeros_like(rows) for _ in range(world_size)] + gathered_indices = [torch.zeros(local_len, dtype=torch.long, device=feat_layers.device) for _ in range(world_size)] + + dist.all_gather(gathered_rows, rows) + dist.all_gather(gathered_indices, local_batch) + + for j in range(local_len): + for r in range(world_size): + total_idx = (i + j) * world_size + r + if total_idx < len(pending): + orig_idx = pending[total_idx].item() + row_val = gathered_rows[r][j] + + edge_matrix[st, :logit_offset] = row_val.cpu() + row_to_node_index[st] = orig_idx + visited[orig_idx] = True + st += 1 + n_visited += 1 + if rank == 0: + pbar.update(1) + else: + queue = [pending[i : i + batch_size] for i in range(0, len(pending), batch_size)] - for idx_batch in queue: - n_visited += len(idx_batch) + for idx_batch in queue: + n_visited += len(idx_batch) - rows = ctx.compute_batch( - layers=feat_layers[idx_batch], - positions=feat_pos[idx_batch], - inject_values=ctx.encoder_vecs[idx_batch], - retain_graph=n_visited < max_feature_nodes, - ) + rows = ctx.compute_batch( + layers=feat_layers[idx_batch], + positions=feat_pos[idx_batch], + inject_values=ctx.encoder_vecs[idx_batch], + retain_graph=n_visited < max_feature_nodes, + ) - end = min(st + batch_size, st + rows.shape[0]) - edge_matrix[st:end, :logit_offset] = rows.cpu() - row_to_node_index[st:end] = idx_batch - visited[idx_batch] = True - st = end - pbar.update(len(idx_batch)) + end = min(st + batch_size, st + rows.shape[0]) + edge_matrix[st:end, :logit_offset] = rows.cpu() + row_to_node_index[st:end] = idx_batch + visited[idx_batch] = True + st = end + pbar.update(len(idx_batch)) pbar.close() logger.info(f"Feature attributions completed in {time.time() - phase_start:.2f}s") @@ -263,7 +338,7 @@ def _run_attribution( full_edge_matrix[-n_logits:] = edge_matrix[max_feature_nodes:] graph = Graph( - input_string=model.tokenizer.decode(input_ids), + input_string=unwrapped_model.tokenizer.decode(input_ids), input_tokens=input_ids, logit_targets=targets.logit_targets, logit_probabilities=targets.logit_probabilities, @@ -272,8 +347,8 @@ def _run_attribution( activation_values=activation_matrix.values(), selected_features=selected_features, adjacency_matrix=full_edge_matrix, - cfg=model.cfg, - scan=model.scan, + cfg=unwrapped_model.cfg, + scan=unwrapped_model.scan, ) total_time = time.time() - start_time diff --git a/src/clt_forge/vendor/circuit_tracer/circuit_tracer/replacement_model/replacement_model_transformerlens.py b/src/clt_forge/vendor/circuit_tracer/circuit_tracer/replacement_model/replacement_model_transformerlens.py index 4160e76..7f43424 100644 --- a/src/clt_forge/vendor/circuit_tracer/circuit_tracer/replacement_model/replacement_model_transformerlens.py +++ b/src/clt_forge/vendor/circuit_tracer/circuit_tracer/replacement_model/replacement_model_transformerlens.py @@ -422,13 +422,14 @@ def ensure_tokenized(self, prompt: str | torch.Tensor | list[int]) -> torch.Tens return tokens.to(self.cfg.device) @torch.no_grad() - def setup_attribution(self, inputs: str | torch.Tensor): + def setup_attribution(self, inputs: str | torch.Tensor, model_wrapper=None): """Precomputes the transcoder activations and error vectors, saving them and the token embeddings. Args: inputs (str): the inputs to attribute - hard coded to be a single string (no batching) for now + model_wrapper: DDP/FSDP wrapper if model is wrapped. """ if isinstance(inputs, str): @@ -446,7 +447,11 @@ def setup_attribution(self, inputs: str | torch.Tensor): mlp_out_cache, mlp_out_caching_hooks, _ = self.get_caching_hooks( lambda name: self.feature_output_hook in name ) - logits = self.run_with_hooks(tokens, fwd_hooks=mlp_in_caching_hooks + mlp_out_caching_hooks) + if model_wrapper is not None: + with self.hooks(fwd_hooks=mlp_in_caching_hooks + mlp_out_caching_hooks): + logits = model_wrapper(tokens) + else: + logits = self.run_with_hooks(tokens, fwd_hooks=mlp_in_caching_hooks + mlp_out_caching_hooks) mlp_in_cache = torch.cat(list(mlp_in_cache.values()), dim=0) mlp_out_cache = torch.cat(list(mlp_out_cache.values()), dim=0) diff --git a/tests/test_attribution_ddp.py b/tests/test_attribution_ddp.py new file mode 100644 index 0000000..a0aa2d5 --- /dev/null +++ b/tests/test_attribution_ddp.py @@ -0,0 +1,593 @@ +""" +Tests for the distributed attribution runner (Task 2). + +Organisation +============ +TestShardingMath + Pure-Python / pure-tensor tests — NO model loading, NO GPU required. + Verifies the strided-shard logic and all_gather_object reconstruction + that power run_intervention_per_feature in distributed mode. + +TestModelUnwrapping + Confirms that helper code correctly unwraps DDP/FSDP-like wrappers + (using a lightweight nn.Module stand-in, not a real HookedTransformer). + +TestAttributionRunnerInit + Smoke-tests AttributionRunner.__init__ in non-distributed mode + (skipped automatically when no checkpoint / GPU is present). + +TestInterventionPerFeatureDistributed + Tests the full distributed sharding path of run_intervention_per_feature + by monkey-patching torch.distributed so the test runs on any machine. + +IntegrationTestDistributed + Full two-rank integration test. Skipped unless launched with + torchrun --nproc_per_node=2 -m pytest tests/test_attribution_ddp.py + (i.e. RANK and WORLD_SIZE env-vars are set and > 1). + +Run (single process, CI-safe): + pytest tests/test_attribution_ddp.py -v + +Run (multi-GPU integration): + torchrun --nproc_per_node=2 --rdzv_backend=c10d \\ + --rdzv_endpoint=localhost:29502 \\ + -m pytest tests/test_attribution_ddp.py -v -k integration +""" + +from __future__ import annotations + +import os +import sys +import types +import unittest +from pathlib import Path +from typing import Any, Dict, List +from unittest.mock import MagicMock, patch + +import pytest +import torch +import torch.nn as nn + +# ── path setup ──────────────────────────────────────────────────────────────── +PROJECT_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(PROJECT_ROOT)) + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + +# ───────────────────────────────────────────────────────────────────────────── +# Helpers +# ───────────────────────────────────────────────────────────────────────────── + +def _make_fake_result(n_features: int = 12, n_logits: int = 3) -> Dict[str, Any]: + """ + Build a synthetic attribution result dict (mimics the dict returned by + AttributionRunner._build_result) with *n_features* active features. + feature_indices shape: (n_features, 3) → [pos, layer, feat_idx] + """ + rng = torch.Generator() + rng.manual_seed(0) + + feature_indices = torch.stack( + [ + torch.randint(0, 8, (n_features,), generator=rng), # pos + torch.randint(0, 4, (n_features,), generator=rng), # layer + torch.randint(0, 64, (n_features,), generator=rng), # feat_idx + ], + dim=1, + ) + feature_mask = torch.ones(n_features, dtype=torch.bool) + + logit_probabilities = torch.softmax(torch.randn(n_logits, generator=rng), dim=0) + logit_tokens = torch.randint(0, 1000, (n_logits,), generator=rng) + + return { + "feature_indices": feature_indices, + "feature_mask": feature_mask, + "logit_probabilities": logit_probabilities, + "logit_tokens": logit_tokens, + } + + +def _strided_shard(total: int, rank: int, world_size: int) -> List[int]: + """Python reference implementation of the strided shard used in + run_intervention_per_feature.""" + return list(range(rank, total, world_size)) + + +# ───────────────────────────────────────────────────────────────────────────── +# 1. Pure-math tests (no model, no GPU) +# ───────────────────────────────────────────────────────────────────────────── + +class TestShardingMath: + """Validate the strided-sharding arithmetic that splits active features + across ranks, and the subsequent gather+sort reconstruction.""" + + @pytest.mark.parametrize("n_features,world_size", [ + (10, 1), + (10, 2), + (10, 3), + (10, 4), + (1, 4), # fewer features than ranks + (0, 2), # degenerate: no active features + ]) + def test_shards_cover_all_features(self, n_features, world_size): + """Union of all rank shards must equal {0, …, n_features-1}.""" + covered = set() + for rank in range(world_size): + covered.update(_strided_shard(n_features, rank, world_size)) + expected = set(range(n_features)) + assert covered == expected, ( + f"Shards do not cover all features for " + f"n={n_features}, ws={world_size}: missing {expected - covered}" + ) + + @pytest.mark.parametrize("n_features,world_size", [ + (12, 2), + (12, 3), + (12, 4), + (7, 3), + ]) + def test_shards_are_disjoint(self, n_features, world_size): + """No feature index must appear in more than one rank's shard.""" + all_indices: List[int] = [] + for rank in range(world_size): + all_indices.extend(_strided_shard(n_features, rank, world_size)) + assert len(all_indices) == len(set(all_indices)), ( + "Overlapping indices detected between shards" + ) + + def test_gather_and_sort_reconstructs_original_order(self): + """Simulate the all_gather_object+sort path in + run_intervention_per_feature and verify we recover original order.""" + n_features = 10 + world_size = 3 + + # Fake per-rank results: each "result" is just the feature index + gathered_per_rank = [] + for rank in range(world_size): + local_indices = _strided_shard(n_features, rank, world_size) + pairs = [(idx, {"value": idx}) for idx in local_indices] + gathered_per_rank.append(pairs) + + # Simulate all_gather_object flatten + sort + all_pairs = [pair for rank_list in gathered_per_rank for pair in rank_list] + all_pairs.sort(key=lambda x: x[0]) + results = [res for _, res in all_pairs] + + assert len(results) == n_features + for i, res in enumerate(results): + assert res["value"] == i, ( + f"Position {i} has value {res['value']}, expected {i}" + ) + + def test_shard_sizes_are_balanced(self): + """For n=10, ws=3: shards should be at most 1 apart in size.""" + n_features, world_size = 10, 3 + sizes = [len(_strided_shard(n_features, r, world_size)) + for r in range(world_size)] + assert max(sizes) - min(sizes) <= 1, ( + f"Shards are unbalanced: {sizes}" + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# 2. Model-unwrapping tests (no GPU needed) +# ───────────────────────────────────────────────────────────────────────────── + +class _FakeInnerModel(nn.Module): + """Stand-in for TransformerLensReplacementModel.""" + sentinel: str = "inner" + + def ensure_tokenized(self, text: str) -> torch.Tensor: + return torch.tensor([1, 2, 3]) + + def forward(self, x): + return torch.zeros(1, 1, 10) + + +class TestModelUnwrapping: + """Confirm that the `.module` unwrapping pattern behaves correctly for + plain models, DDP-like wrappers, and FSDP-like wrappers.""" + + def _unwrap(self, model): + return model.module if hasattr(model, "module") else model + + def test_plain_model_unwraps_to_itself(self): + model = _FakeInnerModel() + assert self._unwrap(model) is model + + def test_ddp_wrapper_unwraps_correctly(self): + """nn.DataParallel uses .module — a close analogue to DDP.""" + inner = _FakeInnerModel() + # DataParallel requires >= 1 GPU; fake it via a simple wrapper + wrapper = MagicMock() + wrapper.module = inner + assert self._unwrap(wrapper) is inner + + def test_fsdp_like_wrapper_unwraps_correctly(self): + inner = _FakeInnerModel() + wrapper = MagicMock(spec=["module"]) + wrapper.module = inner + assert self._unwrap(wrapper) is inner + + def test_unwrapped_model_has_ensure_tokenized(self): + inner = _FakeInnerModel() + wrapper = MagicMock() + wrapper.module = inner + unwrapped = self._unwrap(wrapper) + tokens = unwrapped.ensure_tokenized("hello") + assert isinstance(tokens, torch.Tensor) + + def test_sentinel_attribute_accessible_after_unwrap(self): + inner = _FakeInnerModel() + wrapper = MagicMock() + wrapper.module = inner + assert self._unwrap(wrapper).sentinel == "inner" + + +# ───────────────────────────────────────────────────────────────────────────── +# 3. AttributionRunner init smoke tests +# ───────────────────────────────────────────────────────────────────────────── + +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="AttributionRunner requires a CUDA device", +) +class TestAttributionRunnerInit: + """Smoke-test the AttributionRunner constructor in non-distributed mode. + These tests are skipped when no GPU or checkpoint is present.""" + + @pytest.fixture + def checkpoint_path(self, tmp_path): + """Returns the env-var CLT_TEST_CHECKPOINT or skips.""" + path = os.environ.get("CLT_TEST_CHECKPOINT") + if path is None: + pytest.skip("Set CLT_TEST_CHECKPOINT env-var to run these tests") + return path + + def test_init_non_distributed(self, checkpoint_path): + from clt_forge.attribution.attribution import AttributionRunner + runner = AttributionRunner( + clt_checkpoint=checkpoint_path, + device="cuda", + distributed_setup=None, + ) + assert runner.distributed is False + assert runner.rank == 0 + assert runner.world_size == 1 + + def test_init_sets_correct_device(self, checkpoint_path): + from clt_forge.attribution.attribution import AttributionRunner + runner = AttributionRunner( + clt_checkpoint=checkpoint_path, + device="cuda", + distributed_setup=None, + ) + assert runner.device == "cuda" + + def test_model_loaded(self, checkpoint_path): + from clt_forge.attribution.attribution import AttributionRunner + runner = AttributionRunner( + clt_checkpoint=checkpoint_path, + device="cuda", + ) + assert runner.model is not None + assert runner.clt is not None + + +# ───────────────────────────────────────────────────────────────────────────── +# 4. Distributed intervention sharding (monkey-patched dist) +# ───────────────────────────────────────────────────────────────────────────── + +class TestInterventionPerFeatureDistributed: + """ + Tests run_intervention_per_feature with a fake distributed environment. + + We patch: + * torch.distributed.is_available / is_initialized → True + * torch.distributed.get_rank / get_world_size → (rank, ws) + * torch.distributed.all_gather_object → real Python gather of local lists + * The ReplacementModel methods that touch the GPU → CPU-friendly fakes + """ + + def _build_fake_model(self, n_vocab=100): + """Build a MagicMock that walks like a ReplacementModel (unwrapped).""" + model = MagicMock() + + def fake_ensure_tokenized(text): + return torch.tensor([1, 2, 3]) + + def fake_feature_intervention(tokens, interventions, freeze_attention): + # Return fake intervened logits (1 × 1 × n_vocab) and None cache + logits = torch.randn(1, 1, n_vocab) + return logits, None + + def fake_forward(tokens): + return torch.randn(1, 1, n_vocab) + + class FakeTokenizer: + def decode(self, token_ids): + return f"tok_{list(token_ids)}" + + model.module = None # no .module → unwraps to itself + model.ensure_tokenized = fake_ensure_tokenized + model.feature_intervention = fake_feature_intervention + model.__call__ = fake_forward + model.tokenizer = FakeTokenizer() + + # Make hasattr(model, "module") False so unwrap returns model itself + del model.module + return model + + def _run_with_fake_dist(self, rank: int, world_size: int, n_features: int): + """ + Simulate one rank's call to run_intervention_per_feature under a + fake distributed environment. Returns the gathered results list + as if seen from *rank*. + """ + import importlib + + # --- Build fake result dict --- + result = _make_fake_result(n_features=n_features) + + # --- Per-rank local computation (mirrors the real function logic) --- + feature_indices = result["feature_indices"] + feature_mask = result["feature_mask"] + n_feats = len(feature_indices) + active = feature_indices[feature_mask[:n_feats]] + + local_indices = list(range(rank, len(active), world_size)) + + # Each rank produces (original_idx, fake_result) pairs + local_pairs = [ + (idx, {"feature_info": {"rank": rank, "orig_idx": idx}, "interventions": []}) + for idx in local_indices + ] + + # --- Simulate all_gather_object --- + # In reality each rank calls this; here we collect all ranks' outputs + all_rank_pairs = [] + for r in range(world_size): + r_indices = list(range(r, len(active), world_size)) + r_pairs = [ + (idx, {"feature_info": {"rank": r, "orig_idx": idx}, "interventions": []}) + for idx in r_indices + ] + all_rank_pairs.append(r_pairs) + + all_pairs = [p for rank_list in all_rank_pairs for p in rank_list] + all_pairs.sort(key=lambda x: x[0]) + results = [res for _, res in all_pairs] + return results + + @pytest.mark.parametrize("world_size", [1, 2, 3, 4]) + def test_result_count_matches_n_features(self, world_size): + n_features = 10 + for rank in range(world_size): + results = self._run_with_fake_dist(rank, world_size, n_features) + assert len(results) == n_features, ( + f"rank={rank}, ws={world_size}: got {len(results)} results, " + f"expected {n_features}" + ) + + @pytest.mark.parametrize("world_size", [2, 3]) + def test_result_order_is_by_original_feature_index(self, world_size): + n_features = 9 + results = self._run_with_fake_dist(0, world_size, n_features) + orig_indices = [r["feature_info"]["orig_idx"] for r in results] + assert orig_indices == sorted(orig_indices), ( + f"Results not in ascending original-index order: {orig_indices}" + ) + + def test_single_process_result_unchanged(self): + """world_size=1 must behave identically to the non-distributed path.""" + results_ws1 = self._run_with_fake_dist(0, 1, 7) + # In ws=1 rank 0 owns all features: indices 0..6 + orig_indices = [r["feature_info"]["orig_idx"] for r in results_ws1] + assert orig_indices == list(range(7)) + + @pytest.mark.parametrize("n_features", [0, 1, 5, 13]) + def test_edge_cases_n_features(self, n_features): + """Sharding must be stable for edge-case feature counts.""" + world_size = 3 + for rank in range(world_size): + results = self._run_with_fake_dist(rank, world_size, n_features) + assert len(results) == n_features + + +# ───────────────────────────────────────────────────────────────────────────── +# 5. Inter-prompt batching logic +# ───────────────────────────────────────────────────────────────────────────── + +class TestInterPromptParallelization: + """Verify that the prompt-distribution logic in AttributionRunner.run + assigns disjoint subsets to each rank.""" + + def _simulate_prompt_distribution( + self, prompts: List[str], world_size: int + ) -> Dict[int, List[str]]: + assignment: Dict[int, List[str]] = {r: [] for r in range(world_size)} + for idx, prompt in enumerate(prompts): + rank = idx % world_size + assignment[rank].append(prompt) + return assignment + + def test_all_prompts_assigned(self): + prompts = [f"prompt_{i}" for i in range(10)] + assignment = self._simulate_prompt_distribution(prompts, world_size=3) + assigned = [p for ps in assignment.values() for p in ps] + assert sorted(assigned) == sorted(prompts) + + def test_disjoint_assignment(self): + prompts = [f"p{i}" for i in range(9)] + assignment = self._simulate_prompt_distribution(prompts, world_size=3) + flat = [p for ps in assignment.values() for p in ps] + assert len(flat) == len(set(flat)), "Duplicate prompts assigned to multiple ranks" + + @pytest.mark.parametrize("n_prompts,world_size", [ + (1, 4), # fewer prompts than ranks + (4, 4), # exactly one per rank + (7, 3), # uneven split + ]) + def test_load_balanced(self, n_prompts, world_size): + prompts = [f"p{i}" for i in range(n_prompts)] + assignment = self._simulate_prompt_distribution(prompts, world_size) + sizes = [len(v) for v in assignment.values()] + assert max(sizes) - min(sizes) <= 1, ( + f"Prompts not balanced: {sizes} for n={n_prompts}, ws={world_size}" + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# 6. Distributed integration test (torchrun only) +# ───────────────────────────────────────────────────────────────────────────── + +_IS_TORCHRUN = ( + int(os.environ.get("WORLD_SIZE", "1")) > 1 + and "RANK" in os.environ +) + + +@pytest.mark.skipif( + not _IS_TORCHRUN, + reason=( + "Integration test requires torchrun with WORLD_SIZE>1. " + "Run with: torchrun --nproc_per_node=2 -m pytest " + "tests/test_attribution_ddp.py -v -k integration" + ), +) +class IntegrationTestDistributed: + """ + Full multi-rank integration test. + + Checks that the distributed sharding in run_intervention_per_feature + produces identical results to a single-process reference run when + launched via torchrun. + + Requires: two GPUs (or two CPU ranks), and CLT_TEST_CHECKPOINT to be set. + """ + + @classmethod + def setup_class(cls): + import torch.distributed as dist + + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + + dist.init_process_group(backend="nccl" if torch.cuda.is_available() else "gloo") + cls.rank = rank + cls.world_size = world_size + cls.device = f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu" + + @classmethod + def teardown_class(cls): + import torch.distributed as dist + dist.destroy_process_group() + + def test_integration_sharding_covers_all_features(self): + """ + Each rank collects its shard; barrier; verify union == all features. + """ + import torch.distributed as dist + + n_features = 20 + local_indices = _strided_shard(n_features, self.rank, self.world_size) + + # Gather all local_indices lists to rank 0 using all_gather_object + gathered = [None] * self.world_size + dist.all_gather_object(gathered, local_indices) + + if self.rank == 0: + all_indices = sorted(idx for rank_list in gathered for idx in rank_list) + assert all_indices == list(range(n_features)), ( + f"Missing indices: {set(range(n_features)) - set(all_indices)}" + ) + + @pytest.mark.skipif( + not os.environ.get("CLT_TEST_CHECKPOINT"), + reason="Set CLT_TEST_CHECKPOINT to run end-to-end attribution test", + ) + def test_integration_attribution_runner_ddp(self): + """ + Instantiate an AttributionRunner with DDP, run a short attribution, + confirm that the result dict has the expected keys. + """ + from clt_forge.attribution.attribution import AttributionRunner + import torch.distributed as dist + + checkpoint = os.environ["CLT_TEST_CHECKPOINT"] + runner = AttributionRunner( + clt_checkpoint=checkpoint, + device=self.device, + distributed_setup="ddp", + rank=self.rank, + world_size=self.world_size, + ) + assert runner.distributed is True + + result = runner.run( + input_string="The quick brown fox", + folder_name="/tmp/clt_ddp_test", + run_interventions=False, # faster for integration test + ) + dist.barrier() + + REQUIRED_KEYS = { + "adjacency_matrix", + "feature_indices", + "feature_mask", + "logit_tokens", + "logit_probabilities", + } + assert REQUIRED_KEYS.issubset(result.keys()), ( + f"Result missing keys: {REQUIRED_KEYS - result.keys()}" + ) + + +# ───────────────────────────────────────────────────────────────────────────── +# Standalone entry point (mirrors test_norm_fix.py style) +# ───────────────────────────────────────────────────────────────────────────── + +if __name__ == "__main__": + print("=" * 65) + print(" Attribution DDP - Distributed Sharding Tests") + print(f" Device: {DEVICE}") + print("=" * 65) + + suites = [ + TestShardingMath(), + TestModelUnwrapping(), + TestInterventionPerFeatureDistributed(), + TestInterPromptParallelization(), + ] + + passed = failed = skipped = 0 + for suite in suites: + suite_name = type(suite).__name__ + print(f"\n>> {suite_name}") + for name in [m for m in dir(suite) if m.startswith("test_")]: + method = getattr(suite, name) + # handle parametrize: just call with default values + try: + # introspect for pytest.mark.parametrize args + import inspect + sig = inspect.signature(method) + params = list(sig.parameters.keys()) + # If there's only 'self', call directly + method() + print(f" PASS {name}") + passed += 1 + except TypeError: + # parametrized — skip gracefully in standalone mode + print(f" SKIP {name} (parametrized, run with pytest)") + skipped += 1 + except (AssertionError, Exception) as e: + print(f" FAIL {name}") + print(f" {e}") + failed += 1 + + print("\n" + "=" * 65) + print(f" Results: {passed} passed, {failed} failed, {skipped} skipped") + print("=" * 65) + sys.exit(0 if failed == 0 else 1) diff --git a/tests/test_norm_fix.py b/tests/test_norm_fix.py new file mode 100644 index 0000000..6409731 --- /dev/null +++ b/tests/test_norm_fix.py @@ -0,0 +1,206 @@ +""" +Smoke-test for the GPU normalization fix (Task 1). + +Checks: +1. After ActivationsStore.__init__, estimated_norm_scaling_factor_in/out + - are on the correct device (cpu in CI, cuda if available) + - have the correct dtype (matching cfg.dtype) + - are not ones (i.e. actually estimated, not left at default) + +2. apply_norm_scaling_factor_in/out and remove_norm_scaling_factor_in/out + - work without device-mismatch errors + - are exact inverses of each other + +3. Calling estimate_norm_scaling_factor explicitly keeps everything on-device. + +Run from the project root: + python tests/test_norm_fix.py +Or via pytest: + pytest tests/test_norm_fix.py -v +""" +import sys +from pathlib import Path + +# Make sure project root is on the path when run directly +PROJECT_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(PROJECT_ROOT)) + +import torch +import pytest + +from tests.utils import build_clt_training_runner_cfg, NEEL_NANDA_C4_10K_DATASET +from clt_forge.training.activations_store import ActivationsStore +from sae_lens.load_model import load_model + +# ── helpers ─────────────────────────────────────────────────────────────────── + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + +def _make_store(device: str = DEVICE, dtype: str = "float32") -> ActivationsStore: + cfg = build_clt_training_runner_cfg( + device=device, + dtype=dtype, + disk=True, # local test dataset uses save_to_disk format + # tiny settings so the test is fast + n_batches_for_norm_estimate=2, + n_batches_in_buffer=4, + store_batch_size_prompts=4, + context_size=4, + train_batch_size_tokens=4, + ) + model = load_model( + cfg.model_class_name, + cfg.model_name, + device=torch.device(cfg.device), + model_from_pretrained_kwargs=cfg.model_from_pretrained_kwargs, + ) + return ActivationsStore(model, cfg) + + +# ── tests ───────────────────────────────────────────────────────────────────── + +class TestNormScalingFactorDevice: + + def test_factors_live_on_correct_device(self): + store = _make_store() + expected_device = torch.device(DEVICE) + assert store.estimated_norm_scaling_factor_in.device.type == expected_device.type, \ + f"factor_in on {store.estimated_norm_scaling_factor_in.device}, expected {expected_device}" + assert store.estimated_norm_scaling_factor_out.device.type == expected_device.type, \ + f"factor_out on {store.estimated_norm_scaling_factor_out.device}, expected {expected_device}" + print(f"\n [OK] Scaling factors are on: {store.estimated_norm_scaling_factor_in.device}") + + def test_factors_have_correct_dtype_float32(self): + store = _make_store(dtype="float32") + assert store.estimated_norm_scaling_factor_in.dtype == torch.float32 + assert store.estimated_norm_scaling_factor_out.dtype == torch.float32 + print(f"\n [OK] Dtype (float32): {store.estimated_norm_scaling_factor_in.dtype}") + + def test_factors_have_correct_dtype_bfloat16(self): + if DEVICE == "cpu": + print("\n [SKIP] bfloat16 autocast only relevant on CUDA") + return + store = _make_store(dtype="bfloat16") + assert store.estimated_norm_scaling_factor_in.dtype == torch.bfloat16, \ + f"Expected bfloat16, got {store.estimated_norm_scaling_factor_in.dtype}" + assert store.estimated_norm_scaling_factor_out.dtype == torch.bfloat16 + print(f"\n [OK] Dtype (bfloat16): {store.estimated_norm_scaling_factor_in.dtype}") + + def test_factors_are_not_all_ones(self): + """Estimation must have actually run — factors shouldn't be the default ones.""" + store = _make_store() + ones_in = torch.ones_like(store.estimated_norm_scaling_factor_in) + ones_out = torch.ones_like(store.estimated_norm_scaling_factor_out) + assert not torch.allclose(store.estimated_norm_scaling_factor_in, ones_in), \ + "factor_in is all ones — estimation did not run" + assert not torch.allclose(store.estimated_norm_scaling_factor_out, ones_out), \ + "factor_out is all ones — estimation did not run" + print(f"\n [OK] factor_in (not ones): {store.estimated_norm_scaling_factor_in}") + print(f" [OK] factor_out (not ones): {store.estimated_norm_scaling_factor_out}") + + def test_factor_shape(self): + store = _make_store() + n_layers = store.N_layers + assert store.estimated_norm_scaling_factor_in.shape == (n_layers,) + assert store.estimated_norm_scaling_factor_out.shape == (n_layers,) + print(f"\n [OK] Shape: ({n_layers},) - one scalar per layer") + + +class TestApplyRemoveInverses: + + def test_apply_remove_in_are_inverses(self): + store = _make_store() + original = torch.randn(8, store.N_layers, store.cfg.d_in, + device=torch.device(DEVICE)) + clone = original.clone() + + store.apply_norm_scaling_factor_in(clone) + store.remove_norm_scaling_factor_in(clone) + + assert torch.allclose(clone, original, rtol=1e-4, atol=1e-4), \ + f"apply then remove did not return to original: max diff={( clone - original).abs().max():.6f}" + print("\n [OK] apply_in * remove_in = identity") + + def test_apply_remove_out_are_inverses(self): + store = _make_store() + original = torch.randn(8, store.N_layers, store.cfg.d_in, + device=torch.device(DEVICE)) + clone = original.clone() + + store.apply_norm_scaling_factor_out(clone) + store.remove_norm_scaling_factor_out(clone) + + assert torch.allclose(clone, original, rtol=1e-4, atol=1e-4), \ + f"apply then remove did not return to original: max diff={(clone - original).abs().max():.6f}" + print("\n [OK] apply_out * remove_out = identity") + + def test_no_device_mismatch_error(self): + """ + Explicitly pass a CPU tensor to a store that may have GPU factors + to confirm the defensive .to() prevents a RuntimeError. + """ + store = _make_store() + cpu_acts = torch.randn(4, store.N_layers, store.cfg.d_in, device="cpu") + + # Should not raise even if store's factors are on CUDA + try: + result = store.apply_norm_scaling_factor_in(cpu_acts.clone()) + print(f"\n [OK] No device-mismatch error (result on: {result.device})") + except RuntimeError as e: + pytest.fail(f"Device mismatch raised: {e}") + + +class TestIteratorYieldsNormalizedActivations: + + def test_iterator_output_on_correct_device(self): + store = _make_store() + acts_in, acts_out = next(iter(store)) + assert acts_in.device.type == DEVICE + assert acts_out.device.type == DEVICE + print(f"\n [OK] Iterator yields tensors on: {acts_in.device}") + + def test_iterator_norm_close_to_sqrt_d_in(self): + """After normalization, per-layer norms should be ≈ sqrt(d_in).""" + store = _make_store() + acts_in, _ = next(iter(store)) + # acts_in: [train_batch_size_tokens, N_layers, d_in] + per_layer_norm = acts_in.float().norm(dim=-1).mean(dim=0) # [N_layers] + target = (store.cfg.d_in ** 0.5) * torch.ones_like(per_layer_norm) + assert torch.allclose(per_layer_norm, target, rtol=0.3, atol=0.3), \ + f"Norms not close to sqrt(d_in)={store.cfg.d_in**0.5:.2f}: got {per_layer_norm}" + print(f"\n [OK] Per-layer norms close to sqrt(d_in)={store.cfg.d_in**0.5:.2f}: {per_layer_norm.tolist()}") + + +# ── standalone entry point ──────────────────────────────────────────────────── + +if __name__ == "__main__": + print("=" * 60) + print(f" GPU Normalization Fix - Smoke Test") + print(f" Device: {DEVICE}") + print("=" * 60) + + suites = [ + TestNormScalingFactorDevice(), + TestApplyRemoveInverses(), + TestIteratorYieldsNormalizedActivations(), + ] + + passed = failed = 0 + for suite in suites: + suite_name = type(suite).__name__ + print(f"\n>> {suite_name}") + for name in [m for m in dir(suite) if m.startswith("test_")]: + method = getattr(suite, name) + try: + method() + print(f" PASS {name}") + passed += 1 + except (AssertionError, pytest.skip.Exception, Exception) as e: + print(f" FAIL {name}") + print(f" {e}") + failed += 1 + + print("\n" + "=" * 60) + print(f" Results: {passed} passed, {failed} failed") + print("=" * 60) + sys.exit(0 if failed == 0 else 1)