Revert "Metric Logging updates 6/N - Enable SFT metrics / delete old file" (meta-pytorch#489)

Author: felipemello1

apps/sft/llama3_8b.yaml

@@ -46,7 +46,7 @@ parallelism:
 checkpoint:
   enable: true
   folder: ./checkpoint              # The folder to save checkpoints to.
-  initial_load_path: hf://${model_name}  # The path to load the initial checkpoint from. Ignored if `folder` exists.
+  initial_load_path: hf://${model}  # The path to load the initial checkpoint from. Ignored if `folder` exists.
   initial_load_in_hf: true          # If true, interpret initial_load_path as a HuggingFace model repo
   last_save_in_hf: true
   interval: 500
@@ -56,12 +56,6 @@ 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
-
 # profiling:
 #   enable_profiling: false
 

apps/sft/main.py

@@ -27,7 +27,6 @@
 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 forge.util.config import parse
 
 from monarch.actor import current_rank, current_size, endpoint
@@ -78,6 +77,7 @@ def __init__(self, config: DictConfig):
 
         self.current_step = 0
         self.num_training_steps = job_config.training.steps
+        self.metric_logger = None  # TODO: fix this
         self.gradient_accumulation_steps = 1  # Example value, adjust as needed
         self._rank = current_rank().rank
         self._size = math.prod(current_size().values())
@@ -109,22 +109,9 @@ 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):
-        """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,
@@ -247,9 +234,7 @@ def train_step(self, batch) -> None:
         # ) as grad_acc:
         labels = batch.pop("labels")
         loss = self.forward_backward(batch, labels)
-        loss = loss.item()
 
-        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)
@@ -266,25 +251,14 @@ 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", 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,
@@ -296,23 +270,16 @@ async def train(self) -> None:
     async def cleanup(self) -> None:
         if self.checkpointer:
             self.checkpointer.close()
-        if getattr(self, "mlogger", None):
-            await self.mlogger.shutdown.call_one()
+        if self.metric_logger:
+            self.metric_logger.close()
 
     def __repr__(self) -> str:
         return "Trainer"
 
 
 async def run(cfg: DictConfig) -> None:
-
-    logging.info("Spawning recipe...")
+    logging.info("Spawing recipe...")
     process_cfg = cfg.pop("processes")
-
-    # Initialize metric logger in main process
-    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)
-
     recipe = await ForgeSFTRecipe.options(**process_cfg).as_actor(cfg)
 
     logging.info("Created recipe, running setup.")
@@ -323,7 +290,6 @@ async def run(cfg: DictConfig) -> None:
 
     logging.info("Done training. Clean up")
     await recipe.cleanup.call()
-
     await recipe.mesh.stop()
     logging.info("All done!")
 

apps/sft/qwen3_8b.yaml

@@ -45,7 +45,7 @@ parallelism:
 checkpoint:
   enable: true
   folder: ./checkpoint              # The folder to save checkpoints to.
-  initial_load_path: hf://${model_name}  # The path to load the initial checkpoint from. Ignored if `folder` exists.
+  initial_load_path: hf://${model}  # The path to load the initial checkpoint from. Ignored if `folder` exists.
   initial_load_in_hf: true          # If true, interpret initial_load_path as a HuggingFace model repo
   last_save_in_hf: true
   interval: 500
@@ -55,12 +55,6 @@ 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
-
 # profiling:
 #   enable_profiling: false
 

src/forge/data/__init__.py

@@ -4,12 +4,6 @@
 # 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",
-    "MetricTransform",
-    "DefaultDatasetMetricTransform",
-]
+__all__ = ["collate_packed", "CROSS_ENTROPY_IGNORE_IDX"]

src/forge/data/dataset_metrics/__init__.py

@@ -0,0 +1,39 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .metric_agg_handlers import (
+    AggregationHandler,
+    CategoricalCountAggHandler,
+    MaxAggHandler,
+    MeanAggHandler,
+    MetricState,
+    MinAggHandler,
+    StatsAggHandler,
+    SumAggHandler,
+)
+from .metric_aggregator import MetricsAggregator
+from .metric_transform import (
+    AggregationType,
+    DefaultTrainingMetricTransform,
+    Metric,
+    MetricTransform,
+)
+
+__all__ = [
+    "AggregationType",
+    "AggregationHandler",
+    "CategoricalCountAggHandler",
+    "DefaultTrainingMetricTransform",
+    "StatsAggHandler",
+    "MaxAggHandler",
+    "MeanAggHandler",
+    "Metric",
+    "MetricState",
+    "MetricsAggregator",
+    "MetricTransform",
+    "MinAggHandler",
+    "SumAggHandler",
+]