process_group: fix docs with torch==2.6.0

Author: d4l3k

torchft/process_group.py

@@ -20,7 +20,7 @@
 import queue
 import threading
 from datetime import timedelta
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, Union
+from typing import Dict, List, Optional, Tuple, Type, TYPE_CHECKING, Union
 
 import torch
 import torch.distributed as dist
@@ -31,16 +31,16 @@
 from torch.distributed import (
     BroadcastOptions,
     DeviceMesh,
+    get_rank,
+    init_device_mesh,
     PrefixStore,
     ProcessGroup as BaseProcessGroup,
     ProcessGroupGloo as BaseProcessGroupGloo,
     ProcessGroupNCCL as BaseProcessGroupNCCL,
     Store,
     TCPStore,
-    get_rank,
-    init_device_mesh,
 )
-from torch.distributed.distributed_c10d import Work, _world
+from torch.distributed.distributed_c10d import _world, Work
 from torch.futures import Future
 
 if TYPE_CHECKING:
@@ -132,10 +132,20 @@ def allgather(
         input_tensor: List[torch.Tensor],
         opts: object,
     ) -> Work:
+        """
+        Gathers tensors from the whole group in a list.
+
+        See torch.distributed.all_gather for more details.
+        """
         raise NotImplementedError("not implemented")
 
     # pyre-fixme[14]: inconsistent override
     def broadcast(self, tensor_list: List[torch.Tensor], opts: object) -> Work:
+        """
+        Broadcasts the tensor to the whole group.
+
+        See torch.distributed.broadcast for more details.
+        """
         raise NotImplementedError("not implemented")
 
     def broadcast_one(self, tensor: torch.Tensor, root: int) -> Work: