refactor: defer DeepSpeed import and logging configuration until needed

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
-        deepspeed.utils.logging.logger.setLevel(logging_level)
+        # 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
 
         self.remote_device = remote_device
         self.load_full_weights = load_full_weights
@@ -374,6 +374,8 @@ 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,
@@ -601,6 +603,8 @@ 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),
@@ -628,14 +632,20 @@ def _setup_distributed(self) -> None:
         _validate_device_index_selection(self.parallel_devices)
         reset_seed()
         self._set_world_ranks()
-        self._init_deepspeed_distributed()
+        # 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()
         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
@@ -661,6 +671,8 @@ 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"]