Remove push_to_hub.py; generalise cloud scripts (no account-specific defaults)

README.md

@@ -832,23 +832,19 @@ and saves the controller with a JSON summary. With `--init-from`, the
 previous rung's gates are carried onto the grown model and fresh gates are
 scored only for the new layers and channels.
 
-**`scripts/push_to_hub.py`.** Copies the two remote-code files into the
-checkpoint folder with absolute imports, writes an `auto_map` into
-`config.json`, removes internal marker files, and uploads the checkpoint
-(resumable for large folders), the controller and its summary, and with
-`--with-card` this README and the license.
-
 **`cloud/bootstrap.sh`.** The unattended node script used for stages 2 to
 6. Creates a virtual environment, installs the ROCm or CUDA PyTorch build
-to match the node, pulls the staged data and stage 1 checkpoint from a
-private Hub dataset, runs a short 8-rank FSDP smoke test on a tiny model,
-then for each rung grows, pretrains and fits the controller, skipping any
-rung whose `.rung_done` marker exists. Per-rung training flags are read
-from `pretrain_extra_<stage>.txt` under `RIEMANN_ROOT` so settings can be
-tuned between rungs without editing the script. Finishes by pushing the
-final checkpoint and all controllers to the model repository.
-
-**`cloud/launch_node.ps1`.** Creates the GPU droplet with `doctl`, waits
+to match the node, pulls the staged data and stage 1 checkpoint from the
+private Hub dataset named in `STAGING_REPO`, runs a short 8-rank FSDP
+smoke test on a tiny model, then for each rung grows, pretrains and fits
+the controller, skipping any rung whose `.rung_done` marker exists.
+Per-rung training flags are read from `pretrain_extra_<stage>.txt` under
+`RIEMANN_ROOT` so settings can be tuned between rungs without editing the
+script. Leaves the final checkpoint and controllers on the node for
+archiving.
+
+**`cloud/launch_node.ps1`.** Creates the GPU droplet with `doctl` using
+the SSH key name and staging repository given on the command line, waits
 for SSH, ships the code bundle and the Hugging Face token file, and
 starts `bootstrap.sh` under `nohup` so it survives disconnects.
 
@@ -859,9 +855,10 @@ fetches one stage file by file with per-file SHA-256 verification and
 retry, which is the more reliable path for checkpoints above 20 GB.
 
 **`cloud/verify_before_destroy.ps1`.** Checks that the node reports
-completion, every rung and controller is mirrored and verified locally,
-and the Hub holds every stage 6 shard at the expected size, and prints a
-single safe-or-not verdict before the droplet is destroyed.
+completion and that every rung and controller is mirrored and verified
+locally, optionally also comparing shard sizes against a Hub repository
+when `-ModelRepo` is given, and prints a single safe-or-not verdict
+before the droplet is destroyed.
 
 **`cloud/make_tiny_stage.py`.** Writes a two-layer, 64-wide checkpoint in
 the stage 2 slot for the FSDP smoke test.

training/cloud/bootstrap.sh

