🐛: fix optimizer selection logic so that FusedAdam is never loaded wh…

…en CPU offloading is enabled (#122)

Signed-off-by: Oleg S <[email protected]>

src/instructlab/training/main_ds.py

@@ -228,14 +228,17 @@ def make_inputs_require_grad(module, input, output):
             model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
 
     # need to use this only when the CPU offload optimizer is enabled
-    optimizer = FusedAdam(model.parameters(), lr=args.learning_rate, betas=(0.9, 0.95))
     if args.cpu_offload_optimizer:
         print(
-            "\033[33m!!! CPU offload optimizer enabled, switching optimizer to DeepSpeedCPUAdam !!!\033[0m"
+            "\033[33m!!! CPU offload optimizer enabled, using DeepSpeedCPUAdam !!!\033[0m"
         )
         optimizer = DeepSpeedCPUAdam(
             model.parameters(), lr=args.learning_rate, betas=(0.9, 0.95)
         )
+    else:
+        optimizer = FusedAdam(
+            model.parameters(), lr=args.learning_rate, betas=(0.9, 0.95)
+        )
 
     lr_scheduler = get_scheduler(
         name=args.lr_scheduler,