Question about Detail Communication Cost and Memory Cost of Multi-Dimensional Parallelism

[ConnollyLeon]

Dear developers,

I have read the papers of 2D, 2.5D and 3D parallelism, but still have some questions about the communication cost and memory cost of Colossal-AI.

Here is my question list:

In 3D parallelism, I notice that Colossal-AI moves (MN +NK +MK)/p^3 across p devices in a matrix-matrix multiplication using All-gather or Reduce-scatter. And the final communication bandwidth cost is O(1/p^2).

• Q1: Is this communication implemented by RING, whose communication bandwidth cost is specifically (MN +NK +MK)*(p-1)/p^3?
• Q2: All-gathering matrix A and B might temporally create sub-matrix with shape of (M/p, N/p) and (N/p, K/p). Although it can be dropped after reduce-scattering matrix C, I still think the peak memory cost of matrix multiplication turns out to be O(1/p^2) instead of O(1/p^3) in the paper. I wonder if the O(1/p^3) is described for the whole process of an iteration.
• Q3: I am still a green hand in communication field, and I notice that in the paper, you say the communication latency of 3D parallelism is O(log(p)). Could you please tell me how this latency is calculated?
• Q4: How do you organize the communication well since the bandwidth between intra-node devices and inter-node devices are different? Although the communication volume is highly reduced, would this be a bottleneck in 2D or 3D parallelism?

Waiting for your responses, please help us~~

[FrankLeeeee]

Hi @ConnollyLeon , thank you for your question. @kurisusnowdeng is in charge of 3D parallelism. He will try to answer your question. :)

[kurisusnowdeng]

Hi, @ConnollyLeon . Thank you for your interest. I give my anwsers according to my understanding to your questions one by one.

• A1: If using NCCL as the communication backend, the RING algorithm is supposed to be used, and the bandwidth cost you gave is correct.
• A2: Indeed during both communication and multiplication the memory cost is O(1/p^2). However, such memory can be reused by other layers. Compared to the memory that is kept for backward (usually tens of transformer layers), it seems not so significant in terms of peak memory usage. That is why we also partition the activations by p^3.
• A3: If using a tree topology for parallel communication, the latency can be O(log(p)). You can refer to these slides where we first found the conclusion. However, with the RING topology, the latency should be O(p-1). Thank you for correction.
• A4: Communication across different nodes can be costly. We now suggest to use tensor parallelism within a single server, such as 8x GPUs for 2x2x2 3D parallelism. We are also making our efforts to reduce the communication cost in order to improve the throughput.

[ConnollyLeon]

Thanks for your answers, and most of them matches my guesses, but I still have a small question about A4:

• You suggest to use TP within a single server, but in the 3D tensor parallelism paper, you give an experiment result of 64 GPUs (16 servers, each have 4 GPU), which have already outperformed 1D and 2D. It seems like that it already has a good improvement in multi-node environment. But why did you suggest to use TP within a single server instead? Could you please tell us whether there is any performance deterioration (Compared to 1D or 2D) in deploying 3D TP on multi-nodes?

A little suggestion: It would be better to provide some quantized profiles of communication time and volume along different axes (which may have different bandwidth), to help community understand the advantages of Colossal-AI :).

[kurisusnowdeng]

Since the model size is increasing, we may have to use multiple servers to accommodate a huge transformer model, e.g. GPT-3. So, it is important to present the scalability of 3D parallelism even across servers. I mean we suggest not to use too many servers for TP if unnecessary. BTW, thank you so much for the suggestion.

[ConnollyLeon]

Thanks for your replies. I try to analyze the memory of 3D-TP formally, could you please check whether I am right?

Suppose we have l layers (or operators that has weights), and p=q*q*q devices, and the memory of the model is M. The total memory cost would be (1 / q^2 * 1 / l) M + (1 / q^3 * (l - 1) / l) M = (q + l - 1) /( l * q^3)M, which means that there is only one layer that requires to hold (1 / q^2) part during running. When the l is much larger than q, the memory cost become sO(1 / q^3) according to L'Hôpital's rule. When the l equals to 1, the memory cost becomes O(1 / q^2). So the interval of the memory cost of 3D-TP is (1 / q^3 M, 1 / q^2 M]. When given specific model with l layers, and p devices, the memory cost of each device can be easily calculated and is in this interval.

[ConnollyLeon]

@kurisusnowdeng 还有一个小问题，why not use All-gather instead of Broadcast in 2D model parallelism? According to my analysis, the total communication time (bandwidth cost specifically) of using All-gather would be smaller.

[FrankLeeeee]

@kurisusnowdeng 还有一个小问题，why not use All-gather instead of Broadcast in 2D model parallelism? According to my analysis, the total communication time (bandwidth cost specifically) of using All-gather would be smaller.

We have tested using all-gather and broadcast. All-gather leads to higher memory cost. In terms of speed, they are quite close if using async communication.

[ConnollyLeon]

Thanks for your reply. I have read the implementation here.

ColossalAI/colossalai/nn/layer/parallel_2d/_operation.py

Line 390 in f5ca88e

opb[1 - cur] = dist.broadcast(B_list[1 - cur], src=src_b + summa_dim, group=col_group, async_op=True)

It seems like that you replace Broadcast of A for Reduce of C, and uses a local buffer B_list and C_list to support the overlap of communication and computation. Using buffers can help reuse memory when compared to All-gather, I think the memory advantage comes from here. But why do you make this replacement? Would it be better to analyze the size of A and C first to determine a Broadcast of A or a Reduce of C (pick the one with smaller size for less communication amount)?

[FrankLeeeee]

This is a matmul of AB^T. According to SUMMA Paper, reduction of C is needed.

[ConnollyLeon]

Thanks~, It was my fault not reading it carefully. Now I completely understand. 😄

[FrankLeeeee]

Welcome! :)