[BUG] Uneven work distribution caused by get_shard_size changes

https://github.com/microsoft/DeepSpeed/blob/3dd7ccff8103be60c31d963dd2278d43abb68fd1/deepspeed/module_inject/tp_shard.py#L35

This change introduces a new way for AutoTP to handle work when split_shape isn't divisible by num_kv_heads. This was of sharding is done in mlp, lm_head, and embed_out as well.

In the new approach the total size is split into chunks of 64, and for the remainders a block of 64 is always needed which causes new allocations.
In addition to that, the allocations are always sent to the devices based on their ranks. Which means that devices with lower ranks will always get the extra allocations which end up in full memory and allocation failure.

We saw such failures in Llama 2-7B and llama_2_70b, where a lot of those extra allocations happen in MLP which ended up in allocation failures.

I am not sure what is the real motivation for this change, or what was reason behind choosing that number exactly. Were there any experiments with numbers. results, and conclusions for picking it?

It would be better if this number was configurable as well (in a json or an env variable).

The real issue is finding a better way to evenly distribute those allocations, because in the current implementation all those allocations are going to devices with low rank (0, 1, 2) and so on.

Start with deepspeed/module_inject/tp_shard.py at line 35 and compare the behavior introduced by PR 4697. Trace AutoTP sharding across mlp, lm_head, and embed_out, focusing on how remainder allocations are assigned by rank. Done means allocation work is distributed without concentrating extra blocks on low-rank devices, including the reported Llama 2-7B and Llama 2-70B cases.