MLflow for GPU Training Ops, in depth: server topology, rank-0 logging, metric volume, system metrics, checkpoints and the model registry

The MLflow data model, and what each piece costs

MLflow stores four kinds of thing, and they have very different cost profiles. Knowing which is which explains every operational rule later in the article.

Above these sit runs (one training attempt), experiments (a folder of runs), and, in MLflow 3, logged models: first-class model objects that can be logged with name= rather than the deprecated artifact_path= argument. The model registry then gives a logged model a registered name, numbered versions and mutable aliases such as champion.

The important line in that table is metrics. A run that logs 30 keys every step for 200,000 steps writes six million rows. Multiply by every rank that also logs, and by every sweep, and the backend database becomes the bottleneck of your whole platform. Params are cheap but immutable: logging the same key with a different value fails, which matters when a resumed job tries to re-log its config.

Server topology for GPU teams

The local default (an mlruns directory next to the script) is fine for a laptop and wrong for a cluster: nodes cannot see each other's disks, and a pod restart deletes the history. Run a shared tracking server instead, with three separable parts.

• Tracking server. A stateless REST process. Run at least two replicas behind a load balancer; if it is a single pod, a node drain takes down logging for every job at once.
• Backend store. A real database such as PostgreSQL, not SQLite on a shared volume. Give it connection pooling, backups and a retention plan for the metrics table.
• Artifact store. Object storage. With artifact proxying (the --artifacts-destination server option) clients upload through the server and need no bucket credentials; the cost is that every byte crosses the server. For anything large, that is the wrong path.

mlflow server \
  --backend-store-uri postgresql://mlflow:${DB_PASS}@pg.internal:5432/mlflow \
  --artifacts-destination s3://ml-artifacts/mlflow \
  --host 0.0.0.0 --port 5000

Training nodes only need MLFLOW_TRACKING_URI pointing at the load balancer. Put it in the job template, not in the training script, so the same code runs locally and on the cluster.

Rule one: exactly one rank owns the run

Launchers such as torchrun start one process per GPU and set RANK, LOCAL_RANK and WORLD_SIZE in each environment. If the training script calls mlflow.start_run() unconditionally, a 64-GPU job creates 64 runs with identical params and slightly different losses, and nobody can tell which one is the experiment.

In data-parallel training the loss on rank 0 is representative once you all-reduce it, so rank 0 should own the run and every other rank should make no MLflow calls. If another rank needs the run id (for example to tag its own checkpoint shards), broadcast it rather than letting it search.

import os, mlflow, torch, torch.distributed as dist

def setup_tracking(cfg):
    rank = int(os.environ["RANK"])
    run_id = [None]
    if rank == 0:
        mlflow.set_experiment(cfg["experiment"])
        run = mlflow.start_run(run_id=cfg.get("resume_run_id"),
                               run_name=cfg["run_name"])
        if cfg.get("resume_run_id") is None:          # params are immutable: log once
            mlflow.log_params(cfg["hparams"])
        mlflow.set_tags({
            "git_sha": cfg["git_sha"], "image": cfg["image_digest"],
            "world_size": os.environ["WORLD_SIZE"],
            "torch": torch.__version__, "cuda": torch.version.cuda,
            "nccl": ".".join(map(str, torch.cuda.nccl.version())),
        })
        run_id = [run.info.run_id]
    dist.broadcast_object_list(run_id, src=0)       # every rank learns the id
    return rank, run_id[0]

The tags are the part people skip and later regret. A loss curve that changed after an image rebuild is only debuggable if the run records the image digest, the CUDA and NCCL versions and the world size. Log the dataset manifest hash too; a run without its data version is a number without units.

Rule two: metric logging must never stall the GPU

Two separate costs hide inside a single mlflow.log_metric("loss", loss.item(), step=step). The first is .item(): copying a CUDA tensor to a Python float forces the host to wait until every kernel queued before it has finished. That destroys the overlap between the CPU enqueueing step N+1 and the GPU executing step N. The second is the HTTP call itself, which by default blocks until the server has committed the row.

Work through the numbers. Suppose a step takes 400 ms of GPU time, the training loop logs 12 keys with one call each, and a round trip to the tracking server costs 15 ms. That is 180 ms of blocking per step, so the GPUs are idle for roughly a third of the wall clock, on every rank that waits on rank 0 at the next collective. The fix has three parts.

1. Accumulate on the device. Keep running sums as tensors and call .item() only every N steps, once, on a stacked tensor.
2. Batch the call. mlflow.log_metrics(dict, step=step) sends many keys in one request.
3. Go asynchronous. mlflow.log_metric(..., synchronous=False) queues the write; mlflow.config.enable_async_logging() or MLFLOW_ENABLE_ASYNC_LOGGING=true turns it on for all fluent calls. It does not affect MlflowClient, so code that logs through the client stays synchronous.

LOG_EVERY = 50
acc = torch.zeros(3, device="cuda")      # loss, grad_norm, tokens

for step, batch in enumerate(loader, start=start_step):
    loss, grad_norm, ntok = train_step(batch)
    acc += torch.stack([loss.detach(), grad_norm.detach(), ntok.float()])
    if step % LOG_EVERY == 0:
        dist.all_reduce(acc)                       # every rank participates
        if rank == 0:
            loss_sum, gn_sum, tok = (acc / world_size).tolist()   # one sync
            mlflow.log_metrics({
                "train/loss": loss_sum / LOG_EVERY,
                "train/grad_norm": gn_sum / LOG_EVERY,
                "train/tokens_per_s": tok * world_size / elapsed(),
                "train/lr": scheduler.get_last_lr()[0],
            }, step=step, synchronous=False)
        acc.zero_()

Always pass the global optimizer step as step. MLflow accepts out-of-order and repeated steps, which is exactly what you want after a resume, and charts plot by step rather than by arrival time.

System metrics versus DCGM

MLflow can sample hardware itself. Enable it with mlflow.start_run(log_system_metrics=True), mlflow.enable_system_metrics_logging() or MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING=true; it needs psutil and, for NVIDIA GPUs, nvidia-ml-py. It samples every 10 seconds by default (MLFLOW_SYSTEM_METRICS_SAMPLING_INTERVAL) and writes keys under the system/ prefix: GPU utilization, memory used and its percentage, and power. The exact key names, including how devices are numbered, depend on your MLflow release, so read them from a test run before building dashboards on them.

This is useful and limited. It sees one host. If several nodes enable it against the same run, their samples land under the same keys and interleave into nonsense, so either enable it on rank 0 only, or give each node's local rank 0 its own child run (start_run(nested=True) needs the parent active in the same process, so on other nodes create the child explicitly with the mlflow.parentRunId tag). And NVML's utilization counter only says a kernel was running, not that the tensor cores were busy.

For fleet health use the tool built for it, NVIDIA DCGM, which exports SM activity, tensor-pipe activity, XID errors and NVLink counters per GPU into your metrics stack. The practical split: MLflow holds what explains a run (loss, throughput, a coarse utilization trace); DCGM and Prometheus hold what explains a machine.

Checkpoints are not artifacts

The tempting line is mlflow.log_artifacts("ckpt/step_20000"). For a 70B model with sharded optimizer state that directory is several hundred gigabytes, written by every rank at once. Routed through a proxying tracking server it saturates the server and the node NIC that NCCL also needs, and the job blocks while it uploads.

Instead let each rank write its shard straight to object storage with the distributed checkpoint API your framework provides (see PyTorch FSDP, in depth for sharded state), then have rank 0 record only the pointer and a small manifest.

ckpt_uri = f"s3://ckpts/{run_id}/step_{step:07d}"
save_sharded_checkpoint(model, optim, ckpt_uri)          # every rank, direct to storage
dist.barrier()
if rank == 0:
    mlflow.set_tag("last_checkpoint", ckpt_uri)
    mlflow.log_dict({"step": step, "uri": ckpt_uri, "world_size": world_size,
                     "tokens_seen": tokens_seen}, f"checkpoints/step_{step:07d}.json")

Log a model, in the MLflow sense, only for weights you intend to evaluate or serve: a consolidated final checkpoint, or a small LoRA adapter. That keeps the artifact store full of things people will load, and keeps the registry pointing at them.

Preemption and resume

Long GPU jobs are interrupted: spot reclaim, a node failure, a scheduler preemption. The question is whether a restarted job continues the same run or starts a new one. Continuing is usually right, because the loss curve should be one line. Store the run id somewhere the restarted job can find it (the checkpoint manifest is ideal), then call mlflow.start_run(run_id=...) and skip log_params, since params cannot change.

Expect a short overlap: the job resumes from step 20,000 but had logged up to step 20,300 before it died. Because metrics are keyed by step, the chart shows a small double-drawn segment, which is honest. Add a resumes tag counter so reviewers know the run was interrupted. Runs killed without a clean exit stay in RUNNING forever; a nightly sweeper that marks stale runs FAILED keeps the UI truthful.

The registry: aliases, not stages

Once an evaluation job scores a logged model, register it and move an alias. Serving and batch jobs load models:/llm-sft@champion and never hard-code a version number, so promotion and rollback are a single metadata write.

from mlflow import MlflowClient
client = MlflowClient()
mv = mlflow.register_model(f"runs:/{eval_run_id}/model", "llm-sft")
client.set_registered_model_alias("llm-sft", "candidate", mv.version)
# after the canary passes:
client.set_registered_model_alias("llm-sft", "champion", mv.version)

Gate the alias move on evidence recorded in MLflow itself: the evaluation run's metrics, the training run's data manifest tag and the checkpoint URI. Anyone can then answer "what is serving, trained on what" from one page. For the cross-team policy around that gate see the MLOps article linked above; for the end-to-end fine-tuning runbook, Fine-Tuning Ops on GPU.

Failure modes

What to do next

1. Stand up a shared tracking server with a PostgreSQL backend, object storage for artifacts and two replicas.
2. Set MLFLOW_TRACKING_URI in your job template, not your script.
3. Gate every MLflow call on rank 0 and broadcast the run id.
4. Log git SHA, image digest, CUDA, NCCL, torch version, world size and data manifest hash as tags.
5. Replace per-step logging with device-side accumulation and a batched, asynchronous call every N steps; measure tokens per second before and after.
6. Enable system metrics on one rank, read the real key names, and keep fleet telemetry in DCGM.
7. Write checkpoints directly to storage and log only their URI and manifest.
8. Store the run id in the checkpoint manifest and make resume reuse it.
9. Register evaluated models and serve by alias.