Skip to content

Commit 28a908b

Browse files
committed
feat: 反向模式自动微分 + 升 1.0.0
- 自动微分:新增 src/graph/autograd.ts(逆拓扑传播,惰性建图、不碰 GPU), 为内置算子挂上反向公式(src/tensor/ops/backward.ts)。支持 add/sub/mul/div (含标量/广播,经 sumTo 归约)、neg/abs/exp/log/sqrt/square/relu/sigmoid/tanh、 matmul、sum/mean(全局与任意轴)、softmax(末轴,数值稳定)。新增 expand 算子 支撑 sum/softmax 反向的梯度复制。Tensor 增加 withGrad()/.grad/.requiresGrad 与 backward()。 - 文档:README 新增「自动微分」章节,安装段改为已发布,路线图标记 v1.0 与 npm 发布完成;FAQ 更新为「已支持训练」。 - 工程:package.json 升 1.0.0 + publishConfig;新增 .github/workflows/publish.yml, 通过 GitHub OIDC 可信发布到 npm(免长效 token)。 - 测试:新增 tests/unit/autograd.test.ts(形状级,无 GPU)与 tests/gpu/autograd.test.ts(有限差分数值校验,CI/SwiftShader)。 tsc 通过,单元 27 项全绿。 layernorm/embedding/conv1d、pass 融合、Web Worker 运行属 1.0.x 后续规划,本版未做。
1 parent 2e02e7c commit 28a908b

12 files changed

Lines changed: 552 additions & 15 deletions

File tree

‎.github/workflows/publish.yml‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# 通过 GitHub OIDC「可信发布」把包发到 npm —— 全程不需要任何长效 token。
2+
#
3+
# 前置(仅此一次,由仓库拥有者在 npm 网页完成,不传 token):
4+
# 1. 登录 npmjs.com → 进入包 moxwebgpu → Settings → Trusted Publishers
5+
# 2. Add Publisher:
6+
# - Publisher: GitHub Actions
7+
# - Repository: codecloud-dev/moxwebgpu
8+
# - Workflow name: publish.yml
9+
# - Environment: (留空)
10+
# 3. 保存。
11+
# 配置好后,本流程在推送 `v*` tag 时自动 `pnpm publish --provenance`,
12+
# npm 会用 GitHub 临时签发的 OIDC 令牌完成发布(带 provenance 溯源)。
13+
# 若未配置就推了 tag:本流程会失败(无害),去网页配置后「Re-run」该次运行即可。
14+
15+
name: Publish to npm (OIDC)
16+
17+
on:
18+
push:
19+
tags:
20+
- 'v*'
21+
workflow_dispatch:
22+
23+
permissions:
24+
id-token: write # 供 npm OIDC 签发临时令牌
25+
contents: read
26+
27+
jobs:
28+
publish:
29+
runs-on: ubuntu-latest
30+
steps:
31+
- uses: actions/checkout@v4
32+
33+
- uses: pnpm/action-setup@v4
34+
35+
- uses: actions/setup-node@v4
36+
with:
37+
node-version: 20
38+
cache: pnpm
39+
40+
- name: Install
41+
run: pnpm install --frozen-lockfile
42+
43+
- name: Build
44+
run: pnpm build
45+
46+
- name: Publish (provenance, no token)
47+
run: pnpm publish --provenance --no-git-checks --access public

‎README.md‎

Lines changed: 46 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -157,9 +157,9 @@ WebGPU 的 compute pipeline 能力极强,但裸用它做一次向量加法,你
157157

158158
## 安装
159159

