cortexnet 3.2.1__py3-none-any.whl

cortexnet/cortex_block_lite.py

@@ -0,0 +1,221 @@
+"""
+CortexBlockLite — 精简版 CortexNet 块
+
+仅保留三条核心路径 + SwiGLU FFN，参数量与 Transformer 相当：
+  1. Multi-Scale SSM  — O(1) 解码，长序列优势
+  2. Selective Attention — 兼容预训练权重
+  3. Synaptic Memory   — 快速权重系统
+
+相比 CortexBlockV3 移除了：
+  - MultiAgentSystem (23.7%)
+  - CollaborativeMoE → 直接 SwiGLU FFN
+  - CausalReasoning (8.0%)
+  - TaskAdaptiveController (6.3%)
+  - GraphReasoning (5.2%)
+  - MetaLearningAdapter (4.2%)
+  - AdversarialShield (2.1%)
+
+典型参数节省：60-70%
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple, Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+
+try:
+    from .blocks import RMSNorm
+    from .config import CortexNetConfig
+    from .ssm import MultiScaleSSM
+    from .attention import SelectiveSparseAttention
+    from .memory import SynapticMemory
+except ImportError:
+    from cortexnet.blocks import RMSNorm
+    from cortexnet.config import CortexNetConfig
+    from cortexnet.ssm import MultiScaleSSM
+    from cortexnet.attention import SelectiveSparseAttention
+    from cortexnet.memory import SynapticMemory
+
+
+class SwiGLU_FFN(nn.Module):
+    """SwiGLU 前馈网络（与 LLaMA/Qwen 兼容）。
+
+    参数量 = 3 × d_model × intermediate_size（无 bias）。
+    """
+
+    def __init__(self, d_model: int, intermediate_size: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(d_model, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(d_model, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, d_model, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class CortexBlockLite(nn.Module):
+    """精简 CortexNet 块：SSM + Attention + Memory + SwiGLU FFN。
+
+    核心设计：
+      - 三条并行路径（SSM, Attention, Memory）通过自适应融合门控合并
+      - SwiGLU FFN 替代重量级 MoE（参数减少 80%+）
+      - 支持 KV cache 推理和梯度检查点
+      - SSM 优先模式：长序列生成时可跳过 Attention
+
+    Args:
+        config: CortexNetConfig 配置对象
+        layer_idx: 当前层索引（用于 RoPE 偏移计算）
+    """
+
+    def __init__(self, config: CortexNetConfig, layer_idx: int = 0):
+        super().__init__()
+        D = config.hidden_size
+        self.layer_idx = layer_idx
+        self.use_checkpoint = config.use_gradient_checkpointing
+
+        # ═══ 归一化 ═══
+        self.norm1 = RMSNorm(D, eps=config.norm_eps)
+        self.norm2 = RMSNorm(D, eps=config.norm_eps)
+
+        # ═══ 路径 1: 低秩 SSM（O(1) 解码，O(d×rank) 参数） ═══
+        # Lite 模式：SSM 在降维空间运行，大幅减少参数量
+        ssm_rank = getattr(config, 'compat_ssm_rank', max(64, D // 4))
+        ssm_rank = min(ssm_rank, D)  # 不超过 D
+        self.ssm_rank = ssm_rank
+        self.ssm_down = nn.Linear(D, ssm_rank, bias=False) if ssm_rank < D else nn.Identity()
+        self.ssm = MultiScaleSSM(
+            d_model=ssm_rank,
+            num_scales=1,
+            state_size=config.ssm_state_size,
+            expand_factor=1,
+        )
+        self.ssm_up = nn.Linear(ssm_rank, D, bias=False) if ssm_rank < D else nn.Identity()
+
+        # ═══ 路径 2: Selective Sparse Attention ═══
+        self.attention = SelectiveSparseAttention(
+            d_model=D,
+            num_heads=config.num_heads,
+            num_kv_heads=config.num_kv_heads,
+            top_k_ratio=config.top_k_ratio,
+            k_mode=config.attention_k_mode,
+            max_seq_len=config.max_seq_len,
+            rope_theta=config.rope_theta,
+            dropout=config.dropout,
+            sliding_window_size=config.sliding_window_size,
+            use_qk_norm=getattr(config, "use_qk_norm", False),
+            norm_eps=config.norm_eps,
+        )
+
+        # ═══ 路径 3: Synaptic Memory ═══
+        # 大模型用较小 memory_dim 避免参数膨胀
+        mem_dim = min(config.memory_dim, max(32, D // 16))
+        self.memory = SynapticMemory(
+            d_model=D,
+            memory_dim=mem_dim,
+            decay_init=config.memory_decay_init,
+        )
+
+        # ═══ 轻量级三路径融合（3 个标量权重） ═══
+        # 替代 AdaptiveFusionGate 的全连接网络，参数从 O(d²) 降至 3
+        self.fusion_weights = nn.Parameter(torch.ones(3) / 3)
+
+        # ═══ SwiGLU FFN（替代 MoE） ═══
+        ff_dim = config.intermediate_size if config.intermediate_size > 0 else config.expert_ff_dim
+        self.ffn = SwiGLU_FFN(D, ff_dim)
+
+        # ═══ Dropout ═══
+        self.dropout = nn.Dropout(config.dropout) if config.dropout > 0 else nn.Identity()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        past_cache: Optional[Tuple[Any, Any, Any]] = None,
+        use_cache: bool = False,
+        ssm_only: bool = False,
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: (B, L, D)
+            past_cache: (past_attn_kv, past_ssm_state, past_mem_state)
+            use_cache: 是否返回缓存
+            ssm_only: True → 纯 SSM 模式（跳过 Attention，极速解码）
+
+        Returns:
+            use_cache=False: output (B, L, D)
+            use_cache=True: (output, (attn_cache, ssm_cache, mem_cache))
+        """
+        residual = x
+        x_norm = self.norm1(x)
+
+        # 解包缓存
+        past_attn = past_ssm = past_mem = None
+        if past_cache is not None:
+            past_attn, past_ssm, past_mem = past_cache
+
+        # ═══ SSM 路径（低秩空间） ═══
+        ssm_input = self.ssm_down(x_norm)
+        if use_cache:
+            ssm_out_lr, new_ssm_cache = self.ssm(
+                ssm_input, past_state=past_ssm, use_cache=True,
+            )
+        else:
+            ssm_out_lr = self._maybe_checkpoint(self.ssm, ssm_input)
+            new_ssm_cache = None
+        ssm_out = self.ssm_up(ssm_out_lr)
+
+        if ssm_only:
+            # 纯 SSM 模式：跳过 Attention 和 Memory
+            x = residual + self.dropout(ssm_out)
+        else:
+            # ═══ Attention 路径 ═══
+            if use_cache:
+                attn_out, new_attn_cache = self.attention(
+                    x_norm, past_key_value=past_attn, use_cache=True,
+                )
+            else:
+                attn_out = self._maybe_checkpoint(self.attention, x_norm)
+                new_attn_cache = None
+
+            # ═══ Memory 路径 ═══
+            if use_cache:
+                past_memory = past_mem[0] if isinstance(past_mem, tuple) else past_mem
+                past_z = past_mem[1] if isinstance(past_mem, tuple) else None
+                mem_out, new_mem, new_z = self.memory(
+                    x_norm, past_memory=past_memory, past_z=past_z, use_cache=True,
+                )
+                new_mem_cache = (new_mem, new_z)
+            else:
+                mem_out = self._maybe_checkpoint(self.memory, x_norm)
+                new_mem_cache = None
+
+            # ═══ 轻量融合 ═══
+            w = torch.softmax(self.fusion_weights, dim=0)
+            fused = w[0] * ssm_out + w[1] * attn_out + w[2] * mem_out
+            x = residual + self.dropout(fused)
+
+        # ═══ FFN ═══
+        residual = x
+        x = residual + self.dropout(self.ffn(self.norm2(x)))
+
+        if use_cache:
+            cache = (
+                new_attn_cache if not ssm_only else past_attn,
+                new_ssm_cache,
+                new_mem_cache if not ssm_only else past_mem,
+            )
+            return x, cache
+        return x
+
+    def _maybe_checkpoint(self, module, *args):
+        """条件梯度检查点。"""
+        if self.use_checkpoint and self.training:
+            return checkpoint.checkpoint(module, *args, use_reentrant=False)
+        return module(*args)
+
+    def get_aux_loss(self) -> torch.Tensor:
+        """Lite 模式无 MoE，返回零损失。"""
+        return self.fusion_weights.new_zeros(())

