Upload 20 files

components.py

@@ -0,0 +1,387 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Tuple, Optional, Union
+import math
+
+class YARNScaling:
+    """
+    YARN (Yet Another RoPE extensioN) 缩放策略
+    实现参考: https://arxiv.org/abs/2309.00071
+    """
+    @staticmethod
+    def compute_yarn_parameters(
+        original_max_len: int,
+        target_max_len: int=8192,
+        dim: int=128,
+        base: int = 10000,
+        beta_fast: int = 32,
+        beta_slow: int = 1,
+        alpha: float = 1.0,
+        device: Optional[torch.device] = None
+    ) -> Tuple[torch.Tensor, float]:
+        scale = float(target_max_len) / original_max_len
+        mscale = YARNScaling.compute_mscale(scale, alpha)
+
+        # 确保 dim 为 float 以进行除法运算
+        # RoPE 频率是成对的 (0, 2, ..., d-2)
+        freqs_idx = torch.arange(0, dim, 2, dtype=torch.float32, device=device)
+        
+        # 基础频率 (Original RoPE)
+        freq_extra = 1.0 / (base ** (freqs_idx / dim))
+        
+        # 如果不需要缩放，直接返回基础频率
+        if scale <= 1.0:
+            return freq_extra, 1.0
+
+        # 插值频率 (Interpolated for extension)
+        freq_inter = 1.0 / (scale * base ** (freqs_idx / dim))
+        
+        # 计算 YARN 阈值 (基于波长/索引)
+        # 对应 paper 中的 band constraints
+        # 这里的公式将频率索引 i 映射到阈值
+        def get_limit(beta):
+            return dim * math.log(original_max_len / (2 * math.pi * beta)) / (2 * math.log(base))
+            
+        low = max(math.floor(get_limit(beta_fast)), 0)
+        high = min(math.ceil(get_limit(beta_slow)), dim // 2 - 1)
+        
+        # indices: 0, 1, ..., dim/2 - 1
+        indices = torch.arange(0, dim // 2, dtype=torch.float32, device=device)
+        
+        inv_freq = freq_extra.clone()
+        
+        # 1. 低频部分 (Long wavelengths, Indices > high): 使用插值频率
+        # 这些频率对应的波长已经超过了原始上下文长度，需要拉伸
+        mask_low_freq = indices > high
+        inv_freq[mask_low_freq] = freq_inter[mask_low_freq]
+        
+        # 2. 高频部分 (Short wavelengths, Indices < low): 保持原频率 (freq_extra)
+        # 这些部分受旋转不变性保护，不需要插值
+        
+        # 3. 中间部分: 线性平滑混合 (Ramp function)
+        mid_mask = (indices >= low) & (indices <= high)
+        if mid_mask.any():
+            # 避免除以 0
+            denom = max(high - low, 1)
+            t = (indices[mid_mask] - low) / denom
+            inv_freq[mid_mask] = freq_extra[mid_mask] * (1 - t) + freq_inter[mid_mask] * t
+            
+        return inv_freq, float(mscale)
+
+    @staticmethod
+    def compute_mscale(scale: float, alpha: float = 1.0) -> float:
+        """计算注意力缩放因子 (Temperature scaling)"""
+        if scale <= 1.0:
+            return 1.0
+        # 0.1 * ln(scale) + 1.0 是经验公式，用于修正熵值
+        return 0.1 * math.log(scale) + 1.0
+
+class YARNRotaryEmbedding(nn.Module):
+    """
+    集成 YARN 的旋转位置编码
+    修复了精度问题、缓存管理以及 position_ids 越界问题
+    """
+    def __init__(
+        self,
+        dim: int = 64,
+        max_seq_len: int = 8192,
+        original_max_len: int = 4096,
+        base: int = 10000,
+        scaling_factor: float = 1.0, # 预留接口，暂未使用，由 yarn 逻辑控制
+        beta_fast: int = 32,
+        beta_slow: int = 1,
+        alpha: float = 1.0,
+        rope_percentage: float = 1.0,
+        device: Optional[torch.device] = None
+    ):
+        super().__init__()
+        self.dim = dim
+        self.max_seq_len = max_seq_len
+        self.original_max_len = original_max_len
+        self.base = base
+        self.alpha = alpha
+        
+        # 计算实际应用 RoPE 的维度
+        self.rope_dim = int(dim * rope_percentage)
+        # 确保是偶数
+        if self.rope_dim % 2 != 0:
+            self.rope_dim -= 1 
+        
+        # 初始化频率 (Persistent state)
+        self._init_yarn_frequencies(device)
+        
+        # 缓存 cos/sin (Transient state)
+        # persistent=False 意味着不会保存到 state_dict，减少 checkpoint 大小
+        self.register_buffer("cos_cached", None, persistent=False)
+        self.register_buffer("sin_cached", None, persistent=False)
+
+    def _init_yarn_frequencies(self, device: Optional[torch.device] = None):
+        """初始化 YARN 频率"""
+        inv_freq, mscale = YARNScaling.compute_yarn_parameters(
+            self.original_max_len,
+            self.max_seq_len,
+            self.rope_dim,
+            self.base,
+            beta_fast=32, # 这里通常使用默认值或传入参数，此处修正为使用硬编码默认值保持一致，或应改为 self.beta_fast
+            beta_slow=1,
+            alpha=self.alpha,
+            device=device
+        )
+        # 注册 buffer
+        self.register_buffer("inv_freq", inv_freq, persistent=True)
+        self.register_buffer("mscale", torch.tensor(mscale, dtype=torch.float32, device=device), persistent=True)
+
+    def _compute_cos_sin_cache(
+        self,
+        needed_len: int,
+        device: torch.device,
+        dtype: torch.dtype
+    ):
+        """预计算 cos 和 sin 缓存，始终使用 float32 计算以保证精度"""
+        # 至少分配 max_seq_len，如果 needed_len 更大则扩展
+        alloc_len = max(needed_len, self.max_seq_len)
+        
+        # 如果已有缓存且足够大且设备匹配，则不重新计算 (可选优化，这里选择简单逻辑：不够就重算)
+        if (self.cos_cached is not None and 
+            self.cos_cached.shape[2] >= alloc_len and 
+            self.cos_cached.device == device):
+            return
+
+        t = torch.arange(alloc_len, dtype=torch.float32, device=device)
+        
+        # freqs: [alloc_len, dim // 2]
+        # outer product: t[i] * inv_freq[j]
+        freqs = torch.outer(t, self.inv_freq.to(device))
+        
+        # 拼接以匹配 rotate_half 的逻辑: [theta_0, theta_1, ..., theta_0, theta_1, ...]
+        emb = torch.cat((freqs, freqs), dim=-1)
+        
+        # 应用 mscale 并计算 cos/sin
+        # [alloc_len, rope_dim] -> [1, 1, alloc_len, rope_dim] 用于广播
+        cos_cached = (emb.cos() * self.mscale).view(1, 1, alloc_len, self.rope_dim)
+        sin_cached = (emb.sin() * self.mscale).view(1, 1, alloc_len, self.rope_dim)
+        
+        self.cos_cached = cos_cached.to(dtype) # 缓存可以存为半精度以省显存，但计算时建议 float32
+        self.sin_cached = sin_cached.to(dtype)
+
+    @staticmethod
+    def rotate_half(x: torch.Tensor) -> torch.Tensor:
+        """
+        旋转输入的后半部分
+        Input: [..., d] -> Split into x1, x2 -> Output [-x2, x1]
+        """
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+
+    def apply_rotary_pos_emb(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        position_ids: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """应用 RoPE，包含精度修正和边界检查"""
+        bsz, num_heads, seq_len, head_dim = q.shape
+        
+        # 1. 确定需要的缓存长度
+        if position_ids is not None:
+            # 必须覆盖 position_ids 中的最大索引
+            max_pos = position_ids.max().item() + 1
+            needed_len = max(max_pos, seq_len)
+        else:
+            needed_len = seq_len
+
+        # 2. 检查并更新缓存
+        if (self.cos_cached is None or 
+            self.cos_cached.shape[2] < needed_len or 
+            self.cos_cached.device != q.device):
+            self._compute_cos_sin_cache(needed_len, q.device, q.dtype)
+            
+        # 3. 获取对应的 cos/sin
+        # cos_cached: [1, 1, alloc_len, dim]
+        if position_ids is not None:
+            # position_ids: [bs, seq_len]
+            # 选取对应的 pos embedding -> [bs, 1, seq_len, dim]
+            # 注意: cos_cached[0, 0] 形状为 [alloc_len, dim]
+            cos = self.cos_cached[0, 0][position_ids].unsqueeze(1)
+            sin = self.sin_cached[0, 0][position_ids].unsqueeze(1)
+        else:
+            # 默认假设从 0 开始
+            cos = self.cos_cached[:, :, :seq_len, :]
+            sin = self.sin_cached[:, :, :seq_len, :]
+        
+        # 4. 处理部分 RoPE (如果 rope_dim < head_dim)
+        if self.rope_dim < head_dim:
+            q_rot = q[..., :self.rope_dim]
+            q_pass = q[..., self.rope_dim:]
+            k_rot = k[..., :self.rope_dim]
+            k_pass = k[..., self.rope_dim:]
+        else:
+            q_rot = q
+            k_rot = k
+            q_pass = None
+            k_pass = None
+        
+        # 5. 执行旋转 (强制 float32 计算以避免精度溢出)
+        q_rot_float = q_rot.float()
+        k_rot_float = k_rot.float()
+        cos_float = cos.float()
+        sin_float = sin.float()
+        
+        q_embed = (q_rot_float * cos_float) + (self.rotate_half(q_rot_float) * sin_float)
+        k_embed = (k_rot_float * cos_float) + (self.rotate_half(k_rot_float) * sin_float)
+        
+        # 6. 转回原始类型
+        q_embed = q_embed.type_as(q)
+        k_embed = k_embed.type_as(k)
+        
+        if q_pass is not None:
+            q_embed = torch.cat([q_embed, q_pass], dim=-1)
+            k_embed = torch.cat([k_embed, k_pass], dim=-1)
+        
+        return q_embed, k_embed
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        position_ids: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return self.apply_rotary_pos_emb(q, k, position_ids)
+
+    def extra_repr(self) -> str:
+        return (f"dim={self.dim}, rope_dim={self.rope_dim}, "
+                f"max_seq_len={self.max_seq_len}, original_max_len={self.original_max_len}, "
+                f"base={self.base}")
+
+class RMSNorm(nn.Module):
+    """
+    Root Mean Square Layer Normalization
+    包含 float32 强制转换以确保数值稳定性
+    """
+    def __init__(
+        self,
+        dim: int,
+        eps: float = 1e-6,
+        elementwise_affine: bool = True
+    ):
+        super().__init__()
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        
+        if self.elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim))
+        else:
+            self.register_parameter('weight', None)
+
+    def _norm(self, x: torch.Tensor) -> torch.Tensor:
+        # 始终在 float32 下计算 RMS，防止 FP16 下溢或溢出
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # 1. 转换为 float32 进行统计量计算
+        output = self._norm(x.float())
+        
+        # 2. 转回原始类型
+        output = output.type_as(x)
+        
+        # 3. 应用权重 (如果存在)
+        if self.elementwise_affine and self.weight is not None:
+            output = output * self.weight
+        
+        return output
+
+class QKNorm(nn.Module):
+    """
+    Query-Key Normalization (ViT-22B / Scaling Transformer)
+    用于稳定注意力矩阵的 logits
+    """
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.query_norm = RMSNorm(dim, eps=eps)
+        self.key_norm = RMSNorm(dim, eps=eps)
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        q = self.query_norm(q)
+        k = self.key_norm(k)
+        return q, k
+
+class SwiGLU(nn.Module):
+    """
+    SwiGLU 激活前馈网络
+    结构: Down(SiLU(Gate) * Up)
+    """
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: Optional[int] = None,
+        multiple_of: int = 256,
+        ffn_dim_multiplier: Optional[float] = None,
+        dropout: float = 0.0,
+        bias: bool = False
+    ):
+        super().__init__()
+        
+        if hidden_dim is None:
+            if ffn_dim_multiplier is not None:
+                hidden_dim = int(dim * ffn_dim_multiplier)
+            else:
+                # 默认: 2/3 * 4 * dim = 8/3 * dim (LLaMA standard)
+                hidden_dim = int(2 * dim * 4 / 3)
+            
+            # 确保 hidden_dim 是 multiple_of 的倍数 (通常为了 GPU 核心优化)
+            hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+        
+        self.hidden_dim = hidden_dim
+        
+        # W1: Gate, W3: Up, W2: Down (Standard LLaMA naming conventions)
+        self.w1 = nn.Linear(dim, hidden_dim, bias=bias)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=bias)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=bias)
+        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # SwiGLU(x) = (SiLU(W1·x) ⊙ W3·x) · W2
+        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
+
+class ParallelAttentionFFN(nn.Module):
+    """
+    并行注意力与前馈网络 (PaLM / GPT-J 风格)
+    y = x + Attention(LN(x)) + MLP(LN(x))
+    """
+    def __init__(
+        self,
+        dim: int,
+        attn_module: nn.Module,
+        ffn_module: nn.Module,
+        norm_eps: float = 1e-6
+    ):
+        super().__init__()
+        # 注意: 某些架构(如 PaLM)可能共用一个 LayerNorm，
+        # 但这里为了灵活性保留两个独立的 Norm (如 CodeLlama 某些变体)
+        self.attn_norm = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm = RMSNorm(dim, eps=norm_eps)
+        self.attn = attn_module
+        self.ffn = ffn_module
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        **attn_kwargs
+    ) -> torch.Tensor:
+        # 并行计算：从同一个 x (normalize 后) 分叉
+        attn_input = self.attn_norm(x)
+        ffn_input = self.ffn_norm(x)
+        
+        # 计算注意力
+        attn_out = self.attn(attn_input, **attn_kwargs)
+        
+        # 计算 FFN (确保不传递 attn 特定的 kwargs)
+        ffn_out = self.ffn(ffn_input)
+        
+        # 一次性残差连接
+        return x + attn_out + ffn_out

continual_learning.py

