Wandb integration (#6)

* readme change

* wandb flag

* using job_name for logging

* add wandb reporting to all runners

* fmt

* added wandb dependency

* trainer logging and ray init changes

* training loop updates

* logging adds

* peft dir

* deepspeed triton error fix

* config changes

* add wandb

* add first logging step

* wandb init

* set wandb mode

* readme

* wandb settings

* ray logging

* force ray version

* rm tempfile

* clean up runners

Author: alay2shah

README.md

@@ -54,7 +54,32 @@ It uses Ray Train + Accelerate for distributed training.
 
 Unless you overwrote `output_dir`, results will be stored in `outputs/training_type/job_name/`
 
-### 3. Bundle Checkpoint for LEAP
+### 3. (Optional) Experiment Tracking with Weights & Biases
+
+To enable experiment tracking (using [Weights & Biases](https://wandb.ai)):
+
+- Set `wandb_logging=True` in `config.py` in your `user_config` overrides or default configs.
+- **Offline mode (default)**: If no `WANDB_API_KEY` is set, wandb logs locally to `./wandb/` directory. No API key needed!
+- **Online mode**: Set the `WANDB_API_KEY` environment variable to sync to wandb.ai dashboard:
+
+```bash
+export WANDB_API_KEY=your_api_key  # optional; for online syncing to wandb.ai
+```
+
+You can also customize the project name (defaults to `"leap-finetune"`):
+
+```bash
+export WANDB_PROJECT=my-custom-project  # optional; defaults to "leap-finetune"
+```
+
+After training, view your metrics:
+
+- **Online mode**: View at `https://wandb.ai/<your-entity>/<project-name>/runs/<run-name>`
+- **Offline mode**: Sync later with `wandb sync ./wandb/offline-run-*` or view locally
+
+Runs are named after your `job_name` and metrics are reported via TRL/Transformers. Training metrics (loss, learning rate, etc.) are logged every `logging_steps` (default: 10), and evaluation metrics are logged at the end of each epoch.
+
+### 4. Bundle Checkpoint for LEAP
 
 When training is done, you can bundle your output checkpoint with `leap-bundle` to use it directly within LEAP. Checkout our [Quick Start guide](https://leap.liquid.ai/docs/leap-bundle/quick-start?utm_source=github&utm_medium=link&utm_campaign=LEAP&utm_content=general).
 

config.py

@@ -65,6 +65,7 @@
     "per_device_train_batch_size": None,
     "gradient_accumulation_steps": None,
     "learning_rate": None,
+    "wandb_logging": False,  # If set to True, set your wandb API key via 'export WANDB_API_KEY=<your_api_key>'
 }
 
 
@@ -86,7 +87,7 @@
 
 JOB_CONFIG = JobConfig(
     job_name="my_job_name",
-    model_name="LFM2-8B-A1B",
+    model_name="LFM2-1.2B",
     training_type="sft",
     dataset=example_sft_dataset,
     training_config=TrainingConfig.MOE_SFT.override(**user_config),

pyproject.toml

@@ -11,14 +11,15 @@ dependencies = [
     "accelerate>=1.7.0",
     "peft>=0.15.2",
     "deepspeed>=0.17.1",
-    "ray[train]>=2.47.1",
     "torch>=2.7.1",
     "transformers @ git+https://github.com/huggingface/transformers.git@0c9a72e4576fe4c84077f066e585129c97bfd4e6",
     "trl>=0.18.2",
     "rich>=14.1.0",
     "pillow>=11.3.0",
     "mpi4py>=4.1.0",
     "liger-kernel>=0.6.2",
+    "wandb>=0.22.3",
+    "ray==2.48.0",
 ]
 
 [project.scripts]

src/leap_finetune/configs/sft_configs.py

@@ -80,7 +80,8 @@
     "lr_scheduler_type": "linear",
     "warmup_steps": 100,
     "warmup_ratio": 0.2,
-    "logging_steps": 10,
+    "logging_steps": 10,  # Log training metrics every 10 steps
+    "logging_first_step": True,  # Log at step 0 to see initial metrics
     "save_strategy": "epoch",
     "eval_strategy": "epoch",
     "load_best_model_at_end": True,

src/leap_finetune/configs/vlm_sft_config.py

@@ -50,6 +50,7 @@
     "warmup_steps": 100,
     "warmup_ratio": 0.2,
     "logging_steps": 10,
+    "logging_first_step": True,
     "save_strategy": "epoch",
     "eval_strategy": "epoch",
     "load_best_model_at_end": True,

src/leap_finetune/trainer.py

@@ -1,18 +1,23 @@
 import ray
 import os
+import logging
 
 from accelerate.utils import set_seed
 from ray.train import RunConfig, ScalingConfig
 from ray.runtime_env import RuntimeEnv
 from ray.train.torch import TorchTrainer, TorchConfig
-from rich.console import Console
-from rich.panel import Panel
 from torch import cuda
 
 from leap_finetune.utils.constants import RUNTIME_DIR
 from leap_finetune.training_loops.sft_run import sft_run
 from leap_finetune.training_loops.dpo_run import dpo_run
 from leap_finetune.training_loops.vlm_sft_run import vlm_sft_run
+from leap_finetune.utils.logging_utils import (
+    get_ray_env_vars,
+    print_next_steps_panel,
+    select_ray_temp_dir,
+    select_object_spilling_dir,
+)
 
 
 #################################
@@ -38,27 +43,27 @@ def ray_trainer(job_config: dict) -> None:
     os.makedirs(ray_temp_dir, exist_ok=True)
 
     if not ray.is_initialized():
+        # Force local init and avoid accidental cluster connects
+        for key in ("RAY_ADDRESS", "RAY_HEAD_IP", "RAY_HEAD_NODE_ADDRESS", "RAY_PORT"):
+            os.environ.pop(key, None)
+
+        ray_temp_dir = select_ray_temp_dir(os.path.expanduser("~/ray_temp"))
+        spill_dir = select_object_spilling_dir(ray_temp_dir)
+
+        # Reduce Ray logging verbosity
+        ray_logger = logging.getLogger("ray")
+        ray_logger.setLevel(logging.ERROR)  # Only show errors, not INFO/WARNING
+
+        runtime_env = RuntimeEnv(
+            working_dir=str(RUNTIME_DIR),
+            env_vars=get_ray_env_vars(ray_temp_dir),
+        )
+
         ray.init(
-            runtime_env=RuntimeEnv(
-                working_dir=str(RUNTIME_DIR),
-                env_vars={
-                    "TMPDIR": ray_temp_dir,
-                    "TEMP": ray_temp_dir,
-                    "TMP": ray_temp_dir,
-                    "NCCL_IB_DISABLE": "1",
-                    "NCCL_P2P_DISABLE": "1",  # Added to force CPU level communication during synchronization
-                    "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1",
-                    "NCCL_SOCKET_IFNAME": "lo",
-                    "TORCH_NCCL_BLOCKING_WAIT": "1",
-                    "NCCL_TIMEOUT": "300",  # 5 minute safe timeout
-                    "RAY_DISABLE_IMPORT_WARNING": "1",
-                    "RAY_memory_monitor_refresh_ms": "0",
-                    # Suppress Ray Data verbose logging
-                    "RAY_DATA_DISABLE_PROGRESS_BARS": "1",
-                    "RAY_IGNORE_UNHANDLED_ERRORS": "1",
-                },
-            ),
+            address="local",
+            runtime_env=runtime_env,
             _temp_dir=ray_temp_dir,
+            object_spilling_directory=spill_dir,
         )
 
     if training_type == "sft":
@@ -72,6 +77,7 @@ def ray_trainer(job_config: dict) -> None:
 
     train_loop_config = {
         "model_name": job_config["model_name"],
+        "job_name": job_config.get("job_name", "leap-ft-run"),
         "train_config": job_config["training_config"],
         "peft_config": job_config["peft_config"],
         "dataset": job_config["dataset"],
@@ -98,24 +104,9 @@ def ray_trainer(job_config: dict) -> None:
 
     trainer.fit()
 
-    console = Console()
-    quick_start_url = (
-        "https://leap.liquid.ai/docs/leap-bundle/quick-start?utm_source=github"
-        "&utm_medium=link&utm_campaign=LEAP&utm_content=general"
-    )
-
-    cta_message = (
-        "[bold green]Training complete![/bold green]\n\n"
-        f"[bold]Checkpoint directory:[/bold] [cyan]{output_dir}[/cyan]\n\n"
-        "Bundle your output checkpoint with [bold]leap-bundle[/bold] to use it in LEAP:\n"
-        f"[dim]leap-bundle create {output_dir}/[CHECKPOINT_NAME][/dim]\n\n"
-        f"Quick Start: [link={quick_start_url}]{quick_start_url}[/link]"
-    )
-
-    console.print(
-        Panel.fit(
-            cta_message,
-            title="Next Step: Bundle for LEAP",
-            border_style="green",
-        )
-    )
+    print_next_steps_panel(output_dir)
+    # Ensure Ray cleans up resources promptly to avoid post-training hangs
+    try:
+        ray.shutdown()
+    except Exception:
+        pass

src/leap_finetune/training_loops/dpo_run.py

@@ -4,11 +4,13 @@
 from transformers import PreTrainedTokenizerBase
 from trl import DPOConfig, DPOTrainer
 from ray.train.huggingface.transformers import prepare_trainer
+from ray.train import get_context
 
 from leap_finetune.configs.distributed_configs import MOE_FSDP_CONFIG
 from leap_finetune.utils.load_models import load_model
 from leap_finetune.utils.model_utils import is_moe_model_from_name
 from leap_finetune.utils.peft import apply_peft_to_model, merge_and_save_peft_model
+from leap_finetune.utils.logging_utils import init_wandb_if_enabled
 
 
 def dpo_run(training_config: dict) -> None:
@@ -18,29 +20,42 @@ def dpo_run(training_config: dict) -> None:
         tuple[Dataset, Dataset], training_config.get("dataset")
     )
 
-    train_config = training_config.get("train_config")
     peft_config = training_config.get("peft_config")
     model_name = training_config.get("model_name", "")
+    job_name = training_config.get("job_name", "leap-ft-run")
 
     # Check for MoE model
     is_moe = is_moe_model_from_name(model_name)
     use_fsdp = is_moe and peft_config is None
 
-    # Remove non-DPOConfig parameters
-    train_config.pop("training_type", None)
+    # Filter out non-DPOConfig parameters
+    excluded_keys = {"training_type", "wandb_logging"}
+    if use_fsdp:
+        excluded_keys.add("deepspeed")  # Remove deepspeed when using FSDP
+
+    train_config_filtered = {
+        k: v
+        for k, v in training_config.get("train_config").items()
+        if k not in excluded_keys
+    }
 
-    # Apply FSDP for MoE without PEFT
+    # Configure wandb reporting if enabled via config
+    wandb_logging = bool(
+        training_config.get("train_config", {}).get("wandb_logging", False)
+    )
+    init_wandb_if_enabled(job_name, wandb_logging)
+
+    # Build training args
+    config_kwargs = {
+        "report_to": "wandb" if wandb_logging else "none",
+        "run_name": job_name,
+        **train_config_filtered,
+    }
     if use_fsdp:
-        train_config.pop("deepspeed", None)
-        fsdp_config = MOE_FSDP_CONFIG["fsdp_config"].copy()
-        training_args = DPOConfig(
-            **train_config,
-            fsdp=MOE_FSDP_CONFIG["fsdp"],
-            fsdp_config=fsdp_config,
-        )
-    else:
-        # MoE with PEFT or non-MoE: use DeepSpeed (already in config)
-        training_args = DPOConfig(**train_config)
+        config_kwargs["fsdp"] = MOE_FSDP_CONFIG["fsdp"]
+        config_kwargs["fsdp_config"] = MOE_FSDP_CONFIG["fsdp_config"]
+
+    training_args = DPOConfig(**config_kwargs)
 
     # Load model after config is created
     model, tokenizer = load_model(model_name)
@@ -59,8 +74,27 @@ def dpo_run(training_config: dict) -> None:
 
     # Start training
     trainer = prepare_trainer(trainer)
-    trainer.train()
+    try:
+        trainer.train()
+        print("✅ Training completed successfully")
+    except RuntimeError as e:
+        error_msg = str(e)
+        if any(
+            keyword in error_msg.lower()
+            for keyword in ["cuda error", "ecc error", "nccl", "collective", "timeout"]
+        ):
+            print(
+                f"⚠️  Training completed but hit distributed communication error during cleanup: {error_msg}"
+            )
+            print(
+                "✅ Training was successful - error occurred in post-training synchronization"
+            )
+        else:
+            raise e
 
     # Save PEFT model if applicable
     if peft_config:
-        merge_and_save_peft_model(model, tokenizer, training_args.output_dir)
+        ctx = get_context()
+        is_rank_zero = ctx is None or ctx.get_world_rank() == 0
+        if is_rank_zero:
+            merge_and_save_peft_model(model, tokenizer, training_args.output_dir)

src/leap_finetune/training_loops/sft_run.py

@@ -4,11 +4,13 @@
 from transformers import PreTrainedTokenizerBase
 from trl import SFTConfig, SFTTrainer
 from ray.train.huggingface.transformers import prepare_trainer
+from ray.train import get_context
 
 from leap_finetune.configs.distributed_configs import MOE_FSDP_CONFIG
 from leap_finetune.utils.load_models import load_model
 from leap_finetune.utils.model_utils import is_moe_model_from_name
 from leap_finetune.utils.peft import apply_peft_to_model, merge_and_save_peft_model
+from leap_finetune.utils.logging_utils import init_wandb_if_enabled
 
 
 def sft_run(training_config: dict) -> None:
@@ -18,29 +20,42 @@ def sft_run(training_config: dict) -> None:
         tuple[Dataset, Dataset], training_config.get("dataset")
     )
 
-    train_config = training_config.get("train_config")
     peft_config = training_config.get("peft_config")
     model_name = training_config.get("model_name", "")
+    job_name = training_config.get("job_name", "leap-ft-run")
 
     # Check for MoE model
     is_moe = is_moe_model_from_name(model_name)
     use_fsdp = is_moe and peft_config is None
 
-    # Remove non-SFTConfig parameters
-    train_config.pop("training_type", None)
+    # Filter out non-SFTConfig parameters
+    excluded_keys = {"training_type", "wandb_logging"}
+    if use_fsdp:
+        excluded_keys.add("deepspeed")  # Remove deepspeed when using FSDP
+
+    train_config_filtered = {
+        k: v
+        for k, v in training_config.get("train_config").items()
+        if k not in excluded_keys
+    }
 
-    # Apply FSDP for MoE without PEFT
+    # Configure wandb reporting if enabled via config
+    wandb_logging = bool(
+        training_config.get("train_config", {}).get("wandb_logging", False)
+    )
+    init_wandb_if_enabled(job_name, wandb_logging)
+
+    # Build training args
+    config_kwargs = {
+        "report_to": "wandb" if wandb_logging else "none",
+        "run_name": job_name,
+        **train_config_filtered,
+    }
     if use_fsdp:
-        train_config.pop("deepspeed", None)
-        fsdp_config = MOE_FSDP_CONFIG["fsdp_config"].copy()
-        training_args = SFTConfig(
-            **train_config,
-            fsdp=MOE_FSDP_CONFIG["fsdp"],
-            fsdp_config=fsdp_config,
-        )
-    else:
-        # MoE with PEFT or non-MoE: use DeepSpeed (already in config)
-        training_args = SFTConfig(**train_config)
+        config_kwargs["fsdp"] = MOE_FSDP_CONFIG["fsdp"]
+        config_kwargs["fsdp_config"] = MOE_FSDP_CONFIG["fsdp_config"]
+
+    training_args = SFTConfig(**config_kwargs)
 
     # Load model after config is created
     model, tokenizer = load_model(model_name)
@@ -59,8 +74,27 @@ def sft_run(training_config: dict) -> None:
 
     # Start training
     trainer = prepare_trainer(trainer)
-    trainer.train()
+    try:
+        trainer.train()
+        print("✅ Training completed successfully")
+    except RuntimeError as e:
+        error_msg = str(e)
+        if any(
+            keyword in error_msg.lower()
+            for keyword in ["cuda error", "ecc error", "nccl", "collective", "timeout"]
+        ):
+            print(
+                f"⚠️  Training completed but hit distributed communication error during cleanup: {error_msg}"
+            )
+            print(
+                "✅ Training was successful - error occurred in post-training synchronization"
+            )
+        else:
+            raise e
 
     # Save PEFT model if applicable
     if peft_config:
-        merge_and_save_peft_model(model, tokenizer, training_args.output_dir)
+        ctx = get_context()
+        is_rank_zero = ctx is None or ctx.get_world_rank() == 0
+        if is_rank_zero:
+            merge_and_save_peft_model(model, tokenizer, training_args.output_dir)