FSDP: Padded Flat Parameters
Intermediate
The problem, from first principles
The flat parameter buffer is padded so every rank holds an equal shard. The padding is at most world minus one elements, which is negligible for large models.
From theory to code
Implement fsdp_padded_numel(numel, world), the buffer size after padding.
Constraints
- Return an integer multiple of the world size.
Hints
Hint 1
Round the shard size up, then multiply back 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.