refactor(train): config-file-first CLI generated from a typed pydantic schema

Make scripts/train.py config-file-first without changing what any run does and
without touching the model layer.

- New src/speculators/train/config.py is the single source of truth: ~75 tunables
  as typed fields on cohesive pydantic-settings groups (verifier/draft/data/
  generation/loss/optimizer/scheduler/trainer/logging + per-algo dflash/dspark/
  peagle/mtp). The argparse CLI is GENERATED from these fields, so adding an
  argument is one typed field -- train.py has no per-flag add_argument.
- --config run.yaml supplies any argument; precedence is CLI > YAML > coded
  defaults; --config is optional (no file == the old pure-CLI behaviour).
- --dump-config prints the fully-resolved grouped YAML (a valid --config file);
  the resolved run.yaml is also saved next to each checkpoint for reproducibility.
  It runs the draft-init contract first, so a scaffold is never emitted for a
  config that would fail to load.
- Provenance = the argparse.SUPPRESS'd namespace (not model_fields_set, which the
  nested partial-update machinery pollutes). A YAML leaf counts as provided only
  under its owning group block; a misplaced key is reported as unrecognised and
  kept out of provenance so it can't mis-fire the draft-init conflict check.
- Every bad input is a clean argparse error (exit code 2, no traceback): schema
  ValidationError and a broken --config file (missing/unreadable/malformed/
  non-mapping) alike; a missing required field names the flag to set.

Phase 1: the resolved config is flattened back to the vars(args)-shaped dict the
five SpeculatorModel classes consume via **kwargs, so the model layer and Trainer
are untouched; only main()'s TrainerConfig reads the typed groups. Pushing typed
configs into the model layer is deferred to a follow-up PR.

Rebased onto current main. Absorbs the flags added since the branch point and
ports them into the schema so behaviour is unchanged:
- #736 D-PACE: --per-position-loss-weight, --dpace-alpha as typed DFlash fields,
  with the "dpace requires ce loss / alpha in (0,1]" check as a cross-group
  model_validator.
- #651 scheduler: --scheduler-warmup-ratio, and --scheduler-type as a Literal
  (linear/cosine/none), threaded through the typed TrainerConfig build.
- #749 default sliding-window attention: --sliding-window-indices renamed to
  --full-attention-indices (inverted semantics), matching upstream.

The decoder-shaping flag set (used by the draft-init conflict check) is derived
from the schema via a field tag, so it is not a second place to edit when a
shaping knob is added.

Equivalence is guarded by tests/unit/train/test_golden_cli_equivalence.py: the
refactored flat working-dict is asserted byte-value-identical to the pre-refactor
parser across a 17-invocation matrix (expected values in one golden_flat_dict.json
-- a _baseline all-defaults dict + per-invocation deltas), and the generated CLI's
option-string set is frozen against the pre-refactor parser (golden_cli_flags.json)
so no flag can be silently renamed or dropped.

Removes src/speculators/utils/argparse_utils.py and its test: train.py was the
only caller of explicitly_provided_dests, which the SUPPRESS'd-namespace
provenance replaces.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Rahul-Tuli <rtuli@redhat.com>

Author: rahul-tuli

scripts/train.py

@@ -0,0 +1,91 @@
+[
+  "--block-size",
+  "--checkpoint-freq",
+  "--confidence-head-alpha",
+  "--confidence-head-with-markov",
+  "--d2t-path",
+  "--data-path",
+  "--deterministic-cuda",
+  "--dflash-decay-gamma",
+  "--down-sample-ratio",
+  "--down-sample-ratio-min",
+  "--dpace-alpha",
+  "--draft-arch",
+  "--draft-attn-impl",
+  "--draft-config",
+  "--draft-hidden-act",
+  "--draft-vocab-size",
+  "--dry-run",
+  "--embed-requires-grad",
+  "--enable-confidence-head",
+  "--epochs",
+  "--fc-norm",
+  "--from-pretrained",
+  "--full-attention-indices",
+  "--help",
+  "--hidden-states-dtype",
+  "--hidden-states-path",
+  "--legacy-data",
+  "--log-dir",
+  "--log-freq",
+  "--logger",
+  "--loss-fn",
+  "--lr",
+  "--markov-head-type",
+  "--markov-rank",
+  "--mask-token-id",
+  "--max-anchors",
+  "--max-retries",
+  "--muon-adjust-lr-fn",
+  "--muon-lr",
+  "--muon-momentum",
+  "--muon-ns-steps",
+  "--muon-weight-decay",
+  "--no-confidence-head-with-markov",
+  "--no-embed-requires-grad",
+  "--no-enable-confidence-head",
+  "--no-norm-before-fc",
+  "--no-norm-before-residual",
+  "--no-norm-output",
+  "--no-resume-from-checkpoint",
+  "--noise-std",
+  "--norm-before-fc",
+  "--norm-before-residual",
+  "--norm-output",
+  "--num-depths",
+  "--num-layers",
+  "--num-speculative-steps",
+  "--num-workers",
+  "--on-generate",
+  "--on-missing",
+  "--optimizer",
+  "--per-position-loss-weight",
+  "--prefetch-factor",
+  "--request-timeout",
+  "--run-name",
+  "--save-best",
+  "--save-path",
+  "--scheduler-num-cosine-cycles",
+  "--scheduler-total-steps",
+  "--scheduler-type",
+  "--scheduler-warmup-ratio",
+  "--scheduler-warmup-steps",
+  "--seed",
+  "--sliding-window",
+  "--sliding-window-non-causal",
+  "--speculator-type",
+  "--step-weight-beta",
+  "--t2d-path",
+  "--target-layer-ids",
+  "--token-freq-path",
+  "--total-seq-len",
+  "--train-data-ratio",
+  "--trust-remote-code",
+  "--ttt-step-loss-decay",
+  "--ttt-steps",
+  "--use-off-policy-tokens",
+  "--verifier-name-or-path",
+  "--vllm-endpoint",
+  "--weight-decay",
+  "-h"
+]

