Setting up metrics tracking with StatsLogger for W&B or SwanLab integration: The AReaL Guide

To set up metrics tracking in AReaL, instantiate StatsLogger with your BaseExperimentConfig, record distributed metrics via the StatsTracker API (stat, scalar, denominator), then call export() and logger.commit() at each step to push reduced results to Weights & Biases, SwanLab, or TensorBoard.

The inclusionai/areal repository provides a production-grade observability stack designed for distributed LLM training. Its architecture cleanly separates distributed metric aggregation from external service integration, enabling seamless setting up metrics tracking with StatsLogger for W&B or SwanLab integration across multi-node GPU clusters.

The Two-Layer Metrics Architecture

AReaL implements a strict separation of concerns between collecting statistics and logging them:

• Layer 1: Distributed Collection — The StatsTracker class in areal/utils/stats_tracker.py (lines 1-144) handles gathering, reducing, and exporting per-step statistics from every rank in a distributed run.
• Layer 2: External Logging — The StatsLogger class in areal/utils/stats_logger.py (lines 1-96) pushes exported numbers to Weights & Biases, SwanLab, and TensorBoard.

This design ensures that distributed reductions happen efficiently across all ranks while external API calls occur only on rank 0, preventing duplicate uploads and network congestion.

Distributed Metric Collection with StatsTracker

Hierarchical Namespaces and Scopes

The tracker uses a hierarchical namespace system to group related metrics. The scope context manager and scope_func_wrapper decorator let you define logical groups like train/ or eval/, which the system joins via _get_full_key (lines 70-74) to create flat keys such as train/loss/avg.

Recording Functions

Three primary methods capture different metric types, each protected by thread-safe locking (see with self.lock blocks at lines 94-144):

• scalar(**kwargs) — Stores plain Python numbers (e.g., learning rate) that are identical across all ranks.
• denominator(**kwargs) — Registers a Boolean tensor (e.g., an attention mask) used as a denominator for weighted averages.
• stat(denominator, reduce_type=None, **kwargs) — Records float tensors and links them to a previously registered denominator for proper averaging.

Cross-Rank Reduction Types

The ReduceType enum (lines 17-23) defines how tensors aggregate across the process group:

• AVG_MIN_MAX — Computes average, minimum, and maximum.
• AVG, SUM, MIN, MAX — Single reduction operations.
• SCALAR — No reduction (used for values already synchronized).

Exporting Aggregated Metrics

The export() method triggers dist.all_gather_object and reduction ops via _aggregate (lines 85-115) to return a flat dict[str, float]. The dictionary contains entries like loss/avg, loss/min, and loss/max based on the requested reduction type.

External Logging Integration with StatsLogger

Configuration and Initialization

StatsLogger receives a BaseExperimentConfig and FinetuneSpec during construction, extracting the nested stats_logger sub-config (self.config = config.stats_logger) in __init__ (lines 20-31). The init() method (lines 34-37) enforces an early exit on all ranks except rank 0, ensuring only the primary process initializes external APIs.

Service Setup

The initialization sequence handles three backends:

• Weights & Biases — Sets WANDB_BASE_URL and WANDB_API_KEY environment variables, calls wandb.login(), and invokes wandb.init(...) (lines 44-75).
• SwanLab — Executes swanlab.login followed by swanlab.init(...) (lines 77-92).
• TensorBoard — Instantiates a SummaryWriter if a log path is provided (lines 94-96).

The Commit Workflow

The commit() method (lines 25-39) executes after each training step:

1. Filters out automatically-generated __count keys (lines 30-31).
2. Prints a formatted table via tabulate_stats.
3. Sends the cleaned dictionary to W&B (wandb.log) and SwanLab (swanlab.log).
4. Writes scalars to TensorBoard via self.summary_writer.add_scalar.

When training ends, close() finalizes the run by calling wandb.finish(), swanlab.finish(), and self.summary_writer.close().

End-to-End Implementation Examples

Recording Training Metrics

Use the tracker API inside your training loop to capture distributed statistics:

from areal.utils.stats_tracker import stat, scalar, denominator, export
import torch.distributed as dist

