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.