Distributed Training: Device Mesh

Split a model across GPUs with tensor, pipeline, and data parallelism plus ZeRO — and see whether it fits in memory.

Model size: 7B params

GPUs: 8

Tensor parallel (TP): 1

Pipeline parallel (PP): 1

Microbatches: 8

ZeRO / FSDP stage: 2

Precision

8-bit optimizer

What changed

Set a model size and a GPU count, then split it across the mesh: tensor parallelism (splits layers), pipeline parallelism (splits the stack), data parallelism (replicates), and ZeRO (shards the replicated states). Watch per-GPU memory and the pipeline bubble respond.

Loading…

Illustrative model for building intuition, not a live system.

Loading explorer…

Read the theory: distributed-training