Skip to content

Commit 9e73541

Browse files
committed
refactor: migrate v0 prototype to Centaur 3-stage pipeline
Refactors the CRSM repository to move from v0 prototype scripts to the modular "Centaur" training pipeline. Changes: - Archival: Move `train_full_crsm.py`, `hyperparameter_sweep.py`, and `rl_value_head_finetuning.py` to `scripts/archive/v0_prototype/`. - Pipeline: Implement `stage_1_backbone.py`, `stage_2_dynamics.py`, and `stage_3_value_head.py` in `scripts/training/`. - Configuration: Migrate all baseline configs in `configs/` from JSON to YAML and create main `training_config.yaml`. - Evaluation: Add `benchmark_reasoning.py` and `verify_steering.py` to `scripts/eval/`. - Notebooks: Add cloud-ready notebooks in `notebooks/cloud_training/`. - Core: Update `crsm/cli.py` and `crsm/train.py` to support CLI arg parsing for config-based execution.
1 parent 32beccb commit 9e73541

37 files changed

Lines changed: 1311 additions & 2200 deletions

configs/baseline_170m.json

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

configs/baseline_170m.yaml

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
model:
2+
vocab_size: 50257
3+
d_model: 1024
4+
d_state: 64
5+
d_ffn: 4096
6+
num_layers: 24
7+
dropout: 0.1
8+
9+
reasoning:
10+
c_puct: 1.0
11+
n_simulations: 50
12+
temperature: 0.8
13+
injection_rate: 0.05
14+
15+
training:
16+
batch_size: 8
17+
seq_len: 1024
18+
lr: 4e-4
19+
backbone_epochs: 1
20+
finetune_epochs: 1
21+
finetune_lr: 1e-5
22+
grad_accum: 4
23+
use_amp: true
24+
25+
dynamics:
26+
dynamics_samples: 20000
27+
dynamics_epochs: 5
28+
dynamics_lr: 1e-3
29+
30+
data:
31+
data_dir: "data/text_corpus"
32+
traces_path: "data/train_traces.jsonl"
33+
34+
tokenizer: "gpt2"
35+
seed: 42

configs/baseline_27m.json

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

configs/baseline_27m.yaml

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
model:
2+
vocab_size: 50257
3+
d_model: 512
4+
d_state: 64
5+
d_ffn: 2048
6+
num_layers: 8
7+
dropout: 0.1
8+
9+
reasoning:
10+
c_puct: 1.0
11+
n_simulations: 25
12+
temperature: 0.8
13+
injection_rate: 0.05
14+
15+
training:
16+
batch_size: 32
17+
seq_len: 512
18+
lr: 6e-4
19+
backbone_epochs: 2
20+
finetune_epochs: 2
21+
finetune_lr: 2e-5
22+
grad_accum: 1
23+
use_amp: true
24+
25+
dynamics:
26+
dynamics_samples: 10000
27+
dynamics_epochs: 10
28+
dynamics_lr: 1e-3
29+
30+
data:
31+
data_dir: "data/text_corpus"
32+
traces_path: "data/train_traces.jsonl"
33+
34+
tokenizer: "gpt2"
35+
seed: 42

configs/baseline_350m.json

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

configs/baseline_350m.yaml

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
model:
2+
vocab_size: 50257
3+
d_model: 1536
4+
d_state: 64
5+
d_ffn: 6144
6+
num_layers: 32
7+
dropout: 0.1
8+
9+
reasoning:
10+
c_puct: 1.0
11+
n_simulations: 100
12+
temperature: 0.8
13+
injection_rate: 0.05
14+
15+
training:
16+
batch_size: 4
17+
seq_len: 2048
18+
lr: 3e-4
19+
backbone_epochs: 1
20+
finetune_epochs: 1
21+
finetune_lr: 5e-6
22+
grad_accum: 8
23+
use_amp: true
24+
25+
dynamics:
26+
dynamics_samples: 50000
27+
dynamics_epochs: 3
28+
dynamics_lr: 5e-4
29+
30+
data:
31+
data_dir: "data/text_corpus"
32+
traces_path: "data/train_traces.jsonl"
33+
34+
tokenizer: "gpt2"
35+
seed: 42

configs/small.json

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

configs/small.yaml

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
model:
2+
vocab_size: 1000
3+
d_model: 128
4+
d_state: 64
5+
d_ffn: 512
6+
num_layers: 2
7+
dropout: 0.1
8+
9+
reasoning:
10+
c_puct: 1.0
11+
n_simulations: 10
12+
temperature: 1.0
13+
injection_rate: 0.1
14+
15+
training:
16+
batch_size: 8
17+
seq_len: 32
18+
lr: 1e-3
19+
backbone_epochs: 2
20+
finetune_epochs: 2
21+
finetune_lr: 1e-4
22+
grad_accum: 1
23+
use_amp: false
24+
value_training:
25+
rollouts_per_prompt: 4
26+
epochs: 2
27+
lr: 1e-4
28+
29+
dynamics:
30+
dynamics_samples: 100
31+
dynamics_epochs: 2
32+
dynamics_lr: 1e-3
33+
34+
data:
35+
data_dir: "data/text_corpus"
36+
traces_path: "data/train_traces.jsonl"
37+
38+
tokenizer: null
39+
seed: 42
40+
41+
system:
42+
device: "cpu"
43+
seed: 42

configs/test_config.json

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

configs/test_config.yaml

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
model:
2+
vocab_size: 1000
3+
d_model: 32
4+
d_state: 16
5+
d_ffn: 64
6+
num_layers: 2
7+
dropout: 0.0
8+
9+
reasoning:
10+
c_puct: 1.0
11+
n_simulations: 2
12+
temperature: 1.0
13+
injection_rate: 0.1
14+
15+
training:
16+
batch_size: 2
17+
seq_len: 16
18+
lr: 1e-3
19+
backbone_epochs: 1
20+
finetune_epochs: 1
21+
finetune_lr: 1e-4
22+
grad_accum: 1
23+
use_amp: false
24+
25+
dynamics:
26+
dynamics_samples: 10
27+
dynamics_epochs: 1
28+
dynamics_lr: 1e-3
29+
30+
data:
31+
data_dir: "data/text_corpus"
32+
traces_path: "data/train_traces.jsonl"
33+
34+
tokenizer: null
35+
seed: 42

0 commit comments

Comments
 (0)