Skip to content

Commit 1d02939

Browse files
rahul-tuliclaude
rahul-tuli
and
claude
committed
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. - 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. - ValidationError is routed through parser.error (argparse usage, exit code 2). 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. Includes the D-PACE loss knobs from #736 (--per-position-loss-weight, --dpace-alpha) as typed DFlash fields, with the "dpace requires ce loss / alpha in (0,1]" check ported to a cross-group model_validator. 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 15-invocation matrix. Expected values live in one golden_flat_dict.json (a _baseline all-defaults dict + per-invocation deltas). 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>
1 parent 21033a7 commit 1d02939

8 files changed

Lines changed: 1630 additions & 629 deletions

File tree

‎scripts/train.py‎

Lines changed: 103 additions & 542 deletions
Large diffs are not rendered by default.

‎src/speculators/train/config.py‎

Lines changed: 854 additions & 0 deletions
Large diffs are not rendered by default.

‎src/speculators/utils/argparse_utils.py‎

Lines changed: 0 additions & 34 deletions
This file was deleted.

‎tests/unit/train/golden_cli_matrix.py‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
"""Shared CLI invocation matrix for the proof-of-equivalence golden test.
2+
3+
Each entry is ``(name, argv_tail)`` where ``argv_tail`` are the arguments passed
4+
after ``train.py`` (the required ``--verifier-name-or-path`` is prepended by the
5+
harness). The same matrix is used to (a) capture golden ``vars(args)`` fixtures
6+
from the pre-refactor parser and (b) assert the refactored config layer produces
7+
the identical flat working-dict. See ``test_golden_cli_equivalence.py``.
8+
9+
The matrix intentionally exercises every speculator type and the tricky
10+
draft-init / optimizer / loss / normalization branches -- these are the paths
11+
most likely to drift during the refactor.
12+
"""
13+
14+
# A dummy verifier path keeps ``parse_args()`` fully offline (no model download);
15+
# nothing in argument parsing touches the network.
16+
DUMMY_VERIFIER = "dummy-verifier"
17+
18+
INVOCATIONS: list[tuple[str, list[str]]] = [
19+
("eagle3_default", []),
20+
("dflash_default", ["--speculator-type", "dflash"]),
21+
("dspark_default", ["--speculator-type", "dspark"]),
22+
("peagle_default", ["--speculator-type", "peagle"]),
23+
("mtp_default", ["--speculator-type", "mtp"]),
24+
("from_pretrained", ["--from-pretrained", "/tmp/some/ckpt"]),
25+
("draft_config", ["--draft-config", "/tmp/some/decoder"]),
26+
(
27+
"dflash_sliding_window",
28+
[
29+
"--speculator-type",
30+
"dflash",
31+
"--sliding-window",
32+
"1024",
33+
"--sliding-window-indices",
34+
"0",
35+
"2",
36+
],
37+
),
38+
("muon_optimizer", ["--optimizer", "muon", "--muon-lr", "0.01"]),
39+
("compound_loss", ["--loss-fn", '{"ce": 0.1, "tv": 0.9}']),
40+
("dry_run", ["--dry-run"]),
41+
(
42+
"norm_toggles",
43+
["--no-norm-before-fc", "--no-norm-output", "--fc-norm"],
44+
),
45+
(
46+
"numeric_overrides",
47+
[
48+
"--epochs",
49+
"5",
50+
"--lr",
51+
"2e-4",
52+
"--num-layers",
53+
"3",
54+
"--draft-vocab-size",
55+
"32000",
56+
"--total-seq-len",
57+
"4096",
58+
"--checkpoint-freq",
59+
"0.5",
60+
],
61+
),
62+
(
63+
"dspark_heads",
64+
[
65+
"--speculator-type",
66+
"dspark",
67+
"--markov-rank",
68+
"128",
69+
"--markov-head-type",
70+
"gated",
71+
"--no-enable-confidence-head",
72+
"--confidence-head-alpha",
73+
"0.5",
74+
],
75+
),
76+
(
77+
"peagle_cod",
78+
[
79+
"--speculator-type",
80+
"peagle",
81+
"--num-depths",
82+
"4",
83+
"--down-sample-ratio",
84+
"0.6",
85+
"--down-sample-ratio-min",
86+
"0.1",
87+
],
88+
),
89+
]
90+
91+
92+
def full_argv(argv_tail: list[str]) -> list[str]:
93+
"""Prepend the program name and required verifier path to an argv tail."""
94+
return ["train.py", "--verifier-name-or-path", DUMMY_VERIFIER, *argv_tail]

‎tests/unit/train/golden_flat_dict.json‎

