Metric Logging updates 5/N - enable streaming (#363)

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

Author: felipemello1

.meta/mast/qwen3_14b_mast.yaml

@@ -19,9 +19,9 @@ metric_logging:
   wandb:
     project: "grpo-training"
     group: "grpo_exp_${oc.env:USER}"
-    reduce_across_ranks: True
+    logging_mode: global_reduce
   console:
-    reduce_across_ranks: True
+    logging_mode: global_reduce
 
 # Dataset configuration
 dataset:

.meta/mast/qwen3_1_7b_mast.yaml

@@ -19,9 +19,9 @@ metric_logging:
   wandb:
     project: "grpo-training"
     group: "grpo_exp_${oc.env:USER}"
-    reduce_across_ranks: True
+    logging_mode: global_reduce
   console:
-    reduce_across_ranks: True
+    logging_mode: global_reduce
 
 # Dataset configuration
 dataset:

.meta/mast/qwen3_32b_mast.yaml

@@ -19,9 +19,9 @@ metric_logging:
   wandb:
     project: "grpo-training"
     group: "grpo_exp_${oc.env:USER}"
-    reduce_across_ranks: True
+    logging_mode: global_reduce
   console:
-    reduce_across_ranks: True
+    logging_mode: global_reduce
 
 # Dataset configuration
 dataset:

.meta/mast/qwen3_4b_mast.yaml

@@ -19,9 +19,9 @@ metric_logging:
   wandb:
     project: "grpo-training"
     group: "grpo_exp_${oc.env:USER}"
-    reduce_across_ranks: True
+    logging_mode: global_reduce
   console:
-    reduce_across_ranks: True
+    logging_mode: global_reduce
 
 # Dataset configuration
 dataset:

.meta/mast/qwen3_8b_mast.yaml

@@ -19,9 +19,9 @@ metric_logging:
   wandb:
     project: "grpo-training"
     group: "grpo_exp_${oc.env:USER}"
-    reduce_across_ranks: True
+    logging_mode: global_reduce
   console:
-    reduce_across_ranks: True
+    logging_mode: global_reduce
 
 # Dataset configuration
 dataset:

apps/grpo/main.py

@@ -304,7 +304,7 @@ async def main(cfg: DictConfig):
     else:
         provisioner = await init_provisioner()
 
-    metric_logging_cfg = cfg.get("metric_logging", {"console": {"log_per_rank": False}})
+    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)
 

apps/grpo/qwen3_1_7b.yaml

@@ -16,11 +16,11 @@ 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
   console:
-    reduce_across_ranks: True
+    logging_mode: global_reduce
 
 # Dataset configuration
 dataset:

apps/grpo/qwen3_32b.yaml

@@ -19,11 +19,11 @@ rollout_threads: 32 # make this 4x the number of policy replicas seems to work w
 # 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
   console:
-    reduce_across_ranks: True
+    logging_mode: global_reduce
 
 # Dataset configuration
 dataset:

apps/grpo/qwen3_8b.yaml

@@ -12,11 +12,11 @@ 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
   console:
-    reduce_across_ranks: True
+    logging_mode: global_reduce
 
 # Dataset configuration
 dataset:

src/forge/observability/__init__.py

@@ -12,7 +12,9 @@
 from .metrics import (
     BackendRole,
     ConsoleBackend,
+    get_logger_backend_class,
     LoggerBackend,
+    LoggingMode,
     MaxAccumulator,
     MeanAccumulator,
     Metric,
@@ -43,6 +45,7 @@
     "BackendRole",
     # Enums
     "Reduce",
+    "LoggingMode",
     # Utility functions
     "get_proc_name_with_rank",
     # Actor classes