Uniform `_distribute_buffer_sizes`, and move into shampoo_utils.py

[namgyu-youn]

There are duplicate methods for distributed_buffer_sizes at :

• distributed_shampoo/utils/shampoo_hsdp_distributor.py : Ln 292
• distributed_shampoo/utils/shampoo_ddp_distributor.py : Ln 187
• distributed_shampoo/utils/shampoo_hybrid_shard_distributor.py : Ln 296

Is this intended? In my thought, that duplicate code could be uniformed, and moved into distributed_shampoo/utils/shampoo_utils.py.

[tsunghsienlee]

Hi @namgyu-youn ,

Good question, ideally we want to merge distributed_shampoo/utils/shampoo_ddp_distributor.py (DDP) and distributed_shampoo/utils/shampoo_fsdp_distributor.py (FSDP1) into distributed_shampoo/utils/shampoo_hsdp_distributor.py (HSDP1); similarly distributed_shampoo/utils/shampoo_fully_shard_distributor.py (FSDP2) into distributed_shampoo/utils/shampoo_hybrid_shard_distributor.py (HSDP2). We haven't done that yet because we need to make sure DDP and FSDP1 does not degrade in HSDP1 situation; similarly, FSDP2 does not degrade in HSDP2 situation.

This still means both are using distributed_buffer_sizes() for both HSDP1 and HSDP2, and I don't mind we move this two into a shared utility but we just felt the aforementioned work takes priority so we haven't done this yet.

[namgyu-youn]

Hi @namgyu-youn ,

Good question, ideally we want to merge distributed_shampoo/utils/shampoo_ddp_distributor.py (DDP) and distributed_shampoo/utils/shampoo_fsdp_distributor.py (FSDP1) into distributed_shampoo/utils/shampoo_hsdp_distributor.py (HSDP1); similarly distributed_shampoo/utils/shampoo_fully_shard_distributor.py (FSDP2) into distributed_shampoo/utils/shampoo_hybrid_shard_distributor.py (HSDP2). We haven't done that yet because we need to make sure DDP and FSDP1 does not degrade in HSDP1 situation; similarly, FSDP2 does not degrade in HSDP2 situation.

This still means both are using distributed_buffer_sizes() for both HSDP1 and HSDP2, and I don't mind we move this two into a shared utility but we just felt the aforementioned work takes priority so we haven't done this yet.

Thanks for your comment. But a complete merge would be difficult for me right now, so I opened a simple PR about this. Could you review #94 ?