不用 Python,也能训练大模型:两年之后再看 Go 语言机器学习框架 GoMLX

导读：两年前，我们写过一篇文章介绍刚刚起步的 GoMLX——一个试图在 Go 语言里复刻 PyTorch／Jax／TensorFlow 能力的机器学习框架。两年之后再回头看，它已经收获 1.5k 星标，拥有了独立的官方文档站点，还把核心计算引擎拆分成了独立的 compute 仓库，形成清晰的分层架构。这篇文章借着这次“回访”，从架构设计、四大核心抽象、多后端体系、代码实战到生态现状，全面梳理 GoMLX 现在到底发展成了什么样子。

文章要点：

• 架构重组与独立引擎：GoMLX 在 v0.28 版本完成关键蜕变，将底层计算图定义与 JIT 编译逻辑剥离至独立的compute仓库，形成了“高层 API + 可插拔计算引擎”的清晰解耦架构。
• 四大核心抽象：围绕Backend（硬件编译桥梁）、Graph（纯 Go 描述的计算图）、Tensor（数据与显存载体）和Store（参数与作用域管理）四大概念，建立了透明、无“魔法”的工程心智模型。
• 三系多元后端覆盖：提供基于 OpenXLA/PJRT 的高性能xla后端（压榨 GPU/TPU 极限算力）、零 CGO 依赖且支持编译为 WASM 的纯 Gogo后端，以及面向苹果生态的go-darwinml后端。
• 开源生态无缝对接：通过go-huggingface（原生 Tokenizer、safetensors/GGUF 解析）与onnx-gomlx（ONNX 计算图无缝转换与微调），实现了对通用预训练大模型生态的直接消费。
• 前沿算法与工程特性：紧跟 AI 研究前沿，内建了 KAN（Kolmogorov-Arnold 网络）、VNN、梯度检查点（重计算省显存）以及基于 XLA Shardy 的分布式多卡训练支持。

两年之后，GoMLX 走到了哪一步

两年前我写过一篇文章，介绍了刚刚起步的 GoMLX——一个试图在 Go 语言里复刻 PyTorch／Jax／TensorFlow 能力的机器学习框架。那时它更像是一个雏形：API 还在快速变动，文档也不算完善，社区讨论主要集中在“Go 到底能不能做机器学习”这个基础问题上。

两年之后再看，情况已经发生了明显变化：

• 项目在 GitHub 上收获了1.5k 星标、超过 70 次 fork，累计提交数已超过 3500 次；
• 项目拥有了独立的官方文档站点gomlx.github.io，涵盖从入门到进阶的完整教程体系；
• 最核心的变化是：原本内嵌在主仓库里的计算后端，被拆分成了独立的compute仓库，GoMLX 主仓库转型为“上层 ML API＋生态整合”的角色；
• 官方版本已经迭代到 v0.28 系列，并伴随一次较大规模的 API 重组。

这不再是一个“能不能用”的项目，而是一个开始具备清晰架构分层、逐渐向生产可用性迈进的框架。这篇文章就带你完整梳理一遍现在的 GoMLX，究竟是怎么设计的、能做什么、值不值得现在就上手。

GoMLX 是什么：定位与设计哲学

用一句话概括，GoMLX 是 Go 语言版本的 PyTorch／Jax／TensorFlow。它的核心主张只有一个词：不需要 Python。

GoMLX 支持训练、微调、修改和组合机器学习模型，提供了一整套可微分算子，以及训练过程中的可视化、调试工具。它的设计哲学可以归纳为几点：

• 简单且透明。官方文档中明确提到，GoMLX 追求“简单到可以被读懂和推理”，即使代价是代码更冗长，也要让使用者对发生的事情有一个正确、透明的心智模型，这与 Go 语言本身的哲学是一致的。
• 组合性优先于魔法。没有一个“万能训练对象”把一切都包办，backend、model、store、loss、optimizer、训练循环，每一个组件都可以被独立替换，方便研究者尝试新的优化器思路、复杂的正则化项或非常规的多任务训练方案。
• 文档即代码的一部分。官方的态度很直接：如果文档写得不好，就等于代码不存在。错误信息也要求尽量友好，并且总是带上完整的调用栈。