def train_step(model, batch, current_lr):
    # Record hyperparameters (scalar, same on all ranks)

    scalar(lr=current_lr)
    
    # Register a denominator mask for valid tokens

    denominator(mask=batch["attention_mask"].bool())
    
    # Compute loss and record with reduction

    loss = model(batch).loss
    stat(denominator="mask", reduce_type=None, loss=loss)
    
    # Aggregate across ranks (call only on rank 0 for logging)

    if dist.get_rank() == 0:
        metrics = export()
        # Returns: {'loss/avg': 0.342, 'loss/min': 0.331, 'lr': 5e-5}

        return metrics
    return None

Wiring StatsLogger into Your Trainer

Integrate the logger to automatically push metrics to configured backends:

from areal.utils.stats_logger import StatsLogger
from areal.api.cli_args import BaseExperimentConfig
from areal.api.io_struct import FinetuneSpec

# Configuration loaded from CLI or config files

exp_cfg: BaseExperimentConfig = load_config()
ft_spec: FinetuneSpec = load_spec()

# Initialize logger (only rank 0 will active APIs)

logger = StatsLogger(config=exp_cfg, ft_spec=ft_spec)
logger.init()

for epoch in range(ft_spec.total_train_epochs):
    for step, batch in enumerate(dataloader):
        # ... training logic ...

        step_metrics = export()  # Reduced metrics dict

        global_step = epoch * ft_spec.steps_per_epoch + step
        
        # Push to W&B, SwanLab, and TensorBoard

        logger.commit(epoch, step, global_step, step_metrics)

logger.close()

Direct Logging Without the Distributed Tracker

If you already have computed metrics and want to bypass the tracker:

import wandb
import swanlab

# W&B direct integration

wandb.init(project="my_project", mode="online")
wandb.log({"loss": 0.234, "accuracy": 0.91}, step=global_step)
wandb.finish()

# SwanLab direct integration

swanlab.login()
swanlab.init(project="my_project")
swanlab.log({"loss": 0.234}, step=global_step)
swanlab.finish()

Summary

• AReaL uses a two-layer system: StatsTracker handles distributed aggregation in areal/utils/stats_tracker.py, while StatsLogger manages external services in areal/utils/stats_logger.py.
• Thread-safe recording: Use scalar(), denominator(), and stat() to capture metrics from any rank safely.
• Automatic reduction: The export() method applies ReduceType operations (AVG, MIN, MAX, SUM) across the process group.
• Rank-zero guarding: StatsLogger automatically skips initialization and uploads on ranks > 0 to prevent API conflicts.
• Multi-backend support: Single API calls push to Weights & Biases, SwanLab, and TensorBoard simultaneously based on your StatsLoggerConfig.

Frequently Asked Questions

How does AReaL handle metric aggregation across multiple GPUs?

According to the source code in areal/utils/stats_tracker.py, the export() method uses dist.all_gather_object to collect statistics from every rank, then applies reduction operations (average, min, max, sum) based on the ReduceType specified when recording. This ensures that metrics like loss reflect the global training state rather than individual device values.

What is the difference between StatsTracker and StatsLogger?

StatsTracker (layer 1) is a distributed, thread-safe collector that reduces tensors across process groups and exports flat dictionaries. StatsLogger (layer 2) is a service integration layer that receives those dictionaries and forwards them to W&B, SwanLab, or TensorBoard. The tracker operates on all ranks; the logger initializes and commits only on rank 0.

Can I use StatsLogger with only W&B or only SwanLab?

Yes. The initialization logic in areal/utils/stats_logger.py (lines 44-96) checks for configuration presence before starting each service. If you omit SwanLab credentials in your StatsLoggerConfig, the logger skips swanlab.init() but still initializes W&B if configured, and vice versa.

How do I configure the logging directory for TensorBoard?

Pass a log_dir path in your StatsLoggerConfig. The StatsLogger instantiates a SummaryWriter at lines 94-96 of areal/utils/stats_logger.py. Alternatively, use the get_log_path utility (lines 45-61) which constructs a deterministic path under fileroot/logs/<user>/<experiment>/<trial>/ based on your experiment metadata.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}