Megatron张量并行：MLP和Attention怎么切

# 模拟张量并行的核心思想：将权重矩阵切分到多卡import numpy as np W = np.array([[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]], dtype=float)X = np.array([[1,0,1,0]], dtype=float)Y_full = X @ Wprint(f'完整计算 Y = XW: {Y_full}') # 列并行W_col_0 = W[:, :2]W_col_1 = W[:, 2:]Y_col_0 = X @ W_col_0Y_col_1 = X @ W_col_1Y_col = np.concatenate([Y_col_0, Y_col_1], axis=1)print(f'列并行AllGather: {Y_col}')print(f'一致: {np.allclose(Y_full, Y_col)}') # 行并行W_row_0 = W[:2, :]W_row_1 = W[2:, :]Y_row_0 = X[:, :2] @ W_row_0Y_row_1 = X[:, 2:] @ W_row_1Y_row = Y_row_0 + Y_row_1print(f'行并行AllReduce: {Y_row}')print(f'一致: {np.allclose(Y_full, Y_row)}')

# 模拟列并行MLP第一层import numpy as nphidden_size, ffn_size, N = 8, 32, 4np.random.seed(42)W1 = np.random.randn(hidden_size, ffn_size) * 0.1X = np.random.randn(4, hidden_size)Y_full = np.maximum(0, X @ W1)partial_outputs = []for i in range(N):    W1_col = W1[:, i*(ffn_size//N):(i+1)*(ffn_size//N)]    Y_partial = np.maximum(0, X @ W1_col)    partial_outputs.append(Y_partial)Y_gathered = np.concatenate(partial_outputs, axis=1)print(f'一致: {np.allclose(Y_full, Y_gathered)}')

# 模拟行并行MLP第二层import numpy as nphidden_size, ffn_size, N = 8, 32, 4np.random.seed(42)W2 = np.random.randn(ffn_size, hidden_size) * 0.1Y_parts = [np.random.randn(4, ffn_size//N) for _ in range(N)]Y_full = np.concatenate(Y_parts, axis=1)Z_full = Y_full @ W2partial_sums = []for i in range(N):    W2_row = W2[i*(ffn_size//N):(i+1)*(ffn_size//N), :]    partial_sums.append(Y_parts[i] @ W2_row)Z_reduced = sum(partial_sums)print(f'一致: {np.allclose(Z_full, Z_reduced)}')

# 模拟Attention QKV列并行import numpy as nphidden_size, num_heads, N = 16, 4, 2head_dim = hidden_size // num_headsnp.random.seed(42)X = np.random.randn(4, hidden_size)Wq = np.random.randn(hidden_size, hidden_size) * 0.1heads_per_gpu = num_heads // Nfor gpu in range(N):    start = gpu * heads_per_gpu * head_dim    end = (gpu + 1) * heads_per_gpu * head_dim    Q = X @ Wq[:, start:end]    print(f'GPU{gpu} Q形状: {Q.shape}')    # 每卡独立计算注意力无需通信

# 通信量分析import numpy as npdef tp_analysis(bs, seq, h, tp, bw):    comm = 10*bs*seq*h*(tp-1)/tp    flops = 12*h**2*seq*bs/tp    comm_ms = comm/(bw*1e9)*1000    comp_ms = flops/(312e12)*1000    return comm/1e9, comm_ms, comp_ms, comm_ms/(comm_ms+comp_ms)for tp in [2,4,8,16]:    g,ct,cpt,r = tp_analysis(32,2048,12288,tp,600)    print(f'TP={tp}: {g:.0f}GB, 通信{ct:.0f}ms, 计算{cpt:.0f}ms, 占比{r:.1%}')

