Fix comms logger KeyError when log_name is omitted

deepspeed/comm/comm.py

@@ -23,6 +23,7 @@
 
 import torch
 from torch.distributed import GradBucket  # noqa: F401
+import inspect
 import os
 from typing import Any, Optional, TYPE_CHECKING
 
@@ -104,16 +105,23 @@ def configure(
 
 # Logging wrapper for timing ops
 def timed_op(func):
+    default_log_name = get_default_args(func).get('log_name', func.__name__)
+    # Cache the signature to avoid inspecting it on every communication call.
+    func_signature = inspect.signature(func)
 
     def log_wrapper(*args, **kwargs):
+        should_profile = False
         # Add enabled flag so that overhead to each comm op is two if conditions at most
         if comms_logger.enabled:
-            if ('prof' in kwargs
-                    and kwargs['prof']) or comms_logger.prof_all or ('log_name' in kwargs
-                                                                     and kwargs['log_name'] in comms_logger.prof_ops):
-                # Need func args for their defaults
-                func_args = get_default_args(func)
-                func_args.update(kwargs)
+            bound_args = func_signature.bind_partial(*args, **kwargs)
+            bound_args.apply_defaults()
+            func_args = bound_args.arguments
+            selected_log_name = func_args.get('log_name', default_log_name)
+            should_profile = (func_args.get('prof', False) or comms_logger.prof_all
+                              or selected_log_name in comms_logger.prof_ops)
+            if should_profile:
+                # Ops that do not declare a log_name are logged under their own name
+                func_args['log_name'] = selected_log_name
                 msg_size = get_msg_size_from_args(func, *args, **kwargs)
                 log_name = get_debug_log_name(func_args, comms_logger.debug)
                 timers(log_name).start()
@@ -127,8 +135,7 @@ def log_wrapper(*args, **kwargs):
                 # If we're using MPI, we can't simply sync the stream
                 if cdb.using_mpi:
                     cdb.barrier()
-                if ('prof' in kwargs and kwargs['prof']) or comms_logger.prof_all or (
-                        'log_name' in kwargs and kwargs['log_name'] in comms_logger.prof_ops):
+                if should_profile:
                     log_name = get_debug_log_name(func_args, comms_logger.debug)
                     raw_name = func.__name__
                     timers(log_name).stop()
@@ -230,7 +237,13 @@ def broadcast(tensor, src, group=None, async_op=False, prof=False, log_name='bro
 
 
 @timed_op
-def broadcast_object_list(object_list, src, group=None, device=None):
+def broadcast_object_list(object_list,
+                          src,
+                          group=None,
+                          device=None,
+                          prof=False,
+                          log_name='broadcast_object_list',
+                          debug=get_caller_func()):
     global cdb
     return cdb.broadcast_object_list(object_list=object_list, src=src, group=group, device=device)
 
@@ -364,7 +377,13 @@ def all_to_all_single(output,
 
 
 @timed_op
-def all_to_all(output_tensor_list, input_tensor_list, group=None, async_op=False):
+def all_to_all(output_tensor_list,
+               input_tensor_list,
+               group=None,
+               async_op=False,
+               prof=False,
+               log_name='all_to_all',
+               debug=get_caller_func()):
     global cdb
     return cdb.all_to_all(output_tensor_list, input_tensor_list, group=group, async_op=async_op)
 

deepspeed/utils/comms_logging.py

@@ -37,7 +37,7 @@ def calc_bw_log(comm_op, size, duration):
     n = dist.get_world_size()
     tput = 0
     busbw = 0
-    if comm_op == "all_to_all_single":
+    if comm_op == "all_to_all_single" or comm_op == "all_to_all":
         tput = (size / duration)
         busbw = (size / duration) * ((n - 1) / n)
     elif comm_op == "all_gather" or comm_op == "all_gather_into_tensor" or comm_op == "reduce_scatter" or comm_op == "reduce_scatter_tensor":
@@ -47,7 +47,7 @@ def calc_bw_log(comm_op, size, duration):
     elif comm_op == "all_reduce" or comm_op == "all_reduce_coalesced" or comm_op == "inference_all_reduce":
         tput = (size * 2 / duration)
         busbw = (size / duration) * (2 * (n - 1) / n)
-    elif comm_op == "send" or comm_op == "recv" or comm_op == "isend" or comm_op == "irecv" or comm_op == "broadcast" or comm_op == "reduce" or comm_op == "gather" or comm_op == "scatter" or comm_op == "barrier":
+    elif comm_op == "send" or comm_op == "recv" or comm_op == "isend" or comm_op == "irecv" or comm_op == "broadcast" or comm_op == "broadcast_object_list" or comm_op == "reduce" or comm_op == "gather" or comm_op == "scatter" or comm_op == "barrier":
         tput = (size / duration)
         busbw = tput
     else:

tests/unit/comm/test_comms_logger.py

@@ -3,6 +3,12 @@
 
 # DeepSpeed Team
 
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
 from deepspeed.utils.comms_logging import CommsLogger
 
 
@@ -49,3 +55,146 @@ def test_trim_mean_does_not_mutate_its_argument():
     data = [3.0, 1.0, 2.0]
     assert trim_mean(data, 0.1) == 2.0
     assert data == [3.0, 1.0, 2.0]
+
+
+def test_timed_op_falls_back_to_the_op_name_when_log_name_is_missing(monkeypatch):
+    # timed_op looks up func_args['log_name'], so an op whose signature does not
+    # declare log_name used to raise KeyError as soon as profiling was turned on.
+    # Such an op must still be logged, under its own name.
+    from deepspeed.comm import comm
+
+    monkeypatch.setattr(comm, 'comms_logger', CommsLogger())
+    monkeypatch.setattr(
+        comm, 'cdb', SimpleNamespace(using_mpi=False, is_initialized=lambda: True,
+                                     get_world_size=lambda group=None: 1))
+    monkeypatch.setattr(comm, 'get_accelerator', lambda: SimpleNamespace(synchronize=lambda: None))
+
+    @comm.timed_op
+    def barrier():
+        return 'done'
+
+    comm.comms_logger.enabled = True
+    comm.comms_logger.start_profiling_comms()
+
+    assert barrier() == 'done'
+    assert 'barrier' in comm.comms_logger.comms_dict
+
+
+def test_timed_op_profiles_default_log_name_with_prof_ops(monkeypatch):
+    from deepspeed.comm import comm
+
+    monkeypatch.setattr(comm, 'comms_logger', CommsLogger())
+    monkeypatch.setattr(
+        comm, 'cdb', SimpleNamespace(using_mpi=False, is_initialized=lambda: True,
+                                     get_world_size=lambda group=None: 1))
+    monkeypatch.setattr(comm, 'get_accelerator', lambda: SimpleNamespace(synchronize=lambda: None))
+
+    @comm.timed_op
+    def barrier(log_name='barrier'):
+        return 'done'
+
+    comm.comms_logger.enabled = True
+    comm.comms_logger.prof_ops = ['barrier']
+
+    assert barrier() == 'done'
+    assert 'barrier' in comm.comms_logger.comms_dict
+
+
+@pytest.mark.parametrize('op_name', ['broadcast_object_list', 'all_to_all'])
+@pytest.mark.parametrize('positional', [False, True])
+@pytest.mark.parametrize('debug', [False, True])
+@pytest.mark.parametrize('profile_mode', ['prof_all', 'prof', 'prof_ops_default', 'prof_ops_custom', 'unselected'])
+def test_collective_profiling(monkeypatch, op_name, positional, debug, profile_mode):
+    from deepspeed.comm import comm
+
+    backend_op = Mock(return_value='done')
+    backend = SimpleNamespace(using_mpi=False, is_initialized=lambda: True, get_world_size=lambda group=None: 2)
+    setattr(backend, op_name, backend_op)
+    monkeypatch.setattr(comm, 'cdb', backend)
+    monkeypatch.setattr(comm, 'get_accelerator', lambda: SimpleNamespace(synchronize=lambda: None))
+    op_timer = Mock()
+    op_timer.elapsed.return_value = 1.0
+    timers = Mock(return_value=op_timer)
+    monkeypatch.setattr(comm, 'timers', timers)
+
+    monkeypatch.setattr(comm, 'comms_logger', CommsLogger())
+    comm.comms_logger.enabled = True
+    comm.comms_logger.debug = debug
+    comm.comms_logger.prof_all = profile_mode == 'prof_all'
+    log_name = 'custom_collective' if profile_mode in ('prof', 'prof_ops_custom') else op_name
+    comm.comms_logger.prof_ops = [log_name] if profile_mode.startswith('prof_ops') else []
+
+    if op_name == 'broadcast_object_list':
+        object_list = [{'value': 1}]
+        args = (object_list, 0, None, None)
+        expected_size = 0
+    else:
+        input_list = [torch.ones(4), torch.ones(4)]
+        output_list = [torch.empty_like(tensor) for tensor in input_list]
+        args = (output_list, input_list, None, False)
+        expected_size = sum(tensor.numel() * tensor.element_size() for tensor in input_list)
+
+    kwargs = {}
+    if profile_mode in ('prof', 'prof_ops_custom'):
+        prof = profile_mode == 'prof'
+        if positional:
+            args += (prof, log_name)
+        else:
+            kwargs = {'prof': prof, 'log_name': log_name}
+
+    assert getattr(comm, op_name)(*args, **kwargs) == 'done'
+    if op_name == 'broadcast_object_list':
+        backend_op.assert_called_once_with(object_list=object_list, src=0, group=None, device=None)
+    else:
+        backend_op.assert_called_once_with(output_list, input_list, group=None, async_op=False)
+
+    if profile_mode == 'unselected':
+        assert comm.comms_logger.comms_dict == {}
+        timers.assert_not_called()
+    else:
+        record_name, = comm.comms_logger.comms_dict
+        if debug:
+            assert record_name.startswith(log_name + ' | [Caller Func: ')
+        else:
+            assert record_name == log_name
+        record = comm.comms_logger.comms_dict[record_name][expected_size]
+        assert record[0] == 1
+        assert record[1] == [1.0]
+        op_timer.start.assert_called_once_with()
+        op_timer.stop.assert_called_once_with()
+        op_timer.elapsed.assert_called_once_with(reset=False)
+        assert all(call.args == (record_name, ) for call in timers.call_args_list)
+
+
+def test_timed_op_disabled_does_not_access_profiling_state(monkeypatch):
+    from deepspeed.comm import comm
+
+    @comm.timed_op
+    def barrier(prof=False, log_name='barrier'):
+        return 'done'
+
+    # No other logger attributes or backend should be needed on the disabled path.
+    monkeypatch.setattr(comm, 'comms_logger', SimpleNamespace(enabled=False))
+    monkeypatch.setattr(comm, 'cdb', None)
+    timers = Mock()
+    monkeypatch.setattr(comm, 'timers', timers)
+
+    assert barrier(True, 'custom_barrier') == 'done'
+    timers.assert_not_called()
+
+
+@pytest.mark.parametrize('world_size', [1, 2, 4])
+@pytest.mark.parametrize('comm_op', ['broadcast_object_list', 'all_to_all'])
+def test_calc_bw_log_supports_object_and_list_collectives(monkeypatch, comm_op, world_size):
+    import deepspeed.comm as dist
+    from deepspeed.utils.comms_logging import calc_bw_log
+
+    monkeypatch.setattr(dist, 'get_world_size', lambda group=None: world_size)
+
+    tput, busbw = calc_bw_log(comm_op, 1024, 2.0)
+    expected_tput = 1024 / 2.0 * 8 / 1e6
+    expected_busbw = expected_tput
+    if comm_op == 'all_to_all':
+        expected_busbw *= (world_size - 1) / world_size
+    assert tput == pytest.approx(expected_tput)
+    assert busbw == pytest.approx(expected_busbw)