@@ -0,0 +1,294 @@
+"""
+持续学习模块
+支持EWC和经验回放
+修复版本：适配 MultiModalDenseTransformer 和 data_loader.py
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from torch.utils.data import DataLoader
+from collections import deque
+from typing import List, Dict, Any, Optional, Union
+from tqdm import tqdm
+from dataclasses import dataclass
+
+# 假设 model.py 中已有定义，用于类型提示
+# from model import MultiModalDenseTransformer
+
+@dataclass
+class ModalityConfig:
+    name: str
+    modality_id: int
+
+class UnifiedMultiModalPreprocessor(nn.Module):
+    """
+    统一多模态预处理器
+    职责：仅负责将原始Batch数据格式化为 MultiModalDenseTransformer 接受的 'segments' 结构。
+    不再包含编码器，编码工作交由模型自身完成，以确保 EWC 能够捕捉模型参数的梯度。
+    """
+    def __init__(self, model_dim: int = 2048):
+        super().__init__()
+        self.modality_configs = {
+            'text': ModalityConfig('text', 0),
+            'image': ModalityConfig('image', 1),
+            'audio': ModalityConfig('audio', 2),
+            'video': ModalityConfig('video', 3)
+        }
+
+    def process_batch(self, batch_data: Union[torch.Tensor, List[Any]], modality_type: str) -> List[Dict]:
+        """
+        将特定模态的数据封装为 segment 格式
+        """
+        processed_segments = []
+        if modality_type not in self.modality_configs:
+            return processed_segments
+
+        config = self.modality_configs[modality_type]
+        
+        # 确保数据是 Tensor 格式
+        if isinstance(batch_data, list):
+            # 过滤 None
+            valid_data = [x for x in batch_data if x is not None]
+            if not valid_data:
+                return []
+            # 假设 list 中全是 Tensor，且维度一致，进行堆叠
+            # 如果是 list of tensor (B, C, H, W) -> stack -> (B, C, H, W)
+            try:
+                data_tensor = torch.stack(valid_data)
+            except Exception as e:
+                print(f"Error stacking modality data: {e}")
+                return []
+        elif isinstance(batch_data, torch.Tensor):
+            data_tensor = batch_data
+        else:
+            return []
+
+        processed_segments.append({
+            'type': modality_type,
+            'data': data_tensor, # 保持原始数据 (如图片像素)，模型内部会encode
+            'modality_id': config.modality_id
+        })
+        return processed_segments
+
+
+class ExperienceReplayBuffer:
+    """经验回放缓冲区 - 内存安全版"""
+    def __init__(self, max_size: int = 10000):
+        self.buffer = deque(maxlen=max_size)
+
+    def add(self, sample: Dict[str, Any]):
+        """
+        添加样本到buffer
+        关键修复：将数据移至 CPU 并 detach，防止显存泄漏
+        """
+        safe_sample = {}
+        for k, v in sample.items():
+            if isinstance(v, torch.Tensor):
+                safe_sample[k] = v.detach().cpu()
+            elif isinstance(v, list):
+                # 递归处理 list 中的 tensor
+                safe_sample[k] = [x.detach().cpu() if isinstance(x, torch.Tensor) else x for x in v]
+            else:
+                safe_sample[k] = v
+        self.buffer.append(safe_sample)
+
+    def sample(self, batch_size: int) -> List[Any]:
+        """从buffer中采样"""
+        if not self.buffer:
+            return []
+        
+        indices = np.random.choice(
+            len(self.buffer),
+            min(len(self.buffer), batch_size),
+            replace=False
+        )
+        return [self.buffer[i] for i in indices]
+
+    def __len__(self):
+        return len(self.buffer)
+
+    def clear(self):
+        """清空buffer"""
+        self.buffer.clear()
+
+
+class EWC:
+    """弹性权重固化 (Elastic Weight Consolidation)"""
+    def __init__(
+        self,
+        model: nn.Module,
+        dataloader: DataLoader,
+        preprocessor: UnifiedMultiModalPreprocessor,
+        importance: float = 1000.0
+    ):
+        self.model = model
+        self.preprocessor = preprocessor
+        self.importance = importance
+        self.device = next(model.parameters()).device
+        
+        # 冻结当前参数作为参考
+        self.params = {
+            n: p.clone().detach()
+            for n, p in model.named_parameters()
+            if p.requires_grad
+        }
+        
+        self.fisher = self._compute_fisher(dataloader)
+
+    def _compute_fisher(self, dataloader: DataLoader) -> Dict[str, torch.Tensor]:
+        """计算Fisher信息矩阵 (使用 Empirical Fisher)"""
+        fisher = {
+            n: torch.zeros_like(p)
+            for n, p in self.model.named_parameters()
+            if p.requires_grad
+        }
+        
+        self.model.eval()
+        num_samples = 0
+        
+        # 使用 tqdm 稍微简化输出
+        pbar = tqdm(dataloader, desc="Computing Fisher Matrix", leave=False)
+        for batch in pbar:
+            if batch is None: continue
+
+            self.model.zero_grad()
+            
+            # 1. 准备文本输入
+            instruction_ids = batch['instruction'].to(self.device)
+            response_ids = batch['response'].to(self.device)
+            # 拼接: [Instruction, Response]
+            input_ids = torch.cat([instruction_ids, response_ids], dim=1)
+            
+            # 2. 准备多模态输入结构
+            input_data = {'segments': []}
+            
+            # 处理额外的模态数据 (如果有)
+            # 这里的 batch['modality_data'] 可能是 list (由 collate_fn_v2 生成)
+            raw_modality_data = batch.get('modality_data')
+            if raw_modality_data is not None:
+                # 尝试判断模态类型，如果 dataset 中没有明确指定，默认尝试 'image'
+                # 实际应用中建议 dataset 返回 'modality_type'
+                modality_type = batch.get('modality_type', 'image') 
+                if isinstance(modality_type, list): modality_type = modality_type[0]
+                
+                # Preprocessor 处理数据堆叠和格式化
+                mod_segments = self.preprocessor.process_batch(raw_modality_data, modality_type)
+                # 只有在数据有效时才传给 device
+                for seg in mod_segments:
+                    seg['data'] = seg['data'].to(self.device)
+                    input_data['segments'].append(seg)
+            
+            # 添加文本 Segment
+            input_data['segments'].append({
+                'type': 'text',
+                'data': input_ids,
+                'modality_id': 0
+            })
+            
+            # 3. 前向传播
+            output = self.model(input_data)
+            logits = output['logits'] # (B, Seq_Len, Vocab)
+            
+            # 4. 计算 Loss (Standard Causal LM Loss)
+            # Shift logits and labels
+            # input_ids: [I1, I2, R1, R2]
+            # labels:    [I2, R1, R2, EOS]
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = input_ids[:, 1:].contiguous()
+            
+            # 创建 Mask: 只在 Response 部分计算梯度
+            # Instruction 长度
+            inst_len = instruction_ids.shape[1]
+            loss_mask = torch.ones_like(shift_labels, dtype=torch.float)
+            if inst_len > 1:
+                # 掩盖 Instruction 部分 (注意 shift 后的索引偏移)
+                loss_mask[:, :inst_len-1] = 0.0
+
+            # 计算逐个 Token 的 Loss
+            loss_fct = nn.CrossEntropyLoss(reduction='none')
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+            
+            # 应用 Mask 并求平均
+            loss = (loss * loss_mask.view(-1)).sum() / (loss_mask.sum() + 1e-6)
+            
+            # 5. 反向传播累积梯度平方
+            loss.backward()
+            
+            for n, p in self.model.named_parameters():
+                if p.grad is not None and n in fisher:
+                    fisher[n] += p.grad.detach() ** 2
+            
+            num_samples += input_ids.size(0)
+        
+        # 平均化
+        if num_samples > 0:
+            for n in fisher:
+                fisher[n] /= num_samples
+        
+        self.model.train()
+        return fisher
+
+    def penalty(self, model: Optional[nn.Module] = None) -> torch.Tensor:
+        """计算EWC惩罚项"""
+        # 兼容性处理：如果传入了 model 参数，优先使用（通常 self.model 就是同一个）
+        target_model = model if model is not None else self.model
+        
+        loss = torch.tensor(0.0, device=self.device)
+        
+        for n, p in target_model.named_parameters():
+            if n in self.params and p.requires_grad:
+                if n in self.fisher:
+                    loss += (self.fisher[n] * (p - self.params[n]) ** 2).sum()
+        
+        return self.importance * loss
+
+
+class OnlineEWC(EWC):
+    """在线EWC - 支持持续更新Fisher矩阵"""
+    def __init__(
+        self,
+        model: nn.Module,
+        preprocessor: UnifiedMultiModalPreprocessor,
+        importance: float = 1000.0,
+        gamma: float = 0.9
+    ):
+        # 初始时不计算 Fisher，等待 update_fisher 调用
+        self.model = model
+        self.preprocessor = preprocessor
+        self.importance = importance
+        self.gamma = gamma
+        self.device = next(model.parameters()).device
+        
+        self.params = {}
+        self.fisher = {}
+        self.task_count = 0
+
+    def update_fisher(self, dataloader: DataLoader):
+        """更新Fisher信息矩阵"""
+        print(f"Updating Online EWC Fisher Matrix (Task {self.task_count + 1})...")
+        new_fisher = self._compute_fisher(dataloader)
+        
+        if self.task_count == 0:
+            self.fisher = new_fisher
+        else:
+            for n in self.fisher:
+                if n in new_fisher:
+                    # 移动平均更新 Fisher 信息
+                    self.fisher[n] = self.gamma * self.fisher[n] + new_fisher[n]
+        
+        # 更新参考参数为当前任务训练后的参数
+        self.params = {
+            n: p.clone().detach()
+            for n, p in self.model.named_parameters()
+            if p.requires_grad
+        }
+        
+        self.task_count += 1
+        print(f"Online EWC regularizer updated.")
+
+    def penalty(self, model: Optional[nn.Module] = None) -> torch.Tensor:
+        """计算EWC惩罚项"""
+        if self.task_count == 0:
+            return torch.tensor(0.0, device=self.device)
+        return super().penalty(model)

contrastive_learning.py

@@ -0,0 +1,339 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, Optional, Tuple, Union, Literal, List
+import math
+import copy
+
+class CLIPLoss(nn.Module):
+    """CLIP风格的对比学习损失"""
+    def __init__(self, temperature: float = 0.07, max_temperature: float = 100.0):
+        super().__init__()
+        self.temperature = temperature
+        self.max_temperature = max_temperature
+        # 初始化 logit_scale
+        self.logit_scale = nn.Parameter(torch.ones([]) * math.log(1 / temperature))
+
+    def forward(
+        self,
+        image_features: torch.Tensor,
+        text_features: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            image_features: [B, D]
+            text_features: [B, D]
+        """
+        # 归一化
+        image_features = F.normalize(image_features, dim=-1)
+        text_features = F.normalize(text_features, dim=-1)
+        
+        # 限制 logit_scale 防止数值不稳定
+        logit_scale = self.logit_scale.exp().clamp(max=self.max_temperature)
+        
+        # 计算相似度矩阵 [B, B]
+        # 注意：在 DDP 环境下，这里计算的是局部 Batch 的 Loss。
+        # 完整的 DDP 实现需要 gather 所有 GPU 的 features。
+        logits_per_image = logit_scale * image_features @ text_features.T
+        logits_per_text = logits_per_image.T
+        
+        # 标签: 对角线为正样本
+        batch_size = image_features.shape[0]
+        labels = torch.arange(batch_size, device=image_features.device)
+        
+        # 双向交叉熵
+        loss_i2t = F.cross_entropy(logits_per_image, labels)
+        loss_t2i = F.cross_entropy(logits_per_text, labels)
+        
+        total_loss = (loss_i2t + loss_t2i) / 2
+        
+        return total_loss, loss_i2t, loss_t2i
+
+class SigLIPLoss(nn.Module):
+    """
+    SigLIP损失 - 包含可学习的 Bias 和 Scale
+    Paper: Sigmoid Loss for Language Image Pre-Training
+    """
+    def __init__(self, init_temperature: float = 1.0, init_bias: float = -10.0):
+        super().__init__()
+        self.t_prime = nn.Parameter(torch.tensor(math.log(init_temperature))) 
+        self.b = nn.Parameter(torch.tensor(init_bias))
+
+    def forward(
+        self,
+        image_features: torch.Tensor,
+        text_features: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        注意：SigLIP 的标准实现不需要 Gather 全局负样本即可收敛，
+        但这里实现的是 dense pair loss。对于超大 Batch (如 8k+)，
+        构造 [B, B] 的 labels 矩阵会导致显存爆炸，生产环境建议使用 custom kernel 或 block chunking。
+        """
+        # 归一化
+        image_features = F.normalize(image_features, dim=-1)
+        text_features = F.normalize(text_features, dim=-1)
+        
+        batch_size = image_features.shape[0]
+        
+        # Logits = exp(t) * (x @ yT) + b
+        logits = image_features @ text_features.T * self.t_prime.exp() + self.b
+        
+        # 构造标签: 对角线为1，其余为-1
+        labels = -torch.ones(batch_size, batch_size, device=image_features.device)
+        labels += 2 * torch.eye(batch_size, device=image_features.device)
+        
+        # Sigmoid Loss: -log(sigmoid(label * logits))
+        # 当 label=1: -log(sigmoid(z))
+        # 当 label=-1: -log(sigmoid(-z)) = -log(1 - sigmoid(z))
+        # 这就是标准的 Binary Cross Entropy (Summed)
+        
+        # SigLIP 论文中通常建议除以 batch_size (或正样本数量) 进行归一化
+        loss = -F.logsigmoid(labels * logits).sum() / batch_size
+        
+        return loss
+
+class InfoNCELoss(nn.Module):
+    """InfoNCE损失 - 支持显式负样本或 Batch 内负样本"""
+    def __init__(self, temperature: float = 0.07):
+        super().__init__()
+        self.temperature = temperature
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        positive_key: torch.Tensor,
+        negative_keys: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Args:
+            query: [B, D]
+            positive_key: [B, D]
+            negative_keys: [B, N, D] or None.
+        """
+        query = F.normalize(query, dim=-1)
+        positive_key = F.normalize(positive_key, dim=-1)
+        
+        if negative_keys is not None:
+            # 显式负样本
+            # pos_sim: [B]
+            pos_sim = (query * positive_key).sum(dim=-1) / self.temperature
+            
+            negative_keys = F.normalize(negative_keys, dim=-1)
+            # neg_sim: [B, N]
+            neg_sim = (query.unsqueeze(1) * negative_keys).sum(dim=-1) / self.temperature
+            
+            # [B, 1 + N]
+            logits = torch.cat([pos_sim.unsqueeze(1), neg_sim], dim=1)
+            # 正样本在索引0
+            labels = torch.zeros(query.shape[0], dtype=torch.long, device=query.device)
+        else:
+            # Batch内负样本 (类似于 CLIP 的单向 Loss)
+            logits = query @ positive_key.T / self.temperature
+            labels = torch.arange(query.shape[0], dtype=torch.long, device=query.device)
+        
+        loss = F.cross_entropy(logits, labels)
+        return loss
+
+class ProjectionHead(nn.Module):
+    """
+    投影头：处理特征维度变换和形状适配
+    针对 Transformer 输出 (Sequence) 提供了更精细的 Pooling 控制。
+    """
+    def __init__(
+        self, 
+        input_dim: int, 
+        embed_dim: int, 
+        pooling_type: Literal['cls', 'mean', 'max', 'none'] = 'mean',
+        exclude_first_token: bool = False
+    ):
+        super().__init__()
+        self.pooling_type = pooling_type
+        self.exclude_first_token = exclude_first_token
+        
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, embed_dim),
+            nn.GELU(),
+            nn.Linear(embed_dim, embed_dim)
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # 适配 3D 张量 [B, Seq, D] -> [B, D]
+        if x.dim() == 3:
+            if self.pooling_type == 'cls':
+                # 假设索引0是CLS token (Standard ViT / BERT)
+                x = x[:, 0, :]
+            
+            elif self.pooling_type == 'mean':
+                if self.exclude_first_token and x.shape[1] > 1:
+                    # 对于 ViT，如果使用 mean pooling，通常需要排除 CLS token
+                    x = x[:, 1:, :].mean(dim=1)
+                else:
+                    x = x.mean(dim=1)
+            
+            elif self.pooling_type == 'max':
+                if self.exclude_first_token and x.shape[1] > 1:
+                    x = x[:, 1:, :].max(dim=1)[0]
+                else:
+                    x = x.max(dim=1)[0]
+                    
+            elif self.pooling_type == 'none':
+                # 保留序列维度，适用于 Dense Prediction 或细粒度对比
+                # 此时输出为 [B, Seq, embed_dim]
+                pass
+        
+        return self.net(x)
+
+class MultiModalContrastiveLoss(nn.Module):
+    """多模态对比学习损失 - 支持动态模态和异构维度"""
+    def __init__(
+        self,
+        embed_dim: int = 512,
+        input_dims: Union[int, Dict[str, int]] = 2048,
+        temperature: float = 0.07,
+        loss_type: str = 'clip',
+        modality_config: Optional[Dict[str, str]] = None
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.loss_type = loss_type
+        
+        if loss_type == 'clip':
+            self.loss_fn = CLIPLoss(temperature)
+        elif loss_type == 'siglip':
+            self.loss_fn = SigLIPLoss()
+        else:
+            self.loss_fn = InfoNCELoss(temperature)
+        
+        self.projectors = nn.ModuleDict()
+        
+        if modality_config is None:
+            # 默认常用模态配置
+            # 注意：ImprovedVisionTransformer 输出带 CLS，所以图像推荐用 'cls' 或带排除的 'mean'
+            modality_config = {
+                'text': 'cls',   
+                'image': 'cls',  
+                'audio': 'mean', # AudioEncoder 的双流输出已经是 2D，但如果是纯 Transformer 输出则是 3D
+                'video': 'mean'  # VideoEncoder 输出通常是 [B, T, D]
+            }
+        
+        self.modality_config = modality_config
+        
+        # 初始化投影头
+        for mod_name, pool_type in modality_config.items():
+            dim = 0
+            if isinstance(input_dims, dict):
+                dim = input_dims.get(mod_name)
+                # 如果字典里没给这个模态的维度，跳过初始化，避免 crash
+                if dim is None:
+                    continue
+            else:
+                dim = input_dims
+            
+            # 特殊处理：如果是 'mean' 或 'max' 且是 image/text，可能需要排除 CLS
+            # 这里做一个启发式判断，用户也可以手动修改
+            exclude_first = False
+            if mod_name in ['image', 'text'] and pool_type in ['mean', 'max']:
+                exclude_first = True
+
+            self.projectors[mod_name] = ProjectionHead(
+                input_dim=dim,
+                embed_dim=embed_dim,
+                pooling_type=pool_type,
+                exclude_first_token=exclude_first
+            )
+
+    def forward(
+        self,
+        features: Dict[str, torch.Tensor],
+        modality_pairs: Optional[List[Tuple[str, str]]] = None
+    ) -> Dict[str, torch.Tensor]:
+        
+        # 自动生成对比对：将所有非Text模态与Text对比
+        if modality_pairs is None:
+            if 'text' in features:
+                modality_pairs = [
+                    (mod, 'text') for mod in features.keys() if mod != 'text'
+                ]
+            else:
+                return {}
+        
+        losses = {}
+        
+        for mod_a, mod_b in modality_pairs:
+            if mod_a not in features or mod_b not in features:
+                continue
+            
+            if mod_a not in self.projectors or mod_b not in self.projectors:
+                # 记录警告或跳过
+                continue
+
+            feat_a = self.projectors[mod_a](features[mod_a])
+            feat_b = self.projectors[mod_b](features[mod_b])
+            
+            # 计算损失
+            loss_key = f'{mod_a}_{mod_b}_loss'
+            
+            if self.loss_type == 'clip':
+                loss, _, _ = self.loss_fn(feat_a, feat_b)
+            else:
+                loss = self.loss_fn(feat_a, feat_b)
+            
+            losses[loss_key] = loss
+        
+        return losses
+
+class MomentumEncoder(nn.Module):
+    """
+    动量编码器 - 用于MoCo风格的对比学习
+    支持参数和 Buffer (如 BatchNorm stats) 的动量更新
+    """
+    def __init__(self, encoder: nn.Module, momentum: float = 0.999):
+        super().__init__()
+        self.encoder = encoder
+        self.momentum_encoder = self._build_momentum_encoder(encoder)
+        self.momentum = momentum
+
+    def _build_momentum_encoder(self, encoder: nn.Module) -> nn.Module:
+        """构建动量编码器"""
+        momentum_encoder = copy.deepcopy(encoder)
+        
+        # 冻结动量编码器参数
+        for param in momentum_encoder.parameters():
+            param.requires_grad = False
+        
+        return momentum_encoder
+
+    @torch.no_grad()
+    def _update_momentum_encoder(self):
+        """更新动量编码器 (In-place update)"""
+        # 更新参数
+        for param_q, param_k in zip(
+            self.encoder.parameters(),
+            self.momentum_encoder.parameters()
+        ):
+            # EMA Update: k = m * k + (1 - m) * q
+            param_k.data.mul_(self.momentum).add_(param_q.data, alpha=1.0 - self.momentum)
+        
+        # 更新 Buffers (如 BatchNorm running mean/var)
+        # 简单的策略是直接覆盖，或者同样使用 EMA。通常直接覆盖即可，
+        # 因为 Key Encoder 处于 Eval 模式，不追踪 batch stats。
+        for buffer_q, buffer_k in zip(
+            self.encoder.buffers(),
+            self.momentum_encoder.buffers()
+        ):
+            buffer_k.data.copy_(buffer_q.data)
+
+    def forward(self, x: torch.Tensor, use_momentum: bool = False) -> torch.Tensor:
+        """
+        Args:
+            x: 输入数据
+            use_momentum: 如果为 True，使用动量编码器 (通常用于生成 Key/Target)
+        """
+        if use_momentum:
+            with torch.no_grad():
+                self._update_momentum_encoder()
+                # 动量编码器始终处于 eval 模式
+                self.momentum_encoder.eval()
+                return self.momentum_encoder(x)
+        else:
+            return self.encoder(x)

data_augmentation.py

@@ -0,0 +1,366 @@
+"""
+数据增强模块
+针对不同模态的高级数据增强策略
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+import random
+import math
+
+class RandAugment(nn.Module):
+    """RandAugment for images"""
+    def __init__(self, n: int = 2, m: int = 10):
+        super().__init__()
+        self.n = n
+        self.m = m
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """随机应用n个增强操作"""
+        # 确保输入是 [B, C, H, W]，如果是 [C, H, W] 则增加维度
+        is_batched = x.ndim == 4
+        if not is_batched:
+            x = x.unsqueeze(0)
+
+        augmentations = [
+            self._auto_contrast,
+            self._equalize,
+            self._solarize,
+            self._color,
+            self._contrast,
+            self._brightness,
+            self._sharpness,
+        ]
+        
+        # 这里的ops应该是每一轮随机选择，而不是固定
+        for _ in range(self.n):
+            aug = random.choice(augmentations)
+            x = aug(x)
+        
+        if not is_batched:
+            x = x.squeeze(0)
+        
+        return x
+
+    def _auto_contrast(self, x: torch.Tensor) -> torch.Tensor:
+        """自动对比度: 线性拉伸到 [0, 1]"""
+        # 针对每个样本分别计算 min/max
+        # x: [B, C, H, W]
+        B, C, H, W = x.shape
+        x_flat = x.view(B, C, -1)
+        min_val = x_flat.min(dim=2, keepdim=True)[0].view(B, C, 1, 1)
+        max_val = x_flat.max(dim=2, keepdim=True)[0].view(B, C, 1, 1)
+        return (x - min_val) / (max_val - min_val + 1e-8)
+
+    def _equalize(self, x: torch.Tensor) -> torch.Tensor:
+        """直方图均衡化 (简化版:基于每个通道的CDF)"""
+        # 这是一个计算密集型操作，PyTorch原生实现较复杂。
+        # 这里实现一个基于排序的简化版本，模拟均衡化效果
+        B, C, H, W = x.shape
+        # 将像素值缩放到 [0, 255] 离散化以便计算直方图
+        x_int = (x * 255).long().clamp(0, 255)
+        
+        out = torch.zeros_like(x)
+        
+        for b in range(B):
+            for c in range(C):
+                hist = torch.histc(x[b, c].float(), bins=256, min=0, max=1)
+                cdf = hist.cumsum(0)
+                cdf = cdf / cdf[-1] # 归一化
+                # 使用cdf作为查找表
+                out[b, c] = cdf[x_int[b, c]]
+                
+        return out
+
+    def _solarize(self, x: torch.Tensor) -> torch.Tensor:
+        """曝光"""
+        threshold = random.uniform(0.3, 0.7)
+        return torch.where(x < threshold, x, 1.0 - x)
+
+    def _color(self, x: torch.Tensor) -> torch.Tensor:
+        """颜色增强 (饱和度)"""
+        factor = 1.0 + (random.random() - 0.5) * 0.4
+        # RGB转灰度简单近似: mean over channels
+        # x is [B, C, H, W], dim=1 is channels
+        mean = x.mean(dim=1, keepdim=True) 
+        return torch.clamp(mean + factor * (x - mean), 0, 1)
+
+    def _contrast(self, x: torch.Tensor) -> torch.Tensor:
+        """对比度"""
+        factor = 1.0 + (random.random() - 0.5) * 0.4
+        # 计算整张图的均值，保留 Batch 维度
+        # view(B, -1) -> mean(1) -> view(B, 1, 1, 1)
+        mean = x.view(x.size(0), -1).mean(dim=1).view(-1, 1, 1, 1)
+        return torch.clamp(mean + factor * (x - mean), 0, 1)
+
+    def _brightness(self, x: torch.Tensor) -> torch.Tensor:
+        """亮度"""
+        factor = 1.0 + (random.random() - 0.5) * 0.4
+        return torch.clamp(x * factor, 0, 1)
+
+    def _sharpness(self, x: torch.Tensor) -> torch.Tensor:
+        """锐化: 通过混合原图和高斯模糊图实现"""
+        factor = 1.0 + (random.random() - 0.5) * 0.4
+        # 使用 AvgPool 模拟模糊
+        kernel_size = 3
+        pad = kernel_size // 2
+        blurred = F.avg_pool2d(x, kernel_size=kernel_size, stride=1, padding=pad)
+        # 锐化公式: Original + alpha * (Original - Blurred)
+        # 或者简单的混合: Blend(Original, Blurred, factor)
+        # 这里使用 PIL 风格的锐化: 
+        # result = original * factor + blurred * (1 - factor) 
+        # 但要注意 factor>1 时是锐化，factor<1 是模糊
+        # 更标准的锐化掩模: x + factor * (x - blurred)
+        return torch.clamp(x + (factor - 1.0) * (x - blurred), 0, 1)
+
+class MixUp(nn.Module):
+    """MixUp数据增强"""
+    def __init__(self, alpha: float = 1.0, num_classes: Optional[int] = None):
+        super().__init__()
+        self.alpha = alpha
+        self.num_classes = num_classes
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        y: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], float]:
+        
+        if self.alpha > 0:
+            lambda_ = random.betavariate(self.alpha, self.alpha)
+        else:
+            lambda_ = 1.0
+        
+        batch_size = x.shape[0]
+        index = torch.randperm(batch_size, device=x.device)
+        
+        mixed_x = lambda_ * x + (1 - lambda_) * x[index]
+        
+        mixed_y = None
+        if y is not None:
+            # 处理标签混合
+            y_a = y
+            y_b = y[index]
+            
+            # 检查标签是否需要 One-Hot 编码 (如果 y 是 long 类型且维度不对)
+            if y.dtype == torch.long or y.ndim == 1:
+                if self.num_classes is None:
+                    # 如果未提供 num_classes，尝试推断 (可能有风险)
+                    self.num_classes = int(y.max().item()) + 1
+                
+                y_a = F.one_hot(y_a, num_classes=self.num_classes).float()
+                y_b = F.one_hot(y_b, num_classes=self.num_classes).float()
+            
+            mixed_y = lambda_ * y_a + (1 - lambda_) * y_b
+        
+        return mixed_x, mixed_y, lambda_
+
+class CutMix(nn.Module):
+    """CutMix数据增强"""
+    def __init__(self, alpha: float = 1.0, num_classes: Optional[int] = None):
+        super().__init__()
+        self.alpha = alpha
+        self.num_classes = num_classes
+
+    def _rand_bbox(
+        self,
+        size: Tuple[int, ...],
+        lambda_: float
+    ) -> Tuple[int, int, int, int]:
+        """生成随机bbox"""
+        W = size[-1] # 兼容 [B, C, H, W]
+        H = size[-2]
+        cut_rat = math.sqrt(1.0 - lambda_)
+        cut_w = int(W * cut_rat)
+        cut_h = int(H * cut_rat)
+        
+        cx = random.randint(0, W)
+        cy = random.randint(0, H)
+        
+        bbx1 = torch.tensor(cx - cut_w // 2, device='cpu').clamp(0, W).item()
+        bby1 = torch.tensor(cy - cut_h // 2, device='cpu').clamp(0, H).item()
+        bbx2 = torch.tensor(cx + cut_w // 2, device='cpu').clamp(0, W).item()
+        bby2 = torch.tensor(cy + cut_h // 2, device='cpu').clamp(0, H).item()
+        
+        return int(bbx1), int(bby1), int(bbx2), int(bby2)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        y: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], float]:
+        
+        if self.alpha > 0:
+            lambda_ = random.betavariate(self.alpha, self.alpha)
+        else:
+            lambda_ = 1.0
+        
+        batch_size = x.shape[0]
+        index = torch.randperm(batch_size, device=x.device)
+        
+        bbx1, bby1, bbx2, bby2 = self._rand_bbox(x.size(), lambda_)
+        
+        # 克隆防止就地修改影响后续梯度计算 (虽然这里是输入数据处理，通常还好)
+        x = x.clone()
+        x[:, :, bby1:bby2, bbx1:bbx2] = x[index, :, bby1:bby2, bbx1:bbx2]
+        
+        # 调整lambda为精确的像素比例
+        # 注意: 原始代码中宽高的计算顺序可能有歧义，这里统一 H=size[-2], W=size[-1]
+        H, W = x.size()[-2], x.size()[-1]
+        lambda_ = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (H * W))
+        
+        mixed_y = None
+        if y is not None:
+            y_a = y
+            y_b = y[index]
+            
+            if y.dtype == torch.long or y.ndim == 1:
+                if self.num_classes is None:
+                    # 最好在初始化时传入 num_classes
+                    self.num_classes = int(y.max().item()) + 1
+                y_a = F.one_hot(y_a, num_classes=self.num_classes).float()
+                y_b = F.one_hot(y_b, num_classes=self.num_classes).float()
+
+            mixed_y = lambda_ * y_a + (1 - lambda_) * y_b
+        
+        return x, mixed_y, lambda_
+
+class SpecAugment(nn.Module):
+    """SpecAugment for audio spectrograms"""
+    def __init__(
+        self,
+        freq_mask_param: int = 27,
+        time_mask_param: int = 100,
+        num_freq_masks: int = 2,
+        num_time_masks: int = 2
+    ):
+        super().__init__()
+        self.freq_mask_param = freq_mask_param
+        self.time_mask_param = time_mask_param
+        self.num_freq_masks = num_freq_masks
+        self.num_time_masks = num_time_masks
+
+    def forward(self, spec: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            spec: [B, F, T] or [B, C, F, T]
+        """
+        input_ndim = spec.ndim
+        if input_ndim == 3:
+            spec = spec.unsqueeze(1) # [B, 1, F, T]
+        
+        B, C, F, T = spec.shape
+        spec = spec.clone()
+        
+        # 频率遮罩
+        for _ in range(self.num_freq_masks):
+            # 确保 mask 不超过 F
+            f_param = min(self.freq_mask_param, F)
+            f = random.randint(0, f_param)
+            f0 = random.randint(0, max(0, F - f))
+            spec[:, :, f0:f0+f, :] = 0
+        
+        # 时间遮罩
+        for _ in range(self.num_time_masks):
+            # 确保 mask 不超过 T
+            t_param = min(self.time_mask_param, T)
+            t = random.randint(0, t_param)
+            t0 = random.randint(0, max(0, T - t))
+            spec[:, :, :, t0:t0+t] = 0
+        
+        if input_ndim == 3:
+            return spec.squeeze(1)
+        return spec
+
+class TemporalMasking(nn.Module):
+    """视频的时序遮罩"""
+    def __init__(self, mask_ratio: float = 0.15):
+        super().__init__()
+        self.mask_ratio = mask_ratio
+
+    def forward(self, video: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            video: [B, T, C, H, W]
+        """
+        B, T, C, H, W = video.shape
+        num_mask = int(T * self.mask_ratio)
+        if num_mask == 0:
+            return video
+            
+        video = video.clone()
+        
+        for b in range(B):
+            # 随机采样要遮罩的帧索引
+            mask_indices = torch.randperm(T)[:num_mask]
+            video[b, mask_indices] = 0
+        
+        return video
+
+class MultiModalAugmentation(nn.Module):
+    """统一的多模态数据增强"""
+    def __init__(
+        self,
+        image_aug: bool = True,
+        audio_aug: bool = True,
+        video_aug: bool = True,
+        use_mixup: bool = True,
+        use_cutmix: bool = True,
+        num_classes: Optional[int] = None
+    ):
+        super().__init__()
+        self.image_aug = RandAugment() if image_aug else None
+        self.audio_aug = SpecAugment() if audio_aug else None
+        self.video_aug = TemporalMasking() if video_aug else None
+        
+        self.mixup = MixUp(num_classes=num_classes) if use_mixup else None
+        self.cutmix = CutMix(num_classes=num_classes) if use_cutmix else None
+
+    def forward(
+        self,
+        data: torch.Tensor,
+        modality: str,
+        labels: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        Args:
+            data: 输入数据
+            modality: 模态类型 ('image', 'audio', 'video')
+            labels: 标签（可选）
+        """
+        # 1. 模态特定的增强 (Intra-sample augmentation)
+        if modality == 'image' and self.image_aug is not None:
+            data = self.image_aug(data)
+        elif modality == 'audio' and self.audio_aug is not None:
+            data = self.audio_aug(data)
+        elif modality == 'video' and self.video_aug is not None:
+            data = self.video_aug(data)
+        
+        # 2. 混合增强 (Inter-sample augmentation)
+        if self.training and labels is not None:
+            # 随机选择 MixUp 或 CutMix，或者都不选
+            # 策略：如果有 CutMix 且是图片，50%概率 CutMix；否则看有没有 MixUp
+            
+            apply_mixup = False
+            apply_cutmix = False
+            
+            p = random.random()
+            
+            # 简单的互斥逻辑：如果有CutMix且是图像，一半概率CutMix，一半概率MixUp(如果有)
+            if self.cutmix is not None and modality == 'image':
+                if p < 0.5:
+                    apply_cutmix = True
+                elif self.mixup is not None:
+                    apply_mixup = True
+            elif self.mixup is not None:
+                # 非图像或无CutMix，则只考虑MixUp
+                if p < 0.5: # 假设 50% 概率应用 MixUp
+                    apply_mixup = True
+            
+            if apply_cutmix:
+                data, labels, _ = self.cutmix(data, labels)
+            elif apply_mixup:
+                data, labels, _ = self.mixup(data, labels)
+        
+        return data, labels

data_config.py

@@ -0,0 +1,292 @@
+# data_config.py
+"""
+预训练和后训练数据集配置
+"""
+
+PRETRAIN_DATASETS = {
+    # 文本数据集
+    'the_pile': {
+        'type': 'text',
+        'hf_path': 'EleutherAI/pile',
+        'split': 'train',
+        'streaming': True,
+        'text_field': 'text',
+        'weight': 1.0,
+        'description': 'The Pile - 825GB diverse text corpus'
+    },
+    'c4': {
+        'type': 'text',
+        'hf_path': 'allenai/c4',
+        'config': 'en',
+        'split': 'train',
+        'streaming': True,
+        'text_field': 'text',
+        'weight': 0.5,
+        'description': 'C4 - Colossal Clean Crawled Corpus'
+    },
+    'wikipedia': {
+        'type': 'text',
+        'hf_path': 'HuggingFaceFW/fineweb-edu',
+        'config': 'sample-10BT',
+        'split': 'train',
+        'streaming': True,
+        'text_field': 'text',
+        'weight': 0.3,
+        'description': 'FineWeb Edu - High quality educational content'
+    },
+    'bookcorpus': {
+        'type': 'text',
+        'hf_path': 'HuggingFaceTB/smollm-corpus',
+        'config': 'cosmopedia-v2',
+        'split': 'train',
+        'streaming': True,
+        'text_field': 'text',
+        'weight': 0.2,
+        'description': 'Synthetic textbooks and stories'
+    },
+    # 代码数据集
+    'codeparrot': {
+        'type': 'code',
+        'hf_path': 'bigcode/the-stack-smol',
+        'config': 'default',
+        'split': 'train',
+        'streaming': True,
+        'text_field': 'content',
+        'weight': 0.3,
+        'description': 'The Stack Smol - code'
+    },
+    'the_stack': {
+        'type': 'code',
+        'hf_path': 'bigcode/the-stack-dedup',
+        'split': 'train',
+        'streaming': True,
+        'text_field': 'content',
+        'weight': 0.2,
+        'description': 'The Stack - deduplicated code'
+    },
+    # 多模态数据集
+    'laion400m': {
+        'type': 'image_text',
+        'hf_path': 'laion/laion400m',
+        'split': 'train',
+        'streaming': True,
+        'image_field': 'url',
+        'text_field': 'caption',
+        'weight': 0.4,
+        'description': 'LAION-400M image-text pairs'
+    },
+    'conceptual_captions': {
+        'type': 'image_text',
+        'hf_path': 'google-research-datasets/conceptual_captions',
+        'split': 'train',
+        'streaming': False,
+        'image_field': 'image_url',
+        'text_field': 'caption',
+        'weight': 0.2,
+        'description': 'Conceptual Captions 3M'
+    },
+}
+
+# 后训练数据集配置（instruction tuning + alignment）
+POSTTRAIN_DATASETS = {
+    # Instruction Tuning数据集
+    'flan_v2': {
+        'type': 'instruction',
+        'hf_path': 'Muennighoff/flan',
+        'split': 'train',
+        'streaming': True,
+        'instruction_field': 'inputs',
+        'response_field': 'targets',
+        'weight': 1.0,
+        'max_samples': 100000,
+        'description': 'FLAN v2 collection'
+    },
+    'alpaca': {
+        'type': 'instruction',
+        'hf_path': 'tatsu-lab/alpaca',
+        'split': 'train',
+        'streaming': False,
+        'instruction_field': 'instruction',
+        'input_field': 'input',
+        'response_field': 'output',
+        'weight': 0.5,
+        'description': 'Stanford Alpaca 52K'
+    },
+    'dolly': {
+        'type': 'instruction',
+        'hf_path': 'databricks/databricks-dolly-15k',
+        'split': 'train',
+        'streaming': False,
+        'instruction_field': 'instruction',
+        'context_field': 'context',  # Dolly有context字段
+        'response_field': 'response',
+        'weight': 0.3,
+        'description': 'Dolly 15K'
+    },
+    'oasst1': {
+        'type': 'conversation',
+        'hf_path': 'OpenAssistant/oasst1',
+        'split': 'train',
+        'streaming': False,
+        'weight': 0.4,
+        'description': 'OpenAssistant Conversations',
+        # OASST1需要特殊处理，因为它是树形结构
+        # 可能需要自定义预处理
+    },
+    'sharegpt': {
+        'type': 'conversation',
+        'hf_path': 'anon8231489123/ShareGPT_Vicuna_unfiltered',
+        'split': 'train',
+        'streaming': False,
+        'weight': 0.3,
+        'max_samples': 50000,
+        'description': 'ShareGPT conversations'
+    },
+    # Code instruction数据集
+    'code_alpaca': {
+        'type': 'code_instruction',
+        'hf_path': 'sahil2801/CodeAlpaca-20k',
+        'split': 'train',
+        'streaming': False,
+        'instruction_field': 'instruction',
+        'response_field': 'output',
+        'weight': 0.3,
+        'description': 'Code Alpaca 20K'
+    },
+    # 多模态instruction数据集
+    'llava_instruct': {
+        'type': 'multimodal_instruction',
+        'hf_path': 'liuhaotian/LLaVA-Instruct-150K',
+        'split': 'train',
+        'streaming': False,
+        'image_field': 'image',
+        'instruction_field': 'conversations',
+        'weight': 0.5,
+        'description': 'LLaVA visual instruction tuning'
+    },
+    # Preference数据集 (用于RLHF)
+    'hh_rlhf': {
+        'type': 'preference',
+        'hf_path': 'Anthropic/hh-rlhf',
+        'split': 'train',
+        'streaming': False,
+        'chosen_field': 'chosen',
+        'rejected_field': 'rejected',
+        'weight': 1.0,
+        'description': 'Anthropic HH-RLHF'
+    },
+    'ultrafeedback': {
+        'type': 'preference',
+        'hf_path': 'openbmb/UltraFeedback',
+        'split': 'train',
+        'streaming': True,
+        'chosen_field': 'chosen',  # 添加字段配置
+        'rejected_field': 'rejected',
+        'weight': 0.5,
+        'max_samples': 50000,
+        'description': 'UltraFeedback preferences'
+    },
+     'debug_water': {
+        'type': 'instruction',
+        'hf_path': 'json',              # 使用 json 加载器
+        'data_files': 'debug_water.json', # 指向刚才生成的文件
+        'split': 'train',
+        'streaming': False,
+        'instruction_field': 'instruction',
+        'response_field': 'output',
+        'weight': 1.0,
+        'description': 'Overfitting test for water'
+    },
+}
+
+# 轻量级测试数据集（用于快速验证）
+TEST_DATASETS = {
+    'tiny_shakespeare': {
+        'type': 'text',
+        'hf_path': 'tiny_shakespeare',
+        'split': 'train',
+        'streaming': False,
+        'text_field': 'text',
+        'weight': 1.0,
+        'description': 'Tiny Shakespeare for testing'
+    },
+    'gsm8k': {
+        'type': 'instruction',
+        'hf_path': 'gsm8k',
+        'config': 'main',
+        'split': 'train',
+        'streaming': False,
+        'instruction_field': 'question',
+        'response_field': 'answer',
+        'weight': 1.0,
+        'description': 'GSM8K math problems'
+    },
+}
+
+# 数据集混合策略
+PRETRAIN_MIX = {
+    'default': {
+        'datasets': ['c4', 'wikipedia', 'bookcorpus', 'codeparrot'],
+        'weights': [0.5, 0.2, 0.2, 0.1],
+        'description': 'Default pretrain mix'
+    },
+    'code_heavy': {
+        'datasets': ['c4', 'codeparrot', 'the_stack', 'wikipedia'],
+        'weights': [0.3, 0.4, 0.2, 0.1],
+        'description': 'Code-heavy mix'
+    },
+    'multimodal': {
+        'datasets': ['c4', 'wikipedia', 'laion400m', 'conceptual_captions'],
+        'weights': [0.4, 0.2, 0.3, 0.1],
+        'description': 'Multimodal mix'
+    },
+    'text_only': {
+        'datasets': ['c4', 'wikipedia', 'bookcorpus'],
+        'weights': [0.5, 0.3, 0.2],
+        'description': 'Text-only mix for testing'
+    },
+}
+
+POSTTRAIN_MIX = {
+    'default': {
+        'datasets': ['flan_v2', 'alpaca', 'dolly', 'oasst1'],
+        'weights': [0.4, 0.3, 0.2, 0.1],
+        'description': 'Default instruction tuning mix'
+    },
+    'conversation': {
+        'datasets': ['oasst1', 'sharegpt', 'alpaca'],
+        'weights': [0.4, 0.4, 0.2],
+        'description': 'Conversation-focused mix'
+    },
+    'code_instruct': {
+        'datasets': ['code_alpaca', 'alpaca', 'flan_v2'],
+        'weights': [0.5, 0.3, 0.2],
+        'description': 'Code instruction mix'
+    },
+    'simple_instruct': {
+        'datasets': ['alpaca', 'dolly'],
+        'weights': [0.6, 0.4],
+        'description': 'Simple instruction mix for testing'
+    },
+    'debug_mix': {
+        'datasets': ['debug_water'],
+        'weights': [1.0],
+        'description': 'Debug mix for overfitting'
+    },
+}
+
+# 下载和缓存配置
+DATASET_CACHE_DIR = "./dataset_cache"
+HF_CACHE_DIR = "./hf_cache"
+MAX_RETRIES = 3
+DOWNLOAD_TIMEOUT = 300
+
+# 数据处理配置
+PREPROCESSING_CONFIG = {
+    'max_seq_length': 2048,
+    'min_seq_length': 32,
+    'num_workers': 4,
+    'batch_size': 8,
+    'shuffle_buffer_size': 10000,
+    'seed': 42,
+}

data_loader.py

@@ -0,0 +1,832 @@
+# data_loader.py
+"""
+改进的数据加载器 - 支持预训练和后训练数据集
+"""
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader, IterableDataset
+from datasets import load_dataset, concatenate_datasets, interleave_datasets
+from typing import Dict, List, Optional, Any, Union
+import random
+import numpy as np
+from tqdm import tqdm
+import warnings
+from PIL import Image
+import requests
+from io import BytesIO
+from torchvision import transforms
+import logging
+
+# 设置日志
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+warnings.filterwarnings("ignore", category=UserWarning)
+
+from data_config import (
+    PRETRAIN_DATASETS,
+    POSTTRAIN_DATASETS,
+    TEST_DATASETS,
+    PRETRAIN_MIX,
+    POSTTRAIN_MIX,
+    PREPROCESSING_CONFIG,
+    DATASET_CACHE_DIR,
+    HF_CACHE_DIR
+)
+
+# 图像变换
+image_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+class PreTrainDataset(IterableDataset):
+    """预训练数据集 - 支持流式和混合采样"""
+    def __init__(
+        self,
+        mix_name: str = 'default',
+        tokenizer=None,
+        max_length: int = 2048,
+        streaming: bool = True,
+        seed: int = 42,
+        max_samples: Optional[int] = None
+    ):
+        super().__init__()
+        
+        if tokenizer is None:
+            raise ValueError("tokenizer cannot be None")
+        
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.streaming = streaming
+        self.seed = seed
+        self.max_samples = max_samples
+        self.samples_generated = 0
+        
+        # 获取混合配置
+        if mix_name not in PRETRAIN_MIX:
+            raise ValueError(f"Unknown mix: {mix_name}. Available: {list(PRETRAIN_MIX.keys())}")
+        
+        mix_config = PRETRAIN_MIX[mix_name]
+        dataset_names = mix_config.get('datasets', [])
+        weights = mix_config.get('weights', [])
+        
+        if not dataset_names:
+            raise ValueError(f"No datasets found in mix: {mix_name}")
+        
+        logger.info(f"Loading pretrain mix: {mix_name}")
+        logger.info(f"  Datasets: {dataset_names}")
+        logger.info(f"  Weights: {weights}")
+        
+        # 加载数据集
+        self.datasets = []
+        self.probabilities = []
+        
+        for name, weight in zip(dataset_names, weights):
+            if name not in PRETRAIN_DATASETS:
+                logger.warning(f"Dataset {name} not found in PRETRAIN_DATASETS, skipping")
+                continue
+            
+            config = PRETRAIN_DATASETS[name]
+            try:
+                ds = self._load_dataset(config)
+                if ds is not None:
+                    self.datasets.append((name, ds, config))
+                    self.probabilities.append(weight)
+                    logger.info(f"  Successfully loaded {name}")
+            except Exception as e:
+                logger.error(f"Error loading {name}: {e}")
+                continue
+        
+        if not self.datasets:
+            raise ValueError("No datasets loaded successfully")
+        
+        # 归一化概率
+        total = sum(self.probabilities)
+        self.probabilities = [p / total for p in self.probabilities]
+        
+        logger.info(f"Successfully loaded {len(self.datasets)} datasets")
+
+    def _load_dataset(self, config: Dict):
+        """加载单个数据集"""
+        try:
+            load_kwargs = {
+                'path': config['hf_path'],
+                'split': config.get('split', 'train'),
+                'streaming': config.get('streaming', self.streaming),
+                'cache_dir': HF_CACHE_DIR,
+            }
+            
+            # 添加config参数（如果存在）
+            if 'config' in config:
+                load_kwargs['name'] = config['config']
+            
+            ds = load_dataset(**load_kwargs)
+            return ds
+        except Exception as e:
+            logger.error(f"Failed to load {config.get('hf_path', 'unknown')}: {e}")
+            return None
+
+    def _process_text_sample(self, sample: Dict, config: Dict) -> Optional[Dict]:
+        """处理文本样本"""
+        try:
+            text_field = config.get('text_field', 'text')
+            text = sample.get(text_field, '')
+            
+            if not text or not isinstance(text, str):
+                return None
+            
+            text = text.strip()
+            if len(text) < 10:
+                return None
+            
+            # Tokenize
+            encoding = self.tokenizer(
+                text,
+                max_length=self.max_length,
+                truncation=True,
+                padding='max_length',
+                return_tensors='pt'
+            )
+            
+            return {
+                'input_ids': encoding['input_ids'].squeeze(0),
+                'attention_mask': encoding['attention_mask'].squeeze(0),
+                'type': 'text'
+            }
+        except Exception as e:
+            logger.debug(f"Error processing text sample: {e}")
+            return None
+
+    def _process_image_text_sample(self, sample: Dict, config: Dict) -> Optional[Dict]:
+        """处理图像-文本样本"""
+        try:
+            text_field = config.get('text_field', 'caption')
+            image_field = config.get('image_field', 'image')
+            
+            text = sample.get(text_field, '')
+            image = sample.get(image_field)
+            
+            if not text or image is None:
+                return None
+            
+            # 处理图像
+            if isinstance(image, str):
+                # URL - 添加超时和错误处理
+                try:
+                    response = requests.get(image, timeout=5)
+                    image = Image.open(BytesIO(response.content)).convert('RGB')
+                except Exception as img_error:
+                    logger.debug(f"Failed to load image from URL: {img_error}")
+                    return None
+            elif isinstance(image, Image.Image):
+                image = image.convert('RGB')
+            else:
+                return None
+            
+            # 转换图像
+            image_tensor = image_transform(image)
+            
+            # Tokenize文本
+            encoding = self.tokenizer(
+                text,
+                max_length=self.max_length,
+                truncation=True,
+                padding='max_length',
+                return_tensors='pt'
+            )
+            
+            return {
+                'input_ids': encoding['input_ids'].squeeze(0),
+                'attention_mask': encoding['attention_mask'].squeeze(0),
+                'image': image_tensor,
+                'type': 'image_text'
+            }
+        except Exception as e:
+            logger.debug(f"Error processing image-text sample: {e}")
+            return None
+
+    def __iter__(self):
+        """迭代器"""
+        worker_info = torch.utils.data.get_worker_info()
+        if worker_info is not None:
+            # 多worker时设置不同的随机种子
+            random.seed(self.seed + worker_info.id)
+            np.random.seed(self.seed + worker_info.id)
+        else:
+            random.seed(self.seed)
+            np.random.seed(self.seed)
+        
+        # 创建数据集迭代器
+        iterators = [iter(ds) for _, ds, _ in self.datasets]
+        self.samples_generated = 0
+        
+        while True:
+            # 检查是否达到最大样本数
+            if self.max_samples and self.samples_generated >= self.max_samples:
+                break
+            
+            try:
+                # 根据概率选择数据集
+                idx = np.random.choice(len(self.datasets), p=self.probabilities)
+                name, _, config = self.datasets[idx]
+                
+                # 从选中的数据集获取样本
+                sample = next(iterators[idx])
+                
+                # 处理样本
+                processed = None
+                if config.get('type') in ['text', 'code']:
+                    processed = self._process_text_sample(sample, config)
+                elif config.get('type') == 'image_text':
+                    processed = self._process_image_text_sample(sample, config)
+                else:
+                    logger.debug(f"Unknown type: {config.get('type')}")
+                    continue
+                
+                if processed is not None:
+                    self.samples_generated += 1
+                    yield processed
+                    
+            except StopIteration:
+                # 重新创建迭代器
+                try:
+                    iterators[idx] = iter(self.datasets[idx][1])
+                except Exception as e:
+                    logger.error(f"Failed to recreate iterator for dataset {idx}: {e}")
+                    break
+            except Exception as e:
+                logger.debug(f"Error in iterator: {e}")
+                continue
+
+
+class PostTrainDataset(Dataset):
+    """后训练数据集 - Instruction tuning和对话"""
+    def __init__(
+        self,
+        mix_name: str = 'default',
+        tokenizer=None,
+        max_length: int = 2048,
+        max_samples: Optional[int] = None,
+        split: str = 'train'
+    ):
+        super().__init__()
+        
+        if tokenizer is None:
+            raise ValueError("tokenizer cannot be None")
+        
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.split = split
+        
+        # 获取混合配置
+        if mix_name not in POSTTRAIN_MIX:
+            raise ValueError(f"Unknown mix: {mix_name}. Available: {list(POSTTRAIN_MIX.keys())}")
+        
+        mix_config = POSTTRAIN_MIX[mix_name]
+        dataset_names = mix_config.get('datasets', [])
+        weights = mix_config.get('weights', [])
+        
+        if not dataset_names:
+            raise ValueError(f"No datasets found in mix: {mix_name}")
+        
+        logger.info(f"Loading posttrain mix: {mix_name}")
+        logger.info(f"  Datasets: {dataset_names}")
+        
+        # 加载和合并数据集
+        all_datasets = []
+        
+        for name in dataset_names:
+            if name not in POSTTRAIN_DATASETS:
+                logger.warning(f"Dataset {name} not found in POSTTRAIN_DATASETS")
+                continue
+            
+            config = POSTTRAIN_DATASETS[name]
+            try:
+                load_kwargs = {
+                    'path': config['hf_path'],
+                    'split': split,
+                    'streaming': config.get('streaming', False),
+                    'cache_dir': HF_CACHE_DIR,
+                }
+                # [新增] 如果配置里有 data_files，就加进去
+                if 'data_files' in config:
+                    load_kwargs['data_files'] = config['data_files']
+                # 添加config参数（如果存在）
+                if 'config' in config:
+                    load_kwargs['name'] = config['config']
+                
+                ds = load_dataset(**load_kwargs)
+                
+                # 限制样本数
+                if config.get('max_samples'):
+                    if hasattr(ds, 'take'):
+                        ds = ds.take(config['max_samples'])
+                    elif hasattr(ds, 'select'):
+                        ds = ds.select(range(min(len(ds), config['max_samples'])))
+                
+                # 添加数据集标识
+                def add_source(example):
+                    example['_source'] = name
+                    example['_config'] = config
+                    return example
+                
+                ds = ds.map(add_source)
+                all_datasets.append(ds)
+                
+                ds_len = len(ds) if hasattr(ds, '__len__') else 'streaming'
+                logger.info(f"  Loaded {name}: {ds_len} samples")
+                
+            except Exception as e:
+                logger.error(f"Error loading {name}: {e}")
+                continue
+        
+        # 合并数据集
+        if not all_datasets:
+            raise ValueError("No datasets loaded successfully")
+        
+        if len(all_datasets) == 1:
+            self.dataset = all_datasets[0]
+        else:
+            # 交织数据集
+            probabilities = [w / sum(weights[:len(all_datasets)]) 
+                           for w in weights[:len(all_datasets)]]
+            self.dataset = interleave_datasets(
+                all_datasets,
+                probabilities=probabilities,
+                seed=42,
+                stopping_strategy='all_exhausted'
+            )
+        
+        # 限制总样本数
+        if max_samples and hasattr(self.dataset, '__len__'):
+            actual_len = min(len(self.dataset), max_samples)
+            self.dataset = self.dataset.select(range(actual_len))
+        
+        dataset_len = len(self.dataset) if hasattr(self.dataset, '__len__') else 'streaming'
+        logger.info(f"Total samples: {dataset_len}")
+
+    def _format_instruction(self, sample: Dict, config: Dict) -> str:
+        """格式化instruction"""
+        try:
+            data_type = config.get('type', 'instruction')
+            
+            if data_type == 'instruction':
+                instruction_field = config.get('instruction_field', 'instruction')
+                input_field = config.get('input_field', 'input')
+                context_field = config.get('context_field', None)
+                
+                instruction = sample.get(instruction_field, '')
+                input_text = sample.get(input_field, '')
+                context = sample.get(context_field, '') if context_field else ''
+                
+                # 构建prompt
+                prompt_parts = [f"Instruction: {instruction}"]
+                
+                if context:
+                    prompt_parts.append(f"Context: {context}")
+                
+                if input_text:
+                    prompt_parts.append(f"Input: {input_text}")
+                
+                prompt_parts.append("Response:")
+                return "\n".join(prompt_parts)
+            
+            elif data_type == 'conversation':
+                # 处理对话格式 - 支持不同的对话格式
+                if 'conversations' in sample:
+                    # LLaVA格式
+                    conversations = sample['conversations']
+                    if isinstance(conversations, list) and len(conversations) > 0:
+                        dialogue = []
+                        for conv in conversations[:-1]:
+                            role = conv.get('from', 'user')
+                            content = conv.get('value', '')
+                            dialogue.append(f"{role}: {content}")
+                        return "\n".join(dialogue) + "\nassistant:"
+                
+                elif 'messages' in sample:
+                    # 标准消息格式
+                    messages = sample['messages']
+                    if isinstance(messages, list) and len(messages) > 0:
+                        dialogue = []
+                        for msg in messages[:-1]:
+                            role = msg.get('role', 'user')
+                            content = msg.get('content', '')
+                            dialogue.append(f"{role}: {content}")
+                        return "\n".join(dialogue) + "\nassistant:"
+                
+                # 如果没有标准格式，尝试使用text字段
+                return sample.get('text', '')
+            
+            elif data_type == 'code_instruction':
+                # 代码instruction格式
+                instruction_field = config.get('instruction_field', 'instruction')
+                instruction = sample.get(instruction_field, '')
+                return f"### Instruction:\n{instruction}\n### Response:"
+            
+            elif data_type == 'multimodal_instruction':
+                # 多模态instruction
+                instruction_field = config.get('instruction_field', 'conversations')
+                conversations = sample.get(instruction_field, [])
+                if isinstance(conversations, list) and len(conversations) > 0:
+                    # 提取对话历史（除了最后一条回复）
+                    dialogue = []
+                    for conv in conversations[:-1]:
+                        role = conv.get('from', 'user')
+                        content = conv.get('value', '')
+                        dialogue.append(f"{role}: {content}")
+                    return "\n".join(dialogue) + "\nassistant:"
+                return ""
+            
+            else:
+                return sample.get(config.get('instruction_field', 'text'), '')
+        except Exception as e:
+            logger.debug(f"Error formatting instruction: {e}")
+            return ""
+
+    def _get_response(self, sample: Dict, config: Dict) -> str:
+        """获取响应"""
+        try:
+            data_type = config.get('type', 'instruction')
+            
+            if data_type == 'instruction' or data_type == 'code_instruction':
+                response_field = config.get('response_field', 'output')
+                return sample.get(response_field, '')
+            
+            elif data_type == 'conversation':
+                # 从对话中提取最后一条assistant的回复
+                if 'conversations' in sample:
+                    conversations = sample['conversations']
+                    if isinstance(conversations, list) and len(conversations) > 0:
+                        return conversations[-1].get('value', '')
+                
+                elif 'messages' in sample:
+                    messages = sample['messages']
+                    if isinstance(messages, list) and len(messages) > 0:
+                        return messages[-1].get('content', '')
+                
+                return ""
+            
+            elif data_type == 'multimodal_instruction':
+                instruction_field = config.get('instruction_field', 'conversations')
+                conversations = sample.get(instruction_field, [])
+                if isinstance(conversations, list) and len(conversations) > 0:
+                    return conversations[-1].get('value', '')
+                return ""
+            
+            else:
+                response_field = config.get('response_field', 'output')
+                return sample.get(response_field, '')
+        except Exception as e:
+            logger.debug(f"Error getting response: {e}")
+            return ""
+
+    def __len__(self):
+        return len(self.dataset) if hasattr(self.dataset, '__len__') else 0
+
+    def __getitem__(self, idx):
+        try:
+            sample = self.dataset[idx]
+            
+            # 获取配置
+            if '_config' not in sample:
+                logger.warning(f"Sample at index {idx} missing _config")
+                return None
+            
+            config = sample['_config']
+            
+            # 格式化 instruction 和 response
+            instruction_text = self._format_instruction(sample, config)
+            response_text = self._get_response(sample, config)
+            
+            if not instruction_text or not response_text:
+                return None
+            
+            # 确保 pad_token_id 存在
+            pad_token_id = self.tokenizer.pad_token_id 
+            if pad_token_id is None:
+                pad_token_id = self.tokenizer.eos_token_id
+
+            # =======================================================
+            # 1. 处理 Instruction (不需要 EOS，因为后面紧接 Response)
+            # =======================================================
+            instruction_max_len = self.max_length // 2
+            
+            # Tokenize 不做 padding，手动处理
+            instruction_enc = self.tokenizer(
+                instruction_text,
+                truncation=True,
+                max_length=instruction_max_len,
+                add_special_tokens=False, # 手动控制特殊token
+                return_tensors='pt'
+            )
+            instr_ids = instruction_enc['input_ids'].squeeze(0)
+            
+            # Instruction 手动 Padding
+            instr_len = instr_ids.size(0)
+            if instr_len < instruction_max_len:
+                # 左填充或者右填充皆可，通常 SFT 这里的 Instruction 是右填充
+                # padding_tensor = torch.full((instruction_max_len - instr_len,), pad_token_id, dtype=torch.long)
+                # instr_ids = torch.cat([instr_ids, padding_tensor])
+                # 为了保持代码与原逻辑一致，这里使用右填充至固定长度
+                padding = torch.full((instruction_max_len - instr_len,), pad_token_id, dtype=torch.long)
+                instr_ids = torch.cat([instr_ids, padding])
+                
+                # Mask: 真实token为1，pad为0
+                instr_mask = torch.cat([torch.ones(instr_len, dtype=torch.long), torch.zeros(instruction_max_len - instr_len, dtype=torch.long)])
+            else:
+                instr_mask = torch.ones(instruction_max_len, dtype=torch.long)
+
+            # =======================================================
+            # 2. 处理 Response (【核心修复】：必须加 EOS)
+            # =======================================================
+            response_max_len = self.max_length // 2
+            
+            # Tokenize: 预留1个位置给EOS
+            response_enc = self.tokenizer(
+                response_text,
+                truncation=True,
+                max_length=response_max_len - 1, # 关键：留一个位置给 EOS
+                add_special_tokens=False,
+                return_tensors='pt'
+            )
+            resp_ids = response_enc['input_ids'].squeeze(0)
+            
+            # 【强制添加 EOS Token】
+            eos_token = torch.tensor([self.tokenizer.eos_token_id], dtype=torch.long)
+            resp_ids = torch.cat([resp_ids, eos_token])
+            
+            # Response 手动 Padding
+            curr_resp_len = resp_ids.size(0)
+            if curr_resp_len < response_max_len:
+                padding = torch.full((response_max_len - curr_resp_len,), pad_token_id, dtype=torch.long)
+                resp_ids = torch.cat([resp_ids, padding])
+                
+                # Mask: 真实内容+EOS 为1，Pad 为0
+                resp_mask = torch.cat([torch.ones(curr_resp_len, dtype=torch.long), torch.zeros(response_max_len - curr_resp_len, dtype=torch.long)])
+            else:
+                resp_mask = torch.ones(response_max_len, dtype=torch.long)
+
+            # =======================================================
+            # 3. 组装结果
+            # =======================================================
+            result = {
+                'instruction': instr_ids,
+                'response': resp_ids,
+                'instruction_mask': instr_mask,
+                'response_mask': resp_mask,
+                'task': sample.get('_source', 'unknown'),
+                'modality_data': None
+            }
+            
+            # 如果是多模态数据，添加图像
+            if config.get('type') == 'multimodal_instruction' and 'image' in sample:
+                try:
+                    image = sample['image']
+                    if isinstance(image, Image.Image):
+                        image = image.convert('RGB')
+                        image_tensor = image_transform(image)
+                        result['modality_data'] = {'image': image_tensor}
+                except Exception as e:
+                    logger.debug(f"Error processing image: {e}")
+            
+            return result
+        
+        except Exception as e:
+            logger.debug(f"Error getting item at index {idx}: {e}")
+            import traceback
+            traceback.print_exc()
+            return None
+
+
+class PreferenceDataset(Dataset):
+    """偏好数据集 - 用于RLHF"""
+    def __init__(
+        self,
+        dataset_name: str = 'hh_rlhf',
+        tokenizer=None,
+        max_length: int = 1024,
+        max_samples: Optional[int] = None,
+        split: str = 'train'
+    ):
+        super().__init__()
+        
+        if tokenizer is None:
+            raise ValueError("tokenizer cannot be None")
+        
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        
+        if dataset_name not in POSTTRAIN_DATASETS:
+            raise ValueError(f"Unknown dataset: {dataset_name}. Available: {list(POSTTRAIN_DATASETS.keys())}")
+        
+        config = POSTTRAIN_DATASETS[dataset_name]
+        if config.get('type') != 'preference':
+            raise ValueError(f"{dataset_name} is not a preference dataset (type: {config.get('type')})")
+        
+        logger.info(f"Loading preference dataset: {dataset_name}")
+        
+        load_kwargs = {
+            'path': config['hf_path'],
+            'split': split,
+            'cache_dir': HF_CACHE_DIR,
+        }
+        
+        # 添加config参数（如果存在）
+        if 'config' in config:
+            load_kwargs['name'] = config['config']
+        
+        self.dataset = load_dataset(**load_kwargs)
+        
+        self.chosen_field = config.get('chosen_field', 'chosen')
+        self.rejected_field = config.get('rejected_field', 'rejected')
+        
+        if max_samples and len(self.dataset) > max_samples:
+            self.dataset = self.dataset.select(range(max_samples))
+        
+        logger.info(f"Loaded {len(self.dataset)} preference pairs")
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        try:
+            sample = self.dataset[idx]
+            
+            chosen_text = sample.get(self.chosen_field, '')
+            rejected_text = sample.get(self.rejected_field, '')
+            
+            if not chosen_text or not rejected_text:
+                return None
+            
+            # Tokenize
+            chosen_enc = self.tokenizer(
+                chosen_text,
+                max_length=self.max_length,
+                truncation=True,
+                padding='max_length',
+                return_tensors='pt'
+            )
+            
+            rejected_enc = self.tokenizer(
+                rejected_text,
+                max_length=self.max_length,
+                truncation=True,
+                padding='max_length',
+                return_tensors='pt'
+            )
+            
+            return (
+                chosen_enc['input_ids'].squeeze(0),
+                rejected_enc['input_ids'].squeeze(0),
+                chosen_enc['attention_mask'].squeeze(0),
+                rejected_enc['attention_mask'].squeeze(0)
+            )
+        
+        except Exception as e:
+            logger.debug(f"Error getting preference item at index {idx}: {e}")
+            return None
+
+
+def collate_fn_v2(batch):
+    """改进的collate函数"""
+    # 过滤None
+    batch = [item for item in batch if item is not None]
+    
+    if not batch:
+        logger.warning("Empty batch after filtering None values")
+        # 返回一个空的占位batch而不是None
+        return {
+            'input_ids': torch.empty(0),
+            'attention_mask': torch.empty(0)
+        }
+    
+    # 检查是否是preference数据
+    if isinstance(batch[0], tuple):
+        if len(batch[0]) == 4:  # 包含attention_mask
+            chosen = torch.stack([item[0] for item in batch])
+            rejected = torch.stack([item[1] for item in batch])
+            chosen_mask = torch.stack([item[2] for item in batch])
+            rejected_mask = torch.stack([item[3] for item in batch])
+            return {
+                'chosen': chosen,
+                'rejected': rejected,
+                'chosen_mask': chosen_mask,
+                'rejected_mask': rejected_mask
+            }
+        else:  # 旧格式兼容
+            chosen = torch.stack([item[0] for item in batch])
+            rejected = torch.stack([item[1] for item in batch])
+            return {'chosen': chosen, 'rejected': rejected}
+
+    # 普通数据
+    keys = batch[0].keys()
+    collated = {}
+
+    for key in keys:
+        if key in ['instruction', 'response', 'instruction_mask', 
+                   'response_mask', 'input_ids', 'attention_mask']:
+            tensors = [item[key] for item in batch if item.get(key) is not None]
+            if tensors:
+                collated[key] = torch.stack(tensors)
+            else:
+                collated[key] = None
+        elif key == 'modality_data':
+            # 处理多模态数据
+            modality_list = [item[key] for item in batch if item.get(key) is not None]
+            if modality_list and any(m is not None for m in modality_list):
+                # 收集图像
+                images = [m.get('image') for m in modality_list if m and 'image' in m]
+                if images:
+                    collated[key] = {'image': torch.stack(images)}
+                else:
+                    collated[key] = None
+            else:
+                collated[key] = None
+        else:
+            collated[key] = [item[key] for item in batch]
+
+    return collated
+
+
+def create_pretrain_dataloader(
+    mix_name: str = 'default',
+    tokenizer=None,
+    batch_size: int = 8,
+    num_workers: int = 4,
+    max_length: int = 2048,
+    max_samples: Optional[int] = None
+):
+    """创建预训练数据加载器"""
+    dataset = PreTrainDataset(
+        mix_name=mix_name,
+        tokenizer=tokenizer,
+        max_length=max_length,
+        streaming=True,
+        max_samples=max_samples
+    )
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        collate_fn=collate_fn_v2
+    )
+
+
+def create_posttrain_dataloader(
+    mix_name: str = 'default',
+    tokenizer=None,
+    batch_size: int = 8,
+    num_workers: int = 4,
+    max_length: int = 2048,
+    max_samples: Optional[int] = None,
+    split: str = 'train',
+    shuffle: bool = True
+):
+    """创建后训练数据加载器"""
+    dataset = PostTrainDataset(
+        mix_name=mix_name,
+        tokenizer=tokenizer,
+        max_length=max_length,
+        max_samples=max_samples,
+        split=split
+    )
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        collate_fn=collate_fn_v2,
+        pin_memory=True,
+        drop_last=False  # 保留最后一个batch
+    )
+
+
+def create_preference_dataloader(
+    dataset_name: str = 'hh_rlhf',
+    tokenizer=None,
+    batch_size: int = 8,
+    num_workers: int = 4,
+    max_length: int = 1024,
+    max_samples: Optional[int] = None,
+    split: str = 'train',
+    shuffle: bool = True
+):
+    """创建偏好数据加载器"""
+    dataset = PreferenceDataset(
+        dataset_name=dataset_name,
+        tokenizer=tokenizer,
+        max_length=max_length,
+        max_samples=max_samples,
+        split=split
+    )
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        collate_fn=collate_fn_v2,
+        pin_memory=True
+    )

