[BE] fix circular import in torch/distributed/utils.py (#136286)

XilunWu · pytorchmergebot · commit 3406ac24d935 · 2024-09-22T20:54:12.000Z
**Summary** Fix circular import in `torch/distributed/utils.py` found when running internal test, see D62901023. Curious why this wasn't causing any issue. Is this relevant code deprecated and no longer used? Pull Request resolved: #136286 Approved by: https://github.com/Skylion007
diff --git a/torch/distributed/utils.py b/torch/distributed/utils.py
@@ -18,8 +18,6 @@
 import torch
 import torch.distributed as dist
 from torch import nn
-from torch.nn.parallel._functions import _get_stream
-from torch.nn.parallel.scatter_gather import _is_namedtuple
 from torch.nn.utils.rnn import PackedSequence
 
 
@@ -123,6 +121,9 @@ def to_map(obj):
                 device_mod = getattr(torch, device.type, None)
                 if device.type == "cpu" or device_mod is None:
                     return (obj.to(target_device),)
+
+                from torch.nn.parallel._functions import _get_stream
+
                 # Perform CPU -> target_device copies in a background stream. This code is
                 # motivated from similar logic in torch/nn/parallel/_functions.py
                 stream = _get_stream(target_device)
@@ -141,6 +142,9 @@ def to_map(obj):
                         assert isinstance(output, torch.Tensor)
                         output.record_stream(current_stream)  # type: ignore[arg-type]
                 return (output,)
+
+        from torch.nn.parallel.scatter_gather import _is_namedtuple
+
         if _is_namedtuple(obj):
             return [type(obj)(*args) for args in zip(*map(to_map, obj))]
         if isinstance(obj, tuple) and len(obj) > 0:
@@ -228,6 +232,8 @@ def _apply_to_tensors(fn, container):
     """Recursively apply to all tensor in different kinds of container types."""
 
     def apply(x):
+        from torch.nn.parallel.scatter_gather import _is_namedtuple
+
         if isinstance(x, torch.Tensor):
             return fn(x)
         elif hasattr(x, "__dataclass_fields__"):
