Skip to content

Compile reusable GPU training programs and add configurable browser training - #461

Merged
ealmloff merged 228 commits into
mainfrom
backprop
Oct 2, 2026
Merged

ealmloff merged 228 commits into
mainfrom
backprop

Conversation

@ealmloff

@ealmloff ealmloff commented Sep 14, 2026 •

Copy link
Copy Markdown
Member

Add reusable GPU training programs and a configurable browser transformer trainer. A forward/backward/optimizer graph compiles once, then reuses kernels, allocations, and device state across steps. The compiler owns indexing, fusion boundaries, buffer reuse, and synchronization.

Changes:

  • Add TrainingProgram with persistent inputs, simultaneous state feedback, bounded asynchronous submissions, and independent exported snapshots. Logical indexing and workgroup ownership checks determine legal fusion; lifetime interference and size-based packing reuse storage.
  • Extend ordinary graph execution with grouped dispatches, fused map/reduction regions, view absorption, split-K gradient contractions, and shared reduction helpers. Tiny workloads can run in one kernel; larger transformers retain tiled parallel execution.
  • Bind each physical GPU arena once across mixed f32/u32/i32/f16 views, with race-safe packed-half stores. Construct supported schedule combinations before inserting graph variants; preserve contraction output axes across lowering and derive reduction schedules from the actual accumulator footprint. Performance selection and its persistent cache use timings only. Independent invariant checks and numerical member comparisons are enabled only by the compiler-tests feature.
  • Use advertised WebGPU subgroups and experimental f32 subgroup matrices, with capability checks and portable fallbacks. Simplify matrix addressing and pin the wgpu/Naga fork consistently across all three Cargo roots.
  • Add browser controls for blocks, width, heads, feed-forward width, context, batch size, and training token budget, plus held-out evaluation, generation, and attention inspection. Architecture controls accept positive dimensions without preset upper bounds; shape arithmetic and actual device buffer limits govern admission. The default is a 248,352-parameter model trained for 20 million character tokens. Download and verify an approximately 8-million-character TinyStories slice into a local cache.
  • Add native training/parity tests and browser tests for capabilities, fallbacks, configuration, and corpus caching, with recorded performance results.

Measured on Apple M2 Max, f32:

Workload Result
Small MLP 8 to 1 dispatch; warmed Session 219–222 µs/step versus compiled 38–63 µs/step
Tiny transformer, 240,480 parameters, 1,024 tokens/step 216 to 99 dispatches; native compiled approximately 1.63 ms/step and Chrome approximately 2.4 ms/step
Browser default, 248,352 parameters, 2,048 tokens/step 20,000,768 tokens in 72.1 seconds including compilation, evaluation, sampling, and one pause/resume; training average 5.6 ms/step

The Session comparison measures the two execution paths on this branch. Tiny transformer timings exclude initialization, evaluation, sampling, and rendering. These are single-device measurements.

Validation on Apple M2 Max: the retained training/program/library/cache checks pass (30 tests); 58 CPU/GPU matmul and quantized cases, 36 normalization cases, and 32 sampling cases passed. The native BERT integration test passed with differential member checks enabled. Training-state parity also passed with compiler checks disabled. Release Chrome passed 73 targeted matmul/normalization/attention cases and all eight training capability/fallback modes, with matching losses and no GPU errors. Constructor tests cover cooperative scratch combinations, promoted reduction domains, absent CPU schedules, and equivalent output shapes. All 58 targeted convolution/pooling, normalization, and elementwise CPU/GPU cases pass with complete default shape sampling. The production browser configuration test trains a nine-block model and covers generation, attention images, reset, and pause/resume. Tests compare bitmap pixels and persisted cache formats directly; duplicate assertions and print-only diagnostics are removed. Workspace strict Clippy, CPU-only strict Clippy, production WASM checking, and formatting passed.

Model checks compare every parameter and Adam moment with Session across three architectures. A full default browser training run completed at 0.912 held-out loss and 72% next-character accuracy. These local checks do not replace the Linux/macOS/Windows CI results shown on this PR. Training tests are serialized within the Fusor workspace, with a finite timeout suitable for software GPUs.

TrainingProgram requires a GPU and static, nonempty f32/u32/i32 shapes; browser matrix acceleration depends on experimental browser/adapter support. Large schedule domains are sampled in numerical conformance; small domains and cooperative resource-domain construction are checked exhaustively.

…, fold the contraction families onto one setup, move timing into the launcher
…uites, skip the 400-step training test on software adapters
… level into one group after extraction, row-per-workgroup scatter, packed coop slots in groups
@ealmloff
ealmloff merged commit 04e003d into main Oct 2, 2026
57 checks passed
@ealmloff
ealmloff deleted the backprop branch October 2, 2026 15:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant