Skip to content

Latest commit

 

History

History

README.md

toolkit/ — onboard any model to browser fine-tuning

A repeatable, scripted process for taking a new Hugging Face model and getting it to fine-tune cleanly in a browser LoRA training engine. It generalizes this project's two case studies — Granite 4.0 (dense muP) and LFM2.5 (conv/attention hybrid) — into three commands and a decision tree.

Nothing here needs a GPU or the browser: the tools run under plain node (Node ≥ 22) and Vitest. Run all commands below from this project's root.

npm run probe -- <hf-repo>                # 1. what is this model? what does it need?
npm run recipe-sweep -- --repo <hf-repo>   # 2. what LoRA LR trains it without collapse?
npx vitest run toolkit                     # tests for the tools themselves

The process

1. Probe — classify the architecture

npm run probe -- ibm-granite/granite-4.0-350m

Fetches config.json + the safetensors header and prints a verdict. The verdict is one of:

verdict meaning what to do
as-is dense attention (Gemma/Qwen/Llama-like) wire it into your catalog/config as a plain transformer; done
muP-scalars dense but muP-parametrized (Granite) apply the muP scalars in your engine's config normalization; go to step 2
hybrid-dispatch conv/attention interleave (LFM2) follow lfm2-hybrid/; then step 3
gated-deltanet-dispatch Qwen3.5-style linear_attn (Gated DeltaNet)/attention interleave follow qwen3.5-linear-attn/ — built + forward-parity verified; then step 3
mamba-unsupported real SSM/mamba layers out of scope (no scan kernel)
linear-attn-unsupported gated/linear-attention (linear_attention/linear_attn) in a family other than Qwen3.5 — may be a different recurrence than Gated DeltaNet out of scope until verified — check weight names against qwen3.5-linear-attn/RESULTS.md before assuming a from-scratch build is needed
moe-unsupported expert routing out of scope (no routing kernel)

The classifier reads layer_types and num_local_experts, not the class name — so it correctly reports that GraniteMoeHybridForCausalLM dense checkpoints are muP transformers, not hybrids or MoEs. It also unwraps a nested text_config (multimodal wrappers like Qwen3.5-VL put the real text backbone's fields there, not at top level) before reading anything — probing one of those without the unwrap silently reports an empty, wrong as-is. The header probe adds weight-layout facts a config can't tell you (fused gate/up, per-head q/k norm, tied vs separate lm_head, depthwise conv weights, and A_log/dt_bias — the selective-SSM signature that can hide under an innocuous linear_attention layer_type name) and follows a sharded model.safetensors.index.json when there's no single model.safetensors.

2. Recipe-sweep — find the learning rate

npm run recipe-sweep -- --repo ibm-granite/granite-4.0-350m
# or a built-in preset:
npm run recipe-sweep -- --preset granite   # muP
npm run recipe-sweep -- --preset plain     # non-muP

Builds a fast CPU stand-in with the model's real muP scalars and sweeps LR × seed on a 3-fact dataset, printing where training collapses to one token vs converges cleanly, and a recommended LR. For muP models it also runs the LR ÷ residual_multiplier control to demonstrate that the intuitive "muP fix" starves the adapter (the trap that produced Granite's original "collapse to .").

The stand-in maps the shape of the boundary (where undertraining begins), not a real checkpoint's absolute LR — a 2-layer random model sits in a different basin than the real weights. Use it to know which direction to move and that the model is trainable in principle; confirm the actual LR on real weights.

3. Parity — verify before shipping (the credibility gate)

"Loss went down" ≠ "the model is correct" — the Granite bug was a near-zero loss next to garbage output. parity.ts provides two checks:

  • Self-parity (assertSelfParity) — a model's incremental cached decode must emit the identical sequence to a full recompute. Catches KV-cache / conv-state bugs (the #1 hybrid failure). See it used for real in lfm2-hybrid/lfm2.test.ts.
  • Logits-parity (compareLogits) — the engine's next-token logits must match HF transformers within bf16/f32 tolerance on one prompt. Catches what a random stand-in can't: wrong conv split order, misplaced norm, wrong scalar. Generate the reference with the committed python snippet (REFERENCE_SNIPPET), then compare.

Only after logits-parity passes should a model be shipped as fully supported — that's the difference between "the math type-checks" and "produces correct text on the real model."


Files

file what it is
classifier.ts pure architecture classifier (config → verdict)
hub.ts Node HF fetch helpers (config + safetensors header via Range, sharded index)
safetensors-reader.ts Node-only real local safetensors reader (bf16→f32, byte-range tensor reads) — backs a loader with real downloaded weights instead of a mock
probe.ts CLI: classify a live repo
mup-standin.ts compact self-contained muP transformer for the sweep
recipe-sweep.ts CLI: map the LR/collapse boundary, recommend an LR
parity.ts self-parity + logits-parity harness + reference protocol
*.test.ts tests for the tools (run under Vitest)

Worked examples

npm run probe -- ibm-granite/granite-4.0-350m   # → muP-scalars (dense, not hybrid)
npm run probe -- LiquidAI/LFM2.5-1.2B-Instruct  # → hybrid-dispatch (10 conv + 6 attn)
npm run probe -- Qwen/Qwen3-0.6B                # → as-is
npm run probe -- Qwen/Qwen3.5-0.8B              # → gated-deltanet-dispatch (18 linear_attn + 6 attn; A_log/dt_bias = real GatedDeltaNet recurrence, nested text_config — built + verified)
npm run recipe-sweep -- --preset granite        # → recommends the healthy LR; shows the ÷residual trap

Scope / honesty

  • The classifier's mamba-unsupported / linear-attn-unsupported / moe-unsupported verdicts are hard stops: those kernels genuinely don't exist in this engine yet.
  • The stand-in is a mechanism model, not a checkpoint. Every recipe it suggests is a hypothesis to confirm on real weights.
  • Logits-parity needs transformers (Python) for the reference — the one step that leaves the JS/browser world. That's by design: it's the independent oracle.
  • Nested configs (multimodal wrappers) and sharded weight indices are both real shapes real repos ship in — a probe that only handles the simple case silently under-reports instead of failing loudly, which is worse. Qwen3.5 found both gaps; see ../qwen3.5-linear-attn/RESULTS.md.