MLFlow Integration

docs/training/logging.md

@@ -154,6 +154,49 @@ The plugin automatically forwards the `WANDB_API_KEY` and by default injects CLI
 This allows seamless integration of W&B logging into your training workflow without manual configuration.
 
 
+### MLFlow
+
+Megatron Bridge can log metrics and artifacts to MLFlow, following the same pattern as the W&B integration.
+
+#### What Gets Logged
+
+When enabled, MLFlow receives:
+
+- Training configuration as run parameters
+- Scalar metrics (losses, learning rate, batch size, throughput, timers, memory, runtime, norms, energy, etc.)
+- Checkpoint artifacts saved under an experiment-specific artifact path per iteration
+
+#### Enable MLFlow Logging
+
+  1) Install MLFlow (if not already available):
+
+  ```bash
+  pip install mlflow
+  ```
+
+  2) Configure the tracking server (Optional):
+  - Either set `MLFLOW_TRACKING_URI` in the environment, or
+  - Pass an explicit `mlflow_tracking_uri` in the logger config.
+
+  3) Configure logging in your training setup.
+
+  ```python
+  from megatron.bridge.training.config import LoggerConfig
+
+  cfg.logger = LoggerConfig(
+      tensorboard_dir="./runs/tensorboard",
+      mlflow_experiment="my_megatron_experiment",
+      mlflow_run_name="llama32_1b_pretrain_run",
+      mlflow_tracking_uri="http://mlflow:5000",  # optional
+      mlflow_tags={                              # optional
+          "project": "llama32",
+          "phase": "pretrain",
+      },
+  )
+  ```
+
+
+
 #### Progress Log
 
 When `logger.log_progress` is enabled, the framework generates a `progress.txt` file in the checkpoint save directory.

src/megatron/bridge/training/checkpointing.py

@@ -59,7 +59,7 @@
 from megatron.bridge.training.state import GlobalState, TrainState
 from megatron.bridge.training.tokenizers.config import TokenizerConfig
 from megatron.bridge.training.tokenizers.tokenizer import MegatronTokenizer
-from megatron.bridge.training.utils import wandb_utils
+from megatron.bridge.training.utils import mlflow_utils, wandb_utils
 from megatron.bridge.training.utils.checkpoint_utils import (
     checkpoint_exists,
     ensure_directory_exists,
@@ -740,11 +740,21 @@ def wandb_finalize_fn() -> None:
                 wandb_writer=state.wandb_logger,
             )
 
+        def mlflow_finalize_fn() -> None:
+            mlflow_utils.on_save_checkpoint_success(
+                checkpoint_name,
+                save_dir,
+                train_state.step,
+                mlflow_logger=state.mlflow_logger,
+            )
+
         if ckpt_cfg.async_save:
             assert async_save_request is not None
             async_save_request.add_finalize_fn(wandb_finalize_fn)
+            async_save_request.add_finalize_fn(mlflow_finalize_fn)
         else:
             wandb_finalize_fn()
+            mlflow_finalize_fn()
 
     if ckpt_cfg.async_save:
         schedule_async_save(state, async_save_request)
