ColumnParallelLinear 和 RowParallelLinear

ColumnParallelLinear 和 RowParallelLinear 是两种并行线性层，它们的主要区别在于权重矩阵的分割方式和计算过程。具体来说：

ColumnParallelLinear

• 权重矩阵分割方式：

  • 权重矩阵 W 被按列（column）分割成多个子矩阵，每个子矩阵在并行设备上进行计算。
  • 如果权重矩阵 W 的形状是 (output_dim, input_dim)，它会被分割成 P 个子矩阵，每个子矩阵的形状是 (output_dim, input_dim / P)。

• 计算过程：(矩阵块的运算法则与一般矩阵运算相同)

  The linear layer is defined as Y = XA + b. A is parallelized along
      its second dimension as A = [A_1, ..., A_p].

• 优势：

  • 适合输入数据较小但输出维度较大的场景。

RowParallelLinear

• 权重矩阵分割方式：

  • 权重矩阵 W 被按行（row）分割成多个子矩阵，每个子矩阵在并行设备上进行计算。
  • 如果权重矩阵 W 的形状是 (output_dim, input_dim)，它会被分割成 P 个子矩阵，每个子矩阵的形状是 (output_dim / P, input_dim)。

• 计算过程：(矩阵块的运算法则与一般矩阵运算相同)

        

The linear layer is defined as Y = XA + b. A is parallelized along
    its first dimension and X along its second dimension as:
               -   -
              | A_1 |
              | .   |
          A = | .   |        X = [X_1, ..., X_p]
              | .   |
              | A_p |
               -   -

• 优势：

  • 适合输出维度较小但输入数据较大的场景。