Skip to content

lfm3_omni: LFM3-Omni-2B on Metal, speech in, text and speech out (part of #1627) - #1635

Draft
drowzeys wants to merge 3 commits into
ashhart:mainfrom
drowzeys:lfm3-omni-family-metal
Draft

drowzeys wants to merge 3 commits into
ashhart:mainfrom
drowzeys:lfm3-omni-family-metal

Conversation

@drowzeys

Copy link
Copy Markdown
Contributor

What this changes

Adds families/lfm3_omni, LiquidAI's LFM3-Omni-2B (speech in, streamed text and 24 kHz speech out) as a standalone family in Drama 3's shape, with its Metal backend, host tests, kernel-contract registrations, recipe and the checks' clients under tools/lfm3_omni/. Part of #1627 (the Metal half; the CUDA backend is a separate branch stacked on this one). It's a draft because the gate in #1627 isn't agreed yet.

Receipt

  • Environment: this branch's head; M5 Ultra (256 GB), macOS 27.0.1, Xcode 27.1, Zig 0.17.0; reference transformers 5.21.0 with the checkpoint's own code, PyTorch 2.14.1, MPS bf16, same Mac. The 16 synthetic turns of tools/lfm3_omni/utterances.txt (macOS voices), a voice-assistant system text.
  • Stage parity (greedy, identity codebook order, tools/lfm3_omni/parity.sh):
Stage Metal vs MPS bf16
Log-mel, cosine (min), max abs 1.000000, 3.1e-4
Encoder / projector, minimum cosine 0.99999 / 0.99990
Prompt token ids 16/16 identical
Backbone prefill rows, median cosine (min of medians) 0.99982 (0.99977)
Text top-1, teacher-forced 520/524 (99.2%)
Depth decoder on the reference's own states, greedy agreement; logits cosine median / min 95.5%; 0.999994 / 0.99989
Codes, teacher-forced through the whole loop 89.2%
Codec on the reference's codes, waveform SNR median (min) 36.6 dB (25.1 dB)
  • The port refactor (files split under 600 lines, one-line comments, the attention fix below) changes no output: greedy dumps and teacher-forced runs of this branch and of the pre-split build are byte-identical on all 16 turns (208 dump files), and the served replies at the card's settings are byte-identical too.
  • Latency, the same 16 turns through the server with tools/lfm3_omni/speech_chat_client.py (temperature 0.4, penalty 1.2, seeds 1000 + i); the reference through tools/lfm3_omni/reference_speech.py (first audio = first sampled frame plus that frame's codec decode). Two rounds, this branch, the pre-split build and the reference alternated; other services were resident on the Mac but idle:
Median This branch Pre-split build Reference (MPS bf16)
First text 0.034 s 0.034 s n/a
First audio 0.063 s 0.063 s 0.172 / 0.191 s
Whole reply (6.5 s of speech) 0.682 / 0.670 s 0.645 / 0.695 s 2.535 / 2.409 s
  • Quality of the replies (UTMOS22 strong via SpeechMOS; WER = the checkpoint's own ASR on each reply's audio against its streamed text): TensorFold 4.43 median (4.35 min), WER 1.5%; reference 4.44 (4.15), WER 1.2%.
  • Prompt speed: n/a against a previous release (a new family). Prefill is inside "first text" above.

Kernel contract (Metal, M5 Ultra)

All 45 pipelines of lfm3_omni.metallib are registered (69 registrations over the stacks' variants): gates 1-4 against a float64 replay of each kernel's bf16 rounding points, stated rules for the picks and word moves, a planted source change per registration, and gate-5 launch records with M1/M2 limits from --offline.

tf-kernel-contract --only "lfm3 omni":
SUMMARY 69 kernels: 69 pass gates 1-4, 0 fail, 0 known findings; 69 controls: 69 red as planted, 0 missed; 0 limit problems; gates 13690 ms, limits 7242 ms
tf-kernel-contract (every family), exit 0:
SUMMARY 378 kernels: 378 pass gates 1-4, 0 fail, 3 known findings; 414 controls: 414 red as planted, 0 missed; 0 limit problems; gates 2444558 ms, limits 1229 ms

The contract found one real bug, fixed here before the receipt: om_attention_tiles multiplied masked keys' zero weights into the value sum, so a non-finite value row a query cannot see turned its output into NaN. The value loop now covers only the row's visible keys. Finite outputs are unchanged, byte for byte. The full suite ran through tf-kernel-contract --rotary-accuracy with a tf-rotary-sweep receipt, because zig build kernel-contract stops first at Flash Next's UnrecordedRotaryCompiler precondition, which this change doesn't touch.

Two findings that no gate fails, noted for review: in sampled mode om_depth_pick maps h >> 40 == 2^24 - 1 to u = 1 (about 2^-24 per candidate), and om_logmel returns the floor for a NaN log (registered as its stated rule).

Lean checker and tests

  • python3 tools/zig/lean_check.py: main 4571 files, 425 problems; this branch 4600 files, 425 problems. Zero added; the 425 are inherited.
  • zig build test test-golden -Dcpu=apple_m1: 585/585 steps, 8484 passed, 0 failed, 52 skipped (exit 0).
  • zig build test-kernel-contract: 139/139 passed.
  • zig build dev-tools -Dcpu=apple_m1: 215/215 steps (exit 0).
  • zig build lfm3-omni-generate -Doptimize=ReleaseFast -Dcpu=apple_m1: builds.

Not run, and why

From elsewhere

  • om_common.h's erf (for the exact GELU) uses the Chebyshev erfcc coefficients published in Numerical Recipes. Say if you'd rather have a different erf.
  • tools/lfm3_omni/score_replies.py loads UTMOS22 strong from tarepan/SpeechMOS v1.2.0 through torch.hub (a development tool, not engine code).

Checklist

  • A new GPU kernel is registered with the kernel contract; the Metal SUMMARY lines are above.
  • No precision traded for speed; bf16 rounding points follow the reference op for op.
  • Comments and docstrings are one line; no measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths.
  • Lane-round, drafted, resumed and concurrent contracts: n/a for this one-request tool; see LFM3-Omni-2B on Metal and CUDA: speech in, streamed text and speech out (gate and split) #1627 question 2.

…kHz speech out

A standalone family in Drama 3's shape: the four Liquid stacks of the bf16 release (speech
encoder and projector, LFM3 backbone, depth decoder, Liquid codec decoder) over one metallib,
the log-mel front end on the device, each codebook's pick on the GPU, and the codec streamed
between backbone rows. `zig build lfm3-omni-generate` makes one turn from a WAV, writes and
scores stage dumps against a reference's, and serves POST /v1/audio/chat (server-sent text
deltas and 16-bit 24 kHz audio) and GET /health, one request at a time. Products of many rows
take the tensor unit through the prefill GEMM's blocks; other chips a simdgroup-matrix product.
Host tests cover the mel front end, the prompt, the streamed overlap-add, the synthesis basis
and the Metal argument layouts.
…ts bf16 points, a control each

Registers all 45 pipelines of lfm3_omni.metallib (69 registrations over the stacks' variants) in
tf-kernel-contract: gates 1-4 against a float64 replay of each kernel's rounding points with the
spread its fp32 steps allow, stated rules for the picks and the word moves, and a planted source
change per registration. Gate 5 records the launches with their M1/M2 limits from --offline.
… weights' licence

tools/lfm3_omni holds the development clients behind the receipt: the 16 synthetic spoken turns,
the reference's stage dumps and spoken replies (transformers with the checkpoint's own code, on
CUDA, MPS or the CPU), the stage comparison and its table, a streaming client for
/v1/audio/chat, and the UTMOS and speech-vs-text scoring.
@drowzeys

Copy link
Copy Markdown
Contributor Author

atc overlap --json (ATC 0.7.0) against main 3f0a715. Of 72 overlapping open PRs, four conflict, all where both sides append to a shared list:

None touches the same code. Each resolves by keeping both entries, whichever lands first. The rest report unknown because of ATC's pullFilesLimit, so I also ran git merge-tree against every open PR that merges cleanly with main: no further conflicts. #1627 is now claimed through ATC.

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