《GoMLX从入门到精通》第14章：性能调优——从 Profile 到 Kernel Fusion

Chapter 14 — Version 2

14.1 性能优化的金科玉律：先测再改

调优的第一条铁律：没测过，别动手。

在机器学习里，你能"凭感觉"优化的事极其有限。你猜 batch_size=64 挺快，实际 128 更快 30%——因为 GPU 并行度没吃满。你觉得加一层卷积能提精度，实际精度没变、速度跌了。

正确的顺序：

1. 测量 → 2. 定位瓶颈 → 3. 假设原因 → 4. 改一个变量 → 5. 再测量 → 6. 验证假设

Go 自带的性能分析工具完全适用于 GoMLX 代码。

14.2 微观基准：Go 的 Benchmark

import "testing"

func BenchmarkInference(b *testing.B) {
    backend, _ := backends.New()
    exec := graph.MustNewExec(backend, modelFn)

    input := generateInput()
    b.ResetTimer()

    for i := 0; i < b.N; i++ {
        result := exec.MustExec1(input)
        result.FinalizeAll()
    }
}

// 运行：go test -bench=. -benchmem

输出会告诉你每次推理的耗时和内存分配。-benchmem 特别重要——如果每次推理都分配了意外的内存，说明有不该有的数据搬运。

14.3 宏观分析：go tool pprof

import _ "net/http/pprof"

func main() {
    go func() {
        http.ListenAndServe(":6060", nil)
    }()

    // ... 训练代码 ...
}

# 训练时在另一个终端：
go tool pprof http://localhost:6060/debug/pprof/profile?seconds=30
# 进入 pprof 交互界面后：
(pprof) top10     # 查看 CPU 占用最高的函数
(pprof) list myModel  # 查看具体函数的热点

pprof 能告诉你：训练时间的 60% 花在 graph.MustExec 上，20% 花在数据预处理，10% 花在指标计算。然后你就知道该优化哪里。

14.4 常见瓶颈与对策

瓶颈一：数据搬运是性能杀手

现象：每步训练的耗时远超预期，pprof 里数据传输相关的函数开销很大。

原因：训练数据在每一步都从 Go 堆复制到设备内存（或反过来）。

对策：

1. 增大 batch_size——每次搬运的数据量不变，但搬运频次降低
2. 在 Dataset 层预加载数据到设备——GoMLX 支持异步预取
3. 不要在训练循环里频繁调 MaterializeLocal()——只在需要可视化或保存 Checkpoint 时用

瓶颈二：图编译开销

现象：训练刚开始的几步很慢，后面快起来。

原因：JIT 编译需要时间（中等复杂图约 0.1-2 秒）。形状变化频繁 → 反复编译 → 慢。

对策：

1. 固定 batch_size（用 2 的幂次：64, 128, 256）
2. 不足一批的数据用 0 pad + mask（而非用不同 size）
3. 如果形状确实多变，用 Exec.SetMaxCache(64) 增加缓存上限

瓶颈三：GPU 显存不足（OOM）

现象：训练中途崩溃，日志里有 "out of memory"。

对策：

1. 减小 batch_size
2. 用 Float16 替代 Float32——显存减半
3. 及时调 FinalizeAll()——不用的中间结果立即释放
4. 减小模型宽度（隐层神经元数从 512 → 256）

优先调 batch_size。batch_size 对吞吐的影响是直接的——GPU 一次处理的越多，kernel 启动开销摊销得越薄。

瓶颈四：SimpleGo vs XLA 的选择

现象：在 SimpleGo 后端训练大模型太慢。

对策：如果硬件允许（Linux + NVIDIA GPU），切换到 XLA 后端。XLA 的算子融合 + GPU 并行，在大模型上可能带来 3-10 倍加速。

# 强制用 XLA（如果在 Linux 上）
GOMLX_BACKEND=xla go run main.go

但记住第 5 章的结论：小模型上 SimpleGo 和 XLA 的差距不大——编译开销可能吃掉加速收益。

14.5 优化优先级

瓶颈严重性（从高到低）：
    1. 数据搬运过多        → 调 batch_size / 预加载
    2. 形状变化 → 反复编译  → 固定 batch_size + pad
    3. 显存不够             → 减小 batch / Float16 / 及时 Finalize
    4. 后端选择不当         → SimpleGo → XLA（如果硬件支持）
    5. 模型太大             → 剪枝 / 量化 / 减层

规则：从上往下试。每试一个，测一次。不要同时改两个——你不知道哪个有效。

作者：步子哥，资深AI研究员 | 连载中：《GoMLX从入门到精通》

[上一章：第13章 模型保存、导出与 ONNX 互操作] | [下一章：第15章 GoMLX 的边界——WebAssembly、Gemma 推理与未来]