Change process group order to optimize inter-node communication

[cailun01]

To decrease inter-node communication volume, we intend to configure the ranks within the Pipeline Parallel process group to facilitate inter-node communication. Maybe self.grid of ProcessGroupManager should be:

self.grid = torch.arange(self.world_size).view(pp_size, dp_size, cp_size, tp_size)  # PP * DP * CP * TP grid

instead of

# https://github.com/huggingface/picotron/blob/df3ae8a5f0cce213816b6b287b7febc75ab98a53/picotron/process_group_manager.py#L13
self.grid = torch.arange(self.world_size).view(dp_size, pp_size, cp_size, tp_size)  # DP * PP * CP * TP grid

[zzhhjjj]

Is there any particular reason for this order change? I was following the LLaMA 3.1 paper but didn’t really run the benchmarks

[cailun01]

Is there any particular reason for this order change? I was following the LLaMA 3.1 paper but didn’t really run the benchmarks

Llama3.1 uses FSDP but picotron doesn't.

[Hannibal046]

agreed! communication volume: dp > pp, communication speed: intra-node > inter-node

[3outeille]

@cailun01 @Hannibal046 i'll plan to add FSDP soon and change this. Thanks for pointing this out !

[cailun01]

@cailun01 @Hannibal046 i'll plan to add FSDP soon and change this. Thanks for pointing this out !

Great! Which version of FSDP will you add? FSDP1 or FSDP2?