Snapshot before JURECA access ends

Author: rkstgr

.gitignore

@@ -46,6 +46,10 @@ data/
 *.log
 logs/
 
+# JURECA-specific scripts (kept locally; backed up separately)
+*.slurm
+core.*
+
 # OS
 .DS_Store
 Thumbs.db

scripts/analyze.py

@@ -238,6 +238,31 @@ def build_parser() -> argparse.ArgumentParser:
                         help="Truncated LeCun Normal initialization (HRM-style)")
     arith.add_argument("--use-adam-atan2", action="store_true",
                         help="Adam-atan2 scale-invariant optimizer (HRM-style)")
+    # H/L split (Ablation 3)
+    arith.add_argument("--use-hl-split", action="store_true",
+                        help="H/L module split: slow H layers + fast L inner loop")
+    arith.add_argument("--n-h-layers", type=int, default=2,
+                        help="Number of H (slow) layers in H/L split")
+    arith.add_argument("--t-inner", type=int, default=2,
+                        help="L inner loop steps in H/L split")
+    arith.add_argument("--t-outer", type=int, default=2,
+                        help="H outer loop steps in H/L split")
+    # Q-learning ACT (Ablation 5)
+    arith.add_argument("--use-q-act", action="store_true",
+                        help="Q-learning ACT: replace exit gate with Q-head")
+    arith.add_argument("--q-weight", type=float, default=0.1,
+                        help="Weight for Q-learning loss")
+    arith.add_argument("--q-gamma", type=float, default=0.99,
+                        help="Discount factor for Q-learning TD targets")
+    # MoE-layer recurrence (Ablation 6)
+    arith.add_argument("--use-moe-recurrence", action="store_true",
+                        help="MoE-layer recurrence: route to expert layers per step")
+    arith.add_argument("--num-expert-layers", type=int, default=4,
+                        help="Number of expert TransformerBlocks")
+    arith.add_argument("--moe-top-k", type=int, default=2,
+                        help="Experts activated per token")
+    arith.add_argument("--moe-load-balance-weight", type=float, default=0.01,
+                        help="Load balancing loss weight for MoE")
     arith.add_argument("--extrap-eval", nargs="+", type=int, default=None,
                         help="Test at these T values after training (extrapolation)")
     arith.add_argument("--n-eval", type=int, default=500)
@@ -718,6 +743,17 @@ def run_arith(args) -> None:
         use_postnorm=args.use_postnorm,
         use_lecun_init=args.use_lecun_init,
         use_adam_atan2=args.use_adam_atan2,
+        use_hl_split=args.use_hl_split,
+        n_h_layers=args.n_h_layers,
+        t_inner=args.t_inner,
+        t_outer=args.t_outer,
+        use_q_act=args.use_q_act,
+        q_weight=args.q_weight,
+        q_gamma=args.q_gamma,
+        use_moe_recurrence=args.use_moe_recurrence,
+        num_expert_layers=args.num_expert_layers,
+        moe_top_k=args.moe_top_k,
+        moe_load_balance_weight=args.moe_load_balance_weight,
         extrap_eval_steps=args.extrap_eval,
         n_eval=args.n_eval,
         eval_every=args.eval_every,

