Add giant new-run to scaffold a config.toml + run dir ahead of training

Pulled forward from the not-yet-mergeable condor-gpu-train-rollout branch:
`new-run` resolves CLI hyperparameter overrides into a full config.toml and
run dir (reusing the existing default_out_dir_name collision-avoidance and a
newly factored-out router-override helper shared with `train`), so a run can
be prepared and reviewed before `giant train` actually kicks off. Also
brings README up to date with the model/CLI as it actually stands
(physical/embedding conditioning, WGAN/MoE-router modes, giant analyze,
W&B, setup-stage caching), which had drifted back to describing the
Phase-1 proof-of-concept.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

giant/cli.py

@@ -134,6 +134,30 @@ def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]:
     return out
 
 
+def _router_cli_overrides(
+    router: bool | None,
+    router_type: str | None,
+    n_experts: int | None,
+    router_axis: list[str] | None,
+) -> dict[str, object]:
+    """Build the `model.router` override dict from `--router`/`--router-type`/
+    `--n-experts`/`--router-axis` flags (empty if none were given). Shared by
+    `train` and `new-run` so both resolve router overrides identically.
+    """
+    cli_router: dict[str, object] = {
+        k: v
+        for k, v in {
+            "enabled": router,
+            "type": router_type,
+            "n_experts": n_experts,
+        }.items()
+        if v is not None
+    }
+    if router_axis:
+        cli_router.update(_parse_router_axis_flags(router_axis))
+    return cli_router
+
+
 _CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions")
 
 
@@ -504,17 +528,7 @@ def train(
         }.items()
         if v is not None
     }
