AI System Scalability Questions

Covers designing and operating machine learning systems to handle growth in data volume, model complexity, and traffic. Topics include distributed training strategies such as data parallelism, model parallelism, and pipeline parallelism; coordination and orchestration approaches like parameter servers, gradient aggregation, and framework tools such as PyTorch distributed, Horovod, and TensorFlow strategies; data pipeline and I O considerations including sharding, efficient formats, preprocessing bottlenecks, streaming and batch ingestion; serving and inference scaling including model sharding, batching for throughput, autoscaling, request routing, caching, and latency versus throughput tradeoffs. Also includes monitoring, profiling, checkpointing and recovery, reproducibility, cost and resource optimization, and common bottleneck analysis across network, storage, CPU preprocessing, and accelerator utilization.

HardSystem Design

31 practiced

Design an end-to-end ML platform to support training and serving a 1B-parameter multimodal model (image+text), with weekly offline retraining and inference at 10k QPS with P95 under 200ms. Describe data storage and layout, distributed training strategy including model sharding and optimizer state, orchestration for jobs, serving topology including regional caches, monitoring, and rollback strategy.

Sample Answer

Requirements & constraints:- 1B-parameter multimodal (image+text), weekly offline retrain, 10k QPS inference, P95 < 200ms, global users.

Data storage & layout:- Raw ingestion: S3-compatible object store (immutable, partitioned by date/source) for images and raw text; message queue (Kafka) for real-time events.- Processed datasets: Parquet on object store with partition keys {dataset, year, month, shard_id}; store image features in a columnar store (e.g., Feather/Parquet) or a feature-store (Feast) with versioning.- Metadata: Catalog (Glue/Datacatalog) + Postgres for dataset lineage, schemas, and dataset versions. Keep training manifests (lists of URIs + labels + provenance).- Artifacts: Model checkpoints, tokenizer, precomputed embeddings in artifact registry (MLflow/Weights & Biases), tagged by training run.

Distributed training strategy:- Use data-parallel + model-sharding hybrid (ZeRO-style): sharded optimizer/state (ZeRO-3) to fit 1B parameters across GPU fleet. Use mixed precision (AMP) and gradient accumulation to keep batch sizes stable.- Topology: Multi-node with NCCL for all-reduce and NVLink where available. Checkpointing writes sharded checkpoints to object store with manifest for quick restore.- Optimizer state: Partition optimizer and gradient states across workers; periodically consolidate small metadata to store for fast resume.- Training pipeline: Preprocessing on CPU nodes producing TFRecords/Parquet shards; worker nodes stream mini-batches.

Orchestration for jobs:- Use Kubernetes with a custom controller (or Kubeflow/Tekton) to define pipelines: preprocessing → training → evaluation → validation → deployment.- Each training job declares resource profile (GPUs, mem, node labels). Autoscaler scales GPU node pool based on queued jobs.- CI gating: automated eval suite checks fairness, accuracy, A/B metrics, resource usage; approve manual/automated promotion.

Serving topology:- Prediction service implemented in two tiers: 1) Regional GPU-backed model servers (gRPC/HTTP) hosting sharded inference (tensor model parallelism) with batching and async pipelines. Leverage Triton/Custom fast runtime; use model parallel inference so each instance holds model shard(s). 2) Global edge-tier caches (Redis/MemoryDB + precomputed embeddings) per region to answer frequent queries or partial workloads (text-only or cached image-text pairs).- Front door: Global load balancer routes to nearest region; regional API gateway routes to regional caches first, then to GPU inference pool.- Autoscaling: prediction replicas scale by request rate and tail latency; request batching increases throughput while maintaining P95 <200ms (adaptive batching with deadline-aware policies).

Monitoring & observability:- Metrics: latency histograms, QPS, GPU utilization, batch sizes, queue lengths. Model metrics: accuracy, drift (embedding distance), feature distribution skew, calibration.- Logging/tracing: distributed traces for request path; synthetic canaries against gold dataset every deploy.- Alerts: SLA breach, data drift threshold, GPU OOMs, training divergence.- Dashboards: real-time for SRE and model owner; automated weekly reports comparing new model vs production on holdout slices.

Rollback & safety:- Canary deploy: route small % to new model, run shadow traffic and compare business metrics and latencies. Automated rollback on metric regressions.- Fast rollback: keep last N model versions online; use traffic routing to switch to previous model instantaneously. Maintain compatibility layer for tokenizer/schema changes.- Safe checkpoints: immutable sharded checkpoints + manifest allow redeploy training snapshot if retrain fails.- Post-deploy kill-switch: immediate circuit-breaker to divert all traffic to stable region cache or previous model.

