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

> Easily track metrics with StatsLogger for W&B or SwanLab integration in AReaL. Instantiate StatsLogger, record metrics via StatsTracker API, and export results.

- Repository: [inclusionAI/areal](https://github.com/inclusionai/areal)
- Tags: how-to-guide
- Published: 2026-03-04

---

**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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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:

```python
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:

```python
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:

```python
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`](https://github.com/inclusionai/areal/blob/main/areal/utils/stats_tracker.py), while `StatsLogger` manages external services in [`areal/utils/stats_logger.py`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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`](https://github.com/inclusionai/areal/blob/main/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.