diff --git a/build.zig b/build.zig index 08765bd1df..31ce3658bb 100644 --- a/build.zig +++ b/build.zig @@ -192,7 +192,7 @@ pub fn build(b: *std.Build) void { } /// `zig build native -Dcpu=apple_m1`: tensorfold-native with the Metal engines for the Python package's bundle (a native M5 build traps on M1-M4). -fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, metal: *std.Build.Module, engine: *std.Build.Module, api: *std.Build.Module, engines: *std.Build.Module, build_options: *std.Build.Step.Options, test_step: ?*std.Build.Step) *std.Build.Step.Compile { +fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, metal: *std.Build.Module, engine: *std.Build.Module, living: *std.Build.Module, api: *std.Build.Module, engines: *std.Build.Module, build_options: *std.Build.Step.Options, test_step: ?*std.Build.Step) *std.Build.Step.Compile { // the HTTP side keeps its safety checks; the engine below it runs at `optimize` const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); @@ -201,7 +201,7 @@ fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b .target = target, .optimize = .ReleaseSafe, .link_libc = true, - .imports = &.{ .{ .name = "engine_api", .module = api }, .{ .name = "tokenizer", .module = tokenizer }, .{ .name = "template", .module = template }, .{ .name = "native_engines", .module = engines }, .{ .name = "checkpoint_cli", .module = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cli.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true, .imports = &.{ .{ .name = "native_engines", .module = engines }, exlr8_build.cliImport(b, target, .ReleaseSafe, engine) } }) } }, + .imports = &.{ .{ .name = "engine_api", .module = api }, .{ .name = "tokenizer", .module = tokenizer }, .{ .name = "template", .module = template }, .{ .name = "native_engines", .module = engines }, .{ .name = "living", .module = living }, .{ .name = "checkpoint_cli", .module = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cli.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true, .imports = &.{ .{ .name = "native_engines", .module = engines }, exlr8_build.cliImport(b, target, .ReleaseSafe, engine) } }) } }, }) }); exe.root_module.addOptions("build_options", build_options); if (test_step == null) { @@ -219,7 +219,7 @@ fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b .target = target, .optimize = .Debug, // the server's unit tests run with safety checks, as zig_test.sh builds them .link_libc = true, - .imports = &.{ .{ .name = "engine_api", .module = api }, .{ .name = "tokenizer", .module = tokenizer }, .{ .name = "template", .module = template } }, + .imports = &.{ .{ .name = "engine_api", .module = api }, .{ .name = "tokenizer", .module = tokenizer }, .{ .name = "template", .module = template }, .{ .name = "living", .module = living } }, }); test_step.?.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = server_tests })).step); @import("zig/build/cortheon_fuzz.zig").add(b, target, api, tokenizer, template); @@ -466,7 +466,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b glm_checks.targets(b, target, optimize, metal, engine, sources); flashnext_checks.targets(b, target, optimize, metal, engine); - _ = nativeServer(b, target, optimize, metal, engine, api, engines, build_options, test_step); + _ = nativeServer(b, target, optimize, metal, engine, mods.living, api, engines, build_options, test_step); const cluster_metal = b.createModule(.{ .root_source_file = b.path("zig/src/cluster/metal_sink.zig"), @@ -539,7 +539,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b } /// The server and distribution use the same embedded-source Metal engine. -fn metalEngineModules(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, draft_ids: *std.Build.Module) struct { metal: *std.Build.Module, lanes: *std.Build.Module, fabric: *std.Build.Module, sources: *std.Build.Module, engine: *std.Build.Module, qwen27: qwen27_build.Shared } { +fn metalEngineModules(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, draft_ids: *std.Build.Module) struct { metal: *std.Build.Module, lanes: *std.Build.Module, fabric: *std.Build.Module, sources: *std.Build.Module, engine: *std.Build.Module, living: *std.Build.Module, qwen27: qwen27_build.Shared } { const metal = b.createModule(.{ .root_source_file = b.path("zig/src/metal/metal.zig"), .target = target, @@ -553,14 +553,16 @@ fn metalEngineModules(b: *std.Build, target: std.Build.ResolvedTarget, optimize: const lanes = b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = target, .optimize = optimize, .link_libc = true }); const fabric = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/fabric.zig"), .target = target, .optimize = optimize, .link_libc = true }); const sources = b.createModule(.{ .root_source_file = b.path("zig/kernels/metal/sources.zig") }); + // Living Weights' family-neutral learner, which Nemotron's lessons and the server's view of them share + const living = b.createModule(.{ .root_source_file = b.path("zig/src/living/living.zig"), .target = target, .optimize = optimize, .link_libc = true }); const engine = b.createModule(.{ .root_source_file = b.path("zig/src/tensorfold.zig"), .target = target, .optimize = optimize, .link_libc = true, - .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "kernel_sources", .module = sources }, .{ .name = "nemotron_draft_ids", .module = draft_ids }, .{ .name = "lanes", .module = lanes }, .{ .name = "fabric", .module = fabric } }, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "kernel_sources", .module = sources }, .{ .name = "nemotron_draft_ids", .module = draft_ids }, .{ .name = "lanes", .module = lanes }, .{ .name = "fabric", .module = fabric }, .{ .name = "living", .module = living } }, }); - return .{ .metal = metal, .lanes = lanes, .fabric = fabric, .sources = sources, .engine = engine, .qwen27 = qwen27_build.engineModules(b, target, metal, sources, engine) }; + return .{ .metal = metal, .lanes = lanes, .fabric = fabric, .sources = sources, .engine = engine, .living = living, .qwen27 = qwen27_build.engineModules(b, target, metal, sources, engine) }; } pub fn distMetalServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_ids: *std.Build.Module, build_options: *std.Build.Step.Options, sdk: []const u8) *std.Build.Step.Compile { @@ -574,7 +576,7 @@ pub fn distMetalServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_id m.metal.addSystemIncludePath(b.graph.cwdRelativePath(b.pathJoin(&.{ sdk, "usr/include" }))); const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = m.lanes }} }); const engines = b.createModule(.{ .root_source_file = b.path("zig/src/native/metal.zig"), .target = target, .optimize = .fast, .link_libc = true, .imports = &.{ .{ .name = "metal", .module = m.metal }, .{ .name = "engine_api", .module = api }, .{ .name = "tensorfold", .module = m.engine } } }); - return nativeServer(b, target, .fast, m.metal, m.engine, api, engines, build_options, null); + return nativeServer(b, target, .fast, m.metal, m.engine, m.living, api, engines, build_options, null); } test { diff --git a/docs/recipes/README.md b/docs/recipes/README.md index 1e38e07605..8d77574019 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -29,8 +29,9 @@ from its family pages, not measured per card. | MLX, Apple Silicon | GLM-5.3-Flash is written for a 256 GB Mac, about 151 GiB resident. DeepSeek-V4-Flash is the same, about 151 GiB resident. Flash Next's default command sizes to a 128 GiB M4 Max. Qwen3.8-27B on a 32 GB Mac needs more than the default 22.4 GiB. Machine classes, not measured minimums. | New families go on the Zig engine: [adding a Zig family](adding-a-zig-family.md) has the steps in order, what each -one gained, and which code owns it. For the Python engine, frozen at 0.6.6, the contributor guides cover -[adding an MLX family](adding-a-family.md), [adding a CUDA family](adding-a-cuda-family.md) and +one gained, and which code owns it. [Adding Living Weights to a family](adding-living-weights-to-a-family.md) lists +what a family's host and trainer give the learner for `--slide`. For the Python engine, frozen at 0.6.6, the +contributor guides cover [adding an MLX family](adding-a-family.md), [adding a CUDA family](adding-a-cuda-family.md) and [CUDA implementation rules](cuda.md). [EXL3 weights](exl3.md) and [universal EXL3 experts](exl3-universal-experts.md) describe the shared EXL3 module @vcruz305 wrote for the Python engine's CUDA families in [#42](https://github.com/ashhart/TensorFold/pull/42): diff --git a/docs/recipes/adding-living-weights-to-a-family.md b/docs/recipes/adding-living-weights-to-a-family.md new file mode 100644 index 0000000000..65d3ab9b5f --- /dev/null +++ b/docs/recipes/adding-living-weights-to-a-family.md @@ -0,0 +1,201 @@ +# Adding Living Weights to a family + +This is how a family on the Zig engine learns facts with `--slide`. [Living Weights](../living-weights.md) is one +learner for every family (`zig/src/living`): it writes lessons, reads the model, fits a pocket of new neurons and saves +them into the checkpoint's own shards. A family describes itself as data, and its host and trainer give the learner a +few hooks. The steps are in order. It comes from three families: + +- Nemotron 3.5 Lightning on Metal and CUDA, the first family and the only one listed in `living/families.zig` today. + Its relu^2 shared experts carry the pocket at the MoE layers. +- Qwen3.8 Flash Next on CUDA (#895 and #896). A SwiGLU shared expert behind its own gate at every layer, with + hyper-connections, Gated DeltaNet and gated attention below it. +- Qwen3.8-27B on CUDA, our port on #726 (not sent yet). A dense SwiGLU MLP at every layer, with no experts at all. + +Bring the family up on the Zig engine first, exact and at speed ([adding a Zig family](adding-a-zig-family.md)). The +learner trains through the engine's own forward, so everything below assumes that forward is done. + +## How to use it + +1. Do the steps in order. Each ends in a check that has to pass before the next step means anything. +2. Never teach the original folder. Teach a `cp -RL` copy (`cp -cRL` on a Mac), and check the original's sha256 + after each run. +3. Put the receipts at the end of this page in the PR: an untaught folder serves as before, the trainer's forward is + the serving forward, lessons are kept, and a taught folder serves. + +## 1. Describe the family as data + +Add a `Family` to `living/families.zig`. The learner, the shard writer and the server's Living Weights view all read +it, by the checkpoint's `model_type`. + +| Field | Nemotron (`nemotron_h`) | Flash Next (`qwen4_exp`) | Qwen3.8-27B (`qwen3_5`) | +| --- | --- | --- | --- | +| `names.layers`, `names.expert` | `backbone.layers.`, `mixer.shared_experts.` | `language_model.model.layers.`, `mlp.shared_expert.` | `language_model.model.layers.`, `mlp.` | +| `names.inputs` | `up_proj` | `gate_proj`, `up_proj` | `gate_proj`, `up_proj` | +| `names.mtp` | its own file, `mtp-4bit.safetensors` | in the index, `language_model.mtp.layers.0.mlp.shared_expert.` | none | +| `neuron` | relu^2, margin 0.25 | SwiGLU, gated, margin 0.5, sharp 16 | SwiGLU, margin 0.5, sharp 16 | +| `quant` | 4-bit, groups of 64 | 4-bit, groups of 32 | 4-bit, groups of 64 | +| `keys` | `moe_shared_expert_intermediate_size` at the top level | `shared_expert_intermediate_size` in `text_config` | `intermediate_size` in `text_config` | +| `experts` | `.pattern`: the MoE layers of `hybrid_override_pattern` | `.every` | `.every` | + +**Where the pocket goes.** A pocket is 64 neurons a layer (`Family.neurons`, a whole number of quantization groups), +zero until a fact uses one. Each is one more row of every input projection and one more column of the output +projection, so the first save makes the projection 64 neurons wider and any MLX loader reads a slightly wider model. + +- **An MoE layer with a shared expert:** beside the shared expert, as Nemotron and Flash Next do. The routed experts + are left alone: a pocket there would fire only when its expert is picked. +- **A shared expert behind a gate:** set `neuron.gated`. Flash Next scales its shared expert by + `sigmoid(shared_expert_gate x)`, so a pocket neuron adds `gate act(x) v` and its output is fitted through the gate. + The trainer reports the gate at each row (`gateAt`). +- **A dense MLP:** the MLP is the expert (`expert = "mlp."`, `experts = .every`), and the width key is the MLP's + (the 27B's 17,408 neurons become 17,472). + +**The neuron's kind.** A relu^2 neuron is exactly zero below its threshold. A SwiGLU neuron is not: `silu` leaks +below zero. A SwiGLU neuron fitted as `u` is written as gate row `sharp u` and up row `u / sharp` +(`living/neuron.zig`), which is relu(u x)^2 within 1e-7 on its rows, and its margin is 0.5 so that what leaks there +is at most `sigmoid(-8)`, 3e-4, of it. Keep `sharp` a power of two, so `u` comes back exactly from the up row. + +`Family.check()` refuses a description whose inputs do not match the neuron, or whose pocket is not a whole number of +groups. Add your family to the tests in `families.zig` and `family.zig`: its model type is found, and a tensor name of +its expert gives the layer back while a similar one (a gate, the MTP head's) does not. + +## 2. The pocket in every pass + +The engine adds the pocket wherever the model runs: prompt chunks, verify windows and decode rounds. Load it from the +shards as fp32 rows beside the packed weights (Flash Next: `cuda_pocket.zig`; the 27B: its `pocket.zig`), and take +the shard's projection apart at the old width. + +- **An untaught folder's bits do not change.** A layer with no neuron in use launches nothing. Flash Next sums a + pocket layer's MoE output in fp32, adds the pocket and rounds once; every other layer writes back exactly as before. + Check the untaught folder against the last release: the same replies, drafted and `"draft": false`, and the same + speed. +- **Row by row.** The pocket's kernels read each row alone, so drafted still equals `"draft": false` and a concurrent + stream equals its solo run. Flash Next uses the two-kernel form of #787 (`tf_fn_pocket_in`, then `_out`). +- **Untaught streams.** A request with `untaught` set (the lesson writer's and the judge's, in + `server/slide_ask.zig`) reads the model as shipped. Zero the pocket's output on its rows. The 27B missed this at + first: the judge then read the lesson's own neurons, called a reworded "I don't know" a changed answer, and took a + good lesson out. +- **The config check accepts a taught folder.** If the engine checks the checkpoint's widths at load, subtract + `tensorfold_learned_neurons` from the expert's width first. Flash Next's check refused its own taught folder at the + first restart. +- **The MTP head.** The first save widens the MTP head's shared expert with zero neurons, so its tensors stay one + width. Drafts are made without the pocket; the verify rows run it, so drafted replies still equal plain ones. + +## 3. The host + +The host is a type the learner is generic over (`Learner(H)` in `living/learner.zig`). A family's slide file +(Nemotron's `cuda_slide.zig`, Flash Next's in #896) declares it. + +- **Always:** `Trainer`, `max_rows`, `Ctx`, `family`, `layout(ctx)`, `trainer(gpa, io, ctx)`, `free(t, gpa)` and + `seconds(io)`. +- **The pocket, relu^2:** `pocket(ctx, layer, up, down, used)` hands the forward a layer's neurons; + `pocketRead(ctx, layer, up, down) usize` gives back the neurons loaded from the shards; and + `pocketSave(ctx, gpa, io, layers, neurons, up, down) usize` writes every layer and hands the rows back as written. +- **The pocket, SwiGLU:** `pocketSwiglu`, `pocketReadSwiglu` and `pocketSaveSwiglu`, the same three with a `gate` + beside `up`. +- `up` and `gate` are each neuron's input rows as the forward multiplies them (`[neurons, D]` f32), and `down` is its + output. A save hands the rows back rounded as the shards hold them (#877 checks every lesson on those). +- **Optional:** `attach`, `write` and `layers` (Nemotron's routed blocks), and `codes`, `commit` and `fold4` + (weight edits). The learner offers them only when the trainer has `sites`. + +`layout(ctx)` is a `Layout`: the hidden size, the expert's own width with any saved pocket left out, the layer count, +the layers whose expert carries the pocket (`widen`) and the ones lessons put neurons in (`learn`). +`Layout.fromExperts` widens every layer that has an expert and learns at all but the last, whose nudge has no +gradient to read. Read the widths from `config.json` at load, as Flash Next's `layoutOf` does. + +## 4. The trainer + +The trainer runs the learner's passes on the engine: a keys pass (the rows a pocket reads), a loss pass, and a target +pass that runs the loss's gradient back to the lowest nudged layer. + +- **What it gives:** `normed(layer, row)` (a row's input to the expert, as the forward normed it), `pass(seqs, mode, + turn, out)` and `step(ids, start, mode, turn)`, `spans`, `packs()`, `pack`, `count`, `prefixes`, + `prefixSlots(n)`, `pushAt(i, key, layer)` (the gradient's push at a nudged row), `deltaAt(i, key, layer)`, + `residual(layer, row)`, `nudge` and `from`. A gated family also gives `gateAt(layer, row)`. +- **Optional:** `tapRows` (#891) reads a question asked after a conversation through the engine's own prompt chunks, + so a pocket can read it in context. Nemotron's routed blocks, edits and probes also use `sites`, `sketched`, + `score()`, `projected()`, `key()` and `config()`. +- **Sequences side by side.** A pass packs several sequences into one forward (Flash Next 1,024 rows, the 27B 1,024 + rows over up to 16 sequences). Row-wise launches run over every row; a recurrent mixer or attention runs each + sequence from position 0. A sequence's loss in a pack equals its loss alone, exactly. +- **What the forward keeps.** The backward reads what the forward computed: each layer's expert input and its routing, + and, from the lowest layer that learns, every mixer's projections, the norms' inputs, the gates. Keep them in the + engine's forward rather than in a second forward. Flash Next's batched pass keeps them on request (`Keep`): every + layer's expert input, and the rest from `layers / 5` up, as the 27B's does. +- **The backward kernels.** One input-gradient kernel set for each layer type above the lowest learned layer, checked + against float64 with central differences. What another family could use goes in the core, a family's own in its + folder (#895): + - core: a Gated DeltaNet layer (`cuda.gdn_train`), 4-bit products at the checkpoint's group, routed SwiGLU experts + and top-k routing (`cuda.train_q4`), gated attention (`cuda.train_attn`), Nemotron's Mamba (`train.cu`); + - family: Flash Next's hyper-connections and its gated pocket (`flashnext_train.cu`); the 27B's dense SwiGLU MLP and + its ungated pocket. +- **Memory.** On a GB10 the trainer keeps MLX copies of the small linears its backward reads, and its keep lists. It + took about 26 GiB for Nemotron, 10.6 GiB for Flash Next and 23 GiB for the 27B. `open` reports it as + `learn_bytes`, and the native host keeps it out of the streams' room. + +**Checks.** Add a `tf-cuda-test` case for each backward kernel (Flash Next: `gdn-back`, `q4-back`, `moe-back`, +`attn-back`, `fn-hc-back`, `fn-moe-back`, `fn-pocket`), and a tool that runs the trainer on the real checkpoint +(`flashnext-learn`): + +- the trainer's keys and loss passes equal the serving forward byte for byte, untaught and with a pocket neuron; +- a row's scores never depend on later rows; +- a pack of sequences gives each one's loss exactly as alone; +- a pocket neuron moves the forward by exactly `gate act v` at its row. + +A finite difference against the gradient is noisy in bf16. Below about one ulp of the residual the rounding decides, +and a large step leaves the linear range. Flash Next's direct path agreed within a median of 6% in three windows of +four, and its cross-row paths stayed noisy. The 27B's direct rows gave 2-36%. Treat it as a smoke check: the kernels' +float64 tests and the identity checks above are the real ones, and a lesson that is kept is the end-to-end one. + +## 5. The shard writer + +`living/shards.zig` writes any family's pocket: `Io(core).write(gpa, io, dir, family, layout, layers, neurons, inputs, +down)` rounds the rows to the family's affine format (`living/affine.zig`), widens every expert on the first save +(the MTP head's with zeros), records `tensorfold_learned_neurons` in the family's key scope, and writes a later save in +place. Wrap it in the family's folder (Flash Next's `living.zig`, `Shards(core)`). + +Test it on a small fake checkpoint (Nemotron's and Flash Next's `pocket_shards_test.zig`): the first save widens every +expert and tells `config.json`; a later save lands in place; every other tensor's bytes are kept; and the rows read +back are the rows written. + +## 6. `--slide` on the native host + +The native host (`native/cuda.zig`) serves `--slide` for a family whose module declares: + +- `slide`, a type whose `Learner` is `living.learner.Learner(Host)`, with `Model`, `Example`, `Lesson`, `Report` and + `Step`; +- `learnModel(ctx, dir)`, the learner's `Model` for a loaded engine; +- `linked(arena, io, dir)`, the first file a lesson would rewrite (config, index, a shard holding an expert) that is a + link, so a folder of links into the Hugging Face cache is refused at startup. + +`open` with `.slide = true` reserves the trainer's rows in the pass's scratch and reads the layout from `config.json`. +The host already refuses `--slide` off a GB10 (sm_121) and with `--segments`. A family adds its own refusals: Flash +Next refuses a captured kernel set, whose read-out kernels keep nothing a backward reads. Add each refusal's message +to the table in [Living Weights](../living-weights.md), and a line for the family under "What you need". + +## Receipts to produce + +On a fresh copy for each run, greedy, with the checkpoint's revision, the GPU, driver and CUDA version: + +- **Untaught, as before.** Against the last release on the same machine: the same replies drafted and with + `"draft": false`, decode and cold-prompt speed, and the model's memory. +- **Trainer.** The kernel checks and the identity checks of step 4, and the finite-difference numbers as measured. +- **Lessons.** A handful of facts about the user, one learn request each (Flash Next taught five, Nemotron ten). For + each: kept or not, rounds and steps, seconds, and its held-out wordings right after its lesson and after the last. + Ask wordings no lesson trained on. Then a fixed set of unrelated prompts against the untaught replies: #891 used + twelve (capitals, arithmetic, a haiku, a Python one-liner, a translation and "What is my cat called?"). +- **A taught folder serves.** Restart without `--slide` and with it: the same replies as after the last lesson. + Drafted equals `"draft": false`, concurrent equals solo, a resumed prompt equals a fresh one, and the decode cost of + the pocket (Flash Next: 5% drafted and 1% plain with three facts). +- **Tests.** `zig build test-cuda-host` and `zig build test-living` from a fresh cache, `zig fmt --check build.zig zig`, + and `python3 tools/zig/lean_check.py` with no new problems. + +## Traps + +- **A finished lesson can still hurt later ones.** A pocket neuron's output reaches higher layers, so a lower pocket + neuron can push a higher one past its threshold. On the 27B, after 16 facts, seven of 40 wide unrelated prompts gave + garbage tokens: one neuron leaking +0.4 on LaTeX digits set off about 100 others, layer by layer. The rows that set it + off were in no prompt a lesson read. Read unrelated prompts in several languages, code, math and tables while you + teach, and check them after every lesson, not only at the end. +- **The checks only see what they ask.** Lessons ask single-turn questions. A fact can come back cold and miss after a + long conversation (#891). +- **A captured kernel set keeps nothing.** If a family still has a captured comparison set, `--slide` must refuse it. diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index eac681628c..d36ee9c52b 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -78,19 +78,22 @@ fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti return cuda; } -fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module } { +fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, draft_ids: *std.Build.Module) struct { core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, heat: *std.Build.Module, living: *std.Build.Module } { const tokenizer = b.createModule(.{ .root_source_file = b.path("zig/src/core/tokenizer/tokenizer.zig"), .target = target, .optimize = optimize, .link_libc = true }); const core = b.createModule(.{ .root_source_file = b.path("zig/src/core/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); core.addImport("tokenizer", tokenizer); const lanes = b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = target, .optimize = optimize, .link_libc = true }); const heat = b.createModule(.{ .root_source_file = b.path("zig/src/core/heat.zig"), .target = target, .optimize = optimize, .link_libc = true }); + // Living Weights' family-neutral learner: a family's lessons, and the server's view of them, import it + const living = b.createModule(.{ .root_source_file = b.path("zig/src/living/living.zig"), .target = target, .optimize = optimize, .link_libc = true }); const nemotron = b.createModule(.{ .root_source_file = b.path("zig/src/families/nemotron/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true }); nemotron.addImport("cuda", cuda); nemotron.addImport("core", core); nemotron.addImport("lanes", lanes); nemotron.addImport("heat", heat); nemotron.addImport("nemotron_draft_ids", draft_ids); - return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .tokenizer = tokenizer, .heat = heat }; + nemotron.addImport("living", living); + return .{ .core = core, .lanes = lanes, .nemotron = nemotron, .tokenizer = tokenizer, .heat = heat, .living = living }; } /// Linux targets: fatbins (-Dnvcc builds them, -Dfatbins embeds prebuilt ones), `tensorfold` and `tf-cuda-test`. @@ -139,7 +142,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu runner.addImport("exlr8_frozen_dense", b.createModule(.{ .root_source_file = b.path("zig/tests/fixtures/exlr8/frozen_dense.zig"), .target = target, .optimize = optimize })); for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); b.installArtifact(b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner })); - nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, true).root_module.strip = strip; + nativeServer(b, target, optimize, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, mods.living, build_options, true).root_module.strip = strip; } /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families. @@ -161,7 +164,7 @@ fn checkpointCli(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std. } /// `zig build native`: tensorfold-native with the CUDA engines into zig-out/native/bin, as the Metal build makes it. -fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool) *std.Build.Step.Compile { +fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, tokenizer: *std.Build.Module, living: *std.Build.Module, build_options: *std.Build.Step.Options, install_native: bool) *std.Build.Step.Compile { const m = engines(b, target, optimize, cuda, lanes, nemotron); // the HTTP side keeps its safety checks; the engine below it runs at `optimize` (the tokenizer is the family's) const template = b.createModule(.{ .root_source_file = b.path("zig/src/core/template/template.zig"), .target = target, .optimize = .ReleaseSafe, .link_libc = true }); @@ -170,7 +173,7 @@ fn nativeServer(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b .target = target, .optimize = .ReleaseSafe, .link_libc = true, - .imports = &.{ .{ .name = "engine_api", .module = m.api }, .{ .name = "tokenizer", .module = tokenizer }, .{ .name = "template", .module = template }, .{ .name = "native_engines", .module = m.engines }, .{ .name = "checkpoint_cli", .module = checkpointCli(b, target, .ReleaseSafe, null, m.engines, core) } }, + .imports = &.{ .{ .name = "engine_api", .module = m.api }, .{ .name = "tokenizer", .module = tokenizer }, .{ .name = "template", .module = template }, .{ .name = "native_engines", .module = m.engines }, .{ .name = "living", .module = living }, .{ .name = "checkpoint_cli", .module = checkpointCli(b, target, .ReleaseSafe, null, m.engines, core) } }, }) }); exe.root_module.addOptions("build_options", build_options); if (!install_native) { @@ -199,7 +202,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st .target = host, .optimize = .debug, .link_libc = true, - .imports = &.{ .{ .name = "engine_api", .module = native_modules.api }, .{ .name = "tokenizer", .module = mods.tokenizer }, .{ .name = "template", .module = template } }, + .imports = &.{ .{ .name = "engine_api", .module = native_modules.api }, .{ .name = "tokenizer", .module = mods.tokenizer }, .{ .name = "template", .module = template }, .{ .name = "living", .module = mods.living } }, }); const server_test = b.addRunArtifact(b.addTest(.{ .root_module = server_tests })); step.dependOn(&server_test.step); @@ -211,6 +214,15 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.heat, mods.tokenizer, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + // Living Weights: the family-neutral learner, then the pocket in each family's shards over the core's shard files + const living_run = b.addRunArtifact(b.addTest(.{ .root_module = mods.living })); + const shards = b.createModule(.{ .root_source_file = b.path("zig/src/living_shards_test.zig"), .target = host, .optimize = .debug, .link_libc = true, .imports = &.{.{ .name = "living", .module = mods.living }} }); + const shards_run = b.addRunArtifact(b.addTest(.{ .root_module = shards })); + step.dependOn(&living_run.step); + step.dependOn(&shards_run.step); + const living_step = b.step("test-living", "Living Weights' learner and the pocket in each family's shards, on the host"); + living_step.dependOn(&living_run.step); + living_step.dependOn(&shards_run.step); const cli = b.createModule(.{ .root_source_file = b.path("zig/src/cli/cuda_main.zig"), .target = host, .optimize = .debug, .link_libc = true }); cli.addImport("cuda", cuda); cli.addImport("core", mods.core); @@ -246,7 +258,7 @@ pub fn distServer(b: *std.Build, target: std.Build.ResolvedTarget, draft_ids: *s }; const cuda = runtime(b, target, .fast, if (prebuilt != null) &images else &.{}); const mods = family(b, target, .fast, cuda, draft_ids); - return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, build_options, false); + return nativeServer(b, target, .fast, cuda, mods.core, mods.lanes, mods.nemotron, mods.tokenizer, mods.living, build_options, false); } /// SASS targets every release fatbin holds: Ampere (sm_86), RTX 50 and RTX PRO Blackwell (sm_120), GB10 (sm_121). diff --git a/zig/src/core/shard_edit.zig b/zig/src/core/shard_edit.zig index b340a40523..c8b33b11d4 100644 --- a/zig/src/core/shard_edit.zig +++ b/zig/src/core/shard_edit.zig @@ -114,24 +114,44 @@ pub fn grow(gpa: Allocator, io: Io, dir: []const u8, edits: []const Replacement) /// config.json in `dir` with each of `ints` set (added where missing), every other key as it was. pub fn setConfig(gpa: Allocator, io: Io, dir: []const u8, ints: []const struct { []const u8, i64 }) !void { + return setConfigIn(gpa, io, dir, null, ints); +} + +/// config.json in `dir` with each of `ints` set inside the object `scope` names (null: the top level), as setConfig. +pub fn setConfigIn(gpa: Allocator, io: Io, dir: []const u8, scope: ?[]const u8, ints: []const struct { []const u8, i64 }) !void { var arena_state = std.heap.ArenaAllocator.init(gpa); defer arena_state.deinit(); const arena = arena_state.allocator(); const path = try std.fs.path.join(arena, &.{ dir, "config.json" }); var config = try readJson(arena, io, path); if (config != .object) return error.BadConfig; - for (ints) |kv| try config.object.put(arena, kv[0], .{ .integer = kv[1] }); + const o = if (scope) |name| blk: { + const inner = config.object.getPtr(name) orelse return error.BadConfig; + if (inner.* != .object) return error.BadConfig; + break :blk &inner.object; + } else &config.object; + for (ints) |kv| try o.put(arena, kv[0], .{ .integer = kv[1] }); try writeJson(arena, path, config); } /// An integer of `dir`'s config.json, or null when it has none. pub fn configInt(gpa: Allocator, io: Io, dir: []const u8, key: []const u8) !?i64 { + return configIntIn(gpa, io, dir, null, key); +} + +/// An integer of the object `scope` names in `dir`'s config.json (null: the top level), or null when it has none. +pub fn configIntIn(gpa: Allocator, io: Io, dir: []const u8, scope: ?[]const u8, key: []const u8) !?i64 { var arena_state = std.heap.ArenaAllocator.init(gpa); defer arena_state.deinit(); const arena = arena_state.allocator(); const config = try readJson(arena, io, try std.fs.path.join(arena, &.{ dir, "config.json" })); if (config != .object) return error.BadConfig; - const v = config.object.get(key) orelse return null; + const o = if (scope) |name| blk: { + const inner = config.object.get(name) orelse return error.BadConfig; + if (inner != .object) return error.BadConfig; + break :blk inner.object; + } else config.object; + const v = o.get(key) orelse return null; return if (v == .integer) v.integer else error.BadConfig; } @@ -508,6 +528,29 @@ test "a tensor grown in its shard: the shard is rewritten, the index's total mov try std.testing.expectEqualStrings("x", (try field(config, "model_type")).string); } +test "integers set and read inside a nested config object, the top level and every other key as they were" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try tmp.dir.writeFile(io, .{ .sub_path = "config.json", .data = "{\"model_type\":\"x\",\"width\":7,\"text_config\":{\"width\":2,\"name\":\"t\"}}" }); + const dir = try std.fs.path.join(gpa, &.{ ".zig-cache", "tmp", &tmp.sub_path }); + defer gpa.free(dir); + try std.testing.expectEqual(@as(?i64, 2), try configIntIn(gpa, io, dir, "text_config", "width")); + try std.testing.expectEqual(@as(?i64, null), try configIntIn(gpa, io, dir, "text_config", "learned")); + try setConfigIn(gpa, io, dir, "text_config", &.{ .{ "width", 3 }, .{ "learned", 1 } }); + try std.testing.expectEqual(@as(?i64, 3), try configIntIn(gpa, io, dir, "text_config", "width")); + try std.testing.expectEqual(@as(?i64, 1), try configIntIn(gpa, io, dir, "text_config", "learned")); + try std.testing.expectEqual(@as(?i64, 7), try configInt(gpa, io, dir, "width")); + try std.testing.expectEqual(@as(?i64, null), try configInt(gpa, io, dir, "learned")); + try std.testing.expectError(error.BadConfig, configIntIn(gpa, io, dir, "vision_config", "width")); + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + const arena = arena_state.allocator(); + const config = try readJson(arena, io, try std.fs.path.join(arena, &.{ dir, "config.json" })); + try std.testing.expectEqualStrings("t", (try field(try field(config, "text_config"), "name")).string); +} + test "a folder whose shard is a symlink or a hard link is refused for learning; its own files pass" { const gpa = std.testing.allocator; const io = std.testing.io; diff --git a/zig/src/families/nemotron/adapters.zig b/zig/src/families/nemotron/adapters.zig index 3350594214..7d2bac4403 100644 --- a/zig/src/families/nemotron/adapters.zig +++ b/zig/src/families/nemotron/adapters.zig @@ -5,14 +5,14 @@ const cfg = @import("config.zig"); const wts = @import("weights.zig"); const ops = @import("train_ops.zig"); const pl = @import("prefill_launch.zig"); -const subspace = @import("subspace.zig"); +const subspace = @import("living").subspace; const Buffer = mtl.Buffer; const At = pl.At; const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; const shape = @import("learn_shape.zig"); -const route = @import("lesson_route.zig"); +const route = @import("living").lesson_route; pub const block = shape.block; pub const max_rank = shape.max_rank; diff --git a/zig/src/families/nemotron/backend.zig b/zig/src/families/nemotron/backend.zig index f52e5a4a6c..5f2b737a62 100644 --- a/zig/src/families/nemotron/backend.zig +++ b/zig/src/families/nemotron/backend.zig @@ -5,7 +5,7 @@ const lanes = @import("lanes"); const Fence = @import("../../core/fence.zig").Fence; const fwd = @import("forward.zig"); const st = @import("state.zig"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const layers = @import("layers.zig"); const mtp = @import("mtp.zig"); const timing = @import("timing.zig"); diff --git a/zig/src/families/nemotron/cuda.zig b/zig/src/families/nemotron/cuda.zig index cb2fa934a1..58e979c6e5 100644 --- a/zig/src/families/nemotron/cuda.zig +++ b/zig/src/families/nemotron/cuda.zig @@ -40,7 +40,6 @@ test { _ = @import("cuda_prompt_cuts.zig"); _ = @import("cuda_profile_window.zig"); _ = reuse; - _ = @import("learner.zig"); - _ = @import("subspace.zig"); + _ = slide; _ = @import("cuda_train_back.zig"); } diff --git a/zig/src/families/nemotron/cuda_engine.zig b/zig/src/families/nemotron/cuda_engine.zig index f287c4ced4..7d0e6d4518 100644 --- a/zig/src/families/nemotron/cuda_engine.zig +++ b/zig/src/families/nemotron/cuda_engine.zig @@ -15,7 +15,7 @@ const head_fields = @import("cuda_mtp.zig").seq_fields; const sampler = @import("cuda_sampler.zig"); const segs = @import("cuda_segments.zig"); const train_ops = @import("cuda_train_ops.zig"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const heat = @import("heat"); const seq_graphs = @import("cuda_seq_graphs.zig"); const prompt_rows = @import("cuda_prompt_rows.zig"); diff --git a/zig/src/families/nemotron/cuda_forward.zig b/zig/src/families/nemotron/cuda_forward.zig index a0e0452e86..848bec5856 100644 --- a/zig/src/families/nemotron/cuda_forward.zig +++ b/zig/src/families/nemotron/cuda_forward.zig @@ -9,7 +9,7 @@ const weights = @import("cuda_weights.zig"); const state = @import("cuda_state.zig"); const Dump = @import("cuda_dump.zig").Dump; const train_ops = @import("cuda_train_ops.zig"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const kept = @import("cuda_marks.zig"); pub const Delta = enum { none, dense, moe }; diff --git a/zig/src/families/nemotron/cuda_lanes.zig b/zig/src/families/nemotron/cuda_lanes.zig index 203ccd74cb..3f98a86a5c 100644 --- a/zig/src/families/nemotron/cuda_lanes.zig +++ b/zig/src/families/nemotron/cuda_lanes.zig @@ -14,7 +14,7 @@ const prompt_rows = @import("cuda_prompt_rows.zig"); const config = @import("config.zig"); const seq_graphs = @import("cuda_seq_graphs.zig"); const costs = @import("cuda_costs.zig"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const be = lanes.backend; const LogRow = lanes.LogprobRow; diff --git a/zig/src/families/nemotron/cuda_prompt.zig b/zig/src/families/nemotron/cuda_prompt.zig index ba5408db8d..64e07e3f87 100644 --- a/zig/src/families/nemotron/cuda_prompt.zig +++ b/zig/src/families/nemotron/cuda_prompt.zig @@ -14,7 +14,7 @@ const marks = @import("cuda_marks.zig"); const cuts = @import("cuda_prompt_cuts.zig"); const prompt_rows = @import("cuda_prompt_rows.zig"); const reuse = @import("cuda_reuse.zig"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const heat = @import("heat"); const pin_sampled = engine.pin_sampled; const pin_prompt = engine.pin_prompt; diff --git a/zig/src/families/nemotron/cuda_sites.zig b/zig/src/families/nemotron/cuda_sites.zig index c44e37c527..930e5ef0f5 100644 --- a/zig/src/families/nemotron/cuda_sites.zig +++ b/zig/src/families/nemotron/cuda_sites.zig @@ -4,9 +4,9 @@ const cuda = @import("cuda"); const Config = @import("config.zig").Config; const weights = @import("cuda_weights.zig"); const train_ops = @import("cuda_train_ops.zig"); -const subspace = @import("subspace.zig"); +const subspace = @import("living").subspace; const shape = @import("learn_shape.zig"); -const route = @import("lesson_route.zig"); +const route = @import("living").lesson_route; pub const block = shape.block; pub const max_rank = shape.max_rank; diff --git a/zig/src/families/nemotron/cuda_slide.zig b/zig/src/families/nemotron/cuda_slide.zig index e65dd647b2..9c67cf53a0 100644 --- a/zig/src/families/nemotron/cuda_slide.zig +++ b/zig/src/families/nemotron/cuda_slide.zig @@ -1,7 +1,7 @@ //! Nemotron's Sliding Weights learner on CUDA: learner.zig's lessons over the CUDA trainer, the engine and its shards. const std = @import("std"); const cuda = @import("cuda"); -const lessons = @import("learner.zig"); +const lessons = @import("living").learner; const train = @import("cuda_train.zig"); const engine = @import("cuda_engine.zig"); const sites_mod = @import("cuda_sites.zig"); @@ -10,6 +10,7 @@ const state = @import("cuda_state.zig"); const Living = living.Io(@import("core")); const PocketShards = @import("pocket_shards.zig").Io(@import("core")); +const pocket_layout = @import("pocket_layout.zig"); pub const Example = lessons.Example; pub const Lesson = lessons.Lesson; @@ -27,11 +28,17 @@ const Host = struct { pub const Trainer = train.Trainer; pub const max_rows = train.max_rows; pub const Ctx = Model; + pub const family = pocket_layout.family; pub fn layers(m: Ctx) usize { return m.e.c.layers; } + /// Where the pocket sits and learns: every MoE layer, all but the last learning. + pub fn layout(m: Ctx) @import("living").Layout { + return pocket_layout.layout(m.e.c); + } + pub fn trainer(gpa: std.mem.Allocator, io: std.Io, m: Ctx) !*Trainer { return train.Trainer.init(gpa, io, m.e, m.dir, m.living); } diff --git a/zig/src/families/nemotron/cuda_state.zig b/zig/src/families/nemotron/cuda_state.zig index 92bac057b9..e6a7c7caab 100644 --- a/zig/src/families/nemotron/cuda_state.zig +++ b/zig/src/families/nemotron/cuda_state.zig @@ -1,7 +1,7 @@ //! Device buffers of one Nemotron sequence: its caches and recurrent state, and the scratch a window or chunk uses. const std = @import("std"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const cuda = @import("cuda"); const Config = @import("config.zig").Config; const kern = @import("cuda_kernels.zig"); diff --git a/zig/src/families/nemotron/cuda_train.zig b/zig/src/families/nemotron/cuda_train.zig index 771e472886..ab96e317d4 100644 --- a/zig/src/families/nemotron/cuda_train.zig +++ b/zig/src/families/nemotron/cuda_train.zig @@ -1,6 +1,6 @@ //! Learning in the weights on CUDA (train.zig on Metal): sequences side by side forward, their losses, and back. const std = @import("std"); -const train_mode = @import("train_mode.zig"); +const train_mode = @import("living").train_mode; const cuda = @import("cuda"); const core = @import("core"); const Config = @import("config.zig").Config; @@ -10,12 +10,12 @@ const fwd = @import("cuda_forward.zig"); const engine = @import("cuda_engine.zig"); const ops = @import("cuda_train_ops.zig"); const sites_mod = @import("cuda_sites.zig"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const back = @import("cuda_train_back.zig"); const tfw = @import("cuda_train_fwd.zig"); const lay = @import("cuda_train_layout.zig"); const prefix = @import("cuda_train_prefix.zig"); -const Count = @import("pocket_clock.zig").Count; +const Count = @import("living").pocket_clock.Count; const Source = @import("cuda_source.zig").Source; /// Rows a sequence learns from at most: a question and its answer (attention's backward reads at most 256). diff --git a/zig/src/families/nemotron/cuda_train_fwd.zig b/zig/src/families/nemotron/cuda_train_fwd.zig index 3b39c30976..b099ae94e6 100644 --- a/zig/src/families/nemotron/cuda_train_fwd.zig +++ b/zig/src/families/nemotron/cuda_train_fwd.zig @@ -6,7 +6,7 @@ const state = @import("cuda_state.zig"); const weights = @import("cuda_weights.zig"); const fwd = @import("cuda_forward.zig"); const train_ops = @import("cuda_train_ops.zig"); -const Span = @import("train_mode.zig").Span; +const Span = @import("living").train_mode.Span; const Store = @import("cuda_train_prefix.zig").Store; /// The open fact's pocket neurons scaled row by row: `rows` f32 [rows] on the device, at layer l from `from[l]` on. diff --git a/zig/src/families/nemotron/cuda_train_ops.zig b/zig/src/families/nemotron/cuda_train_ops.zig index 209ad9de3b..b5b10bc39f 100644 --- a/zig/src/families/nemotron/cuda_train_ops.zig +++ b/zig/src/families/nemotron/cuda_train_ops.zig @@ -1,7 +1,7 @@ //! The training kernels (train.cu) as calls on one stream: the loss, the change, products and each layer's backward. const std = @import("std"); const cuda = @import("cuda"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; /// A Mamba layer's shape as the mixer kernels read it (train.cu's Shape). pub const Shape = extern struct { rows: c_int, heads: c_int, dh: c_int, groups: c_int, n: c_int, inner: c_int, conv: c_int, proj: c_int }; diff --git a/zig/src/families/nemotron/forward.zig b/zig/src/families/nemotron/forward.zig index 7dc93e2de2..9fa69def84 100644 --- a/zig/src/families/nemotron/forward.zig +++ b/zig/src/families/nemotron/forward.zig @@ -13,7 +13,7 @@ const fullVocabulary = @import("lanes").gpu_full.fullVocabulary; const gpu_rule = @import("lanes").gpu_rule; const row = @import("../../core/row_projection.zig"); const Masks = @import("../../core/grammar_mask.zig").Masks; -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const Buffer = mtl.Buffer; pub const Enc = encoder.Enc; diff --git a/zig/src/families/nemotron/layers.zig b/zig/src/families/nemotron/layers.zig index e2e18fb105..5dc16880d2 100644 --- a/zig/src/families/nemotron/layers.zig +++ b/zig/src/families/nemotron/layers.zig @@ -6,7 +6,7 @@ const st = @import("state.zig"); const fwd = @import("forward.zig"); const kern = @import("kernels.zig"); const tree = @import("tree.zig"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const max_blocks = @import("learn_shape.zig").max_blocks; const Buffer = mtl.Buffer; diff --git a/zig/src/families/nemotron/learn_shape.zig b/zig/src/families/nemotron/learn_shape.zig index 6e4e3bf779..773c71f875 100644 --- a/zig/src/families/nemotron/learn_shape.zig +++ b/zig/src/families/nemotron/learn_shape.zig @@ -1,30 +1,19 @@ -//! The learned change's shape, shared by the Metal and CUDA learners: blocks of ranks, their scale and gates. +//! The learned change's shape on Nemotron: the shared constants (living/learn_shape.zig) and where its blocks sit. const std = @import("std"); const Config = @import("config.zig").Config; - -pub const block = 16; // ranks a lesson adds at each layer -pub const max_rank = 2048; // a hundred and twenty-eight lessons' blocks, each layer's buffers grown to it as lessons come -pub const max_blocks = max_rank / block; -pub const scale: f32 = 10; -pub const avoid_dims = 192; // directions sketched from the inputs a new block must leave alone -pub const candidates = 64; // directions sketched from the fact's inputs, which a new block's are chosen among - -/// Adam's settings for a block's outputs (its gate keeps them off what must stay): rate, moments' decays, epsilon. -pub const hyper = [4]f32{ 3e-4, 0.9, 0.999, 1e-8 }; - -/// Adam's settings at step `step`: the rate halves every `settle` steps, three times at most, so later rounds refine. -pub fn hyperAt(step: u64) [4]f32 { - var h = hyper; - h[0] /= @as(f32, @floatFromInt(@as(u32, 1) << @intCast(@min(step / settle, 3)))); - return h; -} -pub const settle = 60; - -/// Each new input direction's length, as MLX's LoRA draws its down factor's columns (uniform within 1 / sqrt(in)). -pub const a_norm: f32 = 0.577; - -/// A block's directions' squared length: a row's cosine with one is (x . a) / sqrt(unit |x|^2). -pub const unit = a_norm * a_norm; +const shape = @import("living").learn_shape; + +pub const block = shape.block; +pub const max_rank = shape.max_rank; +pub const max_blocks = shape.max_blocks; +pub const scale = shape.scale; +pub const avoid_dims = shape.avoid_dims; +pub const candidates = shape.candidates; +pub const hyper = shape.hyper; +pub const hyperAt = shape.hyperAt; +pub const settle = shape.settle; +pub const a_norm = shape.a_norm; +pub const unit = shape.unit; /// Layer i's output projection as the checkpoint names its module: the one a lesson's change sits after. pub fn moduleName(buf: []u8, c: Config, i: usize) ![]const u8 { diff --git a/zig/src/families/nemotron/nemotron.zig b/zig/src/families/nemotron/nemotron.zig index 5e26d6d4d1..0d521769f5 100644 --- a/zig/src/families/nemotron/nemotron.zig +++ b/zig/src/families/nemotron/nemotron.zig @@ -25,8 +25,8 @@ pub const learned = @import("learned.zig"); pub const train_ops = @import("train_ops.zig"); pub const train_back = @import("train_back.zig"); pub const train = @import("train.zig"); -pub const pocket_fit = @import("pocket_fit.zig"); -pub const lesson_route = @import("lesson_route.zig"); +pub const pocket_fit = @import("living").pocket_fit; +pub const lesson_route = @import("living").lesson_route; pub const snapshot = @import("snapshot.zig"); pub const learned_prompts = @import("learned_prompts.zig"); diff --git a/zig/src/families/nemotron/pocket_layout.zig b/zig/src/families/nemotron/pocket_layout.zig new file mode 100644 index 0000000000..879ddb28d1 --- /dev/null +++ b/zig/src/families/nemotron/pocket_layout.zig @@ -0,0 +1,12 @@ +//! Nemotron as Living Weights reads it: its family (living/families.zig) and the layout its config gives. +const living = @import("living"); +const Config = @import("config.zig").Config; + +pub const family = living.families.nemotron_h; + +/// Every MoE layer's shared expert carries the pocket; all but the last layer's learn. +pub fn layout(c: Config) living.Layout { + var moe: living.Layers = .empty; + for (0..c.layers) |l| if (c.kinds[l] == .moe) moe.set(l); + return .fromExperts(c.hidden, c.shared_width, c.layers, moe); +} diff --git a/zig/src/families/nemotron/pocket_shards.zig b/zig/src/families/nemotron/pocket_shards.zig index 8eaca2c3b5..5cdab00612 100644 --- a/zig/src/families/nemotron/pocket_shards.zig +++ b/zig/src/families/nemotron/pocket_shards.zig @@ -1,178 +1,16 @@ -//! The pocket in the shards: every shared expert widened once (config.json told), later writes in place. +//! The pocket in Nemotron's shards: living/shards.zig over its family and config (up_proj rows, down_proj columns). const std = @import("std"); const cfg = @import("config.zig"); +const living = @import("living"); +const pocket_layout = @import("pocket_layout.zig"); -const group = 64; - -/// The writer over `core`, the backend's own (its shard_edit, safetensors and affine4_host). +/// The writer over `core`, the backend's own (its shard_edit and safetensors). pub fn Io(comptime core: type) type { - const st = core.safetensors; - const se = core.shard_edit; - const host4 = core.affine4_host; + const Shards = living.shards.Io(core); return struct { /// The learned neurons of MoE layers `layers` into `dir`'s shards (widened first if new), as 4-bit; the bytes. pub fn write(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, c: cfg.Config, layers: []const usize, neurons: usize, up: []f32, down: []f32) !usize { - const d = c.hidden; - const own = c.shared_width; - if (neurons % group != 0 or d % group != 0 or own % group != 0) return error.BadPocket; - const has = try se.configInt(gpa, io, dir, cfg.pocket_key); - const widen = has == null; - if (has) |n| if (n != neurons) return error.PocketOfAnotherSize; - var arena_state = std.heap.ArenaAllocator.init(gpa); - defer arena_state.deinit(); - const arena = arena_state.allocator(); - var files: Files = .{ .gpa = gpa, .io = io, .dir = dir, .paths = .init(gpa) }; - defer files.close(); - var edits: std.ArrayList(se.Replacement) = .empty; - var bytes: usize = 0; - for (0..c.layers) |l| { - if (c.kinds[l] != .moe) continue; - const at = std.mem.indexOfScalar(usize, layers, l); - const n = neurons * d; - const u = if (at) |j| up[j * n ..][0..n] else try zeros(arena, n); - const v = if (at) |j| down[j * n ..][0..n] else try zeros(arena, n); - var buf: [96]u8 = undefined; - const base = try std.fmt.bufPrint(&buf, "backbone.layers.{d}.mixer.shared_experts.", .{l}); - bytes += try expert(arena, &files, &edits, base, own, neurons, d, u, v); - } - if (widen) try se.grow(gpa, io, dir, edits.items) else try se.overwrite(gpa, io, dir, edits.items); - if (widen) { - const mtp = try std.fs.path.join(arena, &.{ dir, "mtp-4bit.safetensors" }); - if (hasMtp(gpa, io, mtp)) { - var only: std.ArrayList(se.Replacement) = .empty; - var none: Files = .{ .gpa = gpa, .io = io, .dir = dir, .single = mtp, .paths = .init(gpa) }; - defer none.close(); - const z = try zeros(arena, neurons * d); - bytes += try expert(arena, &none, &only, "layers.1.mixer.shared_experts.", own, neurons, d, z, z); - _ = try se.rewrite(gpa, io, mtp, only.items); - } - try se.setConfig(gpa, io, dir, &.{ .{ "moe_shared_expert_intermediate_size", @intCast(own + neurons) }, .{ cfg.pocket_key, @intCast(neurons) } }); - } - return bytes; - } - - fn hasMtp(gpa: std.mem.Allocator, io: std.Io, path: []const u8) bool { - var f = st.File.open(gpa, io, path) catch return false; - defer f.close(io); - return f.get("layers.1.mixer.shared_experts.up_proj.weight") != null; - } - - /// The shard files a model's tensors are read from: the index's, or one file (the MTP head's). - const Files = struct { - gpa: std.mem.Allocator, - io: std.Io, - dir: []const u8, - single: ?[]const u8 = null, - open: std.StringArrayHashMapUnmanaged(st.File) = .empty, - paths: std.heap.ArenaAllocator, - map: ?std.json.Parsed(std.json.Value) = null, - - fn get(fs: *Files, name: []const u8) !st.Tensor { - const shard = if (fs.single) |x| x else try fs.shardOf(name); - const slot = try fs.open.getOrPut(fs.gpa, shard); - if (!slot.found_existing) slot.value_ptr.* = st.File.open(fs.gpa, fs.io, shard) catch |e| { - _ = fs.open.pop(); - return e; - }; - return slot.value_ptr.get(name) orelse error.MissingTensor; - } - - /// The path of the shard holding `name`, kept until the files close. - fn shardOf(fs: *Files, name: []const u8) ![]const u8 { - const a = fs.paths.allocator(); - if (fs.map == null) { - const path = try std.fs.path.join(a, &.{ fs.dir, "model.safetensors.index.json" }); - const text = std.Io.Dir.cwd().readFileAlloc(fs.io, path, a, .limited(1 << 26)) catch |e| if (e == error.FileNotFound) { - fs.single = try std.fs.path.join(a, &.{ fs.dir, "model.safetensors" }); - return fs.single.?; - } else return e; - fs.map = try std.json.parseFromSlice(std.json.Value, fs.gpa, text, .{ .allocate = .alloc_always }); - } - const root = fs.map.?.value; - if (root != .object) return error.BadIndex; - const weights = root.object.get("weight_map") orelse return error.BadIndex; - if (weights != .object) return error.BadIndex; - const file = weights.object.get(name) orelse return error.MissingTensor; - if (file != .string) return error.BadIndex; - return std.fs.path.join(a, &.{ fs.dir, file.string }); - } - - fn close(fs: *Files) void { - for (fs.open.values()) |*f| f.close(fs.io); - fs.open.deinit(fs.gpa); - if (fs.map) |*x| x.deinit(); - fs.paths.deinit(); - } - }; - - /// One shared expert's up_proj and down_proj with `neurons` learned after its `own`, into `edits`; the bytes. - fn expert(arena: std.mem.Allocator, files: *Files, edits: *std.ArrayList(se.Replacement), base: []const u8, own: usize, neurons: usize, d: usize, u: []f32, v: []f32) !usize { - const wide = own + neurons; - var name: [128]u8 = undefined; - var parts: [6]st.Tensor = undefined; - const leaves = [_][]const u8{ "up_proj.weight", "up_proj.scales", "up_proj.biases", "down_proj.weight", "down_proj.scales", "down_proj.biases" }; - for (&parts, leaves) |*t, leaf| t.* = try files.get(try std.fmt.bufPrint(&name, "{s}{s}", .{ base, leaf })); - const rows_had = parts[0].dim(0); // up_proj's rows: `own`, or `wide` once widened - if (rows_had != own and rows_had != wide) return error.UnexpectedTensor; - // up_proj: the own rows as they are, then a row a learned neuron - const uw = try arena.alloc(u32, wide * d / 8); - const us = try arena.alloc(u16, wide * d / group); - const ub = try arena.alloc(u16, wide * d / group); - @memcpy(uw[0 .. own * d / 8], words(parts[0].bytes)[0 .. own * d / 8]); - @memcpy(us[0 .. own * d / group], halves(parts[1].bytes)[0 .. own * d / group]); - @memcpy(ub[0 .. own * d / group], halves(parts[2].bytes)[0 .. own * d / group]); - for (0..neurons) |k| { - const r = own + k; - const x = u[k * d ..][0..d]; - host4.quantize(x, uw[r * d / 8 ..][0 .. d / 8], us[r * d / group ..][0 .. d / group], ub[r * d / group ..][0 .. d / group]); - host4.dequantize(uw[r * d / 8 ..][0 .. d / 8], us[r * d / group ..][0 .. d / group], ub[r * d / group ..][0 .. d / group], x); - } - // down_proj: each output row's own columns as they are, then a column a learned neuron - const had = rows_had; // down_proj's columns, as wide as up_proj had rows - const dw = try arena.alloc(u32, d * wide / 8); - const ds = try arena.alloc(u16, d * wide / group); - const db = try arena.alloc(u16, d * wide / group); - const col = try arena.alloc(f32, neurons); - for (0..d) |o| { - @memcpy(dw[o * wide / 8 ..][0 .. own / 8], words(parts[3].bytes)[o * had / 8 ..][0 .. own / 8]); - @memcpy(ds[o * wide / group ..][0 .. own / group], halves(parts[4].bytes)[o * had / group ..][0 .. own / group]); - @memcpy(db[o * wide / group ..][0 .. own / group], halves(parts[5].bytes)[o * had / group ..][0 .. own / group]); - for (col, 0..) |*x, k| x.* = v[k * d + o]; - const w8 = dw[o * wide / 8 ..][own / 8 ..][0 .. neurons / 8]; - const s8 = ds[o * wide / group ..][own / group ..][0 .. neurons / group]; - const b8 = db[o * wide / group ..][own / group ..][0 .. neurons / group]; - host4.quantize(col, w8, s8, b8); - host4.dequantize(w8, s8, b8, col); - for (col, 0..) |x, k| v[k * d + o] = x; - } - const made = [_]struct { leaf: []const u8, dtype: st.DType, shape: [2]usize, bytes: []const u8 }{ - .{ .leaf = leaves[0], .dtype = .u32, .shape = .{ wide, d / 8 }, .bytes = std.mem.sliceAsBytes(uw) }, - .{ .leaf = leaves[1], .dtype = .bf16, .shape = .{ wide, d / group }, .bytes = std.mem.sliceAsBytes(us) }, - .{ .leaf = leaves[2], .dtype = .bf16, .shape = .{ wide, d / group }, .bytes = std.mem.sliceAsBytes(ub) }, - .{ .leaf = leaves[3], .dtype = .u32, .shape = .{ d, wide / 8 }, .bytes = std.mem.sliceAsBytes(dw) }, - .{ .leaf = leaves[4], .dtype = .bf16, .shape = .{ d, wide / group }, .bytes = std.mem.sliceAsBytes(ds) }, - .{ .leaf = leaves[5], .dtype = .bf16, .shape = .{ d, wide / group }, .bytes = std.mem.sliceAsBytes(db) }, - }; - var bytes: usize = 0; - for (made) |m| { - try edits.append(arena, .{ .name = try std.fmt.allocPrint(arena, "{s}{s}", .{ base, m.leaf }), .dtype = m.dtype, .shape = try arena.dupe(usize, &m.shape), .bytes = m.bytes }); - bytes += m.bytes.len; - } - return bytes; + return Shards.write(gpa, io, dir, pocket_layout.family, pocket_layout.layout(c), layers, neurons, &.{up}, down); } }; } - -fn zeros(arena: std.mem.Allocator, n: usize) ![]f32 { - const z = try arena.alloc(f32, n); - @memset(z, 0); - return z; -} - -fn words(bytes: []const u8) []align(1) const u32 { - return std.mem.bytesAsSlice(u32, bytes); -} - -fn halves(bytes: []const u8) []align(1) const u16 { - return std.mem.bytesAsSlice(u16, bytes); -} diff --git a/zig/src/families/nemotron/prefill.zig b/zig/src/families/nemotron/prefill.zig index 46c1070408..0bd134b925 100644 --- a/zig/src/families/nemotron/prefill.zig +++ b/zig/src/families/nemotron/prefill.zig @@ -9,7 +9,7 @@ const pk = @import("prefill_kernels.zig"); const pl = @import("prefill_launch.zig"); const mixers = @import("prefill_mixers.zig"); const layers = @import("layers.zig"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const Model = @import("model.zig").Model; const Buffer = mtl.Buffer; diff --git a/zig/src/families/nemotron/slide.zig b/zig/src/families/nemotron/slide.zig index ac6bca25f4..ff62940596 100644 --- a/zig/src/families/nemotron/slide.zig +++ b/zig/src/families/nemotron/slide.zig @@ -1,6 +1,6 @@ //! Nemotron's Sliding Weights learner on Metal: learner.zig's lessons over the Metal trainer, its weights and shards. const std = @import("std"); -const lessons = @import("learner.zig"); +const lessons = @import("living").learner; const train = @import("train.zig"); const backend = @import("backend.zig"); const adapters = @import("adapters.zig"); @@ -9,6 +9,7 @@ const learn_shape = @import("learn_shape.zig"); const weights = @import("weights.zig"); const mtl = @import("metal"); const ckpt = @import("../../core/checkpoint_metal.zig"); +const pocket_layout = @import("pocket_layout.zig"); const shard_edit = @import("../../core/shard_edit.zig"); const Living = living.Io(struct { @@ -31,11 +32,17 @@ const Host = struct { pub const Trainer = train.Trainer; pub const max_rows = train.max_rows; pub const Ctx = *backend.Metal; + pub const family = pocket_layout.family; pub fn layers(b: Ctx) usize { return b.m.config.layers; } + /// Where the pocket sits and learns: every MoE layer, all but the last learning. + pub fn layout(b: Ctx) @import("living").Layout { + return pocket_layout.layout(b.m.config); + } + pub fn trainer(gpa: std.mem.Allocator, _: std.Io, b: Ctx) !*Trainer { return train.Trainer.init(gpa, b); } diff --git a/zig/src/families/nemotron/state.zig b/zig/src/families/nemotron/state.zig index e45b2ca86c..eec5bab7a1 100644 --- a/zig/src/families/nemotron/state.zig +++ b/zig/src/families/nemotron/state.zig @@ -3,7 +3,7 @@ const std = @import("std"); const mtl = @import("metal"); const cfg = @import("config.zig"); const fwd = @import("forward.zig"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const opts = mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked; diff --git a/zig/src/families/nemotron/train.zig b/zig/src/families/nemotron/train.zig index 7de08a1eb6..641dc36681 100644 --- a/zig/src/families/nemotron/train.zig +++ b/zig/src/families/nemotron/train.zig @@ -1,6 +1,6 @@ //! Learning in the weights: sequences side by side forward through every layer, their losses, back to each change. const std = @import("std"); -const train_mode = @import("train_mode.zig"); +const train_mode = @import("living").train_mode; const mtl = @import("metal"); const cfg = @import("config.zig"); const st = @import("state.zig"); @@ -11,9 +11,9 @@ const back = @import("train_back.zig"); const job = @import("train_job.zig"); const adapters = @import("adapters.zig"); const backend = @import("backend.zig"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const prefix = @import("train_prefix.zig"); -const Count = @import("pocket_clock.zig").Count; +const Count = @import("living").pocket_clock.Count; const Metal = backend.Metal; const At = pl.At; diff --git a/zig/src/families/nemotron/train_back.zig b/zig/src/families/nemotron/train_back.zig index 8a37c70ddd..1dfa6c390f 100644 --- a/zig/src/families/nemotron/train_back.zig +++ b/zig/src/families/nemotron/train_back.zig @@ -7,7 +7,7 @@ const pre = @import("prefill.zig"); const pl = @import("prefill_launch.zig"); const ops = @import("train_ops.zig"); const adapters = @import("adapters.zig"); -const Span = @import("train_mode.zig").Span; +const Span = @import("living").train_mode.Span; const Store = @import("train_prefix.zig").Store; const Buffer = mtl.Buffer; diff --git a/zig/src/families/nemotron/train_forward.zig b/zig/src/families/nemotron/train_forward.zig index d3d2e1dbdb..a137ee06e1 100644 --- a/zig/src/families/nemotron/train_forward.zig +++ b/zig/src/families/nemotron/train_forward.zig @@ -5,7 +5,7 @@ const wts = @import("weights.zig"); const pre = @import("prefill.zig"); const pl = @import("prefill_launch.zig"); const layers = @import("layers.zig"); -const Span = @import("train_mode.zig").Span; +const Span = @import("living").train_mode.Span; const Store = @import("train_prefix.zig").Store; const At = pl.At; diff --git a/zig/src/families/nemotron/train_job.zig b/zig/src/families/nemotron/train_job.zig index c209c52d1f..6586ddcad8 100644 --- a/zig/src/families/nemotron/train_job.zig +++ b/zig/src/families/nemotron/train_job.zig @@ -1,6 +1,6 @@ //! A training pass as one command buffer: its sequences forward through every layer, their losses, back to each change. const std = @import("std"); -const train_mode = @import("train_mode.zig"); +const train_mode = @import("living").train_mode; const mtl = @import("metal"); const cfg = @import("config.zig"); const fwd = @import("forward.zig"); @@ -11,7 +11,7 @@ const back = @import("train_back.zig"); const tf = @import("train_forward.zig"); const adapters = @import("adapters.zig"); const backend = @import("backend.zig"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const train = @import("train.zig"); const Trainer = train.Trainer; const pack_rows = train.pack_rows; diff --git a/zig/src/families/nemotron/train_ops.zig b/zig/src/families/nemotron/train_ops.zig index f690c07059..2a80871631 100644 --- a/zig/src/families/nemotron/train_ops.zig +++ b/zig/src/families/nemotron/train_ops.zig @@ -7,7 +7,7 @@ const wts = @import("weights.zig"); const pl = @import("prefill_launch.zig"); const pk = @import("prefill_kernels.zig"); const layers = @import("layers.zig"); -const lesson_route = @import("lesson_route.zig"); +const lesson_route = @import("living").lesson_route; const Buffer = mtl.Buffer; const At = pl.At; diff --git a/zig/src/living/affine.zig b/zig/src/living/affine.zig new file mode 100644 index 0000000000..177b40a4d4 --- /dev/null +++ b/zig/src/living/affine.zig @@ -0,0 +1,86 @@ +//! MLX's affine layout on the host at any group and code width; at 4 bits in groups of 64 it is affine4_host's. +const std = @import("std"); +const Quant = @import("family.zig").Quant; + +pub fn f32of(b: u16) f32 { + return @bitCast(@as(u32, b) << 16); +} + +/// The nearest bf16, ties to even; NaN stays NaN. +pub fn bf16of(x: f32) u16 { + const u: u32 = @bitCast(x); + if (std.math.isNan(x)) return @intCast((u >> 16) | 0x40); + return @intCast((u + 0x7fff + ((u >> 16) & 1)) >> 16); +} + +/// u32 words `n` values take at `q`'s width. +pub fn wordsOf(comptime q: Quant, n: usize) usize { + return n * q.bits / 32; +} + +/// Row-major values from codes, scales and biases: scale * code + bias in f32 (the rows a whole number of groups). +pub fn dequantize(comptime q: Quant, words: []const u32, scales: []const u16, biases: []const u16, out: []f32) void { + const per = 32 / q.bits; + const mask: u32 = (1 << q.bits) - 1; + for (out, 0..) |*o, i| { + const c: u32 = (words[i / per] >> @intCast(q.bits * (i % per))) & mask; + o.* = f32of(scales[i / q.group]) * @as(f32, @floatFromInt(c)) + f32of(biases[i / q.group]); + } +} + +/// Codes, scales and biases for row-major values: a group's minimum is its bias, its range over the top code its scale. +pub fn quantize(comptime q: Quant, values: []const f32, words: []u32, scales: []u16, biases: []u16) void { + comptime std.debug.assert(32 % q.bits == 0); + const per = 32 / q.bits; + const top: f32 = @floatFromInt((1 << q.bits) - 1); + @memset(words, 0); + for (0..values.len / q.group) |g| { + const v = values[g * q.group ..][0..q.group]; + const lo = std.mem.min(f32, v); + const hi = std.mem.max(f32, v); + scales[g] = bf16of((hi - lo) / top); + biases[g] = bf16of(lo); + const s = f32of(scales[g]); + const b = f32of(biases[g]); + for (v, g * q.group..) |x, i| { + const c: u32 = if (s == 0) 0 else @intFromFloat(std.math.clamp(@round((x - b) / s), 0, top)); + words[i / per] |= c << @intCast(q.bits * (i % per)); + } + } +} + +test "a row at 4 bits in groups of 32 lands within half a step of each value, a flat group exactly on its bias" { + const q: Quant = .{ .bits = 4, .group = 32 }; + var values: [3 * 32]f32 = undefined; + for (&values, 0..) |*v, i| v.* = @sin(@as(f32, @floatFromInt(i)) * 0.41) * 0.07; + @memset(values[64..], -0.03); + var words: [wordsOf(q, 3 * 32)]u32 = undefined; + var scales: [3]u16 = undefined; + var biases: [3]u16 = undefined; + quantize(q, &values, &words, &scales, &biases); + var back: [3 * 32]f32 = undefined; + dequantize(q, &words, &scales, &biases, &back); + for (values, back, 0..) |v, b, i| try std.testing.expect(@abs(v - b) <= f32of(scales[i / 32]) * 0.51 + 1e-3); + for (back[64..]) |b| try std.testing.expectEqual(f32of(bf16of(-0.03)), b); + // written again from what it reads back, a row is the same codes: a save's values are a fixed point + var again: [wordsOf(q, 3 * 32)]u32 = undefined; + var s2: [3]u16 = undefined; + var b2: [3]u16 = undefined; + quantize(q, &back, &again, &s2, &b2); + var back2: [3 * 32]f32 = undefined; + dequantize(q, &again, &s2, &b2, &back2); + for (back, back2) |x, y| try std.testing.expectApproxEqAbs(x, y, 1e-6); +} + +test "8-bit codes pack four a word" { + const q: Quant = .{ .bits = 8, .group = 32 }; + var values: [32]f32 = undefined; + for (&values, 0..) |*v, i| v.* = @as(f32, @floatFromInt(i)) / 31; + var words: [8]u32 = undefined; + var s: [1]u16 = undefined; + var b: [1]u16 = undefined; + quantize(q, &values, &words, &s, &b); + var back: [32]f32 = undefined; + dequantize(q, &words, &s, &b, &back); + for (values, back) |v, x| try std.testing.expectApproxEqAbs(v, x, f32of(s[0]) * 0.51 + 1e-3); +} diff --git a/zig/src/families/nemotron/choice.zig b/zig/src/living/choice.zig similarity index 100% rename from zig/src/families/nemotron/choice.zig rename to zig/src/living/choice.zig diff --git a/zig/src/families/nemotron/editor.zig b/zig/src/living/editor.zig similarity index 100% rename from zig/src/families/nemotron/editor.zig rename to zig/src/living/editor.zig diff --git a/zig/src/living/families.zig b/zig/src/living/families.zig new file mode 100644 index 0000000000..a1db5ab468 --- /dev/null +++ b/zig/src/living/families.zig @@ -0,0 +1,79 @@ +//! Every family Living Weights knows, as data: the learner, the shard writer and the server's view read these. +const std = @import("std"); +const family = @import("family.zig"); +const Family = family.Family; +const test_families = @import("test_families.zig"); + +/// Nemotron 3.5 Lightning (nemotron_h): relu^2 shared experts at the pattern's MoE layers, 4-bit groups of 64. +pub const nemotron_h: Family = .{ + .model_type = "nemotron_h", + .names = .{ + .layers = "backbone.layers.", + .expert = "mixer.shared_experts.", + .inputs = &.{"up_proj"}, + .mtp = .{ .file = "mtp-4bit.safetensors", .base = "layers.1.mixer.shared_experts." }, + }, + .neuron = .{ .act = .relu2, .margin = 0.25 }, + .quant = .{ .bits = 4, .group = 64 }, + .keys = .{ .width = "moe_shared_expert_intermediate_size" }, + .experts = .pattern, +}; + +pub const all = [_]Family{nemotron_h}; + +/// The family of a checkpoint's top-level model_type, or null. +pub fn byModelType(model_type: []const u8) ?Family { + inline for (all) |f| if (std.mem.eql(u8, f.model_type, model_type)) return f; + return null; +} + +/// A config's layer letters (M mamba, E MoE, * attention, - MLP, ? other); a family with an expert at each: all E. +pub fn kinds(a: std.mem.Allocator, f: Family, o: std.json.ObjectMap) ![]const u8 { + const n = intField(o, f.keys.layers); + if (f.experts == .every) { + const count = n orelse if (o.get("layer_types")) |v| (if (v == .array) v.array.items.len else 0) else 0; + const out = try a.alloc(u8, count); + @memset(out, 'E'); + return out; + } + if (o.get("hybrid_override_pattern")) |v| if (v == .string and v.string.len > 0) return a.dupe(u8, v.string); + if (o.get("layers_block_type")) |v| if (v == .array) { + const out = try a.alloc(u8, v.array.items.len); + for (v.array.items, out) |b, *k| k.* = letter(if (b == .string) b.string else ""); + return out; + }; + const out = try a.alloc(u8, n orelse 0); + @memset(out, '?'); + return out; +} + +/// The object a family's model keys sit in (its scope, else the top level). +pub fn scoped(f: Family, root: std.json.ObjectMap) ?std.json.ObjectMap { + const s = f.keys.scope orelse return root; + const v = root.get(s) orelse return null; + return if (v == .object) v.object else null; +} + +fn letter(name: []const u8) u8 { + const names = [_]struct { []const u8, u8 }{ .{ "mamba", 'M' }, .{ "moe", 'E' }, .{ "attention", '*' }, .{ "mlp", '-' } }; + for (names) |n| if (std.mem.eql(u8, name, n[0])) return n[1]; + return '?'; +} + +pub fn intField(o: std.json.ObjectMap, key: []const u8) ?usize { + const v = o.get(key) orelse return null; + return if (v == .integer and v.integer >= 0) @intCast(v.integer) else null; +} + +test "each family's model type finds it; a pattern family's letters come from its pattern, an every-layer one's are E" { + try std.testing.expectEqualStrings("nemotron_h", byModelType("nemotron_h").?.model_type); + try std.testing.expect(byModelType("qwen4_exp") == null); + try std.testing.expect(byModelType("glm") == null); + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const n = try std.json.parseFromSliceLeaky(std.json.Value, a, "{\"layers_block_type\": [\"mamba\", \"moe\", \"attention\"]}", .{}); + try std.testing.expectEqualStrings("ME*", try kinds(a, nemotron_h, n.object)); + const q = try std.json.parseFromSliceLeaky(std.json.Value, a, "{\"text_config\": {\"num_hidden_layers\": 3}}", .{}); + try std.testing.expectEqualStrings("EEE", try kinds(a, test_families.swiglu, scoped(test_families.swiglu, q.object).?)); +} diff --git a/zig/src/living/family.zig b/zig/src/living/family.zig new file mode 100644 index 0000000000..e4498b1d3c --- /dev/null +++ b/zig/src/living/family.zig @@ -0,0 +1,135 @@ +//! What a family hands Living Weights: where a pocket sits, how its neurons fire, how its shards hold them. +const std = @import("std"); + +/// What a pocket neuron does with its input x: relu(u x)^2 (Nemotron), or silu(g x) (w x) (a SwiGLU shared expert). +pub const Activation = enum { relu2, swiglu }; + +/// A pocket neuron's kind: its activation, whether a per-row gate scales its expert's output, and its input fit. +pub const Neuron = struct { + act: Activation, + /// The shared expert's output is scaled by a gate the trainer reads at each row, so a neuron adds gate act(x) v. + gated: bool = false, + /// The input fit's margin: rows a neuron stays silent on are fitted to u x at or below -margin. + margin: f32 = 0.25, + /// SwiGLU: a neuron fitted as u is written as gate row sharp * u and up row u / sharp (a power of two keeps u). + sharp: f32 = 16, + + /// Input projections a neuron has a row in: one (up) or two (gate, then up). + pub fn inputs(n: Neuron) usize { + return switch (n.act) { + .relu2 => 1, + .swiglu => 2, + }; + } +}; + +/// The shared expert's storage: MLX affine, `bits` a code in u32 words, a bf16 scale and bias every `group`. +pub const Quant = struct { bits: u8 = 4, group: usize = 64 }; + +/// The MTP head's shared expert, widened with zero neurons when the model's are, so its tensors stay one width. +pub const Mtp = struct { + file: ?[]const u8 = null, // its own safetensors file in the folder; null: in the model's index beside the rest + base: []const u8, // its shared expert's tensors' prefix, ending in "." +}; + +/// Tensor names: layer i's shared expert is `.`, then `.weight|scales|biases`. +pub const Names = struct { + layers: []const u8, + expert: []const u8, + inputs: []const []const u8, // the input projections in Neuron.inputs order: a row a neuron each + output: []const u8 = "down_proj", // the output projection: a column a neuron + mtp: ?Mtp = null, +}; + +/// config.json keys: where the model's own keys sit, the shared expert's width, the pocket's width once written. +pub const Keys = struct { + scope: ?[]const u8 = null, // the object holding the model's keys (Flash Next: "text_config"); null: the top level + width: []const u8, // the shared expert's width, widened by the pocket on its first save + pocket: []const u8 = "tensorfold_learned_neurons", + layers: []const u8 = "num_hidden_layers", + hidden: []const u8 = "hidden_size", +}; + +/// Which layers have a shared expert: those a layer pattern marks MoE, or every layer. +pub const Experts = enum { pattern, every }; + +pub const Family = struct { + model_type: []const u8, + names: Names, + neuron: Neuron, + quant: Quant, + keys: Keys, + experts: Experts, + /// Pocket neurons at each layer: a whole number of quantization groups. + neurons: usize = 64, + + /// Layer `layer`'s shared expert's tensor prefix, into `buf`. + pub fn expertBase(f: Family, buf: []u8, layer: usize) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}{d}.{s}", .{ f.names.layers, layer, f.names.expert }); + } + + /// The layer a tensor name's shared expert sits at (null: not one of a layer's shared expert's tensors). + pub fn expertLayer(f: Family, name: []const u8) ?usize { + if (!std.mem.startsWith(u8, name, f.names.layers)) return null; + const rest = name[f.names.layers.len..]; + const dot = std.mem.indexOfScalar(u8, rest, '.') orelse return null; + if (!std.mem.startsWith(u8, rest[dot + 1 ..], f.names.expert)) return null; + return std.fmt.parseInt(usize, rest[0..dot], 10) catch null; + } + + /// Whether the description holds together: a neuron's rows match its inputs, the pocket a whole number of groups. + pub fn check(f: Family) !void { + if (f.names.inputs.len != f.neuron.inputs()) return error.BadFamily; + if (32 % f.quant.bits != 0 or f.quant.group == 0 or f.neurons == 0 or f.neurons % f.quant.group != 0) return error.BadFamily; + if (f.neuron.act == .swiglu and !(f.neuron.sharp > 0 and std.math.isPowerOfTwo(@as(u32, @intFromFloat(f.neuron.sharp))))) return error.BadFamily; + } +}; + +pub const max_layers = 128; +pub const Layers = std.bit_set.Static(max_layers); + +/// The served checkpoint as Living Weights reads it, built by the family (from its config) for its host. +pub const Layout = struct { + hidden: usize, + width: usize, // the shared expert's own neurons, a pocket written earlier left out + layers: usize, + widen: Layers = .empty, // layers whose shared expert carries the pocket in the shards (each with an expert) + learn: Layers = .empty, // of those, the ones lessons put neurons in + + /// The layers set in `which`, in order, into `out`. + pub fn list(lay: Layout, which: Layers, out: []usize) []usize { + var n: usize = 0; + for (0..lay.layers) |l| if (which.isSet(l)) { + out[n] = l; + n += 1; + }; + return out[0..n]; + } + + /// Every layer `has` names carries one widened; all but the last learn (a nudge after it has no gradient to read). + pub fn fromExperts(hidden: usize, width: usize, layers: usize, has: Layers) Layout { + var lay: Layout = .{ .hidden = hidden, .width = width, .layers = layers, .widen = has, .learn = has }; + if (layers > 0) lay.learn.unset(layers - 1); + return lay; + } +}; + +test "a family finds its shared expert's layer in a tensor name and builds its prefix" { + const f = @import("test_families.zig").swiglu; + try f.check(); + var buf: [128]u8 = undefined; + try std.testing.expectEqualStrings("language_model.model.layers.7.mlp.shared_expert.", try f.expertBase(&buf, 7)); + try std.testing.expectEqual(@as(?usize, 7), f.expertLayer("language_model.model.layers.7.mlp.shared_expert.up_proj.scales")); + try std.testing.expectEqual(@as(?usize, null), f.expertLayer("language_model.model.layers.7.mlp.shared_expert_gate.scales")); + try std.testing.expectEqual(@as(?usize, null), f.expertLayer("language_model.mtp.layers.0.mlp.shared_expert.up_proj.scales")); + const n = @import("families.zig").nemotron_h; + try n.check(); + try std.testing.expectEqual(@as(?usize, 13), n.expertLayer("backbone.layers.13.mixer.shared_experts.down_proj.weight")); + var has: Layers = .empty; + has.set(1); + has.set(3); + const lay = Layout.fromExperts(8, 16, 4, has); + var out: [4]usize = undefined; + try std.testing.expectEqualSlices(usize, &.{ 1, 3 }, lay.list(lay.widen, &out)); + try std.testing.expectEqualSlices(usize, &.{1}, lay.list(lay.learn, &out)); +} diff --git a/zig/src/living/host_test.zig b/zig/src/living/host_test.zig new file mode 100644 index 0000000000..e9cb25b719 --- /dev/null +++ b/zig/src/living/host_test.zig @@ -0,0 +1,215 @@ +//! The learner over a stand-in host and trainer on the CPU: a whole pocket lesson read, fitted, stepped, saved. +const std = @import("std"); +const learner = @import("learner.zig"); +const families = @import("families.zig"); +const fam = @import("family.zig"); +const train_mode = @import("train_mode.zig"); +const lesson_route = @import("lesson_route.zig"); +const clock = @import("pocket_clock.zig"); +const neuron = @import("neuron.zig"); + +const Example = @import("lesson_examples.zig").Example; +const Seq = train_mode.Seq; +const Span = train_mode.Span; +const Result = train_mode.Result; +const d = 64; +const layers = 8; + +/// A trainer with the pocket's hooks and nothing of Nemotron's routed blocks (no `sites`). +const Toy = struct { + spans: [4]Span = undefined, + pack: usize = 64, + count: clock.Count = .{}, + prefixes: struct { + pub const most = 4; + } = .{}, + nudge: train_mode.Nudge = .{}, + from: [layers]u32 = @splat(std.math.maxInt(u32)), + rows: std.ArrayList(u64) = .empty, // each pass row's hash of the tokens through it + gpa: std.mem.Allocator, + scratch: [d]u16 = undefined, + push: [d]f32 = undefined, + delta: [d]f32 = undefined, + gates: usize = 0, // gate scales read + + pub fn packs(_: *const Toy) bool { + return false; + } + + pub fn prefixSlots(_: *Toy, _: usize) !void {} + + fn lay(t: *Toy, seqs: []const Seq) !void { + t.rows.clearRetainingCapacity(); + for (seqs, 0..) |q, i| { + t.spans[i] = .{ .at = t.rows.items.len, .rows = q.ids.len - 1 - q.skip, .skip = q.skip, .slot = q.slot }; + var h = std.hash.Wyhash.init(0); + for (q.ids[0 .. q.ids.len - 1], 0..) |id, r| { + h.update(std.mem.asBytes(&id)); + if (r >= q.skip) try t.rows.append(t.gpa, h.final()); + } + } + t.count.passes += 1; + t.count.rows += t.rows.items.len; + } + + pub fn pass(t: *Toy, seqs: []const Seq, mode: train_mode.Mode, _: lesson_route.Turn, out: []Result) !void { + try t.lay(seqs); + for (out[0..seqs.len]) |*r| r.* = .{ .loss = if (mode == .keys) 0 else 0.5, .recalled = true }; + } + + pub fn step(t: *Toy, ids: []const u32, start: usize, mode: train_mode.Mode, turn: lesson_route.Turn) !Result { + var out: [1]Result = undefined; + try t.pass(&.{.{ .ids = ids, .start = start, .decide = turn.decide }}, mode, turn, &out); + return out[0]; + } + + /// Row `row`'s input at layer `l`: unit-sized noise from the row's hash, so different contexts read apart. + pub fn normed(t: *Toy, l: usize, row: usize) []const u16 { + var prng = std.Random.DefaultPrng.init(t.rows.items[row] ^ (@as(u64, l) *% 0x9e3779b97f4a7c15)); + for (&t.scratch) |*x| x.* = @truncate(@as(u32, @bitCast(1 + prng.random().floatNorm(f32))) >> 16); + return &t.scratch; + } + + pub fn residual(t: *Toy, l: usize, row: usize) []const u16 { + return t.normed(l + 100, row); + } + + pub fn pushAt(t: *Toy, _: usize, key: usize, l: usize) []const f32 { + for (&t.push, 0..) |*x, i| x.* = 0.01 * @as(f32, @floatFromInt((i + key + l) % 7)) - 0.03; + return &t.push; + } + + pub fn deltaAt(t: *Toy, _: usize, _: usize, _: usize) []f32 { + return &t.delta; + } + + pub fn gateAt(t: *Toy, _: usize, _: usize) f32 { + t.gates += 1; + return 0.25; + } +}; + +/// A host over the toy trainer for family `F`, counting what the learner hands it. +fn Host(comptime F: fam.Family) type { + return struct { + pub const Trainer = Toy; + pub const max_rows = 64; + pub const family = F; + pub const Ctx = *Seen; + + pub const Seen = struct { uploads: usize = 0, saved: usize = 0, gate_ok: bool = true, up_ok: bool = true, written: []const usize = &.{} }; + + pub fn layout(_: Ctx) fam.Layout { + var every: fam.Layers = .empty; + for (0..layers) |l| every.set(l); + return .fromExperts(d, 128, layers, every); + } + + pub fn trainer(gpa: std.mem.Allocator, _: std.Io, _: Ctx) !*Toy { + const t = try gpa.create(Toy); + t.* = .{ .gpa = gpa }; + return t; + } + + pub fn free(t: *Toy, gpa: std.mem.Allocator) void { + t.rows.deinit(gpa); + gpa.destroy(t); + } + + pub fn seconds(_: std.Io) f64 { + return 0; + } + + pub fn pocket(s: Ctx, _: usize, _: []const f32, _: []const f32, _: usize) !void { + s.uploads += 1; + } + + pub fn pocketRead(_: Ctx, _: usize, _: []f32, _: []f32) usize { + return 0; + } + + pub fn pocketSave(s: Ctx, _: std.mem.Allocator, _: std.Io, written: []const usize, _: usize, _: []f32, _: []f32) !usize { + s.saved += 1; + s.written = written; + return 1; + } + + /// The forward's rows before a save: each gate row sharp^2 times its up row (g = sharp u, w = u / sharp). + pub fn pocketSwiglu(s: Ctx, _: usize, gate: []const f32, up: []const f32, _: []const f32, _: usize) !void { + s.uploads += 1; + if (s.saved > 0) return; // after a save the rows are as written, each rounded its own way + for (gate, up) |g, w| if (@abs(g - F.neuron.sharp * F.neuron.sharp * w) > 1e-5 * @abs(g) + 1e-30) { + s.up_ok = false; + }; + } + + pub fn pocketReadSwiglu(_: Ctx, _: usize, _: []f32, _: []f32, _: []f32) usize { + return 0; + } + + pub fn pocketSaveSwiglu(s: Ctx, _: std.mem.Allocator, _: std.Io, written: []const usize, _: usize, gate: []f32, up: []f32, _: []f32) !usize { + s.saved += 1; + s.written = written; + for (gate, up) |g, w| if (@abs(g - F.neuron.sharp * F.neuron.sharp * w) > 1e-5 * @abs(g) + 1e-30) { + s.gate_ok = false; + }; + // handed back as a writer would: rounded, so the learner takes u from the up rows as written + for (up) |*w| w.* = std.math.round(w.* * 4096) / 4096; + return 1; + } + }; +} + +/// A fact's question and answer: token `who` picks whose fact it is, so its rows read apart from others'. +fn qa(gpa: std.mem.Allocator, who: u32, k: u32) !Example { + const ids = try gpa.dupe(u32, &.{ 1, 2, who, 40 + k, 5, 90 + who, 91, 92 }); + return .{ .ids = ids, .start = 5, .decide = 4 }; +} + +fn run(comptime F: fam.Family) !void { + const gpa = std.testing.allocator; + const H = Host(F); + var seen: H.Seen = .{}; + var l = learner.Learner(H).init(gpa, std.testing.io, &seen); + defer l.deinit(); + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var train: [6]Example = undefined; + var held: [2]Example = undefined; + var other: [8]Example = undefined; + for (&train, 0..) |*x, k| x.* = try qa(a, 7, @intCast(k)); + for (&held, 0..) |*x, k| x.* = try qa(a, 7, @intCast(10 + k)); + for (&other, 0..) |*x, k| x.* = try qa(a, 20 + @as(u32, @intCast(k)), @intCast(k)); + try l.begin(.{ .train = &train, .held = &held, .keep = &other }); + var report: ?learner.Report = null; + var steps: usize = 0; + while (true) : (steps += 1) { + const s = l.step(); + if (s.report) |r| report = r; + if (s.done) break; + try std.testing.expect(steps < 10_000); + } + try std.testing.expect(report != null and report.? == .learned); + try std.testing.expect(report.?.learned.recalled); + var placed: [64]learner.Learner(H).Placed = undefined; + const got = l.placed(&placed); + try std.testing.expect(got.len > 0); + for (got) |x| try std.testing.expect(x.layer > 0 and x.layer < layers - 1); + try std.testing.expect(seen.uploads > 0); + // the save: through the family's own hook, every learning layer handed over + try l.begin(.{ .save = true }); + while (!l.step().done) {} + try std.testing.expectEqual(@as(usize, 1), seen.saved); + try std.testing.expectEqual(@as(usize, layers - 1), seen.written.len); + try std.testing.expect(seen.gate_ok and seen.up_ok); + if (F.neuron.gated) try std.testing.expect(l.trainer.?.gates > 0); +} + +test "a pocket lesson over a stand-in relu^2 host: read, fitted, stepped, judged, saved" { + try run(families.nemotron_h); +} + +test "a pocket lesson over a stand-in gated SwiGLU host: its rows handed over as gate sharp u and up u / sharp, the gate read" { + try run(@import("test_families.zig").swiglu); +} diff --git a/zig/src/living/learn_shape.zig b/zig/src/living/learn_shape.zig new file mode 100644 index 0000000000..1ea2572a76 --- /dev/null +++ b/zig/src/living/learn_shape.zig @@ -0,0 +1,25 @@ +//! The learned change's shape, shared by the Metal and CUDA learners: blocks of ranks, their scale and gates. + +pub const block = 16; // ranks a lesson adds at each layer +pub const max_rank = 2048; // a hundred and twenty-eight lessons' blocks, each layer's buffers grown to it as lessons come +pub const max_blocks = max_rank / block; +pub const scale: f32 = 10; +pub const avoid_dims = 192; // directions sketched from the inputs a new block must leave alone +pub const candidates = 64; // directions sketched from the fact's inputs, which a new block's are chosen among + +/// Adam's settings for a block's outputs (its gate keeps them off what must stay): rate, moments' decays, epsilon. +pub const hyper = [4]f32{ 3e-4, 0.9, 0.999, 1e-8 }; + +/// Adam's settings at step `step`: the rate halves every `settle` steps, three times at most, so later rounds refine. +pub fn hyperAt(step: u64) [4]f32 { + var h = hyper; + h[0] /= @as(f32, @floatFromInt(@as(u32, 1) << @intCast(@min(step / settle, 3)))); + return h; +} +pub const settle = 60; + +/// Each new input direction's length, as MLX's LoRA draws its down factor's columns (uniform within 1 / sqrt(in)). +pub const a_norm: f32 = 0.577; + +/// A block's directions' squared length: a row's cosine with one is (x . a) / sqrt(unit |x|^2). +pub const unit = a_norm * a_norm; diff --git a/zig/src/families/nemotron/learner.zig b/zig/src/living/learner.zig similarity index 91% rename from zig/src/families/nemotron/learner.zig rename to zig/src/living/learner.zig index 0607eb90c1..ddad1c54b8 100644 --- a/zig/src/families/nemotron/learner.zig +++ b/zig/src/living/learner.zig @@ -1,4 +1,4 @@ -//! Sliding Weights' lessons, one learner for every backend over a host that gives the trainer, weights and files. +//! Living Weights' lessons, one learner for every family and backend over a host that gives the trainer and hooks. const std = @import("std"); const adapters = @import("learn_shape.zig"); const choice = @import("choice.zig"); @@ -42,16 +42,17 @@ const replay_each = 8; // of a lesson's answers, its first (plain questions) joi /// A lesson builds (build) and learns (steps); an edit is written at its subject (edit), more rounds report (again). const Phase = enum { idle, build, steps, undo, reroute, drop, save, edit, again, probe, pocket, retire }; -/// The lessons over host H: its Trainer, max_rows, Ctx, layers, trainer, free, attach, write and seconds. +/// The lessons over host H: its trainer and hooks (family.zig), routed blocks only when its trainer keeps `sites`. pub fn Learner(comptime H: type) type { return struct { const Self = @This(); - const has_edits = @hasDecl(H.Trainer, "key") and @hasDecl(H, "commit"); // keys, a nudged row, codes to write - const has_pocket = @hasDecl(H.Trainer, "normed") and @hasDecl(H, "pocketSave"); // inputs, gradients, a pocket to write + const has_sites = @hasField(H.Trainer, "sites"); // Nemotron's routed blocks of ranks (the LoRA path) + const has_edits = has_sites and @hasDecl(H.Trainer, "key") and @hasDecl(H, "commit"); // keys, a nudged row, codes to write + const has_pocket = @hasDecl(H.Trainer, "normed") and (@hasDecl(H, "pocketSave") or @hasDecl(H, "pocketSaveSwiglu")); // inputs, gradients, a pocket to write const has_whole = has_edits or has_pocket; // lessons that change the model itself const Edits = if (has_edits) editor.Editor(H.Trainer) else void; const Probes = if (has_edits) probe.Probe(H.Trainer) else void; - const Pockets = if (has_pocket) pocket.Pocket(H.Trainer) else void; + const Pockets = if (has_pocket) pocket.Pocket(H.Trainer, H.family) else void; pub const Host = H; gpa: std.mem.Allocator, @@ -96,7 +97,7 @@ pub fn Learner(comptime H: type) type { pocketed: bool = false, // the last lesson put neurons in the pocket pub fn init(gpa: std.mem.Allocator, io: std.Io, b: H.Ctx) Self { - return .{ .gpa = gpa, .io = io, .b = b, .whole = std.c.getenv("TF_SLIDE_ROUTED") == null, .edits = if (has_edits) Edits.init(gpa) else {}, .method = if (has_edits) whole_lessons.method() else .pocket, .probe = if (has_edits) Probes.init(gpa) else {}, .pockets = if (has_pocket) Pockets.init(gpa) else {} }; + return .{ .gpa = gpa, .io = io, .b = b, .whole = !has_sites or std.c.getenv("TF_SLIDE_ROUTED") == null, .edits = if (has_edits) Edits.init(gpa) else {}, .method = if (has_edits) whole_lessons.method() else .pocket, .probe = if (has_edits) Probes.init(gpa) else {}, .pockets = if (has_pocket) Pockets.init(gpa) else {} }; } pub fn deinit(l: *Self) void { @@ -177,9 +178,11 @@ pub fn Learner(comptime H: type) type { return l.start(); } try l.settleLast(lesson.train); - if (l.trainer.?.sites.rank + adapters.block > adapters.max_rank) return error.LearnedChangeFull; - l.trainer.?.sites.clear(); - l.trainer.?.sketched = 0; + if (has_sites) { + if (l.trainer.?.sites.rank + adapters.block > adapters.max_rank) return error.LearnedChangeFull; + l.trainer.?.sites.clear(); + l.trainer.?.sketched = 0; + } l.built = 0; l.bits.clearRetainingCapacity(); l.seen.clearRetainingCapacity(); @@ -189,6 +192,7 @@ pub fn Learner(comptime H: type) type { l.keeps.clearRetainingCapacity(); l.topic = false; if (l.whole) return if (has_whole) whole_lessons.begin(l, lesson) else error.EditsUnsupported; + if (comptime !has_sites) unreachable; // a host without routed blocks learns whole lessons only l.gpa.free(l.answers); l.answers = try std.mem.concat(l.gpa, Example, &.{ lesson.train, lesson.think }); l.phase = .build; @@ -202,18 +206,21 @@ pub fn Learner(comptime H: type) type { fn ready(l: *Self) !void { if (l.trainer != null) return; const t0 = H.seconds(l.io); - l.choice = try choice.Choice.init(l.gpa, H.layers(l.b)); + if (has_sites) l.choice = try choice.Choice.init(l.gpa, H.layers(l.b)); l.trainer = try H.trainer(l.gpa, l.io, l.b); std.log.info("slide: the trainer is ready in {d:.1} s", .{H.seconds(l.io) - t0}); - l.trainer.?.sites.whole = l.whole; - for (&l.asks) |*side| { - side.* = try l.gpa.alloc(std.ArrayList(f32), H.layers(l.b)); - for (side.*) |*rows| rows.* = .empty; + if (has_sites) { + l.trainer.?.sites.whole = l.whole; + for (&l.asks) |*side| { + side.* = try l.gpa.alloc(std.ArrayList(f32), H.layers(l.b)); + for (side.*) |*rows| rows.* = .empty; + } } } /// A round's steps from the open block as it is now, which an undo comes back to. fn start(l: *Self) !void { + if (comptime !has_sites) return error.RoutedUnsupported; try l.schedule(@max(@min(l.lesson.steps, max_steps), 1)); l.trainer.?.sites.keep(); l.can_undo = true; @@ -230,7 +237,7 @@ pub fn Learner(comptime H: type) type { /// One bounded unit: one example sketched or projected, one step of learning (a check every few), or the undo. pub fn step(l: *Self) Step { return l.advance() catch |e| { - if (l.phase == .steps) l.trainer.?.sites.restore(); + if (has_sites and l.phase == .steps) l.trainer.?.sites.restore(); l.phase = .idle; return .{ .done = true, .changed = true, .report = .{ .failed = @errorName(e) } }; }; @@ -242,7 +249,7 @@ pub fn Learner(comptime H: type) type { .undo => { l.phase = .idle; if (has_whole and l.whole and try whole_lessons.undo(l)) return .{ .done = true, .changed = true }; - if (!l.can_undo) return .{ .done = true }; + if (!has_sites or !l.can_undo) return .{ .done = true }; l.trainer.?.sites.restore(); l.can_undo = false; l.kept_rounds -|= 1; @@ -283,6 +290,7 @@ pub fn Learner(comptime H: type) type { if (has_whole and !try whole_lessons.save(l)) try whole_lessons.foldOpen(l); return .{ .done = true, .changed = true }; } + if (comptime !has_sites) return .{ .done = true }; const ranks = t.sites.rank - @as(usize, if (l.opened and l.kept_rounds == 0) adapters.block else 0); if (ranks == 0) return .{ .done = true }; const t0 = H.seconds(l.io); @@ -298,6 +306,7 @@ pub fn Learner(comptime H: type) type { /// A pass of `ex` in `mode` (`acting`: the open lesson acts on its reply); its first reads which kept ones act. fn pass(l: *Self, ex: Example, mode: train_mode.Mode, acting: bool) !f32 { + if (comptime !has_sites) return error.RoutedUnsupported; const t = l.trainer.?; const k = key(ex); const known = l.bits.get(k); @@ -311,6 +320,7 @@ pub fn Learner(comptime H: type) type { /// Every lesson kept before the open one, all acting (an in-place learner's lessons are part of the model). pub fn keptLessons(l: *const Self) Acting { var bits: Acting = .empty; + if (comptime !has_sites) return bits; // no blocks: the pocket is the model's own for (0..l.trainer.?.sites.rank / adapters.block - @intFromBool(l.opened)) |b| bits.set(b); return bits; } @@ -323,12 +333,14 @@ pub fn Learner(comptime H: type) type { /// Layer by layer, a routing prompt's decision row along the candidates, as the last project pass left it. fn ask(l: *Self, r: Asked) !void { + if (comptime !has_sites) return error.RoutedUnsupported; const t = l.trainer.?; for (l.asks[@backingInt(r.kind)], 0..) |*rows, k| try rows.appendSlice(l.gpa, t.projected(k, r.ex.ids.len - 1)[r.ex.decide * adapters.candidates ..][0..adapters.candidates]); } /// One example at a time: sketched, projected, routing rows copied; then the block chosen, its router fitted. fn buildOnce(l: *Self) !Step { + if (comptime !has_sites) return error.RoutedUnsupported; const t = l.trainer.?; const c = &l.choice.?; // an edit of the weights also keeps the questions about others, and reads every wording of the fact's @@ -383,6 +395,7 @@ pub fn Learner(comptime H: type) type { /// The open block's router: if facts don't stand apart it acts on what most clear, learning to leave the rest. fn route(l: *Self) !void { + if (comptime !has_sites) return error.RoutedUnsupported; const t = l.trainer.?; var fit = try l.fitWith(&.{.fact}, &.{ .near, .swap, .other }); defer fit.deinit(l.gpa); @@ -447,6 +460,7 @@ pub fn Learner(comptime H: type) type { /// A taken question: a topic lesson keeps its answer for next round; a strict one refits its router without it. fn rerouteOnce(l: *Self) !Step { if (l.whole) return if (has_whole) whole_lessons.reroute(l) else unreachable; + if (comptime !has_sites) return error.RoutedUnsupported; if (l.topic) { for (l.lesson.kept) |ex| try l.keeps.append(l.gpa, .{ .ids = try l.gpa.dupe(u32, ex.ids), .start = ex.start, .decide = ex.decide }); l.phase = .idle; @@ -513,7 +527,7 @@ pub fn Learner(comptime H: type) type { const t = l.trainer.?; for (held) |ex| { var bits = if (l.whole) l.keptLessons() else l.earlier(ex); - if (open) bits.set(t.sites.first() / adapters.block); + if (has_sites and open) bits.set(t.sites.first() / adapters.block); if (!(try t.step(ex.ids, ex.start, .loss, .{ .decide = ex.decide, .bits = bits, .whole = l.whole })).recalled) return false; } return true; @@ -523,7 +537,7 @@ pub fn Learner(comptime H: type) type { fn settleLast(l: *Self, answers: []const Example) !void { const now = try own(l.gpa, answers); if (has_whole and l.whole) try whole_lessons.settle(l); - if (l.opened and l.kept_rounds == 0) { + if (has_sites and l.opened and l.kept_rounds == 0) { l.trainer.?.sites.needs()[l.trainer.?.sites.first() / adapters.block] = 0; l.trainer.?.sites.close(); try H.attach(l.b, l.trainer.?); diff --git a/zig/src/families/nemotron/lesson_examples.zig b/zig/src/living/lesson_examples.zig similarity index 100% rename from zig/src/families/nemotron/lesson_examples.zig rename to zig/src/living/lesson_examples.zig diff --git a/zig/src/families/nemotron/lesson_route.zig b/zig/src/living/lesson_route.zig similarity index 100% rename from zig/src/families/nemotron/lesson_route.zig rename to zig/src/living/lesson_route.zig diff --git a/zig/src/living/living.zig b/zig/src/living/living.zig new file mode 100644 index 0000000000..e1d75f4163 --- /dev/null +++ b/zig/src/living/living.zig @@ -0,0 +1,50 @@ +//! Living Weights, family-neutral: the learner, its pocket of new neurons, their fits and the shard writer. +pub const family = @import("family.zig"); +pub const families = @import("families.zig"); +pub const Family = family.Family; +pub const Layout = family.Layout; +pub const Layers = family.Layers; + +pub const learner = @import("learner.zig"); +pub const whole = @import("whole.zig"); +pub const choice = @import("choice.zig"); +pub const editor = @import("editor.zig"); +pub const probe = @import("probe.zig"); +pub const router = @import("router.zig"); +pub const subspace = @import("subspace.zig"); +pub const train_mode = @import("train_mode.zig"); +pub const learn_shape = @import("learn_shape.zig"); +pub const lesson_examples = @import("lesson_examples.zig"); +pub const lesson_route = @import("lesson_route.zig"); +pub const neuron = @import("neuron.zig"); +pub const affine = @import("affine.zig"); +pub const shards = @import("shards.zig"); +pub const pocket = @import("pocket.zig"); +pub const pocket_clock = @import("pocket_clock.zig"); +pub const pocket_dump = @import("pocket_dump.zig"); +pub const pocket_fit = @import("pocket_fit.zig"); +pub const pocket_joint = @import("pocket_joint.zig"); +pub const pocket_load = @import("pocket_load.zig"); +pub const pocket_pack = @import("pocket_pack.zig"); +pub const pocket_prefix = @import("pocket_prefix.zig"); +pub const pocket_quiet = @import("pocket_quiet.zig"); +pub const pocket_read = @import("pocket_read.zig"); +pub const pocket_values = @import("pocket_values.zig"); + +test { + _ = family; + _ = families; + _ = neuron; + _ = affine; + _ = shards; + _ = learner; + _ = whole; + _ = subspace; + _ = pocket; + _ = pocket_fit; + _ = router; + _ = choice; + _ = lesson_route; + _ = @import("swiglu_test.zig"); + _ = @import("host_test.zig"); +} diff --git a/zig/src/living/neuron.zig b/zig/src/living/neuron.zig new file mode 100644 index 0000000000..52db123839 --- /dev/null +++ b/zig/src/living/neuron.zig @@ -0,0 +1,55 @@ +//! A pocket neuron's activation for each family's kind, and the gate and up rows a SwiGLU neuron is written as. +const std = @import("std"); +const Neuron = @import("family.zig").Neuron; + +pub fn sigmoid(y: f32) f32 { + return 1 / (1 + @exp(-y)); +} + +pub fn silu(y: f32) f32 { + return y * sigmoid(y); +} + +/// The activation of a neuron of kind `n` at input x: `z` is u x (u as fitted), `zg` its gate row's g x (SwiGLU). +pub fn act(n: Neuron, z: f32, zg: f32) f32 { + return switch (n.act) { + .relu2 => if (z > 0) z * z else 0, + .swiglu => silu(zg) * (z / n.sharp), // its up row is u / sharp + }; +} + +/// A SwiGLU neuron's gate row from its fitted input u: sharp u. +pub fn gateRow(n: Neuron, u: []const f32, out: []f32) void { + for (out, u) |*o, x| o.* = n.sharp * x; +} + +/// The up row a SwiGLU neuron is written as, from u: u / sharp (exact for a power of two). +pub fn upRow(n: Neuron, u: []const f32, out: []f32) void { + for (out, u) |*o, x| o.* = x / n.sharp; +} + +/// u back from an up row as written: sharp w. +pub fn fromUp(n: Neuron, w: []const f32, out: []f32) void { + for (out, w) |*o, x| o.* = n.sharp * x; +} + +test "a SwiGLU neuron written as sharp u and u / sharp fires as relu^2 on its rows and stays near zero past the margin" { + const n: Neuron = .{ .act = .swiglu, .margin = 0.5, .sharp = 16 }; + const r: Neuron = .{ .act = .relu2 }; + // z = u x; the gate row reads sharp z, and the up row's z / sharp times silu is what the forward adds + for ([_]f32{ 1, 1.5, 3 }) |z| { + try std.testing.expectApproxEqRel(act(r, z, 0), act(n, z, n.sharp * z), 1e-6); + } + var most: f32 = 0; + var z: f32 = -n.margin; + while (z > -20) : (z -= 0.01) most = @max(most, @abs(act(n, z, n.sharp * z))); + try std.testing.expect(most < n.margin * n.margin * sigmoid(-n.sharp * n.margin) * 1.01); + try std.testing.expect(most < 1e-4); + // the up row as written gives u back exactly + const u = [_]f32{ 0.123, -7.5e-3, 3.25 }; + var w: [3]f32 = undefined; + var back: [3]f32 = undefined; + upRow(n, &u, &w); + fromUp(n, &w, &back); + try std.testing.expectEqualSlices(f32, &u, &back); +} diff --git a/zig/src/families/nemotron/pocket.zig b/zig/src/living/pocket.zig similarity index 83% rename from zig/src/families/nemotron/pocket.zig rename to zig/src/living/pocket.zig index ff24ac5ee9..aff4a4ad4a 100644 --- a/zig/src/families/nemotron/pocket.zig +++ b/zig/src/living/pocket.zig @@ -8,6 +8,10 @@ const joint = @import("pocket_joint.zig"); const pocket_dump = @import("pocket_dump.zig"); const clock = @import("pocket_clock.zig"); const reading = @import("pocket_read.zig"); +const quieting = @import("pocket_quiet.zig"); +const loading = @import("pocket_load.zig"); +const family = @import("family.zig"); +const neuron_math = @import("neuron.zig"); const Example = examples.Example; const Acting = lesson_route.Acting; @@ -16,15 +20,10 @@ const Fit = pocket_fit.Fit; const fit = pocket_fit.fit; const fitLayer = pocket_fit.fitLayer; const dotBf16 = pocket_fit.dotBf16; -const margin = pocket_fit.margin; const fit_steps = pocket_fit.fit_steps; -/// A pocket's neurons at each MoE layer: one 4-bit group of the shared expert's width. -pub const neurons = 64; const per_fact = 10; // neurons a fact takes, at the layers that tell its questions apart best const value_neurons = 4; // of those, the ones before its answer's value, when its answers show one -const kept_least: f32 = 0.5; // an earlier neuron refitted must still reach this on every one of its own fact's rows -const refit_above: f32 = 0.5; // an earlier neuron is refitted only where it fires above this on a new fact's rows /// Where a neuron fires: the reply's first row, the named subject's last token, or the row before the answer's value. pub const Key = enum(u2) { reply, subject, value }; @@ -85,10 +84,16 @@ fn better(f: []const Fit, a: usize, b: usize) bool { return f[a].fact_least > f[b].fact_least; } -/// A pocket over trainer T (Metal's: input, residual, nudges) and a host that hands each layer's to the forward. -pub fn Pocket(comptime T: type) type { +/// A pocket over trainer T (its inputs, residuals, nudges) for family F, and a host handing each layer to the forward. +pub fn Pocket(comptime T: type, comptime F: family.Family) type { + comptime F.check() catch unreachable; return struct { const Self = @This(); + /// A pocket's neurons at each layer it sits at: a whole number of quantization groups of the expert's width. + pub const neurons = F.neurons; + pub const of = F; // the family it learns into + pub const swiglu = F.neuron.act == .swiglu; + pub const gated = F.neuron.gated; gpa: std.mem.Allocator, mode: Mode, // as asked (TF_SLIDE_POCKET_KEY), auto by default @@ -105,9 +110,13 @@ pub fn Pocket(comptime T: type) type { split: [3]usize = .{ per_fact, 0, 0 }, // the open fact's neurons on each key subjects: std.ArrayList(u32) = .empty, // the open fact's subject tokens leads: std.ArrayList(u32) = .empty, // the tokens before them: after another, the subject is someone else's - moe: std.ArrayList(usize) = .empty, // the MoE layers a pocket sits at, its slots in order - up: std.ArrayList(f32) = .empty, // [slots, neurons, D]: each neuron's input weights + moe: std.ArrayList(usize) = .empty, // the layers a pocket learns at (the layout's), its slots in order + hidden: usize = 0, // D, as the layout gives it + layers: usize = 0, // the model's layers, as the layout gives them + up: std.ArrayList(f32) = .empty, // [slots, neurons, D]: each neuron's input weights u, as fitted down: std.ArrayList(f32) = .empty, // [slots, neurons, D]: each neuron's output + gate: std.ArrayList(f32) = .empty, // SwiGLU: [slots, neurons, D], each neuron's gate row (sharp u, or as written) + scratch: std.ArrayList(f32) = .empty, // SwiGLU: [neurons, D], a slot's up rows (u / sharp) as handed over used: std.ArrayList(usize) = .empty, // [slots]: neurons in use, the first ones taken: std.ArrayList(Taken) = .empty, // the open fact's neurons answered: []const Example = &.{}, // the fact's answers, which its neurons' outputs are fitted on @@ -124,6 +133,7 @@ pub fn Pocket(comptime T: type) type { retired_marks: usize = 0, // retired_lessons before the last retire, which a put-back returns to retired: std.ArrayList([2]usize) = .empty, // the last retire's (slot, neuron), put back if the newer fact fails retired_w: std.ArrayList(f32) = .empty, // [retired, 2, D]: their input and output weights from before + refit_gate: std.ArrayList(f32) = .empty, // SwiGLU: [refits, D], their gate rows from before phase: enum { read, fit, values, check, done } = .done, at: usize = 0, // the next example read, answer of a step, or held-out answer judged back: usize = 0, // held-out answers that came back token for token @@ -141,7 +151,7 @@ pub fn Pocket(comptime T: type) type { } pub fn deinit(p: *Self) void { - inline for (.{ "moe", "up", "down", "used", "taken", "rows", "labels", "keeps", "refits", "refit_up", "retired_lessons", "retired", "retired_w", "joint", "joint_down", "joint_answers", "learned", "pos", "subjects", "leads" }) |f| @field(p, f).deinit(p.gpa); + inline for (.{ "moe", "up", "down", "gate", "scratch", "used", "taken", "rows", "labels", "keeps", "refits", "refit_up", "refit_gate", "retired_lessons", "retired", "retired_w", "joint", "joint_down", "joint_answers", "learned", "pos", "subjects", "leads" }) |f| @field(p, f).deinit(p.gpa); p.read.deinit(p.gpa); p.v.deinit(p.gpa); } @@ -188,23 +198,29 @@ pub fn Pocket(comptime T: type) type { return which[(slot * neurons + n) * d ..][0..d]; } - /// The pocket's neurons as the shards hold them, read once: every MoE layer but the last. + /// Neuron n of `slot`'s activation at input `x` (bf16): relu(u x)^2, or silu(g x) (u x) / sharp (neuron.zig). + pub fn act(p: *const Self, slot: usize, n: usize, x: []const u16) f32 { + const d = p.hidden; + const z = dotBf16(p.neuron(@constCast(p.up.items), slot, n, d), x); + const zg = if (swiglu) dotBf16(p.neuron(@constCast(p.gate.items), slot, n, d), x) else 0; + return neuron_math.act(F.neuron, z, zg); + } + + /// A SwiGLU neuron's gate row set from its input weights as they are now. + pub fn gateFrom(p: *Self, slot: usize, n: usize) void { + if (!swiglu) return; + neuron_math.gateRow(F.neuron, p.neuron(p.up.items, slot, n, p.hidden), p.neuron(p.gate.items, slot, n, p.hidden)); + } + fn load(p: *Self, t: *T, host: anytype) !void { - if (p.moe.items.len > 0) return; - const c = t.config(); - // every MoE layer but the last (a nudge after it has no gradient to read), its neurons as loaded - for (0..c.layers - 1) |l| if (c.kinds[l] == .moe) try p.moe.append(p.gpa, l); - const n = neurons * c.hidden; - try p.up.appendNTimes(p.gpa, 0, p.slots() * n); - try p.down.appendNTimes(p.gpa, 0, p.slots() * n); - try p.used.appendNTimes(p.gpa, 0, p.slots()); - for (p.moe.items, 0..) |l, j| p.used.items[j] = host.pocketRead(l, p.up.items[j * n ..][0..n], p.down.items[j * n ..][0..n]); + _ = t; + return loading.load(p, host); } /// A replaced fact's neurons, (layer, neuron) each, zeroed where they sit; their places stay taken. pub fn retire(p: *Self, t: *T, host: anytype, gone: []const [2]u16) !void { try p.load(t, host); - const d = t.config().hidden; + const d = p.hidden; p.retired.clearRetainingCapacity(); p.retired_w.clearRetainingCapacity(); p.retired_marks = p.retired_lessons.items.len; @@ -225,7 +241,7 @@ pub fn Pocket(comptime T: type) type { /// The last retire undone: its neurons as they were (the fact that replaced theirs was not learned). pub fn unretire(p: *Self, t: *T, host: anytype) !void { try p.load(t, host); - const d = t.config().hidden; + const d = p.hidden; for (p.retired.items, 0..) |r, i| { @memcpy(p.neuron(p.up.items, r[0], r[1], d), p.retired_w.items[2 * i * d ..][0..d]); @memcpy(p.neuron(p.down.items, r[0], r[1], d), p.retired_w.items[(2 * i + 1) * d ..][0..d]); @@ -260,6 +276,7 @@ pub fn Pocket(comptime T: type) type { p.refit_up.clearRetainingCapacity(); p.joint.clearRetainingCapacity(); p.joint_down.clearRetainingCapacity(); + p.refit_gate.clearRetainingCapacity(); p.learned_before = p.learned.items.len; p.pos_before = p.pos.items.len; p.clock = .{}; @@ -303,7 +320,7 @@ pub fn Pocket(comptime T: type) type { /// The open fact's neurons back to zero and free, and the earlier neurons it fitted again as they were. pub fn undo(p: *Self, t: *T, host: anytype) !void { - const d = t.config().hidden; + const d = p.hidden; p.learned.shrinkRetainingCapacity(p.learned_before); p.pos.shrinkRetainingCapacity(p.pos_before); var i = p.taken.items.len; @@ -312,22 +329,32 @@ pub fn Pocket(comptime T: type) type { const k = p.taken.items[i]; @memset(p.neuron(p.up.items, k.slot, k.neuron, d), 0); @memset(p.neuron(p.down.items, k.slot, k.neuron, d), 0); + if (swiglu) @memset(p.neuron(p.gate.items, k.slot, k.neuron, d), 0); if (k.neuron + 1 == p.used.items[k.slot]) p.used.items[k.slot] -= 1; try p.upload(t, host, k.slot); } for (p.refits.items, 0..) |r, j| { @memcpy(p.neuron(p.up.items, r.slot, r.neuron, d), p.refit_up.items[j * d ..][0..d]); + if (swiglu) @memcpy(p.neuron(p.gate.items, r.slot, r.neuron, d), p.refit_gate.items[j * d ..][0..d]); try p.upload(t, host, r.slot); } p.refits.clearRetainingCapacity(); p.refit_up.clearRetainingCapacity(); try joint.putBack(p, t, host, d); + p.refit_gate.clearRetainingCapacity(); p.taken.clearRetainingCapacity(); } + /// Slot `slot`'s neurons into the forward: relu^2 as u, SwiGLU as its gate rows and up rows u / sharp. pub fn upload(p: *Self, t: *T, host: anytype, slot: usize) !void { - const n = neurons * t.config().hidden; - try host.pocket(p.moe.items[slot], p.up.items[slot * n ..][0..n], p.down.items[slot * n ..][0..n], p.used.items[slot]); + _ = t; + const n = neurons * p.hidden; + const up = p.up.items[slot * n ..][0..n]; + const down = p.down.items[slot * n ..][0..n]; + if (swiglu) { + neuron_math.upRow(F.neuron, up, p.scratch.items); + try host.pocketSwiglu(p.moe.items[slot], p.gate.items[slot * n ..][0..n], p.scratch.items, down, p.used.items[slot]); + } else try host.pocket(p.moe.items[slot], up, down, p.used.items[slot]); } /// Example i and its kind: the fact's answers, then its other wordings, held-out answers and the other prompts. @@ -356,8 +383,7 @@ pub fn Pocket(comptime T: type) type { /// Every key's candidate layers fitted in threads, the best taken; then keeps found, earlier neurons refitted. fn fitAll(p: *Self, t: *T, host: anytype) !bool { - const c = t.config(); - const d = c.hidden; + const d = p.hidden; const n = p.slots(); const m = p.labels.items.len; if (std.c.getenv("TF_SLIDE_POCKET_DUMP")) |dir| pocket_dump.write(p, std.mem.span(dir), p.learned_before, d) catch |e| std.log.warn("slide pocket: the read rows were not written: {t}", .{e}); @@ -372,14 +398,17 @@ pub fn Pocket(comptime T: type) type { defer p.gpa.free(fits); @memset(fits, .{ .fires = std.math.maxInt(usize) }); const t0 = clock.now(host.io); - var threads: [3 * 64]?std.Thread = @splat(null); + const threads = try p.gpa.alloc(?std.Thread, 3 * n); + defer p.gpa.free(threads); + @memset(threads, null); var fitted: usize = 0; - for (0..3) |k| for (p.moe.items, 0..) |l, j| if (p.split[k] > 0 and candidate(l, c.layers, @fromBackingInt(@intCast(k)))) { - threads[k * n + j] = std.Thread.spawn(.{}, fitLayer, .{ p.gpa, p.rows.items, kinds[k * m ..][0..m], n, j, d, us[(k * n + j) * d ..][0..d], &fits[k * n + j] }) catch null; + const below = F.neuron.margin; + for (0..3) |k| for (p.moe.items, 0..) |l, j| if (p.split[k] > 0 and candidate(l, p.layers, @fromBackingInt(@intCast(k)))) { + threads[k * n + j] = std.Thread.spawn(.{}, fitLayer, .{ p.gpa, p.rows.items, kinds[k * m ..][0..m], n, j, d, us[(k * n + j) * d ..][0..d], &fits[k * n + j], below }) catch null; fitted += 1; }; - for (0..3) |k| for (p.moe.items, 0..) |l, j| if (p.split[k] > 0 and candidate(l, c.layers, @fromBackingInt(@intCast(k)))) { - if (threads[k * n + j]) |th| th.join() else fitLayer(p.gpa, p.rows.items, kinds[k * m ..][0..m], n, j, d, us[(k * n + j) * d ..][0..d], &fits[k * n + j]); + for (0..3) |k| for (p.moe.items, 0..) |l, j| if (p.split[k] > 0 and candidate(l, p.layers, @fromBackingInt(@intCast(k)))) { + if (threads[k * n + j]) |th| th.join() else fitLayer(p.gpa, p.rows.items, kinds[k * m ..][0..m], n, j, d, us[(k * n + j) * d ..][0..d], &fits[k * n + j], below); }; const t1 = clock.now(host.io); for (0..3) |k| if (p.split[k] > 0) { @@ -387,14 +416,14 @@ pub fn Pocket(comptime T: type) type { var order: std.ArrayList(usize) = .empty; defer order.deinit(p.gpa); // a layer whose pocket is full is passed over - for (p.moe.items, 0..) |l, j| if (candidate(l, c.layers, key) and p.used.items[j] < neurons) try order.append(p.gpa, j); + for (p.moe.items, 0..) |l, j| if (candidate(l, p.layers, key) and p.used.items[j] < neurons) try order.append(p.gpa, j); if (order.items.len == 0) return error.PocketFull; std.mem.sort(usize, order.items, @as([]const Fit, fits[k * n ..][0..n]), better); for (order.items[0..@min(p.split[k], order.items.len)]) |j| try p.take(j, key, us[(k * n + j) * d ..][0..d], fits[k * n + j], kinds[k * m ..][0..m], d); }; p.logSources(d, m, kinds); try p.findKeeps(d); - try p.quiet(t, host); + try quieting.quiet(p, t, host); std.log.info("slide pocket: {d} candidate neurons fitted on {d} read rows in {d:.1} s; earlier neurons checked in {d:.1} s", .{ fitted, m, t1 - t0, clock.now(host.io) - t1 }); try joint.find(p, d); try p.v.begin(p.gpa, p.taken.items.len + p.joint.items.len, d, p.answered.len + p.joint_answers.items.len + p.keeps.items.len); @@ -430,6 +459,7 @@ pub fn Pocket(comptime T: type) type { p.used.items[j] += 1; try p.taken.append(p.gpa, .{ .slot = j, .neuron = k, .key = key }); @memcpy(p.neuron(p.up.items, j, k, d), u); + p.gateFrom(j, k); @memset(p.neuron(p.down.items, j, k, d), 0); std.log.info("slide pocket: layer {d}, neuron {d}, at the {s}: on for {d:.0}% of the fact's questions (least {d:.2}), held-out {d:.0}% (mean {d:.2}); other prompts at most {d:.2}, other rows {d:.2}, {d} above zero", .{ p.moe.items[j], k, @tagName(key), 100 * f.fact_on, f.fact_least, 100 * f.held_on, f.held_mean, f.steady_most, f.other_most, f.fires }); const start = p.pos.items.len / d; @@ -453,8 +483,7 @@ pub fn Pocket(comptime T: type) type { if (l.kind != .none or l.ex < p.answered.len + p.prompts.len + p.held.len) continue; var sum: f32 = 0; for (p.taken.items) |x| if (x.key == .subject) { - const z = dotBf16(p.neuron(p.up.items, x.slot, x.neuron, d), p.rows.items[(i * n + x.slot) * d ..][0..d]); - sum += if (z > 0) z * z else 0; + sum += @max(p.act(x.slot, x.neuron, p.rows.items[(i * n + x.slot) * d ..][0..d]), 0); }; if (sum > most) { most = sum; @@ -468,59 +497,23 @@ pub fn Pocket(comptime T: type) type { std.log.info("slide pocket: {d} other answers name the fact's subject where its neurons fire; they are to stay as they were", .{p.keeps.items.len}); } - /// Earlier neurons firing on this fact's rows refitted silent there, undo copies kept; same-subject ones stay. - fn quiet(p: *Self, t: *T, host: anytype) !void { - const d = t.config().hidden; - const n = p.slots(); - for (p.learned.items[0..p.learned_before]) |e| { - if (e.key == .subject and std.mem.indexOfScalar(u32, p.subjects.items, e.subject) != null) continue; - const u = p.neuron(p.up.items, e.slot, e.neuron, d); - var most: f32 = -std.math.inf(f32); - for (p.labels.items, 0..) |l, i| if (l.kind == .fact or l.kind == .held) { - most = @max(most, dotBf16(u, p.rows.items[(i * n + e.slot) * d ..][0..d])); - }; - if (most <= refit_above) continue; // weak there: left exactly as it is, this fact's outputs fitted with it - try p.refits.append(p.gpa, .{ .slot = e.slot, .neuron = e.neuron }); - try p.refit_up.appendSlice(p.gpa, u); - var xs: std.ArrayList([]const u16) = .empty; - defer xs.deinit(p.gpa); - var kinds: std.ArrayList(Kind) = .empty; - defer kinds.deinit(p.gpa); - for (0..e.count) |r| { - try xs.append(p.gpa, p.pos.items[(e.start + r) * d ..][0..d]); - try kinds.append(p.gpa, .fact); - } - for (p.labels.items, 0..) |l, i| { - try xs.append(p.gpa, p.rows.items[(i * n + e.slot) * d ..][0..d]); - try kinds.append(p.gpa, switch (l.kind) { - .fact, .held => .steady, // this fact's questions are another fact's to it - .none => .other, - else => l.kind, - }); - } - const f = try fit(p.gpa, xs.items, kinds.items, d, u, true, fit_steps / 2, e.count); - if (f.fact_least < kept_least or @max(f.steady_most, f.other_most) > 0) { - // a refit that forgets its fact or spreads is undone; this fact's outputs are fitted with it firing - @memcpy(u, p.refit_up.items[p.refit_up.items.len - d ..]); - p.refits.shrinkRetainingCapacity(p.refits.items.len - 1); - p.refit_up.shrinkRetainingCapacity(p.refit_up.items.len - d); - std.log.info("slide pocket: layer {d}, neuron {d}, an earlier fact's, fires on this fact's questions (up to {d:.2}); fitted again its own would read at least {d:.2} and other rows at most {d:.2}, so it stays as it was", .{ p.moe.items[e.slot], e.neuron, most, f.fact_least, @max(f.steady_most, f.other_most) }); - continue; - } - try p.upload(t, host, e.slot); - std.log.info("slide pocket: layer {d}, neuron {d}, an earlier fact's, fired on this fact's questions (up to {d:.2}); fitted again: there and elsewhere at most {d:.2}, its own questions' least {d:.2}", .{ p.moe.items[e.slot], e.neuron, most, @max(f.steady_most, f.other_most), f.fact_least }); - } - } - /// Whether no neuron is in use at any layer. pub fn empty(p: *const Self) bool { for (p.used.items) |u| if (u > 0) return false; return true; } - /// Learned neurons into the shards, then into the forward as the shards hold them (4-bit); the bytes written. + /// Learned neurons into the shards, then into the forward as the shards hold them (quantized); the bytes. pub fn save(p: *Self, t: *T, host: anytype) !usize { - const bytes = try host.pocketSave(p.moe.items, neurons, p.up.items, p.down.items); + const bytes = if (swiglu) blk: { + // written as gate rows and up rows u / sharp, handed back as written: u from the up rows again + const rows = try p.gpa.alloc(f32, p.up.items.len); + defer p.gpa.free(rows); + neuron_math.upRow(F.neuron, p.up.items, rows); + const got = try host.pocketSaveSwiglu(p.moe.items, neurons, p.gate.items, rows, p.down.items); + neuron_math.fromUp(F.neuron, rows, p.up.items); + break :blk got; + } else try host.pocketSave(p.moe.items, neurons, p.up.items, p.down.items); for (0..p.slots()) |j| if (p.used.items[j] > 0) try p.upload(t, host, j); return bytes; } @@ -559,7 +552,7 @@ pub fn Pocket(comptime T: type) type { } test "a fact keys as many neurons on its subject as the share of its questions that name it, at least half, else none" { - const P = Pocket(struct {}); + const P = Pocket(struct {}, @import("families.zig").nemotron_h); var p: P = .{ .gpa = std.testing.allocator, .mode = .auto }; defer p.deinit(); // ids: a question's tokens, the subject's last token at row 2, the answer from row 4 diff --git a/zig/src/families/nemotron/pocket_clock.zig b/zig/src/living/pocket_clock.zig similarity index 100% rename from zig/src/families/nemotron/pocket_clock.zig rename to zig/src/living/pocket_clock.zig diff --git a/zig/src/families/nemotron/pocket_dump.zig b/zig/src/living/pocket_dump.zig similarity index 100% rename from zig/src/families/nemotron/pocket_dump.zig rename to zig/src/living/pocket_dump.zig diff --git a/zig/src/families/nemotron/pocket_fit.zig b/zig/src/living/pocket_fit.zig similarity index 95% rename from zig/src/families/nemotron/pocket_fit.zig rename to zig/src/living/pocket_fit.zig index 02c9d99493..ab37028d73 100644 --- a/zig/src/families/nemotron/pocket_fit.zig +++ b/zig/src/living/pocket_fit.zig @@ -2,7 +2,7 @@ const std = @import("std"); const subspace = @import("subspace.zig"); -pub const margin: f32 = 0.25; // how far below zero a fact's neuron sits on rows it stays silent on +pub const margin: f32 = 0.25; // relu^2's margin (families.zig): how far below zero a neuron sits on rows it stays silent on pub const fit_steps = 500; // Adam steps fitting a neuron's input weights pub const fit_rate: f32 = 0.01; // their step, the inputs scaled to unit length @@ -21,7 +21,7 @@ pub const Fit = struct { }; /// Slot j's neuron fitted on its read rows, copied side by side first (a failure: zero, its slot last in line). -pub fn fitLayer(gpa: std.mem.Allocator, rows: []const u16, kinds: []const Kind, slots: usize, j: usize, d: usize, u: []f32, out: *Fit) void { +pub fn fitLayer(gpa: std.mem.Allocator, rows: []const u16, kinds: []const Kind, slots: usize, j: usize, d: usize, u: []f32, out: *Fit, below: f32) void { const own = gpa.alloc(u16, kinds.len * d) catch { out.* = .{ .fires = std.math.maxInt(usize) }; return; @@ -36,14 +36,14 @@ pub fn fitLayer(gpa: std.mem.Allocator, rows: []const u16, kinds: []const Kind, @memcpy(own[i * d ..][0..d], rows[(i * slots + j) * d ..][0..d]); x.* = own[i * d ..][0..d]; } - out.* = fit(gpa, xs, kinds, d, u, false, fit_steps, 0) catch blk: { + out.* = fit(gpa, xs, kinds, d, u, false, fit_steps, 0, below) catch blk: { @memset(u, 0); break :blk .{ .fires = std.math.maxInt(usize) }; }; } -/// Fits `u` on `xs` (from `u` if `warm`): at least 1 on fact rows, at most -margin elsewhere; first `fixed` rows kept. -pub fn fit(gpa: std.mem.Allocator, xs: []const []const u16, kinds: []const Kind, d: usize, u: []f32, warm: bool, steps: usize, fixed: usize) !Fit { +/// Fits `u` on `xs` (from `u` if `warm`): at least 1 on fact rows, at most -`below` elsewhere; first `fixed` rows kept. +pub fn fit(gpa: std.mem.Allocator, xs: []const []const u16, kinds: []const Kind, d: usize, u: []f32, warm: bool, steps: usize, fixed: usize, below: f32) !Fit { const n = kinds.len; var scale: f32 = 0; for (xs) |x| scale += @sqrt(squares(x)); @@ -93,10 +93,10 @@ pub fn fit(gpa: std.mem.Allocator, xs: []const []const u16, kinds: []const Kind, wrong[0] += 1; axpyRow(parts[0..d], -2 * (1 - z), x); }, - .steady, .other => if (z > -margin) { + .steady, .other => if (z > -below) { const at: usize = if (k == .steady) 1 else 2; wrong[at] += 1; - axpyRow(parts[at * d ..][0..d], 2 * (z + margin), x); + axpyRow(parts[at * d ..][0..d], 2 * (z + below), x); }, .held, .none => {}, } @@ -263,7 +263,7 @@ test "a neuron fires on its fact's rows and stays at zero on the rest; fitted ag k.* = if (i < n_fact) .fact else .other; } var u: [d]f32 = undefined; - const first = try fit(gpa, xs, kinds, d, &u, false, fit_steps, 0); + const first = try fit(gpa, xs, kinds, d, &u, false, fit_steps, 0, margin); try std.testing.expectEqual(@as(f32, 1), first.fact_on); try std.testing.expectEqual(@as(usize, 0), first.fires); // the new fact's rows, which nothing told it about, are where it is fitted again to stay quiet @@ -277,7 +277,7 @@ test "a neuron fires on its fact's rows and stays at zero on the rest; fitted ag } const before = try gpa.dupe(f32, &u); defer gpa.free(before); - const again = try fit(gpa, all, kinds2, d, &u, true, fit_steps / 2, n_fact); + const again = try fit(gpa, all, kinds2, d, &u, true, fit_steps / 2, n_fact, margin); try std.testing.expectEqual(@as(f32, 1), again.fact_on); try std.testing.expectEqual(@as(usize, 0), again.fires); try std.testing.expect(again.steady_most <= 0); diff --git a/zig/src/families/nemotron/pocket_joint.zig b/zig/src/living/pocket_joint.zig similarity index 100% rename from zig/src/families/nemotron/pocket_joint.zig rename to zig/src/living/pocket_joint.zig diff --git a/zig/src/living/pocket_load.zig b/zig/src/living/pocket_load.zig new file mode 100644 index 0000000000..b76ec8d019 --- /dev/null +++ b/zig/src/living/pocket_load.zig @@ -0,0 +1,30 @@ +//! The pocket's neurons as the shards hold them, read once, at the layers the family's layout learns at. +const std = @import("std"); +const neuron_math = @import("neuron.zig"); + +/// The pocket's neurons as the shards hold them, read once, at the layers the family's layout learns at. +pub fn load(p: anytype, host: anytype) !void { + const P = @TypeOf(p.*); + if (p.moe.items.len > 0) return; + const lay = host.layout(); + if (lay.layers > 64) return error.TooManyLayers; // a step's nudges mark their layers in a u64 + p.hidden = lay.hidden; + p.layers = lay.layers; + for (0..lay.layers) |l| if (lay.learn.isSet(l)) try p.moe.append(p.gpa, l); + const n = P.neurons * lay.hidden; + try p.up.appendNTimes(p.gpa, 0, p.moe.items.len * n); + try p.down.appendNTimes(p.gpa, 0, p.moe.items.len * n); + try p.used.appendNTimes(p.gpa, 0, p.moe.items.len); + if (P.swiglu) { + try p.gate.appendNTimes(p.gpa, 0, p.moe.items.len * n); + try p.scratch.appendNTimes(p.gpa, 0, n); + } + for (p.moe.items, 0..) |l, j| { + const up = p.up.items[j * n ..][0..n]; + const down = p.down.items[j * n ..][0..n]; + if (P.swiglu) { + p.used.items[j] = host.pocketReadSwiglu(l, p.gate.items[j * n ..][0..n], p.scratch.items, down); + neuron_math.fromUp(P.of.neuron, p.scratch.items, up); + } else p.used.items[j] = host.pocketRead(l, up, down); + } +} diff --git a/zig/src/families/nemotron/pocket_pack.zig b/zig/src/living/pocket_pack.zig similarity index 100% rename from zig/src/families/nemotron/pocket_pack.zig rename to zig/src/living/pocket_pack.zig diff --git a/zig/src/families/nemotron/pocket_prefix.zig b/zig/src/living/pocket_prefix.zig similarity index 99% rename from zig/src/families/nemotron/pocket_prefix.zig rename to zig/src/living/pocket_prefix.zig index ca854f38ee..a5239724a4 100644 --- a/zig/src/families/nemotron/pocket_prefix.zig +++ b/zig/src/living/pocket_prefix.zig @@ -88,7 +88,7 @@ pub fn learn(p: anytype, t: anytype, s: *Prefixes, sq: Seq, sp: Span, which: usi const k = &s.list.items[which]; if (k.seen or sq.skip > 0) return; k.seen = true; - const d = t.config().hidden; + const d = p.hidden; var limit = @min(sq.decide, sq.start - 1); for (sq.nudge.layers, sq.nudge.rows) |mask, row| if (mask != 0) { limit = @min(limit, row); diff --git a/zig/src/living/pocket_quiet.zig b/zig/src/living/pocket_quiet.zig new file mode 100644 index 0000000000..a1e2d70364 --- /dev/null +++ b/zig/src/living/pocket_quiet.zig @@ -0,0 +1,59 @@ +//! Earlier facts' neurons that fire on the open fact's questions, fitted again to stay silent there. +const std = @import("std"); +const pocket_fit = @import("pocket_fit.zig"); + +const Kind = pocket_fit.Kind; +const fit = pocket_fit.fit; +const dotBf16 = pocket_fit.dotBf16; +const fit_steps = pocket_fit.fit_steps; + +pub const kept_least: f32 = 0.5; // an earlier neuron refitted must still reach this on every one of its own fact's rows +pub const refit_above: f32 = 0.5; // an earlier neuron is refitted only where it fires above this on a new fact's rows + +/// Earlier neurons firing on this fact's rows refitted silent there, undo copies kept; same-subject ones stay. +pub fn quiet(p: anytype, t: anytype, host: anytype) !void { + const P = @TypeOf(p.*); + const d = p.hidden; + const n = p.moe.items.len; + for (p.learned.items[0..p.learned_before]) |e| { + if (e.key == .subject and std.mem.indexOfScalar(u32, p.subjects.items, e.subject) != null) continue; + const u = p.neuron(p.up.items, e.slot, e.neuron, d); + var most: f32 = -std.math.inf(f32); + for (p.labels.items, 0..) |l, i| if (l.kind == .fact or l.kind == .held) { + most = @max(most, dotBf16(u, p.rows.items[(i * n + e.slot) * d ..][0..d])); + }; + if (most <= refit_above) continue; // weak there: left exactly as it is, this fact's outputs fitted with it + try p.refits.append(p.gpa, .{ .slot = e.slot, .neuron = e.neuron }); + try p.refit_up.appendSlice(p.gpa, u); + if (P.swiglu) try p.refit_gate.appendSlice(p.gpa, p.neuron(p.gate.items, e.slot, e.neuron, d)); + var xs: std.ArrayList([]const u16) = .empty; + defer xs.deinit(p.gpa); + var kinds: std.ArrayList(Kind) = .empty; + defer kinds.deinit(p.gpa); + for (0..e.count) |r| { + try xs.append(p.gpa, p.pos.items[(e.start + r) * d ..][0..d]); + try kinds.append(p.gpa, .fact); + } + for (p.labels.items, 0..) |l, i| { + try xs.append(p.gpa, p.rows.items[(i * n + e.slot) * d ..][0..d]); + try kinds.append(p.gpa, switch (l.kind) { + .fact, .held => .steady, // this fact's questions are another fact's to it + .none => .other, + else => l.kind, + }); + } + const f = try fit(p.gpa, xs.items, kinds.items, d, u, true, fit_steps / 2, e.count, P.of.neuron.margin); + if (f.fact_least < kept_least or @max(f.steady_most, f.other_most) > 0) { + // a refit that forgets its fact or spreads is undone; this fact's outputs are fitted with it firing + @memcpy(u, p.refit_up.items[p.refit_up.items.len - d ..]); + p.refits.shrinkRetainingCapacity(p.refits.items.len - 1); + p.refit_up.shrinkRetainingCapacity(p.refit_up.items.len - d); + if (P.swiglu) p.refit_gate.shrinkRetainingCapacity(p.refit_gate.items.len - d); + std.log.info("slide pocket: layer {d}, neuron {d}, an earlier fact's, fires on this fact's questions (up to {d:.2}); fitted again its own would read at least {d:.2} and other rows at most {d:.2}, so it stays as it was", .{ p.moe.items[e.slot], e.neuron, most, f.fact_least, @max(f.steady_most, f.other_most) }); + continue; + } + p.gateFrom(e.slot, e.neuron); + try p.upload(t, host, e.slot); + std.log.info("slide pocket: layer {d}, neuron {d}, an earlier fact's, fired on this fact's questions (up to {d:.2}); fitted again: there and elsewhere at most {d:.2}, its own questions' least {d:.2}", .{ p.moe.items[e.slot], e.neuron, most, @max(f.steady_most, f.other_most), f.fact_least }); + } +} diff --git a/zig/src/families/nemotron/pocket_read.zig b/zig/src/living/pocket_read.zig similarity index 98% rename from zig/src/families/nemotron/pocket_read.zig rename to zig/src/living/pocket_read.zig index 77d9bcd9dc..246652d6e3 100644 --- a/zig/src/families/nemotron/pocket_read.zig +++ b/zig/src/living/pocket_read.zig @@ -37,6 +37,7 @@ fn state(p: anytype) u64 { h.update(std.mem.sliceAsBytes(p.up.items)); h.update(std.mem.sliceAsBytes(p.down.items)); h.update(std.mem.sliceAsBytes(p.used.items)); + if (@TypeOf(p.*).swiglu) h.update(std.mem.sliceAsBytes(p.gate.items)); return h.final(); } diff --git a/zig/src/families/nemotron/pocket_values.zig b/zig/src/living/pocket_values.zig similarity index 97% rename from zig/src/families/nemotron/pocket_values.zig rename to zig/src/living/pocket_values.zig index 628ce8f265..46b1a96acb 100644 --- a/zig/src/families/nemotron/pocket_values.zig +++ b/zig/src/living/pocket_values.zig @@ -131,8 +131,7 @@ const Drawn = struct { own: bool, joint: bool = false, which: usize, seq: Seq }; /// One output step: its fact answers, then those that are to stay, side by side; then Adam, and every few the judge. pub fn valueOnce(p: anytype, t: anytype, host: anytype) !bool { const s = &p.v; - const c = t.config(); - const d = c.hidden; + const d = p.hidden; const k = p.taken.items.len + p.joint.items.len; // the open fact's neurons, then the earlier ones fitted with them const vals = s.adam.items[0 .. k * d]; const grad = s.adam.items[3 * k * d ..][0 .. k * d]; @@ -253,8 +252,7 @@ pub fn valueOnce(p: anytype, t: anytype, host: anytype) !bool { /// Answer `dr`, the pass's i-th, into the step: its loss, and each neuron's push times its activation at its row. fn absorb(p: anytype, t: anytype, dr: Drawn, i: usize, got: Result, grad: []f32) void { const s = &p.v; - const c = t.config(); - const d = c.hidden; + const d = p.hidden; const masks = dr.seq.nudge.layers; const rows = dr.seq.nudge.rows; const sp = t.spans[i]; @@ -276,8 +274,10 @@ fn absorb(p: anytype, t: anytype, dr: Drawn, i: usize, got: Result, grad: []f32) const key = @backingInt(x.key); if (masks[key] == 0) continue; const l = p.moe.items[x.slot]; - const z = dotBf16(p.neuron(p.up.items, x.slot, x.neuron, d), t.normed(l, sp.row(rows[key]))); - const a = if (z > 0) z * z else 0; + const row = sp.row(rows[key]); + // relu(u x)^2, or a SwiGLU's silu(g x) (u x) / sharp; times its expert's gate there when one scales it + var a = p.act(x.slot, x.neuron, t.normed(l, row)); + if (@TypeOf(p.*).gated) a *= t.gateAt(l, row); subspace.axpy(grad[j * d ..][0..d], dr.seq.scale * a, t.pushAt(i, key, l)); if (dr.own and s.size.items[j] == 0) s.size.items[j] = normBf16(t.residual(l + 1, sp.row(rows[key]))); } @@ -358,7 +358,7 @@ fn judgeOnce(p: anytype, t: anytype, host: anytype, vals: []f32) !bool { /// The outputs into the pocket's neurons and the forward; the held-out answers judged next. fn finish(p: anytype, t: anytype, host: anytype, vals: []const f32) !bool { const s = &p.v; - const d = t.config().hidden; + const d = p.hidden; try put(p, t, host, vals, d); t.from = @splat(std.math.maxInt(u32)); std.log.info("slide pocket: {d} answers' first rows run once, {d} rows the steps did not run again", .{ s.prefixes.slots, s.prefixes.skipped }); @@ -395,9 +395,9 @@ pub fn checkOnce(p: anytype, t: anytype, host: anytype) !bool { return true; } -/// Logs the forward's change after the fact's lowest neuron on `ex` against its relu(u x)^2 v, to spot a bad forward. +/// Logs the forward's change after the fact's lowest neuron on `ex` against its act(x) v, to spot a bad forward. fn verifyForward(p: anytype, t: anytype, host: anytype, ex: Example) !void { - const d = t.config().hidden; + const d = p.hidden; var low: ?std.meta.Elem(@TypeOf(p.taken.items)) = null; var row: usize = 0; for (p.taken.items) |x| { @@ -420,8 +420,8 @@ fn verifyForward(p: anytype, t: anytype, host: anytype, ex: Example) !void { t.nudge = .{}; _ = try t.step(ex.ids, ex.start, .target, turn); for (with, t.residual(l + 1, row)) |*o, b| o.* = @bitCast(@as(u32, b) << 16); - const z = dotBf16(p.neuron(p.up.items, k.slot, k.neuron, d), t.normed(l, row)); - const a = if (z > 0) z * z else 0; + var a = p.act(k.slot, k.neuron, t.normed(l, row)); + if (@TypeOf(p.*).gated) a *= t.gateAt(l, row); const v = p.neuron(p.down.items, k.slot, k.neuron, d); for (want, v) |*o, x| o.* = a * x; @memcpy(saved, v); diff --git a/zig/src/families/nemotron/probe.zig b/zig/src/living/probe.zig similarity index 98% rename from zig/src/families/nemotron/probe.zig rename to zig/src/living/probe.zig index 1175275c7a..1d2089f2ec 100644 --- a/zig/src/families/nemotron/probe.zig +++ b/zig/src/living/probe.zig @@ -1,4 +1,4 @@ -//! A measure, not a lesson: how far vectors fitted at one row of a few layers bring a fact's held-out answers back. +//! A measure, not a lesson: how far vectors fitted at a few layers bring a fact back (Nemotron's Metal trainer only). const std = @import("std"); const examples = @import("lesson_examples.zig"); const lesson_route = @import("lesson_route.zig"); diff --git a/zig/src/families/nemotron/router.zig b/zig/src/living/router.zig similarity index 100% rename from zig/src/families/nemotron/router.zig rename to zig/src/living/router.zig diff --git a/zig/src/living/shards.zig b/zig/src/living/shards.zig new file mode 100644 index 0000000000..f0899b8de6 --- /dev/null +++ b/zig/src/living/shards.zig @@ -0,0 +1,211 @@ +//! The pocket in the shards, for any family: every shared expert widened once (config.json told), then in place. +const std = @import("std"); +const family = @import("family.zig"); +const affine = @import("affine.zig"); +const Family = family.Family; +const Layout = family.Layout; + +/// The writer over `core`, the backend's own (its shard_edit and safetensors). +pub fn Io(comptime core: type) type { + const st = core.safetensors; + const se = core.shard_edit; + return struct { + /// The learned neurons of `layers` into `dir`'s shards, quantized as `F` stores them, handed back; the bytes. + pub fn write(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, comptime F: Family, lay: Layout, layers: []const usize, neurons: usize, inputs: []const []f32, down: []f32) !usize { + const q = F.quant; + const d = lay.hidden; + const own = lay.width; + if (inputs.len != F.names.inputs.len) return error.BadPocket; + if (neurons % q.group != 0 or d % q.group != 0 or own % q.group != 0) return error.BadPocket; + const has = try se.configIntIn(gpa, io, dir, F.keys.scope, F.keys.pocket); + const widen = has == null; + if (has) |n| if (n != neurons) return error.PocketOfAnotherSize; + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + const arena = arena_state.allocator(); + var files: Files = .{ .gpa = gpa, .io = io, .dir = dir, .paths = .init(gpa) }; + defer files.close(); + var edits: std.ArrayList(se.Replacement) = .empty; + var bytes: usize = 0; + const n = neurons * d; + var rows: [F.names.inputs.len][]f32 = undefined; + for (0..lay.layers) |l| { + if (!lay.widen.isSet(l)) continue; + const at = std.mem.indexOfScalar(usize, layers, l); + for (&rows, inputs) |*r, x| r.* = if (at) |j| x[j * n ..][0..n] else try zeros(arena, n); + const v = if (at) |j| down[j * n ..][0..n] else try zeros(arena, n); + var buf: [160]u8 = undefined; + bytes += try expert(F, arena, &files, &edits, try F.expertBase(&buf, l), own, neurons, d, &rows, v); + } + const mtp = if (widen) F.names.mtp else null; + const z = try zeros(arena, n); + var zs: [F.names.inputs.len][]f32 = @splat(z); + if (mtp) |m| if (m.file == null and try files.indexed(try std.fmt.allocPrint(arena, "{s}{s}.weight", .{ m.base, F.names.inputs[0] }))) { + // the MTP head's expert in the index: grown with the model's own, zero neurons + bytes += try expert(F, arena, &files, &edits, m.base, own, neurons, d, &zs, z); + }; + if (widen) try se.grow(gpa, io, dir, edits.items) else try se.overwrite(gpa, io, dir, edits.items); + if (mtp) |m| if (m.file) |name| { + // in its own file, zero neurons: that file rewritten whole, apart from the index + const path = try std.fs.path.join(arena, &.{ dir, name }); + if (hasMtp(F, gpa, io, path, m.base)) { + var only: std.ArrayList(se.Replacement) = .empty; + var single: Files = .{ .gpa = gpa, .io = io, .dir = dir, .single = path, .paths = .init(gpa) }; + defer single.close(); + bytes += try expert(F, arena, &single, &only, m.base, own, neurons, d, &zs, z); + _ = try se.rewrite(gpa, io, path, only.items); + } + }; + if (widen) try se.setConfigIn(gpa, io, dir, F.keys.scope, &.{ .{ F.keys.width, @intCast(own + neurons) }, .{ F.keys.pocket, @intCast(neurons) } }); + return bytes; + } + + fn hasMtp(comptime F: Family, gpa: std.mem.Allocator, io: std.Io, path: []const u8, base: []const u8) bool { + var f = st.File.open(gpa, io, path) catch return false; + defer f.close(io); + var buf: [192]u8 = undefined; + const name = std.fmt.bufPrint(&buf, "{s}{s}.weight", .{ base, F.names.inputs[0] }) catch return false; + return f.get(name) != null; + } + + /// The shard files a model's tensors are read from: the index's, or one file (an MTP head's own). + const Files = struct { + gpa: std.mem.Allocator, + io: std.Io, + dir: []const u8, + single: ?[]const u8 = null, + open: std.StringArrayHashMapUnmanaged(st.File) = .empty, + paths: std.heap.ArenaAllocator, + map: ?std.json.Parsed(std.json.Value) = null, + + fn get(fs: *Files, name: []const u8) !st.Tensor { + const shard = if (fs.single) |x| x else try fs.shardOf(name); + const slot = try fs.open.getOrPut(fs.gpa, shard); + if (!slot.found_existing) slot.value_ptr.* = st.File.open(fs.gpa, fs.io, shard) catch |e| { + _ = fs.open.pop(); + return e; + }; + return slot.value_ptr.get(name) orelse error.MissingTensor; + } + + /// Whether the folder's index (or its one shard) holds `name`. + fn indexed(fs: *Files, name: []const u8) !bool { + _ = fs.get(name) catch |e| return if (e == error.MissingTensor) false else e; + return true; + } + + /// The path of the shard holding `name`, kept until the files close. + fn shardOf(fs: *Files, name: []const u8) ![]const u8 { + const a = fs.paths.allocator(); + if (fs.map == null) { + const path = try std.fs.path.join(a, &.{ fs.dir, "model.safetensors.index.json" }); + const text = std.Io.Dir.cwd().readFileAlloc(fs.io, path, a, .limited(1 << 26)) catch |e| if (e == error.FileNotFound) { + fs.single = try std.fs.path.join(a, &.{ fs.dir, "model.safetensors" }); + return fs.single.?; + } else return e; + fs.map = try std.json.parseFromSlice(std.json.Value, fs.gpa, text, .{ .allocate = .alloc_always }); + } + const root = fs.map.?.value; + if (root != .object) return error.BadIndex; + const weights = root.object.get("weight_map") orelse return error.BadIndex; + if (weights != .object) return error.BadIndex; + const file = weights.object.get(name) orelse return error.MissingTensor; + if (file != .string) return error.BadIndex; + return std.fs.path.join(a, &.{ fs.dir, file.string }); + } + + fn close(fs: *Files) void { + for (fs.open.values()) |*f| f.close(fs.io); + fs.open.deinit(fs.gpa); + if (fs.map) |*x| x.deinit(); + fs.paths.deinit(); + } + }; + + /// One shared expert with `neurons` learned after its `own`: a row a neuron per input, a column out; the bytes. + fn expert(comptime F: Family, arena: std.mem.Allocator, files: *Files, edits: *std.ArrayList(se.Replacement), base: []const u8, own: usize, neurons: usize, d: usize, inputs: []const []f32, v: []f32) !usize { + const q = F.quant; + const g = q.group; + const wide = own + neurons; + const per_word = 32 / q.bits; + var name: [192]u8 = undefined; + var bytes: usize = 0; + var had: ?usize = null; // the input projections' rows: `own`, or `wide` once widened + for (F.names.inputs, inputs) |proj, x| { + var parts: [3]st.Tensor = undefined; + for (&parts, [_][]const u8{ "weight", "scales", "biases" }) |*t, leaf| t.* = try files.get(try std.fmt.bufPrint(&name, "{s}{s}.{s}", .{ base, proj, leaf })); + const rows_had = parts[0].dim(0); + if (rows_had != own and rows_had != wide) return error.UnexpectedTensor; + if (had) |h| if (h != rows_had) return error.UnexpectedTensor; + had = rows_had; + const w = try arena.alloc(u32, wide * d / per_word); + const s = try arena.alloc(u16, wide * d / g); + const b = try arena.alloc(u16, wide * d / g); + @memcpy(w[0 .. own * d / per_word], words(parts[0].bytes)[0 .. own * d / per_word]); + @memcpy(s[0 .. own * d / g], halves(parts[1].bytes)[0 .. own * d / g]); + @memcpy(b[0 .. own * d / g], halves(parts[2].bytes)[0 .. own * d / g]); + for (0..neurons) |k| { + const r = own + k; + const row = x[k * d ..][0..d]; + const rw = w[r * d / per_word ..][0 .. d / per_word]; + const rs = s[r * d / g ..][0 .. d / g]; + const rb = b[r * d / g ..][0 .. d / g]; + affine.quantize(q, row, rw, rs, rb); + affine.dequantize(q, rw, rs, rb, row); + } + bytes += try add(arena, edits, base, proj, .{ wide, d }, w, s, b, q); + } + // the output projection: each output row's own columns as they are, then a column a learned neuron + var parts: [3]st.Tensor = undefined; + for (&parts, [_][]const u8{ "weight", "scales", "biases" }) |*t, leaf| t.* = try files.get(try std.fmt.bufPrint(&name, "{s}{s}.{s}", .{ base, F.names.output, leaf })); + const cols = had.?; // as wide as the input projections had rows + const dw = try arena.alloc(u32, d * wide / per_word); + const ds = try arena.alloc(u16, d * wide / g); + const db = try arena.alloc(u16, d * wide / g); + const col = try arena.alloc(f32, neurons); + for (0..d) |o| { + @memcpy(dw[o * wide / per_word ..][0 .. own / per_word], words(parts[0].bytes)[o * cols / per_word ..][0 .. own / per_word]); + @memcpy(ds[o * wide / g ..][0 .. own / g], halves(parts[1].bytes)[o * cols / g ..][0 .. own / g]); + @memcpy(db[o * wide / g ..][0 .. own / g], halves(parts[2].bytes)[o * cols / g ..][0 .. own / g]); + for (col, 0..) |*x, k| x.* = v[k * d + o]; + const cw = dw[o * wide / per_word ..][own / per_word ..][0 .. neurons / per_word]; + const cs = ds[o * wide / g ..][own / g ..][0 .. neurons / g]; + const cb = db[o * wide / g ..][own / g ..][0 .. neurons / g]; + affine.quantize(q, col, cw, cs, cb); + affine.dequantize(q, cw, cs, cb, col); + for (col, 0..) |x, k| v[k * d + o] = x; + } + bytes += try add(arena, edits, base, F.names.output, .{ d, wide }, dw, ds, db, q); + return bytes; + } + + /// A projection's codes, scales and biases ([rows, cols] values) as three replacements; their bytes. + fn add(arena: std.mem.Allocator, edits: *std.ArrayList(se.Replacement), base: []const u8, proj: []const u8, shape: [2]usize, w: []const u32, s: []const u16, b: []const u16, comptime q: family.Quant) !usize { + const made = [_]struct { leaf: []const u8, dtype: st.DType, shape: [2]usize, bytes: []const u8 }{ + .{ .leaf = "weight", .dtype = .u32, .shape = .{ shape[0], shape[1] * q.bits / 32 }, .bytes = std.mem.sliceAsBytes(w) }, + .{ .leaf = "scales", .dtype = .bf16, .shape = .{ shape[0], shape[1] / q.group }, .bytes = std.mem.sliceAsBytes(s) }, + .{ .leaf = "biases", .dtype = .bf16, .shape = .{ shape[0], shape[1] / q.group }, .bytes = std.mem.sliceAsBytes(b) }, + }; + var bytes: usize = 0; + for (made) |m| { + try edits.append(arena, .{ .name = try std.fmt.allocPrint(arena, "{s}{s}.{s}", .{ base, proj, m.leaf }), .dtype = m.dtype, .shape = try arena.dupe(usize, &m.shape), .bytes = m.bytes }); + bytes += m.bytes.len; + } + return bytes; + } + }; +} + +fn zeros(arena: std.mem.Allocator, n: usize) ![]f32 { + const z = try arena.alloc(f32, n); + @memset(z, 0); + return z; +} + +fn words(bytes: []const u8) []align(1) const u32 { + return std.mem.bytesAsSlice(u32, bytes); +} + +fn halves(bytes: []const u8) []align(1) const u16 { + return std.mem.bytesAsSlice(u16, bytes); +} diff --git a/zig/src/families/nemotron/subspace.zig b/zig/src/living/subspace.zig similarity index 100% rename from zig/src/families/nemotron/subspace.zig rename to zig/src/living/subspace.zig diff --git a/zig/src/living/swiglu_test.zig b/zig/src/living/swiglu_test.zig new file mode 100644 index 0000000000..994a52f8c3 --- /dev/null +++ b/zig/src/living/swiglu_test.zig @@ -0,0 +1,82 @@ +//! A SwiGLU pocket neuron on the host: fitted, written in 4-bit groups of 32 and read back, it fires as relu^2 would. +const std = @import("std"); +const pocket_fit = @import("pocket_fit.zig"); +const neuron = @import("neuron.zig"); +const affine = @import("affine.zig"); +const families = @import("families.zig"); + +const f = @import("test_families.zig").swiglu; +const d = 128; + +fn bf(x: f32) u16 { + return @truncate(@as(u32, @bitCast(x)) >> 16); +} + +fn dot(a: []const f32, b: []const u16) f32 { + return pocket_fit.dotBf16(a, b); +} + +/// `x` written and read back as the family stores its experts. +fn stored(x: []f32) void { + var w: [affine.wordsOf(f.quant, d)]u32 = undefined; + var s: [d / f.quant.group]u16 = undefined; + var b: [d / f.quant.group]u16 = undefined; + affine.quantize(f.quant, x, &w, &s, &b); + affine.dequantize(f.quant, &w, &s, &b, x); +} + +test "a SwiGLU neuron fitted at margin 0.5 and stored at 4 bits in groups of 32 fires on its rows and stays near zero elsewhere" { + const gpa = std.testing.allocator; + var prng = std.Random.DefaultPrng.init(21); + const r = prng.random(); + var common: [d]f32 = undefined; + var lean: [d]f32 = undefined; + for (&common, &lean) |*c, *l| { + c.* = 1 + r.floatNorm(f32) * 0.1; + l.* = r.floatNorm(f32); + } + const n_fact = 16; + const n_rest = 300; + const rows = try gpa.alloc(u16, (n_fact + n_rest) * d); + defer gpa.free(rows); + const xs = try gpa.alloc([]const u16, n_fact + n_rest); + defer gpa.free(xs); + const kinds = try gpa.alloc(pocket_fit.Kind, n_fact + n_rest); + defer gpa.free(kinds); + for (0..n_fact + n_rest) |i| { + for (0..d) |k| { + var v = common[k] + r.floatNorm(f32) * 0.5; + if (i < n_fact) v += 1.5 * lean[k]; + rows[i * d + k] = bf(v); + } + xs[i] = rows[i * d ..][0..d]; + kinds[i] = if (i < n_fact) .fact else .other; + } + var u: [d]f32 = undefined; + const got = try pocket_fit.fit(gpa, xs, kinds, d, &u, false, pocket_fit.fit_steps, 0, f.neuron.margin); + try std.testing.expectEqual(@as(f32, 1), got.fact_on); + try std.testing.expectEqual(@as(usize, 0), got.fires); + try std.testing.expect(@max(got.steady_most, got.other_most) <= -f.neuron.margin + 0.05); + // as the forward and the shards hold it: gate row sharp u, up row u / sharp, each quantized + var g: [d]f32 = undefined; + var w: [d]f32 = undefined; + neuron.gateRow(f.neuron, &u, &g); + neuron.upRow(f.neuron, &u, &w); + stored(&g); + stored(&w); + var least: f32 = std.math.inf(f32); + var most: f32 = 0; + for (xs, kinds) |x, k| { + const a = neuron.silu(dot(&g, x)) * dot(&w, x); + if (k == .fact) { + least = @min(least, a); + // within a hair of relu(u x)^2 on its own rows, but for the rows' quantization + try std.testing.expectApproxEqRel(neuron.act(.{ .act = .relu2 }, dot(&u, x), 0), a, 0.15); + } else most = @max(most, @abs(a)); + } + try std.testing.expect(least >= 0.7); + try std.testing.expect(most <= 1e-3 * least); + // the same rows through act(): u x and g x as the pocket reads them + const x = xs[0]; + try std.testing.expectApproxEqRel(neuron.silu(dot(&g, x)) * dot(&w, x), neuron.act(f.neuron, f.neuron.sharp * dot(&w, x), dot(&g, x)), 1e-5); +} diff --git a/zig/src/living/test_families.zig b/zig/src/living/test_families.zig new file mode 100644 index 0000000000..89a9255cf1 --- /dev/null +++ b/zig/src/living/test_families.zig @@ -0,0 +1,18 @@ +//! A family for tests: a SwiGLU shared expert shaped as Qwen3.8 Flash Next's, which no checkpoint learns into yet. +const Family = @import("family.zig").Family; + +/// A SwiGLU shared expert at every layer, gated by its own router logit, in 4-bit groups of 32, keys under text_config. +pub const swiglu: Family = .{ + .model_type = "swiglu_test", + .names = .{ + .layers = "language_model.model.layers.", + .expert = "mlp.shared_expert.", + .inputs = &.{ "gate_proj", "up_proj" }, + .mtp = .{ .base = "language_model.mtp.layers.0.mlp.shared_expert." }, + }, + // silu(gate) is not zero below it as relu is: a wider margin keeps sigmoid(-16 * 0.5) = 3e-4 of it there + .neuron = .{ .act = .swiglu, .gated = true, .margin = 0.5, .sharp = 16 }, + .quant = .{ .bits = 4, .group = 32 }, + .keys = .{ .scope = "text_config", .width = "shared_expert_intermediate_size" }, + .experts = .every, +}; diff --git a/zig/src/families/nemotron/train_mode.zig b/zig/src/living/train_mode.zig similarity index 100% rename from zig/src/families/nemotron/train_mode.zig rename to zig/src/living/train_mode.zig diff --git a/zig/src/families/nemotron/whole.zig b/zig/src/living/whole.zig similarity index 91% rename from zig/src/families/nemotron/whole.zig rename to zig/src/living/whole.zig index 4fa4a9fc01..f6a5b34ad0 100644 --- a/zig/src/families/nemotron/whole.zig +++ b/zig/src/living/whole.zig @@ -46,6 +46,24 @@ pub fn Hooks(comptime H: type) type { pub fn pocketSave(f: Self, layers: []const usize, neurons: usize, up: []f32, down: []f32) !usize { return H.pocketSave(f.b, f.gpa, f.io, layers, neurons, up, down); } + + /// The served checkpoint's layout: where pockets sit and learn, its width (family.zig). + pub fn layout(f: Self) @import("family.zig").Layout { + return H.layout(f.b); + } + + /// A SwiGLU family's layer `layer` for the forward: gate rows, up rows (u / sharp) and outputs, [neurons, D]. + pub fn pocketSwiglu(f: Self, layer: usize, gate: []const f32, up: []const f32, down: []const f32, used: usize) !void { + return H.pocketSwiglu(f.b, layer, gate, up, down, used); + } + + pub fn pocketReadSwiglu(f: Self, layer: usize, gate: []f32, up: []f32, down: []f32) usize { + return H.pocketReadSwiglu(f.b, layer, gate, up, down); + } + + pub fn pocketSaveSwiglu(f: Self, layers: []const usize, neurons: usize, gate: []f32, up: []f32, down: []f32) !usize { + return H.pocketSaveSwiglu(f.b, f.gpa, f.io, layers, neurons, gate, up, down); + } }; } diff --git a/zig/src/living_shards_test.zig b/zig/src/living_shards_test.zig new file mode 100644 index 0000000000..555e672295 --- /dev/null +++ b/zig/src/living_shards_test.zig @@ -0,0 +1,32 @@ +//! Living Weights' pocket in each family's own shards, rooted here so the tests reach the core's shard files by path. +const std = @import("std"); +const living = @import("living"); +const host4 = @import("core/affine4_host.zig"); + +test { + _ = @import("families/nemotron/pocket_shards_test.zig"); +} + +test "living's affine writer at 4 bits in groups of 64 writes the very codes, scales and biases affine4_host does" { + const q: living.family.Quant = .{ .bits = 4, .group = 64 }; + var prng = std.Random.DefaultPrng.init(3); + const r = prng.random(); + var values: [8 * 64]f32 = undefined; + for (&values, 0..) |*v, i| v.* = if (i / 64 == 5) 0.125 else r.floatNorm(f32) * 0.07; + var w1: [values.len / 8]u32 = undefined; + var w2: [values.len / 8]u32 = undefined; + var s1: [8]u16 = undefined; + var s2: [8]u16 = undefined; + var b1: [8]u16 = undefined; + var b2: [8]u16 = undefined; + host4.quantize(&values, &w1, &s1, &b1); + living.affine.quantize(q, &values, &w2, &s2, &b2); + try std.testing.expectEqualSlices(u32, &w1, &w2); + try std.testing.expectEqualSlices(u16, &s1, &s2); + try std.testing.expectEqualSlices(u16, &b1, &b2); + var o1: [values.len]f32 = undefined; + var o2: [values.len]f32 = undefined; + host4.dequantize(&w1, &s1, &b1, &o1); + living.affine.dequantize(q, &w2, &s2, &b2, &o2); + try std.testing.expectEqualSlices(f32, &o1, &o2); +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 851f919480..4ea912da72 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -7,7 +7,7 @@ const nemotron = @import("nemotron"); const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); -const nemotron_slide = @import("nemotron_slide.zig"); +const slide_adapter = @import("slide_adapter.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. const registry = .{nemotron.native}; @@ -269,11 +269,13 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O return null; }; if (o.slide) { + // a family learns once its module declares `slide` (its learner), `learnModel` and `linked` if (!@hasDecl(F, "slide")) { problem.* = try std.fmt.allocPrint(a, "--slide: {s} checkpoints do not learn on CUDA yet", .{o.model_type}); return null; } - if (try nemotron_slide.refusal(F, a, io, o.dir, "`cp -RL` (a copy that follows links)")) |why| { + comptime if (!@hasDecl(F, "learnModel") or !@hasDecl(F, "linked")) @compileError(F.model_type ++ " declares slide without learnModel and linked"); + if (try slide_adapter.refusal(F, a, io, o.dir, "`cp -RL` (a copy that follows links)")) |why| { problem.* = why; return null; } @@ -365,7 +367,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O gpa.destroy(store); }; if (loaded.learns != null) { - const Slide = nemotron_slide.Adapter(F.slide); + const Slide = slide_adapter.Adapter(F.slide); h.dir = try gpa.dupe(u8, o.dir); const s = try gpa.create(Slide); s.* = Slide.init(gpa, io, F.learnModel(loaded.ctx, h.dir)); diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index 39fe531408..f52361af10 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -10,8 +10,8 @@ const qwen27 = @import("qwen27_served.zig"); const flashnext = @import("flashnext_host.zig"); const glm = @import("glm_host.zig"); const glm_dsa = @import("glm_dsa_host.zig"); -const nemotron_slide = @import("nemotron_slide.zig"); -const Slide = nemotron_slide.Adapter(nemotron.slide); +const slide_adapter = @import("slide_adapter.zig"); +const Slide = slide_adapter.Adapter(nemotron.slide); const cache_fit = @import("cache_fit.zig"); pub const backends: []const []const u8 = &.{"metal"}; @@ -119,7 +119,7 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c problem.* = try std.fmt.allocPrint(a, "the native engine has no backend for {s} checkpoints yet; the Python engine 0.6.6 may serve them: python -m pip install git+https://github.com/ashhart/TensorFold.git@v0.6.6", .{o.model_type}); return null; } - if (o.slide) if (try nemotron_slide.refusal(tf.shard_edit, a, io, o.dir, "`cp -cRL` (an APFS clone that follows links)")) |why| { + if (o.slide) if (try slide_adapter.refusal(tf.shard_edit, a, io, o.dir, "`cp -cRL` (an APFS clone that follows links)")) |why| { problem.* = why; return null; }; diff --git a/zig/src/native/nemotron_slide.zig b/zig/src/native/slide_adapter.zig similarity index 90% rename from zig/src/native/nemotron_slide.zig rename to zig/src/native/slide_adapter.zig index 157c89e425..39ab12037d 100644 --- a/zig/src/native/nemotron_slide.zig +++ b/zig/src/native/slide_adapter.zig @@ -1,4 +1,4 @@ -//! Nemotron's Sliding Weights learner behind the lane host's learn hook: lessons in, learn events out. +//! A family's Living Weights learner behind the lane host's learn hook: lessons in, learn events out. const std = @import("std"); const api = @import("engine_api"); const Allocator = std.mem.Allocator; @@ -6,9 +6,9 @@ const Allocator = std.mem.Allocator; /// Neurons a learn event carries at most; a fact takes about ten. const max_neurons = 64; -/// Why --slide cannot learn into `dir` (`Checks`: its plainShared and linked; `copy`: the copy to make), or null. +/// Why --slide cannot learn into `dir` (`Checks`: linked and plainShared files; `copy`: the copy to make), or null. pub fn refusal(comptime Checks: type, a: Allocator, io: std.Io, dir: []const u8, copy: []const u8) !?[]const u8 { - if (try Checks.plainShared(a, io, dir)) |name| return try std.fmt.allocPrint(a, "--slide adds learned neurons beside every shared expert, and {s} in {s} is unquantized, as TensorFold 1.0.3 to 1.0.6 saved their lessons: serve this folder without --slide (what it learned stays), or teach a fresh copy", .{ name, dir }); + if (@hasDecl(Checks, "plainShared")) if (try Checks.plainShared(a, io, dir)) |name| return try std.fmt.allocPrint(a, "--slide adds learned neurons beside every shared expert, and {s} in {s} is unquantized, as TensorFold 1.0.3 to 1.0.6 saved their lessons: serve this folder without --slide (what it learned stays), or teach a fresh copy", .{ name, dir }); if (try Checks.linked(a, io, dir)) |name| return try std.fmt.allocPrint(a, "--slide rewrites the model's own files, and {s} in {s} is a link to data another file shares (a Hugging Face cache's blob, or a hard link), so learning would change that file too: serve a copy made with {s}", .{ name, dir, copy }); return null; } diff --git a/zig/src/server/slide_model.zig b/zig/src/server/slide_model.zig index 274f9bb59b..db94df3e69 100644 --- a/zig/src/server/slide_model.zig +++ b/zig/src/server/slide_model.zig @@ -1,6 +1,8 @@ -//! The served model as the Sliding Weights view draws it: its layers, its pockets, and the shards that hold them. +//! The served model as the Living Weights view draws it: layers, pockets and shards, named as its family names them. const std = @import("std"); const json = @import("json.zig"); +const families = @import("living").families; +const Family = @import("living").Family; const Allocator = std.mem.Allocator; const Io = std.Io; @@ -16,15 +18,20 @@ pub const Model = struct { shards: []const Shard = &.{}, }; -const prefix = "backbone.layers."; -const expert = "mixer.shared_experts."; +/// The family a config names by its model_type; Nemotron's for one that names none (as configs before families did). +fn familyOf(o: std.json.ObjectMap) ?Family { + const v = o.get("model_type") orelse return families.nemotron_h; + return if (v == .string) families.byModelType(v.string) else null; +} /// The folder's config, shard headers and stats; the pocket's rows read where the shards hold one. pub fn scan(a: Allocator, io: Io, dir: []const u8) !Model { const config = try readJson(a, io, try std.fs.path.join(a, &.{ dir, "config.json" })); if (config != .object) return error.BadConfig; - const kinds = try kindsOf(a, config.object); - const pocket = intField(config.object, "tensorfold_learned_neurons") orelse 0; + const f = familyOf(config.object) orelse return error.UnknownFamily; + const keys = families.scoped(f, config.object) orelse return error.BadConfig; // text_config, or the top level + const kinds = try families.kinds(a, f, keys); + const pocket = intField(keys, f.keys.pocket) orelse 0; var moe: std.ArrayList(u16) = .empty; for (kinds, 0..) |k, l| if (k == 'E') try moe.append(a, @intCast(l)); const used = try a.alloc(u16, moe.items.len); @@ -39,7 +46,7 @@ pub fn scan(a: Allocator, io: Io, dir: []const u8) !Model { const path = try std.fs.path.join(a, &.{ dir, name }); const stat = try Io.Dir.cwd().statFile(io, path, .{}); var shard: Shard = .{ .name = name, .bytes = stat.size, .mtime = stat.mtime.toNanoseconds(), .inode = @intCast(stat.inode), .layers = &.{}, .expert_bytes = 0 }; - readShard(a, io, path, &shard, pocket, moe.items, used) catch {}; + readShard(a, io, path, f, &shard, pocket, moe.items, used) catch {}; try shards.append(a, shard); } std.mem.sort(Shard, shards.items, {}, byName); @@ -50,33 +57,13 @@ fn byName(_: void, x: Shard, y: Shard) bool { return std.mem.lessThan(u8, x.name, y.name); } -/// The layer kinds as letters, from hybrid_override_pattern, else layers_block_type, else num_hidden_layers unknowns. -fn kindsOf(a: Allocator, o: std.json.ObjectMap) ![]const u8 { - if (o.get("hybrid_override_pattern")) |v| if (v == .string and v.string.len > 0) return a.dupe(u8, v.string); - if (o.get("layers_block_type")) |v| if (v == .array) { - const out = try a.alloc(u8, v.array.items.len); - for (v.array.items, out) |b, *k| k.* = letter(if (b == .string) b.string else ""); - return out; - }; - const n = intField(o, "num_hidden_layers") orelse 0; - const out = try a.alloc(u8, n); - @memset(out, '?'); - return out; -} - -fn letter(name: []const u8) u8 { - const names = [_]struct { []const u8, u8 }{ .{ "mamba", 'M' }, .{ "moe", 'E' }, .{ "attention", '*' }, .{ "mlp", '-' } }; - for (names) |n| if (std.mem.eql(u8, name, n[0])) return n[1]; - return '?'; -} - fn intField(o: std.json.ObjectMap, key: []const u8) ?usize { const v = o.get(key) orelse return null; return if (v == .integer and v.integer >= 0) @intCast(v.integer) else null; } /// The shard's header: the layers whose shared expert it holds and their bytes, and the pocket's neurons in use there. -fn readShard(a: Allocator, io: Io, path: []const u8, shard: *Shard, pocket: usize, moe: []const u16, used: []u16) !void { +fn readShard(a: Allocator, io: Io, path: []const u8, f: Family, shard: *Shard, pocket: usize, moe: []const u16, used: []u16) !void { var file = try Io.Dir.cwd().openFile(io, path, .{}); defer file.close(io); var len: [8]u8 = undefined; @@ -91,14 +78,19 @@ fn readShard(a: Allocator, io: Io, path: []const u8, shard: *Shard, pocket: usiz var it = header.object.iterator(); while (it.next()) |kv| { const name = kv.key_ptr.*; - const at = std.mem.indexOf(u8, name, expert) orelse continue; + const expert = f.names.expert; + const at = std.mem.indexOf(u8, name, expert) orelse continue; // an MTP head's expert's bytes count too const range = offsets(kv.value_ptr.*) orelse continue; shard.expert_bytes += range[1] - range[0]; - const layer = layerOf(name[0..at]) orelse continue; + const layer: u16 = @intCast(f.expertLayer(name) orelse continue); if (std.mem.indexOfScalar(u16, layers.items, layer) == null) try layers.append(a, layer); if (pocket == 0) continue; + // a row a neuron: the first input projection's scales or biases (up_proj, or a SwiGLU's gate_proj) const leaf = name[at + expert.len ..]; - if (!std.mem.eql(u8, leaf, "up_proj.scales") and !std.mem.eql(u8, leaf, "up_proj.biases")) continue; + const proj = f.names.inputs[0]; + if (!std.mem.startsWith(u8, leaf, proj) or leaf.len <= proj.len or leaf[proj.len] != '.') continue; + const part = leaf[proj.len + 1 ..]; + if (!std.mem.eql(u8, part, "scales") and !std.mem.eql(u8, part, "biases")) continue; const slot = std.mem.indexOfScalar(u16, moe, layer) orelse continue; used[slot] = @max(used[slot], try inUse(a, io, file, 8 + n + range[0], kv.value_ptr.*, pocket)); } @@ -106,15 +98,6 @@ fn readShard(a: Allocator, io: Io, path: []const u8, shard: *Shard, pocket: usiz shard.layers = layers.items; } -/// A backbone layer's index from a name's start ("backbone.layers.13."), else null. -fn layerOf(head: []const u8) ?u16 { - if (!std.mem.startsWith(u8, head, prefix)) return null; - const rest = std.mem.trimEnd(u8, head[prefix.len..], "."); - const dot = std.mem.indexOfScalar(u8, rest, '.') orelse rest.len; - if (dot != rest.len) return null; - return std.fmt.parseInt(u16, rest, 10) catch null; -} - fn offsets(entry: std.json.Value) ?[2]u64 { if (entry != .object) return null; const o = entry.object.get("data_offsets") orelse return null; diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index 8a85dfccc1..7a6befdab1 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -40,14 +40,8 @@ test { _ = @import("families/nemotron/simd_attention.zig"); // The pre-M5 attention rewrite finds its lines. _ = @import("families/nemotron/weights.zig"); _ = @import("families/nemotron/kernels.zig"); - _ = @import("families/nemotron/subspace.zig"); // a new block's directions _ = @import("families/nemotron/living_test.zig"); // saved lessons and their routers, beside the shards _ = @import("families/nemotron/pocket_shards_test.zig"); // learned neurons in the shards themselves - _ = @import("families/nemotron/pocket.zig"); // which row a fact's neurons fire at - _ = @import("families/nemotron/pocket_fit.zig"); // where a fact's neurons fire, and stay silent - _ = @import("families/nemotron/router.zig"); // when a lesson acts on a reply - _ = @import("families/nemotron/lesson_route.zig"); - _ = @import("families/nemotron/learner.zig"); _ = @import("families/flashnext/tp_settings.zig"); // speed-up settings read without a typed JSON parse _ = @import("families/flashnext/follow.zig"); // rank 1's reply hash _ = @import("families/flashnext/dense.zig"); // The core projection adapter keeps format, cuts and tiles.