From 5a4987b7702255cd0e8cb21b50b06834fd4c076b Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 02:53:02 -0500 Subject: [PATCH 1/3] Living Weights: the learner, pocket and fits move to zig/src/living, unchanged The learner, its whole-lesson and routed paths, the pocket and its fits and helpers move from families/nemotron to a `living` module that both builds import, as 19 renames with no change to their contents. The family-neutral half of learn_shape.zig (block sizes, Adam settings, direction lengths) moves with them; Nemotron's learn_shape.zig keeps where its blocks sit and names the shared constants from the module. Nemotron's files import the moved ones through the module. The learner's tests now also run on Linux, as `zig build test-living` and under test-cuda-host. No behaviour change. --- build.zig | 8 ++-- zig/build/cuda.zig | 11 +++++- zig/src/families/nemotron/adapters.zig | 4 +- zig/src/families/nemotron/backend.zig | 2 +- zig/src/families/nemotron/cuda.zig | 2 - zig/src/families/nemotron/cuda_engine.zig | 2 +- zig/src/families/nemotron/cuda_forward.zig | 2 +- zig/src/families/nemotron/cuda_lanes.zig | 2 +- zig/src/families/nemotron/cuda_prompt.zig | 2 +- zig/src/families/nemotron/cuda_sites.zig | 4 +- zig/src/families/nemotron/cuda_slide.zig | 2 +- zig/src/families/nemotron/cuda_state.zig | 2 +- zig/src/families/nemotron/cuda_train.zig | 6 +-- zig/src/families/nemotron/cuda_train_fwd.zig | 2 +- zig/src/families/nemotron/cuda_train_ops.zig | 2 +- zig/src/families/nemotron/forward.zig | 2 +- zig/src/families/nemotron/layers.zig | 2 +- zig/src/families/nemotron/learn_shape.zig | 39 +++++++------------ zig/src/families/nemotron/nemotron.zig | 4 +- zig/src/families/nemotron/prefill.zig | 2 +- zig/src/families/nemotron/slide.zig | 2 +- zig/src/families/nemotron/state.zig | 2 +- zig/src/families/nemotron/train.zig | 6 +-- zig/src/families/nemotron/train_back.zig | 2 +- zig/src/families/nemotron/train_forward.zig | 2 +- zig/src/families/nemotron/train_job.zig | 4 +- zig/src/families/nemotron/train_ops.zig | 2 +- .../{families/nemotron => living}/choice.zig | 0 .../{families/nemotron => living}/editor.zig | 0 zig/src/living/learn_shape.zig | 25 ++++++++++++ .../{families/nemotron => living}/learner.zig | 0 .../nemotron => living}/lesson_examples.zig | 0 .../nemotron => living}/lesson_route.zig | 0 zig/src/living/living.zig | 30 ++++++++++++++ .../{families/nemotron => living}/pocket.zig | 0 .../nemotron => living}/pocket_clock.zig | 0 .../nemotron => living}/pocket_dump.zig | 0 .../nemotron => living}/pocket_fit.zig | 0 .../nemotron => living}/pocket_joint.zig | 0 .../nemotron => living}/pocket_pack.zig | 0 .../nemotron => living}/pocket_prefix.zig | 0 .../nemotron => living}/pocket_read.zig | 0 .../nemotron => living}/pocket_values.zig | 0 .../{families/nemotron => living}/probe.zig | 0 .../{families/nemotron => living}/router.zig | 0 .../nemotron => living}/subspace.zig | 0 .../nemotron => living}/train_mode.zig | 0 .../{families/nemotron => living}/whole.zig | 0 zig/src/tensorfold.zig | 6 --- 49 files changed, 114 insertions(+), 69 deletions(-) rename zig/src/{families/nemotron => living}/choice.zig (100%) rename zig/src/{families/nemotron => living}/editor.zig (100%) create mode 100644 zig/src/living/learn_shape.zig rename zig/src/{families/nemotron => living}/learner.zig (100%) rename zig/src/{families/nemotron => living}/lesson_examples.zig (100%) rename zig/src/{families/nemotron => living}/lesson_route.zig (100%) create mode 100644 zig/src/living/living.zig rename zig/src/{families/nemotron => living}/pocket.zig (100%) rename zig/src/{families/nemotron => living}/pocket_clock.zig (100%) rename zig/src/{families/nemotron => living}/pocket_dump.zig (100%) rename zig/src/{families/nemotron => living}/pocket_fit.zig (100%) rename zig/src/{families/nemotron => living}/pocket_joint.zig (100%) rename zig/src/{families/nemotron => living}/pocket_pack.zig (100%) rename zig/src/{families/nemotron => living}/pocket_prefix.zig (100%) rename zig/src/{families/nemotron => living}/pocket_read.zig (100%) rename zig/src/{families/nemotron => living}/pocket_values.zig (100%) rename zig/src/{families/nemotron => living}/probe.zig (100%) rename zig/src/{families/nemotron => living}/router.zig (100%) rename zig/src/{families/nemotron => living}/subspace.zig (100%) rename zig/src/{families/nemotron => living}/train_mode.zig (100%) rename zig/src/{families/nemotron => living}/whole.zig (100%) diff --git a/build.zig b/build.zig index 08765bd1df..5e7677cba7 100644 --- a/build.zig +++ b/build.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, qwen27: qwen27_build.Shared, living: *std.Build.Module } { 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' learner, pocket and fits, which Nemotron's lessons import + 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 { diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index eac681628c..be4526406c 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' learner, pocket and fits, which a family's lessons import + 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`. @@ -211,6 +214,10 @@ 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' learner on the host, beside the family that imports it + const living_run = b.addRunArtifact(b.addTest(.{ .root_module = mods.living })); + step.dependOn(&living_run.step); + b.step("test-living", "Living Weights' learner, pocket and fits, on the host").dependOn(&living_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); 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..c6495fa7fd 100644 --- a/zig/src/families/nemotron/cuda.zig +++ b/zig/src/families/nemotron/cuda.zig @@ -40,7 +40,5 @@ test { _ = @import("cuda_prompt_cuts.zig"); _ = @import("cuda_profile_window.zig"); _ = reuse; - _ = @import("learner.zig"); - _ = @import("subspace.zig"); _ = @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..acc0452082 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"); 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/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..d062d2f40a 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"); 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/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/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 100% rename from zig/src/families/nemotron/learner.zig rename to zig/src/living/learner.zig 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..9b9d1d1afe --- /dev/null +++ b/zig/src/living/living.zig @@ -0,0 +1,30 @@ +//! Living Weights: the lesson learner, its pocket of new neurons beside each shared expert, their fits and helpers. +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 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_pack = @import("pocket_pack.zig"); +pub const pocket_prefix = @import("pocket_prefix.zig"); +pub const pocket_read = @import("pocket_read.zig"); +pub const pocket_values = @import("pocket_values.zig"); + +test { + _ = learner; + _ = subspace; + _ = pocket; + _ = pocket_fit; + _ = router; + _ = lesson_route; +} diff --git a/zig/src/families/nemotron/pocket.zig b/zig/src/living/pocket.zig similarity index 100% rename from zig/src/families/nemotron/pocket.zig rename to zig/src/living/pocket.zig 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 100% rename from zig/src/families/nemotron/pocket_fit.zig rename to zig/src/living/pocket_fit.zig 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/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 100% rename from zig/src/families/nemotron/pocket_prefix.zig rename to zig/src/living/pocket_prefix.zig diff --git a/zig/src/families/nemotron/pocket_read.zig b/zig/src/living/pocket_read.zig similarity index 100% rename from zig/src/families/nemotron/pocket_read.zig rename to zig/src/living/pocket_read.zig diff --git a/zig/src/families/nemotron/pocket_values.zig b/zig/src/living/pocket_values.zig similarity index 100% rename from zig/src/families/nemotron/pocket_values.zig rename to zig/src/living/pocket_values.zig diff --git a/zig/src/families/nemotron/probe.zig b/zig/src/living/probe.zig similarity index 100% rename from zig/src/families/nemotron/probe.zig rename to zig/src/living/probe.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/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/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 100% rename from zig/src/families/nemotron/whole.zig rename to zig/src/living/whole.zig 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. From 75978a1c5c05d9c2aff212a0e2a7d2fd3938145a Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 03:03:27 -0500 Subject: [PATCH 2/3] Living Weights learns into any family described as data: the family's layout, neuron kind and shard format A family now describes itself once, as data (living/families.zig): where its shared experts sit and their tensor names, the neuron's activation (relu^2, or SwiGLU with an optional per-row gate), the quantization, the config keys and the object they live in, and its MTP head's expert. Its host gives the learner a layout and the pocket hooks (living/family.zig). Nemotron is the one family listed, and its lessons, routed blocks, edits and probes are unchanged. - The pocket reads its layers and width from the family's layout, not from Nemotron's config, and the learner offers routed blocks, edits and probes only to a trainer with `sites` (Nemotron's). - living/shards.zig widens any family's experts at its group and width; Nemotron's writer is now a wrapper over it. shard_edit sets and reads nested config keys (text_config). - A SwiGLU neuron fitted as u is written as gate row 16u and up row u/16: relu(u x)^2 on its fact's rows, at most 8e-5 of it at a 0.5 margin; a gated family's outputs fit gate * act(x). No listed family uses it yet; the tests use a stand-in family shaped as Qwen3.8 Flash Next's shared expert. - The server's view reads layers, pockets and shards through the family its config's model_type names, and native/nemotron_slide.zig is native/slide_adapter.zig. - pocket.zig's quiet refit and its loading move to pocket_quiet.zig and pocket_load.zig, for the 600-line rule. - Tests: the generic writer equals affine4_host at 4 bits in groups of 64, Nemotron's save test over the wrapper, a SwiGLU neuron fitted and stored at 4 bits in groups of 32, and a whole pocket lesson over a stand-in host for each activation (test-living, also under test-cuda-host). --- build.zig | 14 +- zig/build/cuda.zig | 21 +- zig/src/core/shard_edit.zig | 47 +++- zig/src/families/nemotron/cuda.zig | 1 + zig/src/families/nemotron/cuda_slide.zig | 7 + zig/src/families/nemotron/pocket_layout.zig | 12 + zig/src/families/nemotron/pocket_shards.zig | 174 +------------- zig/src/families/nemotron/slide.zig | 7 + zig/src/living/affine.zig | 86 +++++++ zig/src/living/families.zig | 79 +++++++ zig/src/living/family.zig | 135 +++++++++++ zig/src/living/host_test.zig | 215 ++++++++++++++++++ zig/src/living/learner.zig | 50 ++-- zig/src/living/living.zig | 22 +- zig/src/living/neuron.zig | 55 +++++ zig/src/living/pocket.zig | 159 +++++++------ zig/src/living/pocket_fit.zig | 18 +- zig/src/living/pocket_load.zig | 30 +++ zig/src/living/pocket_prefix.zig | 2 +- zig/src/living/pocket_quiet.zig | 59 +++++ zig/src/living/pocket_read.zig | 1 + zig/src/living/pocket_values.zig | 22 +- zig/src/living/probe.zig | 2 +- zig/src/living/shards.zig | 211 +++++++++++++++++ zig/src/living/swiglu_test.zig | 82 +++++++ zig/src/living/test_families.zig | 18 ++ zig/src/living/whole.zig | 18 ++ zig/src/living_shards_test.zig | 32 +++ zig/src/native/cuda.zig | 8 +- zig/src/native/metal.zig | 6 +- .../{nemotron_slide.zig => slide_adapter.zig} | 6 +- zig/src/server/slide_model.zig | 61 ++--- 32 files changed, 1303 insertions(+), 357 deletions(-) create mode 100644 zig/src/families/nemotron/pocket_layout.zig create mode 100644 zig/src/living/affine.zig create mode 100644 zig/src/living/families.zig create mode 100644 zig/src/living/family.zig create mode 100644 zig/src/living/host_test.zig create mode 100644 zig/src/living/neuron.zig create mode 100644 zig/src/living/pocket_load.zig create mode 100644 zig/src/living/pocket_quiet.zig create mode 100644 zig/src/living/shards.zig create mode 100644 zig/src/living/swiglu_test.zig create mode 100644 zig/src/living/test_families.zig create mode 100644 zig/src/living_shards_test.zig rename zig/src/native/{nemotron_slide.zig => slide_adapter.zig} (90%) diff --git a/build.zig b/build.zig index 5e7677cba7..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, living: *std.Build.Module } { +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,7 +553,7 @@ 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' learner, pocket and fits, which Nemotron's lessons import + // 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"), @@ -576,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/zig/build/cuda.zig b/zig/build/cuda.zig index be4526406c..d36ee9c52b 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -84,7 +84,7 @@ fn family(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin 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' learner, pocket and fits, which a family's lessons import + // 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); @@ -142,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. @@ -164,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 }); @@ -173,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) { @@ -202,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); @@ -214,10 +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' learner on the host, beside the family that imports it + // 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); - b.step("test-living", "Living Weights' learner, pocket and fits, on the host").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); @@ -253,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/cuda.zig b/zig/src/families/nemotron/cuda.zig index c6495fa7fd..58e979c6e5 100644 --- a/zig/src/families/nemotron/cuda.zig +++ b/zig/src/families/nemotron/cuda.zig @@ -40,5 +40,6 @@ test { _ = @import("cuda_prompt_cuts.zig"); _ = @import("cuda_profile_window.zig"); _ = reuse; + _ = slide; _ = @import("cuda_train_back.zig"); } diff --git a/zig/src/families/nemotron/cuda_slide.zig b/zig/src/families/nemotron/cuda_slide.zig index acc0452082..9c67cf53a0 100644 --- a/zig/src/families/nemotron/cuda_slide.zig +++ b/zig/src/families/nemotron/cuda_slide.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/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/slide.zig b/zig/src/families/nemotron/slide.zig index d062d2f40a..ff62940596 100644 --- a/zig/src/families/nemotron/slide.zig +++ b/zig/src/families/nemotron/slide.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/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/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/learner.zig b/zig/src/living/learner.zig index 0607eb90c1..ddad1c54b8 100644 --- a/zig/src/living/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/living/living.zig b/zig/src/living/living.zig index 9b9d1d1afe..e1d75f4163 100644 --- a/zig/src/living/living.zig +++ b/zig/src/living/living.zig @@ -1,4 +1,10 @@ -//! Living Weights: the lesson learner, its pocket of new neurons beside each shared expert, their fits and helpers. +//! 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"); @@ -10,21 +16,35 @@ 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/living/pocket.zig b/zig/src/living/pocket.zig index ff24ac5ee9..aff4a4ad4a 100644 --- a/zig/src/living/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/living/pocket_fit.zig b/zig/src/living/pocket_fit.zig index 02c9d99493..ab37028d73 100644 --- a/zig/src/living/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/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/living/pocket_prefix.zig b/zig/src/living/pocket_prefix.zig index ca854f38ee..a5239724a4 100644 --- a/zig/src/living/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/living/pocket_read.zig b/zig/src/living/pocket_read.zig index 77d9bcd9dc..246652d6e3 100644 --- a/zig/src/living/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/living/pocket_values.zig b/zig/src/living/pocket_values.zig index 628ce8f265..46b1a96acb 100644 --- a/zig/src/living/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/living/probe.zig b/zig/src/living/probe.zig index 1175275c7a..1d2089f2ec 100644 --- a/zig/src/living/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/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/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/living/whole.zig b/zig/src/living/whole.zig index 4fa4a9fc01..f6a5b34ad0 100644 --- a/zig/src/living/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; From f47c203cea701cd206dfd4b9990650ae58fdf40c Mon Sep 17 00:00:00 2001 From: ZackSample <134018356+ZackSample@users.noreply.github.com> Date: Thu, 15 Oct 2026 18:27:20 -0500 Subject: [PATCH 3/3] Docs: adding Living Weights to a family, the recipe beside adding a Zig family What a family gives the learner in zig/src/living, in order: its description as data (the pocket beside an MoE layer's shared expert, behind a gated one, or in a dense MLP), the pocket in every pass, the host's hooks, the trainer's API with its keeps and backward kernels, the shard writer, --slide on the native host, the receipts a PR should carry, and the traps found bringing up Flash Next and the Qwen3.8-27B. The recipe book links it. --- docs/recipes/README.md | 5 +- .../adding-living-weights-to-a-family.md | 201 ++++++++++++++++++ 2 files changed, 204 insertions(+), 2 deletions(-) create mode 100644 docs/recipes/adding-living-weights-to-a-family.md 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.