BERT on the DSL path: complete bertNetwork(), BertEncoderRuntime (eager + traced), one-call BertEmbeddingModel, legacy runtime removed#222
Open
michalharakal wants to merge 6 commits into
Conversation
…ithCosSin, kgemma FunctionGemma) Pre-existing on develop: both were merged without running apiDump, so apiCheck fails on a pristine checkout. Isolated commit so the BERT work that follows starts from a green gate. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… bertNetwork() BertEmbeddings module: word gather + position rows via narrow(table, 0, 0, L) + token-type row 0 broadcast ([L,dim]+[dim], the Linear-bias pattern), then LayerNorm. Index-free position/type lookups keep the traced graph at exactly one non-parameter leaf. Single-segment by design. bertNetwork() now returns a numerically complete tokens→hidden-states encoder. BertSafeTensorsNameResolver (llm-core copy) gains position/token_type cases; BertNetworkLoader rewritten: loadWeightTensors(loaders) collection, bare-name→ bert.-prefix normalization, fromWeightTensors with strict mapped==total and shape-fallback off. bert module gains compile-dag/opt deps for the upcoming traced encoder runtime. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…tion/L2, post-norm fix BertEncoderRuntime wraps bertNetwork(): full-sequence forward, masked mean pooling, optional sentence-transformers dense projection, L2 norm — encode() keeps the legacy contract. createBertEncoderRuntime maps checkpoint tensors strictly (shape-fallback off) and pulls the 2_Dense projection pair. Two correctness fixes the old stack hid: - bertNetwork()'s single-block encoder layer wired the FFN residual to the pre-LayerNorm value (the blocks' residual rule fits pre-norm decoders). BERT is post-norm; each layer is now two blocks so every residual segment starts at the right value. Verified against the PyTorch-validated eager runtime on real mdbr-leaf-mt: hidden-state parity ≤2.2e-6. - LEAF's 2_Dense head is bias-free; the legacy runtime required weight AND bias and silently skipped the projection (returning 384-dim vectors while advertising 1024). The new runtime applies bias-free projection, verified ≤6e-8 against independent recomputation. BertModelConfig + MDBR_LEAF_IR_CONFIG moved to BertConfig.kt (same package) with a regex BertConfigParser (config.json + 2_Dense/config.json). Resolver learns sub-block layer paths (encoder.layer.N.attn/.ffn). BertNumericalAccuracyTest retargeted to the DSL path; BertLegacyParityTest bridges old→new until BertRuntime is deleted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…O export gate BertEncoderRuntime.OPTIMIZED traces the encoder (tape → toComputeGraph → LLM optimization pipeline → ComputeGraphExecutor), shape-specialized per sequence length in a small LRU cache; DIRECT stays the default. Graph leaves split by module-parameter identity; the token input is the unique Int32 leaf (BertEmbeddings' index-free position/type lookups guarantee that), other non-param leaves (eager weight transposes, attention constants) ride along as trace-time statics. Fixes a latent executor replay bug this path exposed: the builtin dispatch replayed 'permute' as a last-two-dims transpose, dropping the recorded axes — wrong for MHA's [1,0,2] heads/seq swap on any multi-token trace (decode's seqLen-1 shortcut never hits it). LLMFusedOpHandlers now registers an axes-aware permute handler (registry precedes the builtin); remove once the engine executor honors 'axes' upstream. Gates: DIRECT vs OPTIMIZED bit-exact on real mdbr-leaf-mt and synthetic models (BertTraceParityTest); bertNetwork lowers to StableHLO with gather/dot/add intact (BertStableHloExportTest, IREE execution explicitly out of scope). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…s + fromHuggingFace BertEmbeddingModel.fromSafeTensors(dir) auto-detects weights file, config, 2_Dense projection head, and tokenizer (vocab.txt → tokenizer.json), builds the DSL encoder, and returns the neutral EmbeddingModel SPI. fromHuggingFace(repoId) downloads missing snapshot files via the engine's skainet-data-source (hf:// URIs, HF_TOKEN-aware) into ~/.cache/skainet/models/<owner>_<repo> — streamed, .part + atomic rename, optional files (bias-free heads) skip gracefully — then delegates. llm-providers gains data-source + kotlinx-io deps (BOM-aligned). Consumers retargeted off the deprecated eager runtime: - SkaiNetEmbeddingModel now adapts BertEncoderRuntime (suppression gone). - KBertJava/KBertSession: method surface unchanged (Java smoke test passes untouched); internals on the DSL path — also picks up 2_Dense heads the legacy facade silently ignored. - kbert-cli Main/Demo rewritten onto the factory; Main accepts an HF repo id directly (kbert MongoDB/mdbr-leaf-mt "query" "doc"). Gates: BertEmbeddingModelTest (semantic ranking + HF download) on real mdbr-leaf-mt; BertLeafReferenceSmokeTest green on rewritten KBertJava; CLI e2e produces 1024-dim projected embeddings, paraphrase cosine 0.92. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Deleted: BertRuntime, BertRuntimeWeights/BertLayerWeights, loadBertWeights/ BertWeightMapper/BertTensorNames (BertWeightLoader.kt), BertIngestion, and their tests (BertRuntimeTest superseded by BertEncoderRuntimeTest; BertLegacyParityTest served as the migration bridge and dies with the class it validated). BertModelConfig + MDBR_LEAF_IR_CONFIG live on in BertConfig.kt, same package. PooledExecutionContext and HuggingFaceTokenizer stay. BCV dumps regenerated (jvm via apiDump; android reconciled by the same class-level delta — commonMain-only changes, verified structurally). CHANGELOG gains the full Unreleased entry incl. the BREAKING migration notes. README BERT row upgraded to 'verified'. Docs pages (tutorials/embeddings, explanation/embeddings, dsl-vs-handcoded, specs/spring-ai-adapter) rewritten onto the real API — the previous examples referenced loaders that never existed (loadGGUF, fromGguf). git grep BertRuntime now hits only historical prose in llm-core KDoc. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
Author
|
Engine-side permute fix SKaiNET-developers/SKaiNET#803 is now merged into the engine's develop (merge commit No change needed in this PR: it builds against the released engine 0.35.0, where the executor bug still exists, so the axes-aware 🤖 Generated with Claude Code |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Completes the BERT migration to the DSL philosophy (network definition → eager DIRECT execution or traced/optimized ComputeGraph → StableHLO export), adds a one-call consumption API, and removes the deprecated hand-coded eager stack in the same release.
The gap this closes
BertRuntimewas deprecated pointing at a migration guide that was never written and aReplaceWithfactory that doesn't exist, while the "new path" was half-built:bertNetwork()omitted position/token-type embeddings (numerically incomplete on its own),OptimizedLLMRuntimeis generation-shaped (no full-sequence hidden states, no pooling), andSkaiNetEmbeddingModel+KBertJavawere still hard-wired to the deprecated class.What's new
BertEmbeddings— complete embeddings block (word gather + position rows vianarrow(table, 0, 0, L)+ token-type row 0 broadcast, LayerNorm). Index-free position/type lookups keep the traced graph at exactly one non-parameter leaf. Single-segment by design.BertEncoderRuntime— full-sequence encoder + masked mean pooling + optional2_Denseprojection + L2 norm.DIRECT(default) runs the module tree eagerly;OPTIMIZEDexecutes a traced, LLM-pipeline-optimized ComputeGraph, shape-specialized per sequence length (small LRU).exportTape(...)feeds the StableHLO path.BertEmbeddingModel.fromHuggingFace("MongoDB/mdbr-leaf-mt")/fromSafeTensors(dir)(llm-providers) — one-call factories behind the neutralEmbeddingModelSPI, with built-in Hub download via the engine'sskainet-data-source(hf://URIs,HF_TOKEN-aware,.part+ atomic rename, offline-safe after first run).kbert-cliaccepts a repo id directly.SkaiNetEmbeddingModel(adapter),KBertJava/KBertSession(Java method surface unchanged — smoke test passes untouched),kbert-cli.Bugs found and fixed on the way
ResidualAdd) fits pre-norm decoders; BERT is post-norm, so the FFN residual grabbed the pre-LayerNorm value. Each encoder layer is now two blocks (attn/ffn).2_Denseheads were silently dropped — the legacy runtime required weight AND bias; LEAF shipsbias=false, so it returned unprojected 384-dim vectors while advertising 1024. The new runtime applies bias-free projections;KBertJavapicks up heads it previously ignored.permuteaxes (engine executor bug, decode never hits it) — worked around here with an axes-aware"permute"handler inLLMFusedOpHandlers; the proper engine fix is fix(compile-dag): ComputeGraphExecutor replays permute with its recorded axes SKaiNET#803, and the override gets dropped when that release is consumed.Removed (BREAKING)
BertRuntime,BertRuntimeWeights/BertLayerWeights,loadBertWeights,BertWeightMapper,BertTensorNames,BertIngestion. Migration:createBertEncoderRuntime(config, tensors, ctx)orBertEmbeddingModel.from*(...).BertModelConfig+MDBR_LEAF_IR_CONFIGmoved file (same package — imports unaffected). BCV dumps regenerated; CHANGELOG carries the full BREAKING notes.Verification
MongoDB/mdbr-leaf-mt: hidden states ≤ 2.2e-6, embeddings ≤ 9e-8 (bridge test lived until the deletion commit).BertLeafReferenceSmokeTest(Java surface) passes unmodified; new factory tests incl. Hub download; full jvm suites +apiCheckgreen.First commit regenerates two stale BCV dumps that pre-exist on develop (
transformer-coreRoPE,kgemmaFunctionGemma) so the branch starts from a greenapiCheck.🤖 Generated with Claude Code