@@ -1,133 +1,124 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-export RIEMANN_ROOT="${RIEMANN_ROOT:-/workspace/RiemannModel}"
-export HF_HUB_ENABLE_HF_TRANSFER=1
-export TOKENIZERS_PARALLELISM=false
-export NCCL_DEBUG=WARN
-export PYTHONUNBUFFERED=1
-PROJECT="${PROJECT:-/workspace/riemann}"
-STAGING_REPO="OpenPeerAI/riemann-staging"
-MODEL_REPO="OpenPeerAI/My-Name-is-Riemann-Prerelease"
-LOGDIR="$RIEMANN_ROOT/logs"
-mkdir -p "$RIEMANN_ROOT"/{data,checkpoints,ntk_controllers,hf_cache,staging} "$LOGDIR"
-cd "$PROJECT"
-
-log() { echo "$(date '+%F %T') [bootstrap] $*"; }
-gate() { echo; echo "================ GATE: $* ================"; echo; }
-
-gate "preflight"
-log "host: $(hostname)  kernel: $(uname -r)"
-df -h "$RIEMANN_ROOT" | tail -1
-nproc; free -g | head -2
-if command -v rocm-smi >/dev/null 2>&1; then rocm-smi --showproductname | grep -i "card series" | head -8 || true; fi
-if command -v nvidia-smi >/dev/null 2>&1; then nvidia-smi -L || true; fi
-
-gate "python environment"
-VENV="/workspace/venv"
-if [ ! -x "$VENV/bin/python" ]; then
-  log "installing python3-venv/pip and creating venv"
-  export DEBIAN_FRONTEND=noninteractive
-  apt-get update -qq && apt-get install -y -qq python3-venv python3-pip git >/dev/null
-  python3 -m venv "$VENV"
-fi
-export PATH="$VENV/bin:$PATH"
-PY="$VENV/bin/python"
-$PY -m pip install -q --upgrade pip
-if [ "$($PY -c 'import torch;print(torch.cuda.is_available())' 2>/dev/null)" != "True" ]; then
-  log "installing torch GPU build"
-  if command -v rocm-smi >/dev/null 2>&1 || command -v amd-smi >/dev/null 2>&1; then
-    $PY -m pip install -q torch --index-url "${TORCH_INDEX:-https://download.pytorch.org/whl/rocm6.4}"
-  else
-    $PY -m pip install -q torch --index-url https://download.pytorch.org/whl/cu124
-  fi
-fi
-$PY -m pip install -q "transformers>=4.56,<5" "safetensors>=0.4" "huggingface_hub>=0.24" hf_transfer "numpy>=1.26" "pandas>=2.0" "pyarrow>=15" "datasets>=2.20" "accelerate>=0.33" "git+https://github.com/leochlon/ntkmirror.git"
-$PY - <<'EOF'
-import torch, transformers
-print("torch", torch.__version__, "| hip:", getattr(torch.version, "hip", None), "| cuda:", torch.version.cuda)
-n = torch.cuda.device_count()
-print("gpu count:", n)
-for i in range(n):
-    p = torch.cuda.get_device_properties(i)
-    print(f"  [{i}] {p.name}  {p.total_memory/2**30:.0f} GiB")
-assert n == 8, f"expected 8 GPUs, found {n}"
-print("transformers", transformers.__version__)
-import ntkmirror; print("ntkmirror ok")
-EOF
-
-gate "pull staging data"
-if [ ! -f "$RIEMANN_ROOT/data/pretrain_corpus.txt" ]; then
-  hf download "$STAGING_REPO" --repo-type dataset --local-dir "$RIEMANN_ROOT/staging"
-  log "decompressing corpus"
-  gunzip -c "$RIEMANN_ROOT/staging/pretrain_corpus.txt.gz" > "$RIEMANN_ROOT/data/pretrain_corpus.txt"
-  gunzip -c "$RIEMANN_ROOT/staging/ntk_finetune.jsonl.gz" > "$RIEMANN_ROOT/data/ntk_finetune.jsonl"
-  cp -r "$RIEMANN_ROOT/staging/checkpoints/stage1_279m" "$RIEMANN_ROOT/checkpoints/"
-  cp "$RIEMANN_ROOT/staging/ntk_controllers/"* "$RIEMANN_ROOT/ntk_controllers/"
-fi
-wc -l "$RIEMANN_ROOT/data/pretrain_corpus.txt" "$RIEMANN_ROOT/data/ntk_finetune.jsonl"
-ls -la "$RIEMANN_ROOT/checkpoints/stage1_279m" "$RIEMANN_ROOT/ntk_controllers"
-
-if [ "${SKIP_SMOKE:-0}" != "1" ]; then
-gate "FSDP smoke test (tiny model, 8 ranks, ~2 min)"
-$PY cloud/make_tiny_stage.py "$RIEMANN_ROOT/_smoke/checkpoints"
-torchrun --nproc_per_node 8 scripts/pretrain_cuda.py --stage stage2_538m \
-  --checkpoint-root "$RIEMANN_ROOT/_smoke/checkpoints" \
-  --max-steps 6 --batch-size 2 --grad-accum 2 --seq-len 128 --log-every 2 --save-every 3 --num-workers 0 \
-  2>&1 | tee "$LOGDIR/smoke.log" | grep -E "INFO|Error|Traceback" | tail -20
-$PY -c "import sys; sys.path.insert(0,'.'); from riemann.modeling_riemann import RiemannForCausalLM; RiemannForCausalLM.from_pretrained('$RIEMANN_ROOT/_smoke/checkpoints/stage2_538m'); print('smoke checkpoint reloads OK')"
-rm -rf "$RIEMANN_ROOT/_smoke"
-fi
-
-run_rung() {
-  local from="$1" to="$2" epochs="$3"
-  local marker="$RIEMANN_ROOT/checkpoints/$to/.rung_done"
-  if [ -f "$marker" ]; then
-    log "rung $to already complete, skipping"
-    return 0
-  fi
-  gate "rung $from -> $to (epochs=$epochs)"
-  local t0; t0=$(date +%s)
-  if [ ! -f "$RIEMANN_ROOT/checkpoints/$to/config.json" ]; then
-    $PY scripts/grow_and_continue.py --from-stage "$from" --to-stage "$to" 2>&1 | tee "$LOGDIR/grow_$to.log" | grep -E "INFO|Error" | tail -5
-  fi
-  local extra=""
-  if [ -f "$RIEMANN_ROOT/pretrain_extra_$to.txt" ]; then
-    extra="$(cat "$RIEMANN_ROOT/pretrain_extra_$to.txt")"
-  elif [ -f "$RIEMANN_ROOT/pretrain_extra.txt" ]; then
-    extra="$(cat "$RIEMANN_ROOT/pretrain_extra.txt")"
-  fi
-  log "pretrain extra flags for $to: ${extra:-<none>}"
-  torchrun --nproc_per_node 8 scripts/pretrain_cuda.py --stage "$to" --epochs "$epochs" \
-    --num-workers 6 --log-every 20 --save-every 250 $extra \
-    2>&1 | tee "$LOGDIR/pretrain_$to.log" | grep -E "INFO \[rank 0\]|Error|Traceback"
-  $PY scripts/fit_ntk.py --stage "$to" --init-from "$from" --gates 8192 --batch-size 8 --max-length 512 \
-    2>&1 | tee "$LOGDIR/ntk_$to.log" | grep -E "INFO|Error|Traceback"
-  touch "$marker"
-  local t1; t1=$(date +%s)
-  log "rung $to finished in $(( (t1 - t0) / 60 )) min"
-}
-
-gate "stage 2 alone (throughput check)"
-run_rung stage1_279m stage2_538m 1
-log "stage 2 done. Inspect $LOGDIR/pretrain_stage2_538m.log tok/s before continuing."
-if [ "${STOP_AFTER_STAGE2:-0}" = "1" ]; then
-  log "STOP_AFTER_STAGE2=1, exiting for review"
-  exit 0
-fi
-
-run_rung stage2_538m stage3_1.19b 1
-run_rung stage3_1.19b stage4_2.65b 1
-run_rung stage4_2.65b stage5_6.66b 1
-run_rung stage5_6.66b stage6_12.85b 2
-
-gate "push weights and controllers to private model repo (no card)"
-$PY scripts/push_to_hub.py --stage stage6_12.85b --repo-id "$MODEL_REPO" --private 2>&1 | tee "$LOGDIR/push.log" | grep -E "INFO|Error"
-for s in stage2_538m stage3_1.19b stage4_2.65b stage5_6.66b; do
-  hf upload "$MODEL_REPO" "$RIEMANN_ROOT/ntk_controllers/$s.pt" "ntk_controller/$s.pt" --private || true
-  hf upload "$MODEL_REPO" "$RIEMANN_ROOT/ntk_controllers/$s.json" "ntk_controller/$s.json" --private || true
-done
-hf upload "$MODEL_REPO" "$LOGDIR" "training_logs" --private || true
-
-gate "ALL DONE"
-log "Prerelease weights are in $MODEL_REPO (private). Safe to destroy this node after confirming the upload."
+#!/usr/bin/env bash
+set -euo pipefail
+
+export RIEMANN_ROOT="${RIEMANN_ROOT:-/workspace/RiemannModel}"
+export HF_HUB_ENABLE_HF_TRANSFER=1
+export TOKENIZERS_PARALLELISM=false
+export NCCL_DEBUG=WARN
+export PYTHONUNBUFFERED=1
+PROJECT="${PROJECT:-/workspace/riemann}"
+STAGING_REPO="${STAGING_REPO:?set STAGING_REPO to the private Hub dataset holding data/, the stage 1 checkpoint and controller}"
+LOGDIR="$RIEMANN_ROOT/logs"
+mkdir -p "$RIEMANN_ROOT"/{data,checkpoints,ntk_controllers,hf_cache,staging} "$LOGDIR"
+cd "$PROJECT"
+
+log() { echo "$(date '+%F %T') [bootstrap] $*"; }
+gate() { echo; echo "================ GATE: $* ================"; echo; }
+
+gate "preflight"
+log "host: $(hostname)  kernel: $(uname -r)"
+df -h "$RIEMANN_ROOT" | tail -1
+nproc; free -g | head -2
+if command -v rocm-smi >/dev/null 2>&1; then rocm-smi --showproductname | grep -i "card series" | head -8 || true; fi
+if command -v nvidia-smi >/dev/null 2>&1; then nvidia-smi -L || true; fi
+
+gate "python environment"
+VENV="/workspace/venv"
+if [ ! -x "$VENV/bin/python" ]; then
+  log "installing python3-venv/pip and creating venv"
+  export DEBIAN_FRONTEND=noninteractive
+  apt-get update -qq && apt-get install -y -qq python3-venv python3-pip git >/dev/null
+  python3 -m venv "$VENV"
+fi
+export PATH="$VENV/bin:$PATH"
+PY="$VENV/bin/python"
+$PY -m pip install -q --upgrade pip
+if [ "$($PY -c 'import torch;print(torch.cuda.is_available())' 2>/dev/null)" != "True" ]; then
+  log "installing torch GPU build"
+  if command -v rocm-smi >/dev/null 2>&1 || command -v amd-smi >/dev/null 2>&1; then
+    $PY -m pip install -q torch --index-url "${TORCH_INDEX:-https://download.pytorch.org/whl/rocm6.4}"
+  else
+    $PY -m pip install -q torch --index-url https://download.pytorch.org/whl/cu124
+  fi
+fi
+$PY -m pip install -q "transformers>=4.56,<5" "safetensors>=0.4" "huggingface_hub>=0.24" hf_transfer "numpy>=1.26" "pandas>=2.0" "pyarrow>=15" "datasets>=2.20" "accelerate>=0.33" "git+https://github.com/leochlon/ntkmirror.git"
+$PY - <<'EOF'
+import torch, transformers
+print("torch", torch.__version__, "| hip:", getattr(torch.version, "hip", None), "| cuda:", torch.version.cuda)
+n = torch.cuda.device_count()
+print("gpu count:", n)
+for i in range(n):
+    p = torch.cuda.get_device_properties(i)
+    print(f"  [{i}] {p.name}  {p.total_memory/2**30:.0f} GiB")
+assert n == 8, f"expected 8 GPUs, found {n}"
+print("transformers", transformers.__version__)
+import ntkmirror; print("ntkmirror ok")
+EOF
+
+gate "pull staging data"
+if [ ! -f "$RIEMANN_ROOT/data/pretrain_corpus.txt" ]; then
+  hf download "$STAGING_REPO" --repo-type dataset --local-dir "$RIEMANN_ROOT/staging"
+  log "decompressing corpus"
+  gunzip -c "$RIEMANN_ROOT/staging/pretrain_corpus.txt.gz" > "$RIEMANN_ROOT/data/pretrain_corpus.txt"
+  gunzip -c "$RIEMANN_ROOT/staging/ntk_finetune.jsonl.gz" > "$RIEMANN_ROOT/data/ntk_finetune.jsonl"
+  cp -r "$RIEMANN_ROOT/staging/checkpoints/stage1_279m" "$RIEMANN_ROOT/checkpoints/"
+  cp "$RIEMANN_ROOT/staging/ntk_controllers/"* "$RIEMANN_ROOT/ntk_controllers/"
+fi
+wc -l "$RIEMANN_ROOT/data/pretrain_corpus.txt" "$RIEMANN_ROOT/data/ntk_finetune.jsonl"
+ls -la "$RIEMANN_ROOT/checkpoints/stage1_279m" "$RIEMANN_ROOT/ntk_controllers"
+
+if [ "${SKIP_SMOKE:-0}" != "1" ]; then
+gate "FSDP smoke test (tiny model, 8 ranks, ~2 min)"
+$PY cloud/make_tiny_stage.py "$RIEMANN_ROOT/_smoke/checkpoints"
+torchrun --nproc_per_node 8 scripts/pretrain_cuda.py --stage stage2_538m \
+  --checkpoint-root "$RIEMANN_ROOT/_smoke/checkpoints" \
+  --max-steps 6 --batch-size 2 --grad-accum 2 --seq-len 128 --log-every 2 --save-every 3 --num-workers 0 \
+  2>&1 | tee "$LOGDIR/smoke.log" | grep -E "INFO|Error|Traceback" | tail -20
+$PY -c "import sys; sys.path.insert(0,'.'); from riemann.modeling_riemann import RiemannForCausalLM; RiemannForCausalLM.from_pretrained('$RIEMANN_ROOT/_smoke/checkpoints/stage2_538m'); print('smoke checkpoint reloads OK')"
+rm -rf "$RIEMANN_ROOT/_smoke"
+fi
+
+run_rung() {
+  local from="$1" to="$2" epochs="$3"
+  local marker="$RIEMANN_ROOT/checkpoints/$to/.rung_done"
+  if [ -f "$marker" ]; then
+    log "rung $to already complete, skipping"
+    return 0
+  fi
+  gate "rung $from -> $to (epochs=$epochs)"
+  local t0; t0=$(date +%s)
+  if [ ! -f "$RIEMANN_ROOT/checkpoints/$to/config.json" ]; then
+    $PY scripts/grow_and_continue.py --from-stage "$from" --to-stage "$to" 2>&1 | tee "$LOGDIR/grow_$to.log" | grep -E "INFO|Error" | tail -5
+  fi
+  local extra=""
+  if [ -f "$RIEMANN_ROOT/pretrain_extra_$to.txt" ]; then
+    extra="$(cat "$RIEMANN_ROOT/pretrain_extra_$to.txt")"
+  elif [ -f "$RIEMANN_ROOT/pretrain_extra.txt" ]; then
+    extra="$(cat "$RIEMANN_ROOT/pretrain_extra.txt")"
+  fi
+  log "pretrain extra flags for $to: ${extra:-<none>}"
+  torchrun --nproc_per_node 8 scripts/pretrain_cuda.py --stage "$to" --epochs "$epochs" \
+    --num-workers 6 --log-every 20 --save-every 250 $extra \
+    2>&1 | tee "$LOGDIR/pretrain_$to.log" | grep -E "INFO \[rank 0\]|Error|Traceback"
+  $PY scripts/fit_ntk.py --stage "$to" --init-from "$from" --gates 8192 --batch-size 8 --max-length 512 \
+    2>&1 | tee "$LOGDIR/ntk_$to.log" | grep -E "INFO|Error|Traceback"
+  touch "$marker"
+  local t1; t1=$(date +%s)
+  log "rung $to finished in $(( (t1 - t0) / 60 )) min"
+}
+
+gate "stage 2 alone (throughput check)"
+run_rung stage1_279m stage2_538m 1
+log "stage 2 done. Inspect $LOGDIR/pretrain_stage2_538m.log tok/s before continuing."
+if [ "${STOP_AFTER_STAGE2:-0}" = "1" ]; then
+  log "STOP_AFTER_STAGE2=1, exiting for review"
+  exit 0
+fi
+
+run_rung stage2_538m stage3_1.19b 1
+run_rung stage3_1.19b stage4_2.65b 1
+run_rung stage4_2.65b stage5_6.66b 1
+run_rung stage5_6.66b stage6_12.85b 2
+
+gate "ALL DONE"
+log "Final checkpoint: $RIEMANN_ROOT/checkpoints/stage6_12.85b. Controllers: $RIEMANN_ROOT/ntk_controllers. Archive these before destroying the node."

