[manuelgloeckler]

Summary

These are all refactoring/merge regressions, not new bugs. Every fix here restores behaviour that the icml branch already had and that was lost when that code was ported into this repo. No new functionality, no intentional behaviour changes.

Found while investigating why updated_all_3_6_8_128_model_prior_rician_fixed_from_scratch trained far more spikily than the equivalent icml run.

1. DMRITokenizerPPP leaked each mask bit into its own decoder input

The port dropped the embed_cfgs override that DMRITokenizerPPP carries on the icml branch, so it inherited DMRITokenizer.embed_cfgs and its idx_tokens * model_mask masking.

That masking is correct for the base tokenizer but wrong here. embed_model_mask packs cfg token j into decoder position j, and loss_fn aligns logits[i] with model_mask[i] — so zeroing an inactive component's index token writes bit j into the exact position that must predict it. The model-selection head could read its target off its own input, and the shortcut broke at sample/log_prob time where the mask isn't known.

Flipping one mask bit at a time, measuring which decoder positions change:

	flip bit 0	flip bit 1	flip bit 2	flip bit 3
before	[0, 1] leak	[1, 2] leak	[2, 3] leak	[3, 4] leak
after	[1]	[2]	[3]	[4]
icml	[1]	[2]	[3]	[4]

Restoring the override reproduces icml exactly. Added a regression test that asserts the causality directly; confirmed it fails when the bug is reinstated (mask bit 0 leaks into position(s) [0, 1]).

2. continue_training could never find a checkpoint

checkpoint_dir was derived from output_dir, the per-launch Hydra directory (results/<name>/<timestamp>/0/checkpoints). Every relaunch gets a new timestamp, so continue_training: true looked in an empty directory, logged "No existing checkpoints found; starting fresh" and silently restarted from zero.

configure_environment already computes and logs output_super_dir for exactly this, but _main discarded it with log, output_dir, _ = .... Now used, restoring icml's results/<name>/checkpoints layout.

3. RNG restored with the wrong sharding

ArrayRestore was called without restore args, so the saved key came back with the sharding it was written with rather than the running job's. Passes the live key's sharding explicitly.

4. Config presets

Ports the four all_*_model_prior presets to the schema-v2 tree, with composition tests pinning the frozen architectures.

Two pre-existing assertions in test_config_composition.py were already failing on this branch before these changes (they expected run names without the _from_scratch suffix, and restart_optimizer is True); updated to match the configs.

updated_all_3_4_8_128_model_prior.yaml configures a 3/6/8 architecture despite its filename — kept for continuity with the icml preset it was ported from.

Validation

Verified the icml checkpoints restore cleanly into the current model before relaunching: 701/701 leaves (non-conv) and 733/733 (conv), no missing keys, no shape mismatches.

Both runs are now training from the icml checkpoints with the fix active — non-conv resumed at step 2,000,000, conv at 3,380,000.

Note the icml checkpoints predate the add_rician_noise RNG-split fix, so these runs resume across that distribution shift as well as the tokenizer fix.

Test suite was run before the changes were finalised; it should be re-run on a compute node rather than the login node.

🤖 Generated with Claude Code