fix(omnia-training): serialize Burn RNG across the training test binary

Breadway commented

2026-10-20 09:26:50 +08:00

Owner

OmniAstra::new draws weights from Burn's process-global RNG. training_reduces_loss_and_is_seeded_deterministic reseeds that RNG and compares two runs, but four other tests in the same binary constructed models without taking RNG_LOCK (checkpoint_round_trip_and_resume_vs_weights_only, frozen_backbone_does_not_change_but_adapter_does, model_moved_from_inference_to_autodiff_trains_after_enable_grad, empty_batch_is_a_no_op_not_a_panic). A parallel test can consume draws between the two seeded runs.

Reproduction, release binary, default --test-threads=16, eight background sha256sum /dev/zero loops:

Unpatched (0710de2): 54 pass / 1 fail in 55 runs (~12 min). The failure is the existing assertion at training.rs:96: seeded runs diverged: 2.0957823 vs 2.1332889.
Patched (this branch): 47 pass / 0 fail in 47 runs (~12 min), same command and same load.

Commands:

cargo test --release -p omnia-training --test training --no-run
# then the built target/release/deps/training-* binary, looped:
<bin> --test-threads=16

The 1e-4 tolerance was not weakened. Other model-init tests already held RNG_LOCK; this extends that lock to the four that did not.

Refs OMNI-29

OmniAstra::new draws weights from Burn's process-global RNG. training_reduces_loss_and_is_seeded_deterministic reseeds that RNG and compares two runs, but four other tests in the same binary constructed models without taking RNG_LOCK (checkpoint_round_trip_and_resume_vs_weights_only, frozen_backbone_does_not_change_but_adapter_does, model_moved_from_inference_to_autodiff_trains_after_enable_grad, empty_batch_is_a_no_op_not_a_panic). A parallel test can consume draws between the two seeded runs. Reproduction, release binary, default --test-threads=16, eight background `sha256sum /dev/zero` loops: - Unpatched (0710de2): 54 pass / 1 fail in 55 runs (~12 min). The failure is the existing assertion at training.rs:96: seeded runs diverged: 2.0957823 vs 2.1332889. - Patched (this branch): 47 pass / 0 fail in 47 runs (~12 min), same command and same load. Commands: ``` cargo test --release -p omnia-training --test training --no-run # then the built target/release/deps/training-* binary, looped: <bin> --test-threads=16 ``` The 1e-4 tolerance was not weakened. Other model-init tests already held RNG_LOCK; this extends that lock to the four that did not. Refs OMNI-29

Breadway added 1 commit

2026-10-20 09:26:50 +08:00

fix(omnia-training): hold RNG_LOCK around every model init in the training tests 9561c6bd4c

OmniAstra::new draws weights from Burn's process-global RNG. The seeded
determinism test reseeds that RNG, but four other tests in the same binary
constructed models without taking RNG_LOCK, so a parallel test could consume
draws between the two seeded runs.

Under 8 sha256sum loaders, the unpatched release binary failed
training_reduces_loss_and_is_seeded_deterministic once in 55 runs
(2.0957823 vs 2.1332889, tolerance 1e-4). The same soak with this lock
passed 47/47. Tolerance is unchanged.

Refs OMNI-29

Breadway referenced this pull request

2026-10-20 09:26:59 +08:00

training_reduces_loss_and_is_seeded_deterministic failed once in a full workspace run #29

Breadway merged commit b77c421eba into main

2026-10-20 09:48:27 +08:00

Breadway referenced this pull request from a commit

2026-10-20 09:48:29 +08:00

Merge pull request #31: serialize Burn RNG across the training test binary

Rows
Columns

fix(omnia-training): serialize Burn RNG across the training test binary #31