training/cloud/launch_node.ps1

@@ -2,8 +2,9 @@ param(
     [string]$Region = "nyc2",
     [string]$Size = "gpu-mi325x8-2048gb",
     [string]$Image = "gpu-amd-base",
-    [string]$Name = "riemann-climb",
-    [string]$SshKeyName = "github-mass-merger-key",
+    [string]$Name = "riemann-node",
+    [Parameter(Mandatory = $true)][string]$SshKeyName,
+    [Parameter(Mandatory = $true)][string]$StagingRepo,
     [switch]$StopAfterStage2
 )
 $ErrorActionPreference = "Stop"
@@ -48,7 +49,7 @@ ssh root@$ip "cd /workspace/riemann && tar -xzf /workspace/riemann_bundle.tar.gz
 
 $stopFlag = if ($StopAfterStage2) { "STOP_AFTER_STAGE2=1" } else { "STOP_AFTER_STAGE2=0" }
 Write-Host "Starting bootstrap under nohup ($stopFlag) ..."
-ssh root@$ip "cd /workspace/riemann && mkdir -p /workspace/RiemannModel/logs && $stopFlag nohup bash cloud/bootstrap.sh > /workspace/RiemannModel/logs/bootstrap.log 2>&1 &"
+ssh root@$ip "cd /workspace/riemann && mkdir -p /workspace/RiemannModel/logs && STAGING_REPO=$StagingRepo $stopFlag nohup bash cloud/bootstrap.sh > /workspace/RiemannModel/logs/bootstrap.log 2>&1 &"
 Start-Sleep 5
 ssh root@$ip "tail -5 /workspace/RiemannModel/logs/bootstrap.log"
 Write-Host ""

