Training Infrastructure Cost Control: Where ML Spend Actually Goes

The teams that are surprised by their training bills are not teams that miscalculated GPU hours. They’re teams that didn’t account for the other 40% of training infrastructure cost.

A training run involves more than the GPU cluster. Storage for training data, checkpoints, and logs. Egress charges when data moves across availability zones. Experiment tracking metadata. The idle time between runs when instances aren’t terminated. And the experiment runs that were started, ran partway, produced bad results, and should have been stopped earlier.

This is the cost audit that most ML teams haven’t done.

The components of training infrastructure cost

Compute (GPU/TPU): The visible cost. Most teams have visibility into this. The less visible version: GPU utilization rate. A training job that saturates the GPU at 90% utilization is using what it pays for. A job that sits at 40% utilization because the data pipeline is the bottleneck is wasting 60 cents of every dollar.

Storage: Training datasets, model checkpoints (every N steps), logs, and artifact storage. For large-model training, checkpoint storage alone can run into hundreds of gigabytes per experiment. Over dozens of experiments, this accumulates.

Data pipeline I/O: Egress charges between storage and compute. In AWS, data movement within a region between services (S3 to EC2) is free; cross-region is not. If your training data is in us-east-1 and your training cluster is in us-west-2, you’re paying for every training step.

Idle compute: GPU instances that aren’t terminated between runs. The experiment ends; the engineer goes to lunch; the instance runs for two hours. At $10/hour for an A100-class instance, this is a $20 oversight per occurrence.

Experiment waste: The wasted cost of experiments that ran too long before a problem was identified. This is the hardest to measure and the most impactful to reduce.

Audit step 1: GPU utilization

Pull GPU utilization metrics for the last 30 days of training runs. (CloudWatch for AWS, Cloud Monitoring for GCP, or any GPU metrics integration.) Calculate average utilization per training job.

If average utilization is below 70%, you have a data pipeline bottleneck. The GPU is ready to process the next batch before the data loader has it available. Solutions:

• Prefetching in the data loader (PyTorch’s num_workers and prefetch_factor)
• Storing training data in a format optimized for random access (LMDB, WebDataset)
• Moving data closer to compute (same region, or instance-local SSD)

Audit step 2: Checkpoint retention policy

Look at your checkpoint storage. How many checkpoints do you retain per run? How many of those checkpoints do you ever actually load?

A reasonable policy:

• Keep the last 3 checkpoints (for crash recovery)
• Keep the best checkpoint by validation metric
• Delete everything else at run completion

Most teams keep everything indefinitely. The storage cost on long training runs is non-trivial.

Automated cleanup:

def cleanup_checkpoints(run_dir: str, keep_last: int = 3):
    checkpoints = sorted(
        glob.glob(f"{run_dir}/checkpoint-*"),
        key=lambda x: int(x.split('-')[-1])
    )
    
    # Always keep the best checkpoint
    best = find_best_checkpoint(run_dir)
    to_delete = [c for c in checkpoints[:-keep_last] if c != best]
    
    for ckpt in to_delete:
        shutil.rmtree(ckpt)
        print(f"Deleted checkpoint: {ckpt}")

Audit step 3: Spot/preemptible instance usage

On-demand GPU instances are expensive. Spot instances (AWS) and preemptible VMs (GCP) are 60-80% cheaper. For training jobs that support checkpointing and restart (which all serious training jobs should), this is not a reliability tradeoff — it’s a correctness tradeoff in your budget allocation.

Implementation requirements:

• Checkpointing at regular intervals (every 1000 steps, or every 20 minutes of wall clock time)
• Automatic resume on restart (detect latest checkpoint on job startup)
• Signal handling for SIGTERM to trigger a final checkpoint before termination

Teams that aren’t using spot instances for interruptible training jobs are paying 3-5x for the same compute.

Audit step 4: Experiment early stopping

The most impactful cost reduction is experiments you stop before wasting their full budget. This requires:

• Evaluation during training (not just at the end)
• Early stopping criteria based on validation metrics
• A culture where stopping an experiment is acceptable, not a failure

A basic early stopping criterion:

class EarlyStopper:
    def __init__(self, patience: int = 5, min_delta: float = 0.001):
        self.patience = patience
        self.min_delta = min_delta
        self.best_metric = float('inf')
        self.counter = 0

    def should_stop(self, validation_loss: float) -> bool:
        if validation_loss < self.best_metric - self.min_delta:
            self.best_metric = validation_loss
            self.counter = 0
            return False
        else:
            self.counter += 1
            return self.counter >= self.patience

Run validation every 500-1000 steps. Plot the training and validation curves. Stop when the curve has plateaued.

Audit step 5: Instance lifecycle management

Automated shutdown of idle instances is not optional — it’s the floor of reasonable cost management.

In practice: configure a watchdog that monitors GPU utilization. If a GPU instance’s utilization drops below 5% for more than 10 minutes and there’s no active job, terminate the instance.

Cloud platforms provide this through auto-scaling groups with scale-in policies. Use them.

What platforms get right (and wrong) about cost

Databricks: Good cost visibility through the cost tracking dashboard, but Databricks clusters are expensive base configurations. Watch the worker node counts.

SageMaker: Managed termination of training jobs on completion. The managed warm pool feature is genuinely useful for reducing cold-start times without idle costs. The cost of the managed orchestration layer itself is non-trivial.

Vertex AI: Good integration with GCP billing, reasonable spot preemptible support for training. Custom training jobs on preemptible accelerators work well.

Self-managed clusters: Maximum cost control at maximum operational overhead. The right choice for teams with dedicated ML infrastructure engineers; wrong for teams that want to focus on models.

The cost optimization conversation is inseparable from the MLOps platform selection question ↗. Platform overhead costs money; operational savings at scale can more than offset it. Model the break-even for your workload.