metric logger - backend (1/N)

[felipemello1]

Implemented it on a toy script. Will add real metrics to grpo main.py.

To log from anywhere:

from forge.observability.metrics import record_metric, ReductionType
await record_metric("generate/tokens", value, ReductionType.SUM)

Testing

• 2 actors, each with 2 replicas and 2 ranks.
• Used 2 backends at the same time (wandb and console).
• Also tried all 3 distributed methods suggested by wandb.
• Logged 2 metrics, one with sum reduction and another with mean reduction.
• API and abstractions can probably can use some polishing, but they should be enough to work for different user needs.

TLDR 3 ways of doing wandb distributed training

https://docs.wandb.ai/guides/track/log/distributed-training/

a) Track a single process:

• only rank 0 instantiates wandb and logs.
• User can manually reduce and send to rank 0.

Pros:

• Customizations, e.g. std, count, filtering, if we have 2 dp, 8tp, sum_num_tokens should be done for dp only.
• Easy single run to visualize

Cons:

• Needs cross rank reduction logic

b) Track multiple processes:

• Every rank instantiates wandb.
• Runs can be grouped in UI by group name, e.g. “my_group_name”, defined during instantiation.

Pros:

• Maximum visibility per rank

Cons:

• Crowded: 32 ranks = 32 wandb runs
• Reduction is somewhat manual (per graph config - but perhaps we can set via code?) and limited (mean, max, min)

c) Track all processes to a single wandb run ID (experimental):

• Every rank instantiates wandb.
• Primary rank generates a run_id and sends it to other ranks.

pros:

• Not sure

Cons:

• Can't use same x-axis
• Doesn't guarantee logging order

Requirements:

1) Every process should log to the same train step

for train_step in range(2):
	for generation_step in range(3):
		time_generation_step = generation_actor.choose()
		logger.log(time_generation_step, train_step)
	time_train_step = train_actor.choose()
	logger.log(time_train_step, train_step)

For each train_step we have 3 generation steps. They should all belong to the same train_step x-axis.

2) We need a local reduction per train step

If we want the mean time_generation_step per train_step, we have to reduce it ourselves. wandb only keeps the latests

3) We need a global and a local logger

Local loggers are responsible for:
i) per-rank wandb instantiation;
ii) collecting local data.

Global loggers are responsible for communicating:
i) train step + flush command: necessary for all 3 wandb modes;
ii) shared wandb run_id: necessary for option C;
iii) optionally collecting and reducing to rank 0: Necessary for option A

Q1: Can we get rid of local_logger and just use global_logger for everything?
A1: Probably not. For options (b,c) we need to instantiate per rank. Perhaps we could instantiate N instances in a single rank, but this doesn’t make much sense to me. For option (a), we could use just a global logger, but it seems easier to manage it per rank.

4) User must have easy access to a dashboard with a general view of the system

This means that we should NOT require internal tools or external infra setup for basic metric logging. WandB seems to be the best fit for this.

Recommendation

We should support (a) Track a single process with wandb and (b) Track multiple processes with jsonl dump per rank.

Design

These 3 steps are independent:
a) collecting metrics
b) reducing metrics (we can decide to reduce, log per rank, log only on rank 0)
c) flushing metrics to wandb, scuba, jsonl, etc

Controller has a GlobalLoggingActor, every rank has a LocalFetcherActor + MetricLoggerSingleton

a. When a service is spawned, we also spawn a LocalFetcherActor per rank.
b. After every training step, a GlobalLoggingActor calls every LocalFetcherActor which calls MetricLoggerSingleton.

  ┌─────────────────────────────────────────────────────┐
  │                   Controller Process                │
  │ ┌───────────────────────────────────────────────┐   │
  │ │              GlobalLoggingActor               │   │
  │ │  • Coordinates all local loggers              │   │
  │ │  • Applies cross-process reductions           │   │
  │ │  • Dispatches to backends                     │   │
  │ └───────────────────────────────────────────────┘   │
  └─────────────────────────────────────────────────────┘
                           │
            ┌──────────────┼──────────────┐
            │                             │
     ┌──────▼───────┐               ┌─────▼───────┐
     │  Process A   │               │  Process B  │
     │ ┌───────────┐│               │┌───────────┐│
     │ │ Local     ││               ││ Local     ││
     │ │ Fetcher   ││               ││ Fetcher   ││
     │ │ Actor     ││               ││ Actor     ││
     │ └───────────┘│               │└───────────┘│
     │ ┌───────────┐│               │┌───────────┐│
     │ │ Metric    ││               ││ Metric    ││
     │ │ Logger    ││               ││ Logger    ││
     │ │ Singleton ││               ││ Singleton ││
     │ └───────────┘│               │└───────────┘│
     │ ┌───────────┐│               │┌───────────┐│
     │ │ Training  ││               ││ Policy    ││
     │ │ Actors    ││               ││ Actors    ││
     │ └───────────┘│               │└───────────┘│
     └──────────────┘               └─────────────┘

Why?

Each rank can just call the singleton anywhere. This is easy and clean, but a bit counterintuitive at first, i.e. "push metrics to where?" The singleton buffers per-key/train_step, applies local reduction on fetch call.

import push_metrics

def push_metrics(key, value, reduction_type):
	metric_logger = MetricLoggerSingleton()
	metric_logger.append(key, value, reduction_type)

def my_random_actor_or_function():
	push_metrics('key', 1, ReductionTypes.SUM)

Other alternative is to instantiate a class at every actor and pass it around. But I am not sure how to fetch it from GlobalLoggingActor.

def my_random_actor_or_function(my_metric_logger):
	my_metric_logger.push_metrics('key', 1, ReductionTypes.SUM)

Q1 : Why not just use the LocalFetcherActor and drop the Singleton?
A1: Because i couldn't find a way to find the LocalFetcherActor for each rank from the context

Q2: And how is the GlobalMetricLogger aware of LocalFetcherActor and LocalLoggingSingleton?
A2: in get_proc_mesh

# For each rank
local_fetcher_actor = await procs.spawn("local_fetcher_actor", LocalFetcherActor)

# Register with global logger
global_logger = await get_or_spawn_controller("global_logger", GlobalLoggingActor)
await global_logger.register_fetcher.call_one(local_fetcher_actor, process_name)

Then in the main loop:

for train_step in range(3):
	....
	await global_logger.flush.call_one(train_step)

[felipemello1]

tempted to move all backend related things to logger_backends.py

[allenwang28]

looks great!