diff --git a/src/instructlab/training/main_ds.py b/src/instructlab/training/main_ds.py
index c5cdb2b..cf8cd8c 100644
--- a/src/instructlab/training/main_ds.py
+++ b/src/instructlab/training/main_ds.py
@@ -1,4 +1,4 @@
-# SPDX-License-Identifier: Apache-2.0
+ # SPDX-License-Identifier: Apache-2.0
 
 # Standard
 from copy import deepcopy
@@ -12,8 +12,18 @@
 
 # Third Party
 from accelerate import Accelerator
-from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam
-from deepspeed.runtime.zero.utils import ZeRORuntimeException
+try:
+    from deepspeed.ops.adam import DeepSpeedCPUAdam, FusedAdam
+except ImportError:
+    DeepSpeedCPUAdam = None
+    FusedAdam = None
+    print("DeepSpeed Adam optimizers are not available. Some features may be unavailable.")
+
+try:
+    from deepspeed.runtime.zero.utils import ZeRORuntimeException
+except ImportError:
+    ZeRORuntimeException = None
+    print("ZeRORuntimeException is not available. Some features may be unavailable.")
 
 # pylint: disable=no-name-in-module
 from instructlab.dolomite.hf_models import GPTDolomiteForCausalLM