仓库地址

今天来介绍一下我的私货，一个用vulkan进行AI训练的库，背景，我当初学写shader，为了更加方便加shader，而发展出来的。现在还有好些调用还是没有实现完，用了会崩溃。昨天用AI暴力对拍，写出了这个玩具demo。证明了这个通过组织shader的方式训练AI可行，而且能比pytorch快很多，因为控制流透明，可以做很多优化。用上了c++23。

大概率其他机子都跑不通，因为这个协作矩阵和wave工作组的大小大概率不一样，有RTX3060的可以让AI稍微改一下配置，下载一下libpng和vulkansdk。

关于未来畅想，让AI模型设计能够工程化，比如使用节点编辑器根据说明，来规模化组织预制算子集，让着色器资产化，配上各种已验证机器特化，节点编辑器能够检测环境自动应用。

下面是AI介绍

conv_train 训练架构笔记

主题: 一个只用 Vulkan compute 的 ConvSIREN 单图 INR 训练器——为什么它能训、
它的"以 shader 为中心"的数据流抽象长什么样、构建系统怎么把 81 趟着色器装配起来。
依据: CMakeLists.txt + conv_train.cpp + vulkan/ 框架头 + slang/conv/ 着色器源。

1. 这是什么

conv_train.cpp 在没有 PyTorch、没有 CPU 侧训练循环的前提下, 用 81 个 compute pass
在 GPU 上完成单张 512×512 图像的 SIREN 隐式神经表示拟合:

前向 19 趟 (kern/stem/每层4趟/head) → 存全部激活
损失 3 趟 (lossgrad + SSIM×2) → dL/dOut
反向 37 趟 (head + 每层9趟 + stem) → GW (与 WBuf 同布局的梯度)
更新 1 趟 Adam + 17 趟 Muon → 原地改 WBuf
收尾 2 趟 末 N 轮权重平均 (wavg)

每步训练的全部张量——激活、梯度、Adam 的 M/V、Muon 动量——都常驻显存,
跨步不回主机。host 每步只做三件事: 写一次 uniform (lr/β/偏差修正)、
读回一次 Out (786432 float) 算 MSE/PSNR 以驱动 lr 调度、按 ep 开关个别趟。
控制面在 CPU, 数据面全程在 GPU——这是它"能训"的第一层原因。

2. 文件构成与分工

依赖方向干净: computer_base ← acer ← computer_do ← (main / graphic_todo),
上屏组件借宿主句柄而不被宿主认识 (computer_do.hpp 头注释: 唯一接缝是
extraQueueCreateInfos)。

3. 构建: 一份源, 两个相, 三层目标

3.1 同源双宏

同一个 conv_train.cpp 被编译两次, 靠 ACER_SHADER 宏在编译期二选一:

• bulid 目标 (开发期工具): 不定义 ACER_SHADER → 只编译 shaderCode 分支,
  链接 libslang, 把 81 趟 slang 编译成 spv 写进 res_train/;
• conv_train 目标 (运行期): 定义 ACER_SHADER → 只编译 vulkanCode 分支,
  链接 Vulkan, 从内嵌进二进制的 spv 直接建 pipeline, 不碰 slang。

cmake -B build_cmake -DMAIN_SRC=conv_train.cpp -DACER_TRAIN=ON

MAIN_SRC 选主, ACER_TRAIN 选相 (训练 res_train/ / 生产 res/ 目录分开),
两个目标必须同值, 否则 cook 出的趟和运行期要的趟对不上。

3.2 bake 链 (CMake custom command 两级)

cook ① bulid slang → res_train/passN.spv (内容寻址, 完成落 .cook 戳)
cook ② embed.py spv+bin → resources.S/.cpp (.incbin, 零拷贝 .rodata)
链接 conv_train = main + resources.{S,cpp} + Vulkan ← 运行期只链数据

关键性质: 运行期二进制自包含。spv 和权重种子走同一张 resTable
("权重就是这个 shader 的另一种资源"), acer::shaderSource::setRes 把
span 直接挂到 binding 上——字节躺在 .rodata 里, 一字节都不拷。
第一次训练不必先有种子 (embed.py 对缺席的 bin 不出条目)。

4. 为什么能训练

4.1 前向存激活, 反向逐层回传

自动微分的两个前提它都手工满足了:

1. 前向保留中间量。 stem_train 存 h0 和 pre-activation z;
   每层 mix_unpack 存 h_{l+1} 和 m_l (sin 激活的输入);
   dw 直接产出打包好的 cat16_l (反向 mix 权重 GEMM 的 B 矩阵, 复用!)。
   这些 buffer 在反向趟里通过 share 标志被再次绑定——框架据此插屏障。
2. 手写反传。 每层 9 趟:
   bwd_sin_dm(双路打包 dmX16/dmT16) → pack_wmix/gemm_mixw/unpack_dm
   (bwd_mix_x, tensor-core GEMM 传到下层激活) → gemm_mixw/reduce_mixw
   (bwd_mix_w, K 维切 43 片各写 Cpart 再归约——像素数 29 万当 K 维太长,
   分片部分和) → bwd_dw_x/bwd_dw_k(深度卷积两个输入的梯度, 后者 K 维分
   DKP=5 块) → bwd_kern(k-net 梯度)。

梯度 GW 的布局与 WBuf 逐元素一致 (134320 float, 见 upd.slang 头注释
的槽位表), 于是更新趟就是逐元素遍历, 不需要任何索引映射。

4.2 损失与超参: uniform 当控制面