从落地场景看，GoMLX 的目标不止是“能跑通一个 Demo”，而是想成为一个可生产化、可用于研究和教学的完整 ML 平台，具体包括支持现代加速器硬件（GPU/TPU）、支持从 HuggingFace 导入预训练模型进行微调，以及未来可以把模型编译成二进制或 WebAssembly，在任意语言环境里被消费。

四大基石：Backend / Graph / Tensor / Store

理解 GoMLX，本质上是理解四个抽象概念，它们环环相扣，构成了整个框架的心智模型。

1. Backend：连接硬件的桥梁

compute.Backend是你的 Go 进程与底层硬件（CPU、GPU、TPU）之间的连接。它负责把计算图即时编译（JIT）成可执行代码，并管理主机内存与设备内存之间的数据搬运。通常一个进程只创建一个 backend，在程序全局复用：

import (
"github.com/gomlx/compute"
 _ "github.com/gomlx/gomlx/backends/default"// 引入默认后端
)

backend, err := compute.New() // 自动选择最合适的后端
fmt.Printf("Backend: %s\n", backend.Description())

compute.New()会按照 CUDA GPU → Metal（Apple）→ CPU 的顺序自动选择最优后端，也可以通过环境变量GOMLX_BACKEND或者compute.NewWithConfig("go")显式指定。

2. Graph：用纯 Go 函数描述计算

Graph（计算图）是一个用*graph.Node及其相互连接的运算组成的纯函数，用来描述一次计算。GoMLX 提供了丰富的高层 API 来构建这些图，构建完成后交给后端做 JIT 编译并高效执行。

addFn := func(a, b *Node) *Node {
return Add(a, b)
}
addExec, err := NewExec1(backend, addFn)
v1, err := addExec.Call(1.0, 1.0)

这里有两个反直觉但很关键的设计：

• 节点是“未来值”，不是具体张量。在图构建阶段，你无法读取*graph.Node的具体内容，只有在真正调用.Call()执行之后才能拿到结果。节点携带的是形状（shape）和数据类型信息，用于在构建阶段就能检查出维度或类型不匹配的问题。
• 图构建阶段用 panic，而不是返回 error。如果每一个数学算子都要求你手动判断并返回 error，数学公式会写得非常凌乱难读。所以 GoMLX 选择在构建阶段用 panic（异常）来处理形状/类型不匹配，Exec包装器会自动捕获这些 panic，并在.Call()被调用时转换成带完整堆栈信息的标准 Go error。

这种“图先构建、再编译、再执行”的模式，和 JAX 的@jax.jit、TensorFlow 的@tf.function本质上是同一套思路：把计算的全貌交给后端去看，才能做算子融合、内存布局优化等激进优化。

这里也有一个需要留意的“坑”：JIT 编译是和输入的静态形状绑死的。如果每次调用传入的 batch size 或序列长度都不一样，GoMLX 会为每一种新形状重新编译一次图，而编译的代价远高于执行本身，所以官方建议尽量固定输入形状，或者把变长输入 pad 到几个固定的桶（bucket）里复用已编译的图。

3. Tensor：承载数据的具体值

Tensor 是图计算的输入输出，代表一个具体的多维数组（也可以是标量），由形状（shapes.Shape）和数据类型（dtypes.DType）共同定义：

t := tensors.FromValue([][]float32{{1.0, 2.0}, {3.0, 4.0}})
fmt.Printf("Tensor shape: %s\n", t.Shape())
// 输出: Tensor shape: (Float32)[2, 2]

Tensor 内部会同时维护主机内存（本地 CPU）和设备内存（GPU/TPU）两份缓存，数据搬运是惰性触发的，只有在真正需要时才发生，以此减少不必要的拷贝开销。因为 Go 的垃圾回收器无法感知加速器设备上的内存，长期持有大量张量时，建议显式调用FinalizeAll()释放设备内存，而不是完全依赖 GC。

4. Store：管理可训练参数

