Commit Graph

5 Commits

Author SHA1 Message Date
e44e50ef78 data: full TinyStories + tokenized-id cache, val loss, CLI arch
- Corpus::load_cached: tokenize the (large) corpus ONCE, cache the id stream to
  <corpus>.u16.bin (gpt2 vocab 50257 < 65536 → exact u16), read cache on reruns.
- Corpus::split_tail: hold out a tail slice as a validation corpus.
- train(): take an optional valid corpus + eval_every/eval_batches; periodic
  deterministic val-loss eval that checkpoints the BEST val model; returns
  TrainResult{train_losses, evals, best_val}. T6 fixed-cadence path preserved.
- bin/train + bin/export_safetensors: read architecture (--heads/--head-dim/
  --layers/--ffn) + opt knobs (--steps/--batch/--seq/--max-lr/--val-tokens/
  --eval-every) from CLI flags; defaults reproduce the v0-baseline tiny config.
- gitignore the multi-GB corpus + *.u16.bin caches + *.ckpt (dash5-only).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-15 18:34:48 +08:00
7a4f69e430 model: add per-head QK-norm (Qwen3-compat) for xserv export
xserv's Qwen3 forward unconditionally applies per-head RMSNorm to Q and K
(q_norm/k_norm, shape [head_dim]) before RoPE — even gamma=1 is a real RMS
divide, not identity. xtrain never had this, so an exact xserv<->xtrain loop
was structurally impossible. Add it (reusing the 2D rms_norm op on the
[seq*nh, hd] head rows, inserted between reshape and rope to mirror
qwen3.rs's order) so the trained model is genuinely Qwen3-compatible.

params() inserts q_norm,k_norm after wv; num_params() counts them; the
PyTorch parity refs (parity.py / adamw_parity.py) + their name lists add the
same step so the dumps stay self-consistent.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-15 17:33:19 +08:00
5df1d4d57b test: resolve real_training corpus default via CARGO_MANIFEST_DIR
cargo runs tests with cwd = crate dir, so the bare relative default
data/tinystories-valid-3mb.txt didn't resolve. Anchor it to the repo root via
CARGO_MANIFEST_DIR so the test runs out of the box (still overridable with
XTRAIN_CORPUS).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-15 16:41:12 +08:00
2f8118fda9 test: tighten AdamW parity (f32 reference, 10 steps, allclose tol)
The loss trajectory already matched torch.optim.AdamW (worst relerr ~2e-4),
but the float64 torch reference diverged per-weight from the f32 GPU training
after the model memorised the batch (flat region: weights underdetermined,
loss identical). Fixes: run the torch reference in float32 (match engine
precision), shorten to 10 steps (weights still well-determined), and compare
final params with an allclose-style rtol+atol metric (a pure relative metric is
misleading on near-zero weights).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-15 16:34:18 +08:00
22b7434b23 test: AdamW PyTorch parity + checkpoint round-trip + real training
Acceptance tests (GPU-gated not(no_cuda), run on dash5):
- adamw_parity_dump.rs + adamw_parity.py: build the tiny model with fixed init,
  run N AdamW steps on a fixed batch, dump the loss trajectory + final params;
  the Python side rebuilds the identical model and runs torch.optim.AdamW with
  matched lr/wd/betas/eps, comparing trajectory + final params within rtol.
- checkpoint_roundtrip.rs: train a few steps, save, load into a fresh model with
  a DIFFERENT init, assert identical logits/loss on a fixed input.
- real_training.rs (#[ignore], --release): train on TinyStories for a bounded
  budget; assert loss drops substantially and print greedy samples.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-06-15 16:30:06 +08:00