Design distributed training for a 10B to 1T parameter model

In this article

1. Before you start
2. Why interviewers ask this
3. Distributed training from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page belongs to the machine learning operations (MLOps) learning path.

Why interviewers ask this

A design question about training a large model is a test of whether a candidate can start from arithmetic, not from product names. A candidate who opens with "I would use DeepSpeed" has named a tool without showing the reason for it. A candidate who opens with "the state is 16 bytes per parameter, so it cannot live on one GPU" has already earned the follow-ups.

The question separates weak, strong and Staff answers by three things: whether the memory is counted before the compute, whether each kind of parallelism is matched to the link it can afford, and whether failure is part of the design rather than an afterthought. It often leads into machine learning system design questions on cost and into serving questions such as the key-value (KV) cache.

Distributed training from first principles

Pipeline parallelism

Pipeline parallelism is a way of splitting a model by depth, so that each GPU or group of GPUs holds a consecutive block of layers and passes activations to the next block.

Model floating-point operations (FLOPs) utilization

Model FLOPs utilization (MFU) is the share of a GPU's theoretical peak arithmetic speed, counted in floating-point operations (FLOPs) per second, that a training run turns into useful model computation after communication and waiting are paid.

In plain words

Think of printing an encyclopedia too big for one press. Give each press a full plan and different pages (data parallelism). Cut each page into strips and give each press a strip (tensor parallelism). Line the presses up, each doing one stage (pipeline parallelism). The analogy breaks in one place: the presses must agree on every correction after each batch, and that agreement, not the printing, costs time at scale.

In one precise sentence: split the model only as much as the training state requires, then spend the remaining GPUs on speed with data parallelism. The design starts at the memory count in Step 1 and ends at the repair plan in Step 6.

Step 1: Count the memory before the compute

