5.1 Tensor并行与算子拆分

摘要

在大模型推理中，单GPU显存容量成为限制大规模模型部署的关键瓶颈。Tensor并行技术通过将模型权重和计算拆分到多个GPU上，实现了模型规模的有效扩展。本文档深度探讨Tensor并行的理论基础、实现策略和优化技术，涵盖算子拆分、通信优化、负载均衡等关键议题，并结合具体的代码示例和性能数据，为大规模模型推理提供一套完整的并行化解决方案。

1. Tensor并行基础理论

1.1 模型并行化的数学基础

张量分解与并行计算

import torch
import torch.nn as nn
import torch.distributed as dist
import numpy as np
from typing import List, Dict, Tuple, Optional
import math

class TensorParallelism:
    """张量并行化的数学基础实现"""
    
    @staticmethod
    def tensor_parallel_split(tensor: torch.Tensor, 
                            split_dim: int, 
                            num_partitions: int) -> List[torch.Tensor]:
        """
        沿指定维度分割张量
        
        Args:
            tensor: 输入张量
            split_dim: 分割维度
            num_partitions: 分割数量
            
        Returns:
            分割后的张量列表
        """
        if tensor.size(split_dim) % num_partitions != 0:
            raise ValueError(f"Dimension {split_dim} size {tensor.size(split_dim)} "
                           f"must be divisible by {num_partitions}")
        
        split_size = tensor.size(split_dim) // num_partitions
        splits = torch.split(tensor, split_size, dim=split_dim)
        
        return splits
    
    @staticmethod
    def tensor_parallel_gather(tensor_parts: List[torch.Tensor], 
                             gather_dim: int) -> torch.Tensor:
        """
        沿指定维度聚合分割后的张量
        
        Args:
            tensor_parts: 分割后的张量列表
            gather_dim: 聚合维度
            
        Returns:
            聚合后的完整张量
        """
        return torch.cat(tensor_parts, dim=gather_dim)
    
    @staticmethod
    def tensor_parallel_row_split(weight: torch.Tensor, 
                                num_partitions: int) -> List[torch.Tensor]:
        """
        对权重矩阵进行行分割（用于线性层的前向传播）
        
        Args:
            weight: 权重矩阵 [out_features, in_features]
            num_partitions: 分割数量
            
        Returns:
            行分割后的权重列表
        """
        return TensorParallelism.tensor_parallel_split(weight, 0, num_partitions)
    
    @staticmethod
    def tensor_parallel_column_split(weight: torch.Tensor, 
                                  num_partitions: int) -> List[torch.Tensor]:
        """
        对权重矩阵进行列分割（用于线性层的前向传播）
        
        Args:
            weight: 权重矩阵 [out_features, in_features]
            num_partitions: 分割数量
            
        Returns:
            列分割后的权重列表
        """
        return TensorParallelism.tensor_parallel_split(weight, 1, num_partitions)

# 演示张量分割与聚合
def demonstrate_tensor_splitting():
    """演示张量分割操作"""
    print("张量分割演示")
    print("-" * 50)
    
    # 创建一个权重矩阵
    weight = torch.randn(128, 256)  # 128输出特征，256输入特征
    print(f"原始权重矩阵形状: {weight.shape}")
    
    # 行分割（输出维度分割）
    row_partitions = TensorParallelism.tensor_parallel_row_split(weight, 4)
    print(f"行分割为4份，每份形状: {[w.shape for w in row_partitions]}")
    
    # 聚合分割后的张量
    gathered_weight = TensorParallelism.tensor_parallel_gather(row_partitions, 0)
    print(f"聚合后权重矩阵形状: {gathered_weight.shape}")
    print(f"重建误差: {torch.norm(weight - gathered_weight):.6f}")
    
    # 列分割（输入维度分割）
    col_partitions = TensorParallelism.tensor_parallel_column_split(weight, 4)
    print(f"列分割为4份，每份形状: {[w.shape for w in col_partitions]}")
    
    # 聚合分割后的张量
    gathered_weight_col = TensorParallelism.tensor_parallel_gather(col_partitions, 1)
    print(f"聚合后权重矩阵形状: {gathered_weight_col.shape}")
    print(f"重建误差: {torch.norm(weight - gathered_weight_col):.6f}")

# 运行演示
demonstrate_tensor_splitting()
​

线性层的张量并行化

