Transformer架构

import torch
import torch.nn as nn
import math

class TransformerBlock(nn.Module):
    """
    标准Transformer块实现
    """
    def __init__(self, d_model=512, n_heads=8, d_ff=2048, dropout=0.1):
        super().__init__()
        
        # 多头注意力
        self.attention = nn.MultiheadAttention(
            d_model, n_heads, dropout=dropout
        )
        
        # 前馈网络
        self.feed_forward = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(d_ff, d_model)
        )
        
        # 层归一化
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, x, mask=None):
        # 自注意力 + 残差连接
        attn_out, _ = self.attention(x, x, x, attn_mask=mask)
        x = self.norm1(x + self.dropout(attn_out))
        
        # 前馈网络 + 残差连接
        ff_out = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_out))
        
        return x

class ModernTransformer(nn.Module):
    """
    包含2024年优化的Transformer
    """
    def __init__(self, config):
        super().__init__()
        
        # 使用RMSNorm代替LayerNorm (Llama风格)
        self.norm = RMSNorm(config.d_model)
        
        # 分组查询注意力 (GQA)
        self.attention = GroupedQueryAttention(
            d_model=config.d_model,
            n_heads=config.n_heads,
            n_kv_heads=config.n_kv_heads,  # KV头数量少于Q头
            use_flash_attn=True  # 使用Flash Attention
        )
        
        # SwiGLU激活的FFN (GPT-4/Llama风格)
        self.feed_forward = SwiGLUFFN(
            d_model=config.d_model,
            d_ff=config.d_ff
        )
        
        # 旋转位置编码
        self.rope = RotaryPositionalEncoding(
            d_model=config.d_model,
            max_seq_len=config.max_seq_len
        )
    
    def forward(self, x, kv_cache=None):
        # 预归一化 (Pre-LN)
        x_norm = self.norm(x)
        
        # 应用RoPE
        x_with_pos = self.rope(x_norm)
        
        # GQA with Flash Attention
        attn_out, new_kv_cache = self.attention(
            x_with_pos, 
            kv_cache=kv_cache,
            use_flash=True
        )
        
        # 残差连接
        x = x + attn_out
        
        # FFN with pre-norm
        x = x + self.feed_forward(self.norm(x))
        
        return x, new_kv_cache

# 1. 原始正弦位置编码 (2017)
def sinusoidal_position_encoding(seq_len, d_model):
    pe = torch.zeros(seq_len, d_model)
    position = torch.arange(seq_len).unsqueeze(1)
    div_term = torch.exp(torch.arange(0, d_model, 2) * 
                        -(math.log(10000.0) / d_model))
    
    pe[:, 0::2] = torch.sin(position * div_term)
    pe[:, 1::2] = torch.cos(position * div_term)
    return pe

# 2. 旋转位置编码 RoPE (2024主流)
class RotaryPositionalEncoding:
    def __init__(self, dim, max_seq_len=8192, base=10000):
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.base = base
        
    def forward(self, q, k):
        # 应用旋转矩阵编码相对位置
        seq_len = q.shape[1]
        
        # 计算旋转角度
        theta = self.compute_theta(seq_len)
        
        # 旋转查询和键
        q_rot = self.rotate(q, theta)
        k_rot = self.rotate(k, theta)
        
        return q_rot, k_rot

class MultiModalTransformer:
    """
    GPT-4o风格的多模态Transformer
    """
    def __init__(self, config):
        # 统一的Transformer主干
        self.transformer = ModernTransformer(config)
        
        # 不同模态的编码器
        self.text_encoder = TextTokenizer(config.vocab_size)
        self.image_encoder = VisionTransformer(config.image_size)
        self.audio_encoder = AudioTransformer(config.audio_dim)
        
        # 统一的嵌入空间
        self.project_text = nn.Linear(config.text_dim, config.d_model)
        self.project_image = nn.Linear(config.image_dim, config.d_model)
        self.project_audio = nn.Linear(config.audio_dim, config.d_model)
    
    def forward(self, text=None, image=None, audio=None):
        embeddings = []
        
        if text is not None:
            text_emb = self.project_text(self.text_encoder(text))
            embeddings.append(text_emb)
        
        if image is not None:
            image_emb = self.project_image(self.image_encoder(image))
            embeddings.append(image_emb)
        
        if audio is not None:
            audio_emb = self.project_audio(self.audio_encoder(audio))
            embeddings.append(audio_emb)
        
        # 拼接所有模态
        multi_modal_input = torch.cat(embeddings, dim=1)
        
        # 通过统一的Transformer处理
        output = self.transformer(multi_modal_input)
        
        return output