cortexnet/distributed.py

@@ -0,0 +1,213 @@
+"""
+CortexNet 分布式训练支持 (Distributed Training)
+
+提供 FSDP (Fully Sharded Data Parallel) 策略和基础分布式工具，
+为大规模训练做准备。
+
+用法:
+    from cortexnet.distributed import setup_distributed, wrap_fsdp, cleanup_distributed
+
+    setup_distributed()
+    model = CortexNet(config)
+    model = wrap_fsdp(model, config)
+    # ... training loop ...
+    cleanup_distributed()
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Optional, Any
+
+import torch
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+
+def setup_distributed(backend: str = "nccl") -> bool:
+    """初始化分布式训练环境。
+
+    自动检测是否在分布式环境中运行（通过环境变量），
+    如果是则初始化 process group。
+
+    Args:
+        backend: 通信后端 ("nccl" for GPU, "gloo" for CPU)
+
+    Returns:
+        True 如果成功初始化分布式环境
+    """
+    if torch.distributed.is_initialized():
+        logger.info("Distributed already initialized")
+        return True
+
+    # 检查环境变量
+    rank = os.environ.get("RANK")
+    world_size = os.environ.get("WORLD_SIZE")
+    local_rank = os.environ.get("LOCAL_RANK")
+
+    if rank is None or world_size is None:
+        logger.info("Not in distributed environment, skipping initialization")
+        return False
+
+    rank = int(rank)
+    world_size = int(world_size)
+    local_rank = int(local_rank) if local_rank is not None else rank
+
+    # 自动选择后端
+    if backend == "nccl" and not torch.cuda.is_available():
+        backend = "gloo"
+        logger.info("CUDA not available, falling back to gloo backend")
+
+    # 设置 GPU
+    if torch.cuda.is_available():
+        torch.cuda.set_device(local_rank)
+
+    torch.distributed.init_process_group(
+        backend=backend,
+        rank=rank,
+        world_size=world_size,
+    )
+
+    logger.info(
+        f"Distributed initialized: rank={rank}/{world_size}, "
+        f"local_rank={local_rank}, backend={backend}"
+    )
+    return True
+
+
+def cleanup_distributed():
+    """清理分布式训练环境。"""
+    if torch.distributed.is_initialized():
+        torch.distributed.destroy_process_group()
+        logger.info("Distributed process group destroyed")
+
+
+def get_rank() -> int:
+    """获取当前进程的 rank（非分布式返回 0）。"""
+    if torch.distributed.is_initialized():
+        return torch.distributed.get_rank()
+    return 0
+
+
+def get_world_size() -> int:
+    """获取总进程数（非分布式返回 1）。"""
+    if torch.distributed.is_initialized():
+        return torch.distributed.get_world_size()
+    return 1
+
+
+def is_main_process() -> bool:
+    """是否为主进程。"""
+    return get_rank() == 0
+
+
+def wrap_fsdp(
+    model: nn.Module,
+    config: Optional[Any] = None,
+    mixed_precision: Optional[str] = None,
+    activation_checkpointing: bool = False,
+) -> nn.Module:
+    """使用 FSDP 包装模型。
+
+    自动为 CortexNet 的不同组件选择合适的分片策略。
+
+    Args:
+        model: CortexNet 模型实例
+        config: CortexNetConfig（可选，用于读取配置）
+        mixed_precision: "fp16", "bf16", 或 None
+        activation_checkpointing: 是否启用激活检查点
+
+    Returns:
+        FSDP 包装后的模型
+    """
+    if not torch.distributed.is_initialized():
+        logger.warning("Distributed not initialized, returning model as-is")
+        return model
+
+    try:
+        from torch.distributed.fsdp import (
+            FullyShardedDataParallel as FSDP,
+            MixedPrecision,
+            ShardingStrategy,
+        )
+        from torch.distributed.fsdp.wrap import (
+            size_based_auto_wrap_policy,
+        )
+    except ImportError:
+        logger.error("FSDP not available in this PyTorch version")
+        return model
+
+    # 混合精度策略
+    mp_policy = None
+    if mixed_precision == "fp16":
+        mp_policy = MixedPrecision(
+            param_dtype=torch.float16,
+            reduce_dtype=torch.float16,
+            buffer_dtype=torch.float16,
+        )
+    elif mixed_precision == "bf16":
+        mp_policy = MixedPrecision(
+            param_dtype=torch.bfloat16,
+            reduce_dtype=torch.bfloat16,
+            buffer_dtype=torch.bfloat16,
+        )
+
+    # 自动分片策略：按参数数量决定分片粒度
+    min_params = 1_000_000  # 超过 1M 参数的子模块独立分片
+    auto_wrap_policy = size_based_auto_wrap_policy(min_num_params=min_params)
+
+    model = FSDP(
+        model,
+        auto_wrap_policy=auto_wrap_policy,
+        mixed_precision=mp_policy,
+        sharding_strategy=ShardingStrategy.FULL_SHARD,
+        device_id=torch.cuda.current_device() if torch.cuda.is_available() else None,
+    )
+
+    # 激活检查点
+    if activation_checkpointing:
+        try:
+            from torch.distributed.fsdp import apply_activation_checkpointing
+            # 对所有 CortexBlock 启用激活检查点
+            apply_activation_checkpointing(model)
+            logger.info("Activation checkpointing enabled")
+        except ImportError:
+            logger.warning("Activation checkpointing not available")
+
+    logger.info(f"Model wrapped with FSDP (mixed_precision={mixed_precision})")
+    return model
+
+
+def wrap_ddp(
+    model: nn.Module,
+    find_unused_parameters: bool = False,
+) -> nn.Module:
+    """使用 DDP (DistributedDataParallel) 包装模型。
+
+    对于不需要 FSDP 级别分片的场景，使用更简单的 DDP。
+
+    Args:
+        model: 模型
+        find_unused_parameters: 是否查找未使用参数
+
+    Returns:
+        DDP 包装后的模型
+    """
+    if not torch.distributed.is_initialized():
+        logger.warning("Distributed not initialized, returning model as-is")
+        return model
+
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    if torch.cuda.is_available():
+        model = model.to(local_rank)
+
+    model = nn.parallel.DistributedDataParallel(
+        model,
+        device_ids=[local_rank] if torch.cuda.is_available() else None,
+        find_unused_parameters=find_unused_parameters,
+    )
+
+    logger.info(f"Model wrapped with DDP (device={local_rank})")
+    return model

cortexnet/graph_reasoning.py

@@ -0,0 +1,207 @@
+"""
+图推理模块 (Graph Reasoning Module)
+
+核心创新：
+  将 token 序列视为图结构，通过消息传递神经网络 (MPNN)
+  进行多步关系推理。与注意力机制的关键区别：
+
+  ┌─────────────────────────────────────────────────────┐
+  │           注意力 vs 图推理                           │
+  ├─────────────────────────────────────────────────────┤
+  │ 注意力:                                             │
+  │  • 单步信息聚合                                     │
+  │  • 全局/稀疏的软关联                                │
+  │  • 无显式关系建模                                   │
+  │                                                     │
+  │ 图推理:                                             │
+  │  • 多步迭代推理（信息逐步传播）                     │
+  │  • 结构化的邻域关系                                 │
+  │  • 显式边特征建模 token 间关系                      │
+  │  • 门控线性单元保持推理记忆                         │
+  └─────────────────────────────────────────────────────┘
+
+  通过 K 步消息传递，信息可以在图中传播 K 跳距离，
+  实现类似 "思维链" 的逐步推理能力。
+
+  复杂度: O(L · k · D) per iteration，其中 k = 邻居数
+
+  优化 (v3.2):
+    - 用门控线性单元 (GLU) 替代 GRUCell，完全消除顺序依赖，
+      所有节点可并行更新。GRU 的 (B*L, D) 独立 cell 调用本身
+      无序列依赖，但 GLU 更轻量且编译器友好
+    - 添加了 dropout 支持
+"""
+
+import logging
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+logger = logging.getLogger(__name__)
+
+
+class GraphReasoningModule(nn.Module):
+    """图推理模块：基于消息传递的关系推理。
+
+    架构流程（每次迭代）：
+        1. 构建稀疏图（局部 + 跨步邻居）
+        2. 计算边特征（节点对的关系表示）
+        3. 生成消息（边特征 × 邻居特征）
+        4. 聚合消息（加权求和）
+        5. 门控线性单元更新节点状态（替代 GRU，更易并行化）
+
+    多次迭代使信息在图中逐步传播，实现多跳推理。
+
+    Args:
+        d_model: 节点特征维度
+        num_neighbors: 每个节点的邻居数
+        num_iterations: 消息传递迭代次数
+        edge_dim: 边特征维度
+        dropout: Dropout 比率
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_neighbors: int = 16,
+        num_iterations: int = 2,
+        edge_dim: int = 64,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.num_neighbors = num_neighbors
+        self.num_iterations = num_iterations
+
+        # 边特征计算：从节点对中提取关系
+        self.edge_encoder = nn.Sequential(
+            nn.Linear(d_model * 2, edge_dim),
+            nn.GELU(),
+            nn.Linear(edge_dim, edge_dim),
+        )
+
+        # 消息生成：边特征调制邻居信息
+        self.message_fn = nn.Sequential(
+            nn.Linear(d_model + edge_dim, d_model),
+            nn.GELU(),
+        )
+
+        # 消息注意力权重
+        self.attn_proj = nn.Linear(edge_dim, 1)
+
+        # 节点更新: 门控线性单元 (GLU)
+        # 替代 GRU，完全可并行，不需要顺序执行
+        self.node_update = nn.Linear(d_model * 2, d_model * 2)  # 输出一半是gate，一半是value
+        self.update_norm = nn.LayerNorm(d_model)
+
+        # 输出投影
+        self.out_proj = nn.Linear(d_model, d_model, bias=False)
+        self.norm = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def _build_neighbor_indices(
+        self, seq_len: int, device: torch.device
+    ) -> torch.Tensor:
+        """构建混合邻居索引：局部 + 跨步。
+
+        不使用 O(L²) 的全局相似度计算，而是使用：
+        - 局部邻居：捕获相邻 token 的关系
+        - 跨步邻居：捕获长距离 token 的关系
+
+        Returns:
+            neighbors: (1, L, k) 邻居索引
+        """
+        k = min(self.num_neighbors, seq_len - 1)
+        if k <= 0:
+            return torch.zeros(1, seq_len, 1, dtype=torch.long, device=device)
+
+        k_local = max(k // 2, 1)
+        k_stride = k - k_local
+
+        indices = torch.arange(seq_len, device=device)
+
+        # 局部邻居
+        local_offsets = torch.arange(
+            -(k_local // 2), k_local // 2 + 1, device=device
+        )
+        local_offsets = local_offsets[local_offsets != 0][:k_local]
+        if len(local_offsets) < k_local:
+            extra = torch.arange(1, k_local - len(local_offsets) + 1, device=device)
+            local_offsets = torch.cat([local_offsets, extra])
+        local_nb = (indices.unsqueeze(1) + local_offsets.unsqueeze(0)).clamp(
+            0, seq_len - 1
+        )  # (L, k_local)
+
+        # 跨步邻居
+        if k_stride > 0 and seq_len > k_local + 1:
+            stride = max(1, seq_len // (k_stride + 1))
+            stride_positions = torch.arange(0, seq_len, stride, device=device)
+            if len(stride_positions) > k_stride:
+                stride_positions = stride_positions[:k_stride]
+            else:
+                pad = torch.zeros(
+                    k_stride - len(stride_positions),
+                    dtype=torch.long,
+                    device=device,
+                )
+                stride_positions = torch.cat([stride_positions, pad])
+            stride_nb = stride_positions.unsqueeze(0).expand(
+                seq_len, -1
+            )  # (L, k_stride)
+            neighbors = torch.cat([local_nb, stride_nb], dim=-1)  # (L, k)
+        else:
+            neighbors = local_nb  # (L, k_local)
+
+        return neighbors.unsqueeze(0)  # (1, L, k)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: (batch, seq_len, d_model) 节点特征
+        Returns:
+            output: (batch, seq_len, d_model) 更新后的节点特征
+        """
+        B, L, D = x.shape
+
+        if L <= 1:
+            # 单 token 时无法构建图，但仍通过投影和归一化保持表达力
+            return self.out_proj(self.norm(x))
+
+        # 构建邻居索引
+        neighbors = self._build_neighbor_indices(L, x.device)  # (1, L, k)
+        k = neighbors.shape[-1]
+        neighbors = neighbors.expand(B, -1, -1)  # (B, L, k)
+
+        h = x  # 初始节点状态
+
+        for iteration in range(self.num_iterations):
+            # 收集邻居特征
+            nb_flat = neighbors.reshape(B, -1)  # (B, L*k)
+            nb_expanded = nb_flat.unsqueeze(-1).expand(-1, -1, D)  # (B, L*k, D)
+            h_neighbors = h.gather(1, nb_expanded).view(B, L, k, D)  # (B, L, k, D)
+
+            # 计算边特征
+            h_expanded = h.unsqueeze(2).expand(-1, -1, k, -1)  # (B, L, k, D)
+            edge_input = torch.cat([h_expanded, h_neighbors], dim=-1)  # (B,L,k,2D)
+            edge_features = self.edge_encoder(edge_input)  # (B, L, k, edge_dim)
+
+            # 生成消息
+            msg_input = torch.cat([h_neighbors, edge_features], dim=-1)
+            messages = self.message_fn(msg_input)  # (B, L, k, D)
+
+            # 注意力加权聚合
+            attn_logits = self.attn_proj(edge_features).squeeze(-1)  # (B, L, k)
+            attn_weights = F.softmax(attn_logits, dim=-1)
+            agg_messages = (messages * attn_weights.unsqueeze(-1)).sum(
+                dim=2
+            )  # (B, L, D)
+
+            # 门控线性单元节点更新（全并行，替代 GRU）
+            update_input = torch.cat([h, agg_messages], dim=-1)  # (B, L, 2D)
+            update_raw = self.node_update(update_input)  # (B, L, 2D)
+            gate_val, value = update_raw.chunk(2, dim=-1)  # 各 (B, L, D)
+            gate_val = torch.sigmoid(gate_val)
+            h = gate_val * self.update_norm(value) + (1 - gate_val) * h
+
+        output = self.out_proj(self.norm(h))
+        return self.dropout(output)