Remediation: router-bias wiring, EOS packing, numerical stability, trainer hardening, tests

[supernavyl]

Summary

Comprehensive remediation of issues found during a multi-dimensional audit of the codebase. Six logical commits that can be reviewed independently:

1. fix(moe+attn) — vectorized dispatch, wired router bias, numerical stability
2. fix(tokenizer) — correct vocab sizing + EOS-aware encode_with_eos
3. refactor — move moda.py to open_mythos.experimental
4. feat(training) — harden FineWeb-Edu pretraining script
5. chore(deps) — reconcile pyproject / requirements.txt / training/requirements.txt
6. test — cover config validation, router bias update, vectorized dispatch

Headline fixes

1. router_bias was silently inert (correctness bug)

Upstream defines MoEFFN.router_bias as a buffer of zeros but no code path ever updated it — the advertised DeepSeek-V3 aux-loss-free load balancing was not happening. This PR adds:

• MoEFFN.expert_load buffer + bincount accumulation on each training forward.
• MoEFFN.update_bias(speed) with a sign-based update (magnitude bounded by speed regardless of imbalance, per DeepSeek-V3 Eq. 16). Flushes load after each call.
• OpenMythos.update_router_biases(speed=None, ddp=False) that walks every MoE layer and all-reduces load across FSDP ranks before updating.
• Trainer call site: model.update_router_biases(ddp=ddp) after every successful optimizer.step().

2. No EOS between packed documents (training-quality bug)

The FineWebEduDataset packer used encode() which does not append EOS, so concatenated documents flowed into each other with no boundary marker — the model learned spurious cross-document attention.

• New MythosTokenizer.encode_with_eos() appends EOS (with eos -> bos -> all_special_ids[0] -> None fallback).
• Dataset now uses it.

3. Vectorized MoE dispatch (performance)

The naive per-token loop was O(topk * n_experts) Python iterations per forward pass. Replaced with a stable argsort + bincount + index_add_ pipeline — O(n_experts) dense ops. Semantic equivalence verified against a reference loop implementation at 1e-5 tolerance (see tests/test_moe_router.py::test_vectorized_matches_naive_dispatch).

4. Numerical stability

• Softmax upcast to fp32 in GQAttention, MLAttention, and MoEFFN routers. bf16/fp16 softmax quantizes the tail and collapses attention to one-hot or uniform at long context.
• clamp_min(1e-9) on the gate-weight renormalization denominator so a fully underflowed set of topk scores does not divide by zero.
• LTI clamp tightened: log_dt + log_A clamp lower bound raised from -20 to -10. At -20, exp(-exp(-20)) rounds to exactly 1.0 in fp32 and breaks the strict ρ(A) < 1 invariant. At -10 it saturates at ~1 - 4.5e-5, safely below 1 in fp32.

5. vocab_size trap (correctness bug)

MythosTokenizer.vocab_size was self.tokenizer.vocab_size, which for HF tokenizers is the base vocab excluding added specials. Any added special token would index past the model's embedding matrix and trigger a CUDA device-side assert deep into pretraining.

• Now returns len(self.tokenizer), rounded up to vocab_multiple_of (default 128 for tensor-core-friendly widths).

6. Config validation at construction time

MythosConfig.__post_init__ now validates 12+ invariants (attn_type enum, dim/n_heads divisibility, GQA grouping, MLA rope even, loop_dim even, MoE sizing, ACT/dropout ranges, LoRA positivity, vocab_size/max_seq_len positivity). Bad configs fail at MythosConfig(...) instead of mid-step hours into a run.

7. Trainer hardening for long runs

• Per-microstep NaN/Inf loss guard (poisoning Adam moments is unrecoverable).
• Post-clip non-finite grad_norm guard.
• ShardedGradScaler on the fp16 path.
• Full RNG seeding (python/numpy/torch/cuda), persisted through the checkpoint.
• Checkpoint also persists scaler state, torch/cuda versions, seed.
• SIGTERM / SIGINT cooperative shutdown with final atomic save, then exit 130.
• loguru rotating file sink (100 MB / 7-day retention, gz compressed, per-rank).
• On-device loss accumulation (single .item() per log window).
• Directory fsync after atomic rename for power-loss durability.
• cfg.__post_init__() re-run after mutating vocab_size / max_seq_len.

8. moda.py moved to open_mythos.experimental

MoDA (Mixture-of-Depths Attention) is a parallel research line, not part of the canonical Prelude/Recurrent/Coda model. Subpackage boundary makes stability guarantees explicit. Commented-out smoke-test block deleted (dead scaffold, not a test).

9. Defensive freqs_cis[:T] slice in attention modules

Standalone callers (tests, ad-hoc scripts) no longer have to pre-slice the full precomputed RoPE buffer to match the current T. Fixes a pre-existing test_main.py::TestGQAttention::test_output_shape crash.

Test plan

• pytest test_main.py tests/ — 103 passed (0 failed).
• test_moe_router.py::test_vectorized_matches_naive_dispatch verifies the fast dispatch path matches a reference loop at 1e-5.
• test_moe_router.py::test_update_bias_magnitude_bounded_by_speed verifies the sign-based update even on 1M-imbalance load.
• test_config_validation.py covers 11 distinct __post_init__ guards.
• test_tokenizer.py::test_encode_ids_within_vocab — embedding-safety invariant.
• python -m py_compile on every changed file.
• Multi-node FSDP smoke test — not run (no cluster access on review author's side).
• Full pretraining run — not run (30B-token scale).

Breaking changes

• MythosConfig now validates in __post_init__. Configs that relied on silently invalid hyperparameter combinations will now raise ValueError at construction.
• MythosTokenizer.vocab_size may return a larger value than before (base vocab + added specials, rounded to 128). Any code that hardcoded the old value will need to rebuild its embedding matrix.
• open_mythos.moda import path → open_mythos.experimental.moda (or from open_mythos.experimental import MoDAModel). The old import path no longer works.

Honest caveats

• Tests run CPU-only. I did not have GPU access for a real FSDP run, so the trainer changes are validated by static review + unit tests on the model code, not by a live pretraining run.
• The vectorized MoE dispatch is verified equivalent at 1e-5 on small shapes in eval mode; at train time expert_load bookkeeping differs (intentionally, since the fast path accumulates it and the naive reference does not). That's tested separately in test_expert_load_accumulates_on_forward.