parameters = 1_000_000_000
workers = 4
bytes_per_parameter = {
    "parameters": 2,
    "gradients": 2,
    "optimizer": 4 + 4 + 4,  # Master weights and two moments.
}
partitioned = {
    0: set(),
    1: {"optimizer"},
    2: {"optimizer", "gradients"},
    3: {"optimizer", "gradients", "parameters"},
}
totals = []
for stage in range(4):
    parts = {}
    for name, width in bytes_per_parameter.items():
        divisor = workers if name in partitioned[stage] else 1
        parts[name] = parameters * width / divisor / 1_000_000_000
    total = sum(parts.values())
    totals.append(total)
    print(f"stage={stage}: params={parts['parameters']:.1f} GB, "
          f"grads={parts['gradients']:.1f} GB, "
          f"optimizer={parts['optimizer']:.1f} GB, total={total:.1f} GB")
assert totals == [16.0, 7.0, 5.5, 4.0]
assert totals[-1] * workers == totals[0]