Bind positional prof/log_name args and cover with regression tests

timed_op previously read prof/log_name only from kwargs, so callers
that passed them positionally (e.g. dist.all_to_all(out, inp, None,
False, True, 'name')) had those values silently ignored. Bind args
and kwargs against func's signature with inspect.signature(...).
bind_partial() instead, so positional values are picked up the same
as keyword ones. This makes the earlier keyword-only ('*') marker on
broadcast_object_list/all_to_all unnecessary, so it is reverted.

The signature binding stays inside the 'if comms_logger.enabled'
guard to keep the disabled fast path at its previous cost.

Add regression tests for reading prof/log_name from positional args
and for calc_bw_log supporting broadcast_object_list/all_to_all.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Jin, Youzhi <youzhi.jin@intel.com>

Author: jinyouzhi

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
 
@@ -110,8 +111,12 @@ 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:
-            selected_log_name = kwargs.get('log_name', default_log_name)
-            should_profile = (('prof' in kwargs and kwargs['prof']) or comms_logger.prof_all
+            # prof/log_name may be passed positionally, so bind args/kwargs to
+            # func's signature rather than only looking at kwargs
+            bound_args = inspect.signature(func).bind_partial(*args, **kwargs)
+            bound_args.apply_defaults()
+            selected_log_name = bound_args.arguments.get('log_name', default_log_name)
+            should_profile = (bound_args.arguments.get('prof', False) or comms_logger.prof_all
                               or selected_log_name in comms_logger.prof_ops)
             if should_profile:
                 # Need func args for their defaults
@@ -238,7 +243,6 @@ def broadcast_object_list(object_list,
                           src,
                           group=None,
                           device=None,
-                          *,
                           prof=False,
                           log_name='broadcast_object_list',
                           debug=get_caller_func()):
@@ -379,7 +383,6 @@ 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()):

tests/unit/comm/test_comms_logger.py

@@ -94,3 +94,40 @@ def barrier(log_name='barrier'):
 
     assert barrier() == 'done'
     assert 'barrier' in comm.comms_logger.comms_dict
+
+
+def test_timed_op_reads_prof_and_log_name_from_positional_args(monkeypatch):
+    # prof/log_name are ordinary parameters, so callers may pass them positionally
+    # (e.g. dist.all_to_all(out, inp, None, False, True, "my_all_to_all")). timed_op
+    # used to look them up only in kwargs, silently ignoring positional values.
+    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(prof=False, log_name='barrier'):
+        return 'done'
+
+    comm.comms_logger.enabled = True
+
+    assert barrier(True, 'custom_barrier') == 'done'
+    assert 'custom_barrier' in comm.comms_logger.comms_dict
+
+
+def test_calc_bw_log_supports_object_and_list_collectives(monkeypatch):
+    # broadcast_object_list and all_to_all gained profiling support but were
+    # missing from calc_bw_log, so profiling them hit the "wrong comm_op
+    # specified" fallback and exited the process instead of logging bandwidth.
+    import deepspeed.comm as dist
+    from deepspeed.utils.comms_logging import calc_bw_log
+
+    monkeypatch.setattr(dist, 'get_world_size', lambda group=None: 2)
+
+    for comm_op in ('broadcast_object_list', 'all_to_all'):
+        tput, busbw = calc_bw_log(comm_op, 1024, 1.0)
+        assert tput > 0
+        assert busbw > 0