Commit Graph
12 Commits
Author SHA1 Message Date
dela 5a7d949b01 Skip the cold-start SFT best ckpt and free CUDA cache after eval
Step 0 generate was writing a 5GB success=0 snapshot and leaving the
32GB card fragmented, so the next Adam step OOM'd after batch-32 eval.
Only promote _best after step 0 and empty_cache when eval returns.
2026-08-26 10:08:31 +08:00
dela 9652a9a7eb Save SFT last/best checkpoints during training and on interrupt
train_sft used to torch.save only after the full epoch budget, so Ctrl+C
dropped all translation weights. Write _last every --ckpt-every steps
and on KeyboardInterrupt; write _best when frozen eval (success, chrF)
improves; --resume continues from _last.
2026-08-26 10:00:11 +08:00
dela 071dfaf42c Sanitize SwanLab env before login so 0.9 nested project does not crash
OpenBayes sets SWANLAB_PROJECT as a string; swanlab>=0.9 parses that as
ProjectSettings and raises QuoteAwareEnvSettingsSource. Drop it, keep
SWANLAB_PROJ_NAME, and share run-id extraction with train_k3.
2026-08-26 10:00:06 +08:00
dela 9a4862a866 Resume the same SwanLab run from the id stored in the ckpt
swanlab.init always opened a new experiment on --resume. Save the run
id in the checkpoint and pass resume=True, id=... on the next start.
--swanlab-id overrides; --swanlab-new forces a fresh experiment.
2026-08-25 22:08:15 +08:00
dela 47c72e5bb8 Warn and reset chunk_index when resume changes batch or seq_len
Trying a larger micro-batch on an existing 0.5b run re-packs wiki
chunks; keep tokens/opt_step and restart the data cursor.
2026-08-25 22:00:28 +08:00
dela e7185cbf49 Pull OPUS-100 en-zh for SFT instead of a checked-in jsonl
train_sft --data opus-100 streams Helsinki-NLP/opus-100, writes both
directions, and skips frozen eval sentences. Runtime cache stays under
data/sft/ (gitignored).
2026-08-25 21:39:58 +08:00
dela 53d0f4b17a Cut 1B-run I/O: rarer SwanLab, ckpt, and generate
Every micro-step was hitting SwanLab, and every 100 steps wrote a 5GB
ckpt plus greedy decode. 0.5b now logs every 20, eval/held-out every 500,
saves _last every 1000, samples every 2000.
2026-08-25 20:59:29 +08:00
dela 8442f92c58 Keep LatentMoE routed bmm in activation dtype under bf16 autocast
Python float scales and fp32 expert weights promoted SiTU outputs to
fp32, so index_add mixed BFloat16 dest with Float source and crashed
the 0.5b run. Cast packed weights and gate scalars to z.dtype.
2026-08-25 20:22:51 +08:00
dela 49aede9cb2 Fit 0.5b training on 32GB: SDPA MLA, block checkpoint, chunked CE
Whole-mixer checkpoint plus T×T MLA scores OOM'd a 31GB GPU on backward.
Checkpoint each AttnRes block, run absorbed MLA through SDPA, and compute
CE in vocab chunks so [B,T,V] logits are never materialized.

--max-tokens is now the training budget; default --steps 2000 no longer
caps a 1B-token run at 250 optimizer steps.
2026-08-25 20:09:27 +08:00
dela 7a12f61de1 LatentMoE: K3 sigmoid routing and Switch aux/z-loss
Route with σ(W_r x), Top-k(s+b), then L1-normalize over the selected set.
Add Switch/GShard aux and router z-loss into train_k3 and train_sft.
Wiki parquet URLs honor HF_ENDPOINT for mirrored downloads.
2026-08-25 19:50:07 +08:00
dela d1da0816f2 LatentMoE: sparse permute-dispatch + padded bmm
Replace dense all-expert forward (16 experts × all tokens) with
permute-dispatch: sort token-expert pairs by expert id, pad to
[R, C, ℓ] (C = max tokens per expert), run 3 bmm calls for the
batched SiTU-GLU activation, then scatter-add weighted results back.

Routed expert FLOPs drop from R·N to R·C (C ≈ N·k/R under uniform
routing). SiTU parameter structure unchanged; checkpoint compatible.

Tests: sparse-vs-dense fwd/bwd equivalence, unselected expert zero
grad, last_capacity tracking.
2026-08-25 17:47:44 +08:00
dela 584f7e9e73 Initial K3 snapshot: 0.5B KDA/MLA/MoE train path
Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias
backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.
2026-08-25 14:43:17 +08:00