-    cli_router: dict[str, object] = {
-        k: v
-        for k, v in {
-            "enabled": router,
-            "type": router_type,
-            "n_experts": n_experts,
-        }.items()
-        if v is not None
-    }
-    if router_axis:
-        cli_router.update(_parse_router_axis_flags(router_axis))
+    cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
     if cli_router:
         cli_model["router"] = cli_router
     cfg = gconfig.merge_cli_overrides(
@@ -552,13 +566,9 @@ def train(
         # Name only encodes what's non-default (see default_out_dir_name), so
         # two runs with identical hyperparams in the same to-the-minute
         # timestamp would otherwise collide on this name — which also
-        # doubles as the W&B run id (giant.train) — hence the suffix loop.
-        base_name = gconfig.default_out_dir_name(cfg)
-        out_dir = Path("checkpoints") / base_name
-        suffix = 2
-        while out_dir.exists():
-            out_dir = Path("checkpoints") / f"{base_name}_{suffix}"
-            suffix += 1
+        # doubles as the W&B run id (giant.train) — hence the suffix loop in
+        # resolve_default_out_dir.
+        out_dir = gconfig.resolve_default_out_dir(cfg)
 
     typer.echo(f"device:  {_device}")
     typer.echo(f"out_dir: {out_dir}")
@@ -577,6 +587,144 @@ def train(
     )
 
 
+@app.command("new-run")
+def new_run(
+    config: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--config",
+            "-c",
+            help="Base TOML to start from (default: built-in defaults)",
+        ),
+    ] = None,
+    mode: Annotated[Optional[Mode], typer.Option("--mode", "-m")] = None,
+    epochs: Annotated[Optional[int], typer.Option("--epochs", "-e")] = None,
+    batch_size: Annotated[Optional[int], typer.Option("--batch-size", "-b")] = None,
+    lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None,
+    hidden_dim: Annotated[Optional[int], typer.Option("--hidden-dim", "-H")] = None,
+    n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
+    emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
+    dropout: Annotated[Optional[float], typer.Option("--dropout", "-d")] = None,
+    conditioning: Annotated[
+        Optional[Conditioning], typer.Option("--conditioning")
+    ] = None,
+    router: Annotated[Optional[bool], typer.Option("--router/--no-router")] = None,
+    router_type: Annotated[Optional[str], typer.Option("--router-type")] = None,
+    n_experts: Annotated[Optional[int], typer.Option("--n-experts")] = None,
+    router_axis: Annotated[Optional[list[str]], typer.Option("--router-axis")] = None,
+    out: Annotated[
+        Optional[Path],
+        typer.Option("--out", "-o", help="Run dir (default: auto from hyperparams)"),
+    ] = None,
+    comment: Annotated[
+        Optional[str],
+        typer.Option(
+            "--comment", help="Free-text note recorded in config.toml's meta section"
+        ),
+    ] = None,
+    data: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--data",
+            help="Dataset path to fill in the printed next-step command "
+            "(not stored in the config)",
+        ),
+    ] = None,
+    force: Annotated[
+        bool,
+        typer.Option(
+            "--force",
+            help="Overwrite config.toml even if --out already has checkpoints",
+        ),
+    ] = False,
+    dry_run: Annotated[
+        bool,
+        typer.Option(
+            "--dry-run", help="Print the resolved config without writing anything"
+        ),
+    ] = False,
+) -> None:
+    """Scaffold a new training run: resolve hyperparams to a config.toml and lay out its run dir.
+
+    This is the config-file-first counterpart to hand-editing a TOML: start
+    from a base --config (or built-in defaults), override a few hyperparams
+    inline, and this resolves+writes the full `config.toml` into a fresh (or
+    explicit --out) run dir — the same file `giant train --config ...` reads.
+    `giant train` itself overwrites this file in place once it actually runs
+    (with the full dataset-derived meta section), so this scaffold's meta
+    section is just a placeholder recording what was asked for and when.
+    """
+    cli_train = {
+        k: v
+        for k, v in {
+            "mode": mode.value if mode is not None else None,
+            "epochs": epochs,
+            "batch_size": batch_size,
+            "lr": lr,
+        }.items()
+        if v is not None
+    }
+    cli_model: dict[str, object] = {
+        k: v
+        for k, v in {
+            "hidden_dim": hidden_dim,
+            "n_blocks": n_blocks,
+            "emb_dim": emb_dim,
+            "dropout": dropout,
+            "conditioning": conditioning.value if conditioning is not None else None,
+        }.items()
+        if v is not None
+    }
+    cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
+    if cli_router:
+        cli_model["router"] = cli_router
+
+    cfg = gconfig.merge_cli_overrides(
+        gconfig.DEFAULT_CONFIG, config, cli_train, cli_model
+    )
+    run_dir = (out or gconfig.resolve_default_out_dir(cfg)).resolve()
+
+    if not force:
+        existing = [n for n in ("last.pt", "best.pt") if (run_dir / n).exists()]
+        if existing:
+            typer.echo(
+                f"error: {run_dir} already has {', '.join(existing)} — pass "
+                "--force to overwrite its config.toml anyway",
+                err=True,
+            )
+            raise typer.Exit(1)
+
+    typer.echo(f"run dir: {run_dir}")
+
+    if dry_run:
+        typer.echo("dry-run: not writing anything. Resolved config:")
+        for section in ("train", "model"):
+            typer.echo(f"[{section}]")
+            for k, v in cfg[section].items():
+                if k == "router":
+                    continue
+                typer.echo(f"  {k} = {v}")
+        return
+
+    meta = {
+        "git_hash": gconfig.git_hash(),
+        "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+        "created_by": "giant new-run",
+    }
+    if comment:
+        meta["comment"] = comment
+
+    run_dir.mkdir(parents=True, exist_ok=True)
+    gconfig.save_config(cfg, run_dir, meta)
+    config_path = run_dir / "config.toml"
+    typer.echo(f"wrote {config_path}")
+
+    data_arg = str(data) if data is not None else "<data.parquet>"
+    typer.echo("")
+    typer.echo("next:")
+    typer.echo(f"  giant train {data_arg} --config {config_path} --out {run_dir}")
+
+
 @app.command()
 def predict(
     data: Annotated[

giant/config.py

@@ -372,6 +372,21 @@ def default_out_dir_name(cfg: dict, now: datetime | None = None) -> str:
     return name
 
 
+def resolve_default_out_dir(cfg: dict, base: Path = Path("checkpoints")) -> Path:
+    """Auto-derived out dir from cfg's hyperparams (see `default_out_dir_name`),
+    with a numeric suffix loop so two runs whose name collides (same
+    non-default hyperparams, same to-the-minute timestamp) don't clobber each
+    other's directory. Shared by `giant train` and `giant new-run`.
+    """
+    base_name = default_out_dir_name(cfg)
+    out_dir = base / base_name
+    suffix = 2
+    while out_dir.exists():
+        out_dir = base / f"{base_name}_{suffix}"
+        suffix += 1
+    return out_dir
+
+
 def seed_everything(seed: int) -> None:
     random.seed(seed)
     np.random.seed(seed)