损失 = MSE + chroma + grad + SSIM (train.py band_mse 口径), 全在 GPU:
lossgrad 算基础项 → ssim_fwd 算局部系数场 → ssim_add 把 SSIM 梯度
+= 进 DlOut。增益系数 (如 gradScale = GRAD_LAMBDA/den, den 只依赖
target) 由 host 一次性算好写进 LossCfg uniform。

训练循环里 host 每步改的只有 uniform 标量:

cfg.lr = lr;  mcfg.lr = lr;
cfg.invbc1 = 1/(1-0.9^ep);  cfg.invbc2 = 1/(1-0.999^ep);
shaders[79].setSkip(ep <= target - avgN);   // 末 N 轮才开始累积权重平均
list.recordAndSubmit();

warmup / plateau 衰减 / EMA loss 状态全部在 host (UpdCfg 结构体),
Adam 的逐元素更新公式在 upd.slang 里严格复刻 torch。

4.3 优化器也全是 shader

• Adam (upd 趟): W -= lr·(m/√v), M/V 是两块 device buffer, 起始 0;
• Muon (muon_prep + 5×3 趟 Newton-Schulz + muon_apply): 对 mix 权重
  梯度做正交化, gram/poly/upd 全是 Cooperate-Matrix GEMM, X/X2 乒乓;
  host 只给 lr (lr_scale=25 烙进 shader);
• 权重平均 (wavg_acc/wavg_apply): 末 AVG_WINDOW=10 轮累积 W,
  收尾一步 W = A/N, 消掉零均值抖动。

4.4 续训: checkpoint 是 dump 位的副产品

凡是训练状态都标 dump (buffer 级) 或 uniform_bits::dump:

命名规则 -.bin 由读和写共用, 所以读回来的一定是写出去的
那个。最巧的是 uniform 的续训: uniform 是 VK_EXT_external_memory_host
零拷贝导入的 host 常驻内存, binIO 读回时把文件字节原地强转回结构体,
于是 ep / plateau_scale / ema_loss 这些"纯 host 状态"搭着 uniform
的车一起持久化——没有任何 sidecar 文件, lr 调度在中断点连续。

时序上有一处精确的讲究: checkpoint 存原始权重, 必须在 wavg_apply
(它会把 WBuf 原地改成均值) 之前 io.store() 并撤下自动写回, 否则续训
轨迹就被权重平均污染了; 而交付给推理的模型用的是平均权重——两个产物,
两条路, 各走各的。

5. 以 shader 为中心的数据流抽象

这是整套框架最核心的设计: 一切声明围绕"趟" (pass) 展开,
数据只是挂在趟的端口上的资源。

5.1 端口即身份

shaders[19].buffer.setKey(0, 1, 2);        // 这趟有 0/1/2 三个 buffer 端口
shaders[19].buffer[0].setFlag(share);      // 端口 0: 外部共享 (别人建的, 我只绑)
shaders[19].buffer[1].setData(tgt);        // 端口 1: host 上传 (target 图)
shaders[19].buffer[2].setSize(NOUT);// 端口 2: 设备-local, 框架分配
shaders[19].uniform.setKey(3);             // uniform 端口 3

一块数据的"身份"是 (趟槽号, 端口号), 而不是全局变量名。setData/setSize
的区别就是谁拥有: host 种子 vs 设备输出。share 位声明"这块由别的趟
建, 我只是再绑一次"——同时是框架插内存屏障的依据。

5.2 alias: 把数据流图画成绑定图

acer::alias(shaders[18].buffer[2]).to(shaders[19].buffer[0],
                                     shaders[20].buffer[0],
                                     shaders[21].buffer[1]);   // Out 三处消费

alias(src).to(dst...) 让多个端口共享同一个 VkBuffer——数据流的边
就是绑定的重合。重复扇出用 bcast 收成两段:
bc(4,2).to(8,12,16) (同端口号广播)、bc(4,2).to_at(0, 5, 9, 13, 17)
(错位: 源端口 2 → 目标端口 0)。WBuf 这种被 30+ 趟引用的权重, 声明一遍即可。

两层位标志各管一件事, 不混:

• buffer 级 share / dump / dmabuf —— "这块数据怎么处理";
• shader 级 skip / infer —— "这整趟参不参与" (跳帧不跑 / 进不进推理链)。

5.3 常量注入: 注解改写, 编译期定形

slang 侧用注释标注可注入点:

// model
namespace model { static const int NF = 0; ... }
using namespace model;

C++ 侧 SlangCompiler::setVar 找到标注行, 校验字段名/个数/重复,
把 = 0 改写成真值, 再喂给 slang 编 SPV。常用常量打包成参数包整包注入:

s[0].writeSlang();          // 8 个布局常量一次进
s[q+0].writeSlang("dims", KPAD);         // 逐趟变的仍逐条写
s[p+7].writeSlang("cfg", 1<

for (int ep = epStart+1; ep <= target; ++ep) { ...; list.recordAndSubmit(); }
// 收尾复评: 只开 0..18 前向
for (uint32_t i = 0; i <= 18; ++i) shaders[i].setSkip(false);
for (uint32_t i = 19; i <= 80; ++i) shaders[i].setSkip(true);

source /maple/share/vulkan/setup-env.sh
cmake -B build_cmake -DMAIN_SRC=conv_train.cpp -DACER_TRAIN=ON
cmake --build build_cmake -j$(nproc)
VK_ICD_FILENAMES=/usr/share/vulkan/icd.d/nvidia_icd.json \
VK_LAYER_PATH=/maple/share/vulkan/explicit_layer.d \
  ./build_cmake/conv_train 600        # 参数 = 本次继续推进的轮数 (增量)
# CONV_TIME=1 可出逐趟耗时归因表; nStep=0 只回读+重导模型包