Skip to content

Basalt 3.9.0 — 张量生命周期交给 Rust

Choose a tag to compare

View all tags
@github-actions github-actions released this 31 Oct 01:27
· 6 commits to main since this release
v3.9.0
a4d6fd9

Basalt 3.9.0

这一版只做一件事：把张量的生命周期从 Java 的 GC 手里拿回来，交给 Rust。

  • 坐标 — io.github.9178nekomc:basalt:3.9.0
  • 在线文档 — https://9178nekomc.github.io/basalt/
  • 运行要求 — Java 21，启动加 --add-modules jdk.incubator.vector
  • 规模 — 243 个公开类型 / 1568 个公开成员，全部带逐参数说明

问题

一个字符级训练任务（每轮 2000 个窗口、batch 32、2.68 M 参数）在第 15 轮左右崩溃：

memory allocation of 4194304 bytes failed

不是常见的泄漏。每个原生张量的回收挂在一个约 40 字节的 Java 句柄上，通过 Cleaner 触发——
Java 堆几乎没有压力，GC 于是根本不来，而原生侧每一批输入、每个前向中间量、损失张量都还活着。

先试过两个补丁，都只把崩溃点往后推：

补丁 结果
在 Trainer 里显式 close() 四个可见张量 从第 15 轮推到第 30 轮
每 128 步 System.gc() 同样到第 30 轮

两个都是在给「GC 什么时候来」打补丁，没有触及设计本身。

改动

张量在作用域内由 Rust 持有：

  • Context 增加帧栈，每次分配把地址记进栈顶帧；begin_scope 压栈，end_scope 弹栈并把帧内张量一次性释放，返回实际释放个数。
  • 三个 JNI 入口，对应 NativeContext.beginScope() / endScope() / scopeDepth()。
  • Trainer 把每个训练步包进作用域：批次张量、前向中间量、损失都在这一步结束时回收。参数、输入、优化器状态在作用域外创建，不受影响。
  • 顺序有约束并在文档里写明：end_scope 必须在 backward 与 clear_graph 之后，否则计算图仍引用着已释放的张量。
context.beginScope();
try {
    Tensor logits = model.forward(inputs);
    Tensor loss = engine.crossEntropy().loss(logits, labels);
    loss.backward();
    optimizer.step();
} finally {
    context.endScope();   // Rust 一次回收整帧，不等 GC
}

验证

同一套配置改造后跑过 30 轮以上，没有再出现分配失败——这是本次修复的直接判据。
训练在 30 轮处主动停止，没有跑到收敛：这一版的目的是证明内存不再无界增长，不是产出一个模型。

  • ./verify.sh 全绿（7 组 legacy 套件 + JUnit 5 + 6 个编译期拒绝用例，-Xlint:all 零告警）
  • TextEngineTest 8/8、UnifiedModelTest 8/8、AdvancedApiTest 9/9
  • 文档：19 页 / 243 类型 / 1568 个成员中 1568 个有说明 / 0 占位符 / 0 断链

产物

文件 说明
basalt-3.9.0.jar 标准 jar，含四平台原生库
basalt-3.9.0-all.jar Fat jar，无外部依赖
basalt-3.9.0-sources.jar 源码
basalt-docs-3.9.0.zip 离线文档，19 页，解压双击 index.html

前几版留下的

  • 3.8.0 — TextEngine.stream / streamSampled：逐 token 回调，解码结果与批量接口逐位一致
  • 3.7.0 — 文档与基准报告 HTML 移出版本库，由 CI 构建并发布，单次推送从 17.1 MiB 降到 5.8 MiB
  • 3.6.0 — basalt.model 统一模型层，Auto.sequence 可筛序列模型
Assets 6
Loading