如果只是做数学计算，Backend 加 Graph 就够了；但要训练模型，就需要一个地方持久化存放权重、偏置这类可训练变量（Variable），以及模型的超参数——这就是model.Store的职责。

model.Store是真正存放张量数据的容器，而model.Scope则是指向 Store 内某个路径（类似“当前目录”）的轻量指针。层函数接受一个*model.Scope，并在当前作用域内声明或查找变量，从而避免命名冲突：

funcdenseLayer(scope *model.Scope, x *Node, outputDims int) *Node {
 g := x.Graph()
 dtype := x.DType()
 inputDims := x.Shape().Dimensions[1]

 weights := scope.VariableWithShape("weights", shapes.Make(dtype, inputDims, outputDims)).NodeValue(g)
 biases := scope.VariableWithShape("biases", shapes.Make(dtype, 1, outputDims)).NodeValue(g)

return Add(Dot(x, weights).Product(), biases)
}

modelFn := func(scope *model.Scope, x *Node) *Node {
 h := denseLayer(scope.In("layer1"), x, 3) // 变量路径: /layer1/weights, /layer1/biases
 y := denseLayer(scope.In("layer2"), h, 1) // 变量路径: /layer2/weights, /layer2/biases
return y
}

scope.In("layer1")这类调用，本质上是在 Store 里划分出一棵树状的命名空间，让不同层的权重互不冲突，也方便后续统一遍历、保存、加载。

把这四个抽象放到一起看，会更容易理解它们之间的分工与数据流向：

如上图所示：输入张量流入计算图参与运算，Store 中的变量也会注入计算图；计算图交给 Backend 做 JIT 编译并执行，产出结果张量；训练过程中更新后的权重，再写回 Store，形成一个闭环。图中灰色代表数据（Tensor），青色代表框架的核心组件（Store／Graph／Backend）。

架构重组：compute 独立仓库意味着什么

v0.28 版本里，GoMLX 做了一次比较大的“外科手术”：把backends、dtypes、shapes、distributed这些偏底层的包，整体搬到了新的独立仓库github.com/gomlx/compute中。

compute 仓库的定位描述得很清楚：提供一个模块化的 API，用于定义和执行带有可插拔后端的多维计算图。它对外暴露一个compute.Backend接口（一组接口的集合），可以用来构建计算图、JIT 编译、在主机和设备之间搬运缓冲区（张量的原始数值）、执行已编译的计算。

这次拆分背后的用意值得展开说说：

• 职责更清晰。compute 只关心“如何定义和执行计算图”，追求的是正确和最小化，而不是易用性；GoMLX 主仓库则专注于在此之上构建符合人体工程学的高层 API：自动微分、层库、训练循环、数据集、checkpoint 管理等等。文档里原话是：如果你只是要做复杂计算和自动微分，还是应该用 GoMLX，而不是直接用 compute。
• 后端实现随之迁移。原本内置的“go”后端，现在实现在github.com/gomlx/compute/gobackend里，并且做了大量性能改进；“xla”后端则被移到了另一个独立仓库github.com/gomlx/go-xla下的compute/xla包中。
• 为第三方实现新后端铺路。compute 仓库文档里专门给出了实现自定义后端的步骤：继承notimplemented默认实现、实现缓冲区搬运、按需实现具体算子，最后用backendtest.RunAll(t, myBackend)跑一遍标准合规测试即可。换句话说，任何人都可以按照这套接口规范，为 GoMLX 生态贡献一个新的执行后端，而不需要动到上层 API 一行代码。

这种“上层框架＋独立可插拔计算引擎”的架构，和 PyTorch 的 ATen/Dispatcher，或者 JAX 的 XLA 层，思路是相通的：把“怎么算”和“算什么”彻底解耦。

三种后端：xla、go、go-darwinml

目前compute.Backend接口已经有三种实现，分别覆盖了不同的部署场景：

值得一提的是官方给出的一个真实例子：有人把 GoMLX 通过 go 后端编译成 WASM，用于给一个叫 Hive 的桌面游戏跑 AlphaZero 风格的 AI 对手，整个推理过程直接在浏览器里完成，不依赖任何服务端。这也印证了纯 Go 后端“哪里能跑 Go，哪里就能跑 GoMLX”的定位。

