Experiment Tracking

Experiment tracking records every training run — hyperparameters, metrics, artifacts — enabling comparison, reproducibility, and debugging. Without it, you’re guessing which combination of learning rate, data mix, and model variant produced your best checkpoint.

---

1. What to Track

Category	Examples	Why
Hyperparameters	LR, batch size, model size, optimizer, warmup steps	Reproduce and compare runs
Metrics	Loss, eval scores (MMLU, HumanEval), MFU, gradient norm	Track progress and quality
Artifacts	Checkpoints, configs, eval outputs, tokenizer	Recover and inspect
System metrics	GPU utilization, memory, throughput, communication time	Identify bottlenecks
Data config	Dataset version, mixing ratios, tokenizer version	Full reproducibility
Code version	Git SHA, branch, diff	Link results to exact code

---

2. Tool Comparison

Tool	Hosted?	Open source?	Strengths	Limitations
Weights & Biases	Yes (cloud)	No	Best UX, real-time dashboards, team collaboration	Vendor lock-in, cost at scale
MLflow	Self-hosted or Databricks	Yes	Open ecosystem, model registry built-in	UI less polished, setup overhead
Neptune	Yes (cloud)	No	Excellent experiment management	Smaller ecosystem
TensorBoard	Self-hosted	Yes	Quick visualization, PyTorch native	No experiment comparison, no artifact management
Aim	Self-hosted	Yes	Fast queries over many runs	Smaller community

---

3. Weights & Biases Integration

import wandb

# Initialize run with full config
wandb.init(
    project="llama-3-finetune",
    name="sft-v3-lr3e5-bs128",
    config={
        "model": "meta-llama/Llama-3.1-70B",
        "learning_rate": 3e-5,
        "batch_size": 128,
        "optimizer": "adamw",
        "warmup_steps": 200,
        "max_steps": 5000,
        "dataset": "internal/sft-v3",
        "dataset_version": "sha256:abc123",
        "quantization": "bf16",
        "tp_size": 4,
        "dp_size": 8,
    },
    tags=["sft", "70b", "production-candidate"],
)

# Log metrics every step
for step in range(max_steps):
    loss = train_step()
    
    wandb.log({
        "train/loss": loss,
        "train/learning_rate": scheduler.get_lr(),
        "train/gradient_norm": grad_norm,
        "system/gpu_utilization": get_gpu_util(),
        "system/gpu_memory_pct": get_gpu_mem_pct(),
        "system/throughput_tokens_per_sec": throughput,
    }, step=step)
    
    # Log eval metrics periodically
    if step % eval_interval == 0:
        eval_results = run_eval(model, eval_set)
        wandb.log({
            "eval/mmlu": eval_results["mmlu"],
            "eval/humaneval": eval_results["humaneval"],
            "eval/loss": eval_results["loss"],
        }, step=step)

# Save artifacts
wandb.save("configs/training_config.yaml")
artifact = wandb.Artifact("model-checkpoint", type="model")
artifact.add_dir("checkpoints/step-5000/")
wandb.log_artifact(artifact)

wandb.finish()

---

4. MLflow Integration

import mlflow

mlflow.set_tracking_uri("http://mlflow-server:5000")
mlflow.set_experiment("llama-3-finetune")

with mlflow.start_run(run_name="sft-v3-lr3e5") as run:
    # Log parameters
    mlflow.log_params({
        "model": "meta-llama/Llama-3.1-70B",
        "learning_rate": 3e-5,
        "batch_size": 128,
        "dataset_version": "sha256:abc123",
    })
    
    # Log metrics
    for step in range(max_steps):
        loss = train_step()
        mlflow.log_metric("train_loss", loss, step=step)
    
    # Log eval results
    mlflow.log_metrics({
        "eval_mmlu": 0.82,
        "eval_humaneval": 0.67,
    })
    
    # Log model artifact
    mlflow.log_artifacts("checkpoints/step-5000/", artifact_path="model")
    
    # Register model if quality meets threshold
    if eval_results["mmlu"] > 0.80:
        mlflow.register_model(
            f"runs:/{run.info.run_id}/model",
            "llama-3-70b-sft"
        )

---

5. Comparing Runs

The real value: comparing dozens of experiments to find what matters:

# W&B: create comparison table
api = wandb.Api()
runs = api.runs("team/llama-3-finetune", 
                filters={"tags": "sft", "state": "finished"})

comparison = []
for run in runs:
    comparison.append({
        "name": run.name,
        "lr": run.config["learning_rate"],
        "batch_size": run.config["batch_size"],
        "dataset": run.config["dataset"],
        "final_loss": run.summary.get("train/loss"),
        "mmlu": run.summary.get("eval/mmlu"),
        "humaneval": run.summary.get("eval/humaneval"),
        "gpu_hours": run.summary.get("_wandb/runtime") / 3600,
    })

# Sort by MMLU score to find best run
best = sorted(comparison, key=lambda x: x.get("mmlu", 0), reverse=True)

What to Compare

Comparison	What it reveals
Learning rate sweep (fixed data)	Optimal LR for this model/data
Data mix variants (fixed LR)	Which data composition works best
Model size comparison	Scaling law validation
Checkpoint progression	When quality plateaus (early stopping)
System metrics across configs	Which setup has best MFU

---

6. Best Practices

Practice	Why it matters
Log everything automatically	Don’t rely on remembering to add metrics
Tag runs meaningfully	sft-v3-70b-production not run_042
Link to Git SHA	Exact code reproducibility
Save eval outputs (not just scores)	Debug quality issues
Set up alerts	Notify on loss spike, NaN gradient
Archive artifacts for important runs	Checkpoints, configs, eval outputs
Use consistent metric names	eval/mmlu not sometimes mmlu_score

For managing trained models after experiments, see Model Registry.

---

Self-Check

• List the 6 categories of experiment metadata to track
• Integrate W&B or MLflow into a training loop with metrics and artifacts
• Compare 5+ training runs to identify the best hyperparameter configuration
• Set up automated alerts for training anomalies (loss spikes, NaN)
• Link experiment runs to Git commits for full reproducibility

---

References

• Weights & Biases Documentation
• MLflow Documentation
• Aim: Open-Source Experiment Tracking