parallel_mode: str = 'layer'):
"""
初始化张量并行模型

    Args:
        vocab_size: 词汇表大小
        d_model: 模型维度
        nhead: 注意力头数
        num_layers: 层数
        d_ff: 前馈网络维度
        num_partitions: 并行分区数
        parallel_mode: 并行模式 ('layer', 'tensor')
    """
    super().__init__()
    self.vocab_size = vocab_size
    self.d_model = d_model
    self.nhead = nhead
    self.num_layers = num_layers
    self.num_partitions = num_partitions
    self.parallel_mode = parallel_mode
    
    # 嵌入层
    self.embedding = nn.Embedding(vocab_size, d_model)
    
    # 位置编码
    self.pos_encoding = PositionalEncoding(d_model)
    
    # Transformer层
    self.layers = nn.ModuleList([
        TensorParallelTransformerLayer(
            d_model, nhead, d_ff, num_partitions, 
            parallel_mode='attention' if i % 2 == 0 else 'ffn'
        ) for i in range(num_layers)
    ])
    
    # 输出层
    self.output_proj = TensorParallelLinear(d_model, vocab_size, num_partitions)

def forward(self, input_ids: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
    """模型前向传播"""
    # 嵌入和位置编码
    x = self.embedding(input_ids)
    x = self.pos_encoding(x)
    
    # Transformer层
    for layer in self.layers:
        x = layer(x, mask)
    
    # 输出投影
    output = self.output_proj(x)
    
    return output
​

class PositionalEncoding(nn.Module):
"""位置编码"""

def __init__(self, d_model: int, max_len: int = 5000):
    super().__init__()
    pe = torch.zeros(max_len, d_model)
    position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2).float() * 
                       (-math.log(10000.0) / d_model))
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    pe = pe.unsqueeze(0).transpose(0, 1)
    self.register_buffer('pe', pe)

def forward(self, x: torch.Tensor) -> torch.Tensor:
    return x + self.pe[:x.size(0), :]
​

演示张量并行模型

def demonstrate_tensor_parallel_model():
"""演示张量并行模型"""
print("\n张量并行模型演示")
print("-" * 50)

# 模型参数
vocab_size = 10000
d_model = 512
nhead = 8
num_layers = 6
d_ff = 2048
num_partitions = 4
batch_size = 16
seq_len = 128

# 创建输入
input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))

# 创建张量并行模型
model = TensorParallelModel(
    vocab_size=vocab_size,
    d_model=d_model,
    nhead=nhead,
    num_layers=num_layers,
    d_ff=d_ff,
    num_partitions=num_partitions,
    parallel_mode='layer'
)

# 前向传播
output = model(input_ids)

print(f"输入形状: {input_ids.shape}")
print(f"输出形状: {output.shape}")
print(f"模型参数数量: {sum(p.numel() for p in model.parameters()):,}")

# 测试不同并行模式
print("\n不同并行模式:")
for mode in ['layer', 'tensor']:
    model_mode = TensorParallelModel(
        vocab_size=vocab_size,
        d_model=d_model,
        nhead=nhead,
        num_layers=num_layers,
        d_ff=d_ff,
        num_partitions=num_partitions,
        parallel_mode=mode
    )
    output_mode = model_mode(input_ids)
    print(f"模式 {mode}: 输出形状 {output_mode.shape}")
​

运行演示

demonstrate_tensor_parallel_model()

## 2. 算子拆分与优化

### 2.1 算子级并行策略

**矩阵乘法的高效拆分**

