Upload modular/code/train.py with huggingface_hub

modular/code/train.py

@@ -0,0 +1,274 @@
+"""
+Modular arithmetic SoRL training — matches Nanda et al. (2023) architecture.
+
+Architecture: 1L / 4H / 128d / d_mlp=512  (Nanda's exact setup)
+Dataset:      all p²=12769 pairs, 30% train fixed (seed=42)
+
+Usage:
+    python -m arithmetic.modular.training.train --mode baseline
+    python -m arithmetic.modular.training.train --mode sorl --K 1 --abs_vocab 30
+"""
+import sys
+import json
+import argparse
+from dataclasses import dataclass, asdict
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, TensorDataset
+import matplotlib.pyplot as plt
+
+try:
+    import wandb
+    WANDB_AVAILABLE = True
+except ImportError:
+    WANDB_AVAILABLE = False
+
+WANDB_PROJECT = "sorl-modular"
+WANDB_ENTITY  = "nlp_and_interpretability"
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[3]))
+
+from transformers import Qwen3Config
+from sorl.sorl_wrapper import SorlModelWrapper
+from sorl.sorl_trainer import sorl_search, SoRLLoss
+from arithmetic.modular.data.modular import (
+    get_train_set, get_eval_set,
+    VOCAB_SIZE, PAD, PROMPT_LEN, P,
+)
+from arithmetic.modular.training.evaluate import ModularEvaluator
+
+
+@dataclass
+class ModularConfig:
+    # Task
+    p:    int  = P
+    mode: str  = "sorl"   # "baseline" or "sorl"
+
+    # Architecture — Nanda's exact setup
+    n_layer: int = 1
+    n_head:  int = 4
+    n_embd:  int = 128
+    d_mlp:   int = 512
+
+    # SoRL
+    K:               int   = 1
+    abs_vocab:       int   = 30
+    alpha_info_gain: float = 10.0
+    alpha_abs:       float = 0.1
+    alpha_soft_zipf: float = 1.0
+    n_rollouts:      int   = 2
+
+    # Optimizer
+    lr:           float = 1e-3
+    weight_decay: float = 0.1
+    num_epochs:   int   = 500
+    batch_size:   int   = 256
+
+    # Eval / logging
+    eval_every: int  = 50
+    log_every:  int  = 10
+    device:     str  = "cuda"
+    seed:       int  = 42
+    job_name:   str  = ""
+    out_dir:    str  = ""     # if empty, defaults to arithmetic/modular/runs/<job_name>
+    no_wandb:   bool = False
+
+
+def make_model(cfg: ModularConfig) -> SorlModelWrapper:
+    config = Qwen3Config(
+        hidden_size=cfg.n_embd,
+        num_hidden_layers=cfg.n_layer,
+        num_attention_heads=cfg.n_head,
+        num_key_value_heads=cfg.n_head,
+        intermediate_size=cfg.d_mlp,
+        vocab_size=VOCAB_SIZE,
+        max_position_embeddings=32,
+    )
+    abs_v = cfg.abs_vocab if cfg.mode == "sorl" else 1
+    return SorlModelWrapper.from_scratch(config, [VOCAB_SIZE, abs_v], PAD)
+
+
+def make_loader(examples, batch_size: int, shuffle: bool = True):
+    tokens = torch.tensor([ex.tokens for ex in examples], dtype=torch.long)
+    bs = len(examples) if batch_size == 0 else batch_size
+    return DataLoader(TensorDataset(tokens), batch_size=bs, shuffle=shuffle)
+
+
+def compute_base_traj_loss(model, ids: torch.Tensor, attn: torch.Tensor) -> torch.Tensor:
+    """CE on result token only, no abstract tokens."""
+    out    = model(input_ids=ids, attention_mask=attn, memory_span_abs=512, memory_span_traj=512)
+    base_v = int(model.vocab_sizes[0].item())
+    return nn.CrossEntropyLoss()(out.logits[:, PROMPT_LEN - 1, :base_v], ids[:, PROMPT_LEN])
+
+
+def save_curves(history: dict, out_dir: Path):
+    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
+
+    axes[0].plot(history["epoch"], history["train_loss"], label="train loss")
+    axes[0].set_xlabel("epoch"); axes[0].set_ylabel("loss"); axes[0].set_title("Training Loss")
+    axes[0].legend()
+
+    axes[1].plot(history["eval_epoch"], history["test_acc"], color="green", label="test acc")
+    axes[1].set_xlabel("epoch"); axes[1].set_ylabel("accuracy"); axes[1].set_title("Test Accuracy")
+    axes[1].set_ylim(0, 1); axes[1].legend()
+
+    plt.tight_layout()
+    plt.savefig(out_dir / "curves.png", dpi=100)
+    plt.close()
+
+
+def train(cfg: ModularConfig):
+    torch.manual_seed(cfg.seed)
+    device = torch.device(cfg.device)
+
+    out_dir = Path(cfg.out_dir) if cfg.out_dir else (
+        Path(__file__).resolve().parents[2] / "runs" / (cfg.job_name or f"{cfg.mode}_K{cfg.K}")
+    )
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    train_examples = get_train_set(p=cfg.p, seed=cfg.seed)
+    test_examples      = get_eval_set(p=cfg.p, seed=cfg.seed)
+    loader = make_loader(train_examples, cfg.batch_size)
+
+    model     = make_model(cfg).to(device)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
+    evaluator = ModularEvaluator(model, device=cfg.device, K=cfg.K)
+
+    sorl_loss_fn = SoRLLoss(
+        abs_vocab_size=model.vocab_sizes[-1],
+        zipf_alpha=cfg.alpha_soft_zipf,
+    ).to(device) if cfg.mode == "sorl" else None
+
+    history = {"epoch": [], "train_loss": [], "eval_epoch": [], "test_acc": []}
+    best_acc = 0.0
+
+    use_wandb = WANDB_AVAILABLE and not cfg.no_wandb
+    if use_wandb:
+        wandb.init(
+            project=WANDB_PROJECT, entity=WANDB_ENTITY,
+            name=cfg.job_name or f"{cfg.mode}_K{cfg.K}_abs{cfg.abs_vocab}",
+            config=asdict(cfg),
+        )
+
+    print(f"Training {cfg.mode} | p={cfg.p} | {len(train_examples)} train | {len(test_examples)} test")
+    print(f"Model: {cfg.n_layer}L/{cfg.n_head}H/{cfg.n_embd}d | K={cfg.K} | abs_vocab={cfg.abs_vocab}")
+    print(f"Output: {out_dir}")
+
+    for epoch in range(1, cfg.num_epochs + 1):
+        model.train()
+        epoch_loss = 0.0
+
+        for (ids,) in loader:
+            ids  = ids.to(device)
+            attn = torch.ones_like(ids)
+            pl   = torch.full((ids.shape[0],), PROMPT_LEN, dtype=torch.long, device=device)
+            optimizer.zero_grad()
+
+            if cfg.mode == "baseline":
+                out    = model(input_ids=ids, attention_mask=attn, memory_span_abs=512, memory_span_traj=512)
+                base_v = int(model.vocab_sizes[0].item())
+                loss   = nn.CrossEntropyLoss()(out.logits[:, PROMPT_LEN - 1, :base_v], ids[:, PROMPT_LEN])
+            else:
+                btl = compute_base_traj_loss(model, ids, attn)
+                with torch.no_grad():
+                    best_data, _, _, exp_mask, exp_pl = sorl_search(
+                        model, ids, attn, pl, PAD,
+                        n=cfg.n_rollouts, K=cfg.K,
+                        max_iterations=2, memory_span_abs=512, memory_span_traj=512,
+                    )
+                info_loss, abs_loss, zipf_loss = sorl_loss_fn(
+                    best_data, model, btl.detach(), exp_mask, 512, 512, prompt_len=exp_pl,
+                )
+                loss = (btl
+                        + cfg.alpha_info_gain * info_loss
+                        + cfg.alpha_abs      * abs_loss
+                        + cfg.alpha_soft_zipf * zipf_loss)
+
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item()
+
+        avg_loss = epoch_loss / len(loader)
+        history["epoch"].append(epoch)
+        history["train_loss"].append(avg_loss)
+
+        if use_wandb:
+            wandb.log({"train/loss": avg_loss, "epoch": epoch})
+
+        if epoch % cfg.log_every == 0:
+            print(f"  epoch {epoch:5d} | loss {avg_loss:.4f}")
+
+        if epoch % cfg.eval_every == 0:
+            acc = evaluator.run(test_examples, max_examples=1000)
+            history["eval_epoch"].append(epoch)
+            history["test_acc"].append(acc)
+            print(f"  epoch {epoch:5d} | test_acc {acc:.3f}")
+            if use_wandb:
+                wandb.log({"eval/accuracy": acc, "epoch": epoch})
+            save_curves(history, out_dir)
+            with open(out_dir / "history.json", "w") as f:
+                json.dump(history, f, indent=2)
+            if acc > best_acc:
+                best_acc = acc
+                best_dir = out_dir / "best"
+                best_dir.mkdir(exist_ok=True)
+                torch.save(model.state_dict(), best_dir / "model_state_dict.pt")
+                with open(best_dir / "sorl_config.json", "w") as f:
+                    json.dump({"K": cfg.K, "abs_vocab": cfg.abs_vocab, "p": cfg.p,
+                               "n_layer": cfg.n_layer, "n_head": cfg.n_head,
+                               "n_embd": cfg.n_embd, "d_mlp": cfg.d_mlp,
+                               "best_epoch": epoch, "best_acc": acc}, f)
+
+    final_acc = evaluator.run(test_examples, max_examples=2000)
+    print(f"\nFinal test accuracy: {final_acc:.4f} ({int(final_acc * len(test_examples))}/{len(test_examples)})")
+
+    history["final_acc"] = final_acc
+    with open(out_dir / "history.json", "w") as f:
+        json.dump(history, f, indent=2)
+    save_curves(history, out_dir)
+    with open(out_dir / "config.json", "w") as f:
+        json.dump(asdict(cfg), f, indent=2)
+    (out_dir / "final").mkdir(exist_ok=True)
+    torch.save(model.state_dict(), out_dir / "final" / "model_state_dict.pt")
+    with open(out_dir / "final" / "sorl_config.json", "w") as f:
+        json.dump({"K": cfg.K, "abs_vocab": cfg.abs_vocab, "p": cfg.p,
+                   "n_layer": cfg.n_layer, "n_head": cfg.n_head,
+                   "n_embd": cfg.n_embd, "d_mlp": cfg.d_mlp}, f)
+    print(f"Model saved to {out_dir / 'final'}")
+
+    if use_wandb:
+        wandb.log({"eval/final_accuracy": final_acc})
+        wandb.finish()
+
+    return model, final_acc
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--mode",          default="sorl", choices=["baseline", "sorl"])
+    p.add_argument("--K",             type=int,   default=1)
+    p.add_argument("--abs_vocab",     type=int,   default=30)
+    p.add_argument("--num_epochs",    type=int,   default=500)
+    p.add_argument("--batch_size",    type=int,   default=256)
+    p.add_argument("--lr",            type=float, default=1e-3)
+    p.add_argument("--weight_decay",  type=float, default=0.1)
+    p.add_argument("--eval_every",    type=int,   default=50)
+    p.add_argument("--log_every",     type=int,   default=10)
+    p.add_argument("--n_layer",       type=int,   default=1)
+    p.add_argument("--n_head",        type=int,   default=4)
+    p.add_argument("--n_embd",        type=int,   default=128)
+    p.add_argument("--d_mlp",         type=int,   default=512)
+    p.add_argument("--device",        default="cuda")
+    p.add_argument("--job_name",      default="")
+    p.add_argument("--out_dir",       default="")
+    p.add_argument("--no_wandb",      action="store_true")
+    args = p.parse_args()
+
+    cfg = ModularConfig(**vars(args))
+    train(cfg)
+
+
+if __name__ == "__main__":
+    main()