training/cloud/mirror_node.ps1

@@ -1,6 +1,6 @@
 param(
     [string]$NodeIp = (Get-Content "$env:RIEMANN_ROOT\cloud_state\node_ip.txt").Trim(),
-    [string]$LocalRoot = "C:\RiemannModel",
+    [string]$LocalRoot = $env:RIEMANN_ROOT,
     [string]$RemoteRoot = "/workspace/RiemannModel",
     [int]$PollSeconds = 300
 )

training/cloud/pull_stage.ps1

@@ -1,7 +1,7 @@
 param(
     [string]$NodeIp = (Get-Content "$env:RIEMANN_ROOT\cloud_state\node_ip.txt").Trim(),
     [string]$Stage = "stage6_12.85b",
-    [string]$LocalRoot = "C:\RiemannModel",
+    [string]$LocalRoot = $env:RIEMANN_ROOT,
     [string]$RemoteRoot = "/workspace/RiemannModel"
 )
 $ErrorActionPreference = "Continue"

training/cloud/verify_before_destroy.ps1

@@ -1,7 +1,7 @@
 param(
     [string]$NodeIp = (Get-Content "$env:RIEMANN_ROOT\cloud_state\node_ip.txt").Trim(),
-    [string]$LocalRoot = "C:\RiemannModel",
-    [string]$ModelRepo = "OpenPeerAI/My-Name-is-Riemann-Prerelease"
+    [string]$LocalRoot = $env:RIEMANN_ROOT,
+    [string]$ModelRepo = ""
 )
 $ErrorActionPreference = "Continue"
 $ok = $true
