Add `self.group`

deep_ep/buffer.py

@@ -54,13 +54,14 @@ class Buffer:
             explicitly_destroy: If this flag is set to True, you need to explicitly call `destroy()` to release resources;
                 otherwise, the resources will be released by the destructor.
                 Note: Releasing resources in the destructor may cause Python's exception handling process to hang.
-            comm: the mpi4py.MPI.Comm communicator to use in case the group parameter is absent.
+            comm: the `mpi4py.MPI.Comm` communicator to use in case the group parameter is absent.
         """
         check_nvlink_connections(group)
 
         # Initialize the CPP runtime
         if group is not None:
             self.rank = group.rank()
+            self.group = group
             self.group_size = group.size()
 
             def all_gather_object(obj):
@@ -69,6 +70,7 @@ class Buffer:
                 return object_list
         elif comm is not None:
             self.rank = comm.Get_rank()
+            self.group = comm
             self.group_size = comm.Get_size()
 
             def all_gather_object(obj):