1 基础的权重切分
1.1 按行切分权重
x = np.array([[1.16,0.23],[0.57,1.36],[4.41,-2.16]])
w = np.array([[0.22,0.41],[0.17,-0.51]])
print("不切分下的矩阵乘积:")
print(x@w)
x1 = x[:, [0]] # 形状为(3,)
x2 = x[:, [1]] # 形状为(3,)
# 切分第一行
w1 = w[0:1, :] # 形状为(1,2)
# 切分第二行
w2 = w[1:2, :] # 形状为(1,2)
print("切分下的矩阵乘积:")
print(x1@w1 + x2@w2)对于输入x,和参数w进行矩阵计算，如果对于w进行按行切分，x按列切分，各自相乘再相加就会得到和不切分下一模一样的计算结果。对TP并行最直观的一种体现。
1.2 按列切分权重
x = np.array([[1.16,0.23],[0.57,1.36],[4.41,-2.16]])
w = np.array([[0.22,0.41],[0.17,-0.51]])
print("不切分下的矩阵乘积:")
print(x@w)
# 切分第一列
w1 = w[:, [0]] # 形状为(1,2)
# 切分第二列
w2 = w[:, [1]] # 形状为(1,2)
print("切分下的矩阵乘积:")
print(np.concatenate([x@w1, x@w2], axis=-1))相同的我们将参数进行按列切开，x不变，那么各自计算之后的矩阵通过连接就会得到一个和不切分下相同的输出。
2 Transformer 模型的TP计算
import numpy as np
import math
def softmax(matrix):
exp_matrix = np.exp(matrix - np.max(matrix, axis=1, keepdims=True)) # 防止溢出
return exp_matrix / np.sum(exp_matrix, axis=1, keepdims=True)
arr = np.arange(16) # 1x12数组
reshaped_x = arr.reshape(4, 4) # 转为3行4列
reshaped_q = arr.reshape(4, 4) # 转为3行4列
reshaped_v = arr.reshape(4, 4) # 转为3行4列
reshaped_k = arr.reshape(4, 4) # 转为3行4列
#计算k的转置换
y1 = softmax(reshaped_q@reshaped_k.T/math.sqrt(4))@reshaped_v
b_test = [[1,2],[3,4],[5,6],[7,8]]
print("不切分下最终输出")
print(y1@b_test)
#计算多头（这里会将qlen和头置换一下位置
new_x = np.swapaxes(reshaped_x.reshape(4,2,2), 0, 1)
new_q = np.swapaxes(reshaped_q.reshape(4,2,2), 0, 1)
new_v = np.swapaxes(reshaped_v.reshape(4,2,2), 0, 1)
new_k = np.swapaxes(reshaped_k.reshape(4,2,2), 0, 1)
head1=softmax(new_q[0]@new_k[0].T/math.sqrt(2))@new_v[0]
head2=softmax(new_q[1]@new_k[1].T/math.sqrt(2))@new_v[1]
test1 = head1@[[1,2],[3,4]]
test2 = head2@[[5,6],[7,8]]
print("切分下的最终输出")
print(test1+test2)
#可以看到就是最后进行一次allreduce就可以了，相同的反向计算的时候也是进行一次all_reduce操作就行了这段代码首先定义x和QVK矩阵，然后使用transformer的公式求解输出，再进行一次输出矩阵的相乘得到最终的输出。然后我们将QVK增加一个维度，相当于将其进行了切分，然后各自进行计算，对b矩阵进行切分各自相乘，最后只需要一次相加(all_reduce集合通信操作)就可以得到最终和不切分相同的输出了。在transformer中还有一个MLP层，但是那个比较简单，结合上述的行/列切分和链接文章可以很好理解，本处省略其Python代码解释。
3 输入层embeding
import numpy as np
import math
#当batch = 1 ,s(句子单词数) = 2, h(单词维度) = 3, v（词表总数) = 4， N(GPU的数量，TP并行数量) = 2
words_input = [0,3] #表示这个输入是0,3号token
words_table = np.array([[0,4,8],[3,5,18],[5,6,3],[6,7,1]])
#行切分，形成两个单词table
w1 = words_table[0:2, :] # 形状为(1,2)
w2 = words_table[2:4, :] # 形状为(1,2)
enbeding_1 = np.array([[0,4,8],[0,0,0]]) #words_input 和 w1 check的时候得到enbeding_1
enbeding_2 = np.array([[0,0,0],[6,7,1]]) #words_input 和 w2 check的时候得到enbeding_2
print("输入层:enbeding输出:")
print(enbeding_1+enbeding_2) #allreduce操作假设输入是0号单词和3号单词，那么需要在整体的单词表中找到0,3号的token的数组([0.4.8],[6,7,1])，我们将整个大的单词分成两份，将words_input放在第一份中只能找到[0,4,8],第二份中会找到[6,7,1]，然后两者一次相加就会得到([0.4.8],[6,7,1])
4 输出层的embeding
import numpy as np
import math
words_table = np.array([[0,4,8],[3,5,18],[18,6,3],[6,7,1]])
#行切分，形成两个单词table
w1 = words_table[0:2, :] # 形状为(1,2)
w2 = words_table[2:4, :] # 形状为(1,2)
#当模型前向计算之后输出的一个预测的列表的时候，我们假设预测的值变成了如下列表
predicts = np.array([[0,4,8],[6,7,1]])
#使用差异值方式求解预测值和字符表里面所有单词的差异，predicts和w1.T, w2.T之间的点积。
#首先对矩阵进行标准化，然后矩阵的相乘的值就是余炫相似度，约大的值表示越相似
#最后两个矩阵结合，再对每一行进行soft_max得到每个token的预测概率
#用这个概率和真实的值,这里就是[[1,0,0,0,],[0,0,0,1]]使用cross_Entropy函数即可求出loss值。
row_norms = np.linalg.norm(predicts, axis=1, keepdims=True) # 计算每行的L2范数
row_normalized = predicts / row_norms # 广播除法
col_norms1 = np.linalg.norm(w1.T, axis=0, keepdims=True) # 计算每列的L2范数
col_normalized1 = w1.T / col_norms1 # 广播除法
col_norms2 = np.linalg.norm(w2.T, axis=0, keepdims=True) # 计算每列的L2范数
col_normalized2 = w2.T / col_norms2 # 广播除法
total = np.concatenate([row_normalized@col_normalized1, row_normalized@col_normalized2], axis=-1) #all-gather操作
def softmax(matrix):
exp_matrix = np.exp(matrix - np.max(matrix, axis=1, keepdims=True)) # 防止溢出
return exp_matrix / np.sum(exp_matrix, axis=1, keepdims=True)
print(softmax(total))单词表还是和输入端一样切分，然后经过attention+mlp层之后的输出是[0,4,8],[6,7,1]那么代表预测的值是0.3号单词。为了简单和方便，我们将所有的预测值和单词表里面的vector进行正则化，这样如果预测的值就可以用点积的方式求和每个单词的余炫相似度，每一份的点积自己对应的单词表，然后合在一起就会得到预测值和整个单词表中每个单词的相似度，然后使用softmax处理一下，就可以看到预测值和哪一个单词的相似度最高。但是All-Gather会产生额外的通讯量 bsv。当词表v很大时，这个通讯开销也不容忽视。针对这种情况，可以做如下优化：
import numpy as np
import math
words_table = np.array([[0,4,8],[3,5,18],[18,6,3],[6,7,1]])
#行切分，形成两个单词table
w1 = words_table[0:2, :] # 形状为(1,2)
w2 = words_table[2:4, :] # 形状为(1,2)
#当模型前向计算之后输出的一个预测的列表的时候，我们假设预测的值变成了如下列表
predicts = np.array([[0,4,8],[6,7,1]])
#使用差异值方式求解预测值和字符表里面所有单词的差异，predicts和w1.T, w2.T之间的点积。
#首先对矩阵进行标准化，然后矩阵的相乘的值就是余炫相似度，约大的值表示越相似
#最后两个矩阵结合，再对每一行进行soft_max得到每个token的预测概率
#用这个概率和真实的值,这里就是[[1,0,0,0,],[0,0,0,1]]使用cross_Entropy函数即可求出loss值。
row_norms = np.linalg.norm(predicts, axis=1, keepdims=True) # 计算每行的L2范数
row_normalized = predicts / row_norms # 广播除法
col_norms1 = np.linalg.norm(w1.T, axis=0, keepdims=True) # 计算每列的L2范数
col_normalized1 = w1.T / col_norms1 # 广播除法
col_norms2 = np.linalg.norm(w2.T, axis=0, keepdims=True) # 计算每列的L2范数
col_normalized2 = w2.T / col_norms2 # 广播除法
# math.exp(col_normalized1)
Y1 = row_normalized@col_normalized1
Y2 = row_normalized@col_normalized2
# print(np.sum(np.exp(Y1),axis = 1))
# print(np.sum(np.exp(Y2),axis = 1))
sum_e = np.sum(np.exp(Y1),axis = 1) + np.sum(np.exp(Y2),axis = 1)
result1 = np.exp(Y1) / sum_e[:, np.newaxis]
result2 = np.exp(Y2) / sum_e[:, np.newaxis]
#这就是两个分裂开的soft_max值矩阵，相对于之前的那种方式就是通过提前交流soft_max的分母，然后各自在本地计算了soft_max分子
print("new soft max")
print(result1)
print(result2)
#和各自的真值进行corss Entropy计算，例如第一个[0,4,8]的GPU0/1的预测概率是
#[0.33070998 0.32063926],[0.16087279 0.18777798]
#真值是[1,0],[0,0]
#通过计算各自的loss值之后，all_reduce交换一次loss值就可以了。
print("old soft max")
# 之前说到的使用all-gather的通信方式
total = np.concatenate([row_normalized@col_normalized1, row_normalized@col_normalized2], axis=-1) #all-gather操作
def softmax(matrix):
exp_matrix = np.exp(matrix - np.max(matrix, axis=1, keepdims=True)) # 防止溢出
return exp_matrix / np.sum(exp_matrix, axis=1, keepdims=True)
print(softmax(total))每块GPU上，我们可以先按行求和，得到各自GPU上的GPU_sum(e) 2 将每块GPU上结果做AllReduce，得到每行最终的sum(e)，也就softmax中的分母。此时的通讯量为 3 在每块GPU上，即可计算各自维护部分的e/sum(e)，将其与真值做cross-entropy，得到每行的loss，按行加总起来以后得到GPU上scalar Loss。 4 将GPU上的scalar Loss做AllReduce，得到总Loss。此时通讯量为N。
这样，我们把原先的通讯量从 bsv 大大降至 b*s + N。
其中old soft max是刚刚提到的通过all-gather获取的值,而new soft max是优化后的值，可以看到两者在最后的输出一模一样的。