后端选择既可以自动完成，也可以通过环境变量精确控制，例如：

export GOMLX_BACKEND=xla:cuda   # 使用 XLA + Nvidia CUDA
export GOMLX_BACKEND=xla:cpu    # 使用 XLA + CPU
export GOMLX_BACKEND=go         # 使用纯 Go 后端

对于 XLA 后端，GoMLX 还提供了PJRT 插件自动安装能力：首次运行时会自动把对应硬件（CPU/GPU/TPU）所需的 PJRT 插件下载安装到用户本地目录，免去了手动配置的麻烦；如果需要制作精简的生产镜像，也可以通过--tags=pjrt_cpu_static等方式做静态链接。

实战：用不到 60 行代码训练一个神经网络

光讲架构比较抽象，我们来看一段官方示例的完整代码：训练一个多层感知机（MLP），学习把归一化后的像素坐标(x, y)映射为对应的 RGB 颜色，本质上是让网络“记住”一张图片。

// 1. 准备训练数据：把 (x, y) 坐标映射到 (r, g, b) 颜色
inputs := make([][]float32, 0, width*height)
labels := make([][]float32, 0, width*height)

for y := range height {
for x := range width {
  nx := float32(x)/float32(width)*2.0 - 1.0
  ny := float32(y)/float32(height)*2.0 - 1.0
  inputs = append(inputs, []float32{nx, ny})

  r, g, b, _ := img.At(bounds.Min.X+x, bounds.Min.Y+y).RGBA()
  labels = append(labels, []float32{float32(r) / 65535.0, float32(g) / 65535.0, float32(b) / 65535.0})
 }
}

backend := compute.MustNew()
store := model.NewStore()

// 2. 构建内存数据集
ds, err := dataset.InMemoryFromData(backend, "image_pixels", []any{inputs}, []any{labels})
ds.BatchSize(512, false).Shuffle().Infinite(true)

// 3. 定义模型结构（3 层 MLP）
modelFn := func(scope *model.Scope, spec any, inputs []*Node) []*Node {
 x := inputs[0]

 h := denseLayer(scope.In("layer1"), x, 64)
 h = activation.Relu(h)
 h = denseLayer(scope.In("layer2"), h, 64)
 h = activation.Relu(h)
 h = denseLayer(scope.In("layer3"), h, 64)
 h = activation.Relu(h)

 y := Sigmoid(denseLayer(scope.In("output"), h, 3))
return []*Node{y}
}

// 4. 配置训练器：Adam 优化器 + MSE 损失
trainer := train.NewTrainer(
 backend, store, modelFn,
 loss.MeanSquaredError,
 optimizer.Adam().LearningRate(0.003).Done(),
nil, nil,
)

// 5. 运行训练循环
loop := train.NewLoop(trainer)
train.EveryNSteps(loop, 1000, "log_metrics", 0, func(l *train.Loop, metrics []*tensors.Tensor)error {
 fmt.Printf("Step %5d: MSE Loss = %.6f\n", l.LoopStep, metrics[0].Value())
returnnil
})

_, err = loop.RunSteps(ds, 5000)

训练日志大致是这样的：

Starting training loop...
Step   999: MSE Loss = 0.000038 (moving average = 0.000040)
Step  1999: MSE Loss = 0.000067 (moving average = 0.000027)
Step  2999: MSE Loss = 0.000012 (moving average = 0.000030)
Step  3999: MSE Loss = 0.000010 (moving average = 0.000019)
Step  4999: MSE Loss = 0.000013 (moving average = 0.000018)
Training finished!

这段代码把前面讲的四大抽象串到了一起：backend 负责编译执行、store 负责持久化权重、graph（modelFn）负责描述前向计算、trainer/loop 负责组织训练流程。更重要的是每一块都可以单独替换：换个优化器只改一行，换后端只改一个环境变量，模型结构改动完全不影响训练循环的写法——这正是 GoMLX 反复强调的“组合性”设计理念在实际代码里的体现。