encoders.py

@@ -0,0 +1,559 @@
+"""
+改进的多模态编码器 - SOTA级别（修复版）
+集成最新的视觉、音频、视频编码技术
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Tuple, Optional
+from components import RMSNorm, SwiGLU
+from transformer import OptimizedTransformerBlock
+import math
+
+class LayerScale(nn.Module):
+    """LayerScale - 改进训练稳定性"""
+    def __init__(self, dim: int, init_values: float = 1e-5):
+        super().__init__()
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.gamma
+
+class StochasticDepth(nn.Module):
+    """随机深度 - Drop Path"""
+    def __init__(self, drop_prob: float = 0.0):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not self.training or self.drop_prob == 0.0:
+            return x
+        
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+        random_tensor.floor_()
+        return x.div(keep_prob) * random_tensor
+
+class ImprovedPatchEmbedding(nn.Module):
+    """改进的图像分块嵌入 - 支持重叠patch和多尺度"""
+    def __init__(
+        self,
+        patch_size: int = 14,
+        in_channels: int = 3,
+        embed_dim: int = 2048,
+        overlap: int = 0
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        stride = patch_size - overlap
+        self.proj = nn.Conv2d(
+            in_channels,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=stride,
+            padding=overlap // 2
+        )
+        
+        self.norm = RMSNorm(embed_dim)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Tuple[int, int]]:
+        B, C, H, W = x.shape
+        x = self.proj(x)
+        grid_size = (x.shape[2], x.shape[3])
+        x = x.flatten(2).transpose(1, 2)
+        x = self.norm(x)
+        return x, grid_size
+
+class ImprovedVisionBlock(nn.Module):
+    """改进的Vision Transformer Block"""
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        dropout: float = 0.0,
+        drop_path: float = 0.0,
+        use_adapter: bool = False,
+        adapter_dim: int = 64,
+        use_layer_scale: bool = True,
+        layer_scale_init: float = 1e-5
+    ):
+        super().__init__()
+        self.norm1 = RMSNorm(dim)
+        self.attn = nn.MultiheadAttention(
+            dim, n_heads, dropout=dropout, batch_first=True
+        )
+        
+        self.norm2 = RMSNorm(dim)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * 4, dim),
+            nn.Dropout(dropout)
+        )
+        
+        self.drop_path = StochasticDepth(drop_path) if drop_path > 0 else nn.Identity()
+        
+        if use_layer_scale:
+            self.ls1 = LayerScale(dim, layer_scale_init)
+            self.ls2 = LayerScale(dim, layer_scale_init)
+        else:
+            self.ls1 = nn.Identity()
+            self.ls2 = nn.Identity()
+        
+        # 修复：使用简单的adapter实现，避免外部依赖
+        if use_adapter:
+            self.adapter = nn.Sequential(
+                nn.Linear(dim, adapter_dim),
+                nn.GELU(),
+                nn.Linear(adapter_dim, dim)
+            )
+        else:
+            self.adapter = None
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # 注意力
+        normx = self.norm1(x)
+        attn_out, _ = self.attn(normx, normx, normx)
+        x = x + self.drop_path(self.ls1(attn_out))
+        
+        # MLP
+        x = x + self.drop_path(self.ls2(self.mlp(self.norm2(x))))
+        
+        # Adapter
+        if self.adapter is not None:
+            x = x + self.adapter(x)
+        
+        return x
+
+class ImprovedVisionTransformer(nn.Module):
+    """
+    改进的视觉Transformer
+    - LayerScale
+    - Stochastic Depth
+    - 改进的位置编码
+    - 可选的Register tokens
+    """
+    def __init__(
+        self,
+        img_size: int = 224,
+        patch_size: int = 14,
+        in_channels: int = 3,
+        embed_dim: int = 2048,
+        depth: int = 24,
+        n_heads: int = 16,
+        dropout: float = 0.0,
+        drop_path_rate: float = 0.1,
+        use_register_tokens: bool = True,
+        num_register_tokens: int = 4,
+        use_adapter: bool = False,
+        adapter_dim: int = 64,
+        use_layer_scale: bool = True,
+        layer_scale_init: float = 1e-5
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+        self.use_register_tokens = use_register_tokens
+        self.num_register_tokens = num_register_tokens if use_register_tokens else 0
+        
+        # Patch embedding
+        self.patch_embed = ImprovedPatchEmbedding(
+            patch_size, in_channels, embed_dim, overlap=0
+        )
+        
+        self.pretrain_img_size = img_size
+        n_patches_pretrain = (img_size // patch_size) ** 2
+        
+        # CLS token
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        
+        # Register tokens (DINOv2启发)
+        if use_register_tokens:
+            self.register_tokens = nn.Parameter(
+                torch.zeros(1, num_register_tokens, embed_dim)
+            )
+        
+        # 修复：位置编码总数 = 1(CLS) + n_patches + register_tokens
+        total_tokens = 1 + n_patches_pretrain + self.num_register_tokens
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, total_tokens, embed_dim)
+        )
+        self.pos_drop = nn.Dropout(dropout)
+        
+        # Stochastic depth
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+        
+        # Transformer blocks
+        self.blocks = nn.ModuleList([
+            ImprovedVisionBlock(
+                embed_dim,
+                n_heads,
+                dropout,
+                drop_path=dpr[i],
+                use_adapter=use_adapter,
+                adapter_dim=adapter_dim,
+                use_layer_scale=use_layer_scale,
+                layer_scale_init=layer_scale_init
+            )
+            for i in range(depth)
+        ])
+        
+        self.norm = RMSNorm(embed_dim)
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.trunc_normal_(self.cls_token, std=0.02)
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+        if self.use_register_tokens:
+            nn.init.trunc_normal_(self.register_tokens, std=0.02)
+        
+        self.apply(self._init_module_weights)
+
+    def _init_module_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Conv2d):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, RMSNorm):
+            if hasattr(m, 'weight') and m.weight is not None:
+                nn.init.ones_(m.weight)
+
+    def _interpolate_pos_encoding(
+        self,
+        patch_tokens: torch.Tensor,
+        grid_size: Tuple[int, int]
+    ) -> torch.Tensor:
+        """
+        修复：改进的位置编码插值
+        只对patch位置编码进行插值，CLS和register token位置编码保持不变
+        """
+        pretrain_grid_h = self.pretrain_img_size // self.patch_size
+        pretrain_grid_w = pretrain_grid_h
+        
+        # 如果尺寸匹配，直接返回原始位置编码
+        if grid_size[0] == pretrain_grid_h and grid_size[1] == pretrain_grid_w:
+            return self.pos_embed
+        
+        # 分离不同部分的位置编码
+        # pos_embed结构: [CLS(1), register_tokens(n), patches(H*W)]
+        num_extra_tokens = 1 + self.num_register_tokens
+        cls_register_pos = self.pos_embed[:, :num_extra_tokens, :]  # [1, 1+n, dim]
+        patch_pos_embed = self.pos_embed[:, num_extra_tokens:, :]  # [1, H*W, dim]
+        
+        # 2D插值patch位置编码
+        patch_pos_embed = patch_pos_embed.reshape(
+            1, pretrain_grid_h, pretrain_grid_w, -1
+        ).permute(0, 3, 1, 2)
+        
+        patch_pos_embed = F.interpolate(
+            patch_pos_embed,
+            size=grid_size,
+            mode='bicubic',
+            align_corners=False
+        )
+        
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).flatten(1, 2)
+        
+        # 拼接回去
+        return torch.cat([cls_register_pos, patch_pos_embed], dim=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B = x.shape[0]
+        
+        # Patch embedding
+        x, grid_size = self.patch_embed(x)
+        
+        # 添加CLS token
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        
+        # 修复：正确组装tokens序列
+        if self.use_register_tokens:
+            register_tokens = self.register_tokens.expand(B, -1, -1)
+            # 顺序: [CLS, register_tokens, patches]
+            x = torch.cat([cls_tokens, register_tokens, x], dim=1)
+        else:
+            x = torch.cat([cls_tokens, x], dim=1)
+        
+        # 位置编码（插值以适应不同尺寸）
+        pos_embed = self._interpolate_pos_encoding(x, grid_size)
+        x = self.pos_drop(x + pos_embed)
+        
+        # Transformer blocks
+        for block in self.blocks:
+            x = block(x)
+        
+        x = self.norm(x)
+        
+        # 返回所有tokens（调用者可以选择使用CLS token或全局池化）
+        return x
+
+class ImprovedAudioEncoder(nn.Module):
+    """
+    改进的音频编码器
+    - 时序建模
+    - 频率建模
+    - 双流架构
+    """
+    def __init__(
+        self,
+        n_mels: int = 128,
+        target_length: int = 1024,
+        embed_dim: int = 2048,
+        depth: int = 12,
+        n_heads: int = 16,
+        patch_size: int = 16,
+        dropout: float = 0.1,
+        use_adapter: bool = False,
+        adapter_dim: int = 64,
+        use_dual_stream: bool = True
+    ):
+        super().__init__()
+        self.use_dual_stream = use_dual_stream
+        self.patch_size = patch_size
+        
+        # 主编码器
+        self.patch_embed = nn.Conv2d(
+            1, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+        
+        # 修复：计算实际的patch数量
+        self.n_patches_h = n_mels // patch_size
+        self.n_patches_w = target_length // patch_size
+        n_patches = self.n_patches_h * self.n_patches_w
+        
+        self.pos_embed = nn.Parameter(torch.zeros(1, n_patches, embed_dim))
+        self.pos_drop = nn.Dropout(dropout)
+        
+        # Transformer blocks
+        self.blocks = nn.ModuleList([
+            OptimizedTransformerBlock(
+                embed_dim, n_heads, None, None, dropout,
+                use_adapter=use_adapter, adapter_dim=adapter_dim
+            )
+            for _ in range(depth)
+        ])
+        
+        # 双流：时间流和频率流
+        if use_dual_stream:
+            # 修复：使用正确的池化维度
+            self.temporal_pool = nn.AdaptiveAvgPool1d(1)
+            self.frequency_pool = nn.AdaptiveAvgPool1d(1)
+            
+            self.temporal_proj = nn.Linear(embed_dim, embed_dim)
+            self.frequency_proj = nn.Linear(embed_dim, embed_dim)
+            
+            self.fusion = nn.Linear(embed_dim * 2, embed_dim)
+        
+        self.norm = RMSNorm(embed_dim)
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+        self.apply(self._init_module_weights)
+
+    def _init_module_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Conv2d):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+
+    def forward(self, mel_spec: torch.Tensor) -> torch.Tensor:
+        if mel_spec.ndim == 3:
+            mel_spec = mel_spec.unsqueeze(1)
+        
+        # Patch embedding
+        x = self.patch_embed(mel_spec)  # [B, C, H, W]
+        x = x.flatten(2).transpose(1, 2)  # [B, H*W, C]
+        x = self.pos_drop(x + self.pos_embed)
+        
+        # Transformer encoding
+        for block in self.blocks:
+            x, _, _ = block(x)
+        
+        x = self.norm(x)
+        
+        # 修复：双流处理
+        if self.use_dual_stream:
+            B, N, C = x.shape
+            
+            # 重塑为2D网格
+            x_2d = x.transpose(1, 2).reshape(B, C, self.n_patches_h, self.n_patches_w)
+            
+            # 时间流：沿频率维度池化（保留时间）
+            temporal = x_2d.mean(dim=2)  # [B, C, W]
+            temporal = self.temporal_pool(temporal).squeeze(-1)  # [B, C]
+            temporal = self.temporal_proj(temporal).unsqueeze(1)  # [B, 1, C]
+            
+            # 频率流：沿时间维度池化（保留频率）
+            frequency = x_2d.mean(dim=3)  # [B, C, H]
+            frequency = self.frequency_pool(frequency).squeeze(-1)  # [B, C]
+            frequency = self.frequency_proj(frequency).unsqueeze(1)  # [B, 1, C]
+            
+            # 融合
+            x = self.fusion(torch.cat([temporal, frequency], dim=-1))
+        else:
+            # 简单全局平均池化
+            x = x.mean(dim=1, keepdim=True)
+        
+        return x
+
+class ImprovedVideoEncoder(nn.Module):
+    """
+    改进的视频编码器
+    - 因果时序建模
+    - 时空分离注意力
+    - 可选的3D卷积
+    """
+    def __init__(
+        self,
+        img_size: int = 224,
+        patch_size: int = 14,
+        in_channels: int = 3,
+        embed_dim: int = 2048,
+        spatial_depth: int = 12,
+        temporal_depth: int = 4,
+        n_heads: int = 16,
+        num_frames: int = 16,
+        dropout: float = 0.1,
+        use_adapter: bool = False,
+        adapter_dim: int = 64,
+        use_3d_conv: bool = False
+    ):
+        super().__init__()
+        self.num_frames = num_frames
+        self.use_3d_conv = use_3d_conv
+        self.patch_size = patch_size
+        self.img_size = img_size
+        
+        if use_3d_conv:
+            # 3D卷积处理时空信息
+            self.patch_embed = nn.Conv3d(
+                in_channels,
+                embed_dim,
+                kernel_size=(2, patch_size, patch_size),
+                stride=(2, patch_size, patch_size)
+            )
+            # 修复：计算3D卷积后的尺寸
+            self.n_temporal_patches = num_frames // 2
+            self.n_spatial_patches = (img_size // patch_size) ** 2
+        else:
+            # 2D卷积 + 时序建模
+            self.patch_embed = ImprovedPatchEmbedding(
+                patch_size, in_channels, embed_dim
+            )
+            self.n_spatial_patches = (img_size // patch_size) ** 2
+        
+        # 空间位置编码
+        self.spatial_pos_embed = nn.Parameter(
+            torch.zeros(1, self.n_spatial_patches, embed_dim)
+        )
+        self.spatial_pos_drop = nn.Dropout(dropout)
+        
+        # 空间编码器
+        self.spatial_blocks = nn.ModuleList([
+            OptimizedTransformerBlock(
+                embed_dim, n_heads, None, None, dropout,
+                use_adapter=use_adapter, adapter_dim=adapter_dim
+            )
+            for _ in range(spatial_depth)
+        ])
+        
+        # 时间位置编码
+        if use_3d_conv:
+            self.temporal_pos_embed = nn.Parameter(
+                torch.zeros(1, self.n_temporal_patches, embed_dim)
+            )
+        else:
+            self.temporal_pos_embed = nn.Parameter(
+                torch.zeros(1, num_frames, embed_dim)
+            )
+        self.temporal_pos_drop = nn.Dropout(dropout)
+        
+        # 时序编码器
+        self.temporal_blocks = nn.ModuleList([
+            OptimizedTransformerBlock(
+                embed_dim, n_heads, None, None, dropout,
+                use_adapter=use_adapter, adapter_dim=adapter_dim
+            )
+            for _ in range(temporal_depth)
+        ])
+        
+        self.norm = RMSNorm(embed_dim)
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.trunc_normal_(self.spatial_pos_embed, std=0.02)
+        nn.init.trunc_normal_(self.temporal_pos_embed, std=0.02)
+        self.apply(self._init_module_weights)
+
+    def _init_module_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, (nn.Conv2d, nn.Conv3d)):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, T, C, H, W = x.shape
+        
+        if self.use_3d_conv:
+            # 修复：3D卷积路径
+            x = x.transpose(1, 2)  # [B, C, T, H, W]
+            x = self.patch_embed(x)  # [B, embed_dim, T', H', W']
+            
+            # 重塑: [B, D, T', H'*W'] -> [B, T', H'*W', D]
+            B, D, T_new, H_new, W_new = x.shape
+            x = x.view(B, D, T_new, -1).permute(0, 2, 3, 1)  # [B, T', H'*W', D]
+            
+            # 空间位置编码（每帧独立）
+            x = x + self.spatial_pos_embed.unsqueeze(1)
+            
+            # 逐帧空间编码
+            x_flat = x.reshape(B * T_new, -1, D)
+            for block in self.spatial_blocks:
+                x_flat, _, _ = block(x_flat)
+            
+            # 重塑回时序维度
+            x = x_flat.view(B, T_new, -1, D)
+            
+            # 修复：时序聚合 - 使用平均池化而非取第一个token
+            x = x.mean(dim=2)  # [B, T', D]
+            
+        else:
+            # 2D卷积 + 分离时空建模
+            x_flat = x.view(B * T, C, H, W)
+            x_patched, grid_size = self.patch_embed(x_flat)
+            
+            # 空间位置编码
+            x_patched = self.spatial_pos_drop(x_patched + self.spatial_pos_embed)
+            
+            # 空间编码
+            for block in self.spatial_blocks:
+                x_patched, _, _ = block(x_patched)
+            
+            # 修复：时序聚合 - 全局平均池化而非仅mean(dim=2)
+            _, N, D = x_patched.shape
+            x_spatial = x_patched.view(B, T, N, D)
+            x = x_spatial.mean(dim=2)  # [B, T, D] - 对每帧的所有patch取平均
+        
+        # 时序位置编码
+        x = self.temporal_pos_drop(x + self.temporal_pos_embed)
+        
+        # 时序编码
+        for block in self.temporal_blocks:
+            x, _, _ = block(x)
+        
+        return self.norm(x)

gradio1.py

@@ -0,0 +1,228 @@
+"""
+Gradio 推理界面 - 多模态 Dense Transformer (适配 Qwen Tokenizer 版)
+
+用法:
+pip install -r requirements.txt
+# requirements.txt 至少包含:
+# torch>=1.12, transformers, pillow, gradio
+python app_gradio.py --checkpoint /path/to/final_model.pt --tokenizer Qwen/Qwen2.5-7B-Instruct --port 7860 --share False
+"""
+
+import os
+import argparse
+from pathlib import Path
+import json
+from typing import Optional
+
+import torch
+from PIL import Image
+from transformers import AutoTokenizer
+
+# UI
+import gradio as gr
+
+# 本项目代码引用（按你的工程结构调整）
+from model import MultiModalDenseTransformer
+from continual_learning import UnifiedMultiModalPreprocessor
+
+# 设置国内镜像（如需要）
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+
+# ---- 与你原来保持一致的图像预处理 ----
+from torchvision import transforms
+image_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225]),
+])
+
+# -------- ModelInference 类（轻微改写） --------
+class ModelInference:
+    def __init__(self, checkpoint_path: str, tokenizer_name: str, config_path: Optional[str] = None, device: str = 'cuda' if torch.cuda.is_available() else 'cpu'):
+        self.device = torch.device(device)
+        print(f"Using device: {self.device}")
+        print(f"Loading tokenizer: {tokenizer_name}...")
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, use_fast=True, trust_remote_code=True)
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
+        if config_path and Path(config_path).exists():
+            with open(config_path, 'r') as f:
+                self.config = json.load(f)
+        else:
+            # 采用你原始脚本中的默认 config（可按需调整）
+            self.config = {
+                'model_dim': 1536,
+                'vocab_size': len(self.tokenizer),
+                'n_layers': 12,
+                'n_heads': 12,
+                'n_kv_heads': 4,
+                'head_dim': None,
+                'max_seq_len': 512,
+                'dropout': 0.0,
+                'use_moe': False,
+                'use_adapter': False,
+                'use_lora': False,
+                'rope_scaling_type': "yarn",
+                'use_multimodal_fusion': False,
+                'use_contrastive': False
+            }
+
+        # init model + preprocessor
+        print("Initializing model architecture...")
+        self.model = MultiModalDenseTransformer(**self.config)
+        self.preprocessor = UnifiedMultiModalPreprocessor(model_dim=self.config['model_dim'])
+
+        print(f"Loading checkpoint from {checkpoint_path}...")
+        checkpoint = torch.load(checkpoint_path, map_location=self.device)
+        # 支持 checkpoint 包含 'model_state_dict' 的情况
+        state_dict = checkpoint.get('model_state_dict', checkpoint) if isinstance(checkpoint, dict) else checkpoint
+
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            if k.startswith('module.'):
+                new_state_dict[k[7:]] = v
+            else:
+                new_state_dict[k] = v
+
+        missing, unexpected = self.model.load_state_dict(new_state_dict, strict=False)
+        if missing:
+            print(f"Warning: Missing keys: {len(missing)}")
+        if unexpected:
+            print(f"Warning: Unexpected keys: {len(unexpected)}")
+
+        self.model.to(self.device)
+        self.preprocessor.to(self.device)
+        self.model.eval()
+        print("Model loaded successfully!")
+        print(f"Total parameters: {sum(p.numel() for p in self.model.parameters())/1e6:.2f}M")
+
+    @torch.no_grad()
+    def generate_text(self, prompt: str, max_new_tokens: int = 128, temperature: float = 0.7, top_k: int = 10, top_p: float = 0.9, repetition_penalty: float = 1.2, image: Optional[Image.Image] = None) -> str:
+        formatted_prompt = f"Instruction: {prompt}\nResponse:"
+        inputs = self.tokenizer(formatted_prompt, return_tensors="pt")
+        input_ids = inputs['input_ids'].to(self.device)
+
+        input_data = {'segments': []}
+        if image is not None:
+            try:
+                if image.mode != 'RGB':
+                    image = image.convert('RGB')
+                image_tensor = image_transform(image).unsqueeze(0).to(self.device)
+                mod_segments = self.preprocessor.process_batch(image_tensor, 'image')
+                for seg in mod_segments:
+                    input_data['segments'].append(seg)
+            except Exception as e:
+                print(f"Warning: Image processing skipped due to error: {e}")
+
+        input_data['segments'].append({
+            'type': 'text',
+            'data': input_ids,
+            'modality_id': 0
+        })
+
+        try:
+            generated_ids = self.model.generate(
+                input_data,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                do_sample=True,
+                eos_token_id=self.tokenizer.eos_token_id,
+                pad_token_id=self.tokenizer.pad_token_id
+            )
+
+            full_output = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+            # 提取 Response 后的部分并做 stop 处理
+            if "Response:" in full_output:
+                answer = full_output.split("Response:")[-1].strip()
+            else:
+                answer = full_output
+
+            stop_words = ["Instruction", "Input", "###", "Response", "User:", "Assistant:", "\n\n"]
+            for sw in stop_words:
+                if sw in answer:
+                    answer = answer.split(sw)[0].strip()
+
+            # 去掉可能的 echo
+            lines = answer.split('\n')
+            if len(lines) > 0 and prompt.lower() in lines[0].lower():
+                answer = "\n".join(lines[1:]).strip()
+            return answer
+        except Exception as e:
+            import traceback
+            traceback.print_exc()
+            return f"Error: {e}"
+
+# -------- Gradio UI 部分 --------
+def build_ui(model_instance):
+    with gr.Blocks(title="MultiModal Dense Transformer - Gradio", css="""
+        .gradio-container { max-width: 900px; margin: auto; }
+    """) as demo:
+        gr.Markdown("## 🚀 多模态在线推理（文本 + 图片）")
+        with gr.Row():
+            with gr.Column(scale=3):
+                txt = gr.Textbox(label="Prompt (Instruction)", placeholder="请输入指令或问题...", lines=5)
+                img = gr.Image(type="pil", label="(可选) 上传图片（支持多模态）")
+                btn = gr.Button("生成 (Generate)")
+            with gr.Column(scale=2):
+                max_tokens = gr.Slider(label="Max New Tokens", minimum=16, maximum=1024, step=1, value=128)
+                temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.5, step=0.01, value=0.7)
+                top_k = gr.Slider(label="Top-k", minimum=0, maximum=200, step=1, value=40)
+                top_p = gr.Slider(label="Top-p", minimum=0.0, maximum=1.0, step=0.01, value=0.9)
+                rep_pen = gr.Slider(label="Repetition Penalty", minimum=0.5, maximum=2.0, step=0.01, value=1.1)
+                status = gr.Textbox(label="Status", value="Ready", interactive=False)
+        output = gr.Textbox(label="Output", lines=12, interactive=False)
+
+        def gr_generate(prompt, image, max_tokens_v, temp_v, topk_v, topp_v, rep_v):
+            if not prompt or str(prompt).strip() == "":
+                return "", "请输入 Prompt", ""
+            status_msg = "Generating..."
+            # call model
+            out = model_instance.generate_text(prompt=prompt,
+                                               max_new_tokens=int(max_tokens_v),
+                                               temperature=float(temp_v),
+                                               top_k=int(topk_v),
+                                               top_p=float(topp_v),
+                                               repetition_penalty=float(rep_v),
+                                               image=image)
+            return out, "Done", ""
+
+        btn.click(fn=gr_generate, inputs=[txt, img, max_tokens, temperature, top_k, top_p, rep_pen], outputs=[output, status, gr.State()])
+
+        demo.launch(share=True)
+
+    return demo
+
+# -------- CLI / main --------
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", type=str, default="/root/multimodal/checkpoints/posttrain/final_model.pt")
+    parser.add_argument("--tokenizer", type=str, default="Qwen/Qwen2.5-7B-Instruct")
+    parser.add_argument("--config", type=str, default=None)
+    parser.add_argument("--port", type=int, default=7860)
+    parser.add_argument("--share", type=lambda x: x.lower() in ("true","1","yes"), default=True)
+    args = parser.parse_args()
+
+    # 如果 default 的 final_model 不存在，尝试寻找最近 step
+    if not Path(args.checkpoint).exists():
+        possible = list(Path("checkpoints/pretrain").glob("step_*.pt"))
+        if possible:
+            args.checkpoint = str(possible[-1])
+            print(f"未找到 final_model.pt，使用最新 checkpoint: {args.checkpoint}")
+        else:
+            raise FileNotFoundError(f"找不到检查点: {args.checkpoint}")
+
+    global model_instance
+    model_instance = ModelInference(args.checkpoint, args.tokenizer, args.config)
+
+    # 启动 Gradio（使用 share 参数决定是否创建公网链接）
+    demo = build_ui(model_instance)
+    demo.launch(server_port=args.port, share=args.share)
+
+if __name__ == "__main__":
+    main()

grpo.py

@@ -0,0 +1,630 @@
+"""
+改进的 GRPO (Group Relative Policy Optimization) 训练器
+修复了所有已知问题
+"""
+
+import torch
+import torch.optim as optim
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, TensorDataset
+from typing import Dict, List, Tuple, Optional
+from tqdm import tqdm
+import numpy as np
+import gc
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class GRPOTrainer:
+    """
+    GRPO训练器 - Group Relative Policy Optimization
+    参考 DeepSeekMath/DeepSeek-V3 策略
+    
+    主要修复：
+    1. 修复了 generate() 返回格式问题
+    2. 修复了 reward_model 输出处理
+    3. 添加了完整的混合精度训练支持
+    4. 改进了 KL 散度计算的数值稳定性
+    5. 修复了 past_key_values 的使用逻辑
+    6. 改进了内存管理和错误处理
+    """
+    
+    def __init__(
+        self,
+        actor_model,
+        reward_model,
+        ref_model,
+        tokenizer,
+        learning_rate: float = 1e-6,
+        kl_coef: float = 0.04,
+        group_size: int = 4,
+        clip_epsilon: float = 0.2,
+        max_grad_norm: float = 1.0,
+        grpo_epochs: int = 1,
+        update_batch_size: int = 4,
+        use_amp: bool = True,
+        value_clip: bool = False,
+        entropy_coef: float = 0.01,
+        advantage_normalization: str = 'group',  # 'group', 'global', 'none'
+        kl_estimation_method: str = 'forward'  # 'forward', 'reverse', 'symmetric'
+    ):
+        """
+        初始化GRPO训练器
+        
+        Args:
+            actor_model: 要训练的策略模型
+            reward_model: 奖励模型（冻结）
+            ref_model: 参考模型（冻结）
+            tokenizer: 分词器
+            learning_rate: 学习率
+            kl_coef: KL散度惩罚系数
+            group_size: 每个prompt生成的样本数
+            clip_epsilon: PPO clip范围
+            max_grad_norm: 梯度裁剪阈值
+            grpo_epochs: 每批经验的训练轮数
+            update_batch_size: 更新时的mini-batch大小
+            use_amp: 是否使用混合精度训练
+            value_clip: 是否对value进行clip（当前未使用value网络）
+            entropy_coef: 熵正则化系数
+            advantage_normalization: 优势函数归一化方式
+            kl_estimation_method: KL散度估计方法
+        """
+        self.actor = actor_model
+        self.reward_model = reward_model
+        self.ref_model = ref_model
+        self.tokenizer = tokenizer
+        
+        self.kl_coef = kl_coef
+        self.group_size = group_size
+        self.clip_epsilon = clip_epsilon
+        self.max_grad_norm = max_grad_norm
+        self.grpo_epochs = grpo_epochs
+        self.update_batch_size = update_batch_size
+        self.use_amp = use_amp
+        self.entropy_coef = entropy_coef
+        self.advantage_normalization = advantage_normalization
+        self.kl_estimation_method = kl_estimation_method
+        
+        self.device = next(actor_model.parameters()).device
+        
+        # 冻结参考模型和奖励模型
+        self.ref_model.eval()
+        self.ref_model.requires_grad_(False)
+        self.reward_model.eval()
+        self.reward_model.requires_grad_(False)
+        
+        # 优化器配置
+        self.optimizer = optim.AdamW(
+            filter(lambda p: p.requires_grad, actor_model.parameters()),
+            lr=learning_rate,
+            weight_decay=0.01,
+            betas=(0.9, 0.95),
+            eps=1e-8
+        )
+        
+        # 混合精度训练 - 修复：添加 GradScaler
+        self.scaler = torch.amp.GradScaler('cuda', enabled=self.use_amp)
+        
+        # 训练统计
+        self.training_stats = {
+            'iterations': 0,
+            'total_samples': 0,
+            'avg_rewards': [],
+            'avg_kl': [],
+            'policy_losses': []
+        }
+        
+        logger.info(f"GRPO Trainer initialized:")
+        logger.info(f"  Group Size: {group_size}")
+        logger.info(f"  KL Coef: {kl_coef}")
+        logger.info(f"  Clip Epsilon: {clip_epsilon}")
+        logger.info(f"  Learning Rate: {learning_rate}")
+        logger.info(f"  Update Batch Size: {update_batch_size}")
+        logger.info(f"  Mixed Precision: {use_amp}")
+        logger.info(f"  KL Estimation: {kl_estimation_method}")
+
+    def _compute_kl_divergence(
+        self,
+        log_probs: torch.Tensor,
+        ref_log_probs: torch.Tensor,
+        mask: torch.Tensor
+    ) -> torch.Tensor:
+        """
+        计算KL散度（改进数值稳定性）
+        
+        Args:
+            log_probs: 当前策略的log概率
+            ref_log_probs: 参考策略的log概率
+            mask: 有效token的mask
+            
+        Returns:
+            KL散度（标量）
+        """
+        if self.kl_estimation_method == 'forward':
+            # KL(π||π_ref) = Σ π * log(π/π_ref)
+            # ≈ Σ exp(log_π) * (log_π - log_π_ref)
+            # 为了数值稳定，使用 log_π - log_π_ref
+            kl = log_probs - ref_log_probs
+        elif self.kl_estimation_method == 'reverse':
+            # KL(π_ref||π) = Σ π_ref * log(π_ref/π)
+            kl = ref_log_probs - log_probs
+        else:  # symmetric
+            # 对称KL散度
+            forward_kl = log_probs - ref_log_probs
+            reverse_kl = ref_log_probs - log_probs
+            kl = 0.5 * (forward_kl + reverse_kl)
+        
+        # 应用mask并求和
+        kl_penalty = (kl * mask).sum(dim=-1)
+        return kl_penalty
+
+    @torch.no_grad()
+    def generate_experience(
+        self,
+        prompts_dataloader: DataLoader,
+        max_gen_len: int,
+        temperature: float = 1.0,
+        top_p: float = 0.9
+    ) -> Dict:
+        """
+        生成经验数据：采样 -> 计算 LogProbs -> 计算 Rewards(含KL)
+        
+        修复：
+        1. 正确处理 generate() 的返回值
+        2. 修复 reward_model 的输出处理
+        3. 改进数值稳定性
+        """
+        self.actor.eval()
+        
+        all_sequences = []
+        all_log_probs = []
+        all_advantages = []
+        all_prompt_lens = []
+        all_rewards = []
+        
+        logger.info("Generating experience...")
+        
+        for prompts in tqdm(prompts_dataloader, desc="Generating Experience"):
+            try:
+                # 处理不同的输入格式
+                if isinstance(prompts, (list, tuple)):
+                    prompts = prompts[0]
+                
+                prompts = prompts.to(self.device)
+                batch_size = prompts.shape[0]
+                
+                # 扩展prompts以生成group_size个样本
+                prompts_repeated = prompts.repeat_interleave(self.group_size, dim=0)
+                prompt_len = prompts_repeated.shape[1]
+                
+                input_data = {
+                    'segments': [{
+                        'type': 'text',
+                        'data': prompts_repeated,
+                        'modality_id': 0
+                    }]
+                }
+                
+                # 1. 采样生成（修复：generate只返回新生成的tokens）
+                with torch.amp.autocast('cuda', enabled=self.use_amp):
+                    response_ids = self.actor.generate(
+                        input_data,
+                        max_new_tokens=max_gen_len,
+                        do_sample=True,
+                        temperature=temperature,
+                        top_p=top_p,
+                        eos_token_id=self.tokenizer.eos_token_id,
+                        pad_token_id=self.tokenizer.pad_token_id,
+                        use_cache=True  # 使用缓存加速生成
+                    )
+                
+                # 修复：拼接完整序列（prompt + response）
+                sequences = torch.cat([prompts_repeated, response_ids], dim=1)
+                
+                # 检查序列长度
+                if sequences.shape[1] <= prompt_len:
+                    logger.warning("Generated sequence too short, skipping batch")
+                    continue
+                
+                full_input_data = {
+                    'segments': [{
+                        'type': 'text',
+                        'data': sequences,
+                        'modality_id': 0
+                    }]
+                }
+                
+                # 2. 计算当前策略和参考策略的 LogProbs
+                with torch.amp.autocast('cuda', enabled=self.use_amp):
+                    actor_out = self.actor(full_input_data)
+                    ref_out = self.ref_model(full_input_data)
+                
+                logits = actor_out['logits'][:, :-1, :]
+                ref_logits = ref_out['logits'][:, :-1, :]
+                targets = sequences[:, 1:]
+                
+                # 计算log probabilities（改进数值稳定性）
+                log_probs = F.log_softmax(logits, dim=-1)
+                ref_log_probs = F.log_softmax(ref_logits, dim=-1)
+                
+                # 提取对应token的log概率
+                per_token_log_probs = torch.gather(
+                    log_probs, -1, targets.unsqueeze(-1)
+                ).squeeze(-1)
+                per_token_ref_log_probs = torch.gather(
+                    ref_log_probs, -1, targets.unsqueeze(-1)
+                ).squeeze(-1)
+                
+                # 3. 计算 KL 散度 (只针对response部分)
+                response_mask = torch.arange(
+                    sequences.size(1) - 1, device=self.device
+                ) >= (prompt_len - 1)
+                response_mask = response_mask.unsqueeze(0).expand_as(per_token_log_probs)
+                response_mask = response_mask.float()
+                
+                # 使用改进的KL计算
+                kl_penalty = self._compute_kl_divergence(
+                    per_token_log_probs,
+                    per_token_ref_log_probs,
+                    response_mask
+                )
+                
+                # 4. 计算环境奖励（修复：正确处理reward_model输出）
+                with torch.amp.autocast('cuda', enabled=self.use_amp):
+                    reward_output = self.reward_model(full_input_data)
+                
+                # reward_model返回 (batch_size, seq_len)，取最后一个位置的奖励
+                if reward_output.dim() == 2:
+                    raw_rewards = reward_output[:, -1]
+                else:
+                    raw_rewards = reward_output.squeeze(-1)
+                
+                # 5. 组合总奖励: R_total = R_env - β * KL
+                total_rewards = raw_rewards - self.kl_coef * kl_penalty
+                
+                # 6. 计算组内相对优势 (Group Relative Advantage)
+                rewards_grouped = total_rewards.view(batch_size, self.group_size)
+                
+                if self.advantage_normalization == 'group':
+                    # 组内标准化
+                    mean_grouped = rewards_grouped.mean(dim=1, keepdim=True)
+                    std_grouped = rewards_grouped.std(dim=1, keepdim=True) + 1e-8
+                    advantages = (rewards_grouped - mean_grouped) / std_grouped
+                elif self.advantage_normalization == 'global':
+                    # 全局标准化
+                    advantages = (rewards_grouped - rewards_grouped.mean()) / (
+                        rewards_grouped.std() + 1e-8
+                    )
+                else:  # 'none'
+                    advantages = rewards_grouped - rewards_grouped.mean(dim=1, keepdim=True)
+                
+                advantages = advantages.view(-1)
+                
+                # 保存数据
+                all_sequences.append(sequences.cpu())
+                all_log_probs.append(per_token_log_probs.detach().cpu())
+                all_advantages.append(advantages.detach().cpu())
+                all_prompt_lens.append(
+                    torch.full((sequences.size(0),), prompt_len, dtype=torch.long)
+                )
+                all_rewards.append(total_rewards.detach().cpu())
+                
+                # 清理中间变量
+                del logits, ref_logits, actor_out, ref_out
+                del log_probs, ref_log_probs, reward_output
+                
+            except Exception as e:
+                logger.error(f"Error generating experience for batch: {e}")
+                import traceback
+                traceback.print_exc()
+                continue
+            
+            finally:
+                torch.cuda.empty_cache()
+        
+        if not all_sequences:
+            raise RuntimeError("No valid sequences generated")
+        
+        # 合并所有数据
+        experience = {
+            'sequences': torch.cat(all_sequences, dim=0),
+            'log_probs': torch.cat(all_log_probs, dim=0),
+            'advantages': torch.cat(all_advantages, dim=0),
+            'prompt_lengths': torch.cat(all_prompt_lens, dim=0),
+            'rewards': torch.cat(all_rewards, dim=0)
+        }
+        
+        # 统计信息
+        logger.info(f"Generated {len(experience['sequences'])} sequences")
+        logger.info(f"Avg Reward: {experience['rewards'].mean().item():.4f}")
+        logger.info(f"Reward Std: {experience['rewards'].std().item():.4f}")
+        logger.info(f"Avg Advantage: {experience['advantages'].mean().item():.4f}")
+        
+        return experience
+
+    def grpo_step(
+        self,
+        dataset: TensorDataset
+    ) -> Dict[str, float]:
+        """
+        执行 GRPO 优化步骤
+        
+        修复：
+        1. 使用 GradScaler 进行混合精度训练
+        2. 改进损失计算
+        3. 更好的统计信息收集
+        """
+        self.actor.train()
+        
+        dataloader = DataLoader(
+            dataset,
+            batch_size=self.update_batch_size,
+            shuffle=True,
+            drop_last=False
+        )
+        
+        epoch_stats = {
+            'total_loss': 0.0,
+            'policy_loss': 0.0,
+            'entropy': 0.0,
+            'approx_kl': 0.0,
+            'clip_fraction': 0.0,
+            'steps': 0
+        }
+        
+        for batch_data in dataloader:
+            sequences, old_log_probs, advantages, prompt_lens = batch_data
+            
+            sequences = sequences.to(self.device)
+            old_log_probs = old_log_probs.to(self.device)
+            advantages = advantages.to(self.device)
+            
+            input_data = {
+                'segments': [{
+                    'type': 'text',
+                    'data': sequences,
+                    'modality_id': 0
+                }]
+            }
+            
+            # 修复：使用 GradScaler 进行混合精度训练
+            with torch.amp.autocast('cuda', enabled=self.use_amp):
+                outputs = self.actor(input_data)
+                logits = outputs['logits'][:, :-1, :]
+                
+                # 计算新的log probabilities
+                targets = sequences[:, 1:]
+                log_probs_dist = F.log_softmax(logits, dim=-1)
+                new_log_probs = torch.gather(
+                    log_probs_dist, -1, targets.unsqueeze(-1)
+                ).squeeze(-1)
+                
+                # 构建response mask
+                mask = torch.zeros_like(new_log_probs)
+                for i, pl in enumerate(prompt_lens):
+                    mask[i, pl-1:] = 1.0
+                
+                # 计算概率比率
+                ratio = torch.exp(new_log_probs - old_log_probs)
+                
+                # 扩展advantages到序列维度
+                adv_expanded = advantages.unsqueeze(-1).expand_as(new_log_probs)
+                
+                # PPO clip损失
+                surr1 = ratio * adv_expanded
+                surr2 = torch.clamp(
+                    ratio,
+                    1.0 - self.clip_epsilon,
+                    1.0 + self.clip_epsilon
+                ) * adv_expanded
+                
+                # 策略损失（最小化负目标）
+                policy_loss = -torch.min(surr1, surr2)
+                policy_loss = (policy_loss * mask).sum() / (mask.sum() + 1e-8)
+                
+                # 熵奖励（鼓励探索）
+                probs = F.softmax(logits, dim=-1)
+                entropy = -(probs * log_probs_dist).sum(dim=-1)
+                entropy_bonus = (entropy * mask).sum() / (mask.sum() + 1e-8)
+                
+                # 总损失
+                loss = policy_loss - self.entropy_coef * entropy_bonus
+                
+                # 统计信息
+                with torch.no_grad():
+                    log_ratio = new_log_probs - old_log_probs
+                    approx_kl = ((ratio - 1) - log_ratio) * mask
+                    approx_kl = approx_kl.sum() / (mask.sum() + 1e-8)
+                    
+                    clip_fraction = ((ratio > 1 + self.clip_epsilon) | 
+                                   (ratio < 1 - self.clip_epsilon)).float()
+                    clip_fraction = (clip_fraction * mask).sum() / (mask.sum() + 1e-8)
+            
+            # 修复：使用 GradScaler 进行反向传播
+            self.optimizer.zero_grad()
+            self.scaler.scale(loss).backward()
+            
+            # 梯度裁剪
+            self.scaler.unscale_(self.optimizer)
+            grad_norm = torch.nn.utils.clip_grad_norm_(
+                self.actor.parameters(),
+                self.max_grad_norm
+            )
+            
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+            
+            # 累积统计
+            epoch_stats['total_loss'] += loss.item()
+            epoch_stats['policy_loss'] += policy_loss.item()
+            epoch_stats['entropy'] += entropy_bonus.item()
+            epoch_stats['approx_kl'] += approx_kl.item()
+            epoch_stats['clip_fraction'] += clip_fraction.item()
+            epoch_stats['steps'] += 1
+        
+        # 计算平均值
+        for key in epoch_stats:
+            if key != 'steps':
+                epoch_stats[key] /= max(epoch_stats['steps'], 1)
+        
+        return epoch_stats
+
+    def train(
+        self,
+        prompt_dataloader: DataLoader,
+        num_iterations: int = 1,
+        max_gen_len: int = 50,
+        temperature: float = 1.0,
+        save_every: int = 5,
+        save_path: str = "checkpoints"
+    ):
+        """
+        完整的GRPO训练循环
+        
+        Args:
+            prompt_dataloader: 提供prompts的数据加载器
+            num_iterations: 训练迭代次数
+            max_gen_len: 最大生成长度
+            temperature: 采样温度
+            save_every: 每N次迭代保存一次checkpoint
+            save_path: checkpoint保存路径
+        """
+        logger.info(f"\n{'='*80}")
+        logger.info(f"Starting GRPO Training")
+        logger.info(f"  Iterations: {num_iterations}")
+        logger.info(f"  Max Gen Length: {max_gen_len}")
+        logger.info(f"  Temperature: {temperature}")
+        logger.info(f"{'='*80}\n")
+        
+        for iteration in range(num_iterations):
+            try:
+                # 1. 生成经验
+                experience = self.generate_experience(
+                    prompt_dataloader,
+                    max_gen_len,
+                    temperature
+                )
+                
+                dataset = TensorDataset(
+                    experience['sequences'],
+                    experience['log_probs'],
+                    experience['advantages'],
+                    experience['prompt_lengths']
+                )
+                
+                # 2. 策略优化（多个epoch）
+                logger.info(f"Optimizing policy for {self.grpo_epochs} epochs...")
+                all_epoch_stats = []
+                
+                for epoch in range(self.grpo_epochs):
+                    stats = self.grpo_step(dataset)
+                    all_epoch_stats.append(stats)
+                    
+                    logger.info(
+                        f"  Epoch {epoch+1}/{self.grpo_epochs} | "
+                        f"Loss: {stats['total_loss']:.4f} | "
+                        f"KL: {stats['approx_kl']:.4f} | "
+                        f"Clip%: {stats['clip_fraction']*100:.1f}"
+                    )
+                
+                # 3. 汇总统计
+                avg_stats = {
+                    key: np.mean([s[key] for s in all_epoch_stats])
+                    for key in all_epoch_stats[0].keys()
+                }
+                
+                self.training_stats['iterations'] += 1
+                self.training_stats['total_samples'] += len(experience['sequences'])
+                self.training_stats['avg_rewards'].append(
+                    experience['rewards'].mean().item()
+                )
+                self.training_stats['avg_kl'].append(avg_stats['approx_kl'])
+                self.training_stats['policy_losses'].append(avg_stats['policy_loss'])
+                
+                # 4. 打印进度
+                logger.info(f"\n{'='*80}")
+                logger.info(f"Iteration {iteration+1}/{num_iterations} Complete")
+                logger.info(f"  Avg Reward: {experience['rewards'].mean():.4f}")
+                logger.info(f"  Avg Advantage: {experience['advantages'].mean():.4f}")
+                logger.info(f"  Policy Loss: {avg_stats['policy_loss']:.4f}")
+                logger.info(f"  Approx KL: {avg_stats['approx_kl']:.4f}")
+                logger.info(f"  Entropy: {avg_stats['entropy']:.4f}")
+                logger.info(f"  Clip Fraction: {avg_stats['clip_fraction']*100:.1f}%")
+                logger.info(f"{'='*80}\n")
+                
+                # 5. 保存checkpoint
+                if (iteration + 1) % save_every == 0:
+                    self.save_checkpoint(
+                        f"{save_path}/grpo_iter_{iteration+1}.pt"
+                    )
+                
+                # 6. 清理内存
+                del experience, dataset
+                gc.collect()
+                torch.cuda.empty_cache()
+                
+            except Exception as e:
+                logger.error(f"Error in iteration {iteration+1}: {e}")
+                import traceback
+                traceback.print_exc()
+                continue
+        
+        logger.info("GRPO Training Complete!")
+        self.print_training_summary()
+
+    def save_checkpoint(self, path: str):
+        """保存训练checkpoint"""
+        import os
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        
+        checkpoint = {
+            'actor_state_dict': self.actor.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'scaler_state_dict': self.scaler.state_dict(),  # 修复：保存scaler状态
+            'training_stats': self.training_stats,
+            'config': {
+                'kl_coef': self.kl_coef,
+                'group_size': self.group_size,
+                'clip_epsilon': self.clip_epsilon,
+            }
+        }
+        
+        torch.save(checkpoint, path)
+        logger.info(f"Checkpoint saved to {path}")
+
+    def load_checkpoint(self, path: str):
+        """加载训练checkpoint"""
+        checkpoint = torch.load(path, map_location=self.device)
+        
+        self.actor.load_state_dict(checkpoint['actor_state_dict'])
+        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        
+        # 修复：加载scaler状态
+        if 'scaler_state_dict' in checkpoint and self.use_amp:
+            self.scaler.load_state_dict(checkpoint['scaler_state_dict'])
+        
+        self.training_stats = checkpoint['training_stats']
+        
+        logger.info(f"Checkpoint loaded from {path}")
+
+    def print_training_summary(self):
+        """打印训练摘要"""
+        logger.info("\n" + "="*80)
+        logger.info("Training Summary")
+        logger.info("="*80)
+        logger.info(f"Total Iterations: {self.training_stats['iterations']}")
+        logger.info(f"Total Samples: {self.training_stats['total_samples']}")
+        
+        if self.training_stats['avg_rewards']:
+            logger.info(
+                f"Final Avg Reward: "
+                f"{self.training_stats['avg_rewards'][-1]:.4f}"
+            )
+            logger.info(
+                f"Reward Improvement: "
+                f"{self.training_stats['avg_rewards'][-1] - self.training_stats['avg_rewards'][0]:.4f}"
+            )
+        
+        logger.info("="*80 + "\n")

infer.py

@@ -0,0 +1,372 @@
+"""
+Flask推理界面 - 多模态Dense Transformer (适配 Qwen Tokenizer 版)
+"""
+
+import os
+import torch
+import torch.nn.functional as F
+from flask import Flask, render_template, request, jsonify
+from transformers import AutoTokenizer
+from PIL import Image
+import json
+import io
+import base64
+from pathlib import Path
+from typing import Optional
+
+# 确保引入路径正确，根据你之前的文件结构
+from model import MultiModalDenseTransformer
+# 注意：UnifiedMultiModalPreprocessor 之前是在 continual_learning.py 中定义的
+# 如果你移动了它，请修改这里的导入路径
+from continual_learning import UnifiedMultiModalPreprocessor
+# 如果没有 image_transform，我们需要在这里定义或导入
+from torchvision import transforms
+
+# 设置国内镜像
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+
+# 定义图像预处理 (与 training 保持一致)
+image_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+class ModelInference:
+    """模型推理类"""
+    
+    def __init__(
+        self,
+        checkpoint_path: str,
+        tokenizer_name: str,
+        config_path: Optional[str] = None,
+        device: str = 'cuda' if torch.cuda.is_available() else 'cpu'
+    ):
+        self.device = torch.device(device)
+        print(f"Using device: {self.device}")
+        
+        # 1. 加载 Tokenizer (与预训练一致)
+        print(f"Loading tokenizer: {tokenizer_name}...")
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                tokenizer_name, 
+                use_fast=True, 
+                trust_remote_code=True
+            )
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+                self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        except Exception as e:
+            print(f"Error loading tokenizer: {e}")
+            raise e
+        
+        # 2. 配置模型参数 (必须与 pretrain.py 中的配置完全一致)
+        if config_path and Path(config_path).exists():
+            with open(config_path, 'r') as f:
+                self.config = json.load(f)
+        else:
+            # [CRITICAL] 这里使用了你在 pretrain.py 中使用的参数
+            self.config = {
+                'model_dim': 1536,           # 预训练设置
+                'vocab_size': len(self.tokenizer), # 自动适配 Qwen (约 151665)
+                'n_layers': 12,             # 预训练设置
+                'n_heads': 12,              # 预训练设置
+                'n_kv_heads': 4,            # 预训练设置
+                'head_dim': None,           # 自动计算
+                'max_seq_len': 512,        # 预训练设置
+                'dropout': 0.0,             # 推理时关闭 dropout
+                'use_moe': False,           # 预训练设置
+                'use_adapter': False,       # 预训练未开启 Adapter
+                'use_lora': False,          # 预训练未开启 LoRA
+                'rope_scaling_type': "yarn" # 预训练设置
+            }
+        
+        # 3. 初始化模型结构
+        print("Initializing model architecture...")
+        try:
+            self.model = MultiModalDenseTransformer(**self.config)
+            self.preprocessor = UnifiedMultiModalPreprocessor(
+                model_dim=self.config['model_dim']
+            )
+            
+            # 4. 加载权重
+            print(f"Loading checkpoint from {checkpoint_path}...")
+            # weights_only=False 是为了支持加载完整的 checkpoint 字典
+            checkpoint = torch.load(
+                checkpoint_path, 
+                map_location=self.device,
+                weights_only=False 
+            )
+            
+            # 提取 state_dict
+            if 'model_state_dict' in checkpoint:
+                print("Found 'model_state_dict' in checkpoint.")
+                state_dict = checkpoint['model_state_dict']
+            else:
+                state_dict = checkpoint
+            
+            # 处理可能的键名不匹配 (如 DDP 训练产生的 'module.' 前缀)
+            new_state_dict = {}
+            for k, v in state_dict.items():
+                if k.startswith('module.'):
+                    new_state_dict[k[7:]] = v
+                else:
+                    new_state_dict[k] = v
+            
+            # 加载权重 (strict=False 允许忽略一些非关键的不匹配，如 loss 缓存等)
+            missing, unexpected = self.model.load_state_dict(new_state_dict, strict=False)
+            if missing:
+                print(f"Warning: Missing keys: {len(missing)}")
+            if unexpected:
+                print(f"Warning: Unexpected keys: {len(unexpected)}")
+            
+            self.model.to(self.device)
+            self.preprocessor.to(self.device)
+            self.model.eval()
+            
+            print("Model loaded successfully!")
+            print(f"Total parameters: {sum(p.numel() for p in self.model.parameters()) / 1e6:.2f}M")
+            
+        except Exception as e:
+            print(f"Error initializing model: {e}")
+            raise e
+    
+    @torch.no_grad()
+    def generate_text(
+        self,
+        prompt: str,
+        max_new_tokens: int = 128,
+        temperature: float = 0.7,
+        top_k: int = 40,
+        top_p: float = 0.9,
+        repetition_penalty: float = 1.1,
+        image: Optional[Image.Image] = None
+    ) -> str:
+        """生成文本"""
+        
+        # 编码输入
+        inputs = self.tokenizer(prompt, return_tensors="pt")
+        input_ids = inputs['input_ids'].to(self.device)
+        
+        # 构建 MultiModalDenseTransformer 需要的输入格式
+        input_data = {'segments': []}
+        
+        # 处理图像
+        if image is not None:
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+            # 简单的图像处理
+            image_tensor = image_transform(image).unsqueeze(0).to(self.device)
+            # 这里假设预处理器能处理这种输入
+            try:
+                # process_batch 接受 (batch_data, modality_type) 并返回 segments 列表
+                mod_segments = self.preprocessor.process_batch(image_tensor, 'image')
+                # 将返回的 segment 列表合并到 input_data
+                for seg in mod_segments:
+                    input_data['segments'].append(seg)
+            except Exception as e:
+                print(f"Warning: Image processing skipped due to error: {e}")
+        
+        # 添加文本段
+        input_data['segments'].append({
+            'type': 'text',
+            'data': input_ids,
+            'modality_id': 0
+        })
+        
+        # 生成
+        try:
+            # 使用模型自带的 generate 方法
+            generated_ids = self.model.generate(
+                input_data,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                do_sample=True,
+                eos_token_id=self.tokenizer.eos_token_id,
+                pad_token_id=self.tokenizer.pad_token_id
+            )
+            
+            # 解码
+            # 注意：生成的 ids 可能包含原始输入，或者只包含新生成的 token
+            # MultiModalDenseTransformer.generate 通常返回完整的序列
+            generated_text = self.tokenizer.decode(
+                generated_ids[0],
+                skip_special_tokens=True
+            )
+            
+            # 如果包含 prompt，可以选择移除它只显示新内容
+            # if generated_text.startswith(prompt):
+            #     generated_text = generated_text[len(prompt):]
+            
+            return generated_text
+            
+        except Exception as e:
+            print(f"Generation error: {e}")
+            import traceback
+            traceback.print_exc()
+            return f"Error: {str(e)}"
+
+
+# 全局模型实例
+model_instance = None
+app = Flask(__name__)
+app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
+
+@app.route('/')
+def index():
+    display_config = model_instance.config.copy() if model_instance else {}
+    return render_template('index.html', config=display_config)
+
+@app.route('/generate', methods=['POST'])
+def generate():
+    try:
+        data = request.json
+        prompt = data.get('prompt', '')
+        if not prompt.strip():
+            return jsonify({'error': '请输入提示文本'}), 400
+        
+        max_tokens = int(data.get('max_tokens', 100))
+        temperature = float(data.get('temperature', 0.7))
+        top_k = int(data.get('top_k', 40))
+        top_p = float(data.get('top_p', 0.9))
+        repetition_penalty = float(data.get('repetition_penalty', 1.1))
+        
+        image = None
+        if 'image' in data and data['image']:
+            try:
+                image_data = base64.b64decode(data['image'].split(',')[1])
+                image = Image.open(io.BytesIO(image_data))
+            except Exception as e:
+                print(f"Image load error: {e}")
+        
+        output = model_instance.generate_text(
+            prompt, max_tokens, temperature, top_k, top_p, repetition_penalty, image
+        )
+        return jsonify({'output': output})
+    
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+def create_html_template():
+    """写入HTML模板"""
+    html_content = '''
+<!DOCTYPE html>
+<html lang="zh-CN">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Model Inference</title>
+    <style>
+        body { font-family: sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; background: #f0f2f5; }
+        .container { background: white; padding: 30px; border-radius: 12px; box-shadow: 0 4px 6px rgba(0,0,0,0.1); }
+        h1 { color: #1a73e8; text-align: center; }
+        textarea { width: 100%; padding: 10px; border: 1px solid #ddd; border-radius: 8px; margin: 10px 0; min-height: 100px; }
+        .controls { display: grid; grid-template-columns: 1fr 1fr; gap: 15px; margin: 20px 0; background: #f8f9fa; padding: 15px; border-radius: 8px; }
+        button { background: #1a73e8; color: white; border: none; padding: 12px 24px; border-radius: 6px; cursor: pointer; width: 100%; font-size: 16px; transition: background 0.3s; }
+        button:hover { background: #1557b0; }
+        button:disabled { background: #ccc; }
+        #output { margin-top: 20px; padding: 20px; background: #f8f9fa; border-radius: 8px; white-space: pre-wrap; min-height: 100px; border: 1px solid #e0e0e0; }
+        .loading { color: #666; font-style: italic; }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <h1>🚀 模型在线推理</h1>
+        
+        <div>
+            <label><strong>提示词 (Prompt):</strong></label>
+            <textarea id="prompt" placeholder="请输入你的问题..."></textarea>
+        </div>
+
+        <div class="controls">
+            <div>
+                <label>Max Tokens: <span id="maxTokensVal">128</span></label>
+                <input type="range" id="maxTokens" min="32" max="1024" value="128" style="width:100%" oninput="document.getElementById('maxTokensVal').innerText=this.value">
+            </div>
+            <div>
+                <label>Temperature: <span id="tempVal">0.7</span></label>
+                <input type="range" id="temperature" min="0.1" max="1.5" step="0.1" value="0.7" style="width:100%" oninput="document.getElementById('tempVal').innerText=this.value">
+            </div>
+        </div>
+
+        <button id="btn" onclick="generate()">生成 (Generate)</button>
+        
+        <div id="output">结果将显示在这里...</div>
+    </div>
+
+    <script>
+        async function generate() {
+            const prompt = document.getElementById('prompt').value;
+            if(!prompt) return alert("请输入内容");
+            
+            const btn = document.getElementById('btn');
+            const out = document.getElementById('output');
+            
+            btn.disabled = true;
+            btn.innerText = "生成中...";
+            out.innerHTML = '<div class="loading">正在思考中...</div>';
+            
+            try {
+                const res = await fetch('/generate', {
+                    method: 'POST',
+                    headers: {'Content-Type': 'application/json'},
+                    body: JSON.stringify({
+                        prompt: prompt,
+                        max_tokens: parseInt(document.getElementById('maxTokens').value),
+                        temperature: parseFloat(document.getElementById('temperature').value)
+                    })
+                });
+                const data = await res.json();
+                if(data.error) out.innerText = "Error: " + data.error;
+                else out.innerText = data.output;
+            } catch(e) {
+                out.innerText = "请求失败: " + e;
+            } finally {
+                btn.disabled = false;
+                btn.innerText = "生成 (Generate)";
+            }
+        }
+    </script>
+</body>
+</html>
+    '''
+    
+    Path('templates').mkdir(exist_ok=True)
+    with open('templates/index.html', 'w', encoding='utf-8') as f:
+        f.write(html_content)
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    # 默认指向 pretrain 保存的 checkpoint 路径
+    parser.add_argument("--checkpoint", type=str, default="/root/multimodal/checkpoints/pretrain_fixed/step_10000.pt") 
+    parser.add_argument("--tokenizer", type=str, default="Qwen/Qwen2.5-7B-Instruct")
+    parser.add_argument("--port", type=int, default=5001)
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    args = parser.parse_args()
+    
+    if not Path(args.checkpoint).exists():
+        # 尝试找最近的 step checkpoint
+        steps = list(Path("checkpoints/pretrain").glob("step_*.pt"))
+        if steps:
+            print(f"未找到 final_model.pt，尝试使用最新的 checkpoint: {steps[-1]}")
+            args.checkpoint = str(steps[-1])
+        else:
+            print(f"错误: 找不到检查点文件: {args.checkpoint}")
+            return
+
+    create_html_template()
+    
+    global model_instance
+    model_instance = ModelInference(args.checkpoint, args.tokenizer)
+    
+    print(f"\n服务已启动: http://{args.host}:{args.port}")
+    app.run(host=args.host, port=args.port,
+    debug=True,  # 开启调试模式
+    use_reloader=False)
+
+if __name__ == "__main__":
+    main()

infer_sft.py

@@ -0,0 +1,407 @@
+"""
+Flask推理界面 - 多模态Dense Transformer (适配 Qwen Tokenizer 版)
+"""
+
+import os
+import torch
+import torch.nn.functional as F
+from flask import Flask, render_template, request, jsonify
+from transformers import AutoTokenizer
+from PIL import Image
+import json
+import io
+import base64
+from pathlib import Path
+from typing import Optional
+
+# 确保引入路径正确，根据你之前的文件结构
+from model import MultiModalDenseTransformer
+# 注意：UnifiedMultiModalPreprocessor 之前是在 continual_learning.py 中定义的
+# 如果你移动了它，请修改这里的导入路径
+from continual_learning import UnifiedMultiModalPreprocessor
+# 如果没有 image_transform，我们需要在这里定义或导入
+from torchvision import transforms
+
+# 设置国内镜像
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+
+# 定义图像预处理 (与 training 保持一致)
+image_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+class ModelInference:
+    """模型推理类"""
+    
+    def __init__(
+        self,
+        checkpoint_path: str,
+        tokenizer_name: str,
+        config_path: Optional[str] = None,
+        device: str = 'cuda' if torch.cuda.is_available() else 'cpu'
+    ):
+        self.device = torch.device(device)
+        print(f"Using device: {self.device}")
+        
+        # 1. 加载 Tokenizer (与预训练一致)
+        print(f"Loading tokenizer: {tokenizer_name}...")
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                tokenizer_name, 
+                use_fast=True, 
+                trust_remote_code=True
+            )
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+                self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        except Exception as e:
+            print(f"Error loading tokenizer: {e}")
+            raise e
+        
+        # 2. 配置模型参数 (必须与 pretrain.py 中的配置完全一致)
+        if config_path and Path(config_path).exists():
+            with open(config_path, 'r') as f:
+                self.config = json.load(f)
+        else:
+            # [CRITICAL] 这里使用了你在 pretrain.py 中使用的参数
+            self.config = {
+                'model_dim': 1536,           # 预训练设置
+                'vocab_size': len(self.tokenizer), # 自动适配 Qwen (约 151665)
+                'n_layers': 12,             # 预训练设置
+                'n_heads': 12,              # 预训练设置
+                'n_kv_heads': 4,            # 预训练设置
+                'head_dim': None,           # 自动计算
+                'max_seq_len': 512,        # 预训练设置
+                'dropout': 0.0,             # 推理时关闭 dropout
+                'use_moe': False,           # 预训练设置
+                'use_adapter': False,       # 预训练未开启 Adapter
+                'use_lora': False,          # 预训练未开启 LoRA
+                'rope_scaling_type': "yarn", # 预训练设置
+                'use_multimodal_fusion': False, 
+                'use_contrastive': False
+            }
+        
+        # 3. 初始化模型结构
+        print("Initializing model architecture...")
+        try:
+            self.model = MultiModalDenseTransformer(**self.config)
+            self.preprocessor = UnifiedMultiModalPreprocessor(
+                model_dim=self.config['model_dim']
+            )
+            
+            # 4. 加载权重
+            print(f"Loading checkpoint from {checkpoint_path}...")
+            # weights_only=False 是为了支持加载完整的 checkpoint 字典
+            checkpoint = torch.load(
+                checkpoint_path, 
+                map_location=self.device,
+                weights_only=False 
+            )
+            
+            # 提取 state_dict
+            if 'model_state_dict' in checkpoint:
+                print("Found 'model_state_dict' in checkpoint.")
+                state_dict = checkpoint['model_state_dict']
+            else:
+                state_dict = checkpoint
+            
+            # 处理可能的键名不匹配 (如 DDP 训练产生的 'module.' 前缀)
+            new_state_dict = {}
+            for k, v in state_dict.items():
+                if k.startswith('module.'):
+                    new_state_dict[k[7:]] = v
+                else:
+                    new_state_dict[k] = v
+            
+            # 加载权重 (strict=False 允许忽略一些非关键的不匹配，如 loss 缓存等)
+            missing, unexpected = self.model.load_state_dict(new_state_dict, strict=False)
+            if missing:
+                print(f"Warning: Missing keys: {len(missing)}")
+            if unexpected:
+                print(f"Warning: Unexpected keys: {len(unexpected)}")
+            
+            self.model.to(self.device)
+            self.preprocessor.to(self.device)
+            self.model.eval()
+            
+            print("Model loaded successfully!")
+            print(f"Total parameters: {sum(p.numel() for p in self.model.parameters()) / 1e6:.2f}M")
+            
+        except Exception as e:
+            print(f"Error initializing model: {e}")
+            raise e
+    
+    @torch.no_grad()
+    def generate_text(
+        self,
+        prompt: str,
+        max_new_tokens: int = 128,
+        temperature: float = 0.7,
+        top_k: int = 10,
+        top_p: float = 0.9,
+        repetition_penalty: float = 1.2,
+        image: Optional[Image.Image] = None
+    ) -> str:
+        """生成文本"""
+        formatted_prompt = f"Instruction: {prompt}\nResponse:"
+        # 编码输入
+        inputs = self.tokenizer(formatted_prompt, return_tensors="pt")
+        
+        # 编码输入
+        #inputs = self.tokenizer(prompt, return_tensors="pt")
+        input_ids = inputs['input_ids'].to(self.device)
+        
+        # 构建 MultiModalDenseTransformer 需要的输入格式
+        input_data = {'segments': []}
+        
+        # 处理图像
+        if image is not None:
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+            # 简单的图像处理
+            image_tensor = image_transform(image).unsqueeze(0).to(self.device)
+            # 这里假设预处理器能处理这种输入
+            try:
+                # process_batch 接受 (batch_data, modality_type) 并返回 segments 列表
+                mod_segments = self.preprocessor.process_batch(image_tensor, 'image')
+                # 将返回的 segment 列表合并到 input_data
+                for seg in mod_segments:
+                    input_data['segments'].append(seg)
+            except Exception as e:
+                print(f"Warning: Image processing skipped due to error: {e}")
+        
+        # 添加文本段
+        input_data['segments'].append({
+            'type': 'text',
+            'data': input_ids,
+            'modality_id': 0
+        })
+        
+        # 生成
+        try:
+            # 使用模型自带的 generate 方法
+            generated_ids = self.model.generate(
+                input_data,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                repetition_penalty=repetition_penalty,
+                do_sample=True,
+                eos_token_id=self.tokenizer.eos_token_id,
+                pad_token_id=self.tokenizer.pad_token_id
+            )
+            
+            # 3. 解码
+            full_output = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+            print(f"\n====== [DEBUG 原始输出] ======\n{full_output}\n==============================\n")
+            # 4. [关键修改] 截断逻辑 (Stop Logic)
+            # 提取 Response 之后的部分
+            if "Response:" in full_output:
+                answer = full_output.split("Response:")[-1].strip()
+            else:
+                answer = full_output
+
+            # 定义停止词列表 (根据你的图，模型喜欢生成 Instructions: 或 Ingredients:)
+            stop_words = [
+                "Instruction", "Input", "###", "Response", 
+                "User:", "Assistant:", "\n\n" # 双换行通常意味着一段结束
+            ]
+            
+            for stop_word in stop_words:
+                if stop_word in answer:
+                    answer = answer.split(stop_word)[0].strip()
+           
+            # 3. [新增] 强制去除首行重复 (解决 Echo 问题)
+            # 如果模型第一句就是重复 Prompt，去掉它
+            lines = answer.split('\n')
+            if len(lines) > 0 and prompt.lower() in lines[0].lower():
+                answer = "\n".join(lines[1:]).strip()
+            return answer
+            
+        except Exception as e:
+            print(f"Generation error: {e}")
+            import traceback
+            traceback.print_exc()
+            return f"Error: {str(e)}"
+
+
+# 全局模型实例
+model_instance = None
+app = Flask(__name__)
+app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024
+
+@app.route('/')
+def index():
+    display_config = model_instance.config.copy() if model_instance else {}
+    return render_template('index.html', config=display_config)
+
+@app.route('/generate', methods=['POST'])
+def generate():
+    try:
+        data = request.json
+        prompt = data.get('prompt', '')
+        if not prompt.strip():
+            return jsonify({'error': '请输入提示文本'}), 400
+        
+        max_tokens = int(data.get('max_tokens', 100))
+        temperature = float(data.get('temperature', 0.7))
+        top_k = int(data.get('top_k', 40))
+        top_p = float(data.get('top_p', 0.9))
+        repetition_penalty = float(data.get('repetition_penalty', 1.1))
+        
+        image = None
+        if 'image' in data and data['image']:
+            try:
+                image_data = base64.b64decode(data['image'].split(',')[1])
+                image = Image.open(io.BytesIO(image_data))
+            except Exception as e:
+                print(f"Image load error: {e}")
+        
+        output = model_instance.generate_text(
+            prompt, max_tokens, temperature, top_k, top_p, repetition_penalty, image
+        )
+        return jsonify({'output': output})
+    
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+def create_html_template():
+    """写入HTML模板"""
+    html_content = '''
+<!DOCTYPE html>
+<html lang="zh-CN">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Model Inference</title>
+    <style>
+        body { font-family: sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; background: #f0f2f5; }
+        .container { background: white; padding: 30px; border-radius: 12px; box-shadow: 0 4px 6px rgba(0,0,0,0.1); }
+        h1 { color: #1a73e8; text-align: center; }
+        textarea { width: 100%; padding: 10px; border: 1px solid #ddd; border-radius: 8px; margin: 10px 0; min-height: 100px; }
+        .controls { display: grid; grid-template-columns: 1fr 1fr; gap: 15px; margin: 20px 0; background: #f8f9fa; padding: 15px; border-radius: 8px; }
+        button { background: #1a73e8; color: white; border: none; padding: 12px 24px; border-radius: 6px; cursor: pointer; width: 100%; font-size: 16px; transition: background 0.3s; }
+        button:hover { background: #1557b0; }
+        button:disabled { background: #ccc; }
+        #output { margin-top: 20px; padding: 20px; background: #f8f9fa; border-radius: 8px; white-space: pre-wrap; min-height: 100px; border: 1px solid #e0e0e0; }
+        .loading { color: #666; font-style: italic; }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <h1>🚀 模型在线推理</h1>
+        
+        <div>
+            <label><strong>提示词 (Prompt):</strong></label>
+            <textarea id="prompt" placeholder="请输入你的问题..."></textarea>
+        </div>
+
+        <div class="controls">
+            <div>
+                <label>Max Tokens: <span id="maxTokensVal">128</span></label>
+                <input type="range" id="maxTokens" min="32" max="1024" value="128" style="width:100%" oninput="document.getElementById('maxTokensVal').innerText=this.value">
+            </div>
+            <div>
+                <label>Temperature: <span id="tempVal">0.7</span></label>
+                <input type="range" id="temperature" min="0.1" max="1.5" step="0.1" value="0.7" style="width:100%" oninput="document.getElementById('tempVal').innerText=this.value">
+            </div>
+        </div>
+
+        <button id="btn" onclick="generate()">生成 (Generate)</button>
+        
+        <div id="output">结果将显示在这里...</div>
+    </div>
+
+    <script>
+        async function generate() {
+            const prompt = document.getElementById('prompt').value;
+            if(!prompt) return alert("请输入内容");
+            
+            const btn = document.getElementById('btn');
+            const out = document.getElementById('output');
+            
+            btn.disabled = true;
+            btn.innerText = "生成中...";
+            out.innerHTML = '<div class="loading">正在思考中...</div>';
+            
+            try {
+                const res = await fetch('/generate', {
+                    method: 'POST',
+                    headers: {'Content-Type': 'application/json'},
+                    body: JSON.stringify({
+                        prompt: prompt,
+                        max_tokens: parseInt(document.getElementById('maxTokens').value),
+                        temperature: parseFloat(document.getElementById('temperature').value)
+                    })
+                });
+                const data = await res.json();
+                if(data.error) out.innerText = "Error: " + data.error;
+                else out.innerText = data.output;
+            } catch(e) {
+                out.innerText = "请求失败: " + e;
+            } finally {
+                btn.disabled = false;
+                btn.innerText = "生成 (Generate)";
+            }
+        }
+    </script>
+</body>
+</html>
+    '''
+    
+    Path('templates').mkdir(exist_ok=True)
+    with open('templates/index.html', 'w', encoding='utf-8') as f:
+        f.write(html_content)
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    # 默认指向 pretrain 保存的 checkpoint 路径
+    parser.add_argument("--checkpoint", type=str, default="/root/multimodal/checkpoints/posttrain/final_model.pt") 
+    parser.add_argument("--tokenizer", type=str, default="Qwen/Qwen2.5-7B-Instruct")
+    parser.add_argument("--port", type=int, default=5001)
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    args = parser.parse_args()
+    
+    if not Path(args.checkpoint).exists():
+        # 尝试找最近的 step checkpoint
+        steps = list(Path("checkpoints/pretrain").glob("step_*.pt"))
+        if steps:
+            print(f"未找到 final_model.pt，尝试使用最新的 checkpoint: {steps[-1]}")
+            args.checkpoint = str(steps[-1])
+        else:
+            print(f"错误: 找不到检查点文件: {args.checkpoint}")
+            return
+    # ----------------- 新增部分开始 -----------------
+    try:
+        from pyngrok import ngrok, conf
+        
+        # 如果你在国内，ngrok 连接慢，可以配置 region='ap' (亚太) 或 'au' (澳洲)
+        # conf.get_default().region = "ap" 
+        
+        # 建立隧道，映射 5001 端口
+        public_url = ngrok.connect(args.port).public_url
+        print(f"\n========================================")
+        print(f"🎉 公网访问地址 (发给朋友): {public_url}")
+        print(f"========================================\n")
+    except ImportError:
+        print("未安装 pyngrok，无法自动生成公网链接。")
+        print("提示: pip install pyngrok")
+    except Exception as e:
+        print(f"Ngrok 启动失败: {e}")
+    # ----------------- 新增部分结束 -----------------
+    create_html_template()
+    
+    global model_instance
+    model_instance = ModelInference(args.checkpoint, args.tokenizer)
+    
+    print(f"\n服务已启动: http://{args.host}:{args.port}")
+    app.run(host=args.host, port=args.port,
+    debug=True,  # 开启调试模式
+    use_reloader=False)
+
+if __name__ == "__main__":
+    main()

model.py

@@ -0,0 +1,505 @@
+"""
+改进的多模态Dense Transformer主模型
+整合所有SOTA改进
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import List, Dict, Optional, Tuple
+import math
+from components import RMSNorm
+from transformer import OptimizedTransformerBlock
+from multimodel_fusion import MultiModalFusionModule
+from encoders import (
+    ImprovedVisionTransformer,
+    ImprovedAudioEncoder,
+    ImprovedVideoEncoder
+)
+
+class MultiModalDenseTransformer(nn.Module):
+    """
+    改进的统一多模态Dense Transformer
+    主要改进:
+    1. 深度跨模态融合
+    2. 模态特定的优化编码器
+    3. 对比学习对齐
+    4. 改进的位置编码和注意力机制
+    5. 更好的训练稳定性
+    """
+    def __init__(
+        self,
+        model_dim: int = 2048,
+        vocab_size: int = 30000,
+        n_layers: int = 48,
+        n_heads: int = 32,
+        n_kv_heads: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        max_seq_len: int = 8192,
+        dropout: float = 0.0,
+        attn_dropout: float = 0.0,
+        
+        # MoE配置
+        use_moe: bool = False,
+        num_experts: int = 8,
+        moe_top_k: int = 2,
+        moe_layers: Optional[List[int]] = None,
+        
+        # PEFT配置
+        use_adapter: bool = False,
+        adapter_dim: int = 64,
+        use_lora: bool = False,
+        lora_rank: int = 8,
+        
+        # 训练配置
+        use_gradient_checkpointing: bool = False,
+        use_parallel_residual: bool = False,
+        
+        # 位置编码
+        rope_scaling_factor: float = 1.0,
+        rope_scaling_type: str = "yarn",
+        sliding_window: Optional[int] = None,
+        
+        # 规范化
+        norm_eps: float = 1e-6,
+        initializer_range: float = 0.02,
+        ffn_dim_multiplier: Optional[float] = None,
+        tie_word_embeddings: bool = True,
+        
+        # 多模态配置
+        use_multimodal_fusion: bool = True,
+        fusion_layers: int = 4,
+        use_contrastive: bool = True,
+        vision_depth: int = 24,
+        audio_depth: int = 12,
+        video_spatial_depth: int = 12,
+        video_temporal_depth: int = 4
+    ):
+        super().__init__()
+        
+        self.model_dim = model_dim
+        self.vocab_size = vocab_size
+        self.n_layers = n_layers
+        self.max_seq_len = max_seq_len
+        self.use_gradient_checkpointing = use_gradient_checkpointing
+        self.tie_word_embeddings = tie_word_embeddings
+        self.use_multimodal_fusion = use_multimodal_fusion
+        
+        # Token embedding
+        self.token_embedding = nn.Embedding(vocab_size, model_dim)
+        self.modality_embedding = nn.Embedding(4, model_dim)
+        self.embed_dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        
+        # 改进的模态编码器
+        self.vision_encoder = ImprovedVisionTransformer(
+            embed_dim=model_dim,
+            depth=vision_depth,
+            n_heads=n_heads,
+            dropout=dropout,
+            use_adapter=use_adapter,
+            adapter_dim=adapter_dim
+        )
+        
+        self.audio_encoder = ImprovedAudioEncoder(
+            embed_dim=model_dim,
+            depth=audio_depth,
+            n_heads=n_heads,
+            dropout=dropout,
+            use_adapter=use_adapter,
+            adapter_dim=adapter_dim
+        )
+        
+        self.video_encoder = ImprovedVideoEncoder(
+            embed_dim=model_dim,
+            spatial_depth=video_spatial_depth,
+            temporal_depth=video_temporal_depth,
+            n_heads=n_heads,
+            dropout=dropout,
+            use_adapter=use_adapter,
+            adapter_dim=adapter_dim
+        )
+        
+        # 多模态融合模块
+        if use_multimodal_fusion:
+            self.fusion_module = MultiModalFusionModule(
+                dim=model_dim,
+                num_fusion_layers=fusion_layers,
+                n_heads=n_heads,
+                dropout=dropout,
+                use_contrastive=use_contrastive
+            )
+        
+        # Transformer layers
+        if moe_layers is None and use_moe:
+            moe_layers = list(range(n_layers // 2, n_layers))
+        elif moe_layers is None:
+            moe_layers = []
+        
+        self.layers = nn.ModuleList([
+            OptimizedTransformerBlock(
+                dim=model_dim,
+                n_heads=n_heads,
+                n_kv_heads=n_kv_heads,
+                head_dim=head_dim,
+                dropout=dropout,
+                attn_dropout=attn_dropout,
+                use_moe=(use_moe and i in moe_layers),
+                num_experts=num_experts,
+                moe_top_k=moe_top_k,
+                use_adapter=use_adapter,
+                adapter_dim=adapter_dim,
+                use_lora=use_lora,
+                lora_rank=lora_rank,
+                use_parallel_residual=use_parallel_residual,
+                norm_eps=norm_eps,
+                sliding_window=sliding_window,
+                ffn_dim_multiplier=ffn_dim_multiplier,
+                layer_idx=i
+            )
+            for i in range(n_layers)
+        ])
+        
+        self.norm = RMSNorm(model_dim, eps=norm_eps)
+        self.lm_head = nn.Linear(model_dim, vocab_size, bias=False)
+        
+        if tie_word_embeddings:
+            self.lm_head.weight = self.token_embedding.weight
+        
+        self.initializer_range = initializer_range
+        self.apply(self._init_weights)
+        
+        if not tie_word_embeddings:
+            self._init_lm_head()
+        
+        self.n_params = sum(p.numel() for p in self.parameters())
+        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        
+        print(f"\n{'='*80}")
+        print(f"Improved Model Configuration:")
+        print(f"  Model Dimension: {model_dim}")
+        print(f"  Vocab Size: {vocab_size}")
+        print(f"  Layers: {n_layers}")
+        print(f"  Attention Heads: {n_heads}")
+        print(f"  KV Heads: {n_kv_heads if n_kv_heads else n_heads}")
+        print(f"  Max Sequence Length: {max_seq_len}")
+        print(f"  Multimodal Fusion: {use_multimodal_fusion}")
+        print(f"  Contrastive Learning: {use_contrastive}")
+        print(f"  MoE: {use_moe} (Experts: {num_experts}, Top-K: {moe_top_k})")
+        print(f"  Total Parameters: {self.n_params / 1e9:.2f}B")
+        print(f"  Trainable Parameters: {trainable_params / 1e9:.2f}B")
+        print(f"{'='*80}\n")
+
+    def _init_weights(self, module):
+        """权重初始化"""
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=self.initializer_range)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=self.initializer_range)
+            if hasattr(module, 'padding_idx') and module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    def _init_lm_head(self):
+        """初始化LM head"""
+        std = self.initializer_range / math.sqrt(2 * self.n_layers)
+        torch.nn.init.normal_(self.lm_head.weight, mean=0.0, std=std)
+
+    def _encode_modality(self, segment: Dict) -> torch.Tensor:
+        """编码单个模态"""
+        seg_type = segment['type']
+        seg_data = segment['data']
+        
+        if seg_type == 'image':
+            return self.vision_encoder(seg_data)
+        elif seg_type == 'audio':
+            return self.audio_encoder(seg_data)
+        elif seg_type == 'video':
+            return self.video_encoder(seg_data)
+        elif seg_type == 'text':
+            return self.token_embedding(seg_data)
+        else:
+            return seg_data
+
+    def forward(
+        self,
+        input_data: Dict,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        return_hidden: bool = False,
+        use_cache: bool = False,
+        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        compute_contrastive: bool = False
+    ) -> Dict:
+        """前向传播"""
+        device = self.token_embedding.weight.device
+        
+        # 编码每个模态
+        encoded_segments = []
+        for segment in input_data.get('segments', []):
+            encoded = self._encode_modality(segment)
+            
+            # 添加模态嵌入
+            modality_id = segment.get('modality_id', 0)
+            modality_embeds = self.modality_embedding(
+                torch.tensor([modality_id], device=device)
+            ).expand(encoded.shape[0], encoded.shape[1], -1)
+            
+            encoded_segments.append({
+                'type': segment['type'],
+                'data': encoded + modality_embeds,
+                'modality_id': modality_id
+            })
+        
+        # 多模态融合
+        contrastive_losses = {}
+        if self.use_multimodal_fusion and len(encoded_segments) > 1:
+            fusion_output = self.fusion_module(
+                encoded_segments,
+                compute_contrastive=compute_contrastive
+            )
+            x = fusion_output['fused_features']
+            contrastive_losses = fusion_output.get('contrastive_losses', {})
+        else:
+            # 简单拼接
+            all_embeddings = [seg['data'] for seg in encoded_segments]
+            x = torch.cat(all_embeddings, dim=1) if all_embeddings else torch.zeros(
+                1, 1, self.model_dim, device=device
+            )
+        
+        x = self.embed_dropout(x)
+        # 如果��有传入 position_ids，我们需要根据历史长度生成它
+        if position_ids is None:
+            if past_key_values is not None:
+                # 缓存的长度 (KV cache 的 shape 是 [B, H, SeqLen, D])
+                past_length = past_key_values[0][0].size(2)
+                # 当前输入的长度
+                seq_length = x.shape[1]
+                # 生成正确的位置索引: [past_length, past_length + 1, ...]
+                position_ids = torch.arange(
+                    past_length, past_length + seq_length, dtype=torch.long, device=device
+                ).unsqueeze(0).expand(x.shape[0], -1)
+            else:
+                # 如果没有缓存，从 0 开始
+                seq_length = x.shape[1]
+                position_ids = torch.arange(
+                    0, seq_length, dtype=torch.long, device=device
+                ).unsqueeze(0).expand(x.shape[0], -1)
+        # Transformer层
+        present_key_values = [] if use_cache else None
+        all_hidden_states = [] if output_hidden_states else None
+        all_attentions = [] if output_attentions else None
+        moe_aux_loss = torch.tensor(0.0, device=device)
+        
+        for idx, layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states.append(x)
+            
+            past_kv = past_key_values[idx] if past_key_values is not None else None
+            
+            if self.use_gradient_checkpointing and self.training:
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(
+                            inputs[0],
+                            attention_mask=inputs[1],
+                            position_ids=inputs[2],
+                            use_cache=False,
+                            past_kv=None,
+                            output_attentions=False
+                        )
+                    return custom_forward
+                
+                import torch.utils.checkpoint as checkpoint
+                layer_outputs = checkpoint.checkpoint(
+                    create_custom_forward(layer),
+                    x,
+                    attention_mask,
+                    position_ids,
+                    use_reentrant=False
+                )
+                x = layer_outputs[0]
+                present_kv = None
+                attn_weights = None
+            else:
+                layer_outputs = layer(
+                    x,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    use_cache=use_cache,
+                    past_kv=past_kv,
+                    output_attentions=output_attentions
+                )
+                x, present_kv, attn_weights = layer_outputs
+            
+            if use_cache:
+                present_key_values.append(present_kv)
+            
+            if output_attentions:
+                all_attentions.append(attn_weights)
+            
+            if hasattr(layer, 'moe_aux_loss'):
+                moe_aux_loss += layer.moe_aux_loss
+        
+        hidden_states = self.norm(x)
+        logits = self.lm_head(hidden_states)
+        
+        if output_hidden_states:
+            all_hidden_states.append(hidden_states)
+        
+        # 组装输出
+        outputs = {
+            'logits': logits,
+            'moe_aux_loss': moe_aux_loss,
+            'contrastive_losses': contrastive_losses
+        }
+        
+        if use_cache:
+            outputs['past_key_values'] = present_key_values
+        
+        if output_hidden_states:
+            outputs['hidden_states'] = all_hidden_states
+        
+        if output_attentions:
+            outputs['attentions'] = all_attentions
+        
+        if return_hidden:
+            outputs['last_hidden_state'] = hidden_states
+        
+        return outputs
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_data: Dict,
+        max_new_tokens: int = 100,
+        temperature: float = 1.0,
+        top_k: int = 50,
+        top_p: float = 0.9,
+        eos_token_id: int = 2,
+        pad_token_id: Optional[int] = None,
+        use_cache: bool = True,
+        repetition_penalty: float = 1.0,
+        length_penalty: float = 1.0,
+        min_length: int = 0,
+        do_sample: bool = True,
+        num_beams: int = 1
+    ) -> torch.Tensor:
+        """改进的生成方法"""
+        self.eval()
+        device = next(self.parameters()).device
+        
+        if pad_token_id is None:
+            pad_token_id = eos_token_id
+        
+        initial_text_tokens = input_data['segments'][0]['data'].to(device)
+        batch_size = initial_text_tokens.shape[0]
+        
+        if 'attention_mask' in input_data:
+            attention_mask = input_data['attention_mask'].to(device)
+        else:
+            attention_mask = torch.ones_like(initial_text_tokens)
+        initial_seq_len = initial_text_tokens.shape[1]
+        position_ids = torch.zeros((batch_size,initial_seq_len),dtype=torch.long,device=device)
+        
+        for i in range(batch_size):
+            non_pad_mask = attention_mask[i].bool()
+            if non_pad_mask.any():
+                positions = torch.cumsum(non_pad_mask.long(),dim=0) -1
+                position_ids[i]=positions * non_pad_mask.long()
+            
+
+
+        generated_tokens = []
+        past_key_values = None
+        current_tokens = initial_text_tokens
+        unfinished_sequences = torch.ones(batch_size, dtype=torch.long, device=device)
+        
+        for step in range(max_new_tokens):
+            current_input_data = {
+                'segments': [{'type': 'text', 'data': current_tokens, 'modality_id': 0}]
+            }
+            
+            if step > 0 and use_cache:
+                # 添加当前 token 的 mask (1)
+                new_mask = torch.ones(batch_size,1,dtype=torch.long,device=device)
+                attention_mask = torch.cat([attention_mask, new_mask], dim=1)
+                current_positions = (attention_mask.sum(dim=1 , keepdim=True) -1).clamp(min=0)
+                current_positions_ids=current_positions
+            else:
+                current_positions_ids=position_ids
+            outputs = self.forward(
+                current_input_data,
+                attention_mask=attention_mask, # <--- 传入 Mask
+                position_ids=current_positions_ids,
+                use_cache=use_cache,
+                past_key_values=past_key_values
+            )
+            
+            logits = outputs['logits']
+            if use_cache:
+                past_key_values = outputs['past_key_values']
+            
+            next_token_logits = logits[:, -1, :] / max(temperature, 1e-5)
+            
+            # Repetition penalty
+            if repetition_penalty != 1.0 and len(generated_tokens) > 0:
+                prev_generated = torch.cat(generated_tokens, dim=1)
+                score = torch.gather(next_token_logits, 1, prev_generated)
+                score = torch.where(
+                    score < 0,
+                    score * repetition_penalty,
+                    score / repetition_penalty
+                )
+                next_token_logits.scatter_(1, prev_generated, score)
+            
+            # Min length constraint
+            if step < min_length:
+                next_token_logits[:, eos_token_id] = float('-inf')
+            
+            # Sampling
+            if do_sample:
+                if top_k > 0:
+                    top_k_vals, _ = torch.topk(next_token_logits, top_k)
+                    min_val_to_keep = top_k_vals[:, -1].unsqueeze(-1)
+                    next_token_logits[next_token_logits < min_val_to_keep] = float('-inf')
+                
+                if top_p < 1.0:
+                    sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                    sorted_indices_to_remove[..., 0] = 0
+                    indices_to_remove = torch.zeros_like(next_token_logits, dtype=torch.bool)
+                    indices_to_remove.scatter_(1, sorted_indices, sorted_indices_to_remove)
+                    next_token_logits[indices_to_remove] = float('-inf')
+                
+                probs = F.softmax(next_token_logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+            
+            # Apply unfinished mask
+            next_token = next_token * unfinished_sequences[:, None] + pad_token_id * (1 - unfinished_sequences[:, None])
+            
+            generated_tokens.append(next_token)
+            
+            if not use_cache:
+                initial_text_tokens = torch.cat([initial_text_tokens, next_token], dim=1)
+                current_tokens = initial_text_tokens
+            else:
+                current_tokens = next_token
+            
+            # Update unfinished sequences
+            unfinished_sequences = unfinished_sequences.mul(
+                (next_token.squeeze(-1) != eos_token_id).long()
+            )
+            
+            if unfinished_sequences.max() == 0:
+                break
+        
+        if not generated_tokens:
+            return torch.empty(batch_size, 0, dtype=torch.long, device=device)
+        
+        return torch.cat(generated_tokens, dim=1)

moe.py

@@ -0,0 +1,460 @@
+"""
+优化的混合专家系统 (Mixture of Experts)
+基于Mixtral、Switch Transformer、GLaM的最佳实践
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Tuple, Optional, List
+import math
+
+class Expert(nn.Module):
+    """
+    单个专家网络
+    使用SwiGLU激活函数以获得更好的性能
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+        dropout: float = 0.0,
+        bias: bool = False
+    ):
+        super().__init__()
+        self.w1 = nn.Linear(dim, hidden_dim, bias=bias)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=bias)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=bias)
+        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        
+        self._init_weights()
+
+    def _init_weights(self):
+        """改进的权重初始化"""
+        for module in [self.w1, self.w2, self.w3]:
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        前向传播
+        SwiGLU: (Swish(W1·x) ⊙ W3·x) W2
+        """
+        return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
+
+class TopKRouter(nn.Module):
+    """
+    Top-K路由器 - 改进版
+    改进点:
+    1. 专家容量管理
+    2. 负载均衡
+    3. 训练时的噪声注入
+    4. Z-loss防止logits爆炸
+
+    参考:
+    - Switch Transformer
+    - Mixtral 8x7B
+    - ST-MoE
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_experts: int,
+        top_k: int = 2,
+        capacity_factor: float = 1.25,
+        noise_std: float = 1.0,
+        use_expert_capacity: bool = True,
+        router_z_loss_coef: float = 0.001,
+        router_aux_loss_coef: float = 0.01
+    ):
+        super().__init__()
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.capacity_factor = capacity_factor
+        self.noise_std = noise_std
+        self.use_expert_capacity = use_expert_capacity
+        self.router_z_loss_coef = router_z_loss_coef
+        self.router_aux_loss_coef = router_aux_loss_coef
+        
+        self.gate = nn.Linear(dim, num_experts, bias=False)
+        
+        nn.init.normal_(self.gate.weight, mean=0.0, std=0.02)
+
+    def _compute_routing_weights(
+        self,
+        logits: torch.Tensor,
+        use_noise: bool = True
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        计算路由权重
+        
+        Args:
+            logits: 路由logits [batch*seq_len, num_experts]
+            use_noise: 是否添加噪声
+        
+        Returns:
+            top_k_gates: Top-K门控值 [batch*seq_len, top_k]
+            top_k_indices: Top-K专家索引 [batch*seq_len, top_k]
+        """
+        if use_noise and self.training:
+            noise = torch.randn_like(logits) * self.noise_std
+            logits = logits + noise
+        
+        top_k_logits, top_k_indices = torch.topk(logits, self.top_k, dim=-1)
+        
+        top_k_gates = F.softmax(top_k_logits, dim=-1)
+        
+        return top_k_gates, top_k_indices
+
+    def _compute_auxiliary_loss(
+        self,
+        logits: torch.Tensor,
+        top_k_indices: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        计算辅助损失
+        
+        包括:
+        1. 负载均衡损失(确保专家被均匀使用)
+        2. Z-loss(防止logits过大)
+        
+        Args:
+            logits: 路由logits [batch*seq_len, num_experts]
+            top_k_indices: 选中的专家索引 [batch*seq_len, top_k]
+        
+        Returns:
+            load_balance_loss: 负载均衡损失
+            z_loss: Z-loss
+        """
+        num_tokens = logits.shape[0]
+        
+        router_probs = F.softmax(logits, dim=-1)
+        
+        expert_probs = router_probs.mean(dim=0)
+        
+        expert_mask = F.one_hot(top_k_indices, self.num_experts).float()
+        expert_freq = expert_mask.sum(dim=[0, 1]) / (num_tokens * self.top_k)
+        
+        load_balance_loss = self.num_experts * torch.sum(expert_probs * expert_freq)
+        
+        z_loss = torch.mean(logits ** 2)
+        
+        return load_balance_loss, z_loss
+
+    def forward(
+        self,
+        x: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        前向传播
+        
+        Args:
+            x: 输入 [batch*seq_len, dim]
+        
+        Returns:
+            top_k_gates: 门控权重 [batch*seq_len, top_k]
+            top_k_indices: 专家索引 [batch*seq_len, top_k]
+            auxiliary_loss: 辅助损失(标量)
+        """
+        logits = self.gate(x)
+        
+        top_k_gates, top_k_indices = self._compute_routing_weights(
+            logits, use_noise=self.training
+        )
+        
+        if self.training:
+            load_balance_loss, z_loss = self._compute_auxiliary_loss(logits, top_k_indices)
+            auxiliary_loss = (
+                self.router_aux_loss_coef * load_balance_loss +
+                self.router_z_loss_coef * z_loss
+            )
+        else:
+            auxiliary_loss = torch.tensor(0.0, device=x.device)
+        
+        return top_k_gates, top_k_indices, auxiliary_loss
+
+class MixtureOfExperts(nn.Module):
+    """
+    混合专家层 - 优化版
+    改进点:
+    1. 高效的token分发和聚合
+    2. 专家容量管理
+    3. 改进的负载均衡
+    4. 支持专家并行
+
+    参考:
+    - Mixtral 8x7B
+    - Switch Transformer
+    - GShard
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_experts: int = 8,
+        expert_hidden_dim: Optional[int] = None,
+        top_k: int = 2,
+        dropout: float = 0.0,
+        capacity_factor: float = 1.25,
+        use_expert_capacity: bool = True,
+        router_z_loss_coef: float = 0.001,
+        router_aux_loss_coef: float = 0.01,
+        noise_std: float = 1.0,
+        ffn_dim_multiplier: Optional[float] = None
+    ):
+        super().__init__()
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.capacity_factor = capacity_factor
+        self.use_expert_capacity = use_expert_capacity
+        
+        if expert_hidden_dim is None:
+            if ffn_dim_multiplier is not None:
+                expert_hidden_dim = int(dim * ffn_dim_multiplier)
+            else:
+                expert_hidden_dim = int(2 * dim * 4 / 3)
+            expert_hidden_dim = 256 * ((expert_hidden_dim + 255) // 256)
+        
+        self.experts = nn.ModuleList([
+            Expert(dim, expert_hidden_dim, dropout, bias=False)
+            for _ in range(num_experts)
+        ])
+        
+        self.router = TopKRouter(
+            dim=dim,
+            num_experts=num_experts,
+            top_k=top_k,
+            capacity_factor=capacity_factor,
+            noise_std=noise_std,
+            use_expert_capacity=use_expert_capacity,
+            router_z_loss_coef=router_z_loss_coef,
+            router_aux_loss_coef=router_aux_loss_coef
+        )
+
+    def _compute_expert_capacity(self, num_tokens: int) -> int:
+        """计算每个专家的容量"""
+        if not self.use_expert_capacity:
+            return num_tokens
+        
+        capacity = int(
+            (num_tokens / self.num_experts) * self.capacity_factor * self.top_k
+        )
+        return max(capacity, 1)
+
+    def forward(
+        self,
+        x: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        前向传播
+        
+        Args:
+            x: 输入 [batch, seq_len, dim]
+        
+        Returns:
+            output: 输出 [batch, seq_len, dim]
+            auxiliary_loss: 辅助损失
+        """
+        B, T, D = x.shape
+        num_tokens = B * T
+        
+        x_flat = x.view(-1, D)
+        
+        top_k_gates, top_k_indices, auxiliary_loss = self.router(x_flat)
+        
+        output = torch.zeros_like(x_flat)
+        
+        expert_capacity = self._compute_expert_capacity(num_tokens)
+        
+        for expert_idx, expert in enumerate(self.experts):
+            expert_mask = (top_k_indices == expert_idx)
+            
+            token_indices, topk_positions = torch.where(expert_mask)
+            
+            if len(token_indices) == 0:
+                continue
+            
+            if self.use_expert_capacity and len(token_indices) > expert_capacity:
+                perm = torch.randperm(len(token_indices), device=x.device)[:expert_capacity]
+                token_indices = token_indices[perm]
+                topk_positions = topk_positions[perm]
+            
+            expert_input = x_flat[token_indices]
+            expert_gates = top_k_gates[token_indices, topk_positions]
+            
+            expert_output = expert(expert_input)
+            
+            expert_output = expert_output * expert_gates.unsqueeze(-1)
+            
+            output.index_add_(0, token_indices, expert_output)
+        
+        output = output.view(B, T, D)
+        
+        return output, auxiliary_loss
+
+class SparseDispatcher:
+    """
+    稀疏分发器 - 用于高效的MoE计算
+    管理tokens到专家的分配和聚合
+    这是一个可选的辅助类,用于更高效的实现
+    """
+
+    def __init__(
+        self,
+        num_experts: int,
+        gates: torch.Tensor,
+        expert_indices: torch.Tensor
+    ):
+        """
+        Args:
+            num_experts: 专家数量
+            gates: 门控权重 [batch_size, num_experts]
+            expert_indices: 专家索引 [batch_size]
+        """
+        self.num_experts = num_experts
+        self._gates = gates
+        self._expert_indices = expert_indices
+        
+        self._expert_masks = []
+        for i in range(num_experts):
+            self._expert_masks.append((expert_indices == i).nonzero(as_tuple=True)[0])
+
+    def dispatch(self, inp: torch.Tensor) -> List[torch.Tensor]:
+        """
+        将输入分发给各个专家
+        
+        Args:
+            inp: 输入张量 [batch_size, dim]
+        
+        Returns:
+            expert_inputs: 每个专家的输入列表
+        """
+        expert_inputs = []
+        for mask in self._expert_masks:
+            if len(mask) > 0:
+                expert_inputs.append(inp[mask])
+            else:
+                expert_inputs.append(
+                    torch.empty(0, inp.size(-1), device=inp.device, dtype=inp.dtype)
+                )
+        return expert_inputs
+
+    def combine(self, expert_outputs: List[torch.Tensor]) -> torch.Tensor:
+        """
+        组合专家输出
+        
+        Args:
+            expert_outputs: 每个专家的输出列表
+        
+        Returns:
+            output: 组合后的输出 [batch_size, dim]
+        """
+        output_shape = (self._gates.size(0), expert_outputs[0].size(-1))
+        output = torch.zeros(
+            output_shape,
+            device=self._gates.device,
+            dtype=expert_outputs[0].dtype
+        )
+        
+        for expert_idx, expert_out in enumerate(expert_outputs):
+            mask = self._expert_masks[expert_idx]
+            if len(mask) > 0:
+                weighted_output = expert_out * self._gates[mask, expert_idx].unsqueeze(-1)
+                output[mask] += weighted_output
+        
+        return output
+
+    def expert_to_gates(self) -> List[torch.Tensor]:
+        """
+        返回每个专家对应的门控权重
+        
+        Returns:
+            gates_per_expert: 每个专家的门控权重列表
+        """
+        gates_per_expert = []
+        for expert_idx in range(self.num_experts):
+            mask = self._expert_masks[expert_idx]
+            if len(mask) > 0:
+                gates_per_expert.append(self._gates[mask, expert_idx])
+            else:
+                gates_per_expert.append(torch.empty(0, device=self._gates.device))
+        return gates_per_expert
+
+class MoELayer(nn.Module):
+    """
+    MoE层的另一种实现方式
+    使用SparseDispatcher进行更高效的计算
+    """
+    def __init__(
+        self,
+        dim: int,
+        num_experts: int = 8,
+        expert_hidden_dim: Optional[int] = None,
+        top_k: int = 2,
+        dropout: float = 0.0,
+        capacity_factor: float = 1.25
+    ):
+        super().__init__()
+        self.num_experts = num_experts
+        self.top_k = top_k
+        
+        if expert_hidden_dim is None:
+            expert_hidden_dim = int(2 * dim * 4 / 3)
+            expert_hidden_dim = 256 * ((expert_hidden_dim + 255) // 256)
+        
+        self.experts = nn.ModuleList([
+            Expert(dim, expert_hidden_dim, dropout)
+            for _ in range(num_experts)
+        ])
+        
+        self.gate = nn.Linear(dim, num_experts, bias=False)
+        nn.init.normal_(self.gate.weight, std=0.02)
+        
+        self.capacity_factor = capacity_factor
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        前向传播使用SparseDispatcher
+        
+        Args:
+            x: 输入 [batch, seq_len, dim]
+        
+        Returns:
+            output: 输出 [batch, seq_len, dim]
+            aux_loss: 辅助损失
+        """
+        B, T, D = x.shape
+        x_flat = x.view(-1, D)
+        
+        gates = F.softmax(self.gate(x_flat), dim=-1)
+        
+        top_k_gates, top_k_indices = torch.topk(gates, self.top_k, dim=-1)
+        top_k_gates = F.softmax(top_k_gates, dim=-1)
+        
+        expert_probs = gates.mean(dim=0)
+        expert_counts = F.one_hot(top_k_indices, self.num_experts).float().sum(dim=[0, 1])
+        expert_counts = expert_counts / (B * T * self.top_k)
+        aux_loss = self.num_experts * torch.sum(expert_probs * expert_counts)
+        
+        output = torch.zeros_like(x_flat)
+        
+        for expert_idx, expert in enumerate(self.experts):
+            expert_mask = (top_k_indices == expert_idx)
+            token_indices, topk_positions = torch.where(expert_mask)
+            
+            if len(token_indices) == 0:
+                continue
+            
+            expert_input = x_flat[token_indices]
+            expert_gates = top_k_gates[token_indices, topk_positions]
+            
+            expert_output = expert(expert_input)
+            expert_output = expert_output * expert_gates.unsqueeze(-1)
+            
+            output.index_add_(0, token_indices, expert_output)
+        
+        output = output.view(B, T, D)
+        
+        return output, aux_loss

multimodel_fusion.py

@@ -0,0 +1,522 @@
+"""
+跨模态融合模块 - SOTA级别
+支持深度跨模态交互、对比学习、模态对齐
+修复版本：解决了所有接口不匹配和潜在bug
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, List, Optional, Tuple, Union
+from components import RMSNorm
+from transformer import GroupedQueryAttention
+import math
+from contrastive_learning import MultiModalContrastiveLoss
+
+
+class CrossModalAttention(nn.Module):
+    """跨模态注意力 - 允许不同模态之间的信息交互"""
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int = 16,
+        dropout: float = 0.1,
+        qkv_bias: bool = True
+    ):
+        super().__init__()
+        self.dim = dim
+        self.n_heads = n_heads
+        self.head_dim = dim // n_heads
+        self.scale = self.head_dim ** -0.5
+        
+        assert dim % n_heads == 0, f"dim {dim} must be divisible by n_heads {n_heads}"
+        
+        self.q_proj = nn.Linear(dim, dim, bias=qkv_bias)
+        self.k_proj = nn.Linear(dim, dim, bias=qkv_bias)
+        self.v_proj = nn.Linear(dim, dim, bias=qkv_bias)
+        self.o_proj = nn.Linear(dim, dim)
+        
+        self.attn_dropout = nn.Dropout(dropout)
+        self.resid_dropout = nn.Dropout(dropout)
+        
+        self.norm_q = RMSNorm(dim)
+        self.norm_k = RMSNorm(dim)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Args:
+            query: [B, T_q, D] - 查询模态
+            key: [B, T_k, D] - 键模态
+            value: [B, T_v, D] - 值模态 (通常与key相同)
+        """
+        B, T_q, D = query.shape
+        T_k = key.shape[1]
+        
+        # 归一化
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+        
+        # 投影并重塑
+        q = self.q_proj(query).view(B, T_q, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(key).view(B, T_k, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(value).view(B, T_k, self.n_heads, self.head_dim).transpose(1, 2)
+        
+        # 使用Flash Attention或手动实现
+        if hasattr(F, 'scaled_dot_product_attention'):
+            dropout_p = self.attn_dropout.p if self.training else 0.0
+            attn_output = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=attention_mask,
+                dropout_p=dropout_p,
+                is_causal=False
+            )
+        else:
+            attn_scores = (q @ k.transpose(-2, -1)) * self.scale
+            if attention_mask is not None:
+                attn_scores = attn_scores + attention_mask
+            attn_weights = F.softmax(attn_scores, dim=-1)
+            attn_weights = self.attn_dropout(attn_weights)
+            attn_output = attn_weights @ v
+        
+        # 重塑并投影输出
+        attn_output = attn_output.transpose(1, 2).contiguous().view(B, T_q, D)
+        output = self.resid_dropout(self.o_proj(attn_output))
+        
+        return output
+
+
+class ModalityProjector(nn.Module):
+    """模态投影器 - 将不同模态投影到统一空间"""
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        hidden_dim: Optional[int] = None,
+        num_layers: int = 2,
+        use_layer_norm: bool = True
+    ):
+        super().__init__()
+        if hidden_dim is None:
+            hidden_dim = (input_dim + output_dim) // 2
+        
+        layers = []
+        for i in range(num_layers):
+            if i == 0:
+                layers.append(nn.Linear(input_dim, hidden_dim))
+            elif i == num_layers - 1:
+                layers.append(nn.Linear(hidden_dim, output_dim))
+            else:
+                layers.append(nn.Linear(hidden_dim, hidden_dim))
+            
+            if i < num_layers - 1:
+                if use_layer_norm:
+                    layers.append(RMSNorm(hidden_dim))
+                layers.append(nn.GELU())
+        
+        self.projector = nn.Sequential(*layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.projector(x)
+
+
+class ModalityAdapter(nn.Module):
+    """模态适配器 - 为每个模态学习特定的适配参数"""
+    def __init__(
+        self,
+        dim: int,
+        bottleneck_dim: int = 64,
+        num_modalities: int = 4
+    ):
+        super().__init__()
+        self.adapters = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(dim, bottleneck_dim),
+                nn.GELU(),
+                nn.Linear(bottleneck_dim, dim)
+            )
+            for _ in range(num_modalities)
+        ])
+        # 初始化为零，确保开始时是恒等映射
+        for adapter in self.adapters:
+            nn.init.zeros_(adapter[-1].weight)
+            nn.init.zeros_(adapter[-1].bias)
+
+    def forward(self, x: torch.Tensor, modality_id: int) -> torch.Tensor:
+        if modality_id >= len(self.adapters):
+            return x
+        return x + self.adapters[modality_id](x)
+
+
+class CrossModalFusionLayer(nn.Module):
+    """跨模态融合层"""
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int = 16,
+        dropout: float = 0.1,
+        use_adapter: bool = True,
+        adapter_dim: int = 64
+    ):
+        super().__init__()
+        self.dim = dim
+        self.use_adapter = use_adapter
+        
+        # 自注意力
+        self.self_attn = GroupedQueryAttention(
+            dim=dim,
+            n_heads=n_heads,
+            dropout=dropout,
+            attn_dropout=dropout
+        )
+        
+        # 跨模态注意力
+        self.cross_attn = CrossModalAttention(
+            dim=dim,
+            n_heads=n_heads,
+            dropout=dropout
+        )
+        
+        # 前馈网络
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, dim * 4),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * 4, dim),
+            nn.Dropout(dropout)
+        )
+        
+        # 归一化层
+        self.norm1 = RMSNorm(dim)
+        self.norm2 = RMSNorm(dim)
+        self.norm3 = RMSNorm(dim)
+        
+        # 模态适配器
+        if use_adapter:
+            self.adapter = ModalityAdapter(dim, adapter_dim)
+        else:
+            self.adapter = None
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        context: Optional[torch.Tensor] = None,
+        modality_id: Optional[int] = None,
+        attention_mask: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: 当前模态特征 [B, T, D]
+            context: 其他模态的上下文 [B, T_ctx, D]
+            modality_id: 模态ID（用于adapter）
+            attention_mask: 注意力掩码
+        """
+        # 自注意力 - 返回 (output, present_kv, attention_weights)
+        attn_out = self.self_attn(
+            self.norm1(x), 
+            attention_mask=attention_mask
+        )[0]  # 只取输出
+        x = x + attn_out
+        
+        # 跨模态注意力（如果有上下文）
+        if context is not None:
+            cross_attn_out = self.cross_attn(
+                self.norm2(x),
+                context,
+                context,
+                attention_mask=None
+            )
+            x = x + cross_attn_out
+        
+        # 前馈网络
+        x = x + self.ffn(self.norm3(x))
+        
+        # 模态适配器
+        if self.use_adapter and modality_id is not None and self.adapter is not None:
+            x = self.adapter(x, modality_id)
+        
+        return x
+
+
+class PerceiverResampler(nn.Module):
+    """Perceiver Resampler - 压缩模态特征到固定数量的tokens"""
+    def __init__(
+        self,
+        dim: int,
+        depth: int = 6,
+        num_latents: int = 64,
+        n_heads: int = 16,
+        dropout: float = 0.0
+    ):
+        super().__init__()
+        self.num_latents = num_latents
+        self.latents = nn.Parameter(torch.randn(num_latents, dim))
+        
+        self.layers = nn.ModuleList([
+            CrossModalFusionLayer(
+                dim=dim,
+                n_heads=n_heads,
+                dropout=dropout,
+                use_adapter=False
+            )
+            for _ in range(depth)
+        ])
+        
+        self.norm = RMSNorm(dim)
+        
+        # 初始化latents
+        nn.init.trunc_normal_(self.latents, std=0.02)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: [B, T, D] - 输入特征
+        Returns:
+            [B, num_latents, D] - 压缩后的特征
+        """
+        B = x.shape[0]
+        latents = self.latents.unsqueeze(0).expand(B, -1, -1)
+        
+        # 通过多层交叉注意力处理
+        for layer in self.layers:
+            latents = layer(latents, context=x)
+        
+        return self.norm(latents)
+
+
+class MultiModalFusionModule(nn.Module):
+    """多模态融合模块 - 整合所有融合策略"""
+    def __init__(
+        self,
+        dim: int = 2048,
+        num_fusion_layers: int = 4,
+        n_heads: int = 16,
+        dropout: float = 0.1,
+        use_perceiver: bool = True,
+        num_latents: int = 64,
+        use_contrastive: bool = True,
+        contrastive_loss_type: str = 'siglip',
+        contrastive_embed_dim: int = 512
+    ):
+        super().__init__()
+        self.dim = dim
+        self.use_perceiver = use_perceiver
+        self.use_contrastive = use_contrastive
+        
+        # 模态投影器
+        self.modality_projectors = nn.ModuleDict({
+            'image': ModalityProjector(dim, dim),
+            'audio': ModalityProjector(dim, dim),
+            'video': ModalityProjector(dim, dim),
+            'text': ModalityProjector(dim, dim)
+        })
+        
+        # 跨模态融合层
+        self.fusion_layers = nn.ModuleList([
+            CrossModalFusionLayer(
+                dim=dim,
+                n_heads=n_heads,
+                dropout=dropout,
+                use_adapter=True
+            )
+            for _ in range(num_fusion_layers)
+        ])
+        
+        # Perceiver Resampler
+        if use_perceiver:
+            self.perceiver = PerceiverResampler(
+                dim=dim,
+                depth=4,
+                num_latents=num_latents,
+                n_heads=n_heads,
+                dropout=dropout
+            )
+        
+        # 对比学习模块
+        if use_contrastive:
+            # 定义每个模态的输入维度和池化类型
+            modality_config = {
+                'text': 'cls',
+                'image': 'cls',
+                'audio': 'mean',
+                'video': 'mean'
+            }
+            
+            input_dims = {k: dim for k in modality_config.keys()}
+            
+            self.contrastive_module = MultiModalContrastiveLoss(
+                embed_dim=contrastive_embed_dim,
+                input_dims=input_dims,
+                temperature=0.07,
+                loss_type=contrastive_loss_type,
+                modality_config=modality_config
+            )
+        
+        self.final_norm = RMSNorm(dim)
+
+    def _pool_features(self, features: torch.Tensor) -> torch.Tensor:
+        """池化特征到单一向量 [B, T, D] -> [B, D]"""
+        if features.dim() == 3:
+            return features.mean(dim=1)
+        return features
+
+    def forward(
+        self,
+        segments: List[Dict],
+        compute_contrastive: bool = False
+    ) -> Dict:
+        """
+        Args:
+            segments: 列表，每个元素包含 {'type', 'data', 'modality_id'}
+                - type: str, 模态类型 ('image', 'audio', 'video', 'text')
+                - data: Tensor [B, T, D], 模态数据
+                - modality_id: int, 模态ID (0-3)
+            compute_contrastive: 是否计算对比学习损失
+        
+        Returns:
+            Dict containing:
+                - fused_features: 融合后的特征序列
+                - modality_features: 各模态的特征字典
+                - contrastive_losses: 对比学习损失字典
+        """
+        # 分离不同模态
+        modality_features = {}
+        modality_ids = {}
+        
+        for seg in segments:
+            mod_type = seg['type']
+            mod_data = seg['data']
+            mod_id = seg['modality_id']
+            
+            # 检查数据维度
+            if mod_data.dim() != 3:
+                raise ValueError(
+                    f"Expected 3D tensor [B, T, D] for modality {mod_type}, "
+                    f"got shape {mod_data.shape}"
+                )
+            
+            # 投影到统一空间
+            if mod_type in self.modality_projectors:
+                projected = self.modality_projectors[mod_type](mod_data)
+            else:
+                projected = mod_data
+            
+            # 使用Perceiver压缩（可选，非text模态）
+            if self.use_perceiver and mod_type != 'text':
+                projected = self.perceiver(projected)
+            
+            modality_features[mod_type] = projected
+            modality_ids[mod_type] = mod_id
+        
+        # 跨模态融合
+        fused_features = {}
+        
+        for mod_type, features in modality_features.items():
+            # 创建不包含当前模态的上下文
+            if len(modality_features) > 1:
+                other_features = torch.cat([
+                    f for k, f in modality_features.items() if k != mod_type
+                ], dim=1)
+            else:
+                other_features = None
+            
+            # 通过融合层
+            fused = features
+            for layer in self.fusion_layers:
+                fused = layer(
+                    fused,
+                    context=other_features,
+                    modality_id=modality_ids[mod_type]
+                )
+            
+            fused_features[mod_type] = self.final_norm(fused)
+        
+        # 计算对比学习损失（如果需要）
+        contrastive_losses = {}
+        if compute_contrastive and self.use_contrastive:
+            # 准备特征字典 - 保持3D格式供投影头处理
+            pooled_features = fused_features  # 不池化，让ProjectionHead处理
+            
+            # 定义需要对比的模态对
+            modality_pairs = []
+            if 'text' in pooled_features:
+                for mod in pooled_features.keys():
+                    if mod != 'text':
+                        modality_pairs.append((mod, 'text'))
+            
+            # 调用对比学习模块
+            if modality_pairs:
+                contrastive_losses = self.contrastive_module(
+                    pooled_features,
+                    modality_pairs=modality_pairs
+                )
+        
+        # 拼接所有融合后的特征
+        fused_sequence = torch.cat(list(fused_features.values()), dim=1)
+        
+        return {
+            'fused_features': fused_sequence,
+            'modality_features': fused_features,
+            'contrastive_losses': contrastive_losses
+        }
+
+
+class EarlyFusionModule(nn.Module):
+    """早期融合 - 在浅层就融合模态"""
+    def __init__(self, dim: int = 2048):
+        super().__init__()
+        self.fusion_proj = nn.Linear(dim, dim)
+        self.norm = RMSNorm(dim)
+
+    def forward(self, segments: List[Dict]) -> torch.Tensor:
+        """简单拼接所有模态"""
+        all_features = [seg['data'] for seg in segments]
+        fused = torch.cat(all_features, dim=1)
+        fused = self.fusion_proj(fused)
+        return self.norm(fused)
+
+
+class LateFusionModule(nn.Module):
+    """晚期融合 - 在深层才融合模态"""
+    def __init__(
+        self,
+        dim: int = 2048,
+        num_modalities: int = 4,
+        fusion_method: str = 'concat'  # 'concat', 'attention', 'average'
+    ):
+        super().__init__()
+        self.fusion_method = fusion_method
+        
+        if fusion_method == 'concat':
+            self.fusion_proj = nn.Linear(dim * num_modalities, dim)
+        elif fusion_method == 'attention':
+            self.attention_weights = nn.Linear(dim, 1)
+        
+        self.norm = RMSNorm(dim)
+
+    def forward(self, modality_outputs: List[torch.Tensor]) -> torch.Tensor:
+        """
+        Args:
+            modality_outputs: 每个模态独立处理后的输出列表 [B, T, D]
+        """
+        if self.fusion_method == 'concat':
+            # 拼接并投影
+            pooled = [x.mean(dim=1) for x in modality_outputs]
+            fused = torch.cat(pooled, dim=-1)
+            fused = self.fusion_proj(fused)
+        
+        elif self.fusion_method == 'attention':
+            # 注意力加权
+            stacked = torch.stack([x.mean(dim=1) for x in modality_outputs], dim=1)
+            weights = F.softmax(self.attention_weights(stacked), dim=1)
+            fused = (stacked * weights).sum(dim=1)
+        
+        else:  # average
+            stacked = torch.stack([x.mean(dim=1) for x in modality_outputs], dim=1)
+            fused = stacked.mean(dim=1)
+        
+        return self.norm(fused)