160-
> ⚠️ **npm / CDN 尚未发布**:`moxwebgpu` 目前还没有发布到 npm,下方「方式一(CDN)」「方式二(npm)」的 `unpkg` 与 `npm install` **暂不可用**(在线会 404)。请先走「方式三:从源码构建」,或直接在浏览器体验[在线演示](#浏览器演示)。发布后本节会更新。
160+
> ✅ **已发布到 npm**:`moxwebgpu@1.0.0` 已可在 npm 安装,CDN 走 unpkg 即可。本节的「方式一」「方式二」现已可用;发布通过 GitHub OIDC 可信发布完成,无需任何长效 token(见仓库 `.github/workflows/publish.yml`)。
161161
162-
### 方式一:浏览器 `<script>`(零构建,即将发布)
162+
### 方式一:浏览器 `<script>`(零构建)
163163

164164
```html
165165
<script src="https://unpkg.com/moxwebgpu/dist/moxwebgpu.browser.js"></script>
@@ -171,7 +171,7 @@ WebGPU 的 compute pipeline 能力极强,但裸用它做一次向量加法,你
171171

172172
IIFE 产物暴露全局 `MoxWebGPU`,内含 `mox`、`MoxContext`、`Tensor`、`Kernel` 等全部导出。
173173

174-
### 方式二:npm(Node / 打包器,即将发布)
174+
### 方式二:npm(Node / 打包器)
175175

176176
```bash
177177
npm install moxwebgpu
@@ -312,6 +312,43 @@ t.destroy(); // 单个张量释放(通常不需要,池会自动回收中间
312312

313313
---
314314

315+
### 10. 自动微分(反向模式 / 训练)
316+
317+
moxwebgpu 1.0 起内置**反向模式自动微分**,直接叠加在现有惰性计算图上——`OpDef` 只是纯元数据 + codegen,每个算子附带一个 `backward`,反向图天然复用同一套算子,只有当你读取梯度时才上 GPU。
318+
319+
```ts
320+
import { mox } from 'moxwebgpu';
321+
322+
const gpu = await mox.init();
323+
const x = gpu.tensor([1, 2, 3, 4]).withGrad(); // 标记可微叶子
324+
const loss = x.mul(x).sum(); // L = Σ x²
325+
326+
loss.backward(); // 反向传播
327+
console.log(await x.grad.toArray()); // [2, 4, 6, 8] = dL/dx
328+
```
329+
330+
三个要点:
331+
332+
- **`.withGrad()`**:标记该叶子张量需要梯度(其余叶子默认不计)。
333+
- **`.backward()`**:从标量 loss 出发,逆拓扑序传播,把梯度填进每个 `requiresGrad` 叶子的 **`.grad`**。整个过程只建惰性图,**不碰 GPU**,直到你 `await x.grad.toArray()`。
334+
- **广播正确**:二元算子的梯度会用 `sumTo` 自动归约回较小操作数的形状;`sum/mean/softmax` 的梯度用 `expand` 复制回去。
335+
336+
```ts
337+
// 一个最小线性层的前向 + 反向
338+
const W = gpu.tensor([[0.1, 0.2], [0.3, 0.4]]).withGrad();
339+
const inp = gpu.tensor([1, 1]);
340+
const pred = W.matmul(inp); // [2]
341+
const loss = pred.mul(pred).sum(); // Σ pred²
342+
loss.backward();
343+
console.log(await W.grad.toArray()); // dL/dW,形状与 W 一致
344+
```
345+
346+
**已支持反向的算子**:加/减/乘/除(含标量、含广播)、neg/abs/exp/log/sqrt/square/relu/sigmoid/tanh、matmul、sum/mean(全局与任意轴)、softmax(沿末轴,数值稳定)。
347+
348+
> 路线图中 layernorm / embedding / conv1d、pass 融合、Web Worker 运行属于 **1.0.x 后续规划**,本版未做——给后续更新留空间。
349+
350+
---
351+
315352
## 统一 API 参考
316353

317354
### `mox.init(options?)` → `Promise<MoxContext>`
@@ -603,11 +640,11 @@ moxwebgpu/
603640

604641
- [x] **v0.1** —— 三层架构、30+ 算子、27 个真实 GPU e2e、液态玻璃演示
605642
- [x] **v0.2** —— 任意轴归约(任意 rank/负轴)、整型归约修复、调度器 temp 回收加固
606-
- [ ] 自动微分(反向图叠加在现有 lazy graph 上)
607-
- [ ] 更多 NN 算子:layernorm / embedding / conv1d
608-
- [ ] 相邻逐元素算子的 pass 融合(进一步压 dispatch 次数)
609-
- [ ] Web Worker / OffscreenCanvas 内运行
610-
- [ ] npm 首次发布(即将发布)
643+
- [x] **v1.0** —— 反向模式自动微分(叠加在现有 lazy graph 上)、matmul/softmax/sum/mean 反向、OIDC 免 token 发布到 npm
644+
- [ ] 更多 NN 算子:layernorm / embedding / conv1d(规划中,1.0.x)
645+
- [ ] 相邻逐元素算子的 pass 融合(进一步压 dispatch 次数,规划中)
646+
- [ ] Web Worker / OffscreenCanvas 内运行(规划中)
647+
- [x] npm 首次发布
611648
- [x] 文档站上线
612649

613650
> 有想要的功能或方向,欢迎提 Issue 一起讨论。
@@ -626,7 +663,7 @@ A:99% 是 WGSL 编译失败(静默 no-op)。moxwebgpu 已把编译错误打印
626663
A:它们是「模型中心」:面向推理预置模型。moxwebgpu 是「算子中心」:给你 NumPy 式的原始计算能力 + 逃生舱,恰好可以作为它们没有的那层「通用 GPGPU 地基」。
627664

628665
**Q:支持训练(反向传播)吗?**
629-
A:暂不支持,自动微分在路线图第一位(现有 lazy graph 天然适合叠加反向图)。
666+
A:支持。1.0.0 起内置反向模式自动微分:对叶子张量调 `.withGrad()`,正向建图后调 `.backward()`,梯度即填到 `.grad`(惰性,需要时才上 GPU)。详见[自动微分](#10-自动微分反向模式--训练)。
630667

631668
**Q:为什么我的机器跑 `pnpm test:gpu` 也能过?我没有 N 卡。**
632669
A:因为 SwiftShader——Chrome 自带的纯软件 Vulkan 实现。moxwebgpu 的 GPU 测试配方**不要求真显卡**,CI 上也一样。

‎package.json‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
{
22
"name": "moxwebgpu",
3-
"version": "0.2.0",
4-
"description": "⚡ A WebGPU compute (GPGPU) framework — typed tensors, raw kernels and a lazy compute graph, in one tiny package.",
3+
"version": "1.0.0",
4+
"description": "⚡ A WebGPU compute (GPGPU) framework — typed tensors, raw kernels and a lazy compute graph, with reverse-mode autodiff, in one tiny package.",
55
"type": "module",
66
"main": "./dist/index.cjs",
77
"module": "./dist/index.js",
@@ -15,8 +15,13 @@
1515
"./browser": "./dist/moxwebgpu.browser.js"
1616
},
1717
"files": [
18-
"dist"
18+
"dist",
19+
"README.md",
20+
"LICENSE"
1921
],
22+
"publishConfig": {
23+
"access": "public"
24+
},
2025
"scripts": {
2126
"build": "tsup",
2227
"test": "vitest run --config vitest.config.ts",

‎src/graph/autograd.ts‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
/**
2+
* Reverse-mode automatic differentiation.
3+
*
4+
* The gradient graph is built lazily on top of the existing lazy DAG: every
5+
* `BackwardFn` returns new lazy Tensors via the normal op set, so no GPU work
6+
* happens until a gradient is materialized with `.toArray()` / `.toBuffer()`.
7+
*
8+
* `backward(loss)` fills `.grad` on every leaf Tensor whose `requiresGrad` is
9+
* true. The loss is treated as a scalar (seed = 1); pass a pre-reduced scalar
10+
* loss for the usual "dL/dparams" semantics.
11+
*/
12+
13+
import { Tensor } from '../tensor/tensor.js';
14+
import type { MoxContext } from '../core/context.js';
15+
import type { LazyNode } from './lazy.js';
16+
import { numElements } from '../core/dtype.js';
17+
18+
/** Reverse topological order (parents before children); does not mutate state. */
19+
function reverseTopo(root: LazyNode): LazyNode[] {
20+
const order: LazyNode[] = [];
21+
const seen = new Set<LazyNode>();
22+
const visit = (n: LazyNode): void => {
23+
if (seen.has(n)) return;
24+
seen.add(n);
25+
for (const inp of n.inputs) {
26+
if (inp.kind !== 'leaf') visit(inp);
27+
}
28+
order.push(n);
29+
};
30+
visit(root);
31+
return order;
32+
}
33+
34+
/** A lazy ones-tensor shaped like `t` (no GPU upload until materialized). */
35+
function onesLike(t: Tensor): Tensor {
36+
return t.toFloat().mul(0).add(1);
37+
}
38+
39+
/**
40+
* Reduce `t` to `shape` by summing out broadcasted / leading dimensions.
41+
* Used to push a gradient back to a smaller broadcast operand's shape.
42+
*/
43+
export function sumTo(t: Tensor, shape: number[]): Tensor {
44+
let cur = t;
45+
let curShape = cur.shape;
46+
while (curShape.length > shape.length) {
47+
cur = cur.sum(0);
48+
curShape = cur.shape;
49+
}
50+
const outShape: number[] = [];
51+
for (let d = 0; d < shape.length; d++) {
52+
if (curShape[d] !== shape[d]) {
53+
cur = cur.sum(d);
54+
curShape = cur.shape;
55+
outShape.push(1);
56+
} else {
57+
outShape.push(curShape[d]);
58+
}
59+
}
60+
return cur.reshape(...outShape);
61+
}
62+
63+
function wrap(node: LazyNode, ctx: MoxContext): Tensor {
64+
return new Tensor(ctx, node.shape, node.dtype, node, null);
65+
}
66+
67+
export function backward(loss: Tensor): void {
68+
const root = loss.node;
69+
if (!root) return; // leaf loss: nothing to differentiate
70+
71+
const order = reverseTopo(root);
72+
const gradOf = new Map<object, Tensor>();
73+
gradOf.set(root, onesLike(loss));
74+
75+
for (let i = order.length - 1; i >= 0; i--) {
76+
const node = order[i];
77+
const g = gradOf.get(node);
78+
if (!g) continue;
79+
80+
const inputs: Tensor[] = node.inputs.map((inp) => {
81+
if (inp.kind === 'leaf') return inp.tensor as Tensor;
82+
return wrap(inp as LazyNode, loss.ctx);
83+
});
84+
const output = wrap(node, loss.ctx);
85+
86+
const fn = node.op.backward;
87+
const grads: (Tensor | null)[] = fn ? fn(inputs, output, g, node.attrs) : inputs.map(() => null);
88+
89+
for (let j = 0; j < node.inputs.length; j++) {
90+
const gi = grads[j];
91+
if (!gi) continue;
92+
let target: object;
93+
if (node.inputs[j].kind === 'leaf') {
94+
target = (node.inputs[j] as { kind: 'leaf'; tensor: Tensor }).tensor;
95+
} else {
96+
target = node.inputs[j] as LazyNode;
97+
}
98+
const prev = gradOf.get(target);
99+
gradOf.set(target, prev ? prev.add(gi) : gi);
100+
}
101+
}
102+
103+
for (const [key, g] of gradOf) {
104+
if (key instanceof Tensor && key.requiresGrad) {
105+
key.grad = g;
106+
}
107+
}
108+
}

‎src/graph/lazy.ts‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,23 @@
99

1010
import type { DType } from '../core/dtype.js';
1111
import type { GpuDataBuffer } from '../core/buffer.js';
12+
import type { Tensor } from '../tensor/tensor.js';
13+
14+
/**
15+
* Reverse-mode gradient of an op.
16+
* - `inputs` : the original input Tensors (leaves or node-backed), in order
17+
* - `output` : a Tensor view of this node's output
18+
* - `gradOutput`: upstream gradient w.r.t. this node's output
19+
* - returns one gradient Tensor (or null) per input, matching input order.
20+
* Gradients are themselves lazy tensors, so the whole backward graph reuses
21+
* the existing op set and only touches the GPU when materialized.
22+
*/
23+
export type BackwardFn = (
24+
inputs: Tensor[],
25+
output: Tensor,
26+
gradOutput: Tensor,
27+
attrs: any,
28+
) => (Tensor | null)[];
1229

1330
/**
1431
* One GPU compute pass inside an op's execution plan.
@@ -58,6 +75,8 @@ export interface OpDef {
5875
outDtype?: (inputDtypes: DType[], attrs: any) => DType;
5976
/** Build the GPU execution plan for this node. */
6077
build: (inputShapes: number[][], inputDtypes: DType[], attrs: any) => OpPlan;
78+
/** Optional reverse-mode gradient (see BackwardFn). Absent = no gradient. */
79+
backward?: BackwardFn;
6180
}
6281

6382
/** A downstream reference: either another lazy node or a leaf tensor. */

‎src/index.ts‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,19 +22,25 @@ export {
2222

2323
// ---- tensor ----
2424
export { Tensor, type TensorLike, type TensorOptions } from './tensor/tensor.js';
25-
export type { OpDef, KernelStep, LazyNode, OpPlan, BindingSpec } from './graph/lazy.js';
25+
export type { OpDef, KernelStep, LazyNode, OpPlan, BindingSpec, BackwardFn } from './graph/lazy.js';
2626
export { createNode, topoSort } from './graph/lazy.js';
2727
export { Scheduler } from './graph/scheduler.js';
2828
export { PipelineCache } from './graph/pipelineCache.js';
2929

30+
// ---- autograd ----
31+
export { backward, sumTo } from './graph/autograd.js';
32+
3033
// ---- codegen (for building custom ops) ----
3134
export * from './tensor/codegen.js';
3235

3336
// ---- version ----
34-
export const MOXWEBGPU_VERSION = '0.2.0';
37+
export const MOXWEBGPU_VERSION = '1.0.0';
3538

3639
import { Tensor } from './tensor/tensor.js';
3740
import { installTensorOps } from './tensor/ops/index.js';
41+
import { installBackward } from './tensor/ops/backward.js';
3842

3943
// Install the unified chainable op API (add/matmul/sum/softmax/...) once.
4044
installTensorOps(Tensor.prototype);
45+
// Attach reverse-mode gradients to the built-in ops.
46+
installBackward();

0 commit comments

Comments
 (0)