Lines changed: 168 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,168 @@
1+
{
2+
"_baseline": {
3+
"block_size": 8,
4+
"checkpoint_freq": 1.0,
5+
"confidence_head_alpha": 1.0,
6+
"confidence_head_with_markov": true,
7+
"d2t_path": null,
8+
"data_path": "./output",
9+
"deterministic_cuda": false,
10+
"dflash_decay_gamma": 4.0,
11+
"down_sample_ratio": 0.7,
12+
"down_sample_ratio_min": 0.2,
13+
"dpace_alpha": 0.5,
14+
"draft_arch": "llama",
15+
"draft_attn_impl": "simple_flex_attention",
16+
"draft_config": "",
17+
"draft_hidden_act": "silu",
18+
"draft_vocab_size": null,
19+
"dry_run": false,
20+
"embed_requires_grad": false,
21+
"enable_confidence_head": true,
22+
"epochs": 20,
23+
"fc_norm": false,
24+
"from_pretrained": "",
25+
"hidden_states_dtype": "bfloat16",
26+
"hidden_states_path": null,
27+
"legacy_data": false,
28+
"log_dir": "./logs",
29+
"log_freq": 1,
30+
"logger": "",
31+
"loss_fn": "kl_div",
32+
"lr": 0.0001,
33+
"markov_head_type": "vanilla",
34+
"markov_rank": 256,
35+
"mask_token_id": null,
36+
"max_anchors": 3072,
37+
"max_retries": 3,
38+
"muon_adjust_lr_fn": "match_rms_adamw",
39+
"muon_lr": 0.001,
40+
"muon_momentum": 0.95,
41+
"muon_ns_steps": 5,
42+
"muon_weight_decay": 0.1,
43+
"no_resume_from_checkpoint": false,
44+
"noise_std": 0.05,
45+
"norm_before_fc": true,
46+
"norm_before_residual": true,
47+
"norm_output": true,
48+
"num_depths": 8,
49+
"num_layers": 1,
50+
"num_speculative_steps": 3,
51+
"num_workers": 12,
52+
"on_generate": "delete",
53+
"on_missing": "generate",
54+
"optimizer": "muon",
55+
"per_position_loss_weight": "fixed-exp-decay",
56+
"prefetch_factor": 4,
57+
"request_timeout": 120,
58+
"run_name": null,
59+
"save_best": false,
60+
"save_path": "./output/checkpoints",
61+
"scheduler_num_cosine_cycles": 0.5,
62+
"scheduler_total_steps": null,
63+
"scheduler_type": "linear",
64+
"scheduler_warmup_steps": null,
65+
"seed": 42,
66+
"sliding_window": 2048,
67+
"sliding_window_indices": [],
68+
"sliding_window_non_causal": false,
69+
"speculator_type": "eagle3",
70+
"step_weight_beta": 0.6,
71+
"t2d_path": null,
72+
"target_layer_ids": null,
73+
"token_freq_path": null,
74+
"total_seq_len": 8192,
75+
"train_data_ratio": 0.9,
76+
"trust_remote_code": false,
77+
"ttt_step_loss_decay": 1.0,
78+
"ttt_steps": 3,
79+
"use_off_policy_tokens": false,
80+
"verifier_name_or_path": "dummy-verifier",
81+
"vllm_endpoint": "http://localhost:8000/v1",
82+
"weight_decay": 0.01
83+
},
84+
"compound_loss": {
85+
"loss_fn": "{\"ce\": 0.1, \"tv\": 0.9}"
86+
},
87+
"dflash_default": {
88+
"draft_arch": "qwen3",
89+
"norm_before_fc": false,
90+
"norm_output": false,
91+
"speculator_type": "dflash"
92+
},
93+
"dflash_sliding_window": {
94+
"draft_arch": "qwen3",
95+
"norm_before_fc": false,
96+
"norm_output": false,
97+
"sliding_window": 1024,
98+
"sliding_window_indices": [
99+
0,
100+
2
101+
],
102+
"speculator_type": "dflash"
103+
},
104+
"draft_config": {
105+
"draft_config": "/tmp/some/decoder"
106+
},
107+
"dry_run": {
108+
"dry_run": true
109+
},
110+
"dspark_default": {
111+
"draft_arch": "qwen3",
112+
"norm_before_fc": false,
113+
"norm_output": false,
114+
"speculator_type": "dspark"
115+
},
116+
"dspark_heads": {
117+
"confidence_head_alpha": 0.5,
118+
"draft_arch": "qwen3",
119+
"enable_confidence_head": false,
120+
"markov_head_type": "gated",
121+
"markov_rank": 128,
122+
"norm_before_fc": false,
123+
"norm_output": false,
124+
"speculator_type": "dspark"
125+
},
126+
"eagle3_default": {},
127+
"from_pretrained": {
128+
"from_pretrained": "/tmp/some/ckpt"
129+
},
130+
"mtp_default": {
131+
"draft_arch": "qwen3",
132+
"norm_before_fc": false,
133+
"norm_output": false,
134+
"speculator_type": "mtp"
135+
},
136+
"muon_optimizer": {
137+
"muon_lr": 0.01
138+
},
139+
"norm_toggles": {
140+
"fc_norm": true,
141+
"norm_before_fc": false,
142+
"norm_output": false
143+
},
144+
"numeric_overrides": {
145+
"checkpoint_freq": 0.5,
146+
"draft_vocab_size": 32000,
147+
"epochs": 5,
148+
"lr": 0.0002,
149+
"muon_lr": 0.002,
150+
"num_layers": 3,
151+
"total_seq_len": 4096
152+
},
153+
"peagle_cod": {
154+
"down_sample_ratio": 0.6,
155+
"down_sample_ratio_min": 0.1,
156+
"draft_arch": "qwen3",
157+
"norm_before_fc": false,
158+
"norm_output": false,
159+
"num_depths": 4,
160+
"speculator_type": "peagle"
161+
},
162+
"peagle_default": {
163+
"draft_arch": "qwen3",
164+
"norm_before_fc": false,
165+
"norm_output": false,
166+
"speculator_type": "peagle"
167+
}
168+
}

0 commit comments

Comments
 (0)