Add Sample-level Logging API

apps/grpo/main.py

@@ -29,7 +29,7 @@
 from forge.controller.provisioner import init_provisioner, shutdown
 from forge.data.rewards import MathReward, ThinkingReward
 from forge.observability.metric_actors import get_or_create_metric_logger
-from forge.observability.metrics import record_metric, Reduce
+from forge.observability.metrics import record_episode_sample, record_metric, Reduce
 from forge.observability.perf_tracker import Tracer
 
 from forge.types import LauncherConfig, ProvisionerConfig
@@ -55,6 +55,7 @@ class Episode:
     response_tokens: list[int] | None = None
     ref_logprobs: torch.Tensor | None = None
     reward: float | None = None
+    reward_breakdown: dict[str, float] | None = None
     advantage: float | None = None
 
     @property
@@ -169,8 +170,11 @@ class RewardActor(ForgeActor):
     reward_functions: list[Callable]
 
     @endpoint
-    async def evaluate_response(self, prompt: str, response: str, target: str) -> float:
+    async def evaluate_response(
+        self, prompt: str, response: str, target: str
+    ) -> dict[str, float]:
         total_rewards = 0.0
+        reward_breakdown = {}  # reward breakdown by function
         for reward_fn in self.reward_functions:
             reward = reward_fn(prompt, response, target)
             total_rewards += reward
@@ -179,6 +183,7 @@ async def evaluate_response(self, prompt: str, response: str, target: str) -> fl
             reward_fn_name = getattr(
                 reward_fn, "__name__", reward_fn.__class__.__name__
             )
+            reward_breakdown[reward_fn_name] = reward
             # per function reward
             record_metric(
                 f"reward/evaluate_response/sum_{reward_fn_name}_reward",
@@ -211,7 +216,8 @@ async def evaluate_response(self, prompt: str, response: str, target: str) -> fl
             )
 
         avg_reward = total_rewards / len(self.reward_functions)
-        return avg_reward
+        reward_breakdown["reward"] = avg_reward
+        return reward_breakdown
 
 
 @dataclass
@@ -321,7 +327,7 @@ async def main(cfg: DictConfig):
             )
         )
     metric_logging_cfg = cfg.get("metric_logging", {"console": {"log_per_rank": False}})
-    mlogger = await get_or_create_metric_logger()
+    mlogger = await get_or_create_metric_logger(process_name="Controller")
     await mlogger.init_backends.call_one(metric_logging_cfg)
     await ts.initialize(strategy=ts.ControllerStorageVolumes())
 
@@ -403,9 +409,10 @@ async def continuous_rollouts():
                 episode.response = response.text
                 input_ids[i, :max_req_tokens] = episode.request_tensor
                 input_ids[i, max_req_tokens:] = episode.response_tensor
-                episode.reward = await reward_actor.evaluate_response.route(
+                episode.reward_breakdown = await reward_actor.evaluate_response.route(
                     prompt=prompt, response=response.text, target=target
                 )
+                episode.reward = episode.reward_breakdown["reward"]
 
             t.step("reward_evaluation")
 
@@ -424,6 +431,7 @@ async def continuous_rollouts():
             for episode, advantage in zip(group.episodes, advantages):
                 episode.advantage = advantage
                 await replay_buffer.add.call_one(episode)
+                record_episode_sample("rollout/sample", episode)
 
             # Log metrics
             rollout_count += 1

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:

src/forge/controller/provisioner.py

@@ -21,6 +21,7 @@
 from forge.controller.launcher import BaseLauncher, get_launcher
 
 from forge.observability.metric_actors import get_or_create_metric_logger
+from forge.observability.utils import detect_actor_name_from_call_stack
 
 from forge.types import ProcessConfig, ProvisionerConfig
 
@@ -259,8 +260,10 @@ def bootstrap(env: dict[str, str]):
 
             self._proc_host_map[procs] = host_mesh
 
-        # Spawn local logging actor on each process and register with global logger
-        _ = await get_or_create_metric_logger(procs)
+        # Detect actor name and spawn local logging actor on each process
+        process_name = detect_actor_name_from_call_stack()
+        _ = await get_or_create_metric_logger(procs, process_name=process_name)
+
         return procs
 
     async def host_mesh_from_proc(self, proc_mesh: ProcMesh):

src/forge/env_constants.py

@@ -11,7 +11,7 @@
 
 # Force all timing methods in forge.observability.perf_tracker.py to use
 # CPU timer if False or GPU timer if True. If unset, defaults to the assigned value to the function.
-METRIC_TIMER_USES_CUDA = "METRIC_TIMER_USES_CUDA"
+METRIC_TIMER_USES_GPU = "METRIC_TIMER_USES_GPU"
 
 # Makes forge.observability.metrics.record_metric a no-op
 FORGE_DISABLE_METRICS = "FORGE_DISABLE_METRICS"

src/forge/observability/__init__.py

@@ -10,39 +10,47 @@
     LocalFetcherActor,
 )
 from .metrics import (
+    BackendRole,
     ConsoleBackend,
-    # Utility functions
     get_actor_name_with_rank,
     get_logger_backend_class,
-    # Backend classes
     LoggerBackend,
+    LoggingMode,
     MaxAccumulator,
     MeanAccumulator,
-    # Accumulator classes
+    Metric,
     MetricAccumulator,
     MetricCollector,
     MinAccumulator,
+    record_episode_sample,
     record_metric,
     Reduce,
     reduce_metrics_states,
+    SampleAccumulator,
     StdAccumulator,
     SumAccumulator,
+    TopBottomKFilter,
     WandbBackend,
 )
 from .perf_tracker import trace, Tracer
 
 __all__ = [
     # Main API functions
     "record_metric",
+    "record_episode_sample",
     "reduce_metrics_states",
     "get_actor_name_with_rank",
     "get_logger_backend_class",
     "get_or_create_metric_logger",
     # Performance tracking
     "Tracer",
     "trace",
+    # Data classes
+    "Metric",
+    "BackendRole",
     # Enums
     "Reduce",
+    "LoggingMode",
     # Actor classes
     "GlobalLoggingActor",
     "LocalFetcherActor",
@@ -59,4 +67,7 @@
     "MaxAccumulator",
     "MinAccumulator",
     "StdAccumulator",
+    "SampleAccumulator",
+    # Filter classes
+    "TopBottomKFilter",
 ]