class TensorParallelLinear(nn.Module):
    """张量并行化的线性层"""
    
    def __init__(self, 
                 input_size: int, 
                 output_size: int, 
                 num_partitions: int,
                 parallel_mode: str = 'row'):
        """
        初始化张量并行线性层
        
        Args:
            input_size: 输入特征维度
            output_size: 输出特征维度
            num_partitions: 并行分区数
            parallel_mode: 并行模式 ('row' 或 'col')
        """
        super().__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.num_partitions = num_partitions
        self.parallel_mode = parallel_mode
        
        # 验证维度可分割性
        if parallel_mode == 'row':
            if output_size % num_partitions != 0:
                raise ValueError("output_size must be divisible by num_partitions for row splitting")
            self.partition_size = output_size // num_partitions
        else:  # column splitting
            if input_size % num_partitions != 0:
                raise ValueError("input_size must be divisible by num_partitions for column splitting")
            self.partition_size = input_size // num_partitions
        
        # 初始化权重
        self.weight = nn.Parameter(torch.randn(self.partition_size, input_size))
        self.bias = nn.Parameter(torch.randn(self.partition_size))
    
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """
        前向传播
        
        Args:
            input: 输入张量 [batch_size, input_size]
            
        Returns:
            输出张量 [batch_size, partition_size]
        """
        if self.parallel_mode == 'row':
            # 行分割：每个GPU处理部分输出
            output = torch.matmul(input, self.weight.t()) + self.bias
        else:  # column splitting
            # 列分割：每个GPU处理部分输入
            # 实际的矩阵乘法需要跨GPU通信
            output = torch.matmul(self.weight, input.t()).t() + self.bias
        
        return output
    
    def get_communication_pattern(self) -> Dict:
        """获取通信模式信息"""
        return {
            'parallel_mode': self.parallel_mode,
            'num_partitions': self.num_partitions,
            'partition_size': self.partition_size,
            'communication_required': self.parallel_mode == 'col'
        }

class ColumnParallelLinear(nn.Module):
    """列并行线性层（需要AllReduce通信）"""
    
    def __init__(self, 
                 input_size: int, 
                 output_size: int, 
                 num_partitions: int):
        super().__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.num_partitions = num_partitions
        
        # 验证输入维度可分割性
        if input_size % num_partitions != 0:
            raise ValueError("input_size must be divisible by num_partitions")
        
        self.partition_size = input_size // num_partitions
        
        # 初始化权重
        self.weight = nn.Parameter(torch.randn(output_size, self.partition_size))
        self.bias = nn.Parameter(torch.randn(output_size))
    
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """
        列并行前向传播
        
        Args:
            input: 输入张量 [batch_size, input_size]
            
        Returns:
            输出张量 [batch_size, output_size]
        """
        # 将输入分割到各个GPU
        input_parts = TensorParallelism.tensor_parallel_split(input, 1, self.num_partitions)
        
        # 各GPU计算部分结果
        outputs = []
        for i in range(self.num_partitions):
            output_i = torch.matmul(input_parts[i], self.weight.t()) + self.bias
            outputs.append(output_i)
        
        # 使用AllReduce聚合结果
        # 实际实现应该使用分布式通信
        aggregated_output = torch.cat(outputs, dim=1)
        
        return aggregated_output

class RowParallelLinear(nn.Module):
    """行并行线性层（需要AllReduce通信）"""
    
    def __init__(self, 
                 input_size: int, 
                 output_size: int, 
                 num_partitions: int):
        super().__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.num_partitions = num_partitions
        
        # 验证输出维度可分割性
        if output_size % num_partitions != 0:
            raise ValueError("output_size must be divisible by num_partitions")
        
        self.partition_size = output_size // num_partitions
        
        # 初始化权重
        self.weight = nn.Parameter(torch.randn(self.partition_size, input_size))
        self.bias = nn.Parameter(torch.randn(self.partition_size))
    
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """
        行并行前向传播
        
        Args:
            input: 输入张量 [batch_size, input_size]
            
        Returns:
            输出张量 [batch_size, partition_size]
        """
        # 每个GPU计算部分结果
        output = torch.matmul(input, self.weight.t()) + self.bias
        
        # 如果需要完整的输出，需要广播到其他GPU
        # 这里简化处理，实际需要通信
        return output

# 演示张量并行线性层
def demonstrate_tensor_parallel_linear():
    """演示张量并行线性层"""
    print("\n张量并行线性层演示")
    print("-" * 50)
    
    input_size = 256
    output_size = 128
    batch_size = 32
    num_partitions = 4
    
    # 创建输入数据
    input_data = torch.randn(batch_size, input_size)
    
    # 测试列并行线性层
    print("\n列并行线性层:")
    col_parallel_layer = ColumnParallelLinear(input_size, output_size, num_partitions)
    col_output = col_parallel_layer(input_data)
    print(f"输入形状: {input_data.shape}")
    print(f"输出形状: {col_output.shape}")
    
    # 测试行并行线性层
    print("\n行并行线性层:")
    row_parallel_layer = RowParallelLinear(input_size, output_size, num_partitions)
    row_output = row_parallel_layer(input_data)
    print(f"输入形状: {input_data.shape}")
    print(f"输出形状: {row_output.shape}")
    
    # 验证通信模式
    print("\n通信模式:")
    print("列并行层通信模式:", col_parallel_layer.get_communication_pattern())
    print("行并行层通信模式:", row_parallel_layer.get_communication_pattern())

