Repository: microsoft/DeepSpeed
File: deepspeed/runtime/engine.py (Python)

Code before:
            for module in self.module.modules():
                module._parameters._in_forward = False

        self._stop_timers(self.engine_timers.forward_timers)

        flops_profiler_active = (self.flops_profiler_enabled()
                                 and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0)

        if flops_profiler_active:
            self.flops_profiler.stop_profile()

        if not self.autotuning_profile_model_info():
            see_memory_usage("Engine after forward", force=self.memory_breakdown())

    @instrument_w_nvtx
    def forward(self, *inputs, **kwargs):
        r"""Execute forward propagation
        Arguments:
            *inputs: Variable length input list
            **kwargs: variable length keyword arguments
        """
        if self.autotuning_profile_model_info():
            ma = get_ma_status()

        if (self.is_deepcompile_enabled() and not self.is_deepcompile_active() and not self.is_compiled
                and not self._deepcompile_noop_warning_emitted):
            log_dist(
                "DeepCompile is enabled but engine.compile() has not been called; executing without DeepCompile until compile() runs.",
                ranks=[0])
            self._deepcompile_noop_warning_emitted = True

        if self.is_deepcompile_active() and hasattr(self, "launch_compile_passes"):
            # We can't have this in forward prologue as the compiler compiles hooks including the forward prologue.
            self.launch_compile_passes(self.global_steps)

        with autocast_if_enabled(self):
            loss = self.module(*inputs, **kwargs)

        if self.autotuning_profile_model_info():
            activation_mem = get_ma_status() - ma
            self.autotuning_model_info["activation_mem_per_gpu"] = activation_mem
            print_json_dist(self.autotuning_model_info, [0], path=self.autotuning_model_info_path())
            exit()

        return loss

    def _cast_inputs_half(self, inputs):
        if isinstance(inputs, (list, tuple)):
            new_inputs = []
            for v in inputs:
                new_inputs.append(self._cast_inputs_half(v))

Review comment (suggestion):
There is a `warning_once` API from the logging module which allows us to get rid of that "warning_emiited" member. But that interface does not support the "print on rank 0" semantics yet.

I have an issue (#7578) related to logging APIs as well. Do we have any plan to revisit and unify the logging APIs? If not, I can take a look and submit an RFC PR. Of course that's a separate topic and does not block this one.

Code after:
                module._parameters._in_forward = False

        self._stop_timers(self.engine_timers.forward_timers)

        flops_profiler_active = (self.flops_profiler_enabled()
                                 and self.global_steps == self.flops_profiler_profile_step() and self.global_rank == 0)

        if flops_profiler_active:
            self.flops_profiler.stop_profile()

        if not self.autotuning_profile_model_info():
            see_memory_usage("Engine after forward", force=self.memory_breakdown())

    @instrument_w_nvtx
    def forward(self, *inputs, **kwargs):
        r"""Execute forward propagation
        Arguments:
            *inputs: Variable length input list
            **kwargs: variable length keyword arguments
        """
        if self.autotuning_profile_model_info():
            ma = get_ma_status()

        if self.is_deepcompile_enabled() and not self.is_deepcompile_active() and not self.is_compiled:
            log_dist_once(
                "DeepCompile is enabled but engine.compile() has not been called; executing without DeepCompile until compile() runs.",
                ranks=[0])

        if self.is_deepcompile_active() and hasattr(self, "launch_compile_passes"):
            # We can't have this in forward prologue as the compiler compiles hooks including the forward prologue.
            self.launch_compile_passes(self.global_steps)

        with autocast_if_enabled(self):
            loss = self.module(*inputs, **kwargs)

        if self.autotuning_profile_model_info():
            activation_mem = get_ma_status() - ma
            self.autotuning_model_info["activation_mem_per_gpu"] = activation_mem
            print_json_dist(self.autotuning_model_info, [0], path=self.autotuning_model_info_path())
            exit()

        return loss

    def _cast_inputs_half(self, inputs):
        if isinstance(inputs, (list, tuple)):
            new_inputs = []
            for v in inputs:
                new_inputs.append(self._cast_inputs_half(v))
            return inputs.__class__(new_inputs)
        elif isinstance(inputs, dict):
            new_inputs = {}