自动微分方面，GoMLX 使用graph.Gradient(loss, targets...)在图构建阶段自动完成符号微分，把反向传播所需的运算直接追加进计算图。需要注意的是，目前它只支持对标量损失求梯度，不直接提供雅可比矩阵或黑塞矩阵，如果需要高阶导数，需要手动对梯度节点再次求导来实现。

生态拼图：从 HuggingFace 到 ONNX

单靠核心框架很难覆盖真实业务场景，GoMLX 这两年里补齐了不少生态组件：

• go-huggingface：让 Go 程序可以直接对接 HuggingFace 生态——共用 Python 版本一致的模型/数据集缓存机制；提供纯 Go 实现的多种分词器（tokenizer），可直接从 HuggingFace 下载；支持读取 GGUF 或 safetensors 格式的模型参数；还提供了配套的 transformer 库做模型转换，甚至包含了等价于sentence_transformers的句子嵌入能力。像 Tencent 出品、在 RAG 场景中排名靠前的 KaLM-Gemma3（12B 参数）句子编码器，就是通过这套工具在 GoMLX 上跑起来的。
• onnx-gomlx：把 ONNX 模型转换成 GoMLX 可执行的计算图，既可以作为onnxruntime的替代方案（复用 XLA 的加速能力），也可以用来对模型做进一步微调。官方示例里已经跑通了 Gemma 3 270M 文本生成、BERT-base 命名实体识别、MixedBread Reranker 等真实模型。
• 数值互操作：github.com/gomlx/gomlx/core/tensors/numpy包支持直接读取 Numpy 数组，方便和 Python 生态做数据交换。
• Docker + JupyterLab：官方维护了一个预装 GoMLX、JupyterLab 和 GoNB（Go 语言的 Jupyter 内核）的镜像，还内置了 Nvidia CUDA 运行时，一条docker run命令就能拉起一个带 GPU 支持的交互式笔记本环境，对于想快速试用的人来说几乎是零配置。

从这些拼图可以看出一个清晰的信号：GoMLX 团队没有打算“重新发明一切”，而是选择尽可能与 HuggingFace、ONNX 这些既有生态对接，把主要精力放在 Go 语言侧的执行效率和工程体验上。

最新亮点：梯度检查点、KAN、VNN、分布式训练

层库（layers）是 GoMLX 里更新最频繁的部分之一，目前已经覆盖了相当完整的现代神经网络组件：

• 基础层：FFN 前馈层、多种激活函数、Layer Norm/Batch Norm、卷积、池化、Dropout；
• 序列与注意力：LSTM、多头注意力（Multi-Head Attention，用于 Transformer）；
• 新型架构探索：KAN（基于 B-样条的 Kolmogorov-Arnold 网络，也支持 GR-KAN/KAT、离散 KAN、分段线性 KAN 等变体）、可学习的有理函数激活、用于 SO(3) 等变/不变任务的 VNN（Vector Neural Networks）；
• 归一化新选择：刚刚加入的 DyT（Dynamic Tanh Normalizer），来自较新的研究工作，作为传统归一化层的替代方案；
• 性能相关：新加入的梯度检查点（Gradient Checkpointing），用极简的 API 支持“用重计算换显存”，这是训练大模型时常用的一项工程手段，此前在纯 Go 生态里并不常见。

在工程侧，分布式执行是目前官方标注为“仍在积极改进中”的实验特性：基于 XLA Shardy（GSPMD 分布式方案的演进版本）实现跨多 GPU/TPU 的分布式训练，用户只需要配置好分布式数据集，训练器会自动接管剩下的工作,官方也坦诚这部分还在打磨阶段，欢迎社区反馈问题。

另外还有一个专门的命令行工具gomlx_checkpoints，可以检查训练中/训练完的模型 checkpoint，并用 Plotly 生成损失曲线和评估指标的可视化图表，甚至支持把多个模型的训练曲线放在一起对比——对于需要做大量调参实验的场景很实用。

GoMLX 与 Python 生态的取舍

坦白讲，GoMLX 不是要把 Python 生态“拉下马”，它解决的是另一类问题：当你的服务本身就是用 Go 写的，或者你需要一个单文件、无需安装 Python 解释器和一堆依赖的推理/训练程序时，GoMLX 提供了一条不需要跨语言胶水层的路径。

