[logging][ez] Add timer logging for pickling and unpickle for object based collective (#139757)

fduwjj · pytorchmergebot · commit 41e4d88584c4 · 2024-11-05T17:40:27.000Z
Summary: As discussed, we want to measure the time spent during pickling and unpickle. Test Plan: CI Reviewed By: wz337 Differential Revision: D65462767 Pull Request resolved: #139757 Approved by: https://github.com/awgu, https://github.com/Skylion007, https://github.com/fegin, https://github.com/c-p-i-o
diff --git a/torch/distributed/distributed_c10d.py b/torch/distributed/distributed_c10d.py
@@ -2812,6 +2812,7 @@ def reduce(tensor, dst, op=ReduceOp.SUM, group=None, async_op=False):
         work.wait()
 
 
+@_time_logger
 def _object_to_tensor(obj, device, group):
     f = io.BytesIO()
     _pickler(f).dump(obj)
@@ -2831,6 +2832,7 @@ def _object_to_tensor(obj, device, group):
     return byte_tensor, local_size
 
 
+@_time_logger
 def _tensor_to_object(tensor, tensor_size, group):
     if get_debug_level() == DebugLevel.DETAIL and is_nccl_available():
         backend = get_backend(group)