peft_.py

@@ -0,0 +1,213 @@
+"""
+参数高效微调 (PEFT) 模块
+支持LoRA和Adapter
+"""
+import torch
+import torch.nn as nn
+import math
+
+class LoRALayer(nn.Module):
+    """低秩适应层 (LoRA)"""
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        rank: int = 8,
+        alpha: float = 16.0,
+        dropout: float = 0.0
+    ):
+        super().__init__()
+        self.rank = rank
+        self.alpha = alpha
+        self.scaling = alpha / rank
+        
+        self.lora_A = nn.Parameter(torch.zeros(in_features, rank))
+        self.lora_B = nn.Parameter(torch.zeros(rank, out_features))
+        
+        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        
+        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+        nn.init.zeros_(self.lora_B)
+        
+        self.merged = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """前向传播"""
+        result = x @ self.lora_A @ self.lora_B
+        result = self.dropout(result)
+        return result * self.scaling
+
+class LinearWithLoRA(nn.Module):
+    """带LoRA的线性层"""
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        use_lora: bool = False,
+        lora_rank: int = 8,
+        lora_alpha: float = 16.0,
+        lora_dropout: float = 0.0
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.use_lora = use_lora
+        
+        self.base_linear = nn.Linear(in_features, out_features, bias=bias)
+        
+        if use_lora:
+            self.lora = LoRALayer(
+                in_features,
+                out_features,
+                lora_rank,
+                lora_alpha,
+                lora_dropout
+            )
+            self.merged = False
+        else:
+            self.lora = None
+            self.merged = False
+
+    def merge(self):
+        """将LoRA权重合并到基础权重中"""
+        if self.use_lora and not self.merged:
+            lora_weight = (self.lora.lora_A @ self.lora.lora_B) * self.lora.scaling
+            self.base_linear.weight.data += lora_weight.T
+            self.merged = True
+
+    def unmerge(self):
+        """取消合并LoRA权重"""
+        if self.use_lora and self.merged:
+            lora_weight = (self.lora.lora_A @ self.lora.lora_B) * self.lora.scaling
+            self.base_linear.weight.data -= lora_weight.T
+            self.merged = False
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """前向传播"""
+        output = self.base_linear(x)
+        
+        if self.use_lora and self.lora is not None and not self.merged:
+            output = output + self.lora(x)
+        
+        return output
+
+class AdapterLayer(nn.Module):
+    """Adapter层 - 轻量级微调"""
+    def __init__(
+        self,
+        dim: int,
+        bottleneck_dim: int = 64,
+        dropout: float = 0.1,
+        activation: str = 'gelu',
+        residual_scale: float = 1.0
+    ):
+        super().__init__()
+        self.residual_scale = residual_scale
+        
+        self.down_proj = nn.Linear(dim, bottleneck_dim)
+        
+        if activation == 'gelu':
+            self.activation = nn.GELU()
+        elif activation == 'relu':
+            self.activation = nn.ReLU()
+        elif activation == 'silu':
+            self.activation = nn.SiLU()
+        else:
+            self.activation = nn.GELU()
+        
+        self.up_proj = nn.Linear(bottleneck_dim, dim)
+        self.dropout = nn.Dropout(dropout)
+        
+        from components import RMSNorm
+        self.layer_norm = RMSNorm(dim)
+        
+        self._init_weights()
+
+    def _init_weights(self):
+        """初始化权重"""
+        nn.init.kaiming_uniform_(self.down_proj.weight, a=math.sqrt(5))
+        nn.init.zeros_(self.up_proj.weight)
+        if self.down_proj.bias is not None:
+            nn.init.zeros_(self.down_proj.bias)
+        if self.up_proj.bias is not None:
+            nn.init.zeros_(self.up_proj.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """前向传播"""
+        residual = x
+        
+        x = self.layer_norm(x)
+        x = self.down_proj(x)
+        x = self.activation(x)
+        x = self.dropout(x)
+        x = self.up_proj(x)
+        x = self.dropout(x)
+        
+        return residual + x * self.residual_scale
+
+class PrefixTuning(nn.Module):
+    """Prefix Tuning"""
+    def __init__(
+        self,
+        num_layers: int,
+        num_tokens: int,
+        dim: int,
+        num_heads: int
+    ):
+        super().__init__()
+        self.num_layers = num_layers
+        self.num_tokens = num_tokens
+        self.dim = dim
+        self.num_heads = num_heads
+        
+        head_dim = dim // num_heads
+        self.prefix = nn.Parameter(
+            torch.randn(num_layers, 2, num_tokens, num_heads, head_dim)
+        )
+        
+        nn.init.normal_(self.prefix, std=0.02)
+
+    def forward(self, layer_idx: int, batch_size: int) -> torch.Tensor:
+        """获取指定层的prefix"""
+        prefix = self.prefix[layer_idx]
+        prefix = prefix.unsqueeze(1).expand(
+            2, batch_size, self.num_heads, self.num_tokens, -1
+        )
+        
+        return prefix
+
+class PromptTuning(nn.Module):
+    """Prompt Tuning"""
+    def __init__(
+        self,
+        num_tokens: int,
+        dim: int,
+        init_from_vocab: bool = False,
+        vocab_embeddings: nn.Embedding = None
+    ):
+        super().__init__()
+        self.num_tokens = num_tokens
+        self.dim = dim
+        
+        self.prompt_embeddings = nn.Parameter(torch.randn(num_tokens, dim))
+        
+        if init_from_vocab and vocab_embeddings is not None:
+            indices = torch.randint(0, vocab_embeddings.num_embeddings, (num_tokens,))
+            self.prompt_embeddings.data = vocab_embeddings.weight[indices].clone()
+        else:
+            nn.init.normal_(self.prompt_embeddings, std=0.02)
+
+    def forward(self, batch_size: int) -> torch.Tensor:
+        """获取prompt embeddings"""
+        return self.prompt_embeddings.unsqueeze(0).expand(batch_size, -1, -1)
+
+class IALayer(nn.Module):
+    """(IA)³层"""
+    def __init__(self, dim: int):
+        super().__init__()
+        self.scale = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """应用缩放"""
+        return x * self.scale

post.py

@@ -0,0 +1,532 @@
+# posttrain.py
+"""
+后训练脚本 - Instruction tuning和对齐
+"""
+import os
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+from pathlib import Path
+import logging
+from tqdm import tqdm
+import json
+from datetime import datetime
+import copy
+from model import MultiModalDenseTransformer
+
+from data_loader import (
+    create_posttrain_dataloader,
+    create_preference_dataloader
+)
+from data_config import POSTTRAIN_MIX
+from reward_model import RewardModel, RewardModelTrainer
+from grpo import GRPOTrainer
+from typing import Optional
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+
+class PostTrainer:
+    """后训练器 - Supervised Fine-Tuning"""
+    def __init__(
+        self,
+        model: MultiModalDenseTransformer,
+        tokenizer,
+        learning_rate: float = 1e-5,
+        weight_decay: float = 0.01,
+        num_epochs: int = 3,
+        gradient_accumulation_steps: int = 1,
+        max_grad_norm: float = 1.0,
+        log_interval: int = 10,
+        eval_interval: int = 500,
+        save_interval: int = 1000,
+        checkpoint_dir: str = "checkpoints/posttrain"
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        self.model.to(self.device)
+        
+        # 优化器
+        self.optimizer = torch.optim.AdamW(
+            model.parameters(),
+            lr=learning_rate,
+            weight_decay=weight_decay,
+            betas=(0.9, 0.95),
+            eps=1e-8
+        )
+        
+        # 混合精度
+        self.use_amp = torch.cuda.is_available()
+        self.scaler = torch.amp.GradScaler('cuda', enabled=self.use_amp)
+        
+        # 训练参数
+        self.num_epochs = num_epochs
+        self.gradient_accumulation_steps = gradient_accumulation_steps
+        self.max_grad_norm = max_grad_norm
+        self.log_interval = log_interval
+        self.eval_interval = eval_interval
+        self.save_interval = save_interval
+        
+        # Checkpoint管理
+        self.checkpoint_dir = Path(checkpoint_dir)
+        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        
+        # 训练状态
+        self.global_step = 0
+        self.best_eval_loss = float('inf')
+        
+        logger.info(f"PostTrainer initialized:")
+        logger.info(f"  Device: {self.device}")
+        logger.info(f"  Learning Rate: {learning_rate}")
+        logger.info(f"  Num Epochs: {num_epochs}")
+        logger.info(f"  Gradient Accumulation: {gradient_accumulation_steps}")
+
+    def train_step(self, batch: dict) -> dict:
+        """单步训练"""
+        instruction_ids = batch['instruction'].to(self.device)
+        response_ids = batch['response'].to(self.device)
+        
+        # 获取 DataLoader 返回的掩码
+        instruction_mask = batch['instruction_mask'].to(self.device)
+        response_mask = batch['response_mask'].to(self.device)
+        # 拼接输入
+        input_ids = torch.cat([instruction_ids, response_ids], dim=1)
+        attention_mask = torch.cat([instruction_mask, response_mask], dim=1).float()
+        # 创建标签（只计算response部分的损失）
+        labels = input_ids.clone()
+        instr_len = instruction_ids.shape[1]
+        labels[:, :instr_len] = -100
+        labels[attention_mask == 0] = -100
+        
+        
+        # 准备输入数据
+        input_data = {
+            'segments': [{
+                'type': 'text',
+                'data': input_ids,
+                'modality_id': 0
+            }]
+        }
+        
+        # 前向传播
+        with torch.amp.autocast('cuda', enabled=self.use_amp):
+            outputs = self.model(input_data,attention_mask=attention_mask)
+            logits = outputs['logits']
+            
+            # 计算损失
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+            raw_loss = loss.item()
+            loss = loss / self.gradient_accumulation_steps
+        
+        # 反向传播
+        self.scaler.scale(loss).backward()
+        
+        return {
+            'loss': raw_loss
+        }
+
+    def optimizer_step(self):
+        """优化器步骤"""
+        self.scaler.unscale_(self.optimizer)
+        grad_norm = torch.nn.utils.clip_grad_norm_(
+            self.model.parameters(),
+            self.max_grad_norm
+        )
+        
+        self.scaler.step(self.optimizer)
+        self.scaler.update()
+        self.optimizer.zero_grad(set_to_none=True)
+        self.global_step += 1
+        return grad_norm.item()
+
+    @torch.no_grad()
+    def evaluate(self, dataloader, max_batches: int = 50) -> float:
+        """评估"""
+        self.model.eval()
+        total_loss = 0.0
+        num_batches = 0
+        
+        for i, batch in enumerate(dataloader):
+            if i >= max_batches:
+                break
+            
+            if batch is None:
+                continue
+            
+            instruction_ids = batch['instruction'].to(self.device)
+            response_ids = batch['response'].to(self.device)
+            input_ids = torch.cat([instruction_ids, response_ids], dim=1)
+            
+            labels = input_ids.clone()
+            labels[:, :instruction_ids.shape[1]] = -100
+            labels[input_ids == self.tokenizer.pad_token_id] = -100
+            
+            input_data = {
+                'segments': [{
+                    'type': 'text',
+                    'data': input_ids,
+                    'modality_id': 0
+                }]
+            }
+            
+            with torch.amp.autocast('cuda', enabled=self.use_amp):
+                outputs = self.model(input_data)
+                logits = outputs['logits']
+                
+                shift_logits = logits[:, :-1, :].contiguous()
+                shift_labels = labels[:, 1:].contiguous()
+                
+                loss = F.cross_entropy(
+                    shift_logits.view(-1, shift_logits.size(-1)),
+                    shift_labels.view(-1),
+                    ignore_index=-100
+                )
+                
+                total_loss += loss.item()
+                num_batches += 1
+        
+        self.model.train()
+        return total_loss / max(num_batches, 1)
+
+    def train(
+        self,
+        train_dataloader,
+        eval_dataloader=None,
+        resume_from: Optional[str] = None
+    ):
+        """训练循环"""
+        logger.info("\n" + "="*80)
+        logger.info("Starting Post-Training (SFT)")
+        logger.info("="*80 + "\n")
+        
+        if resume_from:
+            self.load_checkpoint(resume_from)
+        
+        self.model.train()
+        
+        for epoch in range(self.num_epochs):
+            logger.info(f"\nEpoch {epoch+1}/{self.num_epochs}")
+            
+            progress_bar = tqdm(train_dataloader, desc=f"Epoch {epoch+1}")
+            running_loss = 0.0
+            step_in_accumulation = 0
+            
+            for batch_idx, batch in enumerate(progress_bar):
+                if batch is None:
+                    continue
+                
+                # 训练步骤
+                stats = self.train_step(batch)
+                running_loss += stats['loss']
+                step_in_accumulation += 1
+                
+                # 优化器更新
+                if step_in_accumulation == self.gradient_accumulation_steps:
+                    grad_norm = self.optimizer_step()
+                    step_in_accumulation = 0
+                    
+                    # 更新进度条
+                    progress_bar.set_postfix({'loss': f"{stats['loss']:.4f}"})
+                    
+                    # 日志
+                    if self.global_step % self.log_interval == 0:
+                        avg_loss = running_loss / self.log_interval
+                        logger.info(
+                            f"Step {self.global_step} | "
+                            f"Epoch {epoch+1} | "
+                            f"Loss: {avg_loss:.4f}"
+                        )
+                        running_loss = 0.0
+                    
+                    # 评估
+                    if eval_dataloader and self.global_step % self.eval_interval == 0:
+                        eval_loss = self.evaluate(eval_dataloader)
+                        logger.info(f"Eval Loss: {eval_loss:.4f}")
+                        
+                        if eval_loss < self.best_eval_loss:
+                            self.best_eval_loss = eval_loss
+                            self.save_checkpoint(
+                                self.checkpoint_dir / "best_model.pt",
+                                is_best=True
+                            )
+                    
+                    # 保存
+                    if self.global_step % self.save_interval == 0:
+                        self.save_checkpoint(
+                            self.checkpoint_dir / f"step_{self.global_step}.pt"
+                        )
+            
+            # Epoch结束评估
+            if eval_dataloader:
+                eval_loss = self.evaluate(eval_dataloader)
+                logger.info(f"\nEpoch {epoch+1} Eval Loss: {eval_loss:.4f}")
+        
+        logger.info("\n" + "="*80)
+        logger.info("Post-Training Complete!")
+        logger.info(f"  Best Eval Loss: {self.best_eval_loss:.4f}")
+        logger.info("="*80 + "\n")
+        
+        self.save_checkpoint(self.checkpoint_dir / "final_model.pt")
+
+    def save_checkpoint(self, path: Path, is_best: bool = False):
+        """保存checkpoint"""
+        checkpoint = {
+            'model_state_dict': self.model.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'scaler_state_dict': self.scaler.state_dict() if self.use_amp else None,
+            'global_step': self.global_step,
+            'best_eval_loss': self.best_eval_loss,
+            'timestamp': datetime.now().isoformat()
+        }
+        
+        torch.save(checkpoint, path)
+        logger.info(f"Checkpoint saved to {path}" + (" (BEST)" if is_best else ""))
+
+    def load_checkpoint(self, path: str):
+        """加载checkpoint"""
+        checkpoint = torch.load(path, map_location=self.device)
+        
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        
+        if self.use_amp and checkpoint.get('scaler_state_dict'):
+            self.scaler.load_state_dict(checkpoint['scaler_state_dict'])
+        
+        self.global_step = checkpoint['global_step']
+        self.best_eval_loss = checkpoint['best_eval_loss']
+        
+        logger.info(f"Checkpoint loaded from {path}")
+
+def main():
+    """主函数"""
+    # 配置
+    config = {
+        # 模型配置
+        'model_dim': 1536,
+        'vocab_size': 151665,
+        'n_layers': 12,
+        'n_heads': 12,
+        'n_kv_heads': 4,
+        'max_seq_len': 512,
+        'dropout': 0.0,
+        'use_moe': False,
+        # 训练配置
+        'batch_size': 2,
+        'gradient_accumulation_steps': 8,
+        'learning_rate': 1e-4,
+        'weight_decay': 0.01,
+        'num_epochs': 1,
+        'max_grad_norm': 1.0,
+        
+        # 数据配置
+        'data_mix': 'debug_mix',
+        'max_samples_train': 1000,
+        'max_samples_eval': 1000,
+        'max_length': 512,
+        'num_workers': 4,
+        
+        # RLHF配置
+        'do_rlhf': False,
+        'preference_dataset': 'hh_rlhf',
+        'grpo_iterations': 3,
+        'grpo_kl_coef': 0.04,
+        'grpo_group_size': 4,
+        
+        # 路径
+        'pretrain_checkpoint': '/root/multimodal/checkpoints/pretrain_fixed/step_10000.pt',
+        'checkpoint_dir': 'checkpoints/posttrain',
+        'log_interval': 50,
+        'eval_interval': 500,
+        'save_interval': 1000,
+    }
+
+    logger.info("Configuration:")
+    logger.info(json.dumps(config, indent=2))
+
+    # 初始化tokenizer
+    logger.info("\nInitializing tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        "Qwen/Qwen2.5-7B-Instruct",
+        use_fast=True,
+        trust_remote_code=True
+    )
+
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    config['vocab_size'] = len(tokenizer)
+
+    # 初始化或加载模型
+    logger.info("\nInitializing model...")
+    model = MultiModalDenseTransformer(
+        model_dim=config['model_dim'],
+        vocab_size=config['vocab_size'],
+        n_layers=config['n_layers'],
+        n_heads=config['n_heads'],
+        n_kv_heads=config['n_kv_heads'],
+        max_seq_len=config['max_seq_len'],
+        dropout=config['dropout'],
+        use_moe=config['use_moe'],
+        use_gradient_checkpointing=False,
+        rope_scaling_type="yarn",
+        use_multimodal_fusion=False,
+        use_contrastive=False
+    )
+
+    # 加载预训练checkpoint（如果有）
+    if config['pretrain_checkpoint']:
+        logger.info(f"Loading pretrain checkpoint: {config['pretrain_checkpoint']}")
+        checkpoint = torch.load(config['pretrain_checkpoint'])
+        model.load_state_dict(checkpoint['model_state_dict'])
+
+    # ===== 阶段1: Supervised Fine-Tuning =====
+    logger.info("\n" + "="*80)
+    logger.info("PHASE 1: Supervised Fine-Tuning")
+    logger.info("="*80)
+
+    # 创建数据加载器
+    train_dataloader = create_posttrain_dataloader(
+        mix_name=config['data_mix'],
+        tokenizer=tokenizer,
+        batch_size=config['batch_size'],
+        num_workers=config['num_workers'],
+        max_length=config['max_length'],
+        max_samples=config['max_samples_train'],
+        split='train',
+        shuffle=True
+    )
+
+    eval_dataloader = create_posttrain_dataloader(
+        mix_name=config['data_mix'],
+        tokenizer=tokenizer,
+        batch_size=config['batch_size'] * 2,
+        num_workers=config['num_workers'],
+        max_length=config['max_length'],
+        max_samples=config['max_samples_eval'],
+        split='train',  # 使用train的后部分作为验证
+        shuffle=False
+    )
+
+    # 创建训练器
+    trainer = PostTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        learning_rate=config['learning_rate'],
+        weight_decay=config['weight_decay'],
+        num_epochs=config['num_epochs'],
+        gradient_accumulation_steps=config['gradient_accumulation_steps'],
+        max_grad_norm=config['max_grad_norm'],
+        log_interval=config['log_interval'],
+        eval_interval=config['eval_interval'],
+        save_interval=config['save_interval'],
+        checkpoint_dir=config['checkpoint_dir']
+    )
+
+    # 开始SFT训练
+    trainer.train(train_dataloader, eval_dataloader)
+
+    # ===== 阶段2: RLHF with GRPO =====
+    if config['do_rlhf']:
+        logger.info("\n" + "="*80)
+        logger.info("PHASE 2: RLHF with GRPO")
+        logger.info("="*80)
+        
+        try:
+            # 训练奖励模型
+            logger.info("\nTraining Reward Model...")
+            
+            reward_base_model = copy.deepcopy(model)
+            reward_model = RewardModel(reward_base_model, use_value_head=True)
+            
+            preference_dataloader = create_preference_dataloader(
+                dataset_name=config['preference_dataset'],
+                tokenizer=tokenizer,
+                batch_size=config['batch_size'],
+                num_workers=config['num_workers'],
+                max_samples=5000,
+                split='train'
+            )
+            
+            reward_trainer = RewardModelTrainer(
+                reward_model=reward_model,
+                learning_rate=1e-5
+            )
+            
+            reward_trainer.train(preference_dataloader, num_epochs=1)
+            
+            # GRPO训练
+            logger.info("\nStarting GRPO Training...")
+            
+            ref_model = copy.deepcopy(model)
+            ref_model.eval()
+            
+            grpo_trainer = GRPOTrainer(
+                actor_model=model,
+                reward_model=reward_model,
+                ref_model=ref_model,
+                tokenizer=tokenizer,
+                learning_rate=1e-6,
+                kl_coef=config['grpo_kl_coef'],
+                group_size=config['grpo_group_size'],
+                update_batch_size=2,
+                use_amp=True
+            )
+            
+            # 准备prompts
+            prompt_dataloader = create_posttrain_dataloader(
+                mix_name=config['data_mix'],
+                tokenizer=tokenizer,
+                batch_size=4,
+                num_workers=2,
+                max_samples=1000,
+                split='train'
+            )
+            
+            # 提取prompts
+            prompts = []
+            for batch in prompt_dataloader:
+                if batch and batch.get('instruction') is not None:
+                    prompts.append(batch['instruction'])
+                if len(prompts) >= 200:
+                    break
+            
+            if prompts:
+                prompt_tensor = torch.cat(prompts[:200], dim=0)
+                from torch.utils.data import TensorDataset, DataLoader
+                prompt_loader = DataLoader(
+                    TensorDataset(prompt_tensor),
+                    batch_size=4
+                )
+                
+                grpo_trainer.train(
+                    prompt_loader,
+                    num_iterations=config['grpo_iterations'],
+                    max_gen_len=50,
+                    save_path=config['checkpoint_dir'] + "/grpo"
+                )
+            
+        except Exception as e:
+            logger.error(f"Error in RLHF: {e}")
+            import traceback
+            traceback.print_exc()
+
+    logger.info("\n" + "="*80)
+    logger.info("All Training Complete!")
+    logger.info("="*80)
+
+if __name__ == "__main__":
+    main()

posttrain.py

@@ -0,0 +1,554 @@
+# posttrain.py
+"""
+后训练脚本 - Instruction tuning和对齐
+"""
+import os
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+from pathlib import Path
+import logging
+from tqdm import tqdm
+import json
+from datetime import datetime
+import copy
+from model import MultiModalDenseTransformer
+
+from data_loader import (
+    create_posttrain_dataloader,
+    create_preference_dataloader
+)
+from data_config import POSTTRAIN_MIX
+from reward_model import RewardModel, RewardModelTrainer
+from grpo import GRPOTrainer
+from typing import Optional
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+
+class PostTrainer:
+    """后训练器 - Supervised Fine-Tuning"""
+    def __init__(
+        self,
+        model: MultiModalDenseTransformer,
+        tokenizer,
+        learning_rate: float = 1e-5,
+        weight_decay: float = 0.01,
+        num_epochs: int = 3,
+        gradient_accumulation_steps: int = 1,
+        max_grad_norm: float = 1.0,
+        log_interval: int = 10,
+        eval_interval: int = 500,
+        save_interval: int = 1000,
+        checkpoint_dir: str = "checkpoints/posttrain"
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        self.model.to(self.device)
+        
+        # 优化器
+        self.optimizer = torch.optim.AdamW(
+            model.parameters(),
+            lr=learning_rate,
+            weight_decay=weight_decay,
+            betas=(0.9, 0.95),
+            eps=1e-8
+        )
+        
+        # 混合精度
+        self.use_amp = torch.cuda.is_available()
+        self.scaler = torch.amp.GradScaler('cuda', enabled=self.use_amp)
+        
+        # 训练参数
+        self.num_epochs = num_epochs
+        self.gradient_accumulation_steps = gradient_accumulation_steps
+        self.max_grad_norm = max_grad_norm
+        self.log_interval = log_interval
+        self.eval_interval = eval_interval
+        self.save_interval = save_interval
+        
+        # Checkpoint管理
+        self.checkpoint_dir = Path(checkpoint_dir)
+        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        
+        # 训练状态
+        self.global_step = 0
+        self.best_eval_loss = float('inf')
+        
+        logger.info(f"PostTrainer initialized:")
+        logger.info(f"  Device: {self.device}")
+        logger.info(f"  Learning Rate: {learning_rate}")
+        logger.info(f"  Num Epochs: {num_epochs}")
+        logger.info(f"  Gradient Accumulation: {gradient_accumulation_steps}")
+
+    def train_step(self, batch: dict) -> dict:
+        """单步训练"""
+        instruction_ids = batch['instruction'].to(self.device)
+        response_ids = batch['response'].to(self.device)
+        
+        # 1. 获取 Mask (这是之前代码里漏掉的)
+        instruction_mask = batch['instruction_mask'].to(self.device)
+        response_mask = batch['response_mask'].to(self.device)
+        
+        # 2. 拼接输入 ID 和 Mask
+        input_ids = torch.cat([instruction_ids, response_ids], dim=1)
+        attention_mask = torch.cat([instruction_mask, response_mask], dim=1)
+        
+        batch_size , seq_len = input_ids.shape
+        position_ids=torch.zeros_like(input_ids)
+
+        for i in range(batch_size):
+            non_pad_mask = attention_mask[i].bool()
+            if non_pad_mask.any():
+                positions=torch.cumsum(non_pad_mask.long(), dim=0) -1 
+                position_ids[i] = positions * non_pad_mask.long()
+
+
+
+
+
+        # 3. 创建标签
+        labels = input_ids.clone()
+        
+        # 屏蔽 Instruction 部分
+        instr_len = instruction_ids.shape[1]
+        labels[:, :instr_len] = -100
+
+        labels[attention_mask == 0] = -100
+ 
+        
+        # 准备输入数据
+        input_data = {
+            'segments': [{
+                'type': 'text',
+                'data': input_ids,
+                'modality_id': 0
+            }]
+        }
+        
+        # 前向传播
+        with torch.amp.autocast('cuda', enabled=self.use_amp):
+            # === 核心修改点 2 ===
+            # 必须传入 attention_mask，否则 transformer 不知道哪里是 padding
+            outputs = self.model(input_data, attention_mask=attention_mask,
+                                 position_ids = position_ids)
+            
+            logits = outputs['logits']
+            
+            # 计算损失
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100
+            )
+            raw_loss = loss.item()
+            loss = loss / self.gradient_accumulation_steps
+        
+        # 反向传播
+        self.scaler.scale(loss).backward()
+        
+        return {
+            'loss': raw_loss
+        }
+
+    def optimizer_step(self):
+        """优化器步骤"""
+        self.scaler.unscale_(self.optimizer)
+        grad_norm = torch.nn.utils.clip_grad_norm_(
+            self.model.parameters(),
+            self.max_grad_norm
+        )
+        
+        self.scaler.step(self.optimizer)
+        self.scaler.update()
+        self.optimizer.zero_grad(set_to_none=True)
+        self.global_step += 1
+        return grad_norm.item()
+
+    @torch.no_grad()
+    def evaluate(self, dataloader, max_batches: int = 50) -> float:
+        """评估"""
+        self.model.eval()
+        total_loss = 0.0
+        num_batches = 0
+        
+        for i, batch in enumerate(dataloader):
+            if i >= max_batches:
+                break
+            
+            if batch is None:
+                continue
+            
+            instruction_ids = batch['instruction'].to(self.device)
+            response_ids = batch['response'].to(self.device)
+            input_ids = torch.cat([instruction_ids, response_ids], dim=1)
+            
+            labels = input_ids.clone()
+            labels[:, :instruction_ids.shape[1]] = -100
+            labels[input_ids == self.tokenizer.pad_token_id] = -100
+            
+            input_data = {
+                'segments': [{
+                    'type': 'text',
+                    'data': input_ids,
+                    'modality_id': 0
+                }]
+            }
+            
+            with torch.amp.autocast('cuda', enabled=self.use_amp):
+                outputs = self.model(input_data)
+                logits = outputs['logits']
+                
+                shift_logits = logits[:, :-1, :].contiguous()
+                shift_labels = labels[:, 1:].contiguous()
+                
+                loss = F.cross_entropy(
+                    shift_logits.view(-1, shift_logits.size(-1)),
+                    shift_labels.view(-1),
+                    ignore_index=-100
+                )
+                
+                total_loss += loss.item()
+                num_batches += 1
+        
+        self.model.train()
+        return total_loss / max(num_batches, 1)
+
+    def train(
+        self,
+        train_dataloader,
+        eval_dataloader=None,
+        resume_from: Optional[str] = None
+    ):
+        """训练循环"""
+        logger.info("\n" + "="*80)
+        logger.info("Starting Post-Training (SFT)")
+        logger.info("="*80 + "\n")
+        
+        if resume_from:
+            self.load_checkpoint(resume_from)
+        
+        self.model.train()
+        
+        for epoch in range(self.num_epochs):
+            logger.info(f"\nEpoch {epoch+1}/{self.num_epochs}")
+            
+            progress_bar = tqdm(train_dataloader, desc=f"Epoch {epoch+1}")
+            running_loss = 0.0
+            step_in_accumulation = 0
+            
+            for batch_idx, batch in enumerate(progress_bar):
+                if batch is None:
+                    continue
+                
+                # 训练步骤
+                stats = self.train_step(batch)
+                running_loss += stats['loss']
+                step_in_accumulation += 1
+                
+                # 优化器更新
+                if step_in_accumulation == self.gradient_accumulation_steps:
+                    grad_norm = self.optimizer_step()
+                    step_in_accumulation = 0
+                    
+                    # 更新进度条
+                    progress_bar.set_postfix({'loss': f"{stats['loss']:.4f}"})
+                    
+                    # 日志
+                    if self.global_step % self.log_interval == 0:
+                        avg_loss = running_loss / self.log_interval
+                        logger.info(
+                            f"Step {self.global_step} | "
+                            f"Epoch {epoch+1} | "
+                            f"Loss: {avg_loss:.4f}"
+                        )
+                        running_loss = 0.0
+                    
+                    # 评估
+                    if eval_dataloader and self.global_step % self.eval_interval == 0:
+                        eval_loss = self.evaluate(eval_dataloader)
+                        logger.info(f"Eval Loss: {eval_loss:.4f}")
+                        
+                        if eval_loss < self.best_eval_loss:
+                            self.best_eval_loss = eval_loss
+                            self.save_checkpoint(
+                                self.checkpoint_dir / "best_model.pt",
+                                is_best=True
+                            )
+                    
+                    # 保存
+                    if self.global_step % self.save_interval == 0:
+                        self.save_checkpoint(
+                            self.checkpoint_dir / f"step_{self.global_step}.pt"
+                        )
+            
+            # Epoch结束评估
+            if eval_dataloader:
+                eval_loss = self.evaluate(eval_dataloader)
+                logger.info(f"\nEpoch {epoch+1} Eval Loss: {eval_loss:.4f}")
+        
+        logger.info("\n" + "="*80)
+        logger.info("Post-Training Complete!")
+        logger.info(f"  Best Eval Loss: {self.best_eval_loss:.4f}")
+        logger.info("="*80 + "\n")
+        
+        self.save_checkpoint(self.checkpoint_dir / "final_model.pt")
+
+    def save_checkpoint(self, path: Path, is_best: bool = False):
+        """保存checkpoint"""
+        checkpoint = {
+            'model_state_dict': self.model.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'scaler_state_dict': self.scaler.state_dict() if self.use_amp else None,
+            'global_step': self.global_step,
+            'best_eval_loss': self.best_eval_loss,
+            'timestamp': datetime.now().isoformat()
+        }
+        
+        torch.save(checkpoint, path)
+        logger.info(f"Checkpoint saved to {path}" + (" (BEST)" if is_best else ""))
+
+    def load_checkpoint(self, path: str):
+        """加载checkpoint"""
+        checkpoint = torch.load(path, map_location=self.device)
+        
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        
+        if self.use_amp and checkpoint.get('scaler_state_dict'):
+            self.scaler.load_state_dict(checkpoint['scaler_state_dict'])
+        
+        self.global_step = checkpoint['global_step']
+        self.best_eval_loss = checkpoint['best_eval_loss']
+        
+        logger.info(f"Checkpoint loaded from {path}")
+
+def main():
+    """主函数"""
+    # 配置
+    config = {
+        # 模型配置
+        'model_dim': 1536,
+        'vocab_size': 151665,
+        'n_layers': 12,
+        'n_heads': 12,
+        'n_kv_heads': 4,
+        'max_seq_len': 512,
+        'dropout': 0.0,
+        'use_moe': False,
+        # 训练配置
+        'batch_size': 2,
+        'gradient_accumulation_steps': 8,
+        'learning_rate': 1e-5,
+        'weight_decay': 0.01,
+        'num_epochs': 3,
+        'max_grad_norm': 1.0,
+        
+        # 数据配置
+        'data_mix': 'simple_instruct',
+        'max_samples_train': 20000,
+        'max_samples_eval': 1000,
+        'max_length': 512,
+        'num_workers': 4,
+        
+        # RLHF配置
+        'do_rlhf': False,
+        'preference_dataset': 'hh_rlhf',
+        'grpo_iterations': 3,
+        'grpo_kl_coef': 0.04,
+        'grpo_group_size': 4,
+        
+        # 路径
+        'pretrain_checkpoint': '/root/multimodal/checkpoints/pretrain_fixed/step_10000.pt',
+        'checkpoint_dir': 'checkpoints/posttrain',
+        'log_interval': 50,
+        'eval_interval': 500,
+        'save_interval': 1000,
+    }
+
+    logger.info("Configuration:")
+    logger.info(json.dumps(config, indent=2))
+
+    # 初始化tokenizer
+    logger.info("\nInitializing tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        "Qwen/Qwen2.5-7B-Instruct",
+        use_fast=True,
+        trust_remote_code=True
+    )
+
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+
+    config['vocab_size'] = len(tokenizer)
+
+    # 初始化或加载模型
+    logger.info("\nInitializing model...")
+    model = MultiModalDenseTransformer(
+        model_dim=config['model_dim'],
+        vocab_size=config['vocab_size'],
+        n_layers=config['n_layers'],
+        n_heads=config['n_heads'],
+        n_kv_heads=config['n_kv_heads'],
+        max_seq_len=config['max_seq_len'],
+        dropout=config['dropout'],
+        use_moe=config['use_moe'],
+        use_gradient_checkpointing=False,
+        rope_scaling_type="yarn",
+        use_multimodal_fusion=False,
+        use_contrastive=False
+    )
+
+    # 加载预训练checkpoint（如果有）
+    if config['pretrain_checkpoint']:
+        logger.info(f"Loading pretrain checkpoint: {config['pretrain_checkpoint']}")
+        checkpoint = torch.load(config['pretrain_checkpoint'])
+        model.load_state_dict(checkpoint['model_state_dict'])
+
+    # ===== 阶段1: Supervised Fine-Tuning =====
+    logger.info("\n" + "="*80)
+    logger.info("PHASE 1: Supervised Fine-Tuning")
+    logger.info("="*80)
+
+    # 创建数据加载器
+    train_dataloader = create_posttrain_dataloader(
+        mix_name=config['data_mix'],
+        tokenizer=tokenizer,
+        batch_size=config['batch_size'],
+        num_workers=config['num_workers'],
+        max_length=config['max_length'],
+        max_samples=config['max_samples_train'],
+        split='train',
+        shuffle=True
+    )
+
+    eval_dataloader = create_posttrain_dataloader(
+        mix_name=config['data_mix'],
+        tokenizer=tokenizer,
+        batch_size=config['batch_size'] * 2,
+        num_workers=config['num_workers'],
+        max_length=config['max_length'],
+        max_samples=config['max_samples_eval'],
+        split='train',  # 使用train的后部分作为验证
+        shuffle=False
+    )
+
+    # 创建训练器
+    trainer = PostTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        learning_rate=config['learning_rate'],
+        weight_decay=config['weight_decay'],
+        num_epochs=config['num_epochs'],
+        gradient_accumulation_steps=config['gradient_accumulation_steps'],
+        max_grad_norm=config['max_grad_norm'],
+        log_interval=config['log_interval'],
+        eval_interval=config['eval_interval'],
+        save_interval=config['save_interval'],
+        checkpoint_dir=config['checkpoint_dir']
+    )
+
+    # 开始SFT训练
+    trainer.train(train_dataloader, eval_dataloader)
+
+    # ===== 阶段2: RLHF with GRPO =====
+    if config['do_rlhf']:
+        logger.info("\n" + "="*80)
+        logger.info("PHASE 2: RLHF with GRPO")
+        logger.info("="*80)
+        
+        try:
+            # 训练奖励模型
+            logger.info("\nTraining Reward Model...")
+            
+            reward_base_model = copy.deepcopy(model)
+            reward_model = RewardModel(reward_base_model, use_value_head=True)
+            
+            preference_dataloader = create_preference_dataloader(
+                dataset_name=config['preference_dataset'],
+                tokenizer=tokenizer,
+                batch_size=config['batch_size'],
+                num_workers=config['num_workers'],
+                max_samples=5000,
+                split='train'
+            )
+            
+            reward_trainer = RewardModelTrainer(
+                reward_model=reward_model,
+                learning_rate=1e-5
+            )
+            
+            reward_trainer.train(preference_dataloader, num_epochs=1)
+            
+            # GRPO训练
+            logger.info("\nStarting GRPO Training...")
+            
+            ref_model = copy.deepcopy(model)
+            ref_model.eval()
+            
+            grpo_trainer = GRPOTrainer(
+                actor_model=model,
+                reward_model=reward_model,
+                ref_model=ref_model,
+                tokenizer=tokenizer,
+                learning_rate=1e-6,
+                kl_coef=config['grpo_kl_coef'],
+                group_size=config['grpo_group_size'],
+                update_batch_size=2,
+                use_amp=True
+            )
+            
+            # 准备prompts
+            prompt_dataloader = create_posttrain_dataloader(
+                mix_name=config['data_mix'],
+                tokenizer=tokenizer,
+                batch_size=4,
+                num_workers=2,
+                max_samples=1000,
+                split='train'
+            )
+            
+            # 提取prompts
+            prompts = []
+            for batch in prompt_dataloader:
+                if batch and batch.get('instruction') is not None:
+                    prompts.append(batch['instruction'])
+                if len(prompts) >= 200:
+                    break
+            
+            if prompts:
+                prompt_tensor = torch.cat(prompts[:200], dim=0)
+                from torch.utils.data import TensorDataset, DataLoader
+                prompt_loader = DataLoader(
+                    TensorDataset(prompt_tensor),
+                    batch_size=4
+                )
+                
+                grpo_trainer.train(
+                    prompt_loader,
+                    num_iterations=config['grpo_iterations'],
+                    max_gen_len=50,
+                    save_path=config['checkpoint_dir'] + "/grpo"
+                )
+            
+        except Exception as e:
+            logger.error(f"Error in RLHF: {e}")
+            import traceback
+            traceback.print_exc()
+
+    logger.info("\n" + "="*80)
+    logger.info("All Training Complete!")
+    logger.info("="*80)
+
+if __name__ == "__main__":
+    main()