@@ -15,19 +15,21 @@ foreach ($s in "stage2_538m", "stage3_1.19b", "stage4_2.65b", "stage5_6.66b", "s
 }
 Check "local training logs present" ((Get-ChildItem "$LocalRoot\cloud_logs" -Filter "pretrain_stage6*.log" -ErrorAction SilentlyContinue).Count -ge 1)
 
-$remoteList = ssh "root@$NodeIp" "cd /workspace/RiemannModel/checkpoints/stage6_12.85b && ls -l model-*.safetensors | awk '{print `$5, `$9}'"
-$hub = python -c "
+if ($ModelRepo) {
+    $remoteList = ssh "root@$NodeIp" "cd /workspace/RiemannModel/checkpoints/stage6_12.85b && ls -l model-*.safetensors | awk '{print `$5, `$9}'"
+    $hub = python -c "
 from huggingface_hub import HfApi
 for f in HfApi().list_repo_tree('$ModelRepo', recursive=True):
     s=getattr(f,'size',None)
     if s and f.path.endswith('.safetensors'): print(s, f.path.split('/')[-1])
 "
-$remoteMap = @{}; foreach ($l in $remoteList) { if ($l -match '^(\d+)\s+(\S+)$') { $remoteMap[$matches[2]] = [long]$matches[1] } }
-$hubMap = @{}; foreach ($l in $hub) { if ($l -match '^(\d+)\s+(\S+)$') { $hubMap[$matches[2]] = [long]$matches[1] } }
-$shardsMatch = ($remoteMap.Count -gt 0) -and ($remoteMap.Count -eq $hubMap.Count)
-foreach ($k in $remoteMap.Keys) { if ($hubMap[$k] -ne $remoteMap[$k]) { $shardsMatch = $false } }
-Check "Hub has all $($remoteMap.Count) stage6 shards with matching sizes" $shardsMatch
-Check "Hub has stage6 controller" (($hub = python -c "from huggingface_hub import HfApi; print(any(f.path=='ntk_controller/stage6_12.85b.pt' for f in HfApi().list_repo_tree('$ModelRepo', recursive=True)))") -eq "True")
+    $remoteMap = @{}; foreach ($l in $remoteList) { if ($l -match '^(\d+)\s+(\S+)$') { $remoteMap[$matches[2]] = [long]$matches[1] } }
+    $hubMap = @{}; foreach ($l in $hub) { if ($l -match '^(\d+)\s+(\S+)$') { $hubMap[$matches[2]] = [long]$matches[1] } }
+    $shardsMatch = ($remoteMap.Count -gt 0) -and ($remoteMap.Count -eq $hubMap.Count)
+    foreach ($k in $remoteMap.Keys) { if ($hubMap[$k] -ne $remoteMap[$k]) { $shardsMatch = $false } }
+    Check "Hub has all $($remoteMap.Count) stage6 shards with matching sizes" $shardsMatch
+    Check "Hub has stage6 controller" ((python -c "from huggingface_hub import HfApi; print(any(f.path=='ntk_controller/stage6_12.85b.pt' for f in HfApi().list_repo_tree('$ModelRepo', recursive=True)))") -eq "True")
+}
 
 ""
