ZeRO: Partition Bounds
Advanced
The problem, from first principles
ZeRO assigns each device a slice of the flattened parameters. Each device needs to know exactly which elements it owns, which is the start and end of its contiguous range.
From theory to code
Implement partition_bounds(total, N, rank), returning the half-open range of elements for one device.
Constraints
- Shard sizes differ by at most one.
Hints
Hint 1
Compute the sizes of all shards with the array-split rule, sum the sizes before this rank for its start, then add its size.
Keywords
research-paperssystemstrainingmemory
Python
Run your code to see results
No results yet
Run your code to check it against the examples, or submit to run every test.