🎓 A1Learner
🔍

PyTorch Model Parallel and Sharding

Model parallel training splits one model across devices instead of copying it. Sharding splits a stored thing, like the weights or the optimizer state, into pieces held by different devices.

The bill for one copy

One million fp32 weights cost 4 megabytes. The shard count divides that bill.

params = 1_000_000
print("Weights in MB:", params * 4 / 1e6)
print("Per shard of 4 in MB:", params * 4 / 4 / 1e6)

What each scheme splits

Data parallel splits the batch and copies the model. These three schemes split the model or its stored state instead.

schemes = {
    "tensor parallel": "one layer's maths",
    "pipeline parallel": "the layers into stages",
    "ZeRO-style sharding": "weights, gradients, optimizer states",
}
for name, split in schemes.items():
    print(name, "splits", split)

Example

params = 500_000_000
bytes_fp32 = 4
weights = params * bytes_fp32
grads = params * bytes_fp32
adam_states = params * bytes_fp32 * 2
total = weights + grads + adam_states
shards = 4
print("Model parameters: 500000000")
print(f"Weights (fp32): {weights / 1e9:.1f} GB")
print(f"Gradients (fp32): {grads / 1e9:.1f} GB")
print(f"Adam states (two per parameter): {adam_states / 1e9:.1f} GB")
print(f"Total on one device: {total / 1e9:.1f} GB")
print(f"Total per device when sharded 4 ways: {total / shards / 1e9:.1f} GB")
print("Arithmetic only: these are counts, not a sharded run")

Key points

  • Tensor parallel: split one layer's maths across devices, with partial results combined during the pass
  • Pipeline parallel: split the layers into stages, with each stage on its own device
  • A pipeline stage waits when its input stage runs late, and that waiting is its main cost
  • ZeRO-style sharding: split the weights, gradients, and optimizer states so no device holds all three in full
  • Adam keeps two state values per parameter, so its states cost twice the weight bytes in fp32
  • fp32: the 4-byte float format the earlier courses train in

In this example, the model size is stated first, at 500 million parameters. Each fp32 weight costs 4 bytes, so the weights alone print as 2.0 GB. The gradients cost the same 2.0 GB at one gradient per parameter. Adam holds two state values per parameter, so its states print as 4.0 GB. The three parts total 8.0 GB on one device. Sharding that total four ways prints 2.0 GB per device. The last line labels the program as counts, not a sharded run. The saving is real arithmetic, and the communication cost of gathering the pieces is the price the prose named.

Note

The closing program is arithmetic on stated counts, and its last line labels it as arithmetic

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

A 500M-parameter fp32 model trains with Adam. How much do the Adam states alone cost?