Fix comms logger KeyError when log_name is omitted

[jinyouzhi]

Summary

Fix a KeyError: 'log_name' raised by the DeepSpeed communication logger
when a wrapped collective is called without an explicit log_name.

"comms_logger": {
    "enabled": true,
    "prof_all": true,
    "debug": true
}

This is exposed by multi-rank AutoTP input consistency checks, which call
broadcast_object_list without passing profiling metadata. Single-rank TP
does not exercise this communication path.

Changes

• Add prof/log_name/debug for broadcast_object_list and all_to_all
• Use the func.__name__ as the default log_name to cover missing status
• Add a regression test

Validation

• python -m pytest -q tests/unit/comm/test_comms_logger.py

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[ebarkhordar]

prof_ops never matches an op that takes its log_name from the signature default. The gate at comm.py:111-113 tests 'log_name' in kwargs, so it fires only when a caller passes the name explicitly, while config-json.md documents "prof_ops": ["all_reduce", "all_gather"] against ordinary calls.

At 1f95164 in a clean container, CPU torch and a stub cdb:

prof_ops = ['all_reduce']  prof_all = False
A. dist.all_reduce(t)               -> comms_dict keys: []
B. dist.all_reduce(t, log_name=...) -> comms_dict keys: ['all_reduce']

Your setdefault is one step from covering this. Resolving the name once per op keeps the per-call fast path at two conditions:

def timed_op(func):
    default_log_name = get_default_args(func).get('log_name', func.__name__)

    def log_wrapper(*args, **kwargs):
        if comms_logger.enabled:
            selected = kwargs.get('log_name', default_log_name)
            if kwargs.get('prof') or comms_logger.prof_all or selected in comms_logger.prof_ops:

then func_args['log_name'] = selected in place of the setdefault, and the same condition in the finally gate. With that, A logs and tests/unit/comm/test_comms_logger.py is still 4 passed. It is a separate bug from the KeyError you are fixing, so it may belong in its own PR.

[jinyouzhi]

prof_ops never matches an op that takes its log_name from the signature default. The gate at comm.py:111-113 tests 'log_name' in kwargs, so it fires only when a caller passes the name explicitly, while config-json.md documents "prof_ops": ["all_reduce", "all_gather"] against ordinary calls.

At 1f95164 in a clean container, CPU torch and a stub cdb:

prof_ops = ['all_reduce']  prof_all = False
A. dist.all_reduce(t)               -> comms_dict keys: []
B. dist.all_reduce(t, log_name=...) -> comms_dict keys: ['all_reduce']

Your setdefault is one step from covering this. Resolving the name once per op keeps the per-call fast path at two conditions:

def timed_op(func):
    default_log_name = get_default_args(func).get('log_name', func.__name__)

    def log_wrapper(*args, **kwargs):
        if comms_logger.enabled:
            selected = kwargs.get('log_name', default_log_name)
            if kwargs.get('prof') or comms_logger.prof_all or selected in comms_logger.prof_ops:

then func_args['log_name'] = selected in place of the setdefault, and the same condition in the finally gate. With that, A logs and tests/unit/comm/test_comms_logger.py is still 4 passed. It is a separate bug from the KeyError you are fixing, so it may belong in its own PR.

Wow, that's a very insightful observation! I agree that prof_ops currently fails to match an operation when log_name comes from the function signature default rather than being explicitly passed in kwargs. Would you prefer that I fix this as part of this PR, or you will handle it separately in a follow-up PR?

[ebarkhordar]

Your call as the author, but I would take it in this PR. The repair replaces the setdefault line you just added, so as two PRs whichever one lands second has to rebase through the other for no benefit. That is the opposite of what I said above about a separate PR, and the overlap on that line is why I changed my mind.

Either way I am not going to open a competing PR for it.

One thing to keep if you do take it: the finally gate at comm.py:133 tests the same 'log_name' in kwargs expression, so it needs the same selected condition. Without it the start gate selects the call, the stop gate does not, and timers(log_name).stop() and comms_logger.append never run, so the op still does not appear.

I re-read comm.py at cbbf47b after your merge from master and both gates are unchanged there, so the differential above still stands at your current head.

[jinyouzhi]

Your call as the author, but I would take it in this PR. The repair replaces the setdefault line you just added, so as two PRs whichever one lands second has to rebase through the other for no benefit. That is the opposite of what I said above about a separate PR, and the overlap on that line is why I changed my mind.

Either way I am not going to open a competing PR for it.

One thing to keep if you do take it: the finally gate at comm.py:133 tests the same 'log_name' in kwargs expression, so it needs the same selected condition. Without it the start gate selects the call, the stop gate does not, and timers(log_name).stop() and comms_logger.append never run, so the op still does not appear.

I re-read comm.py at cbbf47b after your merge from master and both gates are unchanged there, so the differential above still stands at your current head.

Thank you for your patience and guidance. I’ve implemented your suggestions and pushed the changes. Could you please take a look when you have a chance? I’d really appreciate your feedback.

[FU-max-boop]

Independent re-review of exact head 90d37a2bf484749a95ab2674459a5b9f89c5242e against base 715965e027894a2e72ac2e27f2daed2c599e99f0: the two focused behaviors look correct, and tests/unit/comm/test_comms_logger.py is 5/5 green locally on CPU.

One hot-path performance issue remains in this revision. selected_log_name and should_profile are now computed before the comms_logger.enabled gate, so the default-disabled path pays the kwargs lookup, logger attribute reads, and prof_ops membership test on every decorated communication call. That also contradicts the nearby comment that the disabled overhead is at most the enabled check.

