Speed up eflomal training: hot-loop hashing, allocation, and threading

[johnml1135]

Summary

Investigates and improves eflomal Gibbs-sampler training performance across two rounds: whether multithreading actually scales, whether allocator/memory-management behavior was limiting it, what's left for a hypothetical Rust rewrite, and then a follow-up round targeting cache locality, remaining copies, further threading opportunities, and Rust-style ownership discipline. All investigation logs, methodology, and per-change benchmark numbers are folded into this description below (no separate planning docs are included in the diff).

Headline finding (round 1): at default settings (deterministic=false), training already scales close to ideal (~2.8x on 3 threads, ~84% efficiency at samplers=8,threads=4) on a real corpus (7,952 verse-aligned eBible NT sentence pairs). Allocator contention was tested directly with a verified, active mimalloc override and found not to be a bottleneck. The real win came from reducing redundant hashing in sampleSweep's hot loop, not from memory-layout changes.

Top-line results

All numbers below are wall-clock, measured on the same real corpus (7,952 verse-aligned eBible NT sentence pairs, English/Spanish), min-of-N reps per cell (N=2 or 3; see per-change notes for methodology). Only shipped (kept) changes are listed — see "Tried and reverted" further down for attempts that regressed and were backed out.

Change	Scenario	Result	Bit-identical?
Hash-reuse in sampleSweep's hot loop + supporting round-1 changes (chain-init parallelization, phi buffer hoist, flattened fertility arrays)	Training, samplers=1,threads=1	-26.4% wall-clock	Yes
Same round-1 changes	Training, samplers=8,threads=8	-5.8% wall-clock	Yes
Concurrent direct+inverse training via std::async (usage pattern demo, not a library code change)	numSamplers=3	1.74x speedup	N/A — no code change; two independent model instances
Same usage pattern	numSamplers=8	0.98x (no gain — both directions already saturate a shared resource individually)	N/A
Fold accumLexCounts into lexCounts (round 2, item 1)	Training, samplers=1,t=1 / samplers=8,t=8	-7.6% / -4.9%	Yes
Cache decode-path jumpLin/fertRatio per chain (round 2, item 2)	computeTrainingAlignments	-18.0%	Yes
Persistent lexDenomInv member (round 2, item 5)	Training, samplers=1,threads=1	-4.4%	Yes
Parallelize loglikelihoodForPairRange (round 2, item 11)	Scoring all 7,952 pairs, threads=1 vs threads=8	5.83x speedup	Yes — serial and parallel runs produced the exact same log-likelihood value (-859351.89 both, not just statistically close)

What "bit-identical" means here: with deterministic=true (single-threaded, fixed-seed Gibbs sampler), every change marked "Yes" produces trained probability tables and decoded alignments that are byte-for-byte identical to the pre-change code on the same input — not merely similar or statistically close. Every such change is a pure performance/memory-layout change with no change to the sampling algorithm, and this is its actual correctness gate: the existing golden fixed-alignment tests pin exact expected alignments (any arithmetic or ordering drift breaks them immediately), plus dedicated new regression tests (initializeChainParallelIsReproducible, loglikelihoodForPairRangeEdgeCasesAndDeterminism) specifically exercise the newly-parallel code paths for race conditions. Rows marked N/A are usage-pattern demonstrations (no code in this repo changed) rather than correctness-gated optimizations.

Round 1 changes

• Parallelize chain init in startTraining — the per-chain setup loop was serial despite chains being fully independent.
• Hoist the phi scratch buffer out of sampleSweep's and computeFertilityCounts's per-sentence loops, reusing it via assign() instead of reallocating every sentence, every sweep.
• Flatten fertCounts/fertRatioSampled from vector<vector<double>> to a single strided vector<double>.
• Reuse the target word's hash across the remove/candidate-loop/add tsl::robin_map operations in sampleSweep, instead of rehashing the same fixed key up to slen+3 times per token.
• Off-by-default CMake diagnostic (THOT_BENCH_MIMALLOC) and new tests supporting the above.

A second attempted optimization (saturating jump-bucket running counters) was implemented, benchmarked, then found incorrect by a dedicated code review: the counter doesn't stay pinned correctly when it starts a candidate loop already at a saturation boundary, silently corrupting jump-sampling weights for sentences longer than jumpWindow. Reverted before landing; jumpBucketSaturationBoundary was added as regression coverage (see CI fix below for a follow-up to that same test).

Round 2 changes

A follow-up four-angle review (data layout/copies, cache hierarchy, threading headroom, Rust-style ownership discipline) produced 13 candidates, each tested individually with its own correctness gate and same-session A/B benchmark — including two real lessons learned along the way: session-level thermal/system drift is large enough to invert a conclusion if you compare against a stale baseline, and three separate "this should obviously help the compiler" ideas failed under measurement on this toolchain.

Kept (see the top-line table above for numbers):

• Fold accumLexCounts into lexCounts ({live, accum} value instead of a second map).
• Cache decode-path jumpLin/fertRatio per chain instead of rebuilding them from scratch on every decodeMarginalFromTables call. Caught and fixed a real regression along the way: the cached table is sized to the vocab as of last training, but vocab can grow afterward via getBestAlignment's auto-interning of new words, which would have indexed the cache out of bounds. Fixed with a bounds check; added getBestAlignmentWithPostTrainingNewWord as permanent coverage.
• Persistent lexDenomInv member instead of a fresh per-sweep allocation.
• Concurrency-contract documentation — investigated renaming getSrcSent/getTrgSent and hoisting vocab-size getters, found both were solving problems that didn't exist (the naming is an established codebase-wide convention; the getters have no hot-path cost), and instead documented the actual gap: an unstated invariant that vocab is frozen after buildCorpus, which is what makes several omp parallel for regions safe.
• Parallelize loglikelihoodForPairRange — the strongest result of this round: split into a serial vocab pre-conversion pass followed by a race-free omp parallel for reduction. A previously completely-serial function, not a refinement of already-parallel code.

Tried and reverted after honest measurement (each still bit-identical, just not beneficial): a chainAcc allocate-then-merge hoist (no measurable benefit), deduping repeated source words in the candidate loop (measured +24% regression — the O(n²) duplicate-detection scan costs more than the hash lookups it saves), a single-pass cumulative-weights fill (measured +8-16% regression with no clear mechanistic explanation), and hoisting sweep-scope consts for TBAA (genuinely inconclusive even at 5 reps/side).

Deferred, not attempted: an RNG swap (mt19937_64 → xoshiro256++/PCG64) and blocked/AD-LDA-style within-chain Gibbs sampling — both are output-changing and this repo has no gold word-alignment data to validate against; left as follow-up work for a pass that starts by building that infrastructure.

CI fix

The initial push (round 1) failed CI on Linux and both macOS runners — jumpBucketSaturationBoundary pinned an exact expected alignment from a scenario deliberately engineered to stress the jump-bucket saturation boundary (jumpWindow=2 on a 14-token sentence, 6 sweeps). That scenario is legitimately platform-sensitive: removing most of the jump-distance prior's constraining power lets tiny libm (exp/log) and std::uniform_real_distribution implementation differences compound over the sampling chain into a different (but equally valid) final alignment per platform — confirmed empirically, since Windows/MSVC, Linux/glibc, and macOS/libc++ each produced a different alignment for the same input. This is exactly the tradeoff SamplingUtils.h already documents. Fixed by replacing the non-portable exact-value check with a structural validity check plus a same-platform determinism cross-check; the three short-sentence assertions in the same test were unaffected (confirmed stable on all three platforms) and remain pinned.

Known gap left by this fix: the pinned value that was removed was, incidentally, the regression test for round 1's reverted "saturating jump-bucket running counter" bug — the most serious correctness bug found in this effort. The replacement structural/determinism check would not catch a reintroduction of that specific bug. Restoring equivalent coverage portably would mean extracting the candidate-scoring/jump-factor arithmetic into a pure, RNG-free function that can be unit-tested with hand-crafted inputs; left as follow-up work rather than attempted here.

Test plan

• Full test suite passes (116/116) on Windows/MSVC
• Verified compiling cleanly with g++ -Wall -Werror (matching Linux CI's flags) and with GCC 15 more broadly, since local development was Windows-only
• initializeChainParallelIsReproducible, loglikelihoodForPairRangeEdgeCasesAndDeterminism verify the new parallel regions are race-free
• getBestAlignmentWithPostTrainingNewWord guards the OOV bounds-check fix in the decode cache
• Benchmarked on a real corpus (7,952 eBible NT sentence pairs) for every kept/reverted decision — see the top-line table and per-change notes above
• CI (Linux/macOS/Windows matrix, wheels, sdist, nuget) — full matrix green after the portability fix

History is a single commit for review clarity.

🤖 Generated with Claude Code

[ddaspit]

How do these changes impact memory usage? Is the increased performance only realized when deterministic=true?

[johnml1135]

I'll check.

[johnml1135]

Independent re-benchmark: time and memory, before/after

Ran a fresh, independent before/after comparison to answer two questions the PR description doesn't directly cover: does end-to-end wall-clock actually improve outside the isolated hot-loop microbenchmark, and does the new per-chain caching (round 2's jumpLin/fertRatio cache, folded accumLexCounts) grow memory?

Method: git worktree at d0842c2a (parent of this branch's first commit here) as "before", current HEAD as "after". Copied the current DISABLED_realAblation harness into the "before" worktree so both builds run byte-identical benchmark code — only EflomalAlignmentModel.{h,cc} differs. Both built Release/x64/MSVC. Real corpus: eBible NT, English (BSB) ↔ Spanish (RV1909), sliced to the NT range via vref.txt (MAT–REV), 7,939 tokenized sentence pairs (same corpus family as this PR's 7,952; a couple of bracketed/omitted verses differ by translation pair). Machine: i7-12700, 12C/20T. Wall-clock from the harness's own ABLATION line; peak memory via Process.PeakWorkingSet64 polled every 100ms.

Time (full train(), not just the sampleSweep hot loop in isolation):

Config	Train before → after	Δ
samplers=1, threads=1 (default, 12 sweeps)	837 → 782 ms	−6.6%
samplers=1, threads=1, 48 sweeps (fertility-heavy, 3 reps)	3561 → 3294 ms	−7.5%
samplers=8, threads=4	1935 → 1676 ms	−13.4%
samplers=8, threads=20	1591 → 1439 ms	−9.6% (single rep, noisier)
samplers=20, threads=20 (stress)	2304 → 1955 ms	−15.1%

Consistently faster, but more modest than the −26.4% headline — that number is the isolated hash-reuse microbenchmark inside sampleSweep; end-to-end train() also includes corpus loading, chain init, and non-sampleSweep-dominated IBM1/HMM stages, which dilutes it. Both numbers are correct, just different denominators, worth being explicit about for anyone comparing.

One asymmetry: decode/alignment time (accumulateDecodeMarginal over test pairs) was flat-to-slightly-slower after, not faster. Expected — the round-2 decode cache targets computeTrainingAlignments, a different code path than what this harness's alignment phase exercises.

Memory — does it explode? No, it's slightly lower after, in every config tested:

Config	Peak WS before → after	Δ
samplers=1, threads=1	63.7 → 62.8 MB	−1.4%
samplers=1, threads=1, 48 sweeps	65.0 → 63.4 MB	−2.5%
samplers=8, threads=4	139.6 → 134.0 MB	−4.0%
samplers=8, threads=20	142.1 → 136.0 MB	−4.3%
samplers=20, threads=20 (stress)	274.3 → 260.3 MB	−5.1%

The new per-chain caches (decode jumpLin/fertRatio, folded lexCounts) are sized once from vocab, not reallocated per sweep, so they don't grow with iteration count. Meanwhile the round-1 buffer hoists (phi scratch, persistent lexDenomInv) and the flattened fertCounts/fertRatioSampled remove per-sweep/per-sentence allocations, cutting allocator churn and fragmentation — net effect is a small memory win, not the regression one might expect from "add more caches." Memory scales ~linearly with samplers (~13 MB/chain) in both versions, which is inherent to running N independent Gibbs chains, not something this PR changed.

Happy to share the corpus-prep/benchmark scripts if useful for turning this into a permanent (non-disabled) CI benchmark gate.