@@ -1667,6 +1677,7 @@ def _load_checkpoint_from_path(
 
     if not torch.distributed.is_initialized() or is_last_rank():
         wandb_utils.on_load_checkpoint_success(checkpoint_name, load_dir, state.wandb_logger)
+        mlflow_utils.on_load_checkpoint_success(checkpoint_name, load_dir, state.mlflow_logger)
 
     torch.cuda.empty_cache()
 

src/megatron/bridge/training/config.py

@@ -890,6 +890,18 @@ class LoggerConfig:
     wandb_entity: Optional[str] = None
     """The wandb entity name."""
 
+    mlflow_experiment: Optional[str] = None
+    """The MLFlow experiment name."""
+
+    mlflow_run_name: Optional[str] = None
+    """The MLFlow run name."""
+
+    mlflow_tracking_uri: Optional[str] = None
+    """Optional MLFlow tracking URI."""
+
+    mlflow_tags: Optional[dict[str, str]] = None
+    """Optional tags to apply to the MLFlow run."""
+
     logging_level: int = logging.INFO
     """Set default logging level"""
 

src/megatron/bridge/training/state.py

@@ -31,7 +31,7 @@
 from megatron.bridge.training.nvrx_straggler import NVRxStragglerDetectionManager
 from megatron.bridge.training.tokenizers.tokenizer import build_tokenizer
 from megatron.bridge.training.utils.sig_utils import DistributedSignalHandler
-from megatron.bridge.utils.common_utils import get_rank_safe, get_world_size_safe
+from megatron.bridge.utils.common_utils import get_rank_safe, get_world_size_safe, warn_rank_0
 
 
 @dataclass
@@ -124,6 +124,7 @@ def __init__(self) -> None:
         self._tokenizer: Optional[Any] = None
         self._tensorboard_logger: Optional[SummaryWriter] = None
         self._wandb_logger: Optional[Any] = None
+        self._mlflow_logger: Optional[Any] = None
         self._timers: Optional[Timers] = None
         self._train_state: Optional[TrainState] = None
         self.rank_monitor_client: Optional[Any] = None
@@ -234,12 +235,89 @@ def safe_serialize(obj):
                 self._wandb_logger = None
         return self._wandb_logger
 
+    @property
+    def mlflow_logger(self) -> Optional[Any]:
+        """The MLFlow logger instance.
+
+        Uses the configuration under LoggerConfig to create or resume an MLFlow run.
+        Restricted to the last rank to avoid duplicate entries or probably any racing conditions.
+        """
+        if self._mlflow_logger is None:
+            cfg = self.cfg
+            if cfg is None:
+                self._mlflow_logger = None
+                return self._mlflow_logger
+
+            logger_cfg = cfg.logger
+            if logger_cfg.mlflow_experiment and get_rank_safe() == (get_world_size_safe() - 1):
+                if logger_cfg.mlflow_run_name == "":
+                    raise ValueError("Please specify the mlflow_run_name for MLFlow logging!")
+
+                import mlflow
+
+                # set tracking URI
+                if logger_cfg.mlflow_tracking_uri:
+                    mlflow.set_tracking_uri(logger_cfg.mlflow_tracking_uri)
+
+                # Set or get experiment
+                mlflow.set_experiment(logger_cfg.mlflow_experiment)
+
+                # Prepare tags and params
+                def safe_serialize(obj: Any) -> str:
+                    """Safely convert any object to a JSON-serializable string."""
+                    try:
+                        result = str(obj)
+                        if not isinstance(result, str):
+                            return f"<{type(obj).__name__}>"
+                        return result
+                    except Exception:
+                        return f"<{type(obj).__name__}>"
+
+                def _flatten_dict(d: dict[str, Any], parent_key: str = "", sep: str = ".") -> dict[str, Any]:
+                    items: dict[str, Any] = {}
+                    for k, v in d.items():
+                        new_key = f"{parent_key}{sep}{k}" if parent_key else k
+                        if isinstance(v, dict):
+                            items.update(_flatten_dict(v, new_key, sep=sep))
+                        else:
+                            if isinstance(v, (list, tuple)):
+                                v = [safe_serialize(x) for x in v]
+                            items[new_key] = v
+                    return items
+
+                config_dict = cfg.to_dict()
+                sanitized_config = json.loads(json.dumps(config_dict, default=safe_serialize))
+                flat_params = _flatten_dict(sanitized_config)
+
+                # Start or resume a run
+                run_name = logger_cfg.mlflow_run_name
+                tags = logger_cfg.mlflow_tags or {}
+
+                active_run = mlflow.active_run()
+                if active_run is None:
+                    mlflow.start_run(run_name=run_name, tags=tags or None)
+                elif tags:
+                    # If there is already an active run, at least set provided tags
+                    mlflow.set_tags(tags)
+
+                # Log flattened configuration as params (best-effort)
+                stringified_params = {
+                    key: (safe_serialize(value) if not isinstance(value, (int, float, bool, str)) else value)
+                    for key, value in flat_params.items()
+                }
+                mlflow.log_params(stringified_params)
+                self._mlflow_logger = mlflow
+            else:
+                self._mlflow_logger = None
+        return self._mlflow_logger
+
     @property
     def timers(self) -> Timers:
         """The Megatron Timers instance used for tracking execution times."""
         if self._timers is None:
             self._timers = Timers(self.cfg.logger.timing_log_level, self.cfg.logger.timing_log_option)
             self._timers.write_to_wandb = types.MethodType(_timers_write_to_wandb, self._timers)
+            self._timers.write_to_mlflow = types.MethodType(_timers_write_to_mlflow, self._timers)
         return self._timers
 
     @property
@@ -344,6 +422,7 @@ def reset_for_restart(self) -> None:
         self._train_state = None
         self._tensorboard_logger = None
         self._wandb_logger = None
+        self._mlflow_logger = None
         self._energy_monitor = None
         self._energy_monitor_created = False
         self._signal_handler = None
@@ -371,3 +450,28 @@ def _timers_write_to_wandb(
         for name in name_to_min_max_time:
             _, max_time = name_to_min_max_time[name]
             writer.log({name + "-time": max_time}, iteration)
+
+
+def _timers_write_to_mlflow(
+    self: Timers,
+    names: list[str],
+    logger: Any,
+    iteration: int,
+    normalizer: float = 1.0,
+    reset: bool = True,
+    barrier: bool = False,
+) -> None:
+    """Patch to write timers to MLFlow for Megatron Core Timers."""
+    assert normalizer > 0.0
+    name_to_min_max_time = self._get_global_min_max_time(names, reset, barrier, normalizer)
+    if logger is not None:
+        metrics: dict[str, float] = {}
+        for name in name_to_min_max_time:
+            _, max_time = name_to_min_max_time[name]
+            sanitized_name = name.replace("/", "_") + "-time"
+            metrics[sanitized_name] = max_time
+        try:
+            logger.log_metrics(metrics, step=iteration)
+        except Exception:
+            # continue training
+            warn_rank_0("Failed to log timer metrics to MLFlow; continuing without timer metrics.")

src/megatron/bridge/training/utils/mlflow_utils.py

@@ -0,0 +1,72 @@
+from pathlib import Path
+from typing import Any, Optional
+
+from megatron.bridge.utils.common_utils import print_rank_last
+
+
+def on_save_checkpoint_success(
+    checkpoint_path: str,
+    save_dir: str,
+    iteration: int,
+    mlflow_logger: Optional[Any],
+) -> None:
+    """Callback executed after a checkpoint is successfully saved.
+
+    If an MLFlow logger is provided, logs the checkpoint directory as an MLFlow
+    artifact under a structured artifact path that includes the iteration number.
+
+    Args:
+        checkpoint_path: The path to the specific checkpoint file/directory saved.
+        save_dir: The base directory where checkpoints are being saved.
+        iteration: The training iteration at which the checkpoint was saved.
+        mlflow_logger: The MLFlow module (e.g., ``mlflow``) with an active run.
+                       If None, this function is a no-op.
+    """
+    if mlflow_logger is None:
+        return
+
+    try:
+        checkpoint_path = str(Path(checkpoint_path).resolve())
+        base_name = Path(save_dir).name or "checkpoints"
+        artifact_subdir = f"{base_name}/iter_{iteration:07d}"
+        mlflow_logger.log_artifacts(checkpoint_path, artifact_path=artifact_subdir)
+    except Exception as exc:
+        # continue training
+        print_rank_last(f"Failed to log checkpoint artifacts to MLFlow: {exc}")
+
+
+def on_load_checkpoint_success(
+    checkpoint_path: str,
+    load_dir: str,
+    mlflow_logger: Optional[Any],
+) -> None:
+    """Callback executed after a checkpoint is successfully loaded.
+
+    For MLFlow, this emits a simple metric and tag to document which checkpoint
+    was loaded during the run. It does not perform artifact lookups.
+
+    Args:
+        checkpoint_path: The path to the specific checkpoint file/directory loaded.
+        load_dir: The base directory from which the checkpoint was loaded.
+        mlflow_logger: The MLFlow module (e.g., ``mlflow``) with an active run.
+                       If None, this function is a no-op.
+    """
+    if mlflow_logger is None:
+        return
+
+    try:
+        resolved_ckpt = str(Path(checkpoint_path).resolve())
+        resolved_load_dir = str(Path(load_dir).resolve())
+        mlflow_logger.set_tags(
+            {
+                "last_loaded_checkpoint": resolved_ckpt,
+                "checkpoint_base_dir": resolved_load_dir,
+            }
+        )
+    except Exception as exc:
+        print_rank_last(f"Failed to record loaded checkpoint information to MLFlow: {exc}")
+
+
+def _sanitize_mlflow_metrics(metrics: dict[str, Any]) -> dict[str, Any]:
+    """Sanitize all metric names in a dictionary for MLFlow logging."""
+    return {key.replace("/", "_"): value for key, value in metrics.items()}