cortexnet 3.2.1__py3-none-any.whl

cortexnet/config.py

@@ -0,0 +1,234 @@
+"""
+CortexNet 配置模块 (Configuration Module)
+
+定义 CortexNet 架构和训练的所有超参数。
+使用 dataclass 实现类型安全的配置管理。
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CortexNetConfig:
+    """CortexNet 架构配置。
+
+    控制模型的规模、结构和行为。所有 CortexNet 模块共享此配置。
+
+    Attributes:
+        vocab_size: 词汇表大小
+        hidden_size: 隐藏层维度 (d_model)
+        num_layers: CortexBlock 层数
+        num_heads: 注意力头数
+        num_scales: SSM 多尺度数
+        ssm_state_size: SSM 状态维度
+        ssm_expand_factor: SSM 内部扩展倍数
+        top_k_ratio: 稀疏注意力 top-k 比例
+        attention_k_mode: top-k 计算模式 ("ratio", "sqrt", "log", "fixed")
+        max_seq_len: 最大序列长度
+        rope_theta: RoPE 频率基数
+        dropout: Dropout 比率
+        memory_dim: 记忆模块维度
+        memory_decay_init: 记忆衰减初始值
+        expert_ff_dim: 专家 FFN 中间维度
+        num_experts: 专家总数
+        num_active_experts: 每个 token 激活的专家数
+        moe_aux_loss_weight: MoE 辅助损失权重
+        moe_capacity_factor: MoE 容量因子
+        norm_eps: 归一化层 epsilon
+    """
+
+    # ═══ 基础架构参数 ═══
+    vocab_size: int = 32000
+    hidden_size: int = 512
+    num_layers: int = 4
+    num_heads: int = 8
+    max_seq_len: int = 8192
+    dropout: float = 0.0
+    norm_eps: float = 1e-6
+
+    # ═══ SSM 参数 ═══
+    num_scales: int = 4
+    ssm_state_size: int = 16
+    ssm_expand_factor: int = 2
+
+    # ═══ 注意力参数 ═══
+    top_k_ratio: float = 0.25
+    attention_k_mode: str = "ratio"
+    rope_theta: float = 10000.0
+    sliding_window_size: int = 0
+
+    # ═══ 记忆参数 ═══
+    memory_dim: int = 64
+    memory_decay_init: float = 0.95
+    episodic_slots: int = 32
+    semantic_slots: int = 64
+
+    # ═══ MoE 参数 ═══
+    expert_ff_dim: int = 1024
+    num_experts: int = 8
+    num_active_experts: int = 2
+    moe_aux_loss_weight: float = 0.02
+    moe_capacity_factor: float = 1.25  # MoE 路由容量因子
+
+    # ═══ V2/V3 扩展参数 ═══
+    graph_neighbors: int = 16
+    graph_iterations: int = 2
+    num_task_modes: int = 4
+    num_counterfactuals: int = 4
+    num_agents: int = 4
+    use_gradient_checkpointing: bool = False
+    use_mixture_of_depths: bool = False
+    mod_capacity: float = 0.5
+    causal_top_k_ratio: float = 0.25  # 因果推理干预注意力的 top-k 比例
+
+    # ═══ 适配器参数（新增：用于开源模型加载） ═══
+    model_type: str = "cortexnet"  # 源模型类型标识
+    source_model_path: Optional[str] = None  # HuggingFace 模型路径
+    auto_calibrate: bool = True  # 是否在加载后自动校准
+    intermediate_size: int = 0  # 源模型中间层维度（如有）
+    num_kv_heads: int = 0  # GQA 中的 KV 头数（0 = 与 num_heads 相同）
+    use_qk_norm: bool = False  # 注意力中是否启用 per-head Q/K RMSNorm（Qwen2/3）
+    rope_scaling: Optional[dict] = None  # RoPE 缩放配置
+    tie_word_embeddings: bool = True  # 是否绑定嵌入和输出权重
+    compatibility_mode: bool = False  # 兼容大模型的轻量 V3 路径
+    lite: bool = True  # Lite 模式：仅 SSM+Attention+Memory+FFN（参数减少 60-70%）
+    ssm_decode_after: int = 0  # SSM 纯解码阈值（>0 时，超过此 token 数后仅用 SSM）
+    expand_gqa_weights: bool = True  # 是否将 KV 投影扩展到全头（旧结构兼容）
+    compat_ssm_rank: int = 256  # 兼容模式 SSM 低秩维度（控制增量参数量）
+    fusion_long_context_threshold: int = 2048  # 长序列时启用 SSM 最低占比的阈值
+    fusion_long_context_ssm_ratio: float = 0.35  # 长序列时 SSM 最低占比
+    mapped_cache_enabled: bool = False  # 是否启用映射后权重缓存（可选，默认关闭）
+    mapped_cache_dir: Optional[str] = None  # 映射缓存目录（None 使用默认目录）
+    mapped_cache_force_refresh: bool = False  # 是否强制忽略缓存并重建
+    mapped_cache_auto_enable_with_lazy: bool = True  # lazy_device_load 时自动启用映射缓存
+    mapped_cache_fast_init_on_hit: bool = True  # 命中映射缓存时跳过额外自定义重初始化
+    lazy_device_load: bool = False  # 是否启用惰性上设备（from_pretrained 快速返回）
+    lazy_cpu_fallback: bool = True  # 惰性阶段是否允许 CPU 兜底推理
+    lazy_background_warmup: bool = True  # 首次推理后是否后台预热到目标设备
+    lazy_start_warmup_on_load: bool = True  # from_pretrained 返回前立即启动后台预热线程
+    lazy_disable_on_cache_hit: bool = True  # 命中映射缓存时自动关闭 lazy，优先首 token 体验
+
+    def __post_init__(self):
+        # ═══ 默认值推导 ═══
+        if self.num_kv_heads == 0:
+            self.num_kv_heads = self.num_heads
+        if self.intermediate_size == 0:
+            self.intermediate_size = self.expert_ff_dim
+
+        # ═══ 参数验证 ═══
+        self._validate()
+
+    def _validate(self):
+        """全面的参数验证，确保配置合法且防止隐蔽错误。"""
+        # 基础维度检查
+        if self.hidden_size % self.num_heads != 0:
+            raise ValueError(
+                f"hidden_size ({self.hidden_size}) 必须能被 num_heads ({self.num_heads}) 整除"
+            )
+        if self.num_kv_heads > self.num_heads:
+            raise ValueError(
+                f"num_kv_heads ({self.num_kv_heads}) 不能超过 num_heads ({self.num_heads})"
+            )
+        if self.num_heads % self.num_kv_heads != 0:
+            raise ValueError(
+                f"num_heads ({self.num_heads}) 必须能被 num_kv_heads ({self.num_kv_heads}) 整除 (GQA 要求)"
+            )
+
+        # MoE 参数检查
+        if self.num_active_experts > self.num_experts:
+            raise ValueError(
+                f"num_active_experts ({self.num_active_experts}) 不能超过 num_experts ({self.num_experts})"
+            )
+        if self.moe_capacity_factor <= 0:
+            raise ValueError(
+                f"moe_capacity_factor ({self.moe_capacity_factor}) 必须为正数"
+            )
+
+        # 注意力参数检查
+        if not (0 < self.top_k_ratio <= 1.0):
+            raise ValueError(
+                f"top_k_ratio ({self.top_k_ratio}) 必须在 (0, 1.0] 范围内"
+            )
+        _valid_k_modes = {"ratio", "sqrt", "log", "fixed"}
+        if self.attention_k_mode not in _valid_k_modes:
+            raise ValueError(
+                f"attention_k_mode ('{self.attention_k_mode}') 必须是 {_valid_k_modes} 之一"
+            )
+
+        # 因果推理参数检查
+        if not (0 < self.causal_top_k_ratio <= 1.0):
+            raise ValueError(
+                f"causal_top_k_ratio ({self.causal_top_k_ratio}) 必须在 (0, 1.0] 范围内"
+            )
+
+        # 正数检查
+        if self.hidden_size <= 0:
+            raise ValueError(f"hidden_size ({self.hidden_size}) 必须为正整数")
+        if self.num_layers <= 0:
+            raise ValueError(f"num_layers ({self.num_layers}) 必须为正整数")
+        if self.max_seq_len <= 0:
+            raise ValueError(f"max_seq_len ({self.max_seq_len}) 必须为正整数")
+        if self.vocab_size <= 0:
+            raise ValueError(f"vocab_size ({self.vocab_size}) 必须为正整数")
+
+        # 范围检查
+        if not (0 <= self.dropout < 1.0):
+            raise ValueError(
+                f"dropout ({self.dropout}) 必须在 [0, 1.0) 范围内"
+            )
+
+        # 混合精度检查
+        if self.fusion_long_context_ssm_ratio < 0 or self.fusion_long_context_ssm_ratio > 1:
+            raise ValueError(
+                f"fusion_long_context_ssm_ratio ({self.fusion_long_context_ssm_ratio}) 必须在 [0, 1] 范围内"
+            )
+
+        # 软警告
+        if self.hidden_size < 64:
+            logger.warning(f"hidden_size={self.hidden_size} 过小，可能影响模型表达能力")
+        if self.num_experts > 1 and self.num_active_experts < 1:
+            logger.warning("num_active_experts < 1，MoE 路由可能无效")
+
+    @classmethod
+    def from_dict(cls, d: dict) -> "CortexNetConfig":
+        """从字典创建配置，忽略未知字段。"""
+        valid_fields = {f.name for f in cls.__dataclass_fields__.values()}
+        filtered = {k: v for k, v in d.items() if k in valid_fields}
+        return cls(**filtered)
+
+
+@dataclass
+class TrainingConfig:
+    """训练超参数配置。
+
+    Attributes:
+        learning_rate: 学习率
+        weight_decay: 权重衰减
+        num_epochs: 训练轮数
+        batch_size: 批大小
+        gradient_accumulation_steps: 梯度累积步数
+        max_grad_norm: 梯度裁剪最大范数
+        warmup_steps: 学习率预热步数
+        eval_interval: 评估间隔步数
+        save_interval: 保存间隔步数
+        mixed_precision: 混合精度类型 ("no", "fp16", "bf16")
+    """
+
+    learning_rate: float = 3e-4
+    weight_decay: float = 0.01
+    num_epochs: int = 3
+    batch_size: int = 8
+    gradient_accumulation_steps: int = 1
+    max_grad_norm: float = 1.0
+    warmup_steps: int = 100
+    eval_interval: int = 500
+    save_interval: int = 1000
+    mixed_precision: str = "no"
+    seed: int = 42
+    log_interval: int = 10

cortexnet/continual_learning.py

@@ -0,0 +1,256 @@
+"""
+连续学习模块 (Continual Learning Module)
+
+核心创新：
+  解决神经网络在学习新任务时"灾难性遗忘"旧知识的问题。
+  实现两种互补的反遗忘策略：
+
+  1. 弹性权重巩固 (EWC, Elastic Weight Consolidation)
+     - 计算 Fisher 信息矩阵，识别对旧任务重要的参数
+     - 在训练新任务时，惩罚对这些重要参数的大幅修改
+     - 效果：旧知识被"保护"，新知识在不重要的参数空间中学习
+
+  2. 渐进式记忆回放 (Progressive Memory Replay)
+     - 维护一个经验缓冲区，存储旧任务的关键样本
+     - 在训练新任务时，混合旧样本进行回放
+     - 效果：周期性"复习"旧知识，防止遗忘
+
+  类比人类学习：
+    - EWC ≈ 大脑对重要神经连接的保护机制
+    - 回放 ≈ 人类的睡眠巩固和复习
+"""
+
+import torch
+import torch.nn as nn
+from typing import Dict, Optional
+
+
+class ElasticWeightConsolidation:
+    """弹性权重巩固 (EWC)。
+
+    通过 Fisher 信息矩阵估计每个参数对已学任务的重要性，
+    在学习新任务时添加正则化项以保护重要参数。
+
+    数学原理：
+        L_total = L_new_task + λ · Σ_i F_i · (θ_i - θ*_i)²
+
+    其中 F_i 是第 i 个参数的 Fisher 信息（重要性）,
+    θ*_i 是在旧任务上的最优参数值。
+
+    Args:
+        model: CortexNet 模型
+        lambda_ewc: EWC 正则化强度
+    """
+
+    def __init__(self, model: nn.Module, lambda_ewc: float = 1000.0):
+        self.model = model
+        self.lambda_ewc = lambda_ewc
+        self.fisher: Dict[str, torch.Tensor] = {}
+        self.optimal_params: Dict[str, torch.Tensor] = {}
+        self._consolidated = False
+        self._num_consolidations = 0
+
+    def consolidate(self, dataloader, num_samples: int = 200):
+        """计算 Fisher 信息并保存当前最优参数。
+
+        在完成一个任务的训练后调用，将当前知识"巩固"。
+
+        Args:
+            dataloader: 当前任务的数据加载器
+            num_samples: 用于估计 Fisher 的样本数
+        """
+        self.model.train()  # 需要梯度计算
+        fisher = {
+            n: torch.zeros_like(p)
+            for n, p in self.model.named_parameters()
+            if p.requires_grad
+        }
+
+        count = 0
+        for batch in dataloader:
+            if count >= num_samples:
+                break
+
+            input_ids, labels = batch
+            input_ids = input_ids.to(
+                next(self.model.parameters()).device
+            )
+            labels = labels.to(next(self.model.parameters()).device)
+
+            self.model.zero_grad()
+            output = self.model(input_ids, labels=labels)
+
+            # 使用对数似然的梯度平方作为 Fisher 近似
+            output["loss"].backward()
+
+            for n, p in self.model.named_parameters():
+                if p.requires_grad and p.grad is not None:
+                    fisher[n] += p.grad.pow(2)
+
+            count += 1
+
+        # 平均 Fisher 信息
+        for n in fisher:
+            fisher[n] /= max(count, 1)
+
+        # 如果之前已巩固，使用在线更新（累积 Fisher）
+        if self._consolidated:
+            for n in fisher:
+                if n in self.fisher:
+                    fisher[n] = (
+                        self.fisher[n] * self._num_consolidations
+                        + fisher[n]
+                    ) / (self._num_consolidations + 1)
+
+        self.fisher = fisher
+        self.optimal_params = {
+            n: p.clone().detach()
+            for n, p in self.model.named_parameters()
+            if p.requires_grad
+        }
+        self._consolidated = True
+        self._num_consolidations += 1
+
+    def penalty(self) -> torch.Tensor:
+        """计算 EWC 正则化惩罚。
+
+        在训练新任务时添加到损失函数中。
+
+        Returns:
+            penalty: 标量张量
+        """
+        if not self._consolidated:
+            device = next(self.model.parameters()).device
+            return torch.tensor(0.0, device=device)
+
+        loss = torch.tensor(
+            0.0, device=next(self.model.parameters()).device
+        )
+        for n, p in self.model.named_parameters():
+            if n in self.fisher and p.requires_grad:
+                loss = loss + (
+                    self.fisher[n] * (p - self.optimal_params[n]).pow(2)
+                ).sum()
+
+        return self.lambda_ewc * loss
+
+
+class ProgressiveMemoryReplay:
+    """渐进式记忆回放。
+
+    维护一个经验缓冲区，在训练新任务时混合旧样本。
+    使用 reservoir sampling 确保缓冲区均匀覆盖所有旧任务。
+
+    Args:
+        buffer_size: 缓冲区最大样本数
+        replay_ratio: 每批中旧样本的比例
+    """
+
+    def __init__(self, buffer_size: int = 5000, replay_ratio: float = 0.3):
+        self.buffer_size = buffer_size
+        self.replay_ratio = replay_ratio
+        self.buffer_inputs = []
+        self.buffer_labels = []
+        self._count = 0
+
+    def add_samples(
+        self, input_ids: torch.Tensor, labels: torch.Tensor
+    ):
+        """向缓冲区添加样本（使用 reservoir sampling）。"""
+        batch_size = input_ids.shape[0]
+
+        for i in range(batch_size):
+            if len(self.buffer_inputs) < self.buffer_size:
+                self.buffer_inputs.append(input_ids[i].cpu())
+                self.buffer_labels.append(labels[i].cpu())
+            else:
+                # Reservoir sampling
+                idx = torch.randint(0, self._count + 1, (1,)).item()
+                if idx < self.buffer_size:
+                    self.buffer_inputs[idx] = input_ids[i].cpu()
+                    self.buffer_labels[idx] = labels[i].cpu()
+            self._count += 1
+
+    def get_replay_batch(
+        self, batch_size: int, device: torch.device
+    ) -> Optional[tuple]:
+        """获取回放批次。"""
+        if len(self.buffer_inputs) == 0:
+            return None
+
+        num_replay = max(1, int(batch_size * self.replay_ratio))
+        num_replay = min(num_replay, len(self.buffer_inputs))
+
+        indices = torch.randperm(len(self.buffer_inputs))[:num_replay]
+
+        replay_inputs = torch.stack(
+            [self.buffer_inputs[i] for i in indices]
+        ).to(device)
+        replay_labels = torch.stack(
+            [self.buffer_labels[i] for i in indices]
+        ).to(device)
+
+        return replay_inputs, replay_labels
+
+    @property
+    def size(self) -> int:
+        return len(self.buffer_inputs)
+
+
+class ContinualLearningManager:
+    """连续学习管理器：整合 EWC + 记忆回放。
+
+    提供统一接口管理连续学习的各个组件。
+
+    使用方法：
+        manager = ContinualLearningManager(model)
+
+        # 任务 1 训练
+        for batch in task1_loader:
+            loss = model(batch)
+            loss = loss + manager.get_regularization_loss()
+            loss.backward()
+
+        # 巩固任务 1 的知识
+        manager.consolidate_task(task1_loader)
+
+        # 任务 2 训练（自动保护任务 1 的知识）
+        for batch in task2_loader:
+            loss = model(batch)
+            loss = loss + manager.get_regularization_loss()
+            replay = manager.get_replay_batch(batch_size)
+            if replay:
+                replay_loss = model(replay)
+                loss = loss + replay_loss * 0.5
+            loss.backward()
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        lambda_ewc: float = 1000.0,
+        buffer_size: int = 5000,
+        replay_ratio: float = 0.3,
+    ):
+        self.ewc = ElasticWeightConsolidation(model, lambda_ewc)
+        self.replay = ProgressiveMemoryReplay(buffer_size, replay_ratio)
+        self.task_count = 0
+
+    def consolidate_task(self, dataloader, num_samples: int = 200):
+        """巩固当前任务的知识。"""
+        self.ewc.consolidate(dataloader, num_samples)
+        self.task_count += 1
+
+    def get_regularization_loss(self) -> torch.Tensor:
+        """获取防遗忘正则化损失。"""
+        return self.ewc.penalty()
+
+    def add_experience(
+        self, input_ids: torch.Tensor, labels: torch.Tensor
+    ):
+        """记录训练样本到经验缓冲区。"""
+        self.replay.add_samples(input_ids, labels)
+
+    def get_replay_batch(self, batch_size: int, device: torch.device):
+        """获取回放样本。"""
+        return self.replay.get_replay_batch(batch_size, device)