feat(m3): stage 1 alignment training script + smoke (T3.2)
- src/tsmm/train/stage1.py: streams align.jsonl, freeze_llm, trains Encoder+Projector (AdamW lr=1e-4), bs=8/grad_accum=4/ctx=512/BF16, gradient checkpointing + enable_input_require_grads, loss = lm + lambda*InfoNCE (original vs perturbed TS representation), ckpt + tensorboard every 2000 steps. - scripts/train_stage1.sh: wrapper with design defaults. - src/tsmm/data/collator.py: handle JSON null (missing markers) -> NaN -> fill. - .gitignore: /checkpoints/ and data/*.jsonl. - Smoke verified: --max_steps 5 on 64 samples, finite loss, peak 5.52 GB (≤6GB). - 100 tests passing.
This commit is contained in:
Executable
+23
@@ -0,0 +1,23 @@
|
||||
#!/usr/bin/env bash
|
||||
# Stage ① alignment training (T3.2).
|
||||
# Loads align.jsonl (50万 target), freezes LLM, trains Encoder+Projector.
|
||||
set -euo pipefail
|
||||
|
||||
DATA="${1:-data/align.jsonl}"
|
||||
MAX_STEPS="${2:-10000}"
|
||||
CKPT_DIR="${3:-checkpoints/stage1}"
|
||||
|
||||
cd "$(dirname "$0")/.."
|
||||
.venv/bin/python -m tsmm.train.stage1 \
|
||||
--data "$DATA" \
|
||||
--ckpt_dir "$CKPT_DIR" \
|
||||
--max_steps "$MAX_STEPS" \
|
||||
--bs 8 \
|
||||
--grad_accum 4 \
|
||||
--ctx 512 \
|
||||
--lr 1e-4 \
|
||||
--lambda_contrast 0.1 \
|
||||
--perturb_std 0.1 \
|
||||
--log_every 20 \
|
||||
--ckpt_every 2000 \
|
||||
--tensorboard
|
||||
Reference in New Issue
Block a user