nits

Author: Felipe Mello

apps/sft/main.py

@@ -16,7 +16,6 @@
 import math
 import os
 import sys
-import warnings
 from functools import partial
 from typing import Any
 
@@ -117,14 +116,16 @@ async def setup_metric_logger(self):
         return mlogger
 
     def record_batch_metrics(self, data_metrics: list):
-        """Record dataset metrics using the observability system."""
+        """Since the dataloader creates new processes, we dont call `record_metric` in the dataset.
+        Instead, pop the metrics from the batch and record them here."""
         for metric in data_metrics:
             record_metric(metric.key, metric.value, metric.reduction)
 
     @endpoint
     async def setup(self):
         self.train_dataloader = self.setup_data()
         self.mlogger = await self.setup_metric_logger()
+
         # self.train_dataloader = self.setup_data(
         #     self.train_config.train_dataset_config,
         #     self.train_config.train_dataloader_config,
@@ -268,9 +269,7 @@ async def train(self) -> None:
 
             # Pop and record metrics from batch before moving to device
             self.record_batch_metrics(batch.pop("metrics", []))
-            record_metric(
-                "ForgeSFTRecipe/train_step/step", self.current_step, Reduce.MEAN
-            )
+            record_metric("ForgeSFTRecipe/train/step", self.current_step, Reduce.MEAN)
 
             # Move tensors to the appropriate device
             for k, v in batch.items():
@@ -306,23 +305,11 @@ def __repr__(self) -> str:
 
 async def run(cfg: DictConfig) -> None:
 
-    # TODO (allenwang28) Required for metric logging to work. Should be removed when V1 becomes default
-    MONARCH_HOSTMESH_V1 = os.getenv("MONARCH_HOSTMESH_V1")
-    if MONARCH_HOSTMESH_V1 != "1":
-        warnings.warn(
-            "MONARCH_HOSTMESH_V1 is set to {MONARCH_HOSTMESH_V1}. Setting it to '1' for SFT v2 to work properly. ",
-            UserWarning,
-            stacklevel=2,
-        )
-    os.environ["MONARCH_HOSTMESH_V1"] = "1"
-
     logging.info("Spawning recipe...")
     process_cfg = cfg.pop("processes")
 
     # Initialize metric logger in main process
-    metric_logging_cfg = cfg.get(
-        "metric_logging", {"console": {"logging_mode": "global_reduce"}}
-    )
+    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)
 
@@ -337,8 +324,6 @@ async def run(cfg: DictConfig) -> None:
     logging.info("Done training. Clean up")
     await recipe.cleanup.call()
 
-    # Shutdown metric logger
-    await mlogger.shutdown.call_one()
     await recipe.mesh.stop()
     logging.info("All done!")
 

src/forge/data/datasets/hf_dataset.py

@@ -232,8 +232,9 @@ def __iter__(self) -> Iterator[dict[str, Any]]:
                     # .map is applied lazily and the advantage would be to leverage caching.
                     sample = self._apply_transforms(sample)
 
-                    # Track the number of epochs completed for each dataset.
-                    # This is especially useful when interleaving multiple datasets.
+                    # Track the number of epochs completed for each dataset. This is
+                    # especially useful when interleaving multiple datasets, but
+                    # also necessary to track dataset-level metrics.
                     if "metrics" not in sample:
                         sample["metrics"] = []
 

src/forge/observability/metric_actors.py

@@ -28,7 +28,6 @@
     MetricCollector,
     reduce_metrics_states,
 )
-from forge.observability.utils import detect_actor_name_from_call_stack
 
 
 logger = logging.getLogger(__name__)
@@ -84,9 +83,6 @@ async def get_or_create_metric_logger(
         await mlogger.shutdown.call_one()
     """
 
-    if process_name is None:
-        process_name = detect_actor_name_from_call_stack()
-
     # Get or create the singleton global logger
     global _global_logger
 

src/forge/observability/metrics.py

@@ -128,11 +128,11 @@ def reduce_metrics_states(states: list[dict[str, dict[str, Any]]]) -> list[Metri
     states is more precise than merging locally reduced metrics.
 
     Args:
-        states (list[dict[str, dict[str, Any]]]): list of state of one or more metrics,
+        states (list[dict[str, dict[str, Any]]]): List of state of one or more metrics,
             normally retrieved using `forge.observability.metrics.MetricAccumulator.get_state()`.
 
     Returns:
-        list[Metric]: list of reduced metrics
+        list[Metric]: List of reduced metrics
 
     Example:
         states = [

tests/unit_tests/datasets/test_hf.py

@@ -113,7 +113,7 @@ def test_default_dataset_name(self, small_dataset_file):
             split="train",
             # dataset_name not provided - should auto-generate
             seed=SEED,
-            metric_transform=None,  # Now using new observability system
+            metric_transform=DefaultDatasetMetricTransform(),
             num_shards_per_rank=4,
         )
 
@@ -131,7 +131,7 @@ def test_default_dataset_name(self, small_dataset_file):
             dataset_name="my_dataset",
             weight=custom_weight,
             seed=SEED,
-            metric_transform=None,  # Now using new observability system
+            metric_transform=DefaultDatasetMetricTransform(),
             num_shards_per_rank=4,
         )
 
@@ -317,7 +317,7 @@ def create_loader():
                         dataset_name="epoch_test",
                         seed=SEED,
                         shuffle_buffer_size=0,  # No shuffle for determinism
-                        metric_transform=None,  # Now using new observability system
+                        metric_transform=DefaultDatasetMetricTransform(),
                         num_shards_per_rank=2,
                     )
                     loader = StatefulDataLoader(

tests/unit_tests/datasets/test_interleaved.py

@@ -308,7 +308,7 @@ def test_metrics_aggregation(
             if "metrics" in sample:
                 collected_metrics.extend(sample["metrics"])
 
-        # Count metrics by dataset name (using new metric key)
+        # Count metrics by dataset name
         ds1_samples_processed = sum(
             1
             for m in collected_metrics