```python
import torch.nn.functional as F
from torch.nn import init

class OptimizedTensorParallelLinear(nn.Module):
    """优化的张量并行线性层"""
    
    def __init__(self, 
                 input_size: int, 
                 output_size: int, 
                 num_partitions: int,
                 bias: bool = True,
                 gather_output: bool = True):
        super().__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.num_partitions = num_partitions
        self.gather_output = gather_output
        
        # 验证维度
        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.empty(self.partition_size, input_size))
        if bias:
            self.bias = nn.Parameter(torch.empty(self.partition_size))
        else:
            self.register_parameter('bias', None)
        
        # 初始化参数
        self.reset_parameters()
    
    def reset_parameters(self):
        """重置参数"""
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            init.uniform_(self.bias, -bound, bound)
    
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """优化的前向传播"""
        # 输入已经在前一层的AllReduce中完成聚合
        output = F.linear(input, self.weight, self.bias)
        
        # 如果需要完整输出，在这里进行AllReduce
        # 实际实现应该使用分布式通信
        if self.gather_output:
            # 这里简化处理，实际应该调用分布式AllReduce
            # output = distributed_all_reduce(output)
            pass
        
        return output

class ColumnParallelLinearV2(nn.Module):
    """优化的列并行线性层"""
    
    def __init__(self, 
                 input_size: int, 
                 output_size: int, 
                 num_partitions: int,
                 bias: bool = True,
                 gather_output: bool = True):
        super().__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.num_partitions = num_partitions
        self.gather_output = gather_output
        
        # 验证维度
        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.empty(output_size, self.partition_size))
        if bias:
            self.bias = nn.Parameter(torch.empty(output_size))
        else:
            self.register_parameter('bias', None)
        
        # 初始化参数
        self.reset_parameters()
    
    def reset_parameters(self):
        """重置参数"""
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            init.uniform_(self.bias, -bound, bound)
    
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """前向传播"""
        # 输入分割
        input_list = torch.split(input, self.partition_size, dim=-1)
        
        # 各GPU计算
        output_list = [F.linear(input_i, self.weight, self.bias) for input_i in input_list]
        
        # 合并输出
        output = torch.cat(output_list, dim=-1)
        
        # 如果需要，执行AllReduce（这里简化）
        # if self.gather_output:
        #     output = distributed_all_reduce(output)
        
        return output

class RowParallelLinearV2(nn.Module):
    """优化的行并行线性层"""
    
    def __init__(self, 
                 input_size: int, 
                 output_size: int, 
                 num_partitions: int,
                 bias: bool = True,
                 input_is_parallel: bool = True):
        super().__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.num_partitions = num_partitions
        self.input_is_parallel = input_is_parallel
        
        # 验证维度
        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.empty(self.partition_size, input_size))
        if bias:
            self.bias = nn.Parameter(torch.empty(self.partition_size))
        else:
            self.register_parameter('bias', None)
        
        # 初始化参数
        self.reset_parameters()
    
    def reset_parameters(self):
        """重置参数"""
        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
        if self.bias is not None:
            fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
            bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
            init.uniform_(self.bias, -bound, bound)
    
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        """前向传播"""
        # 输入可能需要AllReduce（如果是跨GPU的）
        if not self.input_is_parallel:
            # 这里简化处理，实际应该调用分布式AllReduce
            # input = distributed_all_reduce(input)
            pass
        
        # 矩阵乘法
        output = F.linear(input, self.weight, self.bias)
        
        return output

# 演示优化的张量并行线性层
def demonstrate_optimized_tensor_parallel():
    """演示优化的张量并行线性层"""
    print("\n优化的张量并行线性层演示")
    print("-" * 50)
    
    # 参数设置
    input_size = 1024
    output_size = 512
    batch_size = 32
    num_partitions = 4
    
    # 创建输入数据
    input_data = torch.randn(batch_size, input_size)
    
    # 测试列并行层
    print("\n列并行线性层:")
    col_parallel = ColumnParallelLinearV2(input_size, output_size, num_partitions)
    col_output = col_parallel(input_data)
    print(f"输入形状: {input_data.shape}")
    print(f"输出形状: {col_output.shape}")
    
    # 测试行并行层
    print("\n行并行线性层:")
    row_parallel = RowParallelLinearV2(input_size, output_size, num_partitions)
    row_output = row_parallel(input_data)
    print(f"输入形状: {input_data.shape}")
    print(f"输出形状: {row_output.shape}")
    
    # 性能对比
    print("\n性能测试:")
    import time
    
    # 测试原始实现
    start_time = time.time()
    for _ in range(100):
        _ = col_parallel(input_data)
    col_time = time.time() - start_time
    
    # 测试优化实现
    start_time = time.time()
    for _ in range(100):
        _ = ColumnParallelLinearV2(input_size, output_size, num_partitions)(input_data)
    opt_time = time.time() - start_time
    
    print(f"原始实现时间: {col_time:.4f}s")
    print(f"优化实现时间: {opt_time:.4f}s")
    print(f"性能提升: {col_time/opt_time:.2f}x")

# 运行演示
demonstrate_optimized_tensor_parallel()
​

注意力的并行化优化

class OptimizedTensorParallelAttention(nn.Module):
    """优化的张量并行注意力机制"""
    
    def __init__(self, 
                 d_model: int, 
                 nhead: int, 
                 num_partitions: int,
                 dropout: float = 0.1):
        super().__init__()
        self.d_model = d_model
        self.nhead = nhead
        self.num_partitions = num_partitions
        self.dropout = dropout
        
        # 验证维度
        if d_model % nhead != 0:
            raise ValueError("d_model must be divisible by nhead")
        
        self.d_k = d_model // nhead
        
        # 并行化投影层
        self.q_proj = ColumnParallelLinearV2(d_model, d_model, num_partitions)
        self.k_proj = ColumnParallelLinearV2(d_model, d_model, num_partitions)
        self.v_proj = ColumnParallelLinearV2(d_model, d_model, num_partitions)
        
        # 输出投影
        self.out_proj = RowParallelLinearV2(d_model, d_model, num_partitions)
        
        # Dropout层
        self.dropout_layer = nn.Dropout(dropout)
    
    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_weights = self.dropout_layer(attn_weights)
        
        # 应用注意力权重到值
        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 FlashAttentionOptimized(nn.Module):
    """优化的Flash注意力机制"""
    
    def __init__(self, 
                 d_model: int, 
                 nhead: int, 
                 num_partitions: int,
                 dropout: float = 0.1):
        super().__init__()
        self.d_model = d_model
        self.nhead = nhead
        self.num_partitions = num_partitions
        self.dropout = dropout
        
        # 验证维度
        if d_model % nhead != 0:
            raise ValueError("d_model must be divisible by nhead")
        
        self.d_k = d_model // nhead
        
        # 优化后的并行化投影层
        self.q_proj = OptimizedTensorParallelLinear(d_model, d_model, num_partitions)
        self.k_proj = OptimizedTensorParallelLinear(d_model, d_model, num_partitions)
        self.v_proj = OptimizedTensorParallelLinear(d_model, d_model, num_partitions)
​