pretrain.py

@@ -0,0 +1,502 @@
+# pretrain.py - 完全修复版本
+
+import os
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer
+from pathlib import Path
+import logging
+from tqdm import tqdm
+import json
+from datetime import datetime
+from model import MultiModalDenseTransformer
+from data_loader import create_pretrain_dataloader
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
+
+
+class PreTrainer:
+    """预训练器 - 完全修复版"""
+    def __init__(
+        self,
+        model: MultiModalDenseTransformer,
+        tokenizer,
+        learning_rate: float = 3e-4,
+        weight_decay: float = 0.1,
+        warmup_steps: int = 1000,
+        max_steps: int = 100000,
+        gradient_accumulation_steps: int = 16,
+        max_grad_norm: float = 1.0,
+        log_interval: int = 10,
+        save_interval: int = 1000,
+        checkpoint_dir: str = "checkpoints/pretrain",
+        loss_log_file: str = "checkpoints/pretrain/train_loss.log"
+    ):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        self.model.to(self.device)
+        
+        # 优化器配置 - 使用标准AdamW参数
+        self.optimizer = torch.optim.AdamW(
+            model.parameters(),
+            lr=learning_rate,
+            weight_decay=weight_decay,
+            betas=(0.9, 0.95),
+            eps=1e-8
+        )
+        
+        # 🔧 修复：使用更简单的学习率调度器
+        from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts
+        
+        # Warmup + Cosine Decay
+        self.warmup_steps = warmup_steps
+        self.max_lr = learning_rate
+        self.min_lr = learning_rate * 0.1
+        self.current_step = 0
+        
+        # 混合精度
+        self.use_amp = torch.cuda.is_available()
+        self.scaler = torch.amp.GradScaler('cuda', enabled=self.use_amp)
+        
+        # 训练参数
+        self.gradient_accumulation_steps = gradient_accumulation_steps
+        self.max_grad_norm = max_grad_norm
+        self.max_steps = max_steps
+        self.log_interval = log_interval
+        self.save_interval = save_interval
+        
+        # Checkpoint管理
+        self.checkpoint_dir = Path(checkpoint_dir)
+        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        
+        # 损失日志
+        self.loss_log_file = Path(loss_log_file)
+        self.loss_log_file.parent.mkdir(parents=True, exist_ok=True)
+        
+        # 训练状态
+        self.global_step = 0
+        self.tokens_seen = 0
+        self.running_loss = 0.0
+        self.best_loss = float('inf')
+        
+        logger.info(f"PreTrainer initialized:")
+        logger.info(f"  Device: {self.device}")
+        logger.info(f"  Learning Rate: {learning_rate}")
+        logger.info(f"  Max Steps: {max_steps}")
+        logger.info(f"  Gradient Accumulation: {gradient_accumulation_steps}")
+        logger.info(f"  Effective Batch Size: {gradient_accumulation_steps}")
+        logger.info(f"  Mixed Precision: {self.use_amp}")
+
+    def _get_lr(self) -> float:
+        """手动计算学习率（Warmup + Cosine）"""
+        if self.current_step < self.warmup_steps:
+            # Linear warmup
+            return self.max_lr * (self.current_step / self.warmup_steps)
+        else:
+            # Cosine decay
+            progress = (self.current_step - self.warmup_steps) / (self.max_steps - self.warmup_steps)
+            return self.min_lr + (self.max_lr - self.min_lr) * 0.5 * (1 + torch.cos(torch.tensor(progress * 3.14159)))
+
+    def _set_lr(self, lr: float):
+        """设置学习率"""
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = lr
+
+    def train_step(self, batch: dict) -> dict:
+        """
+        🔧 完全修复的训练步骤
+        关键：不要在loss计算时除以gradient_accumulation_steps
+        """
+        input_ids = batch['input_ids'].to(self.device)
+        attention_mask = batch['attention_mask'].to(self.device)
+        batch_size, seq_len = input_ids.shape
+        position_ids= torch.zeros_like(input_ids)
+        
+        for i in range(batch_size):
+            non_pad_mask = attention_mask[i].bool()
+            if non_pad_mask.any():
+                positions = torch.cumsum(non_pad_mask.long(), dim=0) -1
+                position_ids[i]=positions * non_pad_mask.long()
+
+
+        
+        # 准备输入
+        input_data = {
+            'segments': [{
+                'type': 'text',
+                'data': input_ids,
+                'modality_id': 0
+            }]
+        }
+        
+        # 前向传播
+        with torch.amp.autocast('cuda', enabled=self.use_amp):
+            outputs = self.model(
+                input_data,
+                attention_mask=attention_mask,
+                position_ids=position_ids)
+            logits = outputs['logits']
+            
+            # 计算损失（标准自回归）
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = input_ids[:, 1:].contiguous()
+            shift_attention_mask = attention_mask[:, 1:].contiguous()
+            
+            # 🔧 关键修复：直接计算平均loss，不要除以gradient_accumulation_steps
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                reduction='none'
+            )
+            
+            # 应用mask
+            loss = (loss * shift_attention_mask.view(-1)).sum() / (shift_attention_mask.sum() + 1e-8)
+            
+            # 🔧 重要：为了数值稳定，在这里手动处理梯度累积
+            # 方法：缩放loss用于反向传播，但记录原始loss
+            loss_for_backward = loss / self.gradient_accumulation_steps
+        
+        # 反向传播（使用缩放后的loss）
+        self.scaler.scale(loss_for_backward).backward()
+        
+        # 🔧 关键修复：不在这里累积loss，改在optimizer_step时累积
+        # self.running_loss += loss.item()  # ❌ 移除
+        self.tokens_seen += attention_mask.sum().item()
+        
+        return {
+            'loss': loss.item(),  # 返回真实的、未缩放的loss
+            'lr': self.optimizer.param_groups[0]['lr']
+        }
+
+    def optimizer_step(self):
+        """优化器步骤"""
+        # Unscale梯度
+        self.scaler.unscale_(self.optimizer)
+        
+        # 梯度裁剪
+        grad_norm = torch.nn.utils.clip_grad_norm_(
+            self.model.parameters(),
+            self.max_grad_norm
+        )
+        
+        # 更新参数
+        self.scaler.step(self.optimizer)
+        self.scaler.update()
+        self.optimizer.zero_grad(set_to_none=True)
+        
+        # 更新学习率
+        self.current_step += 1
+        self.global_step += 1
+        lr = self._get_lr()
+        self._set_lr(lr)
+        
+        return grad_norm.item()
+
+    def _write_loss_to_txt(self, step, avg_loss, lr, tokens_seen):
+        """写入损失日志"""
+        log_content = (
+            f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] "
+            f"Step: {step}/{self.max_steps}, "
+            f"Average Loss: {avg_loss:.4f}, "
+            f"Learning Rate: {lr:.2e}, "
+            f"Tokens Seen: {tokens_seen/1e9:.2f}B\n"
+        )
+        with open(self.loss_log_file, 'a', encoding='utf-8') as f:
+            f.write(log_content)
+
+    def train(self, dataloader, resume_from=None):
+        """训练循环"""
+        logger.info("\n" + "="*80)
+        logger.info("Starting Pre-Training (Fixed Version)")
+        logger.info("="*80 + "\n")
+        
+        # 恢复训练
+        if resume_from:
+            self.load_checkpoint(resume_from)
+        
+        # 初始化日志
+        if not self.loss_log_file.exists():
+            with open(self.loss_log_file, 'w', encoding='utf-8') as f:
+                f.write("🚀 Fixed Training Log (Real Loss Values)\n")
+                f.write("="*80 + "\n")
+        
+        self.model.train()
+        progress_bar = tqdm(total=self.max_steps, initial=self.global_step)
+        
+        step_in_accumulation = 0
+        accumulated_loss = 0.0  # 🔧 用于累积一个完整step的loss
+        
+        batches_to_skip = self.global_step * self.gradient_accumulation_steps
+        
+        logger.info(f"Current Global Step: {self.global_step}")
+        if batches_to_skip > 0:
+            logger.info(f"🔄 Resuming: Need to skip {batches_to_skip} batches to restore data state...")
+            logger.info("This might take a while depending on network/disk speed...")
+        
+        # 创建迭代器
+        data_iterator = iter(dataloader)
+        
+        # 1. 执行跳过逻辑
+        skipped = 0
+        if batches_to_skip > 0:
+            with tqdm(total=batches_to_skip, desc="Skipping trained batches", unit="batch") as skip_pbar:
+                while skipped < batches_to_skip:
+                    try:
+                        # 只取数据，不进模型，不计算梯度
+                        _ = next(data_iterator)
+                        skipped += 1
+                        skip_pbar.update(1)
+                    except StopIteration:
+                        logger.error("Dataset exhausted during skipping! Check your dataset size or max_steps.")
+                        return
+        
+        logger.info("✅ Data fast-forward complete. Resuming training...")
+
+        # 2. 正式训练循环
+        try:
+            # 注意：这里不能再用 for batch in dataloader，因为迭代器已经被消费了一部分
+            # 我们继续使用上面创建的 data_iterator
+            while True:
+                try:
+                    batch = next(data_iterator)
+                except StopIteration:
+                    break # 数据耗尽
+
+                if batch is None or batch['input_ids'].size(0) == 0:
+                    continue
+                #print("Sample input:", self.tokenizer.decode(batch['input_ids'][0][:50]))
+                # 训练步骤
+                stats = self.train_step(batch)
+                step_in_accumulation += 1
+                accumulated_loss += stats['loss']  # 🔧 累积当前micro-batch的loss
+                
+                # 梯度累积完成，执行优化器更新
+                if step_in_accumulation >= self.gradient_accumulation_steps:
+                    # 🔧 计算这个完整step的平均loss
+                    avg_step_loss = accumulated_loss / self.gradient_accumulation_steps
+                    
+                    grad_norm = self.optimizer_step()
+                    stats['grad_norm'] = grad_norm
+                    stats['loss'] = avg_step_loss  # 🔧 更新为平均loss
+                    
+                    # 🔧 累积到running_loss（用于日志记录）
+                    self.running_loss += avg_step_loss
+                    
+                    step_in_accumulation = 0
+                    accumulated_loss = 0.0  # 🔧 重置累积器
+                    
+                    # 更新进度条
+                    progress_bar.update(1)
+                    progress_bar.set_postfix({
+                        'loss': f"{stats['loss']:.4f}",
+                        'lr': f"{stats['lr']:.2e}",
+                        'tokens': f"{self.tokens_seen/1e9:.2f}B",
+                        'grad': f"{grad_norm:.2f}"
+                    })
+                    
+                    # 日志记录
+                    if self.global_step % self.log_interval == 0:
+                        avg_loss = self.running_loss / self.log_interval
+                        
+                        logger.info(
+                            f"Step {self.global_step}/{self.max_steps} | "
+                            f"Loss: {avg_loss:.4f} | "
+                            f"LR: {stats['lr']:.2e} | "
+                            f"GradNorm: {grad_norm:.2f} | "
+                            f"Tokens: {self.tokens_seen/1e9:.2f}B"
+                        )
+                        
+                        # 🔧 检测训练异常
+                        if avg_loss > 10.0 and self.global_step > 100:
+                            logger.warning(f"⚠️  Loss异常高 ({avg_loss:.2f})，可能存在问题！")
+                        
+                        if avg_loss < self.best_loss:
+                            self.best_loss = avg_loss
+                            logger.info(f"✨ New best loss: {self.best_loss:.4f}")
+                        
+                        self._write_loss_to_txt(
+                            step=self.global_step,
+                            avg_loss=avg_loss,
+                            lr=stats['lr'],
+                            tokens_seen=self.tokens_seen
+                        )
+                        self.running_loss = 0.0
+                    
+                    # 保存checkpoint
+                    if self.global_step % self.save_interval == 0:
+                        self.save_checkpoint(
+                            self.checkpoint_dir / f"step_{self.global_step}.pt"
+                        )
+                    
+                    # 完成训练
+                    if self.global_step >= self.max_steps:
+                        break
+        
+        except KeyboardInterrupt:
+            logger.info("\n⚠️  Training interrupted by user")
+            self.save_checkpoint(
+                self.checkpoint_dir / f"interrupted_step_{self.global_step}.pt"
+            )
+        
+        finally:
+            progress_bar.close()
+        
+        logger.info("\n" + "="*80)
+        logger.info("Pre-Training Complete!")
+        logger.info(f"  Total Steps: {self.global_step}")
+        logger.info(f"  Total Tokens: {self.tokens_seen/1e9:.2f}B")
+        logger.info(f"  Best Loss: {self.best_loss:.4f}")
+        logger.info("="*80 + "\n")
+        
+        # 保存最终模型
+        self.save_checkpoint(self.checkpoint_dir / "final_model.pt")
+
+    def save_checkpoint(self, path: Path):
+        """保存checkpoint"""
+        checkpoint = {
+            'model_state_dict': self.model.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+            'scaler_state_dict': self.scaler.state_dict() if self.use_amp else None,
+            'global_step': self.global_step,
+            'current_step': self.current_step,
+            'tokens_seen': self.tokens_seen,
+            'best_loss': self.best_loss,
+            'timestamp': datetime.now().isoformat()
+        }
+        
+        torch.save(checkpoint, path)
+        logger.info(f"💾 Checkpoint saved to {path}")
+
+    def load_checkpoint(self, path: str):
+        """加载checkpoint"""
+        checkpoint = torch.load(path, map_location=self.device, weights_only=True)
+        
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        
+        if self.use_amp and checkpoint.get('scaler_state_dict'):
+            self.scaler.load_state_dict(checkpoint['scaler_state_dict'])
+        
+        self.global_step = checkpoint['global_step']
+        self.current_step = checkpoint.get('current_step', self.global_step)
+        self.tokens_seen = checkpoint['tokens_seen']
+        self.best_loss = checkpoint.get('best_loss', float('inf'))
+        
+        logger.info(f"📂 Checkpoint loaded from {path}")
+        logger.info(f"  Resuming from step {self.global_step}")
+        logger.info(f"  Tokens seen: {self.tokens_seen/1e9:.2f}B")
+
+
+def main():
+    """主函数"""
+    # 🔧 优化后的配置
+    config = {
+        # 模型配置
+        'model_dim': 1536,
+        'vocab_size': 151665,
+        'n_layers': 12,
+        'n_heads': 12,
+        'n_kv_heads': 4,
+        'max_seq_len': 512,  # 🔧 减小以提升速度
+        'dropout': 0.1,
+        'use_moe': False,
+        
+        # 🔧 训练配置（关键修复）
+        'batch_size': 4,           # 增加
+        'gradient_accumulation_steps': 8,  # 减少
+        'learning_rate': 3e-4,     # 标准值
+        'weight_decay': 0.1,
+        'warmup_steps': 500,       # 更快warmup
+        'max_steps': 10000,
+        'max_grad_norm': 1.0,
+        
+        # 数据配置
+        'data_mix': 'text_only',
+        'max_length': 512,         # 🔧 与max_seq_len一致
+        'num_workers': 2,          # 🔧 减少避免网络问题
+        
+        # 日志和保存
+        'log_interval': 10,
+        'save_interval': 500,      # 🔧 更频繁保存
+        'checkpoint_dir': 'checkpoints/pretrain_fixed',
+        'loss_log_file': 'checkpoints/pretrain_fixed/train_loss.log'
+    }
+    
+    logger.info("="*80)
+    logger.info("🔧 Fixed Configuration:")
+    logger.info(json.dumps(config, indent=2))
+    logger.info("="*80 + "\n")
+    
+    # 初始化tokenizer
+    logger.info("Initializing tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        "Qwen/Qwen2.5-7B-Instruct",
+        use_fast=True,
+        trust_remote_code=True
+    )
+    
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+    
+    config['vocab_size'] = len(tokenizer)
+    logger.info(f"Vocab size: {config['vocab_size']}\n")
+    
+    # 初始化模型
+    logger.info("Initializing model...")
+    model = MultiModalDenseTransformer(
+        model_dim=config['model_dim'],
+        vocab_size=config['vocab_size'],
+        n_layers=config['n_layers'],
+        n_heads=config['n_heads'],
+        n_kv_heads=config['n_kv_heads'],
+        max_seq_len=config['max_seq_len'],
+        dropout=config['dropout'],
+        use_moe=config['use_moe'],
+        use_gradient_checkpointing=True,
+        rope_scaling_type="yarn",
+        use_multimodal_fusion=False,
+        use_contrastive=False
+    )
+    
+    # 创建数据加载器
+    logger.info(f"\nCreating dataloader (mix: {config['data_mix']})...")
+    dataloader = create_pretrain_dataloader(
+        mix_name=config['data_mix'],
+        tokenizer=tokenizer,
+        batch_size=config['batch_size'],
+        num_workers=config['num_workers'],
+        max_length=config['max_length']
+    )
+    
+    # 创建训练器
+    trainer = PreTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        learning_rate=config['learning_rate'],
+        weight_decay=config['weight_decay'],
+        warmup_steps=config['warmup_steps'],
+        max_steps=config['max_steps'],
+        gradient_accumulation_steps=config['gradient_accumulation_steps'],
+        max_grad_norm=config['max_grad_norm'],
+        log_interval=config['log_interval'],
+        save_interval=config['save_interval'],
+        checkpoint_dir=config['checkpoint_dir'],
+        loss_log_file=config['loss_log_file']
+    )
+    
+    # 🔧 开始训练（从头开始，不要用旧的checkpoint）
+    logger.info("\n🚀 Starting fresh training with fixes...\n")
+    trainer.train(dataloader, resume_from="/root/step_6500.pt")
+    #trainer.train(dataloader)
+
+
+if __name__ == "__main__":
+    main()

