diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index e963ef643677..7661a791c94e 100755 --- a/deepspeed/runtime/zero/stage3.py +++ b/deepspeed/runtime/zero/stage3.py @@ -102,8 +102,8 @@ def _apply_to_tensors_only(module, functional, backward_function, outputs): elif type(outputs) is torch.Tensor: return functional.apply(module, backward_function, outputs) else: - if not is_builtin_type(outputs): - logger.warning( + if not is_builtin_type(outputs) and dist.get_rank() == 0: + logger.debug( f"A module has unknown inputs or outputs type ({type(outputs)}) and the tensors embedded in it cannot be detected. " "The ZeRO-3 hooks designed to trigger before or after backward pass of the module relies on knowing the input and " "output tensors and therefore may not get triggered properly.")