Support tensor parallel/pipeline parallel

[gongel]

Support tensor parallel/pipeline parallel currently?

[karan6181]

Can you please share more details ?

[gongel]

NVIDIA-Megatron team proposed "Tensor Parallelism". When training in "Tensor Parallelism", the rank in same group has same data.
Paper: https://arxiv.org/pdf/2205.05198.pdf
Repo: https://github.com/NVIDIA/Megatron-LM

But in streaming, you only support DDP/FSDP.

[andreamad8]

Any plan to add this?

one easy solution that seems not to be working could be:

os.environ["WORLD_SIZE"] = str(os.environ["WORLD_SIZE"]  // model_parallel_size)
os.environ["RANK"] = str(os.environ["RANK"] // model_parallel_size)

I tried, but seems the code gets stuck after calling something like:

batch = next(batch_iterator)

where batch_iterator a dataloder.

cc: @karan6181

[karan6181]

@snarayan21 Looks like this is being addressed. Is that right?

[huxuan]

Would like to know if there is any example of megatron integration.

[snarayan21]

@andreamad8 @huxuan @gongel please see the replication argument detailed in our docs here.

@huxuan We don't have an explicit example of a megatron integration, but as it's pytorch based, you can simply swap in the dataset / dataloader.

[casper-hansen]

@snarayan21

We can compute replication based on TP and SP as follows. But it seems we are lacking documentation on PP. I think part of this issue was unanswered because of that. How do we modify the parameters to streaming with pipeline parallelism?

TP and SP:

replication = tp * sp

Various configurations I tried just to sanity check. None of these are correct (unexpected high loss).

replication = tp * sp * pp

replication = tp * sp * pp
num_canonical_nodes = world_size // replication

replication = tp * sp
num_canonical_nodes = world_size // (tp * sp * pp)

[casper-hansen]

Opened an issue for this separately:
#827