Mixed-precision training with the Adam optimizer (which keeps running averages of each weight's gradient and of its square) keeps 16 bytes of state for every parameter. Rajbhandari et al. (2019) break this down as 2 bytes for the 16-bit weights, 2 bytes for the 16-bit gradients and 12 bytes for the optimizer, which keeps a 32-bit copy of the weights plus a momentum and a variance value of 4 bytes each. For their 1.5-billion-parameter example they state a need of at least 24 gigabytes against 3 gigabytes for the 16-bit weights alone.

Apply that to the range in the title. A 10-billion-parameter model has 160 GB of state, which is already twice the 80 GB of one H100 (NVIDIA H100 page, as of October 2026). A 1-trillion-parameter model has 16 TB, which needs at least 200 such GPUs for state alone.¹ An activation is an intermediate value saved during the forward pass for use in the backward pass, and activations come on top. What you now know: the first decision is how to split the state.

Step 2: Shard the state with data parallelism

Fully sharded data parallelism keeps data parallelism's simple batch splitting but stores only a slice of the state on each GPU. The idea is ZeRO, the Zero Redundancy Optimizer, which partitions the optimizer state, the gradients and then the weights across the data-parallel GPUs; its full partitioning cuts the state per GPU from 16Ψ to 16Ψ divided by the number of GPUs, where Ψ is the parameter count (Rajbhandari et al., 2019). PyTorch's version is Fully Sharded Data Parallel (FSDP), which its authors report matches the speed of plain data parallelism while supporting much larger models (Zhao et al., 2023).

For the 10-billion-parameter model, 64 GPUs leave 2.5 GB of state each instead of 160 GB.¹ The price is communication: each GPU must gather the weight slices it does not hold before using them. Input: a model that does not fit. Output: a model that fits, still trained with one simple batch-splitting loop. For many models in the 10-billion range, this step alone is the whole design.

Text description of the diagram

• A 10-billion-parameter model has 160 GB of state at 16 bytes per parameter.
• 64 GPUs leave 2.5 GB of state each instead of 160 GB.

Step 3: Split the layers inside a server with tensor parallelism

When one layer's work is too large or too slow for one GPU, tensor parallelism splits each layer's matrix multiplications across the GPUs of one server. It needs a result exchange inside every layer, so it is only affordable on the fastest links.

Narayanan et al. (2021) state the rule as their first takeaway: use tensor parallelism up to the number of GPUs in a server, then use pipeline parallelism to go beyond. The reason they give is that the all-reduce traffic of tensor parallelism (an exchange after which every GPU holds the sum of all GPUs' values) is slower across servers than on NVLink, the high-bandwidth link inside a server, which NVIDIA lists at 900 GB/s on the H100 SXM.

That is why large runs use a tensor-parallel size of 8 on servers of 8 GPUs: Megatron's trillion-parameter run and Meta's Llama 3 405B run both do (Meta Llama Team, 2024, Table 4). Input: a layer split into 8 slices. Output: the same layer result, computed with 8 GPUs and many fast exchanges.

Step 4: Split depth across servers with pipeline parallelism

Pipeline parallelism assigns consecutive blocks of layers to different servers and streams microbatches, small slices of the batch, through them. It sends only the activations at block borders, a small amount of traffic, so it tolerates slower inter-server links (Huang et al., 2018). Its cost is idle time. At the end of every batch the pipeline must flush. The pipeline bubble is the extra idle time per batch, and Narayanan et al. give it as a fraction of the ideal compute time: (p − 1) divided by m, where p is the number of pipeline stages and m the number of microbatches.

With 16 stages and 16 microbatches the bubble is 93.8% of the ideal compute time, which means the stages sit idle 48.4% of each step; with 128 microbatches the bubble is 11.7% of the ideal time.¹ The paper's interleaved schedule, which gives each GPU several small blocks instead of one large block, divides the bubble by the number of blocks per GPU at the price of that many times more communication. Meta's Llama 3 405B run uses 16 pipeline stages and the interleaved schedule. Input: a model deeper than one server's memory. Output: stages with a calculable bubble.

Text description of the diagram

• With 16 stages and 16 microbatches the bubble is 93.8% of the ideal compute time.
• With 128 microbatches the bubble is 11.7%.

Step 5: Estimate the time from 6ND and utilization

Training time follows from the arithmetic: a dense transformer, one that uses all its weights on every token, needs about 6 floating-point operations (FLOPs) per parameter per training token, so total work is 6 × N × D for N parameters and D tokens. The Llama 3 paper gives the check: 405 billion parameters and 15.6 trillion tokens come to 3.8 × 10^25 FLOPs in their report, and the formula gives 3.79 × 10^25.¹ The number of tokens grows with model size; Hoffmann et al. (2022) find that for each doubling of model size the training tokens should also double.

Divide the work by what the cluster really delivers, not its peak. NVIDIA lists 1,979 teraFLOPS of 16-bit tensor-core speed with sparsity (a 2:4 pattern of zeroed weights that doubles the listed speed) for the H100 SXM; the dense figure is half of that, 989.5.¹ Llama 3 reports an MFU of 38% to 43%, so 40% is a defensible planning figure. What you now know: time is a division, and the divisor is mostly utilization, which depends on Steps 2 to 4.

Step 6: Plan for failure with checkpoints

At thousands of GPUs something fails every few hours, so the design must resume from saved state. Meta reports 466 job interruptions in a 54-day period of Llama 3 405B pre-training, 419 of them unexpected and about 78% traced to confirmed or suspected hardware faults, with manual intervention needed only three times (Meta Llama Team, 2024). That is one unexpected stop about every 3.1 hours.¹

A checkpoint is a saved copy of the full training state. Its interval sets the work lost per failure: on average half the interval. The Megatron paper reports a 13.8 TB checkpoint for its trillion-parameter model, written at 273 GB/s, which is about 50 seconds if the storage keeps that speed.¹ Input: a run that will be interrupted. Output: a run that loses minutes, not days.

1 · Count memory2 · Shard state3 · Split layers4 · Split depth5 · Estimate time6 · CheckpointReplay

Step-by-step description

1. 16 bytes of state per parameter.
2. ZeRO or FSDP split the state.
3. Tensor parallelism inside a server.
4. Pipeline stages across servers.
5. 6 x N x D over utilization.
6. Resume after failures.

• Mixed-precision Adam keeps 16 bytes of state for every parameter.
• Tensor parallelism stays inside a server and pipeline parallelism goes across servers.
• Training work is about 6 times parameters times tokens.

Worked example

We sized two ends of the range with the figures above.¹ The 10-billion-parameter case assumes 1 trillion training tokens and 256 H100 GPUs at 40% MFU; both are illustrative choices, not a recommendation. The 1-trillion-parameter case assumes 15 trillion tokens and 16,384 GPUs at 40% MFU, also illustrative.

The layout for the largest case copies the published one: Megatron ran a 1-trillion-parameter configuration on 3,072 A100 GPUs (tensor 8, pipeline 64, data 6, since 3,072 ÷ 512 = 6) at 163 teraFLOP/s per GPU, 52% of the A100's 312 teraFLOP/s peak, and estimated 84 days for 450 billion tokens (Narayanan et al., 2021, Table 1). Those were A100s, an older GPU than the H100 used in the other figures here, so the 163 is not comparable with the H100 peak of 989.5.

Text description of the diagram

• At 7.76 unexpected stops a day, a 30-minute checkpoint interval loses about 8.1% of the run to repeated work.
• A 120-minute interval loses 32.3%.

The failure row shows why Step 6 is not optional. With about 7.8 unexpected stops a day and a checkpoint every 30 minutes (illustrative), each stop loses 15 minutes on average, so the run loses about 8% of its time (chart above) to repeated work before any restart delay.¹ A 120-minute interval would lose 32.3%. Halving the interval halves the loss but doubles the checkpoint traffic, which is the trade-off the interviewer will probe.

Trade-offs and edge cases

Each kind of parallelism spends a different resource. Data parallelism with sharding spends network bandwidth to save memory; tensor parallelism spends the fastest links to save per-layer memory; pipeline parallelism spends idle time to save memory across servers. A design is a choice of which resource is cheapest on your cluster.

Three edge cases matter. First, slower networks push the design toward pipeline stages over sharding, since pipelines send less traffic. Second, the budget of a long context changes memory: activations grow with sequence length, and Llama 3's long-context stage runs with a context-parallel size of 16 (one long sequence split across 16 GPUs) on top of the other three splits (Meta Llama Team, 2024, Table 4).

Third, utilization falls when the batch per data-parallel group shrinks: Llama 3's MFU fell from 43% to 41% when it doubled the GPUs at a fixed global batch.¹ Reported practice also differs on whether to shard weights for the whole run: Meta's Llama 3 run shards optimizer states and gradients but does not re-shard the weights after the forward pass, to avoid an extra gather.

How to answer it in an interview

Say it in this order: memory count, shard first, split by link speed, time estimate, failure plan.

Common mistakes

• "Plain data parallelism is enough, I will just add GPUs" (wrong). Every GPU would hold the full 160 GB of state for a 10-billion-parameter model. (Right: "state is 16 bytes per parameter, so I shard it or split the model.")
• "Tensor parallelism across all 3,072 GPUs" (wrong). Its exchange in every layer would run over slower inter-server links. (Right: "tensor parallelism inside a server, pipeline parallelism across servers.")
• "Time is parameters divided by GPU speed" (wrong). It omits tokens and utilization, the two numbers that move the answer by orders of magnitude. (Right: "six times parameters times tokens, divided by peak times MFU.")
• "I would use DeepSpeed" as the whole answer (wrong). It names a tool and skips the reason. (Right: "ZeRO shards the state; here is the memory it saves.")
• "Failures are rare, so I will checkpoint at the end" (wrong). Meta reports one unexpected stop about every three hours at 16,000 GPUs. (Right: "checkpoint on a schedule I can justify from the failure rate.")

Follow-up questions

Do you know DeepSpeed and Megatron?

DeepSpeed is the library that implements ZeRO, which shards optimizer state, gradients and weights; Megatron is NVIDIA's library that combines tensor, pipeline and data parallelism. Say what each splits, not just the name.

What is a characteristic of distributed training in AI/ML clusters?

It is communication-bound. GPUs synchronise gradients or gather weight slices every step, so network bandwidth and the slowest GPU (a straggler) set the speed more than raw arithmetic does. That is why the design matches each split to a link: on a slow inter-server network, prefer pipeline stages over deeper sharding, since pipelines send only small activations at block borders.

How do you estimate how long training will take?

Work is about 6 × parameters × tokens FLOPs; divide by the number of GPUs times their peak speed times the utilization, 38% to 43% in Meta's Llama 3 run. For a different token count, rerun the division, as the 10B and 1T columns above do.

How often do you checkpoint and why?

Pick the interval from the failure rate: each stop loses on average half an interval. At Meta's rate of one unexpected stop every 3.1 hours, a 30-minute interval loses about 8% of the run and a 10-minute interval about 2.7%, against more checkpoint traffic. When a GPU fails, the scheduler replaces the node and the job restarts from the last checkpoint; Meta handled 419 unexpected interruptions in 54 days with manual help only three times.

What does FSDP shard and when does it gather?

FSDP shards the weights, gradients and optimizer state across the data-parallel GPUs. Before a layer runs, each GPU gathers that layer's full weights from the others, uses them, then frees them; gradients are reduced and re-sharded after the backward pass.

Does the same plan work for fine-tuning the model afterwards?

Usually a smaller one does: fine-tuning needs the same 16 bytes per trained parameter unless you train only a small set of added weights, so the memory count in Step 1 still decides it. For the serving side, see how memory limits batch size there.

For a neighbouring design question from the same learning path, see design a feature store.

Sources

1. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models — arXiv (Rajbhandari et al.), 4 Oct 2019
2. Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM — arXiv (Narayanan et al.), SC 2021, v5 revised 2021-08-23, 9 Apr 2021
3. The Llama 3 Herd of Models — arXiv (Llama Team, Meta), v3 revised 2024-11-23, 31 Jul 2024
4. PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel — arXiv (Zhao et al.), v2 revised 2023-09-12, 21 Apr 2023
5. GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism — arXiv (Huang et al.), 16 Nov 2018
6. Training Compute-Optimal Large Language Models — arXiv (Hoffmann et al.), 29 Mar 2022
7. NVIDIA H100 Tensor Core GPU — NVIDIA (date = read date), 24 Oct 2026