[Metric logging] log config.yaml (#605)

Co-authored-by: Felipe Mello <[email protected]>

Author: felipemello1

apps/grpo/main.py

@@ -15,6 +15,7 @@
 import torch
 import torch.nn.functional as F
 import torchstore as ts
+import yaml
 from datasets import load_dataset
 from forge.actors._torchstore_utils import (
     get_dcp_whole_state_dict_key,
@@ -33,11 +34,14 @@
 from forge.observability.perf_tracker import Tracer
 from forge.types import LauncherConfig, ProvisionerConfig
 from forge.util.config import parse
+from forge.util.logging import get_logger
 from forge.util.ops import compute_logprobs
 from monarch.actor import endpoint
-from omegaconf import DictConfig
+from omegaconf import DictConfig, OmegaConf
 from vllm.transformers_utils.tokenizer import get_tokenizer
 
+logger = get_logger("INFO")
+
 
 @dataclass
 class Episode:
@@ -358,9 +362,14 @@ async def drop_weights(version: int):
 
 async def main(cfg: DictConfig):
     """Main GRPO training loop with rollout and training processes."""
-    group_size = cfg.group_size
-    max_req_tokens = cfg.max_req_tokens
-    max_res_tokens = cfg.max_res_tokens
+    # Convert OmegaConf config to plain dict
+    run_config_for_logging = OmegaConf.to_container(cfg, resolve=True)
+
+    # Log config
+    logger.info("=" * 30 + " CONFIGURATION " + "=" * 30)
+    logger.info(
+        yaml.dump(run_config_for_logging, default_flow_style=False, sort_keys=False)
+    )
 
     # ---- Global setups ---- #
     provisioner = None
@@ -372,8 +381,11 @@ async def main(cfg: DictConfig):
         provisioner = await init_provisioner()
 
     metric_logging_cfg = cfg.get("metric_logging", {})
+
     mlogger = await get_or_create_metric_logger(process_name="Controller")
-    await mlogger.init_backends.call_one(metric_logging_cfg)
+    await mlogger.init_backends.call_one(
+        backend_config=metric_logging_cfg, run_config=run_config_for_logging
+    )
 
     # ---- Setup services ---- #
 
@@ -411,6 +423,10 @@ async def main(cfg: DictConfig):
         ),
     )
 
+    group_size = cfg.group_size
+    max_req_tokens = cfg.max_req_tokens
+    max_res_tokens = cfg.max_res_tokens
+
     # Set max_steps to the configured value, or -1 if not specified or Null
     max_steps = cfg.trainer.training.steps or -1
 

src/forge/observability/metric_actors.py

