Calculate persistent model-state bytes per rank for replicated data parallelism (DDP) and ZeRO stages 1 through 3. Use this pooled-byte model: P, G, and O are the supplied parameter, gradient, and optimizer-state byte counts, and R is the number of ranks. Each M denotes the largest per-rank allocation under that model; ceiling brackets round up to an integer.
MDDP=P+G+O. MZeRO1=P+G+⌈RO⌉. MZeRO2=P+⌈RG+O⌉. MZeRO3=⌈RP+G+O⌉.The sharded categories in each numerator form one byte pool. Activations, transient gathered buffers, and allocator padding are excluded. Return a Python dictionary with ddp, zero1, zero2, and zero3, each an exact Python integer byte count.
Input: parameter_bytes = 400, gradient_bytes = 400, optimizer_bytes = 800, world_size = 4
Output: {"ddp":1600,"zero1":1000,"zero2":700,"zero3":400}
Explanation: DDP replicates all 1600 bytes. The ZeRO stages progressively shard optimizer, gradient, and parameter state across four ranks.
Input: parameter_bytes = 5, gradient_bytes = 7, optimizer_bytes = 11, world_size = 3
Output: {"ddp":23,"zero1":16,"zero2":11,"zero3":8}
Input: parameter_bytes = 100, gradient_bytes = 50, optimizer_bytes = 200, world_size = 1
Output: {"ddp":350,"zero1":350,"zero2":350,"zero3":350}
divmod returns an integer quotient and remainder.
Python integers support exact arithmetic without a fixed-width overflow.
Accepts: number
Accepts: number
Accepts: number
Accepts: number
Calculate persistent model-state bytes per rank for replicated data parallelism (DDP) and ZeRO stages 1 through 3. Use this pooled-byte model: P, G, and O are the supplied parameter, gradient, and optimizer-state byte counts, and R is the number of ranks. Each M denotes the largest per-rank allocation under that model; ceiling brackets round up to an integer.
MDDP=P+G+O. MZeRO1=P+G+⌈RO⌉. MZeRO2=P+⌈RG+O⌉. MZeRO3=⌈RP+G+O⌉.The sharded categories in each numerator form one byte pool. Activations, transient gathered buffers, and allocator padding are excluded. Return a Python dictionary with ddp, zero1, zero2, and zero3, each an exact Python integer byte count.
Input: parameter_bytes = 400, gradient_bytes = 400, optimizer_bytes = 800, world_size = 4
Output: {"ddp":1600,"zero1":1000,"zero2":700,"zero3":400}
Explanation: DDP replicates all 1600 bytes. The ZeRO stages progressively shard optimizer, gradient, and parameter state across four ranks.
Input: parameter_bytes = 5, gradient_bytes = 7, optimizer_bytes = 11, world_size = 3
Output: {"ddp":23,"zero1":16,"zero2":11,"zero3":8}
Input: parameter_bytes = 100, gradient_bytes = 50, optimizer_bytes = 200, world_size = 1
Output: {"ddp":350,"zero1":350,"zero2":350,"zero3":350}
divmod returns an integer quotient and remainder.
Python integers support exact arithmetic without a fixed-width overflow.
Accepts: number
Accepts: number
Accepts: number
Accepts: number
