Add communication volume calculation script

Would be good to model the communication volume in bytes of a given parallelism setup. Situations to model:
- Different parallelism schemes
    - ZeRO-1/2/3, ZeRO++
    - 3D parallelism
- Activation checkpointing
- Different dtypes

Bonus points:
- % volume breakdown separated by collective