wandb

name	wandb
description	Track experiments, log metrics and artifacts, and compare runs with Weights & Biases
version	1.0.0
author	kai-agent
metadata	{"kai":{"tags":["kai","compute","wandb","experiment-tracking","metrics","artifacts","ml"]}}

Weights & Biases (W&B)

Skill for experiment tracking, metric logging, artifact management, and run comparison using Weights & Biases.

Authentication

This skill uses the WANDB_API_KEY environment variable. If it is not set, ask your admin to add it in Agent Settings. You can generate a key at https://wandb.ai/authorize.

Install

pip install wandb

Initialize a run

import wandb

run = wandb.init(
    project="inference-optimization",
    name="vllm-llama3-benchmark",
    config={
        "model": "meta-llama/Llama-3.1-8B-Instruct",
        "quantization": "awq",
        "tensor_parallel": 2,
        "max_batch_size": 64,
    },
)

Set entity (team/org)

run = wandb.init(
    entity="your-team",
    project="inference-optimization",
    name="run-name",
)

Log metrics

Basic metric logging

for step in range(100):
    wandb.log({
        "throughput_tokens_per_sec": throughput,
        "latency_p50_ms": p50,
        "latency_p99_ms": p99,
        "gpu_utilization": gpu_util,
        "memory_used_gb": mem_gb,
    })

Log with explicit step

wandb.log({"loss": 0.42, "accuracy": 0.91}, step=500)

Log summary metrics (final values)

run.summary["best_throughput"] = 1250.0
run.summary["best_latency_p99"] = 45.2

Log tables

Useful for comparing model outputs or benchmark results:

table = wandb.Table(columns=["model", "quantization", "throughput", "latency_p99"])
table.add_data("Llama-3.1-8B", "fp16", 850, 62.1)
table.add_data("Llama-3.1-8B", "awq-int4", 1340, 38.5)
table.add_data("Llama-3.1-8B", "gptq-int4", 1280, 41.2)
wandb.log({"benchmark_results": table})

Log artifacts

Upload a model checkpoint

artifact = wandb.Artifact(
    name="llama3-optimized",
    type="model",
    description="AWQ quantized Llama 3.1 8B with custom calibration",
)
artifact.add_dir("./model-checkpoint/")
run.log_artifact(artifact)

Upload a dataset

artifact = wandb.Artifact(name="calibration-data", type="dataset")
artifact.add_file("./calibration_prompts.jsonl")
run.log_artifact(artifact)

Download an artifact

api = wandb.Api()
artifact = api.artifact("your-team/inference-optimization/llama3-optimized:latest")
local_path = artifact.download(root="./downloaded-model/")

Query and compare runs

List runs in a project

api = wandb.Api()
runs = api.runs(
    path="your-team/inference-optimization",
    filters={"config.quantization": "awq"},
    order="-summary_metrics.best_throughput",
)

for run in runs:
    print(f"{run.name}: throughput={run.summary.get('best_throughput')}")

Get a specific run

run = api.run("your-team/inference-optimization/run-id-here")
print(run.config)
print(run.summary)

Export run history as DataFrame

history_df = run.history()
print(history_df[["_step", "throughput_tokens_per_sec", "latency_p99_ms"]])

Compare runs programmatically

runs = api.runs("your-team/inference-optimization")
for run in runs:
    print(
        f"{run.name} | "
        f"quant={run.config.get('quantization')} | "
        f"throughput={run.summary.get('best_throughput', 'N/A')} | "
        f"p99={run.summary.get('best_latency_p99', 'N/A')}"
    )

Alerts

wandb.alert(
    title="Latency regression detected",
    text=f"p99 latency increased to {p99}ms, exceeding 50ms threshold",
    level=wandb.AlertLevel.WARN,
)

Sweeps (hyperparameter search)

Define a sweep config

sweep_config = {
    "method": "bayes",
    "metric": {"name": "throughput_tokens_per_sec", "goal": "maximize"},
    "parameters": {
        "max_batch_size": {"values": [16, 32, 64, 128]},
        "gpu_memory_utilization": {"min": 0.7, "max": 0.95},
        "quantization": {"values": ["awq", "gptq", "fp16"]},
    },
}
sweep_id = wandb.sweep(sweep_config, project="inference-optimization")

Run the sweep agent

def train():
    run = wandb.init()
    config = wandb.config
    # Run your benchmark with config params...
    wandb.log({"throughput_tokens_per_sec": result_throughput})

wandb.agent(sweep_id, function=train, count=20)

Finishing a run

Always finish the run when done to ensure all data is flushed:

wandb.finish()

Or use a context manager:

with wandb.init(project="inference-optimization") as run:
    wandb.log({"metric": value})
# Automatically finished

Inference optimization patterns

1. Benchmark logging: Log throughput, latency percentiles, GPU utilization, and memory for every serving configuration.
2. Artifact versioning: Version quantized model checkpoints so you can trace which calibration data produced the best results.
3. Sweep over serving params: Use sweeps to find optimal max_batch_size, gpu_memory_utilization, and max_num_seqs for vLLM/TGI.
4. Tables for A/B comparison: Log output quality tables alongside performance metrics to catch quality regressions from quantization.
5. Alerts for production: Set up alerts for latency or throughput regressions in production inference pipelines.

Rules

1. ALWAYS call wandb.finish() or use a context manager to avoid orphaned runs
2. Use meaningful run names that encode the key experiment variables
3. Log config at init time so runs are filterable and comparable
4. Never log secrets or API keys as config or metrics
5. Use artifacts for large files, not wandb.log with file paths