-if ($ok) { "SAFE TO DESTROY: all archive and Hub checks passed." } else { "DO NOT DESTROY: one or more checks failed." }
+if ($ok) { "SAFE TO DESTROY: all archive checks passed." } else { "DO NOT DESTROY: one or more checks failed." }

training/scripts/push_to_hub.py

@@ -1,138 +0,0 @@
-from __future__ import annotations
-import argparse
-import json
-import logging
-import shutil
-import sys
-from pathlib import Path
-
-sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
-from huggingface_hub import HfApi, create_repo, whoami
-from riemann.growth_ladder import get_stage
-from riemann.paths import CHECKPOINT_ROOT, CONTROLLER_ROOT
-
-logger = logging.getLogger("riemann.push_to_hub")
-PROJECT_ROOT = Path(__file__).resolve().parent.parent
-REMOTE_CODE_FILES = ("configuration_riemann.py", "modeling_riemann.py")
-
-
-def stage_remote_code(checkpoint_dir: Path) -> None:
-    for name in REMOTE_CODE_FILES:
-        src = (PROJECT_ROOT / "riemann" / name).read_text(encoding="utf-8")
-        src = src.replace(
-            "from .configuration_riemann import", "from configuration_riemann import"
-        )
-        (checkpoint_dir / name).write_text(src, encoding="utf-8")
-    for stray in (".rung_done", ".mirrored"):
-        p = checkpoint_dir / stray
-        if p.exists():
-            p.unlink()
-    config_path = checkpoint_dir / "config.json"
-    config = json.loads(config_path.read_text(encoding="utf-8"))
-    config["auto_map"] = {
-        "AutoConfig": "configuration_riemann.RiemannConfig",
-        "AutoModelForCausalLM": "modeling_riemann.RiemannForCausalLM",
-    }
-    config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")
-
-
-def main() -> None:
-    logging.basicConfig(
-        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
-    )
-    parser = argparse.ArgumentParser(
-        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
-    )
-    parser.add_argument(
-        "--stage", default="stage0_124m", help="which checkpoint to push"
-    )
-    parser.add_argument(
-        "--repo-id", required=True, help="e.g. your-username/my-name-is-riemann"
-    )
-    parser.add_argument(
-        "--private", action="store_true", help="create/push as a private repo"
-    )
-    parser.add_argument(
-        "--with-card",
-        action="store_true",
-        help="also upload README.md and LICENSE (off by default so the card is only pushed after review)",
-    )
-    parser.add_argument("--checkpoint-root", default=str(CHECKPOINT_ROOT))
-    parser.add_argument("--controller-root", default=str(CONTROLLER_ROOT))
-    args = parser.parse_args()
-    try:
-        who = whoami()
-        logger.info("Logged in to Hugging Face as %s", who.get("name"))
-    except Exception as exc:
-        raise RuntimeError(
-            "Not logged in to Hugging Face. Run `huggingface-cli login` or set the HF_TOKEN environment variable first."
-        ) from exc
-    stage = get_stage(args.stage)
-    checkpoint_dir = Path(args.checkpoint_root) / stage.name
-    if not (checkpoint_dir / "config.json").exists():
-        raise FileNotFoundError(f"No checkpoint found at {checkpoint_dir}")
-    create_repo(args.repo_id, private=args.private, exist_ok=True)
-    api = HfApi()
-    stage_remote_code(checkpoint_dir)
-    logger.info("Uploading checkpoint folder %s -> %s", checkpoint_dir, args.repo_id)
-    total_bytes = sum(
-        p.stat().st_size for p in checkpoint_dir.rglob("*") if p.is_file()
-    )
-    if total_bytes > 10 * 1024**3:
-        api.upload_large_folder(
-            folder_path=str(checkpoint_dir),
-            repo_id=args.repo_id,
-            repo_type="model",
-        )
-    else:
-        api.upload_folder(
-            folder_path=str(checkpoint_dir),
-            repo_id=args.repo_id,
-            commit_message=f"Upload My Name is Riemann checkpoint ({stage.name})",
-        )
-    controller_path = Path(args.controller_root) / f"{stage.name}.pt"
-    if controller_path.exists():
-        logger.info("Also uploading NTK gate controller %s", controller_path)
-        api.upload_file(
-            path_or_fileobj=str(controller_path),
-            path_in_repo=f"ntk_controller/{stage.name}.pt",
-            repo_id=args.repo_id,
-            commit_message=f"Upload NTK gate controller ({stage.name})",
-        )
-        summary_path = controller_path.with_suffix(".json")
-        if summary_path.exists():
-            api.upload_file(
-                path_or_fileobj=str(summary_path),
-                path_in_repo=f"ntk_controller/{stage.name}.json",
-                repo_id=args.repo_id,
-                commit_message=f"Upload NTK controller summary ({stage.name})",
-            )
-    else:
-        logger.info(
-            "No NTK gate controller found at %s yet, skipping that part.",
-            controller_path,
-        )
-    if args.with_card:
-        readme_path = PROJECT_ROOT / "README.md"
-        if readme_path.exists():
-            api.upload_file(
-                path_or_fileobj=str(readme_path),
-                path_in_repo="README.md",
-                repo_id=args.repo_id,
-                commit_message="Upload README",
-            )
-        license_path = PROJECT_ROOT / "LICENSE"
-        if license_path.exists():
-            api.upload_file(
-                path_or_fileobj=str(license_path),
-                path_in_repo="LICENSE",
-                repo_id=args.repo_id,
-                commit_message="Upload LICENSE",
-            )
-    else:
-        logger.info("Skipping README.md and LICENSE (pass --with-card to upload them).")
-    logger.info("Done. https://huggingface.co/%s", args.repo_id)
-
-
-if __name__ == "__main__":
-    main()