A same-host synthetic no-op wrapper benchmark, intended only to isolate Python dispatch overhead rather than claim end-to-end collective latency, produced:

base  715965e: median  88.39 ns/call
head  90d37a2: median 141.39 ns/call
                     +53.00 ns, about +60%

Each result is the median of 9 repeats × 5,000,000 calls under the same Python 3.12.13 / Torch 2.13 environment with enabled=False, prof_all=False, and prof_ops=[]. A non-empty prof_ops list makes the new disabled-path work grow further.

The narrow fix is to initialize should_profile = False, then resolve selected_log_name and the selection expression only inside if comms_logger.enabled:. That retains the default-log-name repair while restoring the disabled fast path. I would keep the existing broader synchronization semantics out of this PR; this finding is only about overhead introduced by the current diff.

[jinyouzhi]

Nice catch! @FU-max-boop Thank you very much, I pushed the fix.

[FU-max-boop]

Re-reviewed exact head 9a92148c4b8a0b807421403c2639c269b5512d30 against base 715965e027894a2e72ac2e27f2daed2c599e99f0.

The disabled-path performance blocker from my previous review is resolved: selected_log_name resolution and the prof_ops membership test now occur only under comms_logger.enabled, while the single should_profile decision is reused by the finally gate.

Validation:

• tests/unit/comm/test_comms_logger.py: 5/5 passed locally on CPU.
• DCO, test collection, and the full Modal CI job: 3/3 passed on this exact head.
• An interleaved same-process synthetic no-op benchmark of the exact timed_op source on Python 3.12.13 (31 ABBA/BAAB rounds, 2,000,000 calls per sample, two samples per variant per round) measured a base median of 128.375 ns/call and a head median of 135.398 ns/call, with a paired median delta of +6.694 ns/call. The host was noisy, so the absolute timing is directional; static disassembly shows that the residual disabled-path difference is only two bytecodes, LOAD_CONST False and STORE_FAST should_profile, rather than the unconditional kwargs lookup and logger membership work in the prior revision. That residual is non-blocking.

git diff --check is clean, and I found no remaining correctness, API, or performance blocker in this revision. No remaining blockers from me.

[tohtana]

Hi @jinyouzhi,
Thank you for the improvement! I think the overall direction of this PR looks good, but there seems to be a few small issues we should fix before merge. I left comments on them.

[tohtana]

This should be added to calc_bw_log.

[tohtana]

prof and log_name will be ignored if they are passed as positional args.

[jinyouzhi]

Fixed — added broadcast_object_list and all_to_all to calc_bw_log. Thanks!

[tohtana]

This should be added to calc_bw_log.

[tohtana]

prof and log_name will be ignored if they are passed as positional args.

[jinyouzhi]

Fixed — positional prof and log_name arguments are now handled correctly. Thanks!

[tohtana]

I noticed the above might be preexisting issues, but it would be great if we could address in this PR.

[jinyouzhi]

Thanks for catching these, @tohtana! Both issues are now addressed:

• Added  broadcast_object_list  and  all_to_all  to  calc_bw_log , using the existing bandwidth formulas for  broadcast  and  all_to_all_single , respectively.
• Updated  timed_op  to resolve positional and keyword arguments using a cached function signature, so positional  prof  and  log_name  are honored without making them keyword-only or changing the existing call signatures.

Argument binding remains inside the  comms_logger.enabled  guard, preserving the disabled-logging fast path.

[tohtana]

Thank you for the update, @jinyouzhi!
Looks good to me.

[jinyouzhi]

Failed with MPS CPUAdam JIT:

FAILED: [code=1] cpu_adam_impl.o 
c++ -MMD -MF cpu_adam_impl.o.d -DTORCH_EXTENSION_NAME=cpu_adam -DTORCH_API_INCLUDE_EXTENSION_H -I/Users/runner/work/DeepSpeed/DeepSpeed/csrc/includes -isystem /Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/torch/include -isystem /Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/torch/include/torch/csrc/api/include -isystem /Library/Frameworks/Python.framework/Versions/3.12/include/python3.12 -fPIC -std=c++20 -O3 -std=c++17 -g -Wno-reorder -D__SCALAR__ -Xpreprocessor -fopenmp -I/opt/homebrew/opt/libomp/include -UC10_USE_GLOG -c /Users/runner/work/DeepSpeed/DeepSpeed/csrc/adam/cpu_adam_impl.cpp -o cpu_adam_impl.o 
In file included from /Users/runner/work/DeepSpeed/DeepSpeed/csrc/adam/cpu_adam_impl.cpp:6:
In file included from /Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/torch/include/torch/extension.h:6:
/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/torch/include/torch/csrc/api/include/torch/all.h:5:2: error: C++20 or later compatible compiler is required to use PyTorch.
    5 | #error C++20 or later compatible compiler is required to use PyTorch.

@PKUWZP Should we modify the CPUAdam opbuilder cpp version for MPS.

Like #8466

DeepSpeed/op_builder/mps/cpu_adam.py

Lines 48 to 53 in 02b5350

	def cxx_args(self):
	args = ['-O3', '-std=c++17', '-g', '-Wno-reorder', '-D__SCALAR__']
	libomp = self._libomp_prefix()
	if libomp is not None:
	args += ['-Xpreprocessor', '-fopenmp', f'-I{libomp}/include']
	return args

:

cpp_standard = '-std=c++20' if (TORCH_MAJOR, TORCH_MINOR) >= (2, 12) else '-std=c++17'