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是優化后的值，可以看到兩者在最后的輸出一模一樣的。