它的取舍很清楚：

• 性能：底层同样是 XLA，与 Jax、TensorFlow 共用一套编译引擎，在很多场景下能做到同等速度；
• 可移植性：纯 Go 后端几乎可以编译到任何 Go 支持的平台，包括浏览器 WASM，这是 Python 生态很难做到的；
• 代价：生态体量依然无法和 PyTorch/Jax 相比，模型库、论文复现速度、社区资源都还在追赶；API 也更“啰嗦”，这是团队主动选择的结果，为了换取更透明的心智模型。

如果你的场景是研究探索、追求最快原型迭代速度，Python 生态目前仍是更稳妥的选择；但如果你要把模型部署进一个 Go 编写的后端服务，或者想要一个体积小、依赖少、可以编译成单一可执行文件的推理程序，GoMLX 提供的价值就非常直接了。

现状数据与路线图

简单盘点一下现在的项目状态：

• GitHub 星标 **1.5k+**，fork 70+，累计提交超过3500 次，最近一次正式发布是 v0.27.3；
• 独立的官方文档站点已经建成，覆盖入门、核心概念、训练与层、生态整合、进阶调试等完整体系；
• compute 仓库已独立运作，成为可被第三方直接依赖的底层计算引擎；
• 一个由社区维护的 Slack 频道#gomlx、Google Groups 讨论组，以及 GitHub Discussions 都在保持活跃。

官方公开的长期目标可以概括为三条主线：

1. 让 Go 成为训练模型的一等公民，而不只是部署阶段的语言，重点是可读性、清晰可组合的 API、及时更新的文档、有效的错误提示；
2. 成为高效的研究和教学平台，重点支持镜像训练与各种形式的分布式训练（模型并行/数据并行），尤其是面向大语言模型这类超大规模训练场景；
3. 成为可靠的生产平台，包括支持 TPU/GPU 等现代加速器、扩展 XLA 之外的更多后端（llama.cpp、WebNN 等）、支持从 HuggingFace Hub 导入并微调预训练模型，以及未来把模型编译成 C 库或 WebAssembly 供任意语言消费。

写在最后

两年时间，GoMLX 从一个“Go 能不能做机器学习”的验证性项目,成长为一个有清晰分层架构、独立后端引擎、覆盖训练到部署全流程、并且开始积累真实生态组件（HuggingFace、ONNX、Docker 镜像）的框架。compute 仓库的独立，某种程度上标志着这个项目已经从“单体原型”走向了“可持续演进的工程体系”。

它依然不完美：分布式训练还在实验阶段，动态形状支持才刚刚起步，生态体量也远不能和 Python 阵营相提并论。但对于长期用 Go 写后端服务、又想在自己的技术栈内完成模型训练和推理的团队和个人开发者来说，GoMLX 已经从“可以关注一下”变成了“值得认真评估”的选项。

如果你想亲自上手,可以从官方仓库github.com/gomlx/gomlx的 README 开始，里面的 Jupyter 教程和 Docker 镜像能让你在几分钟内跑起第一个模型。

参考链接：

• GoMLX 官方文档：https://gomlx.github.io/docs/overview/
• GoMLX 官网：https://gomlx.github.io/
• GoMLX 主仓库：https://github.com/gomlx/gomlx
• compute 独立仓库：https://github.com/gomlx/compute

🔥 还在为写 Agent 框架频频死循环、上下文爆炸而束手无策？我的新专栏《从0 开始构建 Agent Harness》将带你：

• 抛弃臃肿框架，回归“驾驭工程 (Harness Engineering)”的第一性原理
• 用 Go 语言手写 ReAct 循环、并发拦截与上下文压缩引擎等，复刻极简OpenClaw
• 构建坚不可摧的 Safety Middleware 与飞书人工审批防线
• 在底层实现 Token 成本审计、链路追踪与自动化跑分评估
• 从“调包侠”进化为掌控大模型边界的“AI 操作系统架构师”

扫描下方二维码👇，开启从 0 开始构建Agent Harness 的实战之旅。