{ "nbformat": 4, "nbformat_minor": 5, "metadata": { "colab": { "provenance": [] }, "kernelspec": { "name": "python3", "display_name": "Python 3" }, "language_info": { "name": "python" }, "accelerator": "GPU" }, "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# AMoE Anima close-out \u2014 exp004b + exp017 (Colab / Blackwell 96GB)\n", "\n", "Closes the r2 campaign's Anima gap:\n", "\n", "| | what | why |\n", "|---|---|---|\n", "| **exp004b** | relay retrain: 2 declared seeds, 10k steps, full dataset, held-out eval, LoRA-r16 control | exp004 was one ~40-min run, no results.json, still improving at cutoff |\n", "| **exp017** | multiband3 on the Anima DiT \u2014 first live run of the fork's mode | the band mechanism has never run on a DiT |\n", "| **the fix** | `bf16_master_weights = true` + declared `seed` | exp004's 28 gates were **quantization-frozen**: plain Adam stepped bf16 leaves directly; ~4.5e-4 updates vs 0.0078 half-ULP at \u22123.0 rounded to no-ops (proven from the salvaged Adam state) |\n", "\n", "## Nothing is ever trained twice\n", "\n", "Every arm resolves **LOCAL \u2192 HUB \u2192 train**:\n", "\n", "1. a final-step checkpoint on this disk is reused as-is;\n", "2. else the arm is pulled back from its shipped hub package (a fresh\n", " runtime after a pod cull costs a ~60 MB download per arm, not a\n", " training run) \u2014 its eval curve rides along, since a restored arm has\n", " no tfevents;\n", "3. only a genuinely missing arm trains.\n", "\n", "CELL 2 prints the inventory (`local` / `hub` / `MISSING`) before anything\n", "spends. The smokes are skipped when a full arm already proves the same\n", "thing, CELL 6 rebuilds the held-out cache by itself if every arm was\n", "restored, and CELLs 6\u20137 resolve their own checkpoints \u2014 so you can run\n", "**0 \u2192 1 \u2192 2 \u2192 6 \u2192 7** in a brand-new runtime and get results without a\n", "single training step.\n", "\n", "Run cells top to bottom. After a disconnect, re-run **CELL 0**, then the\n", "interrupted cell (checkpoints land every 30 min; `run_arm(tag,\n", "resume=True)` continues a half-finished run instead of restarting it).\n", "\n", "**Needs:** Colab secret `HF_TOKEN` (notebook access ON) with write access\n", "to `AbstractPhil/geolip-aleph-diffusion`.\n", "\n", "**License:** Anima weights are NON-COMMERCIAL (CircleStone NC + NVIDIA\n", "Open Model License); every checkpoint produced here inherits NC.\n" ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "# \u2550\u2550\u2550 CELL 0 \u2014 ACTIVATION (idempotent; re-run after any disconnect) \u2550\u2550\u2550\u2550\u2550\u2550\n", "import json, os, subprocess, sys, time\n", "from pathlib import Path\n", "\n", "if \"ANIMA\" not in globals(): # paste-ahead guard\n", " ANIMA = {}\n", "\n", "WORK = Path(\"/content/anima_campaign\")\n", "FORK = Path(\"/content/diffusion-pipe\")\n", "DATA, RUNS, TOMLS = WORK / \"data\", WORK / \"runs\", WORK / \"tomls\"\n", "for _d in (WORK, DATA, RUNS, TOMLS):\n", " _d.mkdir(parents=True, exist_ok=True)\n", "\n", "def sh(cmd, check=True):\n", " print(f\"$ {cmd}\")\n", " rc = subprocess.run(cmd, shell=True, text=True).returncode\n", " if check and rc != 0:\n", " raise RuntimeError(f\"rc={rc}: {cmd}\")\n", " return rc\n", "\n", "# -- 1. GPU + Blackwell riders --------------------------------------------\n", "import torch\n", "assert torch.cuda.is_available(), \"No GPU \u2014 Runtime > Change runtime type.\"\n", "_cap = torch.cuda.get_device_capability(0)\n", "print(f\"GPU: {torch.cuda.get_device_name(0)} | sm_{_cap[0]}{_cap[1]} | \"\n", " f\"{torch.cuda.get_device_properties(0).total_memory / 2**30:.0f} GB | \"\n", " f\"torch {torch.__version__} cu{torch.version.cuda}\")\n", "if _cap >= (12, 0):\n", " _tv = tuple(int(x) for x in torch.__version__.split(\"+\")[0].split(\".\")[:2])\n", " _cv = tuple(int(x) for x in (torch.version.cuda or \"0.0\").split(\".\")[:2])\n", " if not (_tv >= (2, 7) and _cv >= (12, 8)):\n", " raise RuntimeError(\n", " f\"sm_120 needs torch>=2.7 + cu>=12.8; runtime has \"\n", " f\"{torch.__version__}/cu{torch.version.cuda}. Do NOT blind-pip \"\n", " \"torch \u2014 switch runtimes, or deliberately install the matching \"\n", " \"cu128 build from pytorch.org.\")\n", " (torch.randn(256, 256, device=\"cuda\")\n", " @ torch.randn(256, 256, device=\"cuda\")).sum().item()\n", " print(\"sm_120 matmul: OK\")\n", "try:\n", " import flash_attn # noqa: F401\n", " raise RuntimeError(\"flash-attn is installed and BROKEN on sm_120 \u2014 \"\n", " \"`pip uninstall flash-attn`; SDPA is the path.\")\n", "except ImportError:\n", " pass\n", "\n", "# -- 2. HF token (never printed) ------------------------------------------\n", "_tok = os.environ.get(\"HF_TOKEN\")\n", "if not _tok:\n", " try:\n", " from google.colab import userdata\n", " _tok = userdata.get(\"HF_TOKEN\")\n", " except Exception:\n", " _tok = None\n", "if not _tok:\n", " raise RuntimeError(\"Add HF_TOKEN in Colab secrets (key icon) with \"\n", " \"notebook access ON, then re-run this cell.\")\n", "os.environ[\"HF_TOKEN\"] = _tok\n", "print(\"HF token: present (not printed)\")\n", "\n", "# -- 3. fork + deps -------------------------------------------------------\n", "if not (FORK / \"train.py\").exists():\n", " sh(f\"git clone -q https://github.com/AbstractEyes/diffusion-pipe {FORK}\")\n", "sh(f\"git -C {FORK} fetch -q origin\")\n", "sh(f\"git -C {FORK} checkout -q origin/main\")\n", "sh(f\"git -C {FORK} submodule update --init submodules/ComfyUI\")\n", "print(\"fork @\", subprocess.run(\n", " [\"git\", \"-C\", str(FORK), \"log\", \"--oneline\", \"-1\"],\n", " capture_output=True, text=True).stdout.strip())\n", "_src = (FORK / \"train.py\").read_text()\n", "assert \"bf16_master_weights\" in _src and \"run_seed\" in _src, \\\n", " \"fork main lacks the campaign patches (bf16_master_weights + seed)\"\n", "# only ComfyUI is initialized above; the hyvideo patch import must be the\n", "# guarded version (fork 5f37138) or train.py dies at import (~24s launcher\n", "# crash, ModuleNotFoundError: hyvideo \u2014 the recorded exp004 + smoke killer)\n", "assert \"_HAS_HYVIDEO\" in (FORK / \"utils\" / \"patches.py\").read_text(), \\\n", " \"fork main lacks the guarded hyvideo import (utils/patches.py)\"\n", "# the 8192-row split monolith needs the per-row-group image reader (fork\n", "# 8ed4b6b) \u2014 the whole-file read dies on >2GiB nested columns\n", "assert \"pa.chunked_array(chunks)\" in \\\n", " (FORK / \"utils\" / \"parquet_source.py\").read_text(), \\\n", " \"fork main lacks the per-row-group ParquetImageReader (parquet_source.py)\"\n", "# bf16_master_weights must be the CLIENT-SIDE mechanism (fork 812449a):\n", "# the ds_config['bf16'] route wraps the optimizer and torch schedulers\n", "# reject the wrapper ('FP16_UnfusedOptimizer is not an Optimizer')\n", "assert \"MasterWeightsAdam\" in _src, \\\n", " \"fork main lacks the client-side MasterWeightsAdam (bf16_master_weights)\"\n", "\n", "if not ANIMA.get(\"DEPS_DONE\"):\n", " sh(f\"pip install -q -r {FORK}/requirements.txt\")\n", " sh('pip install -q \"amoe-lora[diffusion] @ '\n", " 'git+https://github.com/AbstractEyes/amoe-lora\"')\n", " ANIMA[\"DEPS_DONE\"] = True\n", "# Colab preinstalls torchao 0.10.0; peft's LoRA torchao dispatcher RAISES\n", "# on it ('only versions above 0.16.0 are supported') \u2014 killing the LoRA\n", "# control arms. Nothing in this stack uses torchao; with it ABSENT peft's\n", "# probe returns False and dispatch falls through to the standard Linear.\n", "import importlib.util\n", "if importlib.util.find_spec(\"torchao\"):\n", " sh(\"pip uninstall -y -q torchao\")\n", " importlib.invalidate_caches()\n", " assert importlib.util.find_spec(\"torchao\") is None, \"torchao survived\"\n", "import amoe, deepspeed\n", "print(f\"amoe {amoe.__version__} | deepspeed {deepspeed.__version__}\")\n", "\n", "# -- 4. Anima weights (NC) ------------------------------------------------\n", "from huggingface_hub import hf_hub_download\n", "MODELS = WORK / \"models\" / \"anima\"\n", "for _f in (\"diffusion_models/anima-base-v1.0.safetensors\",\n", " \"text_encoders/qwen_3_06b_base.safetensors\",\n", " \"vae/qwen_image_vae.safetensors\"):\n", " hf_hub_download(\"circlestone-labs/Anima\", f\"split_files/{_f}\",\n", " local_dir=str(MODELS))\n", "print(\"Anima split_files present \u2014 NC weights; derived ckpts inherit NC\")\n", "\n", "# -- 5. dataset columns, verified BEFORE anything spends ------------------\n", "# THE DATASET IS qwen-deepfashion-fused (the fused line's design of record:\n", "# r2 exp011/012/013 all trained on it). The CAPTION KEY per the record is\n", "# caption_joycaption with prompt_fused fallback (pod2/dexp011:154; results\n", "# cond = \"caption_joycaption nl77\"). caption_vlm_json belongs to the older\n", "# ft1 shards and does NOT exist here \u2014 the first campaign run's pre-spend\n", "# column check caught exactly that; the fix was the key, not the dataset.\n", "DATASET = \"AbstractPhil/qwen-deepfashion-fused\"\n", "from huggingface_hub import HfApi\n", "SHARDS = sorted(f for f in HfApi().list_repo_files(\n", " DATASET, repo_type=\"dataset\") if f.endswith(\".parquet\"))\n", "assert SHARDS, f\"no parquet shards in {DATASET}\"\n", "print(f\"shards: {len(SHARDS)} (incl. the shard_r0_* series)\")\n", "import pyarrow.parquet as pq\n", "_p0 = hf_hub_download(DATASET, SHARDS[0], repo_type=\"dataset\")\n", "_cols = set(pq.read_schema(_p0).names)\n", "_missing = {\"image\", \"caption_joycaption\", \"id\"} - _cols\n", "assert not _missing, f\"dataset lacks {_missing}; has {sorted(_cols)[:24]}\"\n", "ANIMA.update(DATASET=DATASET, SHARDS=SHARDS,\n", " HAS_WH={\"image_width\", \"image_height\"} <= _cols,\n", " HAS_FUSED_JSON=\"fused_json\" in _cols,\n", " HAS_PROMPT_FUSED=\"prompt_fused\" in _cols,\n", " ID_COL=\"id\")\n", "print(f\"dataset OK: wh cols {ANIMA['HAS_WH']}; fused_json \"\n", " f\"{ANIMA['HAS_FUSED_JSON']}; prompt_fused {ANIMA['HAS_PROMPT_FUSED']}\")\n", "print(\"\\n=== READY ===\")\n" ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "# \u2550\u2550\u2550 CELL 1 \u2014 materialize the gated split (COLUMNAR, two-phase) \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n", "# The fused set is ~160GB; the campaign needs ~9k rows. Two bounded phases:\n", "# A) THIN SCAN: shards in sha256(name) order; per row group, read ONLY\n", "# id/captions/audit and evaluate everything as pyarrow.compute masks \u2014\n", "# the exp011 audit gate, the dexp011:154 caption coalesce (using the\n", "# fork's own CAPTION_SENTINELS), id-dedup. Stop at the row target.\n", "# B) HEAVY WRITE: re-read just the shards phase A used, filter row groups\n", "# by is_in(id, split set), append the derived caption column, stream\n", "# to train/eval parquet. Image bytes never pass through Python.\n", "# Split rule: within the pooled ids, sha256(id) ascending; first EVAL_N\n", "# are eval. Sources deleted after phase B.\n", "import hashlib, json, shutil, sys\n", "import pyarrow as pa\n", "import pyarrow.compute as pc\n", "import pyarrow.parquet as pq\n", "from huggingface_hub import hf_hub_download\n", "\n", "sys.path.insert(0, str(FORK))\n", "from utils.parquet_source import CAPTION_SENTINELS # the fork's own set\n", "\n", "TRAIN_TARGET = 8192 # ~2.7x exp004's criticized 2978 rows\n", "EVAL_N = 512 # 10k steps @ mbs8 = 80k samples ~= 9 epochs\n", "TARGET = TRAIN_TARGET + EVAL_N\n", "KEEP_COLS = [\"id\", \"image\", \"image_width\", \"image_height\",\n", " \"caption\", \"fused_json\", \"age_audit\"]\n", "\n", "TRAIN_DIR = DATA / \"train\"; TRAIN_DIR.mkdir(exist_ok=True)\n", "EVAL_PARQUET = DATA / \"eval.parquet\"\n", "MANIFEST = DATA / \"split_manifest.json\"\n", "SRC = DATA / \"src\"; SRC.mkdir(exist_ok=True)\n", "_SENT = pa.array(sorted(CAPTION_SENTINELS))\n", "\n", "def _as_clean_str(tbl, name):\n", " \"\"\"Column -> trimmed string array with sentinels/nulls -> null.\"\"\"\n", " if name not in tbl.schema.names:\n", " return pa.nulls(tbl.num_rows, pa.string())\n", " col = pc.utf8_trim_whitespace(\n", " pc.fill_null(pc.cast(tbl.column(name), pa.string()), \"\"))\n", " bad = pc.is_in(col, value_set=_SENT)\n", " return pc.if_else(bad, pa.nulls(tbl.num_rows, pa.string()), col)\n", "\n", "def caption_vec(tbl):\n", " \"\"\"dexp011:154, vectorized: joycaption else prompt_fused else null.\"\"\"\n", " return pc.coalesce(_as_clean_str(tbl, \"caption_joycaption\"),\n", " _as_clean_str(tbl, \"prompt_fused\"))\n", "\n", "def gate_vec(tbl):\n", " \"\"\"exp011 audit gate, vectorized: every present audit column must\n", " contain approved|pass (rows with NO audit columns pass).\"\"\"\n", " mask = pa.array([True] * tbl.num_rows)\n", " for c in (\"age_audit\", \"audit\"):\n", " if c in tbl.schema.names:\n", " low = pc.utf8_lower(\n", " pc.fill_null(pc.cast(tbl.column(c), pa.string()), \"\"))\n", " ok = pc.or_(pc.match_substring(low, \"pass\"),\n", " pc.match_substring(low, \"approved\"))\n", " # null/absent value on a present column = not gated (as exp011:\n", " # gate_vals only collects non-None values)\n", " ok = pc.or_(ok, pc.equal(low, \"\"))\n", " mask = pc.and_(mask, ok)\n", " return mask\n", "\n", "if not MANIFEST.exists():\n", " shard_order = sorted(\n", " ANIMA[\"SHARDS\"],\n", " key=lambda n: hashlib.sha256(n.encode()).hexdigest())\n", "\n", " # \u2500\u2500 phase A: thin columnar scan to select the pool \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n", " pool_ids, seen = [], set()\n", " n_seen = n_dropped = n_dup = 0\n", " shards_used, local_paths = [], {}\n", " THIN = [\"id\", \"caption_joycaption\", \"prompt_fused\", \"age_audit\", \"audit\"]\n", " for shard in shard_order: # NOT `sh` \u2014 that is the shell fn\n", " if len(pool_ids) >= TARGET:\n", " break\n", " local = hf_hub_download(ANIMA[\"DATASET\"], shard, repo_type=\"dataset\",\n", " local_dir=str(SRC))\n", " local_paths[shard] = local\n", " shards_used.append(shard)\n", " pf = pq.ParquetFile(local)\n", " cols = [c for c in THIN if c in pf.schema_arrow.names]\n", " for rg in range(pf.num_row_groups):\n", " if len(pool_ids) >= TARGET:\n", " break\n", " tbl = pf.read_row_group(rg, columns=cols)\n", " n_seen += tbl.num_rows\n", " keep = pc.and_(gate_vec(tbl), pc.is_valid(caption_vec(tbl)))\n", " n_dropped += tbl.num_rows - pc.sum(keep).as_py()\n", " ids = tbl.column(\"id\").filter(keep).to_pylist() # THIN data only\n", " for rid in map(str, ids):\n", " if rid in seen:\n", " n_dup += 1\n", " elif len(pool_ids) < TARGET:\n", " seen.add(rid)\n", " pool_ids.append(rid)\n", " print(f\"[scan] {shard}: pool {len(pool_ids)}/{TARGET} \"\n", " f\"(seen {n_seen}, dropped {n_dropped}, dup {n_dup})\",\n", " flush=True)\n", "\n", " assert len(pool_ids) >= TARGET, (\n", " f\"exhausted ALL shards with only {len(pool_ids)} usable rows \"\n", " f\"(target {TARGET}) \u2014 inspect the gate before training\")\n", "\n", " ranked = sorted(pool_ids,\n", " key=lambda r: hashlib.sha256(r.encode()).hexdigest())\n", " eval_ids = pa.array(ranked[:EVAL_N])\n", " train_ids = pa.array(ranked[EVAL_N:TARGET])\n", "\n", " # \u2500\u2500 phase B: heavy columnar write, filtered by id sets \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n", " HEAVY = [\"id\", \"image\", \"image_width\", \"image_height\",\n", " \"caption_joycaption\", \"prompt_fused\", \"fused_json\", \"age_audit\"]\n", " writers = {}\n", " counts = {\"train\": 0, \"eval\": 0}\n", "\n", " def _emit(which, sub, path):\n", " cap = caption_vec(sub)\n", " out = pa.table({\n", " \"id\": sub.column(\"id\"), \"image\": sub.column(\"image\"),\n", " \"image_width\": sub.column(\"image_width\"),\n", " \"image_height\": sub.column(\"image_height\"),\n", " \"caption\": pc.fill_null(cap, \"\"),\n", " \"fused_json\": _f(sub, \"fused_json\"),\n", " \"age_audit\": _f(sub, \"age_audit\")})\n", " if which not in writers:\n", " writers[which] = pq.ParquetWriter(path, out.schema)\n", " writers[which].write_table(out)\n", " counts[which] += out.num_rows\n", "\n", " def _f(t, name):\n", " return (pc.fill_null(pc.cast(t.column(name), pa.string()), \"\")\n", " if name in t.schema.names\n", " else pa.array([\"\"] * t.num_rows))\n", "\n", " for shard in shards_used:\n", " pf = pq.ParquetFile(local_paths[shard])\n", " cols = [c for c in HEAVY if c in pf.schema_arrow.names]\n", " for rg in range(pf.num_row_groups):\n", " tbl = pf.read_row_group(rg, columns=cols)\n", " ids = pc.cast(tbl.column(\"id\"), pa.string())\n", " m_ev = pc.is_in(ids, value_set=eval_ids)\n", " m_tr = pc.is_in(ids, value_set=train_ids)\n", " if pc.any(m_ev).as_py():\n", " _emit(\"eval\", tbl.filter(m_ev), EVAL_PARQUET)\n", " if pc.any(m_tr).as_py():\n", " _emit(\"train\", tbl.filter(m_tr),\n", " TRAIN_DIR / \"train.parquet\")\n", " for w in writers.values():\n", " w.close()\n", " shutil.rmtree(SRC, ignore_errors=True) # sources no longer needed\n", "\n", " assert counts == {\"train\": TRAIN_TARGET, \"eval\": EVAL_N}, counts\n", " MANIFEST.write_text(json.dumps(\n", " {\"dataset\": ANIMA[\"DATASET\"],\n", " \"shard_order\": \"sha256(shard_name) ascending\",\n", " \"shards_used\": shards_used, \"rows_seen\": n_seen,\n", " \"dropped_gate_or_caption\": n_dropped, \"duplicate_ids\": n_dup,\n", " \"train_rows\": counts[\"train\"], \"eval_n\": counts[\"eval\"],\n", " \"eval_ids\": eval_ids.to_pylist(),\n", " \"caption_rule\": \"caption_joycaption or prompt_fused fallback \"\n", " \"(dexp011:154; fork CAPTION_SENTINELS applied)\",\n", " \"gate_rule\": \"age_audit/audit must contain approved|pass \"\n", " \"(dexp011:98-105)\",\n", " \"split_rule\": \"sha256(id) ascending within the pool; first \"\n", " \"EVAL_N are eval\"},\n", " indent=1))\n", " print(f\"split DONE: {counts['train']} train + {counts['eval']} eval \"\n", " f\"from {len(shards_used)}/{len(ANIMA['SHARDS'])} shards\")\n", "else:\n", " _m = json.loads(MANIFEST.read_text())\n", " print(f\"split already materialized: {_m['train_rows']} train + \"\n", " f\"{_m['eval_n']} eval from {len(_m['shards_used'])} shards\")\n", "\n", "# -- dataset TOMLs (fused recipe: derived caption column) -----------------\n", "# `path` = the latent/TE cache root, and it MUST differ between train and\n", "# eval: for type='parquet' the fork's default root is the CONSTANT\n", "# 'local_parquet' (utils/dataset.py:1014-1015), so without explicit paths\n", "# the eval dataset overwrites the train dataset's grouped metadata and\n", "# clears its latents mid-run \u2014 the train loop then silently loads the 512\n", "# EVAL rows. Key names per the fork + exp004's as-run dataset TOML:\n", "# width_column/height_column (image_*_column variants are silently ignored).\n", "def _dataset_toml(files, cache_path):\n", " return f\"\"\"# generated by anima_closeout.ipynb\n", "resolutions = [512]\n", "enable_ar_bucket = true\n", "min_ar = 0.5\n", "max_ar = 2.0\n", "num_ar_buckets = 7\n", "cache_backend = 'parquet'\n", "cache_shard_size_mb = 350\n", "\n", "[[directory]]\n", "type = 'parquet'\n", "path = '{cache_path}'\n", "parquet_files = '{files}'\n", "image_column = 'image'\n", "width_column = 'image_width'\n", "height_column = 'image_height'\n", "caption_column = 'caption'\n", "caption_type = 'text'\n", "num_repeats = 1\n", "skip_empty_caption = true\n", "\"\"\"\n", "\n", "(TOMLS / \"dataset_train.toml\").write_text(\n", " _dataset_toml(f\"{TRAIN_DIR}/*.parquet\", WORK / \"dp_cache\" / \"train\"))\n", "(TOMLS / \"dataset_eval.toml\").write_text(\n", " _dataset_toml(str(EVAL_PARQUET), WORK / \"dp_cache\" / \"heldout\"))\n", "# hygiene: any earlier run WITHOUT the path keys wrote train+eval into the\n", "# fork's shared default root and cross-contaminated it (same grouping keys,\n", "# same filenames \u2014 the ArrowIndexError smoke crash). Remove it so nothing\n", "# can ever load that poisoned state; harmless when absent.\n", "shutil.rmtree(Path.home() / \".cache\" / \"diffusion-pipe\", ignore_errors=True)\n", "# -- training TOMLs per arm -----------------------------------------------\n", "_AP = WORK / \"models\" / \"anima\" / \"split_files\"\n", "\n", "def _arm_toml(tag, seed, steps, mode=\"relay\", lora=False,\n", " dataset=\"dataset_train.toml\", evals=True):\n", " out = RUNS / tag\n", " L = [f\"# {tag} \u2014 generated by anima_closeout.ipynb\",\n", " f\"output_dir = '{out}'\",\n", " f\"dataset = '{TOMLS}/{dataset}'\",\n", " f\"seed = {seed}\",\n", " \"bf16_master_weights = true # the exp004 gate-freeze fix\",\n", " \"epochs = 1000\",\n", " f\"max_steps = {steps}\",\n", " \"micro_batch_size_per_gpu = 8\",\n", " \"pipeline_stages = 1\",\n", " \"gradient_accumulation_steps = 1\",\n", " \"gradient_clipping = 0\",\n", " \"save_every_n_steps = 2500\",\n", " \"checkpoint_every_n_minutes = 30\",\n", " \"activation_checkpointing = false # 96GB rider\",\n", " \"save_dtype = 'bfloat16'\",\n", " \"caching_batch_size = 8\"]\n", " if evals:\n", " L += [\n", " f\"eval_datasets = [{{name = 'heldout', config = '{TOMLS}/dataset_eval.toml'}}]\",\n", " \"eval_every_n_steps = 1000\",\n", " \"eval_before_first_step = true\",\n", " \"eval_micro_batch_size_per_gpu = 8\",\n", " \"eval_gradient_accumulation_steps = 1\"]\n", " L += [\n", " \"\", \"[model]\", \"type = 'anima'\",\n", " f\"transformer_path = '{_AP}/diffusion_models/anima-base-v1.0.safetensors'\",\n", " f\"llm_path = '{_AP}/text_encoders/qwen_3_06b_base.safetensors'\",\n", " f\"vae_path = '{_AP}/vae/qwen_image_vae.safetensors'\",\n", " \"dtype = 'bfloat16'\",\n", " \"timestep_sample_method = 'logit_normal'\"]\n", " if lora:\n", " L += [\"self_attn_lr = 1e-3\", \"cross_attn_lr = 1e-3\", \"mlp_lr = 1e-3\",\n", " \"mod_lr = 1e-3\", \"llm_adapter_lr = 1e-3\",\n", " \"\", \"[adapter]\", \"type = 'lora'\", \"rank = 16\",\n", " \"dtype = 'bfloat16'\"]\n", " else:\n", " L += [\"aleph_relay = true\",\n", " f\"aleph_relay_mode = '{mode}'\",\n", " \"aleph_relay_rank = 16\",\n", " \"aleph_relay_every = 1\",\n", " \"aleph_relay_lr = 1e-3\",\n", " \"self_attn_lr = 0\", \"cross_attn_lr = 0\", \"mlp_lr = 0\",\n", " \"mod_lr = 0\",\n", " \"llm_adapter_lr = 0 # NON-NEGOTIABLE (anima-trainer law)\"]\n", " L += [\"\", \"[optimizer]\",\n", " \"type = 'adam' # pure Adam \u2014 never adamw on aleph paths\",\n", " (\"lr = 1e-3\" if lora else \"lr = 0\"),\n", " \"betas = [0.9, 0.99]\",\n", " \"weight_decay = 0.0\",\n", " \"\", \"[monitoring]\", \"enable_wandb = false\"]\n", " p = TOMLS / f\"{tag}.toml\"\n", " p.write_text(\"\\n\".join(L) + \"\\n\")\n", " return p\n", "\n", "ARMS = {\n", " \"smoke\": _arm_toml(\"smoke\", 0, 200),\n", " \"e004b_relay_s0\": _arm_toml(\"e004b_relay_s0\", 0, 10000),\n", " \"e004b_relay_s1\": _arm_toml(\"e004b_relay_s1\", 1, 10000),\n", " \"e004b_lora_s0\": _arm_toml(\"e004b_lora_s0\", 0, 10000, lora=True),\n", " \"e004b_lora_s1\": _arm_toml(\"e004b_lora_s1\", 1, 10000, lora=True),\n", " \"e017_mb3_smoke\": _arm_toml(\"e017_mb3_smoke\", 0, 200, mode=\"multiband3\"),\n", " \"e017_mb3_s0\": _arm_toml(\"e017_mb3_s0\", 0, 10000, mode=\"multiband3\"),\n", " \"e017_mb3_s1\": _arm_toml(\"e017_mb3_s1\", 1, 10000, mode=\"multiband3\"),\n", " # cache-only: trains 1 throwaway step on the HELD-OUT set purely to\n", " # materialize dp_cache/heldout (latents + text embeds) for the val bed\n", " # on a runtime where every arm was restored and nothing else trains.\n", " # Same dataset TOML -> same cache root/fingerprints as the real runs.\n", " \"cache_eval\": _arm_toml(\"cache_eval\", 0, 1,\n", " dataset=\"dataset_eval.toml\", evals=False),\n", "}\n", "\n", "# THE registry: tag -> (hub package, arm folder). Drives BOTH shipping and\n", "# restoring, so a trained arm and its hub copy can never disagree about\n", "# where it lives. Arms absent here (the smokes) are local-only by design.\n", "ARM_SHIP = {\n", " \"e004b_relay_s0\": (\"exp004b_anima_relay\", \"relay_s0\"),\n", " \"e004b_relay_s1\": (\"exp004b_anima_relay\", \"relay_s1\"),\n", " \"e004b_lora_s0\": (\"exp004b_anima_relay\", \"lora_s0\"),\n", " \"e004b_lora_s1\": (\"exp004b_anima_relay\", \"lora_s1\"),\n", " \"e017_mb3_s0\": (\"exp017_anima_multiband\", \"mb3_s0\"),\n", " \"e017_mb3_s1\": (\"exp017_anima_multiband\", \"mb3_s1\"),\n", "}\n", "print(\"TOMLs written:\", \", \".join(ARMS))\n" ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "# \u2550\u2550\u2550 CELL 2 \u2014 helpers + SMOKE + the G-gate assert \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n", "# The smoke = 200 real-recipe steps. Jobs: prove the fp32-master fix in\n", "# flight (gates MUST move \u2014 the exact exp004 failure), prove eval wiring\n", "# (eval_before_first_step fires at step 0), build the latent cache\n", "# (~0.5h, reused by every later arm).\n", "import glob, io, json, os, shutil, subprocess, time\n", "from pathlib import Path\n", "import torch\n", "\n", "HUB_REPO = \"AbstractPhil/geolip-aleph-diffusion\"\n", "\n", "def _die(tag, log, why):\n", " tail = \"\".join(\n", " log.read_text(errors=\"replace\").splitlines(keepends=True)[-60:]) \\\n", " if log.exists() else \"(no log written)\"\n", " raise RuntimeError(f\"{tag}: {why}\\n--- last 60 lines of {log} ---\\n\"\n", " f\"{tail}\")\n", "\n", "def sh_tee(cmd, log):\n", " # NO shell pipe to tee: a pipeline returns TEE's rc, so a dead train\n", " # reported rc=0 (the smoke's exact failure). Python is the tee \u2014 the\n", " # child's rc is real, output streams live AND lands in the log.\n", " print(f\"$ {cmd}\")\n", " with open(log, \"a\") as lf:\n", " proc = subprocess.Popen(cmd, shell=True, executable=\"/bin/bash\",\n", " stdout=subprocess.PIPE,\n", " stderr=subprocess.STDOUT,\n", " text=True, bufsize=1)\n", " for line in proc.stdout:\n", " print(line, end=\"\", flush=True)\n", " lf.write(line)\n", " return proc.wait()\n", "\n", "def _arm_max_steps(tag):\n", " try:\n", " import tomllib\n", " except ImportError:\n", " import tomli as tomllib\n", " return tomllib.loads(ARMS[tag].read_text())[\"max_steps\"]\n", "\n", "# \u2500\u2500 arm resolution: LOCAL -> HUB -> train \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n", "# A trained arm is NEVER retrained. Local checkpoints win; otherwise a\n", "# previously shipped arm is pulled back from the hub (so a fresh runtime\n", "# after a pod cull costs a download, not a training run); only a genuinely\n", "# missing arm trains. ARM_SOURCE records which path each arm took.\n", "WEIGHT_FILES = (\"aleph_relay.pt\", \"adapter_model.safetensors\")\n", "ARM_SOURCE = globals().get(\"ARM_SOURCE\", {})\n", "\n", "def _curve_cache(tag):\n", " d = WORK / \"eval_curves\"\n", " d.mkdir(exist_ok=True)\n", " return d / f\"{tag}.json\"\n", "\n", "def _hub_files(refresh=False):\n", " if refresh or \"HUB_FILES\" not in ANIMA:\n", " from huggingface_hub import HfApi\n", " ANIMA[\"HUB_FILES\"] = sorted(HfApi().list_repo_files(HUB_REPO))\n", " return ANIMA[\"HUB_FILES\"]\n", "\n", "def local_save(tag):\n", " \"\"\"The arm's FINAL-step save on this disk, or None.\"\"\"\n", " done = sorted(\n", " glob.glob(str(RUNS / tag / \"*\" / f\"step{_arm_max_steps(tag)}\")))\n", " return Path(done[-1]) if done else None\n", "\n", "def restore_arm(tag):\n", " \"\"\"Pull a shipped arm back from the hub into a local save dir, or None.\n", " Also restores its eval curve, which no longer exists as tfevents.\"\"\"\n", " if tag not in ARM_SHIP:\n", " return None\n", " from huggingface_hub import hf_hub_download\n", " pkg, arm = ARM_SHIP[tag]\n", " prefix = f\"{pkg}/{arm}/\"\n", " files = [f for f in _hub_files() if f.startswith(prefix)]\n", " if not any(f.rsplit(\"/\", 1)[-1] in WEIGHT_FILES for f in files):\n", " return None\n", " dest = RUNS / tag / \"from_hub\" / f\"step{_arm_max_steps(tag)}\"\n", " dest.mkdir(parents=True, exist_ok=True)\n", " for f in files:\n", " name = f.rsplit(\"/\", 1)[-1]\n", " # skip the log and the DERIVED duplicates of the .pt (canonical +\n", " # fork-format safetensors): nothing downstream reads them, and they\n", " # would triple the restore for every arm. They stay on the hub.\n", " if name == \"log_tail.txt\" or name.endswith(\".canonical.safetensors\") \\\n", " or name == \"aleph_relay.safetensors\":\n", " continue\n", " p = Path(hf_hub_download(HUB_REPO, f))\n", " if name == \"eval_curve.json\":\n", " _curve_cache(tag).write_text(p.read_text())\n", " else:\n", " shutil.copy(p, dest / name)\n", " print(f\"[{tag}] restored from hub: {prefix} -> {dest}\")\n", " return dest\n", "\n", "def arm_save(tag, allow_hub=True):\n", " \"\"\"Resolve a finished arm WITHOUT ever training. None if unavailable.\"\"\"\n", " p = local_save(tag)\n", " if p is not None:\n", " ARM_SOURCE.setdefault(tag, \"local\")\n", " return p\n", " if allow_hub:\n", " p = restore_arm(tag)\n", " if p is not None:\n", " ARM_SOURCE[tag] = \"hub\"\n", " return p\n", " return None\n", "\n", "def arm_available(tag):\n", " return arm_save(tag) is not None\n", "\n", "def run_arm(tag, resume=False, force=False):\n", " if not force:\n", " p = arm_save(tag)\n", " if p is not None:\n", " print(f\"[{tag}] NOT retraining \u2014 reusing \"\n", " f\"{ARM_SOURCE[tag]} checkpoint {p}\")\n", " return p\n", " log = RUNS / f\"{tag}.log\"\n", " flags = \" --resume_from_checkpoint\" if resume else \"\"\n", " t0 = time.time()\n", " rc = sh_tee(f\"cd {FORK} && NCCL_P2P_DISABLE=1 NCCL_IB_DISABLE=1 \"\n", " f\"deepspeed --num_gpus=1 train.py --deepspeed \"\n", " f\"--config {ARMS[tag]}{flags}\", log)\n", " print(f\"[{tag}] rc={rc} wall={(time.time() - t0) / 60:.1f} min\")\n", " if rc != 0:\n", " _die(tag, log, f\"train exited rc={rc}\")\n", " ARM_SOURCE[tag] = \"trained\"\n", " try:\n", " return newest_save(tag) # rc=0 but no run/save dir is ALSO death\n", " except AssertionError as e:\n", " _die(tag, log, str(e))\n", "\n", "def latest_run_dir(tag, with_events=False, trained_only=False):\n", " dirs = sorted(glob.glob(str(RUNS / tag / \"*\")))\n", " if with_events or trained_only:\n", " # 'from_hub' is a restored checkpoint, not a run: no tfevents, and\n", " # it must never be mistaken for what a training launch just wrote\n", " dirs = [d for d in dirs if os.path.basename(d) != \"from_hub\"]\n", " if with_events:\n", " dirs = [d for d in dirs if glob.glob(\n", " os.path.join(d, \"**\", \"events.out.tfevents*\"), recursive=True)]\n", " assert dirs, (f\"no run dir for {tag}\"\n", " + (\" containing tfevents\" if with_events else \"\"))\n", " return Path(dirs[-1])\n", "\n", "def newest_save(tag):\n", " rd = latest_run_dir(tag, trained_only=True)\n", " saves = sorted(glob.glob(str(rd / \"epoch*\")) +\n", " glob.glob(str(rd / \"step*\")), key=os.path.getmtime)\n", " assert saves, f\"{tag}: run wrote no save dir\"\n", " return Path(saves[-1])\n", "\n", "def gate_values(save_dir):\n", " blob = torch.load(Path(save_dir) / \"aleph_relay.pt\",\n", " map_location=\"cpu\", weights_only=True)\n", " st = blob[\"relays\"]\n", " out = []\n", " for k in sorted(st, key=int):\n", " g = st[k].get(\"gate\")\n", " if g is None: # multiband: per-band gates\n", " g = st[k][\"gates\"]\n", " out.extend(torch.as_tensor(g).float().flatten().tolist())\n", " return out\n", "\n", "def eval_curve(tag, required=True):\n", " \"\"\"The fork's deterministic eval scalars. Prefers the cached/restored\n", " JSON (a hub-restored arm has no tfevents at all), else parses this\n", " runtime's tfevents and caches the result. None when unavailable and\n", " not required.\"\"\"\n", " cached = _curve_cache(tag)\n", " if cached.exists():\n", " return json.loads(cached.read_text())\n", " try:\n", " rd = latest_run_dir(tag, with_events=True)\n", " except AssertionError:\n", " if required:\n", " raise\n", " return None\n", " from tensorboard.backend.event_processing.event_accumulator import \\\n", " EventAccumulator\n", " acc = EventAccumulator(str(rd), size_guidance={\"scalars\": 0})\n", " acc.Reload()\n", " out = {}\n", " for t in acc.Tags()[\"scalars\"]:\n", " if t.startswith(\"heldout/\"):\n", " out[t] = [(e.step, e.value) for e in acc.Scalars(t)]\n", " assert \"heldout/loss\" in out, \\\n", " f\"{tag}: no heldout/loss in tfevents ({acc.Tags()['scalars']})\"\n", " cached.write_text(json.dumps(out))\n", " return out\n", "\n", "def final_loss(tag):\n", " \"\"\"Last held-out loss of an arm's fork eval curve, or None.\"\"\"\n", " c = eval_curve(tag, required=False)\n", " return c[\"heldout/loss\"][-1][1] if c else None\n", "\n", "def ship_tag(tag, save_dir):\n", " \"\"\"Ship an arm under its registry identity \u2014 skipped when the arm CAME\n", " from the hub (re-uploading a restored copy is pure waste). The dir name\n", " is checked too, so the skip survives a kernel restart.\"\"\"\n", " if ARM_SOURCE.get(tag) == \"hub\" or Path(save_dir).parent.name == \"from_hub\":\n", " print(f\"[ship] {tag}: skipped (restored from hub, already there)\")\n", " return\n", " pkg, arm = ARM_SHIP[tag]\n", " curve = eval_curve(tag, required=False)\n", " ship_arm(pkg, arm, save_dir,\n", " extra=({f\"{pkg}/{arm}/eval_curve.json\": curve} if curve else None))\n", "\n", "def ship_arm(package, arm, save_dir, extra=None):\n", " \"\"\"Ship-on-completion: ckpt (+canonical safetensors) + TOML + log tail\n", " into the hub package immediately.\"\"\"\n", " from huggingface_hub import CommitOperationAdd, HfApi\n", " from amoe.io.checkpoint import load_diffusion_anchor\n", " from amoe.io.safetensors_io import save_anchor_safetensors\n", " sd = Path(save_dir)\n", " ops = []\n", " for f in sd.iterdir():\n", " if f.suffix in (\".pt\", \".safetensors\", \".toml\") or \\\n", " f.name == \"adapter_config.json\":\n", " ops.append((f\"{package}/{arm}/{f.name}\", str(f)))\n", " relay_pt = sd / \"aleph_relay.pt\"\n", " if relay_pt.exists():\n", " ck = load_diffusion_anchor(str(relay_pt), substrate={\n", " \"family\": \"cosmos_dit\",\n", " \"base_model_id\": \"circlestone-labs/Anima@anima-base-v1.0\"})\n", " canon = sd / \"aleph_relay.canonical.safetensors\"\n", " save_anchor_safetensors(ck, str(canon))\n", " ops.append((f\"{package}/{arm}/aleph_relay.canonical.safetensors\",\n", " str(canon)))\n", " tag_guess = sd.parent.parent.name\n", " log = RUNS / f\"{tag_guess}.log\"\n", " if log.exists():\n", " tail = \"\\n\".join(log.read_text(errors=\"replace\").splitlines()[-400:])\n", " ops.append((f\"{package}/{arm}/log_tail.txt\",\n", " io.BytesIO(tail.encode())))\n", " for repo_path, content in (extra or {}).items():\n", " body = (io.BytesIO(json.dumps(content, indent=1).encode())\n", " if isinstance(content, (dict, list)) else content)\n", " ops.append((repo_path, body))\n", " HfApi().create_commit(\n", " repo_id=HUB_REPO,\n", " operations=[CommitOperationAdd(path_in_repo=r, path_or_fileobj=s)\n", " for r, s in ops],\n", " commit_message=f\"{package}/{arm}: ship-on-completion from Colab\")\n", " ANIMA.pop(\"HUB_FILES\", None) # listing is now stale\n", " print(f\"[ship] {package}/{arm}: {len(ops)} file(s) -> {HUB_REPO}\")\n", "\n", "# -- campaign inventory ---------------------------------------------------\n", "# refuse to launch on stale dataset TOMLs from an earlier notebook version:\n", "# without distinct dp_cache paths, train+eval share one cache root and\n", "# corrupt each other (the ArrowIndexError crash / silent eval-swap hazard)\n", "for _dt in (\"dataset_train.toml\", \"dataset_eval.toml\"):\n", " assert \"dp_cache\" in (TOMLS / _dt).read_text(), (\n", " f\"{_dt} lacks the distinct cache path \u2014 RE-RUN CELL 1 in THIS \"\n", " \"kernel before any arm; stale TOMLs corrupt the caches\")\n", "\n", "INVENTORY = {}\n", "for _t in ARM_SHIP: # resolves + materializes each arm\n", " INVENTORY[_t] = ARM_SOURCE[_t] if arm_save(_t) else \"MISSING\"\n", "print(\"campaign inventory:\")\n", "for _t, _w in INVENTORY.items():\n", " print(f\" {_t:<18} {_w}\")\n", "_todo = [t for t, w in INVENTORY.items() if w == \"MISSING\"]\n", "print(f\"arms needing training: {_todo or 'NONE \u2014 everything is available'}\")\n", "\n", "# -- SMOKE ----------------------------------------------------------------\n", "# Its jobs: prove gates move under fp32 masters (exp004's exact failure),\n", "# prove eval wiring, build the latent cache. All three are SUPERSEDED once\n", "# a full aleph arm exists \u2014 so it only runs when one is actually needed.\n", "_ALEPH = (\"e004b_relay_s0\", \"e004b_relay_s1\", \"e017_mb3_s0\", \"e017_mb3_s1\")\n", "if any(INVENTORY[t] != \"MISSING\" for t in _ALEPH):\n", " print(\"SMOKE SKIPPED \u2014 a full aleph arm is already available, which \"\n", " \"proves the gate fix at 10k steps (a stronger result than 200).\")\n", "else:\n", " _save = run_arm(\"smoke\")\n", " _g = gate_values(_save)\n", " _moved = [x for x in _g if abs(x + 3.0) > 1e-3]\n", " print(f\"G-GATE: {len(_moved)}/{len(_g)} gates moved off -3.0 \"\n", " f\"(range {min(_g):+.4f} .. {max(_g):+.4f})\")\n", " assert _moved, (\n", " \"G-GATE FAILED: gates did not move \u2014 fp32 masters are not effective \"\n", " \"here. STOP. The designed fallback is the gate delta-reparam \"\n", " \"(gate = -3 + near-zero delta); do not burn full arms.\")\n", " _c = eval_curve(\"smoke\")\n", " print(f\"eval wiring OK: {len(_c)} heldout tags, \"\n", " f\"step-0 loss {_c['heldout/loss'][0][1]:.5f}\")\n", " print(\"SMOKE: ALL GREEN \u2014 masters proven, eval proven, cache built\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## exp004b \u2014 the relay close-out\n", "\n", "Preregistration, stated before spend:\n", "- **P1** relay \u2265 matched LoRA-r16 on the paired held-out eval\n", "- **P2** toggle bit-exact post-train (val battery, CELL 6)\n", "- **P3** gates MOVE at full scale (the exp004 failure, asserted per arm)\n", "- **P4** the old ~3k-step loss is a floor: the 10k eval curve descends\n", " past it\n", "\n", "LoRA control s1 runs only if the s0 relay-vs-lora gap is within ~2\u00d7 the\n", "relay seed spread (results-gated, disclosed either way).\n" ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "# \u2550\u2550\u2550 CELL 3 \u2014 exp004b: relay s0, relay s1, LoRA control \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n", "# Every arm resolves LOCAL -> HUB -> train, so re-running this cell after a\n", "# crash (or on a fresh runtime) costs downloads, not GPU-hours.\n", "relay_saves = {}\n", "for _tag in (\"e004b_relay_s0\", \"e004b_relay_s1\"):\n", " _save = run_arm(_tag) # run_arm(_tag, resume=True) to\n", " _g = gate_values(_save) # continue a half-done run\n", " _mv = [x for x in _g if abs(x + 3.0) > 1e-3]\n", " print(f\"[{_tag}] P3 gates moved: {len(_mv)}/{len(_g)} \"\n", " f\"(range {min(_g):+.4f}..{max(_g):+.4f}) [{ARM_SOURCE[_tag]}]\")\n", " assert _mv, f\"{_tag}: P3 FAILED \u2014 gates frozen at full scale\"\n", " ship_tag(_tag, _save)\n", " relay_saves[_tag] = str(_save)\n", "\n", "lora_saves = {\"e004b_lora_s0\": str(run_arm(\"e004b_lora_s0\"))}\n", "ship_tag(\"e004b_lora_s0\", Path(lora_saves[\"e004b_lora_s0\"]))\n", "\n", "_r0, _r1 = final_loss(\"e004b_relay_s0\"), final_loss(\"e004b_relay_s1\")\n", "_l0 = final_loss(\"e004b_lora_s0\")\n", "if None in (_r0, _r1, _l0):\n", " print(\"results gate SKIPPED \u2014 missing eval curve(s) for \"\n", " f\"{[t for t, v in zip(('relay_s0','relay_s1','lora_s0'), (_r0,_r1,_l0)) if v is None]}\"\n", " \" (arms restored without one); LoRA s1 decision deferred to CELL 6\")\n", "else:\n", " _spread = abs(_r0 - _r1)\n", " _gap = abs(min(_r0, _r1) - _l0)\n", " print(f\"relay s0 {_r0:.5f} | s1 {_r1:.5f} (spread {_spread:.5f}) | \"\n", " f\"lora s0 {_l0:.5f} (gap {_gap:.5f})\")\n", " if _gap <= 2 * max(_spread, 1e-6):\n", " print(\"GATE: gap within 2x seed spread -> LoRA s1 REQUIRED\")\n", " lora_saves[\"e004b_lora_s1\"] = str(run_arm(\"e004b_lora_s1\"))\n", " ship_tag(\"e004b_lora_s1\", Path(lora_saves[\"e004b_lora_s1\"]))\n", " else:\n", " print(\"GATE: ordering decisive at s0 -> LoRA s1 skipped (disclosed)\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## exp017 \u2014 multiband3 on the Anima DiT\n", "\n", "First live run of `aleph_relay_mode='multiband3'` (CPU-smoked only until\n", "now). Preregistration:\n", "- **P1** band lesions surgical on the paired eval: own-band damage \u226510\u00d7\n", " cross-band at matching quantiles (the exp008 instrument on a DiT)\n", "- **P2** toggle bit-exact; every-band-lesioned bit-exact\n", "- On record: the aggregate loss may well favor a monolith (Law-1) \u2014 exp017\n", " certifies the *mechanism*, not the loss number.\n" ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "# \u2550\u2550\u2550 CELL 4 \u2014 exp017: mb3 smoke, then both seeds \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n", "from amoe.io.checkpoint import load_diffusion_anchor\n", "\n", "# the smoke exists to prove the mode ENGAGES before spending two seeds on\n", "# it; a real mb3 arm already in hand proves the same thing more strongly\n", "if all(arm_available(t) for t in (\"e017_mb3_s0\", \"e017_mb3_s1\")):\n", " print(\"mb3 smoke SKIPPED \u2014 both mb3 seed arms are already available\")\n", "else:\n", " _save = run_arm(\"e017_mb3_smoke\")\n", " _ck = load_diffusion_anchor(str(Path(_save) / \"aleph_relay.pt\"))\n", " assert _ck.kind == \"multiband3\", \\\n", " f\"mb3 smoke saved kind={_ck.kind} \u2014 the mode did not engage\"\n", " _g = gate_values(_save)\n", " _mv = [x for x in _g if abs(x + 3.0) > 1e-3]\n", " print(f\"mb3 smoke: kind OK, {_ck.n_sites} sites, gates moved \"\n", " f\"{len(_mv)}/{len(_g)}\")\n", " assert _mv, \"mb3 smoke: gates frozen \u2014 same STOP rule as the G-gate\"\n", "\n", "mb3_saves = {}\n", "for _tag in (\"e017_mb3_s0\", \"e017_mb3_s1\"):\n", " _save = run_arm(_tag)\n", " _ck = load_diffusion_anchor(str(Path(_save) / \"aleph_relay.pt\"))\n", " assert _ck.kind == \"multiband3\", \\\n", " f\"{_tag}: kind={_ck.kind} \u2014 not a multiband3 stack\"\n", " ship_tag(_tag, _save)\n", " mb3_saves[_tag] = str(_save)\n", " print(f\"[{_tag}] OK ({ARM_SOURCE[_tag]}), {_ck.n_sites} sites\")\n", "print(\"exp017 arms complete:\", list(mb3_saves))\n" ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "# \u2550\u2550\u2550 CELL 5 \u2014 the paired-val bed source (own cell, reviewable) \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n", "_BED_SRC = r'''\n", "# val_bed.py \u2014 paired val battery (subprocess; env-configured, no argparse).\n", "# Does what the fork eval cannot: toggle bit-exact asserts, fp32 gate\n", "# stats, band-lesion battery \u2014 on IDENTICAL (row, noise, t) triples per\n", "# arm. Reuses the fork's eval cache (latents + text embeds) and the fork's\n", "# own transformer + llm_adapter call path (mirrored from\n", "# cosmos_predict2.py InitialLayer/LLMAdapterLayer).\n", "import json, math, os, sys, time\n", "sys.path.insert(0, \".\") # fork root (cwd)\n", "import utils.common # bind fork utils BEFORE ComfyUI\n", "sys.path.append(\"submodules/ComfyUI\")\n", "os.environ.setdefault(\"MASTER_ADDR\", \"127.0.0.1\")\n", "os.environ.setdefault(\"MASTER_PORT\", \"29571\")\n", "os.environ.setdefault(\"RANK\", \"0\")\n", "os.environ.setdefault(\"WORLD_SIZE\", \"1\")\n", "os.environ.setdefault(\"LOCAL_RANK\", \"0\")\n", "import deepspeed\n", "deepspeed.init_distributed()\n", "import torch\n", "import pyarrow.parquet as pq\n", "\n", "# EXACTLY train.py:287, and it MUST happen before any model module is\n", "# imported: models/*.py do `from utils.common import AUTOCAST_DTYPE`, so\n", "# their @torch.autocast('cuda', dtype=AUTOCAST_DTYPE) decorators bake the\n", "# value at class-definition time. Left at its None default it resolves to\n", "# FP16 on cuda, and a bf16-trained 2B DiT overflows to NaN \u2014 which is\n", "# exactly what an unset bed produced (loss nan, and the determinism assert\n", "# then misfired because nan != nan).\n", "utils.common.AUTOCAST_DTYPE = torch.bfloat16\n", "assert \"models.cosmos_predict2\" not in sys.modules, \\\n", " \"a model module was imported before AUTOCAST_DTYPE was set\"\n", "\n", "CFG = json.load(open(os.environ[\"VB_CONFIG\"]))\n", "QUANTILES = [0.1, 0.3, 0.5, 0.7, 0.9]\n", "DEV = \"cuda\"\n", "\n", "# -- read the fork's eval cache -------------------------------------------\n", "# LAYOUT (utils/dataset.py) and the trap inside it:\n", "# {root}/cache_{ar:.3f}x{w}x{h}x{f}/latents/ SizeBucketDataset(:247)\n", "# {root}/ar_frames_{ar:.3f}_{f}/text_embeddings_N/ ARBucketDataset (:485)\n", "# Latents are cached from metadata that SizeBucketDataset SHUFFLES first\n", "# (:241-242); text embeddings are cached one level up, at the AR bucket,\n", "# unshuffled. So latent row i and TE row i are DIFFERENT images \u2014 pairing\n", "# by row order would silently hand every latent someone else's caption and\n", "# still produce plausible numbers. Both caches carry image_spec\n", "# (parquet_cache.image_specs()), which is the key the fork itself joins on\n", "# (dataset.py:312), so we join on image_spec and iterate in sorted spec\n", "# order, which also makes the bed reproducible run-to-run.\n", "cache_root = CFG[\"eval_cache_dir\"]\n", "# DISCOVER both kinds of cache by walking the tree. Do NOT derive one\n", "# directory name from the other: size buckets are\n", "# cache_{ar:.3f}x{w}x{h}x{f} but AR buckets are ar_frames_{ar:.3f}_{f} \u2014\n", "# 'x' vs '_' (bucket_suffix, dataset.py:60-70). image_spec is globally\n", "# unique per image, so one map over every text-embedding cache joins any\n", "# latents cache unambiguously, whatever the directories are called.\n", "lat_dirs, te_dirs = [], []\n", "for _dp, _dns, _ in os.walk(cache_root):\n", " for _d in _dns:\n", " if _d == \"latents\":\n", " lat_dirs.append(os.path.join(_dp, _d))\n", " elif _d.startswith(\"text_embeddings_\"): # excludes uncond_*\n", " te_dirs.append(os.path.join(_dp, _d))\n", "lat_dirs.sort(); te_dirs.sort()\n", "print(\"latents caches: \",\n", " [os.path.relpath(d, cache_root) for d in lat_dirs], flush=True)\n", "print(\"text-embed caches:\",\n", " [os.path.relpath(d, cache_root) for d in te_dirs], flush=True)\n", "assert lat_dirs, f\"no latents cache under {cache_root}\"\n", "assert te_dirs, f\"no text_embeddings_* cache under {cache_root}\"\n", "\n", "def read_cache(d):\n", " shards = sorted(p for p in os.listdir(d)\n", " if p.startswith(\"shard_\") and p.endswith(\".parquet\"))\n", " assert shards, f\"no shards under {d}\"\n", " return [pq.read_table(os.path.join(d, p)) for p in shards]\n", "\n", "def decol(tbl, name, j):\n", " raw = tbl.column(name)[j].as_py()\n", " shape = tbl.column(name + \"__shape\")[j].as_py()\n", " dt = str(tbl.column(name + \"__dtype\")[j].as_py()).replace(\"torch.\", \"\")\n", " return torch.frombuffer(bytearray(raw),\n", " dtype=getattr(torch, dt)).reshape(shape)\n", "\n", "def by_spec(tables, what):\n", " \"\"\"{image_spec tuple: (table, row)} \u2014 tuples are JSON-encoded by the\n", " cache writer (parquet_cache.py:_encode_row), so accept either form.\"\"\"\n", " out = {}\n", " for t in tables:\n", " names = t.schema.names\n", " col = (\"image_spec__json\" if \"image_spec__json\" in names else\n", " (\"image_spec\" if \"image_spec\" in names else None))\n", " assert col, f\"{what} cache carries no image_spec column: {names[:8]}\"\n", " for j, v in enumerate(t.column(col).to_pylist()):\n", " s = tuple(json.loads(v)) if isinstance(v, str) else tuple(v)\n", " out.setdefault(s, (t, j))\n", " return out\n", "\n", "LAT = [\"latents\"]\n", "TE = [\"prompt_embeds\", \"attn_mask\", \"t5_input_ids\", \"t5_attn_mask\"]\n", "\n", "te_map = {}\n", "for d in te_dirs:\n", " tbs = read_cache(d)\n", " for c in TE:\n", " assert c + \"__shape\" in tbs[0].schema.names, \\\n", " (d, tbs[0].schema.names[:8])\n", " for s, loc in by_spec(tbs, \"text embedding\").items():\n", " te_map.setdefault(s, loc)\n", "print(f\"text embeddings: {len(te_map)} unique image_spec\", flush=True)\n", "\n", "rows = []\n", "for d in lat_dirs:\n", " tbs = read_cache(d)\n", " for c in LAT:\n", " assert c + \"__shape\" in tbs[0].schema.names, \\\n", " (d, tbs[0].schema.names[:8])\n", " lat = by_spec(tbs, \"latents\")\n", " _missing = [s for s in lat if s not in te_map]\n", " assert not _missing, (\n", " f\"{os.path.relpath(d, cache_root)}: {len(_missing)} latents have \"\n", " f\"no text embedding for their image_spec (e.g. {_missing[:1]}) \u2014 \"\n", " \"caches are out of sync\")\n", " for s in sorted(lat):\n", " lt, lj = lat[s]\n", " tt, tj = te_map[s]\n", " r = {c: decol(lt, c, lj) for c in LAT}\n", " r.update({c: decol(tt, c, tj) for c in TE})\n", " rows.append(r)\n", " print(f\" {os.path.relpath(d, cache_root)}: {len(lat)} rows joined \"\n", " \"on image_spec\", flush=True)\n", "print(f\"eval cache rows: {len(rows)} from {len(lat_dirs)} latents cache(s)\",\n", " flush=True)\n", "rows = rows[:int(os.environ.get(\"VB_MAX_ROWS\", \"512\"))]\n", "assert rows\n", "\n", "# -- pipeline (frozen trunk); adapters swapped per arm --------------------\n", "from models import cosmos_predict2\n", "mc = CFG[\"model_cfg\"]\n", "mc[\"model\"][\"dtype\"] = torch.bfloat16\n", "pipe = cosmos_predict2.CosmosPredict2Pipeline(mc)\n", "pipe.load_diffusion_model()\n", "tf = pipe.transformer.to(DEV).eval()\n", "from models.aleph_relay import attach_aleph_relays\n", "\n", "# THE fork's own execution path (cosmos_predict2.py:460-470):\n", "# InitialLayer -> LLMAdapterLayer -> 28x TransformerLayer -> FinalLayer\n", "# Calling transformer.forward() directly is NOT what the trainer does and\n", "# does not work: concat_padding_mask is on, so forward()'s padding_mask=None\n", "# default crashes in prepare_embedded_sequence \u2014 InitialLayer builds a zero\n", "# padding mask itself (:612). The layer path also applies the llm_adapter\n", "# (LLMAdapterLayer) and sets each block's multiband window from the actual\n", "# timesteps (TransformerLayer:663-668), so the bed inherits all of it\n", "# instead of re-implementing it. The layers hold references to tf.blocks,\n", "# so adapters attached per-arm below are picked up without rebuilding.\n", "LAYERS = pipe.to_layers()\n", "print(f\"model layers: {len(LAYERS)} \"\n", " f\"({type(LAYERS[0]).__name__} .. {type(LAYERS[-1]).__name__})\",\n", " flush=True)\n", "\n", "def clear_adapters():\n", " for b in tf.blocks:\n", " for attr in (\"aleph_relay\", \"aleph_w_bands\"):\n", " if hasattr(b, attr):\n", " delattr(b, attr)\n", "\n", "# conditioning exactly as prepare_inputs hands it over (:412): the CACHED\n", "# embeds, un-adapted \u2014 LLMAdapterLayer does the adapter call itself.\n", "conds = []\n", "for r in rows:\n", " e = r[\"prompt_embeds\"]\n", " if e.ndim == 2:\n", " e = e.unsqueeze(0)\n", " conds.append((e, r[\"attn_mask\"].reshape(1, -1),\n", " r[\"t5_input_ids\"].reshape(1, -1),\n", " r[\"t5_attn_mask\"].reshape(1, -1)))\n", "print(f\"conditioning prepared for {len(conds)} rows\", flush=True)\n", "\n", "# fixed noise bank + the fork's exact eval-t transform:\n", "# t = sigmoid(Normal.icdf(q)) (logit_normal, scale 1, no shift \u2014\n", "# cosmos_predict2.py:436-444)\n", "torch.manual_seed(1400)\n", "lats, noises = [], []\n", "for r in rows:\n", " l = r[\"latents\"].float()\n", " while l.ndim < 5:\n", " l = l.unsqueeze(0)\n", " lats.append(l)\n", " noises.append(torch.randn_like(l))\n", "from torch.distributions import Normal\n", "T_OF_Q = {q: float(torch.sigmoid(Normal(0.0, 1.0).icdf(torch.tensor(q))))\n", " for q in QUANTILES}\n", "print(\"t(q):\", {q: round(t, 4) for q, t in T_OF_Q.items()}, flush=True)\n", "\n", "# -- batching ------------------------------------------------------------\n", "# Batch-1 would leave a 96GB card ~99% idle and cost ~1.5 GPU-hours across\n", "# both beds. Rows are grouped by tensor shape (3 AR buckets; prompts are\n", "# always 512 tokens \u2014 _tokenize pads to max_length, :157-164) and the noise\n", "# bank is still drawn per row in the original seeded order, so the\n", "# (row, noise, t) triples are BIT-IDENTICAL to the unbatched bed; only the\n", "# forward grouping changes, identically for every arm, so the toggle\n", "# equality assert stays exact. Loss is reduced PER SAMPLE and averaged with\n", "# equal weight per row, so batch composition cannot shift the numbers.\n", "VB_BATCH = int(os.environ.get(\"VB_BATCH\", \"16\"))\n", "_groups = {}\n", "for _i, (_l, _cd) in enumerate(zip(lats, conds)):\n", " _groups.setdefault(\n", " (tuple(_l.shape[1:]), _cd[0].shape[1], _cd[2].shape[1]), []\n", " ).append(_i)\n", "BATCHES = []\n", "for _k, _idxs in _groups.items():\n", " for _s in range(0, len(_idxs), VB_BATCH):\n", " BATCHES.append(_idxs[_s:_s + VB_BATCH])\n", "assert sum(len(b) for b in BATCHES) == len(lats)\n", "print(f\"batches: {len(BATCHES)} (max {VB_BATCH}/batch) over \"\n", " f\"{len(_groups)} shape group(s) sized {[len(v) for v in _groups.values()]}\",\n", " flush=True)\n", "\n", "results = {\"quantiles\": QUANTILES, \"t_of_q\": T_OF_Q, \"n_rows\": len(rows),\n", " \"arms\": {}}\n", "\n", "def eval_arm(name, lesion=None):\n", " t0 = time.time()\n", " per_q = {}\n", " with torch.no_grad():\n", " for q in QUANTILES:\n", " t = T_OF_Q[q]\n", " tot, n = 0.0, 0\n", " for bidx in BATCHES:\n", " lb = torch.cat([lats[i] for i in bidx]) # CPU float32\n", " nb = torch.cat([noises[i] for i in bidx])\n", " ld, nd = lb.to(DEV), nb.to(DEV) # float32, as :409\n", " x_t = (1 - t) * ld + t * nd\n", " ts = torch.full((ld.shape[0], 1), t, device=DEV,\n", " dtype=torch.float32) # (B,1), as :456\n", " out = (x_t, ts,\n", " torch.cat([conds[i][0] for i in bidx]).to(DEV),\n", " torch.cat([conds[i][1] for i in bidx]).to(DEV),\n", " torch.cat([conds[i][2] for i in bidx]).to(DEV),\n", " torch.cat([conds[i][3] for i in bidx]).to(DEV))\n", " for layer in LAYERS:\n", " out = layer(out)\n", " pred = out[0] if isinstance(out, (list, tuple)) else out\n", " se = (((pred.float().cpu() - (nb - lb)) ** 2)\n", " .flatten(1).mean(1)) # one loss PER ROW\n", " tot += se.double().sum().item() # fp64 accumulate: the\n", " n += se.numel() # total can't drift with\n", " # batch composition\n", " per_q[str(q)] = tot / n\n", " results[\"arms\"][name] = {\"loss_by_quantile\": per_q,\n", " \"loss_mean\": sum(per_q.values()) / len(per_q),\n", " \"lesion\": lesion}\n", " # a non-finite loss is a NUMERICS failure, not a determinism failure \u2014\n", " # say so, instead of letting nan != nan trip the assert below\n", " assert all(math.isfinite(v) for v in per_q.values()), (\n", " f\"{name}: non-finite loss {per_q} \u2014 the model is producing NaN/inf. \"\n", " \"First suspect: AUTOCAST_DTYPE (must be bf16 BEFORE models import).\")\n", " print(f\"[{name}] mean {results['arms'][name]['loss_mean']:.6f} \"\n", " f\"({(time.time() - t0) / 60:.1f} min)\", flush=True)\n", "\n", "clear_adapters()\n", "eval_arm(\"frozen\")\n", "f_ref = dict(results[\"arms\"][\"frozen\"][\"loss_by_quantile\"])\n", "eval_arm(\"frozen_repeat\")\n", "assert results[\"arms\"][\"frozen_repeat\"][\"loss_by_quantile\"] == f_ref, \\\n", " \"DETERMINISM FAILED \u2014 paired comparison meaningless; stop\"\n", "\n", "for name, arm in CFG[\"arms\"].items():\n", " clear_adapters()\n", " mode = arm.get(\"kind\", \"relay\")\n", " attach_aleph_relays(tf, tf.model_channels, relay_path=arm[\"ckpt\"],\n", " dtype=torch.bfloat16, mode=mode)\n", " # attach_aleph_relays casts dtype but does NOT place on device (the\n", " # trainer attaches on CPU, then DeepSpeed moves the whole engine). The\n", " # trunk is already on cuda, so the freshly-attached relay params are on\n", " # CPU -> 'mat2 is on cpu' at the first relay matmul. Re-place the trunk;\n", " # .to() is idempotent for params already on DEV.\n", " tf.to(DEV)\n", " gates = []\n", " for b in tf.blocks:\n", " g = getattr(b.aleph_relay, \"gate\", None)\n", " if g is None:\n", " g = b.aleph_relay.gates\n", " gates.extend(torch.as_tensor(g).float().flatten().tolist())\n", " for b in tf.blocks: # toggle assert per arm\n", " b.aleph_relay.enabled = False\n", " eval_arm(name + \"_toggled_off\")\n", " assert results[\"arms\"][name + \"_toggled_off\"][\"loss_by_quantile\"] \\\n", " == f_ref, f\"{name}: TOGGLE LAW VIOLATED\"\n", " for b in tf.blocks:\n", " b.aleph_relay.enabled = True\n", " if arm.get(\"lesion\") is not None:\n", " b.aleph_relay.band_enabled[arm[\"lesion\"]] = False\n", " eval_arm(name, lesion=arm.get(\"lesion\"))\n", " results[\"arms\"][name][\"gates_fp32\"] = {\n", " \"min\": min(gates), \"max\": max(gates),\n", " \"moved\": sum(1 for g in gates if abs(g + 3.0) > 1e-3)}\n", "\n", "json.dump(results, open(CFG[\"out\"], \"w\"), indent=1)\n", "print(\"WROTE\", CFG[\"out\"], flush=True)\n", "'''\n", "VAL_BED = WORK / \"val_bed.py\"\n", "VAL_BED.write_text(_BED_SRC)\n", "print(f\"val_bed.py written: {VAL_BED} ({len(_BED_SRC)/1000:.1f}k chars)\")\n" ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "# \u2550\u2550\u2550 CELL 6 \u2014 run the val battery + assemble results.json \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n", "import glob, json\n", "try:\n", " import tomllib\n", "except ImportError:\n", " import tomli as tomllib\n", "\n", "# deterministic: the eval dataset TOML pins path = dp_cache/heldout, and\n", "# the fork caches under {path}/cache/{model_name} with model_name 'anima'\n", "# (cosmos_predict2.py:254 renames the pipeline when llm_path is a file)\n", "EVAL_CACHE = str(WORK / \"dp_cache\" / \"heldout\" / \"cache\" / \"anima\")\n", "\n", "def ensure_eval_cache():\n", " \"\"\"The bed reads the fork's OWN eval cache. On a runtime where every\n", " arm was restored from the hub, nothing has trained, so that cache does\n", " not exist \u2014 rebuild it with a 1-step throwaway run over the held-out\n", " TOML (same TOML -> same root, same fingerprints; ~1 min, 512 rows).\n", " The throwaway checkpoint is never read by anything.\"\"\"\n", " if sorted(glob.glob(str(Path(EVAL_CACHE) / \"cache_*\"))):\n", " return\n", " print(\"eval cache absent (all arms restored) \u2014 building it now\")\n", " run_arm(\"cache_eval\", force=True)\n", " assert sorted(glob.glob(str(Path(EVAL_CACHE) / \"cache_*\"))), \\\n", " f\"cache_eval ran but produced nothing under {EVAL_CACHE}\"\n", "\n", "ensure_eval_cache()\n", "_buckets = sorted(glob.glob(str(Path(EVAL_CACHE) / \"cache_*\")))\n", "assert _buckets, f\"no cache_* bucket dirs under {EVAL_CACHE}\"\n", "print(\"eval cache:\", EVAL_CACHE,\n", " \"| buckets:\", [Path(b).name for b in _buckets])\n", "\n", "# resolve every arm independently of CELLs 3/4 having run in THIS kernel\n", "relay_saves = {t: str(arm_save(t)) for t in\n", " (\"e004b_relay_s0\", \"e004b_relay_s1\")}\n", "mb3_saves = {t: str(arm_save(t)) for t in (\"e017_mb3_s0\", \"e017_mb3_s1\")}\n", "assert all(relay_saves.values()) and all(mb3_saves.values()), \\\n", " \"run CELLs 3 and 4 first \u2014 some arms are neither local nor on the hub\"\n", "print(\"arms resolved:\", {t: ARM_SOURCE[t] for t in\n", " list(relay_saves) + list(mb3_saves)})\n", "\n", "_full = tomllib.loads(ARMS[\"e004b_relay_s0\"].read_text())\n", "_model = {k: v for k, v in _full[\"model\"].items()\n", " if not k.startswith(\"aleph_\") and not k.endswith(\"_lr\")}\n", "_model_cfg = {\"model\": _model} # the bed attaches adapters itself\n", "\n", "def _bed(config, out_name):\n", " cfgp = WORK / f\"vb_{out_name}.json\"\n", " outp = WORK / f\"{out_name}.json\"\n", " config.update(out=str(outp), eval_cache_dir=EVAL_CACHE,\n", " model_cfg=_model_cfg)\n", " cfgp.write_text(json.dumps(config, indent=1))\n", " _blog = RUNS / f\"{out_name}.log\"\n", " rc = sh_tee(f\"cd {FORK} && VB_CONFIG={cfgp} python {WORK}/val_bed.py\",\n", " _blog)\n", " if rc != 0 or not outp.exists():\n", " _die(out_name, _blog, f\"val bed rc={rc}, results \"\n", " f\"{'missing' if not outp.exists() else 'present'}\")\n", " return json.loads(outp.read_text())\n", "\n", "def _pt(saves, tag):\n", " return str(Path(saves[tag]) / \"aleph_relay.pt\")\n", "\n", "r4 = _bed({\"arms\": {\n", " \"relay_s0\": {\"kind\": \"relay\", \"ckpt\": _pt(relay_saves, \"e004b_relay_s0\")},\n", " \"relay_s1\": {\"kind\": \"relay\", \"ckpt\": _pt(relay_saves, \"e004b_relay_s1\")},\n", "}}, \"results_exp004b\")\n", "# (the LoRA arm is PEFT-format; its numbers come from the fork eval curve)\n", "\n", "r17 = _bed({\"arms\": {\n", " \"mb3_s0\": {\"kind\": \"multiband3\", \"ckpt\": _pt(mb3_saves, \"e017_mb3_s0\")},\n", " \"mb3_s0_lesion_b0\": {\"kind\": \"multiband3\", \"lesion\": 0,\n", " \"ckpt\": _pt(mb3_saves, \"e017_mb3_s0\")},\n", " \"mb3_s0_lesion_b1\": {\"kind\": \"multiband3\", \"lesion\": 1,\n", " \"ckpt\": _pt(mb3_saves, \"e017_mb3_s0\")},\n", " \"mb3_s0_lesion_b2\": {\"kind\": \"multiband3\", \"lesion\": 2,\n", " \"ckpt\": _pt(mb3_saves, \"e017_mb3_s0\")},\n", " \"mb3_s1\": {\"kind\": \"multiband3\", \"ckpt\": _pt(mb3_saves, \"e017_mb3_s1\")},\n", "}}, \"results_exp017\")\n", "\n", "_curves = {t: eval_curve(t, required=False) for t in\n", " (\"e004b_relay_s0\", \"e004b_relay_s1\",\n", " \"e004b_lora_s0\", \"e004b_lora_s1\")}\n", "RES4 = {\"bed\": r4,\n", " \"fork_eval\": {t: c for t, c in _curves.items() if c},\n", " \"arm_source\": {t: ARM_SOURCE.get(t) for t in ARM_SHIP},\n", " \"split_manifest\": json.loads((DATA / \"split_manifest.json\")\n", " .read_text())}\n", "_rl = min(final_loss(\"e004b_relay_s0\"), final_loss(\"e004b_relay_s1\"))\n", "_l0 = final_loss(\"e004b_lora_s0\")\n", "_curve0 = RES4[\"fork_eval\"][\"e004b_relay_s0\"][\"heldout/loss\"]\n", "_at3k = min(v for s, v in _curve0 if s <= 3000)\n", "RES4[\"verdict\"] = {\n", " # bed loss is the paired instrument; the fork curve decides P1/P4\n", " \"P1_relay_vs_lora\": (\"relay\" if _rl <= _l0 else \"lora\")\n", " if _l0 is not None else \"UNDECIDED (no LoRA control)\",\n", " \"P1_numbers\": {\"relay_best\": _rl, \"lora_s0\": _l0},\n", " \"P2_toggle_bit_exact\": True, # the bed asserted, else it raised\n", " \"P3_gates_moved\": True, # per-arm asserts in CELL 3\n", " \"P4_3k_was_a_floor\": _curve0[-1][1] < _at3k,\n", "}\n", "(WORK / \"results_exp004b_full.json\").write_text(json.dumps(RES4, indent=1))\n", "ship_arm(\"exp004b_anima_relay\", \"results\",\n", " Path(relay_saves[\"e004b_relay_s0\"]),\n", " extra={\"exp004b_anima_relay/results.json\": RES4})\n", "print(\"exp004b verdict:\", RES4[\"verdict\"])\n", "\n", "_q = [str(x) for x in r17[\"quantiles\"]]\n", "_allon = r17[\"arms\"][\"mb3_s0\"][\"loss_by_quantile\"]\n", "_off = r17[\"arms\"][\"mb3_s0_toggled_off\"][\"loss_by_quantile\"]\n", "_bandq = {0: [\"0.1\", \"0.3\"], 1: [\"0.5\", \"0.7\"], 2: [\"0.9\"]}\n", "_surgical = {}\n", "for _b, _own_q in _bandq.items():\n", " _les = r17[\"arms\"][f\"mb3_s0_lesion_b{_b}\"][\"loss_by_quantile\"]\n", " _own = sum(_les[x] - _allon[x] for x in _own_q) / len(_own_q)\n", " _cq = [x for x in _q if x not in _own_q]\n", " _cross = sum(_les[x] - _allon[x] for x in _cq) / len(_cq)\n", " _surgical[_b] = {\"own_damage\": _own, \"cross_damage\": _cross,\n", " \"ratio\": (_own / _cross) if _cross > 0 else None}\n", "# EFFECT SIZE GUARD. The bed is bit-deterministic, so any lesion delta is\n", "# \"real\" \u2014 but a stack whose gates trained shut barely moves the loss at\n", "# all, and then a beautiful ratio is measuring nothing. The shipped mb3\n", "# gates sit at sigmoid(-3.6..-7.8) = 0.027..0.0004, so this is the live\n", "# risk, not a hypothetical: report the adapter's own effect and refuse to\n", "# call a lesion surgical when the stack it lesions is inert.\n", "_effect = sum(_off[x] - _allon[x] for x in _q) / len(_q)\n", "_own_mean = sum(s[\"own_damage\"] for s in _surgical.values()) / len(_surgical)\n", "_vacuous = abs(_effect) < 1e-6 or abs(_own_mean) < 0.01 * abs(_effect)\n", "_clean = all((s[\"ratio\"] is None and s[\"own_damage\"] > 0) or\n", " (s[\"ratio\"] is not None and s[\"ratio\"] >= 10)\n", " for s in _surgical.values())\n", "RES17 = {\"bed\": r17, \"surgical\": _surgical,\n", " \"adapter_effect_mean\": _effect, # toggled_off - all_on\n", " \"own_damage_mean\": _own_mean,\n", " \"P1_surgical_10x\": bool(_clean and not _vacuous),\n", " \"P1_verdict\": (\"VACUOUS \u2014 the stack is near-inert (gates trained \"\n", " \"shut); lesion ratios measure nothing\"\n", " if _vacuous else\n", " (\"SURGICAL\" if _clean else \"NOT SURGICAL\")),\n", " \"P2_toggle_bit_exact\": True}\n", "(WORK / \"results_exp017_full.json\").write_text(json.dumps(RES17, indent=1))\n", "ship_arm(\"exp017_anima_multiband\", \"results\",\n", " Path(mb3_saves[\"e017_mb3_s0\"]),\n", " extra={\"exp017_anima_multiband/results.json\": RES17})\n", "print(\"exp017 surgical:\", json.dumps(_surgical, indent=1))\n", "print(f\"exp017 P1: {RES17['P1_verdict']} | adapter effect \"\n", " f\"{_effect:+.3e} | mean own-band damage {_own_mean:+.3e}\")\n" ] }, { "cell_type": "code", "metadata": {}, "execution_count": null, "outputs": [], "source": [ "# \u2550\u2550\u2550 CELL 7 \u2014 production naming + final ship \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n", "# v2 adapters into amoe_loras/production/anima/, usage cards written from\n", "# the MEASURED verdicts of this run. Local follow-ups: catalog regen,\n", "# ComfyUI KNOWN_ADAPTERS + live verify, write-back to the mind.\n", "import json\n", "from amoe.io.checkpoint import load_diffusion_anchor\n", "from amoe.io.safetensors_io import save_anchor_safetensors\n", "from huggingface_hub import CommitOperationAdd, HfApi\n", "\n", "def _named(save, out_name, display, usage, evidence, seed, caveats,\n", " bands=None):\n", " ck = load_diffusion_anchor(str(Path(save) / \"aleph_relay.pt\"),\n", " substrate={\n", " \"family\": \"cosmos_dit\",\n", " \"base_model_id\": \"circlestone-labs/Anima@anima-base-v1.0\"})\n", " ck.meta.update(display_name=display, usage=usage, evidence=evidence,\n", " tier=\"capability\", seed=seed, caveats=caveats, nc=True,\n", " license=\"NC (CircleStone NC + NVIDIA Open Model License)\",\n", " objective={\"kind\": \"flow\"}, recommended_strength=1.0)\n", " if bands:\n", " ck.meta[\"band_roles\"] = bands\n", " p = WORK / f\"{out_name}.safetensors\"\n", " save_anchor_safetensors(ck, str(p))\n", " return (f\"amoe_loras/production/anima/{out_name}.safetensors\", str(p))\n", "\n", "_BANDS = [\"fidelity/detail (LOW noise)\", \"continuity/semantics (MID)\",\n", " \"diversity/structure (HIGH noise)\"]\n", "# survive a kernel that ran CELL 6 in an earlier session\n", "RES4 = globals().get(\"RES4\") or json.loads(\n", " (WORK / \"results_exp004b_full.json\").read_text())\n", "RES17 = globals().get(\"RES17\") or json.loads(\n", " (WORK / \"results_exp017_full.json\").read_text())\n", "relay_saves = globals().get(\"relay_saves\") or {\n", " t: str(arm_save(t)) for t in (\"e004b_relay_s0\", \"e004b_relay_s1\")}\n", "mb3_saves = globals().get(\"mb3_saves\") or {\n", " t: str(arm_save(t)) for t in (\"e017_mb3_s0\", \"e017_mb3_s1\")}\n", "_v = RES4[\"verdict\"]\n", "_ops = []\n", "for _seed, _tag in ((0, \"e004b_relay_s0\"), (1, \"e004b_relay_s1\")):\n", " _ops.append(_named(\n", " relay_saves[_tag], f\"anima-dit-relay-v2-s{_seed}\",\n", " f\"Anima 2B DiT Relay v2 (seed {_seed}) \u2014 NON-COMMERCIAL\",\n", " \"The exp004b retrain: fp32 master weights (working gates), 10k \"\n", " \"steps, full dataset, declared seed, real held-out eval.\",\n", " f\"exp004b: P1 relay_vs_lora={_v['P1_relay_vs_lora']}, \"\n", " f\"P4 old-3k-was-a-floor={_v['P4_3k_was_a_floor']}; numbers in \"\n", " \"exp004b_anima_relay/results.json.\",\n", " _seed,\n", " [\"Non-commercial (derived from NC weights).\",\n", " \"Needs an Anima checkpoint (28 sites, width 2048).\"]))\n", "for _seed, _tag in ((0, \"e017_mb3_s0\"), (1, \"e017_mb3_s1\")):\n", " _ops.append(_named(\n", " mb3_saves[_tag], f\"anima-dit-multiband-v1-s{_seed}\",\n", " f\"Anima 2B DiT Multiband v1 (seed {_seed}) \u2014 NON-COMMERCIAL\",\n", " \"Three sigma-band experts on the Anima DiT, step-gated \u2014 the \"\n", " \"first multiband stack on a DiT-class trunk.\",\n", " f\"exp017: P1 {RES17['P1_verdict']} (surgical>=10x=\"\n", " f\"{RES17['P1_surgical_10x']}, adapter effect \"\n", " f\"{RES17['adapter_effect_mean']:+.2e}); lesion numbers in \"\n", " \"exp017_anima_multiband/results.json.\",\n", " _seed,\n", " [\"Non-commercial (derived from NC weights).\",\n", " \"Needs an Anima checkpoint; band gating needs a step-gated \"\n", " \"loader (ComfyUI nodes / StepGatedSampler).\"],\n", " bands=_BANDS))\n", "\n", "HfApi().create_commit(\n", " repo_id=HUB_REPO,\n", " operations=[CommitOperationAdd(path_in_repo=r, path_or_fileobj=s)\n", " for r, s in _ops],\n", " commit_message=\"Anima v2 production adapters (exp004b/exp017), usage \"\n", " \"cards from measured verdicts\")\n", "print(\"shipped:\", [r for r, _s in _ops])\n", "print(\"\\nCAMPAIGN COMPLETE \u2014 local follow-ups: catalog regen, ComfyUI \"\n", " \"KNOWN_ADAPTERS + live verify, write-back to the mind.\")\n" ] } ] }