SunnyNet_CSDN
# SunnyNet 实现拆解:一个原创 O(N) 序列架构的 CPU 训练工程实践 > 项目地址:https://github.com/SunnyClaw-cx/SunnyNet (Apache-2.0) > 本文所有性能数字均来自仓库 `docs/EXPERIMENT_REPORT.md`,原始数据在 `experiments/results/*.json`。 --- ## 一、摘要 SunnyNet 是一个**不基于 Transformer / RNN / Mamba / CNN** 的通用序列神经网络架构,设计目标是 **CPU 高效训练**。 | 项目 | 说明 | |---|---| | 核心技术 | 一阶线性递推 + 对数均匀多尺度遗忘 + 组内成对竞争门控 | | 序列复杂度 | 严格 O(N),无 token 两两交互 | | 依赖 | PyTorch ≥ 2.1(**CPU 版即可**)、numpy | | 规模实测 | 1.4M / 20M / 100M | | 许可 | Apache-2.0 | **性能摘要**: | 场景 | 结果 | 条件 | |---|---|---| | 长序列前向 | T=4096 快 **22.69×** | batch=1、纯前向、8 线程、1.4M 模型 | | 长序列训练 | T=1024 快 2.08×、ppl 优 21.2%、内存省 36.8% | 200 步、batch=4 | | 训练步速度 | T=256 **与 RWKV 统计持平** | L=4;L=6 时反慢 30–40% | | T=256 质量 | 七架构**第二**(ppl 57.83) | 600 步未收敛 | --- ## 二、架构实现 ### 2.1 仓库结构 ``` sunnynet/ ops.py 原创算子: poly_sigmoid / softwedge / CenteredScaleNorm / 多尺度衰减 / 对数域并行扫描 / 通用门控扫描 layers.py SunnyBlock: 状态更新单元 + 循环滤波 + 局部卷积 + 非线性支路 linear_baselines.py 同复杂度对照: 线性注意力 / RWKV / Mamba / Hyena / H3 model.py SunnyNet 主体 (状态 (B,L,D), 逐层接力) init.py 谱锚定初始化 tokenizer.py 从零实现的 BPE 分词器 profiling.py 真实性能探针 (峰值内存 / 吞吐) ``` ### 2.2 状态更新:凸组合保证有界 ``` s_t = p_t ⊙ s_{t-1} + (1 - p_t) ⊙ u_t , p_t ∈ [1-δ_i, 1] ``` 两系数非负且和为 1 ⇒ 逐元素凸组合,故 ``` |s_t,i| ≤ max(|s_t-1,i|, |u_t,i|) ⇒ ‖s_t‖∞ ≤ max(‖s_0‖∞, max_k ‖u_k‖∞) ``` **状态天然有界。** 对比 LSTM 的 `c_t = f⊙c_{t-1} + i⊙c̃`——新信息是加法注入、无上界,这是两者数学结构上的根本差异。 ### 2.3 关键实现一:成对竞争门控(替代 softmax) 原始设计用 `softmax(g_t)` 做通道竞争,实测梯度范数随深度按 0.018 倍衰减,**16 层后精确归零**。根因是 softmax 的 Jacobian 含 `1/d` 因子。 修正为组内成对竞争: ```python def _compete(self, gate: torch.Tensor) -> torch.Tensor: """组内成对竞争门控 w = sigmoid(g_i-g_j)*sigmoid(g_i-g_k) in (0,1). Jacobian O(1).""" B, T, D = gate.shape gg = gate.view(B, T, D // self.group, self.group) w = torch.sigmoid(gg - torch.roll(gg, -1, dims=-1)) * \ torch.sigmoid(gg - torch.roll(gg, -2, dims=-1)) return w.reshape(B, T, D) ``` 语义等价(都是"多通道争夺资源"),但 **Jacobian 从 O(1/d) 恢复到 O(1)**。 ### 2.4 关键实现二:分块并行扫描 递推是仿射的,块内闭式解: ``` s_j = exp(P_j) · ( s_in + Σ_{k≤j} g_k · exp(-P_k) ) P_j = Σ_{k≤j} ln p_k ``` 块首状态可由块末精确反解:`s_in = (S_T - (1-p_T)u_T)/p_T`。因此块内可全并行,块间只串行 `T/chunk` 次。 ```python def gated_scan(p, g, s0, chunk=64, p_min_hint=None): """两级: 块内一次 cumsum 覆盖所有块 (并置), 块间串行 nblk 步.""" p, g = torch.broadcast_tensors(p, g) T = p.shape[-2] p_min = p.detach().amin().item() if p_min_hint is None else p_min_hint chunk = _safe_chunk(p_min, chunk) nblk = (T + chunk - 1) // chunk ... lp = torch.log(p).reshape(*lead, nblk, chunk, N) gg = g.reshape(*lead, nblk, chunk, N) P = torch.cumsum(lp, dim=-2) E = torch.exp(P) A = torch.cumsum(gg * torch.exp(-P), dim=-2) Elast, Alast = E[..., :, -1, :], A[..., :, -1, :] s = s0 s_in = [] for c in range(nblk): s_in.append(s) s = Elast[..., c, :] * (s + Alast[..., c, :]) ... ``` > **与逐块串行版本代数等价,非近似**——fp32/fp64 实测 bit-identical。 **与之配套的一处等价优化**:`p_min` 用解析下界 `p_min_bound` 传入(`p_min_hint`),跳过 `p.amin()` 的同步开销。单步 0.596ms → 0.043ms,**取值不变**。 ### 2.5 关键实现三:`_safe_chunk` 的边界处理 ```python def _safe_chunk(p_min: float, requested: int, headroom_bits: int = 64) -> int: """按 p 最小值限制块长, 使 exp(-P) = 1: return min(requested, 64) per_bit = -math.log(p_min) if per_bit **踩坑记录**:最初没写 `isfinite` 判断。对照实验里 Mamba 在某学习率下发散产生 `p_min=NaN`,于是 `int(NaN)` 抛 `ValueError`,**整个对照进程崩溃、无任何结果落盘**。修复后训练器会在 loss/grad_norm 非有限时打印 `[发散]` 并优雅中止,把 `diverged` 写入 summary。 ### 2.6 原始组件 **SoftWedge 激活**: ```python def poly_sigmoid(z): return 0.5 + 0.5 * (z / (1.0 + z.abs())) def softwedge(z, a): return z + a * poly_sigmoid(z) ``` `f'(z) = 1 + a·s'(z) ∈ (1, 1+a/2]`——**下限恒为 1**,任何尺度梯度不消失;上限有界,配合谱锚定可保证多层乘积增益有界。且只需加法、取绝对值、一次除法,**无指数/三角运算**,CPU 上比 GELU/SiLU 便宜。 **CenteredScaleNorm**: ```python class CenteredScaleNorm(nn.Module): def __init__(self, dim, eps=1e-5, s_min=0.02): super().__init__() self.eps, self.s_min = eps, s_min self.gamma = nn.Parameter(torch.ones(dim)) self.center = nn.Parameter(torch.zeros(dim)) def forward(self, x): return F.layer_norm(x, (x.shape[-1],), self.gamma, self.center, self.s_min * self.s_min + self.eps) ``` 两个设计点:① `center` 是**逐维可学习零点**(对反复读写的状态而言,"该通道的零点是哪"比统计均值更合理);② 分母下界 `s_min` 把近常量输入的放大倍数**钳在 1/s_min = 50 倍**。 > **这行 `s_min` 改掉过一次**:旧值 0.25 高于实际激活尺度(各层 norm 输入 std 低至 0.06、嵌入仅 0.02),导致所有输入都走"近常量分支",分母恒为 0.25,输出只是输入的 4× 常数缩放——**归一化从未真正发生**,600 步 ppl 劣化约 26%。定稿取 0.02(扫描显示 ppl 随 s_min 单调下降,平台在 [0.01, 0.025])。 **注**:`forward` 走 `F.layer_norm`(oneDNN 融合核,把 mean/var/sub/mul/add 融成单 kernel)。手写链实测前向慢 19 倍。在 `c=0` 时 CSN 与 gamma-scaled LayerNorm 严格等价。 **谱锚定初始化**: ```python def orthogonal_spectrum_(weight, rho, generator=None): """原地写入谱半径为 rho 的正交切片, 条件数为 1.""" out_f, in_f = weight.shape n = max(out_f, in_f) a = torch.randn(n, n, generator=generator, dtype=weight.dtype, device=weight.device) q, r = torch.linalg.qr(a) sign = torch.sign(torch.diagonal(r)) sign = torch.where(sign == 0, torch.ones_like(sign), sign) q = q * sign.unsqueeze(0) base = q[:out_f, :in_f] if out_f **一处被实测推翻的判断**:最初按 `ρ_l = ρ_0(l+1)^{-β}`、`β=0.15` 做层深衰减,理由是"抵消深堆叠增益"。实测证明**符号设反了**——`β=0.15` 时 24 层/8 层梯度比从 1.28 掉到 **0.017**,是直接制造梯度消失。 --- ## 三、性能优化实战 ### 3.1 唯一可行的融合路径:oneDNN 现成核 `torch.compile` 在本机**不可用**。原因不是没装 MSVC,而是 **MSVC 实际存在(BuildTools 14.39.33519),但 `Windows Kits\10` 缺 `Include`/`Lib`**,`cl /nologo t.c`(含 `stdio.h`)报 `cannot open 'stdio.h'`;且 `vcvars64.bat` 依赖被沙箱阻止的 `reg.exe`。 CPU 上真正可用的融合是 **oneDNN 原生融合算子**: - CSN 改走 `F.layer_norm`:前向 4.535ms → **0.232ms(19.5×)**,反向 11.075 → 2.485ms(4.46×) - 扫描改**两级并行**(bit-identical) **端到端训练步提速**: | T | 256 | 512 | 1024 | 2048 | 4096 | |---|---|---|---|---|---| | 提速 | −19.3% | −24.4% | −18.9% | −29.4% | **−41.2%** | **质量同步提升**(600 步 3 种子均值): | 实现 | 均值 ppl | |---|---| | 旧(手写 CSN + 串行扫描) | 92.04 | | **新(融合 LN + 并行扫描)** | **84.37** | 不是拿质量换速度,两项同时变好。 ### 3.2 减层加宽:修掉自己的公平性缺陷 meta device 真实计数发现**参数对齐下 SunnyNet 层数远超同行**(20M: L10 vs 5–7;100M: L14 vs 7–9),违反项目自定的"层数偏离 ≤±1"纪律。 在栈参数偏差 ±0.1% 前提下重平衡: | 预设 | 配置 | 栈参数 | T=256 ms/tok | 峰值 RSS | |---|---|---|---|---| | proto20m | d512 i896 r320 **L10** | 16,766,514 | 1.262 | 2253 MB | | **proto20mw** | d640 i1280 r384 **L6** | 16,755,230(+0.07%) | **1.110(−12.1%)** | **1935 MB** | 质量也同步改善(400 步,T=256): | 预设 | 100 步 | 200 | 300 | 400 | |---|---|---|---|---| | proto20m (L10) | 121.56 | 88.47 | 69.13 | 61.55 | | **proto20mw (L6)** | **112.75** | **78.79** | **60.94** | **54.82** | **ppl 全程领先,末档 −10.9%**,内存 −14%,训练步 −12.1%。 ### 3.3 收益上界测量:融合 kernel 只值 8.8% 把扫描与竞争门控**各替换成极简算子**(数值不等价,仅为量上界): | 变体 | ms/token | 相对基线 | |---|---|---| | base | 0.1575 | — | | 两者同时归零(上界) | 0.1436 | **−8.8%** | **根因**:反向占训练步约 **65%**,前向扫描占比再高也摊不进大分母。所以**优先级低于减层加宽**(12–17%,且无需编译工具链)。 ### 3.4 被实测否决的优化(勿重犯) | 尝试 | 判定 | 依据 | |---|---|---| | 扫描块长调整(64→128/256/512/1024) | 无收益 | T=512…4096 差异在 ±10% 噪声内 | | `gated_scan` 单块免 `torch.cat` | 否决 | A/B 增益 ±5% 且方向不稳定 | | 竞争门控 `roll`→切片改写 | 否决 | 实测更慢(4.91ms vs 1.95ms) | | CSN 用 `var_mean` 融合 | 否决 | 前向快 46% 但反向更慢,数值差 4.8e-7 超不变量 | | 状态反馈 `_feedback` 免 `cat` | 否决 | 数值不等价(偏差 1.8) | | 竞争门控 `roll`→`index_select` | 否决 | 微基准快(cpu 0.891×/cuda 0.717×)但**端到端无差异** | | 块内转置扫描 | 否决 | 隔离大维度下快 3.4×,**真实 `dim=128` 下反慢 26.2%** | | 多线程 1→8 | **有效** | T=4096:322ms→147ms(2.2×) | > **两条反复出现的教训**: > ① **微基准不得单独作为采纳依据**——算子级优势会被端到端摊薄进噪声(`index_select` 一例:微基准明显更快,端到端两平台均无差异)。 > ② **扫描基准必须用真实块维度**(`dim=128`),不能用隔离测试的大维度(N=512)——两者结论**相反**。 --- ## 四、测量方法论(这部分比代码更值得抄) ### 4.1 进程隔离是硬要求 **发现**:同一模型在同一进程"构建过多个模型之后"测量,耗时可达干净进程的 **2 倍以上**(T=4096: 325ms vs 776–822ms),且**放大倍数随架构而变**。 ``` §2.1(污染) vs §2.4(进程隔离) T=4096 12.63× 22.69× 交叉点 T≈834 T≈384 ``` **修正做法**:每个 `(架构, 序列长度)` 单元格单独起子进程,进程内只构建一个模型。 ### 4.2 参数量对齐用 meta device 计数,不用手推公式 实测手推公式**漏项严重**——RWKV 偏差达 **34%**。改为 `meta device` 实例化真实计数。 ### 4.3 速度结论要用"逐轮配对"而非单一中位数 9 轮配对,每轮两架构**在相邻时刻各跑一次**,抵消轮间漂移。`sunny80 vs rwkv` 三个指标互相矛盾(中位比 vs 配对比值 vs 符号检验),因此写作"统计持平"而非"超越"。 --- ## 五、复现步骤 ```bash # 1. 环境 pip install torch --index-url https://download.pytorch.org/whl/cpu pip install numpy matplotlib pytest # 2. 数据(Project Gutenberg 中文古籍 + 英文文学 + TheAlgorithms/Python) python prepare_data.py python build_corpus.py --vocab-size 8192 # 3. 训练 python train/train.py --preset small --steps 800 --batch-size 8 --seq-len 256 python train/train.py --preset proto20m --steps 5000 --batch-size 4 --seq-len 512 # 4. 评估(困惑度 / 长距离探针 / 生成) python train/evaluate.py --ckpt runs/.pt --gen # 5. 完整实验 bash experiments/battery.sh # 基线 1200 步 + 消融/种子 600 步 + 20M 缩放 bash experiments/battery2.sh # 20M 缩放对照 + 长序列训练对照 python experiments/long_sequence.py --only latency # O(N) 核心验证 # 6. 出图与汇总 python experiments/make_figures.py python experiments/collect_results.py --out docs/_tables.md ``` ### 语料统计(真实采集) | 域 | 训练 token | 验证 token | 训练文件 | 验证文件 | |---|---|---|---|---| | 中文古籍 `zh_ancient` | 39,987,804 | 2,183,213 | 345 | 22 | | 英文文学 `en_literature` | 4,343,482 | 86,836 | 15 | 1 | | Python 代码 `code_python` | 1,254,856 | 79,005 | 1,267 | 80 | | **合计** | **45,586,142** | **2,349,054** | 1,627 | 103 | > 验证集按**文件级**切分(`val_ratio=0.06`),同一文本相邻片段不会跨训练/验证,杜绝信息泄漏。 > **缺失域如实标注**:中文维基在本机网络不可达,因此没有"百科"域,只覆盖 3 个真实域。 --- ## 六、局限(必读) 1. **短序列劣势**:O(N) 优势只在长序列兑现。交叉点 **T≈384**,短于此比 Transformer 慢。**不适合短序列场景。** 2. **长序列对同行也无绝对优势**:T=4096 时 Hyena 仍快 1.9×、H3 基本追平。领先全部同行的区间只有 **512 模型为**实验科研项目**,未经过安全对齐流程。若语料含中医古籍,其输出**不构成任何医疗建议或诊断依据**。 --- ## 七、结语 SunnyNet 的适用条件很明确:**长序列 + 质量优先 + 可接受训练成本**。 它最有价值的部分可能不是那 22.69×,而是**一套完整的、包含负面结论的实测记录**——实验报告的 TL;DR 里有 4 行标注"已由 §X 修正",每个速度数字都带测量条件,每个"第一"都注明步数与收敛状态。 **项目地址**:https://github.com/SunnyClaw-cx/SunnyNet **完整实验报告**:`docs/EXPERIMENT_REPORT.md`(含全部负面结论与复核数据) **架构数学推导**:`docs/ARCHITECTURE.md`