Skip to content

Fix CLT normalization and distributed attribution #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JitUikey28 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Fix CLT normalization and distributed attribution #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
236 changes: 209 additions & 27 deletions src/clt_forge/attribution/attribution.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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 (
Expand All @@ -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__(
Expand All @@ -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()
Expand All @@ -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(),
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -118,31 +286,41 @@ 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,
edge_threshold=edge_threshold,
)

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,
Expand All @@ -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

Expand Down
Loading