diff --git a/deepspeed/utils/comms_logging.py b/deepspeed/utils/comms_logging.py index 4532ac34e7be..32ecff8bf9cc 100644 --- a/deepspeed/utils/comms_logging.py +++ b/deepspeed/utils/comms_logging.py @@ -91,7 +91,7 @@ def start_profiling_comms(self): self.prof_all = True def stop_profiling_comms(self): - self.prof_all = True + self.prof_all = False # E.g. start_profiling_op('all_reduce') def start_profiling_op(self, op_name_list): diff --git a/tests/unit/comm/test_comms_logger.py b/tests/unit/comm/test_comms_logger.py new file mode 100644 index 000000000000..ef4d77b66de3 --- /dev/null +++ b/tests/unit/comm/test_comms_logger.py @@ -0,0 +1,19 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from deepspeed.utils.comms_logging import CommsLogger + + +def test_stop_profiling_comms_disables_prof_all(): + # start_profiling_comms()/stop_profiling_comms() toggle the global comm + # profiling flag prof_all. stop_profiling_comms() must clear it; otherwise + # global comm profiling can never be turned off once it has been started. + comms_logger = CommsLogger() + + comms_logger.start_profiling_comms() + assert comms_logger.prof_all is True + + comms_logger.stop_profiling_comms() + assert comms_logger.prof_all is False