Skip to content

Latest commit

 

History

History
107 lines (87 loc) · 5.26 KB

File metadata and controls

107 lines (87 loc) · 5.26 KB

LFM2.5 hybrid — results

Summary

The browser engine can now train a per-layer conv/attention hybrid end-to-end. Previously only the conv mixer kernel existed (engine/hybrid/lfm2-conv.ts); the wiring that dispatches conv vs attention per layer, runs the SwiGLU FFN, and decodes with a conv-state cache did not. It does now, and it's verified on CPU.

Real-weights forward-parity: PASSED (2026-07-04). The real LiquidAI/LFM2.5-1.2B-Instruct checkpoint (2.34 GB, bf16), loaded via lfm2-loader.ts + a local safetensors reader, produces next-token logits matching a transformers (f32) reference on one prompt within bf16/f32 tolerance: maxAbsErr 1.8e-5, cosine 0.99999999987, argmax and top-5 agree, over the full 65536-way vocab. Conv split order, norm placement, and the SwiGLU role mapping are all correct against real weights, not just a CPU stand-in. LFM2.5 is now verified end-to-end and ready to integrate.

The architecture (verified against the real HF checkpoint)

LiquidAI/LFM2.5-1.2B-Instruct, from the real safetensors header + config (fetched 2026-07):

  • 16 layers, layer_types = conv conv attn conv conv attn conv conv attn conv attn conv attn conv attn conv10 conv + 6 attention (attention at 2,5,8,10,12,14).
  • hidden 2048, 32 heads / 8 kv, head_dim 64, rope θ=1e6.
  • per layer: operator_norm (pre-mixer), ffn_norm (pre-ffn), SwiGLU FFN (feed_forward.w1=gate, w3=up, w2=down, real width 8192).
  • conv layer: conv.in_proj [6144,2048] → split (B,C,x), depthwise conv.conv [2048,1,3] (kernel 3), conv.out_proj [2048,2048].
  • attention layer: GQA q/k/v/out + per-head q/k layernorm (dim 64), like Qwen3.
  • model.embedding_norm = final norm; embeddings tied to the LM head.

What was built

  • lfm2-model.ts — the hybrid Lfm2Model: per-layer conv|attention dispatch, rope + q/k-norm attention, SwiGLU, tied head, LoRA on the attention layers' q/v, and a conv-state generation cache (each conv layer keeps its last K−1 B*x frames; attention layers keep normal KV).
  • lfm2-loader.ts — maps the real safetensors onto the model, with the two non-obvious transforms handled: conv weight [convDim,1,K] → [K,convDim] and the SwiGLU w1/w2/w3 roles.

What was verified (CPU, lfm2.test.ts + lfm2-loader.test.ts)

On a tiny model carrying both mixer types:

  1. forward runs through conv and attention layers, finite outputs.
  2. gradients are finite and nonzero through conv weights (in_proj, conv, out_proj) and attention weights (q/v_proj).
  3. training drives a fact to loss < 0.05 and the model generates it back ([10,11,12]).
  4. generation is faithful — greedy first token == teacher-forced argmax at the last prompt position (no train/decode divergence).
  5. conv-KV-cache is correct — incremental cached decode produces the identical token sequence to full recompute. This is the subtle part of a hybrid (conv needs history, not zero-pad, mid-stream) and it's proven, not assumed.

Loader validation (mock reader, no download): the real tensor names/shapes assemble into a runnable model; conv reshape preserves tap order; ffn width is inferred from w1.

Real-weights forward-parity — how it was checked

verify-parity.ts + ../toolkit/safetensors-reader.ts (new): a Node-only local safetensors reader (handles bf16 → f32 widening exactly, since bf16 is just the top 16 bits of an f32) backs lfm2-loader.ts with the real downloaded checkpoint instead of the mock reader used in unit tests. gen_reference.py (this directory; the committed toolkit/parity.ts snippet) produced one transformers reference — chat-templated prompt, f32 logits at the last position — checked into reference.json so the parity result is reproducible without regenerating it. verify-parity.ts ran the same input_ids through the JS engine and compared with compareLogits.

# 1. reference (already committed as reference.json for this prompt/repo):
python lfm2-hybrid/gen_reference.py LiquidAI/LFM2.5-1.2B-Instruct \
  "What is the capital of France?" > lfm2-hybrid/reference.json

# 2. download the real checkpoint (2.34 GB, not committed — config.json +
#    model.safetensors from LiquidAI/LFM2.5-1.2B-Instruct) into <ckpt-dir>,
#    then:
node --experimental-transform-types lfm2-hybrid/verify-parity.ts <ckpt-dir> lfm2-hybrid/reference.json

(--experimental-transform-types is needed because the model class uses a TS parameter-property constructor that Node's strip-only mode rejects; Vitest's esbuild transform handles it without the flag.)

Result: pass (see Summary). Elementwise error is consistent with bf16-weights-computed-in-f32 rounding, not a structural bug.

What's left

  1. Integrate the hybrid dispatch into whichever product engine's model/config modules need it (accept model_type: "lfm2", add a memory-model entry, flip its catalog status). The runtime is verified; this is wiring, not research.
  2. LoRA-on-conv (in_proj) is a later experiment; q/v on the 6 attention layers works day one and is what training was verified against.
  3. The parity check above is n=1 prompt — cheap to extend to a few more prompts/seeds if the integration step wants more confidence before shipping broadly.