Design a Large-Scale Model Training Pipeline (Foundation Model)

Problem statement

Design a platform to train a foundation model with tens to hundreds of billions of parameters on petabyte-scale heterogeneous data. The platform must ingest raw corpora, clean and tokenize data, create deterministic shuffled training shards, schedule thousands of GPUs, run data-parallel and model-parallel training, checkpoint weights and optimizer state, evaluate on validation sets, recover from partial node failures, and optimize expensive accelerator utilization.

This is not a simple ML pipeline. A training run is a long-running HPC workload with tightly coupled collective communication. One unhealthy GPU, one slow NIC, one corrupted shard, or one checkpoint-storage stall can waste hours of cluster time. The design must separate four planes: data plane, compute plane, control plane, and artifact plane. The data plane turns raw documents into immutable token shards. The compute plane runs the synchronous training loop. The control plane owns experiment metadata, job lifecycle, scheduling, and recovery. The artifact plane owns checkpoints, eval results, manifests, and model releases.

Why the problem is distinctive

Web systems usually tolerate partial failure by retrying idempotent requests. Distributed training is different: step progress is collective. If one rank crashes, the communicator often cannot continue unless the job is restarted from a checkpoint. The architecture therefore needs frequent low-overhead checkpoints, fast node replacement, deterministic data ordering, and strict versioning of code, config, dataset, and model graph.

The problem requires petabyte ingestion, distributed training with model parallelism, checkpointing and validation, and HPC scheduling. The correct answer should go beyond saying Kubernetes and PyTorch. It should explain token budgets, 3D or 4D parallelism, checkpoint sharding, straggler detection, spot-aware scheduling, and recovery semantics.