Trade-offs & reasoning:- ZeRO-style sharding reduces memory but complicates checkpoint/restore; acceptable for 1B scale for cost efficiency.- Regional caches reduce latency for P95 but add consistency complexity—use TTL and versioned cache keys.- Weekly offline retrain balances data freshness and stability; fast rollback + canaries mitigate risk.

This design supports scalable training, low-latency global serving, strong observability, and safe rollouts appropriate for a data-scientist-led ML platform.

EasyTechnical

34 practiced

What is Horovod and when would you choose it over framework-native distributed training such as PyTorch DDP or TensorFlow MirroredStrategy? Mention its dependencies (e.g., NCCL, MPI) and typical production use cases.

HardTechnical

27 practiced

Your distributed training yields inconsistent evaluation metrics across runs due to dataset skew and sampling differences. Propose a concrete plan to ensure statistical correctness including deterministic sharding, seed management, stratified sampling, use of warm-up epochs, and validating evaluation with confidence intervals and variance estimates.

Sample Answer

Plan (goal: make evaluation statistically correct, reproducible, and robust to dataset skew)

1) Clarify requirements & metrics- Define exact eval metric(s) and acceptable CI width / significance level (e.g., 95% CI ±1%).

2) Deterministic sharding & sampling- Implement deterministic sharding by hashing stable sample ID (not file order) mod world_size, so each process gets consistent partitions across runs.- Use epoch-aware shuffling with a fixed global seed: shuffle_index = RNG(seed + epoch) so shuffles are repeatable.

Example (PyTorch-style):

python

# per-worker deterministic shard and seed
global_seed = 12345
rank, world_size = get_dist_info()
ids = all_ids[hash_fn(all_ids) % world_size == rank]  # deterministic shard by id hash
rng = np.random.RandomState(global_seed + epoch)
rng.shuffle(ids)

3) Seed management- Centralize seed handling: set seeds for Python, NumPy, random, and framework RNGs on all workers.- Make sure any non-deterministic ops (cuDNN, async data loaders) are controlled or logged.- Persist seed and metadata with each run.

4) Stratified sampling for skew- If labels or strata are imbalanced, perform stratified sharding so each worker's shard preserves population proportions (or use oversampling/cost-sensitive losses for training).- For time/temporal skew, shard by contiguous time windows but rotate evaluation windows across runs.

5) Warm-up epochs & burn-in- Use warm-up epochs to let distributed sampling buffers and batchnorm statistics stabilize. Exclude first N epochs from final evaluation aggregation.- For BN in distributed env, use synchronized BN or accumulate running stats before eval.

6) Statistical validation- Compute per-run metric and its standard error. For binary/continuous metrics: - SE = sqrt(p(1-p)/n) for proportions; use delta method for derived metrics (AUC).- Use bootstrap or repeated subsampling (B=1000) to get empirical 95% CI and variance estimate.- When comparing runs, run paired tests (e.g., paired bootstrap) to assess significance.

7) Monitoring & checks- Log per-shard counts, class distributions, seeds, shard hashes, and RNG state.- Run deterministic reproducibility checks: rerun with same seed should match exactly.- Run sensitivity experiments varying seed and report metric variance; if too large, increase eval sample size or use stratification.

8) Operationalize- Automate the pipeline to record metadata, CI, and variance in experiment reports.- Require minimum effective sample size for evaluation; if below threshold, aggregate more data.

Why this works:- Deterministic sharding + centralized seed makes runs reproducible.- Stratification reduces variance from skew.- Warm-up stabilizes model state across distributed workers.- Bootstrap/CIs provide formal uncertainty quantification so differences are interpretable.

MediumTechnical

30 practiced

Your organization runs hundreds of hyperparameter tuning jobs on large models and costs are high. Propose a pragmatic strategy to reduce cost while preserving model quality, covering algorithmic, infrastructure, and tooling techniques such as early stopping, multi-fidelity search, asynchronous Hyperband, and transfer learning.

Sample Answer

Approach: treat this as an optimization + operational problem — reduce wasted compute while keeping or improving model quality. I'll propose a layered, pragmatic strategy spanning algorithmic, infrastructure, and tooling changes, with measurable rollout steps.

1) Algorithmic techniques- Multi-fidelity search: run cheap, low-fidelity evaluations first (smaller dataset subset, fewer epochs, lower model width). Promote promising configs to higher fidelity. This preserves signal while cutting average job cost.- Asynchronous Hyperband (ASHA): use aggressive early-stopping with bracketed resource allocation to kill poor trials quickly and scale promising ones without synchronization delays.- Bayesian optimization warm-started with prior experiments: use surrogate models to focus search where gains are likeliest; seed with transfer-learning priors from related tasks/models.

2) Transfer learning & meta-learning- Reuse pretrained weights and only tune small subset (heads or adapters) to reduce training time per trial.- Maintain a meta-dataset of past tuning runs (hyperparams → outcome) to initialize BO priors and prune improbable regions.

