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

Code before:
        self._do_args_sanity_check(args)
        self._configure_with_arguments(args, mpu)
        self._do_sanity_check()
        if self.autotp_size() > 1:
            self._configure_tensor_parallel(model, self.tensor_parallel_config())
        see_memory_usage("DeepSpeed Engine: After args sanity test", force=self.memory_breakdown())
        if mpu is not None:
            if self.elasticity_enabled():
                if not self.is_elastic_model_parallel_supported():
                    assert not self.elasticity_enabled(), ("Elasticity is not currently supported"
                                                           " with model parallelism.")

        self._set_distributed_vars(args)

        dist.configure(self._config)

        self.monitor = MonitorMaster(self._config.monitor_config)

        see_memory_usage(
            "DeepSpeed Engine: Before configure distributed model",
            force=self.memory_breakdown(),
        )

        self.pipeline_parallelism = isinstance(model, PipelineModule)

        self._deepcompile_active = None
        self._deepcompile_noop_warning_emitted = False

        # Configure distributed model
        self._configure_distributed_model(model)

        # These hooks should be disabled later if DeepCompile is not active.
        self.module_forward_pre_hook = self._create_module_forward_pre_hook()
        self.module_forward_post_hook = self._create_module_forward_post_hook()

        # needed for zero_to_fp32 weights reconstruction to remap nameless data to state_dict
        self.param_names = {param: name for name, param in model.named_parameters()}

        self._get_model_parameters()

        see_memory_usage("DeepSpeed Engine: After configure distributed model")

        # Configure wall clock timers
        self.timers = SynchronizedWallClockTimer()
        # Throughput timer
        self.tput_timer = ThroughputTimer(self._config.timers_config,
                                          batch_size=self.train_batch_size(),
                                          steps_per_output=self.steps_per_print(),
                                          monitor_memory=False)

        log_dist(f"DeepSpeed Flops Profiler Enabled: {self.flops_profiler_enabled()}", ranks=[0])

Review comment (refactor):
Is it possible to simplify `_deepcompile_active` to a bool (instead of an `Optional[bool]`)? The only effect of initializing it with `None` is that `_set_deepcompile_active` will try to reinstall the hooks, and that is unnecessary because the hooks are already there.

By design, not calling `model.compile` should always mean DeepCompile is not activated. So it looks reasonable to me to make it a plain bool.

Code after:
        self._do_args_sanity_check(args)
        self._configure_with_arguments(args, mpu)
        self._do_sanity_check()
        if self.autotp_size() > 1:
            self._configure_tensor_parallel(model, self.tensor_parallel_config())
        see_memory_usage("DeepSpeed Engine: After args sanity test", force=self.memory_breakdown())
        if mpu is not None:
            if self.elasticity_enabled():
                if not self.is_elastic_model_parallel_supported():
                    assert not self.elasticity_enabled(), ("Elasticity is not currently supported"
                                                           " with model parallelism.")

        self._set_distributed_vars(args)

        dist.configure(self._config)

        self.monitor = MonitorMaster(self._config.monitor_config)

        see_memory_usage(
            "DeepSpeed Engine: Before configure distributed model",
            force=self.memory_breakdown(),
        )

        self.pipeline_parallelism = isinstance(model, PipelineModule)

        self._deepcompile_active = False

        # Configure distributed model
        self._configure_distributed_model(model)

        # These hooks should be disabled later if DeepCompile is not active.
        self.module_forward_pre_hook = self._create_module_forward_pre_hook()
        self.module_forward_post_hook = self._create_module_forward_post_hook()

        # needed for zero_to_fp32 weights reconstruction to remap nameless data to state_dict
        self.param_names = {param: name for name, param in model.named_parameters()}

        self._get_model_parameters()

        see_memory_usage("DeepSpeed Engine: After configure distributed model")

        # Configure wall clock timers
        self.timers = SynchronizedWallClockTimer()
        # Throughput timer
        self.tput_timer = ThroughputTimer(self._config.timers_config,
                                          batch_size=self.train_batch_size(),
                                          steps_per_output=self.steps_per_print(),
                                          monitor_memory=False)

        log_dist(f"DeepSpeed Flops Profiler Enabled: {self.flops_profiler_enabled()}", ranks=[0])