anima_closeout v19: the bed's NaN was AUTOCAST_DTYPE. utils/common.py defines it as None and train.py sets it at :287 (common.AUTOCAST_DTYPE = config['model']['dtype']) BEFORE importing any model module — models/*.py do `from utils.common import AUTOCAST_DTYPE` and bake the value into their @torch .autocast decorators at class-definition time. The bed never set it, so every layer ran under autocast(dtype=None), which resolves to FP16 on cuda (verified: fast_dtype == float16), and the bf16-trained 2B DiT overflowed to NaN — the 'DETERMINISM FAILED' assert was just nan != nan. Bed now sets bf16 immediately after importing utils.common with an assert that no model module is loaded yet, and eval_arm raises a real numerics error on any non-finite loss instead of misreporting determinism.

colab/anima_closeout.ipynb

@@ -935,7 +935,7 @@
     "# 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, os, sys, time\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",
@@ -949,6 +949,17 @@
     "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",
@@ -1168,6 +1179,11 @@
     "    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",