# 3D并行配置分析def config_analysis(gpus, params_b, seq, tp, pp, dp):    assert tp*pp*dp == gpus    h = int((params_b*1e9/24)**0.5)    mem = params_b*1e9*4/(tp*pp*1e9) + 12*h*seq*4/tp/1e9    tp_comm = 10*1*seq*h*(tp-1)/tp/1e9    return mem, tp_commfor tp,pp,dp in [(8,16,8),(8,128,1),(4,32,8)]:    m,c = config_analysis(1024,175,2048,tp,pp,dp)    print(f'TP={tp},PP={pp},DP={dp}: 显存{m:.1f}GB, TP通信{c:.1f}GB')

import numpy as np np.random.seed(42)hidden_size, ffn_size, num_heads = 16, 64, 4head_dim = hidden_size // num_headsseq_len, batch_size, N = 8, 2, 2tokens = batch_size * seq_lenX = np.random.randn(tokens, hidden_size) W1 = np.random.randn(hidden_size, ffn_size) * 0.01W2 = np.random.randn(ffn_size, hidden_size) * 0.01Wq = np.random.randn(hidden_size, hidden_size) * 0.01Wk = np.random.randn(hidden_size, hidden_size) * 0.01Wv = np.random.randn(hidden_size, hidden_size) * 0.01Wo = np.random.randn(hidden_size, hidden_size) * 0.01 # 完整计算mlp_out = np.maximum(0, X @ W1) @ W2Q, K, V = X @ Wq, X @ Wk, X @ WvQ = Q.reshape(tokens, num_heads, head_dim).transpose(0, 2, 1)K = K.reshape(tokens, num_heads, head_dim).transpose(0, 2, 1)V = V.reshape(tokens, num_heads, head_dim).transpose(0, 2, 1)scores = np.matmul(Q, K.transpose(0, 2, 1)) / np.sqrt(head_dim)attn = np.exp(scores) / np.exp(scores).sum(axis=-1, keepdims=True)attn_out = np.matmul(attn, V).transpose(0, 2, 1).reshape(tokens, hidden_size)layer_out = mlp_out + attn_out @ Wo # TP并行计算W1_cols = [W1[:, i*(ffn_size//N):(i+1)*(ffn_size//N)] for i in range(N)]W2_rows = [W2[i*(ffn_size//N):(i+1)*(ffn_size//N), :] for i in range(N)]mlp_parts = []for i in range(N):    y = np.maximum(0, X @ W1_cols[i])    mlp_parts.append(y @ W2_rows[i])mlp_tp = sum(mlp_parts) hpg = num_heads // NWq_cols = [Wq[:, i*hpg*head_dim:(i+1)*hpg*head_dim] for i in range(N)]Wk_cols = [Wk[:, i*hpg*head_dim:(i+1)*hpg*head_dim] for i in range(N)]Wv_cols = [Wv[:, i*hpg*head_dim:(i+1)*hpg*head_dim] for i in range(N)]Wo_rows = [Wo[i*hpg*head_dim:(i+1)*hpg*head_dim, :] for i in range(N)]attn_parts = []for i in range(N):    qi = (X @ Wq_cols[i]).reshape(tokens, hpg, head_dim).transpose(0, 2, 1)    ki = (X @ Wk_cols[i]).reshape(tokens, hpg, head_dim).transpose(0, 2, 1)    vi = (X @ Wv_cols[i]).reshape(tokens, hpg, head_dim).transpose(0, 2, 1)    si = np.matmul(qi, ki.transpose(0, 2, 1)) / np.sqrt(head_dim)    ai = np.exp(si) / np.exp(si).sum(axis=-1, keepdims=True)    oi = np.matmul(ai, vi).transpose(0, 2, 1).reshape(tokens, hpg*head_dim)    attn_parts.append(oi @ Wo_rows[i])attn_tp = sum(attn_parts) layer_tp = mlp_tp + attn_tpprint(f'完整: {layer_out.shape}, TP={N}: {layer_tp.shape}')print(f'一致: {np.allclose(layer_out, layer_tp, atol=1e-6)}')