Add communication volume calculation script

[Quentin-Anthony]

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