3 files changed, 73 insertions, 0 deletions
diff --git a/srm/config/arch/hrm_v1.yaml b/srm/config/arch/hrm_v1.yaml
new file mode 100644
index 0000000..a5646b8
--- /dev/null
+++ b/srm/config/arch/hrm_v1.yaml
@@ -0,0 +1,21 @@
+name: hrm.hrm_act_v1@HierarchicalReasoningModel_ACTV1
+loss:
+  name: losses@ACTLossHead
+  loss_type: stablemax_cross_entropy
+
+halt_exploration_prob: 0.1
+halt_max_steps: 16
+
+H_cycles: 2
+L_cycles: 2
+
+H_layers: 4
+L_layers: 4
+
+hidden_size: 512
+num_heads: 8  # min(2, hidden_size // 64)
+expansion: 4
+
+puzzle_emb_ndim: ${.hidden_size}
+
+pos_encodings: rope
diff --git a/srm/config/arch/srm_v1.yaml b/srm/config/arch/srm_v1.yaml
new file mode 100644
index 0000000..1b5ece5
--- /dev/null
+++ b/srm/config/arch/srm_v1.yaml
@@ -0,0 +1,21 @@
+name: srm.srm_aol_v1@StableRecursionModel_ACTV1
+loss:
+  name: losses@ACTLossHead
+  loss_type: stablemax_cross_entropy
+
+halt_exploration_prob: 0.1
+halt_max_steps: 16
+
+# SRM-specific
+n_iters: 12           # joint micro-steps per ACT step (≈ HRM's H_cycles·L_cycles+H_cycles = 6 with deeper schedule)
+n_aol_layers: 2       # depth of AOL ψ block (channel + token mix per layer)
+kappa: 0.9            # contraction factor: per-step Lip_P ≤ (1-α)+α·κ = κ
+eta: 1.0              # weighting of L block in P-norm (1.0 = symmetric)
+alpha: 1.0            # damping (1.0 = full step)
+
+hidden_size: 512
+puzzle_emb_ndim: ${.hidden_size}
+
+# Unused (kept so pretrain.py's __pydantic_extra__ doesn't break)
+# pretrain.py's create_model() passes some fields HRM expects; Pydantic 'ignore'
+# (default) drops them silently.
diff --git a/srm/config/cfg_pretrain.yaml b/srm/config/cfg_pretrain.yaml
new file mode 100644
index 0000000..51c55a0
--- /dev/null
+++ b/srm/config/cfg_pretrain.yaml
@@ -0,0 +1,31 @@
+# ARC training config
+
+defaults:
+  - arch: hrm_v1
+  - _self_
+
+hydra:
+  output_subdir: null
+
+# Data path
+data_path: data/arc-aug-1000
+
+# Hyperparams - Training
+global_batch_size: 768
+
+epochs: 100000
+eval_interval: 10000
+checkpoint_every_eval: True
+
+lr: 1e-4
+lr_min_ratio: 1.0
+lr_warmup_steps: 2000
+
+# Standard hyperparameter settings for LM, as used in Llama
+beta1: 0.9
+beta2: 0.95
+weight_decay: 0.1
+puzzle_emb_weight_decay: 0.1
+
+# Hyperparams - Puzzle embeddings training
+puzzle_emb_lr: 1e-2