Tensor Parallelism

• We can horizontally partition the computation for one tensor operation across multiple devices, named Tensor parallelism (TP).
• We write the derivations by splitting in two, however in general, the matrixes are split equally within the GPUs in a given node

Parallelizing a GEMM (General Matrix multiply)

• $X$ is $B\times d_{model}$ , A is $d_{model} \times d_{hidden}$$

• First option (parallelize and aggregate, each “thread” computes a matrix of the same dimension as $XA$) (more memory efficient, but requires all_reduce at the end): - split $A$ along its rows and input $X$ along its columns: - $X=[X_1,X_2]$, $A=[A_1,A_2{]}^T$ - $X_i$ is $B\times d_{model}/2$ and $A_i$ is $d_{model}/2\times d_{hidden}$ - Then, $XA=X_1{A}_1+X_2{A}_2$ (it’s true, I checked) - intuition: - A matrix-matrix mul can be seen as multiple matrix-vector mul concatenated (along the columns of A). - In such matrix-vector mul, each element of a given column in $A$ is responsible for picking out its corresponding column in $X$, mutliply it by itself, and then the result is summed to obtain the new column in $XA$ - Here, we parallelize the computation over these columns, and aggregate at the end

• Second option (parallelize and concatenate, each “thread” produces a slice of $XA$) (less memory efficient but no synchronization):

  • Split $A$ along its columns $A=[A_1,A_2]$
  • $[Y_1,Y_2]=[X{A}_1,X{A}_2]$

MLP

• $X$ is $B\times d_{model}$ , A is $d_{model}\times d_{hidden}$, B is $d_{hidden}\times d_{model}$

• Usual two-layer MLP block is $f(X)=GeLU(XA)B$

  • Y = $GeLU(XA)$
  • i.e. one GEMM (general matrix multiply)
  • one GeLU
  • one GEMM

• Parallelizing the $Y=GeLU(XA)$

  • First option (parallelize and aggregate, each “thread” computes a matrix of the same dimension as $XA$):

  • • GeLU is nonlinear, so $GeLU(X_1{A}_1+X_2{A}_2)\ne GeLU(X_1{A}_1)+GeLU(X_2{A}_2)$
      • Thus we need to sychnronize before the GeLU function

  • Second option (parallelize and concatenate, each “thread” produces a slice of $XA$):

    • Split $A$ along its columns $A=[A_1,A_2]$

  • This partitioning allows the GeLU nonlinearity to be independently applied to the output of each partitioned GEMM $[Y_1,Y_2]=[GeLU(X{A}_1),GeLU(X{A}_2)]$

  • This is advantageous as it removes a synchronization point

• Parallelizing $Z=YB$

  • Given we receive $Y$ = $[Y_1,Y_2]$, split by the columns, we split $B=[B_1,B_2{]}^T$ by its rows
  • Compute $Z_i=Y_i{B}_i$
  • Synchronization
    • Z = all_reduce(Y_iB_i) by summing them
    • Called $g$ in the diagram

• Diagram

  • $g$ is an all-reduce in the forward where the matrix are aggregated by summing, identity (or splitting) in the backward
  • $f$ is an identity (or splitting) in the forward, and an all-reduce in the backward
    • 

Self-Attention

• They exploit inherent parallelism in the multihead attention operation.

  • partitioning the GEMMs associated with key (K), query (Q), and value (V ) in a column parallel fashion such that the matrix multiply corresponding to each attention head is done locally on one GPU

  • This allows us to split per attention head parameters and workload across the GPUs, and doesnt require any immediate communication to complete the self-attention.

  • The subsequent GEMM $Z=YB$ from the output linear layer (after self attention) is parallelized along its rows (i.e. $B=[B_1,B_2{]}^T$), given it receives the self-attention $Y=\text{Self-Attention}(X)$ split by columns, by design (requiring no communication)

  • Finally, we apply $g$, the all_reduce to obtain the result (before dropout)

• Diagram

  •