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

Code before:
            if enable_deepcompile:
                # Restore default hooks if compilation fails before completing.
                self._set_deepcompile_active(False)
            raise

        self._is_compiled = True
        self._is_compiled_autograd_enabled = compiled_autograd_enabled
        self._compile_kwargs = compile_kwargs

    def _set_deepcompile_active(self, active: bool) -> None:
        """Toggle DeepCompile runtime state and manage forward hooks accordingly."""
        if self._deepcompile_active == active:
            return

        if active:
            if self.module_forward_pre_hook is not None:
                self.module_forward_pre_hook.remove()
                self.module_forward_pre_hook = None
            if self.module_forward_post_hook is not None:
                self.module_forward_post_hook.remove()
                self.module_forward_post_hook = None

            if self._is_compiled_autograd_enabled:
                logger.warning("Compiled autograd is not compatible with DeepCompile, disabling compiled autograd.")
                self._is_compiled_autograd_enabled = False
                self._compile_kwargs = None
        else:
            if self.module_forward_pre_hook is None:
                self.module_forward_pre_hook = self._create_module_forward_pre_hook()
            if self.module_forward_post_hook is None:
                self.module_forward_post_hook = self._create_module_forward_post_hook()

        self._deepcompile_active = active

    def get_compile_time(self):
        from deepspeed.compile.backend import opt_pass_times
        return opt_pass_times

    def register_compile_pass(self, pass_name: str, pass_fn: Callable) -> None:
        register_compile_pass(pass_name, pass_fn)

    def is_deepcompile_enabled(self) -> bool:
        return self._config.compile_config.deepcompile

    def is_deepcompile_active(self) -> bool:
        return self._deepcompile_active

    @property
    def is_compiled(self) -> bool:
        return self._is_compiled


Review comment (suggestion):
`self._set_deepcompile_active` is called before `self._is_compiled_autograd_enabled` is set at the end of `compile()`. Thus the assignments here will be overriden.

How about checking `self._deepcompile_active` when setting `self._is_compiled_autograd_enabled` in `compile()`? i.e.,

```
    if compiled_autograd_enabled:
        if not self._deepcompile_active:
            self._is_compiled_autograd_enabled = compiled_autograd_enabled
            self._compile_kwargs = compile_kwargs
        else:
            logger.warning("Compiled autograd is not compatible with DeepCompile, disabling compiled autograd.")
```

Code after:
        if self._optimizer_has_ckpt_event_epilogue():
            self.optimizer.checkpoint_event_epilogue()

        return state_dict

    def save_fp16_model(self, save_dir, save_filename="pytorch_model.bin"):
        """has been renamed to save_16bit_model, keeping this around for backwards
        compatibility"""
        return self.save_16bit_model(save_dir, save_filename)

    def save_16bit_model(self, save_dir, save_filename="pytorch_model.bin", exclude_frozen_parameters=False):
        """
        Save 16bit model weights

        This method saves the 16bit model weights at the desired destination.

        Arguments:
            save_dir: Required. Directory for saving the model
            save_filename: Optional. Filename to save to. Defaults to ``pytorch_model.bin``
            exclude_frozen_parameters: Optional. Exclude frozen parameters from checkpointed state.

        Returns:
            ``True`` when a model has been saved, ``False`` otherwise. It will not be saved if
            stage3_gather_16bit_weights_on_model_save is ``False``.

        Important: all processes must call this method and not just the process with rank 0. It is
        because the processes need to work in sync to gather the weights. This method will hang
        waiting to synchronize with other processes if it's called just for the process with rank 0.

        """

        path = os.path.join(save_dir, save_filename)

        if self.zero_optimization_partition_weights():
            if self.zero_gather_16bit_weights_on_model_save():
                # consolidation is expensive in time and memory and therefore isn't a default
                state_dict = self._zero3_consolidated_16bit_state_dict(
                    exclude_frozen_parameters=exclude_frozen_parameters)
            else:
                # the model will be bogus if not consolidated so don't confuse the user by saving it
                logger.info(
                    f"Did not save the model {path} because stage3_gather_16bit_weights_on_model_save is False")
                return False
        else:
            state_dict = self.module_state_dict(exclude_frozen_parameters=exclude_frozen_parameters)

        tag = f"global_step{self.global_steps}"
        tag = str(tag)
        commit_info = CheckpointCommitInfo(tag=tag, save_dir=save_dir, save_latest=False)
        self.checkpoint_engine.create(commit_info)