FSDP: Elements per Shard
Beginner
The problem, from first principles
PyTorch FSDP (Zhao et al., 2023) shards each group of parameters across ranks in a flat buffer. Every rank stores one equal shard, rounded up so the shards line up across devices.
From theory to code
Implement fsdp_shard_numel(numel, world), the elements stored per rank.
Constraints
- Round up.
Hints
Hint 1
Use integer ceiling division by the world size.
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.