scripts/eval_checkpoints.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python
+"""Evaluate LoopLM checkpoints on lm-eval benchmarks.
+
+Usage:
+    uv run scripts/eval_checkpoints.py \
+        --checkpoint /path/to/step_0008000.pt \
+        --tasks hellaswag,arc_easy \
+        --limit 500
+"""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from src.model.config import LoopLMConfig
+from src.model.looplm import LoopLM
+from src.inference.lm_eval_wrapper import LoopLMLM, run_eval, _extract_acc
+from transformers import AutoTokenizer
+
+
+def main():
+    p = argparse.ArgumentParser(description="Evaluate LoopLM checkpoint")
+    p.add_argument("--checkpoint", required=True, help="Path to .pt checkpoint")
+    p.add_argument("--tasks", default="hellaswag,arc_easy")
+    p.add_argument("--limit", type=int, default=None)
+    p.add_argument("--tokenizer-id", default="HuggingFaceTB/SmolLM2-135M")
+    p.add_argument("--batch-size", type=int, default=4)
+    args = p.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    tasks = [t.strip() for t in args.tasks.split(",") if t.strip()]
+
+    print(f"Loading checkpoint: {args.checkpoint}")
+    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
+
+    model_cfg = LoopLMConfig(**ckpt["model_config"])
+    print(f"Model: L={model_cfg.num_layers}, T={model_cfg.max_recurrent_steps}, "
+          f"h={model_cfg.hidden_size} ({model_cfg.num_parameters()/1e6:.1f}M params)")
+
+    # Detect if model was trained with Q-ACT by checking for q_head in state dict
+    state_dict = ckpt["model_state_dict"]
+    use_q_act = any("q_head" in k for k in state_dict)
+
+    model = LoopLM(model_cfg, use_q_act=use_q_act).to(device)
+    model.load_state_dict(state_dict)
+    model.eval()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id)
+    max_T = model_cfg.max_recurrent_steps
+
+    print(f"\nEvaluating tasks={tasks}, T=1..{max_T}, limit={args.limit}")
+    print(f"{'T':>3} | " + " | ".join(f"{t:^14}" for t in tasks))
+    print("-" * (6 + 17 * len(tasks)))
+
+    all_results = {}
+    for T in range(1, max_T + 1):
+        wrapper = LoopLMLM(model, tokenizer, device, num_steps=T, batch_size=args.batch_size)
+        raw = run_eval(wrapper, tasks, args.limit)
+        row_parts = []
+        for task in tasks:
+            acc = _extract_acc(raw, task)
+            if acc is not None:
+                all_results[f"{task}/T{T}"] = acc
+                row_parts.append(f"{acc:^14.4f}")
+            else:
+                row_parts.append(f"{'N/A':^14}")
+        print(f"{T:>3} | " + " | ".join(row_parts))
+
+    # Save results
+    out_path = Path(args.checkpoint).parent.parent / "eval_results.json"
+    with open(out_path, "w") as f:
+        json.dump(all_results, f, indent=2)
+    print(f"\nResults saved to {out_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -35,6 +35,15 @@ def parse_args():
     # Model
     p.add_argument("--model-config", default="small",
                    choices=["small", "ouro_1_4b", "ouro_2_6b"])
+    # Model config overrides (applied after loading preset)
+    p.add_argument("--hidden-size", type=int, default=None)
+    p.add_argument("--num-layers", type=int, default=None)
+    p.add_argument("--num-heads", type=int, default=None)
+    p.add_argument("--intermediate-size", type=int, default=None)
+    # Ablation flags
+    p.add_argument("--use-q-act", action="store_true")
+    p.add_argument("--q-weight", type=float, default=0.1)
+    p.add_argument("--q-gamma", type=float, default=0.99)
 
     # Data
     p.add_argument("--dataset", default="wikitext")
@@ -91,14 +100,25 @@ def parse_args():
     return p.parse_args()
 
 
-def build_model_config(name: str, seq_len: int) -> LoopLMConfig:
+def build_model_config(args) -> LoopLMConfig:
     configs = {
         "small": LoopLMConfig.small(),
         "ouro_1_4b": LoopLMConfig.ouro_1_4b(),
         "ouro_2_6b": LoopLMConfig.ouro_2_6b(),
     }
-    cfg = configs[name]
-    cfg.max_seq_len = seq_len
+    cfg = configs[args.model_config]
+    cfg.max_seq_len = args.seq_len
+    # Apply overrides
+    if args.hidden_size is not None:
+        cfg.hidden_size = args.hidden_size
+    if args.num_layers is not None:
+        cfg.num_layers = args.num_layers
+    if args.num_heads is not None:
+        cfg.num_heads = args.num_heads
+    if args.intermediate_size is not None:
+        cfg.intermediate_size = args.intermediate_size
+    if args.num_recurrent_steps is not None:
+        cfg.max_recurrent_steps = args.num_recurrent_steps
     return cfg
 
 
@@ -188,7 +208,7 @@ def main():
         rank = dist.get_rank()
         world_size = dist.get_world_size()
 
-    model_cfg = build_model_config(args.model_config, args.seq_len)
+    model_cfg = build_model_config(args)
     eval_tasks = [t.strip() for t in args.eval_tasks.split(",") if t.strip()]
 
     # Build multi-stage list (each non-zero step value enables that stage)
@@ -222,12 +242,25 @@ def main():
         eval_tasks=eval_tasks,
         eval_limit=args.eval_limit,
         tokenizer_id=args.tokenizer_id,
+        q_weight=args.q_weight,
+        q_gamma=args.q_gamma,
     )
 
+    # Model kwargs for ablation flags
+    model_kwargs = {}
+    if args.use_q_act:
+        model_kwargs["use_q_act"] = True
+
     n_params = model_cfg.num_parameters()
     eff_batch_tokens = args.batch_size * args.grad_accum_steps * args.seq_len * world_size
     if rank == 0:
-        print(f"Model: {args.model_config}  ({n_params/1e6:.1f}M params)")
+        overrides = []
+        if args.hidden_size: overrides.append(f"h={args.hidden_size}")
+        if args.num_layers: overrides.append(f"L={args.num_layers}")
+        if args.num_recurrent_steps: overrides.append(f"T={args.num_recurrent_steps}")
+        if args.use_q_act: overrides.append("Q-ACT")
+        override_str = f" [{', '.join(overrides)}]" if overrides else ""
+        print(f"Model: {args.model_config}{override_str}  ({n_params/1e6:.1f}M params)")
         if distributed:
             print(f"Distributed: {world_size} GPUs (DDP)")
         if stages:
@@ -262,7 +295,7 @@ def main():
             max_chunks=args.max_chunks,
         )
 
-    trainer = Trainer(model_cfg, trainer_cfg)
+    trainer = Trainer(model_cfg, trainer_cfg, model_kwargs=model_kwargs)
     num_steps = args.num_recurrent_steps or model_cfg.max_recurrent_steps
 
     if args.gradient_checkpointing: