fix(omnia-model): use rel VOCAB in bench example lookups #30

Merged
Breadway merged 1 commit from fix/omni-28-bench-index into main 2026-10-14 09:55:20 +08:00
Owner

Closes #28

Root cause: the bench example built its diagnostic embedding tables with a hardcoded vocab of 40, but relation ids now reach 41 (REG_REG) since register relations were added (rel::VOCAB is 42). Any batch with n_reg > 0 put ids 40/41 into a dim-40 table, so burn-ndarray select panicked with assertion failed: index < dim at bench.rs:87. Model code was already correct (BiasTable uses rel::VOCAB); no model changes. The example now also builds its batch with the same BatchOpts the trainer uses (mask_ids from use_ids).

Throughput (machine: bos, AMD Ryzen AI 7 350, 16 threads, 30GB RAM):

  • cargo run --release -p omnia-model --example bench -- nano 32: iters print 205, 224, 224, 227 samples/s (last iter 227 samples/s)
  • cargo run --release -p omnia-model --example bench -- tiny 32: iters print 75, 78, 77, 82 samples/s (last iter 82 samples/s)

Checks:

  • cargo test --workspace --release: all pass, including new bench_path_smoke (0.02s) and the OMNI-29 seeded-determinism test on first run
  • cargo clippy -p omnia-model --all-targets -- -D warnings: clean; workspace-wide -D warnings still fails on 4 pre-existing lints in untouched omnia-dsl files (analysis.rs:209 needless_range_loop, ir.rs:271,283 derivable_impl, relabel.rs:67 ptr_arg)
  • cargo fmt --check: fails on pre-existing violations present on main (repo is not rustfmt-clean; e.g. untouched bench.rs:28); the new test file is rustfmt-clean and edited lines match surrounding style
  • No CI workflow exists in the repo (no .forgejo/workflows or .github/workflows), so the smoke run is a #[test] (crates/omnia-model/tests/bench_smoke.rs) instead of a workflow step
Closes #28 Root cause: the bench example built its diagnostic embedding tables with a hardcoded vocab of 40, but relation ids now reach 41 (REG_REG) since register relations were added (rel::VOCAB is 42). Any batch with n_reg > 0 put ids 40/41 into a dim-40 table, so burn-ndarray select panicked with `assertion failed: index < dim` at bench.rs:87. Model code was already correct (BiasTable uses rel::VOCAB); no model changes. The example now also builds its batch with the same BatchOpts the trainer uses (mask_ids from use_ids). Throughput (machine: bos, AMD Ryzen AI 7 350, 16 threads, 30GB RAM): - `cargo run --release -p omnia-model --example bench -- nano 32`: iters print 205, 224, 224, 227 samples/s (last iter 227 samples/s) - `cargo run --release -p omnia-model --example bench -- tiny 32`: iters print 75, 78, 77, 82 samples/s (last iter 82 samples/s) Checks: - `cargo test --workspace --release`: all pass, including new `bench_path_smoke` (0.02s) and the OMNI-29 seeded-determinism test on first run - `cargo clippy -p omnia-model --all-targets -- -D warnings`: clean; workspace-wide `-D warnings` still fails on 4 pre-existing lints in untouched omnia-dsl files (analysis.rs:209 needless_range_loop, ir.rs:271,283 derivable_impl, relabel.rs:67 ptr_arg) - `cargo fmt --check`: fails on pre-existing violations present on main (repo is not rustfmt-clean; e.g. untouched bench.rs:28); the new test file is rustfmt-clean and edited lines match surrounding style - No CI workflow exists in the repo (no .forgejo/workflows or .github/workflows), so the smoke run is a #[test] (`crates/omnia-model/tests/bench_smoke.rs`) instead of a workflow step
The bench example built diagnostic embedding tables with a hardcoded
vocab of 40, but relids now reach 41 (REG_REG) since register relations
were added. Any batch with n_reg > 0 panicked in burn-ndarray select
with assertion failed: index < dim. Look up omnia_observationrelVOCAB
instead, and build the batch with the same BatchOpts the trainer uses
(mask_ids from use_ids) so the two cannot drift apart.

Add a bench_path_smoke test that builds a small batch the way the
example does, asserts every relation id fits relVOCAB, and runs one
forward, backward and optimizer step.
Author
Owner

Review verdict: PASS on head 645e14aefd (base 1913272111). Root cause confirmed: batch.rs writes register rel ids 39/40/41 whenever n_reg>0 (measured max rel_ss=41, rel_as=40 on a real nano batch; rel::VOCAB=42, REG_REG=41) while the example built dim-40 diagnostic tables, so burn-ndarray select panicked with index < dim. Fix uses rel::VOCAB in all three spots and trainer-canonical BatchOpts (mask_ids: !cfg.use_ids, matching trainer.rs batch_opts). Model code already used rel::VOCAB (model.rs:97-99), no model change needed. Verified in isolated worktree: bench nano 32 prints throughput (iters 87/94/78/84 samples/s, rel lookup dims [32,68,68]) and bench tiny 32 prints throughput (iters 31/34/34/34, dims [32,70,70]); new bench_path_smoke passes; cargo test --workspace --release all green (14 suites ok, training 10/10 including OMNI-29 seeded test first run); clippy -p omnia-model --all-targets clean, workspace -D warnings fails only on 4 pre-existing omnia-dsl lints in untouched files (analysis.rs:209, ir.rs:271,283, relabel.rs:67, identical on base); repo has no CI workflows so checks were run manually; fmt drift is repo-wide pre-existing (base bench.rs:28 fails too), new test file is rustfmt-clean. Merged 645e14a to main as 0710de29c1.

Review verdict: PASS on head 645e14aefdd1da71a6281a981a886245369fd48a (base 1913272111594632ea5a9f18d0ea35e2ad0ddbf8). Root cause confirmed: batch.rs writes register rel ids 39/40/41 whenever n_reg>0 (measured max rel_ss=41, rel_as=40 on a real nano batch; rel::VOCAB=42, REG_REG=41) while the example built dim-40 diagnostic tables, so burn-ndarray select panicked with index < dim. Fix uses rel::VOCAB in all three spots and trainer-canonical BatchOpts (mask_ids: !cfg.use_ids, matching trainer.rs batch_opts). Model code already used rel::VOCAB (model.rs:97-99), no model change needed. Verified in isolated worktree: bench nano 32 prints throughput (iters 87/94/78/84 samples/s, rel lookup dims [32,68,68]) and bench tiny 32 prints throughput (iters 31/34/34/34, dims [32,70,70]); new bench_path_smoke passes; cargo test --workspace --release all green (14 suites ok, training 10/10 including OMNI-29 seeded test first run); clippy -p omnia-model --all-targets clean, workspace -D warnings fails only on 4 pre-existing omnia-dsl lints in untouched files (analysis.rs:209, ir.rs:271,283, relabel.rs:67, identical on base); repo has no CI workflows so checks were run manually; fmt drift is repo-wide pre-existing (base bench.rs:28 fails too), new test file is rustfmt-clean. Merged 645e14a to main as 0710de29c152e15efa2e60200f22b7a44fcc5772.
Sign in to join this conversation.
No reviewers
No labels
bug
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
Breadway/omniastra!30
No description provided.