Revert "refactor: defer DeepSpeed import and logging configuration until needed"

This reverts commit 59dda02.

Author: littlebullGit

src/lightning/fabric/strategies/deepspeed.py

@@ -282,10 +282,10 @@ def __init__(
                 sub_group_size=sub_group_size,
             )
 
+        import deepspeed
+
         self._config_initialized = False
-        # Defer importing and configuring DeepSpeed logging until it is actually needed.
-        # Store the desired logging level to be applied on first use.
-        self._logging_level = logging_level
+        deepspeed.utils.logging.logger.setLevel(logging_level)
 
         self.remote_device = remote_device
         self.load_full_weights = load_full_weights
@@ -374,8 +374,6 @@ def module_sharded_context(self) -> AbstractContextManager:
 
         import deepspeed
 
-        deepspeed.utils.logging.logger.setLevel(self._logging_level)
-
         assert self._config_initialized
         return deepspeed.zero.Init(
             enabled=self.zero_stage_3,
@@ -603,8 +601,6 @@ def _initialize_engine(
         """
         import deepspeed
 
-        deepspeed.utils.logging.logger.setLevel(self._logging_level)
-
         model_parameters = filter(lambda p: p.requires_grad, model.parameters())
         deepspeed_engine, deepspeed_optimizer, _, deepspeed_scheduler = deepspeed.initialize(
             args=argparse.Namespace(device_rank=self.root_device.index),
@@ -632,20 +628,14 @@ def _setup_distributed(self) -> None:
         _validate_device_index_selection(self.parallel_devices)
         reset_seed()
         self._set_world_ranks()
-        # Avoid initializing DeepSpeed distributed for single-process runs. This also avoids importing
-        # DeepSpeed in environments where it may not be fully functional (e.g., missing nvcc),
-        # while still allowing configuration and dataloader setup logic to run.
-        if self.world_size > 1:
-            self._init_deepspeed_distributed()
+        self._init_deepspeed_distributed()
         if not self._config_initialized:
             self._format_config()
             self._config_initialized = True
 
     def _init_deepspeed_distributed(self) -> None:
         import deepspeed
 
-        deepspeed.utils.logging.logger.setLevel(self._logging_level)
-
         assert self.cluster_environment is not None
         if platform.system() != "Windows":
             # do not set env variables on windows, allow deepspeed to control setup
@@ -671,8 +661,6 @@ def _set_node_environment_variables(self) -> None:
     def _set_deepspeed_activation_checkpointing(self) -> None:
         import deepspeed
 
-        deepspeed.utils.logging.logger.setLevel(self._logging_level)
-
         assert isinstance(self.config, dict)
         if self.config.get("activation_checkpointing"):
             checkpoint_config = self.config["activation_checkpointing"]