# 运行演示
demonstrate_tensor_parallel_linear()
​

1.2 模型架构的并行化策略

Transformer层的张量并行化

class TensorParallelTransformerLayer(nn.Module):
    """张量并行的Transformer层"""
    
    def __init__(self, 
                 d_model: int, 
                 nhead: int, 
                 d_ff: int,
                 num_partitions: int,
                 parallel_mode: str = 'intermediate'):
        """
        初始化张量并行Transformer层
        
        Args:
            d_model: 模型维度
            nhead: 注意力头数
            d_ff: 前馈网络维度
            num_partitions: 并行分区数
            parallel_mode: 并行模式 ('attention', 'ffn', 'intermediate')
        """
        super().__init__()
        self.d_model = d_model
        self.nhead = nhead
        self.d_ff = d_ff
        self.num_partitions = num_partitions
        self.parallel_mode = parallel_mode
        
        # 根据并行模式设计层结构
        if parallel_mode == 'attention':
            # 注意力层并行
            self.attention = TensorParallelMultiheadAttention(
                d_model, nhead, num_partitions
            )
            self.ffn = nn.Linear(d_model, d_ff)
        elif parallel_mode == 'ffn':
            # 前馈网络并行
            self.attention = nn.MultiheadAttention(d_model, nhead)
            self.ffn = TensorParallelLinear(d_model, d_ff, num_partitions)
        else:  # intermediate
            # 中间层并行
            self.attention = nn.MultiheadAttention(d_model, nhead)
            self.ffn = nn.Linear(d_model, d_ff)
    
    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """前向传播"""
        if self.parallel_mode == 'attention':
            # 注意力层并行
            attn_output = self.attention(x, x, x, mask=mask)
            ffn_output = self.ffn(attn_output)
        elif self.parallel_mode == 'ffn':
            # 前馈网络并行
            attn_output = self.attention(x, x, x, mask=mask)[0]
            ffn_output = self.ffn(attn_output)
        else:
            # 中间层并行
            attn_output = self.attention(x, x, x, mask=mask)[0]
            ffn_output = self.ffn(attn_output)
        
        return ffn_output

class TensorParallelMultiheadAttention(nn.Module):
    """张量并行多头注意力机制"""
    
    def __init__(self, d_model: int, nhead: int, num_partitions: int):
        super().__init__()
        self.d_model = d_model
        self.nhead = nhead
        self.num_partitions = num_partitions
        
        # 验证维度
        if d_model % nhead != 0:
            raise ValueError("d_model must be divisible by nhead")
        
        self.d_k = d_model // nhead
        
        # 并行化查询、键、值矩阵
        self.q_proj = TensorParallelLinear(d_model, d_model, num_partitions, 'col')
        self.k_proj = TensorParallelLinear(d_model, d_model, num_partitions, 'col')
        self.v_proj = TensorParallelLinear(d_model, d_model, num_partitions, 'col')
        
        # 输出投影
        self.out_proj = TensorParallelLinear(d_model, d_model, num_partitions, 'row')
    
    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, 
               mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """多头注意力前向传播"""
        batch_size, seq_len, _ = query.shape
        
        # 查询、键、值投影
        q = self.q_proj(query)  # [batch_size, seq_len, d_model]
        k = self.k_proj(key)    # [batch_size, seq_len, d_model]
        v = self.v_proj(value)  # [batch_size, seq_len, d_model]
        
        # 重塑为多头格式
        q = q.view(batch_size, seq_len, self.nhead, self.d_k).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.nhead, self.d_k).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.nhead, self.d_k).transpose(1, 2)
        
        # 计算注意力分数
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        # 应用mask
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        
        # 计算注意力权重
        attn_weights = torch.softmax(scores, dim=-1)
        
        # 应用注意力权重到值
        attn_output = torch.matmul(attn_weights, v)
        
        # 重塑回原始格式
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, self.d_model)
        
        # 输出投影
        output = self.out_proj(attn_output)
        
        return output

class TensorParallelModel(nn.Module):
    """完整的张量并行模型"""
    
    def __init__(self, 
                 vocab_size: int, 
                 d_model: int, 
                 nhead: int, 
                 num_layers: int,
                 d_ff: int,
                 num_partitions: int,
​