fix(util): move is_primary_rank to distributed utils to avoid circular import

Author: dest1n1s

src/lm_saes/abstract_sae.py

@@ -35,13 +35,10 @@
 from lm_saes.backend.language_model import LanguageModelConfig
 from lm_saes.config import BaseModelConfig
 from lm_saes.utils.auto import PretrainedSAEType, auto_infer_pretrained_sae_type
-from lm_saes.utils.distributed import DimMap, distributed_topk, item, mesh_dim_size
+from lm_saes.utils.distributed import DimMap, distributed_topk, is_primary_rank, item, mesh_dim_size
 from lm_saes.utils.distributed.utils import execute_and_broadcast
 from lm_saes.utils.logging import get_distributed_logger
 from lm_saes.utils.math import topk
-from lm_saes.utils.misc import (
-    is_primary_rank,
-)
 from lm_saes.utils.tensor_specs import TensorSpecs, apply_token_mask
 from lm_saes.utils.timer import timer
 

src/lm_saes/analysis/feature_analyzer.py

@@ -16,9 +16,8 @@
 from lm_saes.crosscoder import CrossCoder
 from lm_saes.sparse_dictionary import SparseDictionary
 from lm_saes.utils.discrete import KeyedDiscreteMapper
-from lm_saes.utils.distributed import DimMap, masked_fill, to_local
+from lm_saes.utils.distributed import DimMap, is_primary_rank, masked_fill, to_local
 from lm_saes.utils.distributed.ops import item
-from lm_saes.utils.misc import is_primary_rank
 from lm_saes.utils.tensor_dict import concat_dict_of_tensor, sort_dict_of_tensor
 
 from .post_analysis import PostAnalysisProcessor, get_post_analysis_processor

src/lm_saes/analysis/post_analysis/lorsa.py

@@ -34,9 +34,9 @@
 from lm_saes.lorsa import LowRankSparseAttention
 from lm_saes.sparse_dictionary import SparseDictionary
 from lm_saes.utils.discrete import KeyedDiscreteMapper
+from lm_saes.utils.distributed import is_primary_rank
 from lm_saes.utils.distributed.ops import item
 from lm_saes.utils.logging import get_distributed_logger
-from lm_saes.utils.misc import is_primary_rank
 
 from .base import PostAnalysisProcessor, register_post_analysis_processor
 

src/lm_saes/runners/topk_to_jumprelu_conversion.py

@@ -14,8 +14,8 @@
 from lm_saes.database import MongoClient, MongoDBConfig
 from lm_saes.resource_loaders import load_dataset, load_model
 from lm_saes.sparse_dictionary import SparseDictionary
+from lm_saes.utils.distributed import is_primary_rank
 from lm_saes.utils.logging import get_distributed_logger, setup_logging
-from lm_saes.utils.misc import is_primary_rank
 from lm_saes.utils.topk_to_jumprelu_conversion import topk_to_jumprelu_conversion
 
 from .utils import PretrainedSAE, load_config

src/lm_saes/runners/train.py

@@ -22,9 +22,8 @@
 from lm_saes.resource_loaders import load_dataset, load_model
 from lm_saes.sparse_dictionary import SparseDictionary, SparseDictionaryConfig
 from lm_saes.trainer import Trainer, TrainerConfig, WandbConfig
-from lm_saes.utils.distributed import mesh_rank
+from lm_saes.utils.distributed import is_primary_rank, mesh_rank
 from lm_saes.utils.logging import get_distributed_logger, setup_logging
-from lm_saes.utils.misc import is_primary_rank
 
 from .utils import PretrainedSAE, load_config
 

src/lm_saes/sparse_dictionary.py

@@ -35,13 +35,10 @@
 from lm_saes.backend.language_model import LanguageModelConfig
 from lm_saes.config import BaseModelConfig
 from lm_saes.utils.auto import PretrainedSAEType, auto_infer_pretrained_sae_type
-from lm_saes.utils.distributed import DimMap, distributed_topk, item, mesh_dim_size
+from lm_saes.utils.distributed import DimMap, distributed_topk, is_primary_rank, item, mesh_dim_size
 from lm_saes.utils.distributed.utils import execute_and_broadcast
 from lm_saes.utils.logging import get_distributed_logger
 from lm_saes.utils.math import topk
-from lm_saes.utils.misc import (
-    is_primary_rank,
-)
 from lm_saes.utils.tensor_specs import TensorSpecs, apply_token_mask
 from lm_saes.utils.timer import timer
 

src/lm_saes/trainer.py

@@ -35,12 +35,12 @@
 )
 from lm_saes.optim import SparseAdam, clip_grad_norm, get_scheduler
 from lm_saes.sparse_dictionary import SparseDictionary
+from lm_saes.utils.distributed import is_primary_rank
 from lm_saes.utils.distributed.ops import item
 from lm_saes.utils.logging import get_distributed_logger, log_metrics
 from lm_saes.utils.misc import (
     convert_str_to_torch_dtype,
     convert_torch_dtype_to_str,
-    is_primary_rank,
 )
 from lm_saes.utils.tensor_specs import apply_token_mask
 from lm_saes.utils.timer import timer

src/lm_saes/utils/distributed/__init__.py

@@ -4,6 +4,7 @@
     all_gather_dict,
     all_gather_list,
     get_process_group,
+    is_primary_rank,
     mesh_dim_rank,
     mesh_dim_size,
     mesh_rank,
@@ -12,6 +13,7 @@
 
 __all__ = [
     "DimMap",
+    "is_primary_rank",
     "distributed_topk",
     "item",
     "masked_fill",

src/lm_saes/utils/distributed/utils.py

@@ -7,7 +7,17 @@
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.tensor import Placement
 
-from lm_saes.utils.misc import is_primary_rank
+
+def is_primary_rank(device_mesh: DeviceMesh | None, dim_name: str = "sweep") -> bool:
+    """Check if the current rank is the primary rank for the given mesh dimension."""
+    if device_mesh is None:
+        return True
+    coord = device_mesh.get_coordinate()
+    mesh_dim_names = device_mesh.mesh_dim_names
+    if coord is None or mesh_dim_names is None:
+        return False
+    coord = [c for i, c in enumerate(coord) if dim_name not in mesh_dim_names or i != mesh_dim_names.index(dim_name)]
+    return all(c == 0 for c in coord)
 
 
 def all_gather_dict(

src/lm_saes/utils/misc.py

@@ -18,17 +18,6 @@ def is_master() -> bool:
     return not dist.is_initialized() or dist.get_rank() == 0
 
 
-def is_primary_rank(device_mesh: DeviceMesh | None, dim_name: str = "sweep") -> bool:
-    if device_mesh is None:
-        return True
-    coord = device_mesh.get_coordinate()
-    mesh_dim_names = device_mesh.mesh_dim_names
-    if coord is None or mesh_dim_names is None:
-        return False
-    coord = [c for i, c in enumerate(coord) if dim_name not in mesh_dim_names or i != mesh_dim_names.index(dim_name)]
-    return all(c == 0 for c in coord)
-
-
 def print_once(
     *values: object,
     sep: str | None = " ",