Alayshah/lea 182 deepspeed for dpo (#2)

* naming for example datasets

* changed setup_training call

* dpo configs updated

* sft configs updated

* logging / deepspeed import changes

---------

Co-authored-by: EdoardoMosca <edoardo@liquid.ai>

Author: alay2shah

config.py

@@ -19,10 +19,14 @@
     subset: Dataset subset to use for HuggingFace datasets, default None
 """
 
-smoltalk_dataset = DatasetLoader(
+example_sft_dataset = DatasetLoader(
     "HuggingFaceTB/smoltalk", "sft", limit=1000, test_size=0.2, subset="all"
 )
 
+example_dpo_dataset = DatasetLoader(
+    "mlabonne/orpo-dpo-mix-40k", "dpo", limit=1000, test_size=0.2, subset="default"
+)
+
 
 #################################
 #       Training Config         #
@@ -72,10 +76,10 @@
 """
 
 JOB_CONFIG = JobConfig(
-    job_name="my_output_dir",
+    job_name="my_job",
     model_name="LFM2-1.2B",
     training_type="sft",
-    dataset=smoltalk_dataset,
+    dataset=example_sft_dataset,
     training_config=training_config,
     peft_config=peft_config,
 )

src/leap_finetune/__init__.py

@@ -1,14 +1,16 @@
 import sys
 
-from leap_finetune.trainer import ray_trainer
 from leap_finetune.utils.logging import setup_training_environment
-from leap_finetune.utils.constants import LEAP_FINETUNE_DIR
 
-# Set Ray Core logging before importing Ray
 setup_training_environment()
 
+from leap_finetune.trainer import ray_trainer  # noqa
+from leap_finetune.utils.constants import LEAP_FINETUNE_DIR  # noqa
+
 
 def main() -> None:
+    # Set Ray Core logging before importing Ray
+
     print("Launching leap-finetune ✅")
 
     # Import config only in main process, not Ray workers
@@ -29,7 +31,7 @@ def main() -> None:
             "JOB_CONFIG not found in config.py. Please define JOB_CONFIG in your config.py file."
         )
     except ValueError as e:
-        raise ValueError(f"Invalid JOB_CONFIG: {e}")
+        raise ValueError(f"Issue with JOB_CONFIG: {e}")
     finally:
         sys.path.remove(root_path)
 

src/leap_finetune/configs/dpo_configs.py

@@ -15,15 +15,6 @@
     "train_micro_batch_size_per_gpu": "auto",
     "gradient_clipping": "auto",
     "gradient_accumulation_steps": "auto",
-    "optimizer": {
-        "type": "AdamW",
-        "params": {
-            "lr": "auto",  # Uses learning_rate from training config
-            "betas": "auto",  # DEFAULT: (0.9, 0.999)
-            "eps": "auto",  # DEFAULT: 1e-8
-            "weight_decay": "auto",  # DEFAULT: 0.01
-        },
-    },
     "bf16": {"enabled": "auto"},
     "activation_checkpointing": {
         "partition_activations": False,
@@ -54,5 +45,6 @@
     "save_strategy": "epoch",
     "eval_strategy": "epoch",
     "load_best_model_at_end": True,
-    # "deepspeed": DEEPSEED_CONFIG,
+    "ddp_find_unused_parameters": False,
+    "deepspeed": DEEPSEED_CONFIG,
 }

src/leap_finetune/configs/sft_configs.py

@@ -54,5 +54,6 @@
     "save_strategy": "epoch",
     "eval_strategy": "epoch",
     "load_best_model_at_end": True,
+    "ddp_find_unused_parameters": False,
     "deepspeed": DEEPSEED_CONFIG,
 }

src/leap_finetune/utils/logging.py

@@ -1,14 +1,48 @@
 import os
+import warnings
+import logging
+
+_ENV_DONE = False
 
 
 def setup_training_environment() -> None:
-    """
-    Setup the training environment with proper logging configuration.
-    Call this in the main trainer before ray.init().
-    """
+    global _ENV_DONE
+    if _ENV_DONE:
+        return
+
+    os.environ.setdefault("DS_DISABLE_CONFIG_PRINT", "1")
+    os.environ.setdefault("DEEPSPEED_LOG_LEVEL", "ERROR")
+    warnings.filterwarnings("ignore")  # keep only tracebacks
+
+    cache = "/dev/shm/triton_cache"
+    try:
+        os.makedirs(cache, exist_ok=True)
+    except OSError:
+        cache = "/tmp/triton_cache"
+        os.makedirs(cache, exist_ok=True)
+    os.environ["TRITON_CACHE_DIR"] = cache
+
+    try:
+        import deepspeed
+        from deepspeed.runtime.bf16_optimizer import BF16_Optimizer
+
+        ds_log = deepspeed.utils.logging.logger
+        ds_log.setLevel(logging.ERROR)
+        for h in ds_log.handlers:
+            h.setLevel(logging.ERROR)
+
+        _orig_ds_destroy = BF16_Optimizer.destroy
+
+        def _safe_ds_destroy(self, *a, **kw):
+            try:
+                _orig_ds_destroy(self, *a, **kw)
+            except IndexError:
+                pass
+
+        BF16_Optimizer.destroy = _safe_ds_destroy
 
-    triton_cache = os.path.expanduser("~/tmp/triton_cache")
-    os.makedirs(triton_cache, exist_ok=True)
-    os.environ["TRITON_CACHE_DIR"] = triton_cache
+    except ImportError:
+        print("Deepspeed not available in environment; nothing to patch")
 
     print("Training environment configured ✅")
+    _ENV_DONE = True