Decompose scripts/train.py:main() into cohesive, testable units

[amazloumi]

What

main() in scripts/train.py is ~1000 lines that mix many responsibilities — config load/validate, distributed setup, model/optimizer/scheduler build, checkpoint & resume, the training loop, teardown, and resilience — in a single scope. Break it into named, single-responsibility units called by a thin main() orchestrator.

Benefits:

• Readability / maintainability — each phase is understandable and changeable in isolation; the shared scaffold lives in one place (smaller bug-fix surface).
• Unit-testability — extracting the loop/teardown into run_training_loop(...) (or a small Trainer) creates a seam to inject collaborators (e.g. a spy CheckpointManager), so save/exit decisions can be unit-tested fast and in default CI. Today the loop is only reachable via the opt-in --e2e tier — e.g. the Fix clean runs not always saving on last step #129 final-checkpoint fix has no per-PR coverage.
• Reuse — clean phases are reusable by other entry points (eval, smoke harnesses) without copy-paste.

Suggested units (illustrative): load_and_validate_config, setup_distributed, build_model_and_optim, setup_checkpointing, and run_training_loop(...) holding the loop + teardown, with the PP/VLM/standard model-call as an injected step function.

Scope

• scripts/train.py — main() becomes a thin orchestrator; phases move to functions.
• Optional new home: kempnerforge/training/loop.py (and siblings).
• tests/unit/ — new unit tests enabled by the seam.
• No config sections change; CLI/UX unchanged.

Non-goals

• Do not intend to split into per-model scripts (train_dense.py / train_moe.py / train_vlm.py). That duplicates the shared scaffold and contradicts the registry/config-driven design ("swap arch without touching scripts/train.py"; loop variants differ only at the model-call site). Keep one entry point; model differences stay behind the registry + a thin step seam.

Backward compatibility

Behavior-preserving refactor — no functional change. Same CLI (uv run torchrun … scripts/train.py configs/train/.toml …), same single entry point, same config/registry dispatch. All branches must stay equivalent: PP / VLM / standard step; NaN-rollback, NaN-continue, NCCL-health, graceful-shutdown exits; phase scheduling; metrics/profiler/eval; and the save → wait order. tests/e2e/test_training_e2e.py must pass unchanged as the safety net.

Acceptance criteria

• main() reduced to a thin orchestrator; each phase a named unit; pyright + ruff clean.
• tests/e2e/ --e2e passes unchanged (proves no behavior change).
• run_training_loop(...) unit-testable fro example with a fake CheckpointManager (no real DCP/subprocess);