fix circular import

src/nanotron/parallel/context.py

@@ -5,7 +5,6 @@
 import torch
 
 import nanotron.distributed as dist
-from nanotron.utils import find_free_port
 
 DistributedBackend = Literal["gloo", "mpi", "nccl"]
 
@@ -50,6 +49,8 @@ def __init__(
         assert backend == "nccl", "Only nccl backend is supported for now."
 
         if not dist.is_initialized():
+            from nanotron.utils import find_free_port
+
             port = find_free_port() if port is None else port
             dist.initialize_torch_distributed(port)
 

tests/helpers/utils.py

@@ -8,7 +8,6 @@
 
 import torch.cuda
 from nanotron.parallel import ParallelContext
-from nanotron.utils import find_free_port
 from packaging import version
 from torch.distributed.launcher import elastic_launch
 
@@ -80,6 +79,8 @@ def __call__(self):
             # NOTE: we use a different random RNG, so that each unit tests don't generate the same port
             seed = random.randint(0, 9999)
             with torch.random.fork_rng(devices=["cuda"], seed=seed):
+                from nanotron.utils import find_free_port
+
                 port = find_free_port()
                 parallel_context = ParallelContext(
                     data_parallel_size=self.dp, pipeline_parallel_size=self.pp, tensor_parallel_size=self.tp, port=port