src/speculators/train/config.py

@@ -0,0 +1,133 @@
+"""Shared CLI invocation matrix for the proof-of-equivalence golden test.
+
+Each entry is ``(name, argv_tail)`` where ``argv_tail`` are the arguments passed
+after ``train.py`` (the required ``--verifier-name-or-path`` is prepended by the
+harness). The same matrix is used to (a) capture golden ``vars(args)`` fixtures
+from the pre-refactor parser and (b) assert the refactored config layer produces
+the identical flat working-dict. See ``test_golden_cli_equivalence.py``.
+
+The matrix intentionally exercises every speculator type and the tricky
+draft-init / optimizer / loss / normalization branches -- these are the paths
+most likely to drift during the refactor.
+"""
+
+# A dummy verifier path keeps ``parse_args()`` fully offline (no model download);
+# nothing in argument parsing touches the network.
+DUMMY_VERIFIER = "dummy-verifier"
+
+INVOCATIONS: list[tuple[str, list[str]]] = [
+    ("eagle3_default", []),
+    ("dflash_default", ["--speculator-type", "dflash"]),
+    ("dspark_default", ["--speculator-type", "dspark"]),
+    ("peagle_default", ["--speculator-type", "peagle"]),
+    ("mtp_default", ["--speculator-type", "mtp"]),
+    ("from_pretrained", ["--from-pretrained", "/tmp/some/ckpt"]),
+    ("draft_config", ["--draft-config", "/tmp/some/decoder"]),
+    (
+        "dflash_full_attention_indices",
+        [
+            "--speculator-type",
+            "dflash",
+            "--sliding-window",
+            "1024",
+            "--full-attention-indices",
+            "0",
+            "2",
+        ],
+    ),
+    (
+        "scheduler_warmup_ratio",
+        [
+            "--scheduler-type",
+            "cosine",
+            "--scheduler-warmup-ratio",
+            "0.1",
+        ],
+    ),
+    ("muon_optimizer", ["--optimizer", "muon", "--muon-lr", "0.01"]),
+    ("compound_loss", ["--loss-fn", '{"ce": 0.1, "tv": 0.9}']),
+    ("dry_run", ["--dry-run"]),
+    (
+        "norm_toggles",
+        ["--no-norm-before-fc", "--no-norm-output", "--fc-norm"],
+    ),
+    (
+        "numeric_overrides",
+        [
+            "--epochs",
+            "5",
+            "--lr",
+            "2e-4",
+            "--num-layers",
+            "3",
+            "--draft-vocab-size",
+            "32000",
+            "--total-seq-len",
+            "4096",
+            "--checkpoint-freq",
+            "0.5",
+        ],
+    ),
+    (
+        "dspark_heads",
+        [
+            "--speculator-type",
+            "dspark",
+            "--markov-rank",
+            "128",
+            "--markov-head-type",
+            "gated",
+            "--no-enable-confidence-head",
+            "--confidence-head-alpha",
+            "0.5",
+        ],
+    ),
+    (
+        "peagle_cod",
+        [
+            "--speculator-type",
+            "peagle",
+            "--num-depths",
+            "4",
+            "--down-sample-ratio",
+            "0.6",
+            "--down-sample-ratio-min",
+            "0.1",
+        ],
+    ),
+    # Kitchen sink: toggle every otherwise-untouched bool (store_true and
+    # BooleanOptionalAction --no- forms), the list field with real values, and a
+    # few non-default Literals -- so the golden proof locks their generated CLI
+    # *shape*, not just their defaults.
+    (
+        "kitchen_sink",
+        [
+            "--speculator-type",
+            "dspark",
+            "--save-best",
+            "--deterministic-cuda",
+            "--no-resume-from-checkpoint",
+            "--use-off-policy-tokens",
+            "--legacy-data",
+            "--trust-remote-code",
+            "--embed-requires-grad",
+            "--no-norm-before-residual",
+            "--no-confidence-head-with-markov",
+            "--target-layer-ids",
+            "2",
+            "3",
+            "5",
+            "--optimizer",
+            "adamw",
+            "--scheduler-type",
+            "cosine",
+            "--markov-head-type",
+            "rnn",
+        ],
+    ),
+]
+
+
+def full_argv(argv_tail: list[str]) -> list[str]:
+    """Prepend the program name and required verifier path to an argv tail."""
+    return ["train.py", "--verifier-name-or-path", DUMMY_VERIFIER, *argv_tail]