From bacf519d538d117945d0ede4499246089e93f516 Mon Sep 17 00:00:00 2001 From: guptaishaan Date: Sun, 25 Oct 2026 17:16:15 -0700 Subject: [PATCH] Skip accelerator sync for unprofiled comm ops timed_op synchronized the accelerator after every wrapped comm op whenever comms logging was enabled, even for ops that were not selected by prof_ops, prof_all or prof. Those ops never get a timer or a record, so the sync only added overhead. Only synchronize when should_profile is set. The MPI barrier is unchanged. Count synchronize() calls in test_collective_profiling for both profiled and unselected ops. Signed-off-by: guptaishaan --- deepspeed/comm/comm.py | 3 ++- tests/unit/comm/test_comms_logger.py | 5 ++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/deepspeed/comm/comm.py b/deepspeed/comm/comm.py index 5a4a6842c521..84016a7408bc 100755 --- a/deepspeed/comm/comm.py +++ b/deepspeed/comm/comm.py @@ -130,7 +130,8 @@ def log_wrapper(*args, **kwargs): finally: if comms_logger.enabled: # Need to make op blocking for accurate logging - get_accelerator().synchronize() + if should_profile: + get_accelerator().synchronize() # If we're using MPI, we can't simply sync the stream if cdb.using_mpi: cdb.barrier() diff --git a/tests/unit/comm/test_comms_logger.py b/tests/unit/comm/test_comms_logger.py index 2e85864e3022..d4228bdd1c2e 100644 --- a/tests/unit/comm/test_comms_logger.py +++ b/tests/unit/comm/test_comms_logger.py @@ -111,7 +111,8 @@ def test_collective_profiling(monkeypatch, op_name, positional, debug, profile_m 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)) + synchronize = Mock() + monkeypatch.setattr(comm, 'get_accelerator', lambda: SimpleNamespace(synchronize=synchronize)) op_timer = Mock() op_timer.elapsed.return_value = 1.0 timers = Mock(return_value=op_timer) @@ -151,7 +152,9 @@ def test_collective_profiling(monkeypatch, op_name, positional, debug, profile_m if profile_mode == 'unselected': assert comm.comms_logger.comms_dict == {} timers.assert_not_called() + synchronize.assert_not_called() else: + synchronize.assert_called_once_with() record_name, = comm.comms_logger.comms_dict if debug: assert record_name.startswith(log_name + ' | [Caller Func: ')