End-to-End ML System Design Questions

Redesign the input path
I would separate expensive preprocessing from the training loop and make the input stream sharded, cached, and deterministic.

Main changes

• Compact many small object-store files into larger shards so reads are sequential instead of chatty.
• Precompute expensive transforms once, then store the transformed output or an intermediate cache.
• Use a manifest that records shard order, sample IDs, and versioned transform code.
• Prefetch data to local disk or memory so GPUs do not wait on network reads.
• Keep lightweight, deterministic transforms in the training worker so results are reproducible.

Worked example
If the dataset is 10,000 files of 2 MB each, that is 20 GB spread across many requests. Repacking into 200 shards of 100 MB keeps the same 20 GB, but cuts file open and listing overhead dramatically.

Why this helps
The GPUs stay busy because decoding and network latency happen ahead of time. Reproducibility comes from versioned manifests and fixed seeds. Debuggability comes from keeping raw sample IDs so a bad example can be replayed exactly.

This is usually the highest leverage change when preprocessing time dominates training time.

Design
I would keep a small pool of warm replicas and use dynamic batching, an autoscaler, and request admission control. The goal is to absorb bursts without letting queueing time blow up p95 latency.

Key components

• API gateway or router that sends requests to healthy replicas
• Inference workers with a batcher that waits only a short time before running a batch
• Autoscaler that looks at queue depth, GPU utilization, and recent latency
• Warm capacity so scale-up is not from zero during traffic spikes
• Priority or deadline-aware routing for urgent requests

How I would control latency
I would set a maximum batch wait time so batching improves throughput without causing long queues. During quiet periods, I would scale down gradually but keep a floor of ready GPUs. During spikes, I would prefer a small amount of extra capacity over missing the p95 target.

Worked example
If a batcher waits up to 5 ms and batches up to 8 requests, it can improve GPU efficiency while still keeping queue delay bounded.

Tradeoff
More batching lowers cost per request, but too much batching hurts tail latency. The right answer is a controlled batch size, warm headroom, and fast scaling signals rather than aggressive scale-to-zero.

Approach
I would make checkpointing a consistent snapshot of everything needed to resume training, not just model weights. That means saving model parameters, optimizer state, learning-rate scheduler state, random number generator seeds, current epoch and batch cursor, and any data sampler state.

Design

• Write checkpoints atomically: first write to temporary storage, validate checksums, then publish a manifest pointer.
• Keep checkpoints incremental when possible, but always make the latest one self-contained for fast restart.
• Trigger an immediate checkpoint on spot interruption notice, then resume from the last good manifest.
• Store code version, config, and data version so the run is reproducible.

Worked example
If I checkpoint every 10 minutes and a preemption happens at minute 37, the restart only redoes at most 7 minutes of work, not the whole job.

Why this works
The manifest guarantees consistency, the RNG and sampler state keep replay deterministic, and the atomic publish prevents half-written checkpoints from being used.

First split the problem into two buckets

1. True optimization issues: the model is mathematically harder to train at larger scale.
2. Distributed systems issues: the parallel run is changing the effective training behavior.

Things I would investigate

• Effective batch size and learning rate scaling. A bigger world size often means a larger batch, which can require LR tuning.
• Data sharding. I would verify that each sample is seen once, not duplicated or skipped.
• Gradient synchronization. I would check that all ranks contribute the same gradients and that there are no stale or dropped updates.
• Precision and numerics. Mixed precision can introduce overflow, underflow, or different rounding behavior at scale.
• Stateful layers like batch normalization, which can behave differently across replicas.

How I would separate them
I would compare small-scale and large-scale runs with the same global batch and fixed seeds, then inspect gradient norms, loss curves, and sample hashes.

Worked example
If 8 workers train with a batch of 256 each, the global batch becomes 2048. If loss gets noisier only after that change, it may be an optimization tuning problem. If the run only degrades when messages cross nodes, it is more likely a communication or synchronization problem.

That distinction tells me whether to tune the optimizer or fix the distributed stack.

Approach
I would trace one request from ingress to response using distributed tracing. A trace is a single request’s timeline, broken into spans, which are timed steps like feature retrieval, preprocessing, model inference, and downstream calls. CPU can look normal because the bottleneck may be waiting on network, locks, queueing, or a slow dependency, not raw compute.

How I would isolate it

1. Start with one trace ID from a slow p99 request.
2. Check span timings for each stage: feature store, preprocessing, batching, model execution, postprocessing, and any outbound calls.
3. Compare slow traces to fast ones. I care about where the extra time accumulates and whether it is wait time or compute time.
4. Add or inspect metrics per stage, such as feature fetch latency, batch queue time, inference time, and downstream timeout rate.
5. If tracing is coarse, I would add sub-spans in the service and temporary structured logs with the trace ID to pinpoint the jump.

Worked example
If median requests are 80 ms and p99 is now 160 ms, I might find: feature retrieval 15 ms, preprocessing 10 ms, batch queue 55 ms, model execution 20 ms, downstream call 45 ms. That tells me the issue is batching or a dependency, not the model itself.

Next checks

• Feature retrieval slow only on cache misses, then look at cache hit rate and feature store latency.
• Batching spikes, then inspect queue depth and batch timeout settings.
• Model execution spikes, then profile the serving container.
• Downstream dependency spikes, then check retries, timeouts, and circuit breakers.

If autoscaling did not trigger, I would also verify the scaling signal. CPU-based scaling often misses latency problems, so I would prefer queue depth, request rate, or custom p95/p99 latency signals for this workload.