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

Code before:
                self.warn_unscaled_loss = False

        return scaled_loss

    @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._config.compile_config.enabled and not self._is_compiled:
            if self._compiler_backend is None:
                self._compiler_backend = get_backend_fn(self._config.compile_config.backend)

            if self._compiler_fn is None:
                compiled_model = torch.compile(self.module,
                                               backend=self._compiler_backend,
                                               **self._config.compile_config.kwargs)
            else:
                compiled_model = self._compiler_fn(self.module)

            self._set_client_model(compiled_model)
            self._is_compiled = True

        if self.autotuning_profile_model_info():
            ma = get_ma_status()
        else:
            see_memory_usage("Engine before forward", force=self.memory_breakdown())

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

        # used to check quantization happens at step 0!
        if self.global_steps == 0 and hasattr(self, "compression_scheduler"):
            self.compression_scheduler.step(step_zero_check=True)
            if self.quantizer:
                tensor_to_quantize = self.optimizer.bit16_groups if self.zero_optimization_stage(
                ) == 2 else self.optimizer.fp16_groups
                if self.compression_scheduler.weight_quantization_enabled:
                    self.quantizer.quantize(
                        tensor_to_quantize,
                        (self.optimizer.overflow if self.fp16_enabled() else False),
                        self.eigenvalue_enabled(),
                        None,
                    )

        if flops_profiler_active:
            self.flops_profiler.start_profile(ignore_list=None)


Review comment (bug):
```suggestion
        if self._config.compile_config.enabled not self._is_compiled and is_compile_supported():
            backend = self._compiler_fn if self._compiler_fn is not None else self._compiler_backend
	    self.module.compile(backend = backend, **self._compile_kwargs)
	    self._is_compiled = True
```
Reasons for the above:

1. Allows passing user defined kwargs even when the user provides custom compiler_fn
2. Type of self.module does not changes since torch.nn.Module.compile is called instead of torch.compile. The module is compiled in-place.
I used to fail on [this line](https://github.com/microsoft/DeepSpeed/blob/2fc702ed9f84f3c88179cb713bb9f4effbbdba33/deepspeed/runtime/pipe/engine.py#L70). It should pass now as we only compile on forward and not in __init__, nevertheless we should consider keeping the type of self.module.



Code after:

    @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()
        else:
            see_memory_usage("Engine before forward", force=self.memory_breakdown())

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

        # used to check quantization happens at step 0!
        if self.global_steps == 0 and hasattr(self, "compression_scheduler"):
            self.compression_scheduler.step(step_zero_check=True)
            if self.quantizer:
                tensor_to_quantize = self.optimizer.bit16_groups if self.zero_optimization_stage(
                ) == 2 else self.optimizer.fp16_groups
                if self.compression_scheduler.weight_quantization_enabled:
                    self.quantizer.quantize(
                        tensor_to_quantize,
                        (self.optimizer.overflow if self.fp16_enabled() else False),
                        self.eigenvalue_enabled(),
                        None,
                    )

        if flops_profiler_active:
            self.flops_profiler.start_profile(ignore_list=None)

        if self.module.training:
            if self.progressive_layer_drop:
                kwargs.update(self.progressive_layer_drop.get_state())

        if self.__class__.__name__ != "PipelineEngine":
            # TODO: The above if condition is a HACK since for PipelineEngine
            # it's difficult to inject argument in forward pass.
            if self.module.training and self.curriculum_enabled_legacy():
                self.curriculum_scheduler_legacy.update_difficulty(self.global_steps + 1)
                if self.curriculum_params_legacy()["curriculum_type"] == "seqlen":
                    kwargs.update({"curriculum_seqlen": self.curriculum_scheduler_legacy.get_current_difficulty()})

        if self.module.training and self.random_ltd_enabled():
            self.random_ltd_scheduler.update_seq(self.global_steps)

        if self.zero_optimization_partition_weights():
            # Enable automated discovery of external parameters by indicating that