FSDP: Memory per Rank
Beginner
The problem, from first principles
Sharding is the point of FSDP: each rank holds only its portion of the parameters. The memory per rank falls roughly as one over the world size.
From theory to code
Implement fsdp_memory_per_rank(numel, world, nbytes), the bytes of one rank's shard.
Constraints
- Returns an integer.
Hints
Hint 1
Compute the shard size and multiply by bytes per element.
Keywords
research-paperssystemstrainingsharding
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.