3) Infrastructure & orchestration- Spot/Preemptible instances for short-lived trials; checkpoint frequently and use elastic autoscaling.- Containerize trials and use a scheduler that supports early preemption and reuse workers (e.g., Ray Tune, KubeFlow Katib).- Cache datasets and precomputed features on fast storage to avoid repeated I/O cost.

4) Tooling and process- Implement metrics and dashboards: cost-per-improvement, time-to-best, and cumulative compute wasted. Use these to set RL targets (e.g., 3x reduction in compute per experiment).- Enforce CI/Gating: run low-cost sanity checks (quick smoke runs) before full tuning.- Incremental rollout: enable multi-fidelity + ASHA for 20% of jobs, measure quality drift, then expand.

Example rollout & expected impact- Phase 1 (30 days): enable dataset subsampling + ASHA on non-critical jobs → expect 4–8x fewer GPU-hours per job for low-quality trials.- Phase 2: integrate BO with historical priors + transfer learning → focus search, reduce number of full-train runs by 50%.- Phase 3: infra optimizations (spot instances, caching) → additional 20–40% cost reduction.

Trade-offs and monitoring- Risk: multi-fidelity may mis-rank some configs. Mitigate by conservative promotion thresholds and validation on full fidelity before deployment.- Always monitor final-model quality vs baseline (A/B tests or validation metrics). Use rollback thresholds.

This combined approach (ASHA + multi-fidelity + BO with priors + transfer learning + infra savings) typically cuts tuning costs substantially while preserving or improving final model quality.

MediumSystem Design

25 practiced

Design a checkpointing policy for long-running training jobs on preemptible or spot instances that minimizes expected recompute time while avoiding excessive checkpoint overhead. Include storage strategy, asynchronous uploads, partial checkpoints, and how you would coordinate checkpoints across a multi-node job.

Sample Answer

Requirements:- Minimize expected recompute when spot/preemptible nodes die- Keep checkpoint overhead low (CPU, I/O, network)- Support multi-node distributed training and fast recovery- Durable storage (S3-compatible) and tolerance for partial failures

High-level design:1. Checkpoint frequency via cost/benefit: estimate mean time to preemption (MTTP) and time-to-checkpoint (TTC). Trigger checkpoint when expected lost work (time since last checkpoint) × preemption hazard > TTC + upload overhead. Practical rule: periodic light checkpoints every N minutes + final full checkpoint on epoch end.

2. Storage strategy:- Local ephemeral disk for transient checkpoint write (fast).- Asynchronous, incremental uploads to S3 (or MinIO) so training isn't blocked.- Use content-addressable blobs and manifest metadata per checkpoint (small JSON): model weights, optimizer state, RNG seed, training step, timestamp.

3. Partial / incremental checkpoints:- Save deltas for model weights (e.g., only changed tensors) or use sharded checkpoints per layer to reduce upload size.- For optimizers or large embeddings, snapshot only frequently-changing parts; use checkpoint versioning and garbage collection.

4. Asynchronous upload and consistency:- Worker writes checkpoint to local disk, marks it as READY locally, then an uploader thread/process streams parts to S3 in the background; upload returns a manifest with part URIs and checksums.- Use atomic manifest publish: only after all parts uploaded, write manifest to durable storage (S3 key /checkpoints/job_id/latest.json). Recovery reads manifests only.

5. Multi-node coordination:- Leader election (rank 0) coordinates global checkpoint windows. Each worker writes its shard locally and starts its background upload. Leader waits for worker ACKs of manifest upload (or times out).- Use staggered checkpoint start times across workers to smooth I/O peaks (e.g., ring order with small delays).- Allow partial recovery: if some worker shards are missing, use last consistent global checkpoint where all shard manifests are present. For large embedding tables, allow independent reload and recovery of partial shards; job restarts with reduced parallelism until shards restored.

6. Recovery:- On preemption, orchestration restarts nodes and reads latest consistent manifest. Download required shards in parallel; leader re-initializes global state (optimizer, RNG).- For very frequent preemptions, fall back to warm-start from smaller metadata (step-only resume) and recompute small amount.

Trade-offs and notes:- More frequent checkpoints reduce recompute but increase upload cost; MTTP-driven scheduling balances this.- Incremental/partial checkpoints complicate logic but considerably cut network and storage costs for large models.- Use checksums and manifest atomicity to prevent corrupt restores.- Leverage framework hooks (PyTorch DDP / TensorFlow CheckpointManager) to integrate seamlessly.

Unlock Full Question Bank

Get access to hundreds of AI System Scalability interview questions and detailed answers.

Join thousands of developers preparing for their dream job.