reward_model.py

@@ -0,0 +1,189 @@
+"""
+奖励模型 - 用于RLHF
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.utils.data import DataLoader
+from collections import defaultdict
+from typing import Dict, Tuple, Union, Optional
+from tqdm import tqdm
+from model import MultiModalDenseTransformer
+
+class RewardModel(nn.Module):
+    """奖励模型 - 用于RLHF"""
+    def __init__(
+        self,
+        base_model: MultiModalDenseTransformer,
+        use_value_head: bool = True
+    ):
+        super().__init__()
+        self.base_model = base_model
+        self.use_value_head = use_value_head
+        
+        self.reward_head = nn.Sequential(
+            nn.Linear(base_model.model_dim, base_model.model_dim // 2),
+            nn.ReLU(),
+            nn.Dropout(0.1),
+            nn.Linear(base_model.model_dim // 2, 1)
+        )
+        
+        if use_value_head:
+            self.value_head = nn.Sequential(
+                nn.Linear(base_model.model_dim, base_model.model_dim // 2),
+                nn.ReLU(),
+                nn.Dropout(0.1),
+                nn.Linear(base_model.model_dim // 2, 1)
+            )
+
+    def forward(
+        self,
+        input_data: Dict,
+        return_values: bool = False
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        """前向传播"""
+        output = self.base_model(input_data, return_hidden=True)
+        hidden_states = output['last_hidden_state']
+        
+        rewards = self.reward_head(hidden_states).squeeze(-1)
+        
+        if return_values and self.use_value_head:
+            values = self.value_head(hidden_states).squeeze(-1)
+            return rewards, values
+        
+        return rewards
+
+class RewardModelTrainer:
+    """奖励模型训练器"""
+    def __init__(
+        self,
+        reward_model: RewardModel,
+        learning_rate: float = 1e-5,
+        margin: float = 0.0
+    ):
+        self.reward_model = reward_model
+        self.margin = margin
+        
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.reward_model.to(self.device)
+        
+        for param in self.reward_model.base_model.parameters():
+            param.requires_grad = False
+        
+        for layer in self.reward_model.base_model.layers[-2:]:
+            for param in layer.parameters():
+                param.requires_grad = True
+        
+        trainable_params = list(self.reward_model.reward_head.parameters())
+        if self.reward_model.use_value_head:
+            trainable_params += list(self.reward_model.value_head.parameters())
+
+        self.optimizer = optim.AdamW(
+            filter(lambda p: p.requires_grad, self.reward_model.parameters()),
+            lr=learning_rate
+        )
+
+    def train_step(self, chosen_batch: Dict, rejected_batch: Dict) -> Dict:
+        """单步训练"""
+        self.reward_model.train()
+        self.optimizer.zero_grad()
+        
+        chosen_rewards = self.reward_model(chosen_batch)[:, -1]
+        rejected_rewards = self.reward_model(rejected_batch)[:, -1]
+        
+        loss = -F.logsigmoid(chosen_rewards - rejected_rewards - self.margin).mean()
+        
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(self.reward_model.parameters(), 1.0)
+        self.optimizer.step()
+        
+        accuracy = (chosen_rewards > rejected_rewards).float().mean().item()
+        
+        return {
+            'loss': loss.item(),
+            'accuracy': accuracy
+        }
+
+    def train(
+        self,
+        dataloader: DataLoader,
+        num_epochs: int = 1,
+        log_interval: int = 10
+    ):
+        """训练循环"""
+        print(f"Starting reward model training on {self.device}...")
+        
+        for epoch in range(num_epochs):
+            total_stats = defaultdict(float)
+            num_steps = 0
+            progress_bar = tqdm(
+                dataloader,
+                desc=f"Reward Model Epoch {epoch+1}/{num_epochs}"
+            )
+            
+            for batch_idx, (chosen_ids, rejected_ids) in enumerate(progress_bar):
+                chosen_batch = {
+                    'segments': [{'type': 'text', 'data': chosen_ids.to(self.device), 'modality_id': 0}]
+                }
+                
+                rejected_batch = {
+                    'segments': [{'type': 'text', 'data': rejected_ids.to(self.device), 'modality_id': 0}]
+                }
+                
+                stats = self.train_step(chosen_batch, rejected_batch)
+                
+                for k, v in stats.items():
+                    total_stats[k] += v
+                num_steps += 1
+                
+                if (batch_idx + 1) % log_interval == 0:
+                    avg_stats = {
+                        k: v / num_steps
+                        for k, v in total_stats.items()
+                    }
+                    progress_bar.set_postfix(avg_stats)
+                    total_stats = defaultdict(float)
+        
+        print("Reward model training complete!")
+
+    def evaluate(self, dataloader: DataLoader) -> Dict[str, float]:
+        """评估奖励模型"""
+        self.reward_model.eval()
+        total_stats = defaultdict(float)
+        num_batches = 0
+        
+        with torch.no_grad():
+            for chosen_ids, rejected_ids in dataloader:
+                chosen_batch = {
+                    'segments': [{'type': 'text', 'data': chosen_ids.to(self.device), 'modality_id': 0}]
+                }
+                
+                rejected_batch = {
+                    'segments': [{'type': 'text', 'data': rejected_ids.to(self.device), 'modality_id': 0}]
+                }
+                
+                chosen_rewards = self.reward_model(chosen_batch)[:, -1]
+                rejected_rewards = self.reward_model(rejected_batch)[:, -1]
+                
+                loss = -F.logsigmoid(chosen_rewards - rejected_rewards - self.margin).mean()
+                accuracy = (chosen_rewards > rejected_rewards).float().mean().item()
+                
+                total_stats['loss'] += loss.item()
+                total_stats['accuracy'] += accuracy
+                num_batches += 1
+        
+        return {k: v / num_batches for k, v in total_stats.items()}
+
+    def save_checkpoint(self, path: str):
+        """保存检查点"""
+        torch.save({
+            'model_state_dict': self.reward_model.state_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+        }, path)
+
+    def load_checkpoint(self, path: str):
+        """加载检查点"""
+        checkpoint = torch.load(path, map_location=self.device)
+        self.reward_model.load_state_dict(checkpoint['model_state_dict'])
+        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

transformer.py

@@ -0,0 +1,335 @@
+"""
+优化的Transformer架构
+支持GQA/MQA、滑动窗口注意力、Flash Attention 2、YARN位置编码
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, List
+import math
+from components import RMSNorm, SwiGLU, YARNRotaryEmbedding, QKNorm
+from peft_ import LinearWithLoRA, AdapterLayer
+from moe import MixtureOfExperts
+
+class GroupedQueryAttention(nn.Module):
+    """分组查询注意力 (GQA) - 优化版 with YARN"""
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        n_kv_heads: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        dropout: float = 0.0,
+        attn_dropout: float = 0.0,
+        use_flash: bool = True,
+        qkv_bias: bool = False,
+        use_lora: bool = False,
+        lora_rank: int = 8,
+        max_seq_len: int = 8192,
+        rope_scaling_factor: float = 1.0,
+        rope_scaling_type: str = "yarn",
+        use_qk_norm: bool = False,
+        sliding_window: Optional[int] = None,
+        use_alibi: bool = False
+    ):
+        super().__init__()
+        
+        self.dim = dim
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads
+        
+        assert n_heads % self.n_kv_heads == 0, \
+            f"n_heads ({n_heads}) must be divisible by n_kv_heads ({self.n_kv_heads})"
+        
+        self.n_rep = n_heads // self.n_kv_heads
+        self.head_dim = head_dim if head_dim is not None else dim // n_heads
+        self.scale = self.head_dim ** -0.5
+        
+        self.use_flash = use_flash and hasattr(F, 'scaled_dot_product_attention')
+        self.sliding_window = sliding_window
+        
+        self.q_proj = LinearWithLoRA(
+            dim, n_heads * self.head_dim,
+            bias=qkv_bias, use_lora=use_lora, lora_rank=lora_rank
+        )
+        self.k_proj = LinearWithLoRA(
+            dim, self.n_kv_heads * self.head_dim,
+            bias=qkv_bias, use_lora=use_lora, lora_rank=lora_rank
+        )
+        self.v_proj = LinearWithLoRA(
+            dim, self.n_kv_heads * self.head_dim,
+            bias=qkv_bias, use_lora=use_lora, lora_rank=lora_rank
+        )
+        self.o_proj = LinearWithLoRA(
+            n_heads * self.head_dim, dim,
+            bias=False, use_lora=use_lora, lora_rank=lora_rank
+        )
+        
+        self.attn_dropout = nn.Dropout(attn_dropout) if attn_dropout > 0 else nn.Identity()
+        self.resid_dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        
+        self.use_qk_norm = use_qk_norm
+        if use_qk_norm:
+            self.q_norm = QKNorm(self.head_dim)
+            self.k_norm = QKNorm(self.head_dim)
+        
+        self.use_alibi = use_alibi
+        if use_alibi:
+            self.register_buffer(
+                "alibi_slopes",
+                self._get_alibi_slopes(n_heads),
+                persistent=False
+            )
+        else:
+            self.rotary_emb = YARNRotaryEmbedding(
+                self.head_dim,
+                max_seq_len=max_seq_len,
+                original_max_len=4096,
+                scaling_factor=rope_scaling_factor,
+                rope_percentage=1.0
+            )
+
+    def _get_alibi_slopes(self, n_heads: int) -> torch.Tensor:
+        """计算ALiBi斜率"""
+        def get_slopes_power_of_2(n):
+            start = 2 ** (-(2 ** -(math.log2(n) - 3)))
+            ratio = start
+            return [start * ratio ** i for i in range(n)]
+        
+        if math.log2(n_heads).is_integer():
+            slopes = get_slopes_power_of_2(n_heads)
+        else:
+            closest_power_of_2 = 2 ** math.floor(math.log2(n_heads))
+            slopes = get_slopes_power_of_2(closest_power_of_2)
+            extra_slopes = get_slopes_power_of_2(2 * closest_power_of_2)[::2]
+            slopes.extend(extra_slopes[:n_heads - closest_power_of_2])
+        
+        return torch.tensor(slopes).view(n_heads, 1, 1)
+
+    def repeat_kv(self, x: torch.Tensor) -> torch.Tensor:
+        """重复KV heads以匹配Q heads"""
+        if self.n_rep == 1:
+            return x
+        
+        B, n_kv_heads, seq_len, head_dim = x.shape
+        return x[:, :, None, :, :].expand(
+            B, n_kv_heads, self.n_rep, seq_len, head_dim
+        ).reshape(B, n_kv_heads * self.n_rep, seq_len, head_dim)
+
+    def _apply_sliding_window_mask(
+        self,
+        attn_scores: torch.Tensor,
+        seq_len: int
+    ) -> torch.Tensor:
+        """应用滑动窗口mask"""
+        if self.sliding_window is None or seq_len <= self.sliding_window:
+            return attn_scores
+        
+        mask = torch.ones(seq_len, seq_len, device=attn_scores.device, dtype=torch.bool)
+        mask = torch.triu(mask, diagonal=-self.sliding_window + 1)
+        mask = torch.tril(mask, diagonal=0)
+        
+        attn_scores = attn_scores.masked_fill(~mask, float('-inf'))
+        return attn_scores
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+        past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        output_attentions: bool = False
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[torch.Tensor]]:
+        """前向传播"""
+        B, T, C = x.shape
+        
+        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        
+        if self.use_qk_norm:
+            q_shape = q.shape
+            k_shape = k.shape
+            q = self.q_norm.query_norm(q.view(-1, self.head_dim)).view(q_shape)
+            k = self.k_norm.key_norm(k.view(-1, self.head_dim)).view(k_shape)
+        
+        if not self.use_alibi:
+            q, k = self.rotary_emb(q, k, position_ids)
+        
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        
+        present_kv = (k, v) if use_cache else None
+        
+        k = self.repeat_kv(k)
+        v = self.repeat_kv(v)
+        
+        seq_len_k = k.size(2)
+        
+        if self.use_flash and not output_attentions and attention_mask is None:
+            dropout_p = self.attn_dropout.p if isinstance(self.attn_dropout, nn.Dropout) and self.training else 0.0
+            attn_output = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=attention_mask,
+                dropout_p=dropout_p,
+                is_causal=True if attention_mask is None else False
+            )
+            attention_weights = None
+        else:
+            attn_scores = (q @ k.transpose(-2, -1)) * self.scale
+            
+            if self.use_alibi:
+                position_bias = self.alibi_slopes.to(x.device) * torch.arange(
+                    seq_len_k, device=x.device
+                ).view(1, 1, -1)
+                attn_scores = attn_scores + position_bias
+            
+            if self.sliding_window is not None:
+                attn_scores = self._apply_sliding_window_mask(attn_scores, seq_len_k)
+            
+            if attention_mask is not None:
+                if attention_mask.dim() == 2:
+                    attention_mask = attention_mask[:, None, None, :]
+                if attention_mask.dtype != torch.float:
+                    # 假设传入的是 1(Keep)/0(Mask)
+                    extended_mask = (1.0 - attention_mask) * torch.finfo(attn_scores.dtype).min
+                else:
+                    # 假设传入的已经是加性 mask (0/-inf)
+                    extended_mask = attention_mask
+            
+                attn_scores = attn_scores + extended_mask
+            
+            is_causal = seq_len_k > 1
+            if is_causal:
+                causal_mask = torch.triu(
+                torch.ones(seq_len_k, seq_len_k, device=x.device, dtype=torch.bool),
+                diagonal=1
+                )
+                causal_mask = causal_mask[-q.shape[2]:, :]#还没懂
+                attn_scores = attn_scores.masked_fill(causal_mask, float('-inf'))
+            
+            attention_weights = F.softmax(attn_scores, dim=-1, dtype=torch.float32).to(q.dtype)
+            attention_weights = self.attn_dropout(attention_weights)
+            
+            attn_output = attention_weights @ v
+        
+        attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, -1)
+        output = self.resid_dropout(self.o_proj(attn_output))
+        
+        return output, present_kv, attention_weights if output_attentions else None
+
+class OptimizedTransformerBlock(nn.Module):
+    """优化的Transformer块"""
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        n_kv_heads: Optional[int] = None,
+        head_dim: Optional[int] = None,
+        dropout: float = 0.0,
+        attn_dropout: float = 0.0,
+        use_moe: bool = False,
+        num_experts: int = 8,
+        moe_top_k: int = 2,
+        use_adapter: bool = False,
+        adapter_dim: int = 64,
+        use_lora: bool = False,
+        lora_rank: int = 8,
+        use_parallel_residual: bool = False,
+        norm_eps: float = 1e-6,
+        sliding_window: Optional[int] = None,
+        ffn_dim_multiplier: Optional[float] = None,
+        layer_idx: int = 0
+    ):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.use_moe = use_moe
+        self.use_adapter = use_adapter
+        self.use_parallel_residual = use_parallel_residual
+        
+        self.attention = GroupedQueryAttention(
+            dim=dim,
+            n_heads=n_heads,
+            n_kv_heads=n_kv_heads,
+            head_dim=head_dim,
+            dropout=dropout,
+            attn_dropout=attn_dropout,
+            use_lora=use_lora,
+            lora_rank=lora_rank,
+            sliding_window=sliding_window,
+            rope_scaling_type="yarn"
+        )
+        
+        if use_moe:
+            self.ffn = MixtureOfExperts(
+                dim=dim,
+                num_experts=num_experts,
+                top_k=moe_top_k,
+                dropout=dropout,
+                ffn_dim_multiplier=ffn_dim_multiplier
+            )
+        else:
+            self.ffn = SwiGLU(
+                dim=dim,
+                dropout=dropout,
+                ffn_dim_multiplier=ffn_dim_multiplier
+            )
+        
+        if use_adapter:
+            self.adapter = AdapterLayer(dim, adapter_dim, dropout)
+        
+        self.attention_norm = RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm = RMSNorm(dim, eps=norm_eps)
+        
+        self.moe_aux_loss = torch.tensor(0.0)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+        past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        output_attentions: bool = False
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[torch.Tensor]]:
+        """前向传播"""
+        
+        attn_out, present_kv, attn_weights = self.attention(
+            self.attention_norm(x),
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            use_cache=use_cache,
+            past_kv=past_kv,
+            output_attentions=output_attentions
+        )
+        
+        if self.use_parallel_residual:
+            ffn_input = self.ffn_norm(x)
+            
+            if self.use_moe:
+                ffn_out, aux_loss = self.ffn(ffn_input)
+                self.moe_aux_loss = aux_loss
+            else:
+                ffn_out = self.ffn(ffn_input)
+                self.moe_aux_loss = torch.tensor(0.0, device=x.device)
+            
+            x = x + attn_out + ffn_out
+        else:
+            x = x + attn_out
+            
+            if self.use_adapter:
+                x = self.adapter(x)
+            
+            ffn_input = self.ffn_norm(x)
+            if self.use_moe:
+                ffn_out, aux_loss = self.ffn(ffn_input)
+                x = x + ffn_out
+                self.moe_aux_loss = aux_loss
+            else:
+                x = x + self.ffn(ffn_input)
+                self.moe_aux_loss = torch.tensor(0.0, device=x.device)
+        
+        return x, present_kv, attn_weights