Metric Logging updates 6/N - Enable SFT metrics / delete old file

apps/grpo/qwen3_1_7b.yaml

@@ -16,11 +16,12 @@ rollout_threads: 1   # Recommended to set equal to policy.num_replicas
 # Observability configuration
 metric_logging:
   wandb:
-    project: "grpo-training"
-    group: "grpo_exp_${oc.env:USER}"
-    reduce_across_ranks: True
+    project: grpo-training
+    group: grpo_exp_${oc.env:USER}
+    logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce
+    per_rank_share_run: False
   console:
-    reduce_across_ranks: True
+    logging_mode: global_reduce
 
 # Dataset configuration
 dataset:

apps/grpo/qwen3_32b.yaml

@@ -19,11 +19,12 @@ rollout_threads: 1   # Recommended to set equal to policy.num_replicas
 # Observability configuration
 metric_logging:
   wandb:
-    project: "grpo-training"
-    group: "grpo_exp_${oc.env:USER}"
-    reduce_across_ranks: True
+    project: grpo-training
+    group: grpo_exp_${oc.env:USER}
+    logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce
+    per_rank_share_run: False
   console:
-    reduce_across_ranks: True
+    logging_mode: global_reduce
 
 # Dataset configuration
 dataset:

apps/grpo/qwen3_8b.yaml

@@ -12,11 +12,12 @@ off_by_n: 1 # Off by one by default
 # Observability configuration
 metric_logging:
   wandb:
-    project: "grpo-training"
-    group: "grpo_exp_${oc.env:USER}"
-    reduce_across_ranks: True
+    project: grpo-training
+    group: grpo_exp_${oc.env:USER}
+    logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce
+    per_rank_share_run: False
   console:
-    reduce_across_ranks: True
+    logging_mode: global_reduce
 
 # Dataset configuration
 dataset:

apps/sft_v2/llama3_8b.yaml

@@ -55,6 +55,13 @@ activation_checkpoint:
   mode: selective
   selective_ac_option: op
 
+metric_logging:
+  wandb:
+    project: sft-training
+    group: sft_exp_${oc.env:USER}
+    logging_mode: global_reduce # global_reduce, per_rank_reduce, per_rank_no_reduce
+    per_rank_share_run: False
+
 # profiling:
 #   enable_profiling: false
 

apps/sft_v2/main.py

@@ -16,6 +16,7 @@
 import math
 import os
 import sys
+import warnings
 from functools import partial
 from typing import Any
 
@@ -28,6 +29,7 @@
 from forge.data.datasets.packed import PackedDataset, TextPacker
 from forge.data.datasets.sft_dataset import AlpacaToMessages, sft_iterable_dataset
 from forge.data.tokenizer import HuggingFaceModelTokenizer
+from forge.observability import get_or_create_metric_logger, record_metric, Reduce
 
 from monarch.actor import current_rank, current_size, endpoint
 from omegaconf import DictConfig, OmegaConf
@@ -109,9 +111,20 @@ def _init_dist(self):
         os.environ.update(env)
         logger.info("env: {}".format(env))
 
+    async def setup_metric_logger(self):
+        """Initialization happens in the main process. Here we just retrieve it"""
+        mlogger = await get_or_create_metric_logger()
+        return mlogger
+
+    def record_batch_metrics(self, data_metrics: list):
+        """Record dataset metrics using the observability system."""
+        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,
@@ -235,6 +248,7 @@ def train_step(self, batch) -> None:
         labels = batch.pop("labels")
         loss = self.forward_backward(batch, labels)
 
+        record_metric("ForgeSFTRecipe/train_step/loss", loss, Reduce.MEAN)
         logger.info(f"{self.current_step} / {self.num_training_steps}|Loss: {loss}")
         # self.pbar.set_description(f"{self.current_step}|Loss: {loss}")
         # self.pbar.update(1)
@@ -251,14 +265,27 @@ async def train(self) -> None:
 
         while self.current_step < self.num_training_steps:
             batch = next(dataloader)
+
+            # 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
+            )
+
             # Move tensors to the appropriate device
             for k, v in batch.items():
                 if isinstance(v, torch.Tensor):
                     batch[k] = v.to("cuda")  # TODO: hardcoded for now
+
             self.train_step(batch)
             # self.profiler.step()
             self.current_step += 1
 
+            # Flush metrics
+            if self._rank == 0:
+                logger.debug(f"Flushing metrics at step {self.current_step}")
+                await self.mlogger.flush.call_one(global_step=self.current_step)
+
             self.checkpointer.save(
                 curr_step=self.current_step,
                 last_step=self.current_step == self.num_training_steps,
@@ -270,16 +297,35 @@ async def train(self) -> None:
     async def cleanup(self) -> None:
         if self.checkpointer:
             self.checkpointer.close()
-        if self.metric_logger:
-            self.metric_logger.close()
+        if hasattr(self, "mlogger") and self.mlogger:
+            await self.mlogger.shutdown.call_one()
 
     def __repr__(self) -> str:
         return "Trainer"
 
 
 async def run(cfg: DictConfig) -> None:
-    logging.info("Spawing recipe...")
+
+    # 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"}}
+    )
+    mlogger = await get_or_create_metric_logger(process_name="Controller")
+    await mlogger.init_backends.call_one(metric_logging_cfg)
+
     recipe = await ForgeSFTRecipe.options(**process_cfg).as_actor(cfg)
 
     logging.info("Created recipe, running setup.")
@@ -290,6 +336,9 @@ 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/__init__.py

@@ -4,6 +4,12 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 from .collate import collate_packed
+from .metric_transform import DefaultDatasetMetricTransform, MetricTransform
 from .utils import CROSS_ENTROPY_IGNORE_IDX
 
-__all__ = ["collate_packed", "CROSS_ENTROPY_IGNORE_IDX"]
+__all__ = [
+    "collate_packed",
+    "CROSS_ENTROPY_IGNORE_IDX",
+    "MetricTransform",
+    "DefaultDatasetMetricTransform",
+]