@@ -169,22 +169,27 @@ async def flush(
     async def init_backends(
         self,
         metadata_per_controller_backend: dict[str, dict[str, Any]],
-        config: dict[str, Any],
+        backend_config: dict[str, Any],
+        run_config: dict[str, Any] | None = None,
         global_step: int = 0,
     ) -> None:
         """Init per-rank logger backends and MetricCollector.
 
         Args:
             metadata_per_controller_backend (dict[str, dict[str, Any]]): Metadata from controller backends for shared state.
-            config (dict[str, Any]): Backend configurations with logging modes and settings.
+            backend_config (dict[str, Any]): Backend configurations with logging modes and settings.
+            run_config (dict[str, Any] | None): Your application's configuration
+                (hyperparameters, dataset, model settings) to log to backends for
+                experiment tracking.
             global_step (int): Initial step for metrics.
         """
         collector = MetricCollector()
         await collector.init_backends(
             metadata_per_controller_backend,
-            config,
+            backend_config,
             global_step,
             process_name=self.process_name,
+            run_config=run_config,
         )
 
     @endpoint
@@ -211,6 +216,7 @@ class GlobalLoggingActor(ForgeActor):
     def __init__(self):
         self.fetchers: dict[str, LocalFetcherActor] = {}
         self.config: dict[str, Any] | None = None
+        self.run_config: dict[str, Any] | None = None
         self.global_logger_backends: dict[str, LoggerBackend] = {}
         self.metadata_per_controller_backend: dict[str, dict[str, Any]] = {}
 
@@ -267,15 +273,17 @@ def _validate_backend_config(
         }
 
     @endpoint
-    async def init_backends(self, config: dict[str, Any]) -> None:
+    async def init_backends(
+        self, backend_config: dict[str, Any], run_config: dict[str, Any] | None = None
+    ) -> None:
         """Sets config in global actor and initializes existing backends and collectors. Later spawned actors
         are initialized in `register_fetcher` endpoint.
 
         Controller backends (instantiated in the controller) can provide metadata to be shared with rank backends,
         e.g. shared run IDs for WandB. For details on logging modes, see `forge.observability.metrics.LoggingMode`.
 
         Args:
-            config (dict[str, Any]): Config for metric logging where keys are backend names.
+            backend_config (dict[str, Any]): Config for metric logging where keys are backend names.
                 Each backend config supports:
                 - logging_mode (str | LoggingMode): Check LoggingMode for options. Defaults to "global_reduce".
                 - per_rank_share_run (bool, default False): For per-rank modes only. Whether ranks
@@ -291,21 +299,23 @@ async def init_backends(self, config: dict[str, Any]) -> None:
                         "project": "my_project",
                     }
                 }
+            run_config (dict[str, Any] | None): Your application's configuration
+                (hyperparameters, dataset, model settings) to log to backends for
+                experiment tracking.
 
         Raises:
             ValueError: If backend config is invalid or missing required fields.
         """
         self.config = {}
+        self.run_config = run_config
 
         # Skip initialization if disabled by environment flag
         if FORGE_DISABLE_METRICS.get_value():
             return
 
         # Validate and normalize each backend config
-        for backend_name, backend_config in config.items():
-            self.config[backend_name] = self._validate_backend_config(
-                backend_name, backend_config
-            )
+        for backend_name, cfg in backend_config.items():
+            self.config[backend_name] = self._validate_backend_config(backend_name, cfg)
 
         # Initialize backends based on logging mode
         for backend_name, backend_config in self.config.items():
@@ -314,7 +324,11 @@ async def init_backends(self, config: dict[str, Any]) -> None:
             backend: LoggerBackend = get_logger_backend_class(backend_name)(
                 **backend_config
             )
-            await backend.init(role=BackendRole.GLOBAL, process_name="global_reduce")
+            await backend.init(
+                role=BackendRole.GLOBAL,
+                process_name="global_reduce",
+                run_config=self.run_config,
+            )
 
             # Extract metadata from controller logger to be shared with per-rank loggers
             if mode != LoggingMode.GLOBAL_REDUCE:
@@ -331,7 +345,7 @@ async def init_backends(self, config: dict[str, Any]) -> None:
         if self.fetchers:
             tasks = [
                 fetcher.init_backends.call(
-                    self.metadata_per_controller_backend, self.config
+                    self.metadata_per_controller_backend, self.config, self.run_config
                 )
                 for fetcher in self.fetchers.values()
             ]
@@ -351,7 +365,7 @@ async def register_fetcher(self, fetcher: LocalFetcherActor, proc_id: str) -> No
         if self.config:
             logger.debug(f"Initializing new LocalFetcherActor for proc_id={proc_id}")
             await fetcher.init_backends.call(
-                self.metadata_per_controller_backend, self.config
+                self.metadata_per_controller_backend, self.config, self.run_config
             )
 
     @endpoint

src/forge/observability/metrics.py

@@ -550,9 +550,10 @@ def __init__(self) -> None:
     async def init_backends(
         self,
         metadata_per_controller_backend: dict[str, dict[str, Any]] | None,
-        config: dict[str, Any],
+        backend_config: dict[str, Any],
         global_step: int = 0,
         process_name: str | None = None,
+        run_config: dict[str, Any] | None = None,
     ) -> None:
         """Initialize per-rank logger backends and MetricCollector state.
 
@@ -563,12 +564,15 @@ async def init_backends(
             metadata_per_controller_backend (Optional[Dict[str, Dict[str, Any]]]): Metadata from controller
                 for backends that require shared state across processes, e.g.,
                 {"wandb": {"shared_run_id": "abc123"}}.
-            config (Dict[str, Any]): Backend configurations where each key is a backend name
+            backend_config (Dict[str, Any]): Backend configurations where each key is a backend name
                 and value contains logging_mode and backend-specific settings.
                 e.g., {"wandb": {"logging_mode": "per_rank_no_reduce", "project": "my_proj"}}
             global_step (int, default 0): Initial step for logging. Can be used when
                 resuming from a checkpoint.
             process_name (str | None): The meaningful process name for logging.
+            run_config (dict[str, Any] | None): Your application's configuration
+                (hyperparameters, dataset, model settings) to log to backends for
+                experiment tracking.
         """
         if self._is_initialized:
             logger.debug(
@@ -583,8 +587,8 @@ async def init_backends(
         self.per_rank_no_reduce_backends: list[LoggerBackend] = []
 
         # Initialize backends based on logging mode
-        for backend_name, backend_config in config.items():
-            mode = backend_config["logging_mode"]
+        for backend_name, cfg in backend_config.items():
+            mode = cfg["logging_mode"]
 
             # sanity check
             if not isinstance(mode, LoggingMode):
@@ -605,13 +609,12 @@ async def init_backends(
                 )
 
             # instantiate local backend
-            backend: LoggerBackend = get_logger_backend_class(backend_name)(
-                **backend_config
-            )
+            backend: LoggerBackend = get_logger_backend_class(backend_name)(**cfg)
             await backend.init(
                 role=BackendRole.LOCAL,
                 controller_logger_metadata=controller_metadata,
                 process_name=self.proc_name_with_rank,
+                run_config=run_config,
             )
 
             # Categorize by logging mode
@@ -781,6 +784,7 @@ async def init(
         role: BackendRole,
         controller_logger_metadata: dict[str, Any] | None = None,
         process_name: str | None = None,
+        run_config: dict[str, Any] | None = None,
     ) -> None:
         """
         Initializes backend, e.g. wandb.run.init().
@@ -791,6 +795,9 @@ async def init(
             controller_logger_metadata (dict[str, Any] | None): From global backend for
                 backend that required shared info, e.g. {"shared_run_id": "abc123"}.
             process_name (str | None): Process name for logging.
+            run_config (dict[str, Any] | None): Your application's configuration
+                (hyperparameters, dataset, model settings) to log to backend for
+                experiment tracking.
 
         Raises: ValueError if missing metadata for shared local init.
         """
@@ -856,6 +863,7 @@ async def init(
         role: BackendRole,
         controller_logger_metadata: dict[str, Any] | None = None,
         process_name: str | None = None,
+        run_config: dict[str, Any] | None = None,
     ) -> None:
         self.process_name = process_name
 
@@ -927,13 +935,15 @@ async def init(
         role: BackendRole,
         controller_logger_metadata: dict[str, Any] | None = None,
         process_name: str | None = None,
+        run_config: dict[str, Any] | None = None,
     ) -> None:
         if controller_logger_metadata is None:
             controller_logger_metadata = {}
 
         # Pop name, if any, to concat to process_name.
         run_name = self.backend_kwargs.pop("name", None)
         self.process_name = process_name
+        self.run_config = run_config
 
         # Format run name based on mode and role
         if self.logging_mode == LoggingMode.GLOBAL_REDUCE:
@@ -964,20 +974,29 @@ async def init(
     async def _init_global(self, run_name: str | None):
         import wandb
 
-        self.run = wandb.init(name=run_name, **self.backend_kwargs)
+        self.run = wandb.init(
+            name=run_name, config=self.run_config, **self.backend_kwargs
+        )
 
     async def _init_per_rank(self, run_name: str):
         import wandb
 
-        self.run = wandb.init(name=run_name, **self.backend_kwargs)
+        self.run = wandb.init(
+            name=run_name, config=self.run_config, **self.backend_kwargs
+        )
 
     async def _init_shared_global(self, run_name: str | None):
         import wandb
 
         settings = wandb.Settings(
             mode="shared", x_primary=True, x_label="controller_primary"
         )
-        self.run = wandb.init(name=run_name, settings=settings, **self.backend_kwargs)
+        self.run = wandb.init(
+            name=run_name,
+            config=self.run_config,
+            settings=settings,
+            **self.backend_kwargs,
+        )
 
     async def _init_shared_local(
         self, run_name: str, shared_id: str, process_name: str
@@ -994,7 +1013,11 @@ async def _init_shared_local(
 
         settings = wandb.Settings(mode="shared", x_primary=False, x_label=process_name)
         self.run = wandb.init(
-            name=run_name, id=shared_id, settings=settings, **self.backend_kwargs
+            name=run_name,
+            id=shared_id,
+            config=self.run_config,
+            settings=settings,
+            **self.backend_kwargs,
         )
 
     async def log_batch(