anima_closeout v20: the bed now RUNS — frozen 0.110923, frozen_repeat identical (determinism holds), relay_s0_toggled_off bit-identical to frozen (TOGGLE LAW passes). Only the ENABLED relay crashed: attach_aleph_relays casts dtype but does not place on device (the trainer attaches on CPU then DeepSpeed moves the engine), and the bed had already moved the trunk to cuda, so fresh relay params sat on CPU -> 'mat2 is on cpu' at the first relay matmul. Fix: tf.to(DEV) after each attach (idempotent for placed params); verified a fresh attribute-assigned submodule auto-registers and .to() follows it.
Browse files
colab/anima_closeout.ipynb
CHANGED
|
@@ -1199,6 +1199,12 @@
|
|
| 1199 |
" mode = arm.get(\"kind\", \"relay\")\n",
|
| 1200 |
" attach_aleph_relays(tf, tf.model_channels, relay_path=arm[\"ckpt\"],\n",
|
| 1201 |
" dtype=torch.bfloat16, mode=mode)\n",
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1202 |
" gates = []\n",
|
| 1203 |
" for b in tf.blocks:\n",
|
| 1204 |
" g = getattr(b.aleph_relay, \"gate\", None)\n",
|
|
|
|
| 1199 |
" mode = arm.get(\"kind\", \"relay\")\n",
|
| 1200 |
" attach_aleph_relays(tf, tf.model_channels, relay_path=arm[\"ckpt\"],\n",
|
| 1201 |
" dtype=torch.bfloat16, mode=mode)\n",
|
| 1202 |
+
" # attach_aleph_relays casts dtype but does NOT place on device (the\n",
|
| 1203 |
+
" # trainer attaches on CPU, then DeepSpeed moves the whole engine). The\n",
|
| 1204 |
+
" # trunk is already on cuda, so the freshly-attached relay params are on\n",
|
| 1205 |
+
" # CPU -> 'mat2 is on cpu' at the first relay matmul. Re-place the trunk;\n",
|
| 1206 |
+
" # .to() is idempotent for params already on DEV.\n",
|
| 1207 |
+
" tf.to(DEV)\n",
|
| 1208 |
" gates = []\n",
|
| 1209 |
" for b in tf.blocks:\n",
|
| 1210 |
" g = getattr(b.aleph_relay, \"gate\", None)\n",
|