Confusing memory usage estimations

The calculations for memory estimation are confusing and seems to be somewhat lacking:

• No estimation for ZeRO stage 1 is available.
• ZeRO stage 2 is not taking into consideration the partitioning of gradients.
• Several of the multipliers seem to be higher than what is theoretically required. Seems like these are redundant memory copies for implementation convenience. However, they are not documented. Further, it seems from various discussions that it is hard to reach a consensus on what they actually are.
• Only considers half-precision formats. As DeepSpeed now supports full-precision, shouldn’t the estimations also be able to handle this?
• There also seems to be some inaccuracies with ZeRO stage 3 with and without optimizer offloading.

I would propose to make the estimations purely theoretical, and use the additional_buffer_factor for taking care of the extra implementation convenience DeepSpeed uses. Or at least, make it clear in the code + docs.

I tried to address these problems in a theoretical calculator I made. It is based on the official Hugging Face memory estimation space. The new/updated estimate_zero{1,2,3}_model_states_mem_needs are available here. Not sure if it is helpful, but I would love some feedback on the correctness of the calculations.

The issue points to docs/code-docs/source/memory.rst and the external parallelism_utils.py calculator. Start by comparing DeepSpeed’s current memory estimators with the theoretical calculator and the listed ZeRO stages, then establish agreed formulas and documentation; done means stage, precision, and offloading cases are covered with their assumptions explained.