diagnostic: record mixed-traffic failure onset

3 files changed, 17690 insertions, 0 deletions

diff --git a/results/kp_innovation_diagnosis/innovation.json b/results/kp_innovation_diagnosis/innovation.json
new file mode 100644
index 0000000..a57277e
--- /dev/null
+++ b/results/kp_innovation_diagnosis/innovation.json
@@ -0,0 +1,7778 @@
+{
+  "first_any_nonfinite_groups": {
+    "bn_running_var": [
+      2,
+      3,
+      5
+    ]
+  },
+  "first_any_nonfinite_step": 37,
+  "first_nonfinite_by_group": {
+    "bn_running_var": {
+      "indices": [
+        2,
+        3,
+        5
+      ],
+      "step": 37
+    },
+    "forward_momentum": {
+      "indices": [
+        0
+      ],
+      "step": 74
+    },
+    "forward_weight": {
+      "indices": [
+        0
+      ],
+      "step": 74
+    }
+  },
+  "first_nonfinite_by_metric": {
+    "raw_apical_rms": 74,
+    "traffic_rms": 74
+  },
+  "first_nonfinite_metrics": [],
+  "first_training_failure_groups": {
+    "forward_momentum": [
+      0
+    ],
+    "forward_weight": [
+      0
+    ]
+  },
+  "first_training_failure_step": 74,
+  "initial_parameter_state": {
+    "bn_running_mean": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "bn_running_var": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 1.0,
+      "nonfinite_indices": []
+    },
+    "feedback_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "feedback_readout": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.4011031687259674,
+      "nonfinite_indices": []
+    },
+    "feedback_readout_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "feedback_weight": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.8331910371780396,
+      "nonfinite_indices": []
+    },
+    "forward_bn_bias": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_bn_bias_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_bn_scale": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 1.0,
+      "nonfinite_indices": []
+    },
+    "forward_bn_scale_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_readout": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.3956821858882904,
+      "nonfinite_indices": []
+    },
+    "forward_readout_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_weight": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 1.1162850856781006,
+      "nonfinite_indices": []
+    },
+    "predictor_bias": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.027761053293943405,
+      "nonfinite_indices": []
+    },
+    "predictor_slope": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.09883774071931839,
+      "nonfinite_indices": []
+    }
+  },
+  "max_steps": 352,
+  "predictor_warmup": {
+    "first_mse": 0.0002822713352582903,
+    "last_mse": 5.667926707839955e-06,
+    "post_warmup_traffic_residual_rms_ratio": 0.13994200341774746,
+    "steps": 20,
+    "task_loader_state_restored": true
+  },
+  "protocol": "kp_mixed_traffic_nonfinite_diagnosis_v1",
+  "provenance": {
+    "git_commit": "e9962317b3d976fe25e28a5693f2757e3758480a",
+    "git_tracked_dirty": false
+  },
+  "rule": "innovation",
+  "scope": "training_only_no_validation_or_test_evaluation",
+  "split": {
+    "dataset": "cifar10",
+    "input_layout": "NCHW",
+    "input_shape": [
+      3,
+      32,
+      32
+    ],
+    "loader_seed": 0,
+    "normalization_mean": [
+      0.49140000343322754,
+      0.4821999967098236,
+      0.4465000033378601
+    ],
+    "normalization_std": [
+      0.24699999392032623,
+      0.2434999942779541,
+      0.26159998774528503
+    ],
+    "split_from_training_only": true,
+    "split_seed": 2027,
+    "test_examples": 10000,
+    "train_examples": 45000,
+    "training_augmentation": "random_crop_32_padding_4_zero_then_horizontal_flip_p0.5",
+    "validation_class_counts": {
+      "0": 500,
+      "1": 500,
+      "2": 500,
+      "3": 500,
+      "4": 500,
+      "5": 500,
+      "6": 500,
+      "7": 500,
+      "8": 500,
+      "9": 500
+    },
+    "validation_examples": 5000,
+    "validation_index_sha256": "8328b206a97c420e49e54e3eca4abe3274c4756b084355784ea3fb8059e4515b"
+  },
+  "test_evaluations": 0,
+  "traffic_calibration": {
+    "instruction_rms": [
+      0.008471464738249779,
+      0.006617916747927666,
+      0.003988116513937712,
+      0.0043865954503417015,
+      0.002713605063036084,
+      0.003158062230795622,
+      0.002004207344725728,
+      0.003757704282179475,
+      0.0027085731271654367,
+      0.0030200635083019733,
+      0.002056022873148322,
+      0.0023899395018815994,
+      0.0016842411132529378,
+      0.0034010226372629404,
+      0.002549725119024515,
+      0.0030542907770723104,
+      0.0020918569061905146,
+      0.002605808200314641,
+      0.0018916188273578882
+    ],
+    "realized_traffic_instruction_rms_ratio": [
+      4.0,
+      4.0,
+      4.0,
+      4.0,
+      4.0,
+      4.0,
+      3.999999523162842,
+      4.0,
+      4.000000476837158,
+      3.999999761581421,
+      4.0,
+      4.0,
+      3.999999761581421,
+      4.0,
+      4.000000476837158,
+      4.000000476837158,
+      4.0,
+      4.0,
+      4.0
+    ],
+    "target_ratio": 4.0,
+    "traffic_gain": [
+      0.0443243607878685,
+      0.035817068070173264,
+      0.013778294436633587,
+      0.022762222215533257,
+      0.007214079145342112,
+      0.017218582332134247,
+      0.004616298712790012,
+      0.02042512409389019,
+      0.007327491417527199,
+      0.01611751690506935,
+      0.004703814629465342,
+      0.012562948279082775,
+      0.003324519144371152,
+      0.01805141195654869,
+      0.006063441745936871,
+      0.016296451911330223,
+      0.0044285994954407215,
+      0.013874167576432228,
+      0.0036007820162922144
+    ],
+    "unscaled_traffic_rms": [
+      0.7644973993301392,
+      0.7390796542167664,
+      1.157796859741211,
+      0.7708553671836853,
+      1.504616141319275,
+      0.7336404919624329,
+      1.736635684967041,
+      0.7358984351158142,
+      1.478581428527832,
+      0.7495108246803284,
+      1.7483876943588257,
+      0.7609485983848572,
+      2.0264477729797363,
+      0.7536302804946899,
+      1.6820316314697266,
+      0.7496824264526367,
+      1.8894071578979492,
+      0.7512690424919128,
+      2.10134220123291
+    ]
+  },
+  "trajectory": [
+    {
+      "batch_loss": 2.9134020805358887,
+      "innovation_rms": 0.0047711752299396705,
+      "instruction_rms": 0.004215397916032357,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.40285658836364746,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9875131845474243,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6892958879470825,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4058851897716522,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6322960257530212,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8331910371780396,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.07263921201229095,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7263920903205872,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1397736072540283,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3977354764938354,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0743976831436157,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3794170618057251,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.632293164730072,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1159934997558594,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": 4.976208295513219e-06,
+      "predictor_updated": true,
+      "raw_apical_rms": 0.017318262531229,
+      "step": 1,
+      "teaching_rms": 0.0047711752299396705,
+      "traffic_rms": 0.01680093257108933
+    },
+    {
+      "batch_loss": 2.4969701766967773,
+      "innovation_rms": 0.002759960898169119,
+      "instruction_rms": 0.0013977544950979999,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7672886252403259,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.755084991455078,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.103839635848999,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.42035770416259766,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7175377607345581,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8331910371780396,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.24586249887943268,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7979724407196045,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4633216857910156,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.2354812622070312,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.103909969329834,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3822507858276367,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7175324559211731,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.114304542541504,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.018933221801795217,
+      "step": 2,
+      "teaching_rms": 0.002759960898169119,
+      "traffic_rms": 0.01888791884750171
+    },
+    {
+      "batch_loss": 2.808241128921509,
+      "innovation_rms": 0.002827276396737058,
+      "instruction_rms": 0.0007422902489447251,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1650590896606445,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18.907119750976562,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.135357141494751,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.45309779047966003,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1354742050170898,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8331910371780396,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4564003646373749,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.46323299407959,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.072580337524414,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.092585563659668,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1353305578231812,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.40342289209365845,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1354973316192627,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1105809211730957,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.022930354797487982,
+      "step": 3,
+      "teaching_rms": 0.002827276396737058,
+      "traffic_rms": 0.02292513507898896
+    },
+    {
+      "batch_loss": 5.85762882232666,
+      "innovation_rms": 0.0032404673695190997,
+      "instruction_rms": 0.0006564077781129235,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6060140132904053,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 55.902095794677734,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.771088719367981,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4172212481498718,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7612570524215698,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8331910371780396,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6743493676185608,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.648998737335205,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1927340030670166,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.201536178588867,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7710490226745605,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4166879653930664,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7613394260406494,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1045191287994385,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.02851504304869197,
+      "step": 4,
+      "teaching_rms": 0.0032404673695190997,
+      "traffic_rms": 0.028530091316332357
+    },
+    {
+      "batch_loss": 6.421552658081055,
+      "innovation_rms": 0.004012600107790864,
+      "instruction_rms": 0.0006237183841124908,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.8697242736816406,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 226.96554565429688,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5972275733947754,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4789431691169739,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6568351984024048,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8331910371780396,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3503779172897339,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.845800399780273,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.1754255294799805,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.826915740966797,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5971802473068237,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4699835777282715,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6569137573242188,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0935394763946533,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.03815050381612367,
+      "step": 5,
+      "teaching_rms": 0.004012600107790864,
+      "traffic_rms": 0.03816605228034207
+    },
+    {
+      "batch_loss": 11.498717308044434,
+      "innovation_rms": 0.0056700982373924765,
+      "instruction_rms": 0.0007255506849147926,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.791121006011963,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 789.2258911132812,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4394961595535278,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5409393310546875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.117813587188721,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8622811436653137,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.933739185333252,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15.833613395690918,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.693208694458008,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.17782974243164,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.038958787918091,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5269163250923157,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.117769241333008,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0828872919082642,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.05595474338981351,
+      "step": 6,
+      "teaching_rms": 0.0056700982373924765,
+      "traffic_rms": 0.05597822669223256
+    },
+    {
+      "batch_loss": 38.18779373168945,
+      "innovation_rms": 0.008904289058286411,
+      "instruction_rms": 0.001225207597366297,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.43674373626709,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3821.960693359375,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.270296096801758,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.643233597278595,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.465946674346924,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2184159755706787,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.028405666351318,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30.946664810180664,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14.941880226135254,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 62.48671340942383,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.27035903930664,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6373922824859619,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.4659318923950195,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3157312870025635,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.08771435650057213,
+      "step": 7,
+      "teaching_rms": 0.008904289058286411,
+      "traffic_rms": 0.0878032313284443
+    },
+    {
+      "batch_loss": 56.086124420166016,
+      "innovation_rms": 0.014898834527524718,
+      "instruction_rms": 0.0013541738213322413,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 141.75970458984375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 567439.5,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.142507553100586,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7134369015693665,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.214983940124512,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.13008975982666,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.1319580078125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 61.035518646240234,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25.95669174194336,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 110.14811706542969,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.142576217651367,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8331893682479858,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.214798927307129,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.225393772125244,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.14463046595067733,
+      "step": 8,
+      "teaching_rms": 0.014898834527524718,
+      "traffic_rms": 0.1447301026798327
+    },
+    {
+      "batch_loss": 88.54937744140625,
+      "innovation_rms": 0.02542648321057176,
+      "instruction_rms": 0.0015879852101984948,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 558.5450439453125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5227687.5,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28.86151695251465,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2487165927886963,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.786975860595703,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.850591659545898,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25.10658073425293,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 129.74623107910156,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 42.01982879638672,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 167.85585021972656,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28.861492156982422,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.559989333152771,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.787158966064453,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.944836139678955,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.24218480431164466,
+      "step": 9,
+      "teaching_rms": 0.02542648321057176,
+      "traffic_rms": 0.24227499140944508
+    },
+    {
+      "batch_loss": 250.72975158691406,
+      "innovation_rms": 0.045763076046941735,
+      "instruction_rms": 0.0032991611071113446,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1637.776611328125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28017968.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22.971668243408203,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8755358457565308,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16.366619110107422,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.398952960968018,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 48.57054901123047,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 234.63967895507812,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 72.15222930908203,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 301.323974609375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41.22135925292969,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9793704748153687,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16.366594314575195,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.492237091064453,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.42875119982644977,
+      "step": 10,
+      "teaching_rms": 0.045763076046941735,
+      "traffic_rms": 0.4289306146135509
+    },
+    {
+      "batch_loss": 847.3389282226562,
+      "innovation_rms": 0.0826060383757554,
+      "instruction_rms": 0.005403889961997733,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4115.26904296875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 117325016.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29.202810287475586,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.9798243045806885,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 52.07854080200195,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.792180061340332,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 91.6338882446289,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 430.63336181640625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 127.03309631347656,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 548.8086547851562,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 145.7760467529297,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.8329885005950928,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 52.078643798828125,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.844808578491211,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.7669757197498637,
+      "step": 11,
+      "teaching_rms": 0.0826060383757554,
+      "traffic_rms": 0.7673285313007404
+    },
+    {
+      "batch_loss": 1456.1527099609375,
+      "innovation_rms": 0.1488487851756446,
+      "instruction_rms": 0.007916429550500146,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9591.8896484375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 388613824.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 132.1781005859375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.874972343444824,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 104.57723236083984,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.992986679077148,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 171.8209686279297,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 801.870849609375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 224.8765411376953,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 978.4344482421875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 132.1780242919922,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.892233848571777,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 104.57722473144531,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26.949443817138672,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.3731121085584415,
+      "step": 12,
+      "teaching_rms": 0.1488487851756446,
+      "traffic_rms": 1.373488620588331
+    },
+    {
+      "batch_loss": 9487.517578125,
+      "innovation_rms": 0.2717326766628019,
+      "instruction_rms": 0.024449094987210283,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21653.3515625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 965573760.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 542.1923828125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15.26036548614502,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 188.550537109375,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 54.816532135009766,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 319.693603515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1478.7261962890625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 400.7353210449219,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1758.587890625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 542.1923828125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15.277626991271973,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 188.55055236816406,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 54.88191223144531,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2.4970706818965436,
+      "step": 13,
+      "teaching_rms": 0.2717326766628019,
+      "traffic_rms": 2.4983659565391805
+    },
+    {
+      "batch_loss": 23258.83203125,
+      "innovation_rms": 0.5039805954439139,
+      "instruction_rms": 0.049479498192594465,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 89271.4140625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 76001730560.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 559.6827392578125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30.550294876098633,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 386.0753479003906,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 102.96351623535156,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 595.093505859375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2753.998779296875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 714.5072631835938,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3137.71923828125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 559.682861328125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30.244033813476562,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 386.0755920410156,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 102.90254974365234,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4.612100141342417,
+      "step": 14,
+      "teaching_rms": 0.5039805954439139,
+      "traffic_rms": 4.613756785731884
+    },
+    {
+      "batch_loss": 74670.375,
+      "innovation_rms": 0.9358941753042994,
+      "instruction_rms": 0.09690040056870634,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 395987.1875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 761320636416.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 526.406005859375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 60.312957763671875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 639.5701904296875,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 146.92283630371094,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1105.8680419921875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5107.7451171875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1274.7078857421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5602.005859375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 526.406005859375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 60.00672149658203,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 639.5703125,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 146.8618621826172,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 8.547812908154881,
+      "step": 15,
+      "teaching_rms": 0.9358941753042994,
+      "traffic_rms": 8.550583656970959
+    },
+    {
+      "batch_loss": 414248.90625,
+      "innovation_rms": 1.7568866063152748,
+      "instruction_rms": 0.2862059157354659,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1185289.625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4690009391104.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2356.9755859375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 96.13926696777344,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1291.882080078125,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 244.150146484375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2066.31103515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9604.4296875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2255.385986328125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9806.7802734375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2356.9755859375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 95.81422424316406,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1291.88232421875,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 244.0204315185547,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 15.939239959056371,
+      "step": 16,
+      "teaching_rms": 1.7568866063152748,
+      "traffic_rms": 15.946682011199735
+    },
+    {
+      "batch_loss": 1068312.25,
+      "innovation_rms": 3.291665090172094,
+      "instruction_rms": 0.5043453838312517,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3012151.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22521078349824.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2665.037841796875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 193.00978088378906,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2394.646240234375,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 457.6437072753906,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3829.5029296875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18279.6796875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4013.245849609375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17578.59765625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2665.037841796875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 192.6847686767578,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2394.64599609375,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 457.51397705078125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": 10.670841003792198,
+      "predictor_updated": true,
+      "raw_apical_rms": 29.85815266250202,
+      "step": 17,
+      "teaching_rms": 3.291665090172094,
+      "traffic_rms": 29.86770369035303
+    },
+    {
+      "batch_loss": 3723487.0,
+      "innovation_rms": 5.560727640277686,
+      "instruction_rms": 1.684154987002244,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6929510.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 144708722491392.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8246.1474609375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 560.7797241210938,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6526.595703125,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 799.1076049804688,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7010.78515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33241.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6917.349609375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29041.03515625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8246.1474609375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 560.6668090820312,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6526.59619140625,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 799.3399047851562,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 56.14805589669773,
+      "step": 18,
+      "teaching_rms": 5.560727640277686,
+      "traffic_rms": 56.14852989731709
+    },
+    {
+      "batch_loss": 18693106.0,
+      "innovation_rms": 11.264667837923133,
+      "instruction_rms": 5.6162971728306115,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16728729.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 824487020855296.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8735.8271484375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1045.246337890625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10032.599609375,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1536.1827392578125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13247.720703125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 62369.3515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11856.8486328125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 49394.9921875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8735.8271484375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1045.1334228515625,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10032.59765625,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1536.2130126953125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 102.05386398229633,
+      "step": 19,
+      "teaching_rms": 11.264667837923133,
+      "traffic_rms": 101.9704956654027
+    },
+    {
+      "batch_loss": 50152412.0,
+      "innovation_rms": 19.86923290857296,
+      "instruction_rms": 8.558431659383785,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41275708.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2451744499957760.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11818.958984375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1874.683837890625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22750.109375,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2191.80078125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25008.5625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 117608.421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20198.33203125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 92129.6953125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11818.958984375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1874.5372314453125,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22750.111328125,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2191.8310546875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 185.32871548353896,
+      "step": 20,
+      "teaching_rms": 19.86923290857296,
+      "traffic_rms": 185.23590846587
+    },
+    {
+      "batch_loss": 255640048.0,
+      "innovation_rms": 48.05750768773175,
+      "instruction_rms": 35.030944314015834,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 91910384.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.993595453793894e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11085.376953125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3506.18017578125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 44249.8125,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2680.5458984375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 47360.125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 223515.609375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36887.7734375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 171533.40625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13111.970703125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3506.069580078125,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 44249.8125,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2680.576171875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 340.6450894148393,
+      "step": 21,
+      "teaching_rms": 48.05750768773175,
+      "traffic_rms": 339.13119495939634
+    },
+    {
+      "batch_loss": 1153638400.0,
+      "innovation_rms": 62.79746044008235,
+      "instruction_rms": 17.288717876638955,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 193690880.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.900145695779717e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11253.15625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6263.9404296875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 65730.3125,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3697.6640625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 89625.859375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 422657.3125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 69086.9375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 321991.625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21184.943359375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6263.82958984375,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 65730.3125,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3697.549072265625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 621.9789944453244,
+      "step": 22,
+      "teaching_rms": 62.79746044008235,
+      "traffic_rms": 622.4596035232903
+    },
+    {
+      "batch_loss": 2905861632.0,
+      "innovation_rms": 114.85880979139532,
+      "instruction_rms": 27.453292305649818,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 431770528.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.42159282749073e+18,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 65725.0078125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10174.228515625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 140268.484375,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7080.2744140625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 169811.421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 801855.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 129315.7421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 602288.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 65725.0078125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10174.2158203125,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 140268.484375,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7080.1181640625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1144.0240451263883,
+      "step": 23,
+      "teaching_rms": 114.85880979139532,
+      "traffic_rms": 1144.6853322528616
+    },
+    {
+      "batch_loss": 9168957440.0,
+      "innovation_rms": 216.55432575292375,
+      "instruction_rms": 65.89381456647692,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1253806976.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3261129458763956e+19,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 336917.53125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18466.134765625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 125239.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34133.0546875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 322526.96875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1527155.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 241303.5625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1119878.25,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 336917.53125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18466.123046875,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 125239.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34132.91015625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2113.345970698597,
+      "step": 24,
+      "teaching_rms": 216.55432575292375,
+      "traffic_rms": 2114.01383406673
+    },
+    {
+      "batch_loss": 39081639936.0,
+      "innovation_rms": 484.7700874188069,
+      "instruction_rms": 300.3864148449255,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26668593152.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.638770194046062e+21,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 303250.71875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39907.36328125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 515960.65625,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 64458.12890625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 612450.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2899236.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 450833.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2095294.5,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 303250.71875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39907.37890625,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 515960.65625,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 64457.984375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3912.9865651501977,
+      "step": 25,
+      "teaching_rms": 484.7700874188069,
+      "traffic_rms": 3905.333899457679
+    },
+    {
+      "batch_loss": 157594779648.0,
+      "innovation_rms": 747.1018399268916,
+      "instruction_rms": 250.17171880558766,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 114262212608.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.2705144332153094e+22,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 272940.03125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86153.640625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1097004.625,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 91752.1328125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1163216.875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5507663.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 840726.5625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3898935.75,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 272940.03125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86153.953125,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1097004.625,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 91751.984375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 7232.222916934155,
+      "step": 26,
+      "teaching_rms": 747.1018399268916,
+      "traffic_rms": 7236.345563390998
+    },
+    {
+      "batch_loss": 560067837952.0,
+      "innovation_rms": 1631.2172332936816,
+      "instruction_rms": 978.367502620053,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 330656382976.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.9535177412728463e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1034668.1875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 163030.421875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1864474.75,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 116283.0078125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2207668.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10444512.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1570884.125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7301576.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1034668.1875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 163030.734375,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1864474.75,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 116282.859375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 13440.838373279787,
+      "step": 27,
+      "teaching_rms": 1631.2172332936816,
+      "traffic_rms": 13421.453980665738
+    },
+    {
+      "batch_loss": 1976891473920.0,
+      "innovation_rms": 3087.4331727937674,
+      "instruction_rms": 1918.8919395712537,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 812596658176.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.895227518213575e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2869032.25,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 294881.75,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3855363.75,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 281528.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4185630.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19779628.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2943418.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13725342.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4698523.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 294881.625,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3855363.75,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 467714.75,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 24919.93985402302,
+      "step": 28,
+      "teaching_rms": 3087.4331727937674,
+      "traffic_rms": 24876.86473811417
+    },
+    {
+      "batch_loss": 5803712446464.0,
+      "innovation_rms": 6279.124074151342,
+      "instruction_rms": 4383.032243602451,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1809261920256.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1143983181150187e+25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6324684.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 543012.25,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6992883.5,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 641724.125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7992836.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38072052.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5352012.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24085944.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16543108.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 543012.125,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6992883.5,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1664542.375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 46314.75344139226,
+      "step": 29,
+      "teaching_rms": 6279.124074151342,
+      "traffic_rms": 46155.83405417223
+    },
+    {
+      "batch_loss": 22621580165120.0,
+      "innovation_rms": 11853.588491694605,
+      "instruction_rms": 8396.55048282417,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5755571798016.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.599761305188379e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10528076.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 952526.4375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10981554.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1206007.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15181742.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 71889056.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9945211.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 45931980.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14883098.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 952526.3125,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10981554.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3147332.25,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 86306.18677405613,
+      "step": 30,
+      "teaching_rms": 11853.588491694605,
+      "traffic_rms": 85998.74901856213
+    },
+    {
+      "batch_loss": 77531323564032.0,
+      "innovation_rms": 17577.036747588798,
+      "instruction_rms": 8040.259958044565,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24388994859008.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.9244792788467573e+27,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15591619.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1603147.75,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21955204.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2077728.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28771596.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 135898528.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18607798.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86625864.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15591619.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1603147.875,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21955204.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4482977.5,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 160617.03800716676,
+      "step": 31,
+      "teaching_rms": 17577.036747588798,
+      "traffic_rms": 160603.02039261052
+    },
+    {
+      "batch_loss": 299608563515392.0,
+      "innovation_rms": 37777.05847468805,
+      "instruction_rms": 24043.1263698353,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 70409546366976.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.555817464972888e+28,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15951830.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3980421.5,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 51300816.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2938746.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 54460780.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 256891856.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34959764.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 163519648.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15951830.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3980421.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 51300816.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5688266.5,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.396739959716797,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 300402.68537443853,
+      "step": 32,
+      "teaching_rms": 37777.05847468805,
+      "traffic_rms": 299825.2448033978
+    },
+    {
+      "batch_loss": 1217121925201920.0,
+      "innovation_rms": 118377.73797099873,
+      "instruction_rms": 105246.46457286323,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 171619192406016.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1472941201993678e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63154732.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7154641.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 92835096.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6475553.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 103047424.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 485866400.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 65701372.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 307416064.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63154732.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7154640.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 92835096.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6782622.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": 3028291639.022103,
+      "predictor_updated": true,
+      "raw_apical_rms": 567804.2183367815,
+      "step": 33,
+      "teaching_rms": 118377.73797099873,
+      "traffic_rms": 558799.8846624728
+    },
+    {
+      "batch_loss": 4140923998961664.0,
+      "innovation_rms": 123726.98881566343,
+      "instruction_rms": 88291.16793684504,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 757464324636672.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.07190001661708e+30,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 138724896.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12627441.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 156503408.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13896475.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 187516832.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 844694080.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 120384992.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 546836224.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 138724896.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12627441.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 156503408.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13896475.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1044458.1622444034,
+      "step": 34,
+      "teaching_rms": 123726.98881566343,
+      "traffic_rms": 1042257.8670527334
+    },
+    {
+      "batch_loss": 1.6409654846291968e+16,
+      "innovation_rms": 173932.8449718454,
+      "instruction_rms": 69630.36318708396,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3115893546549248.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.874320736445638e+31,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 124868800.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18466778.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 273738112.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26383356.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 339647616.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1521307776.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 219000848.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 986158592.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 129638152.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18466778.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 273738112.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26383356.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1881599.1741790557,
+      "step": 35,
+      "teaching_rms": 173932.8449718454,
+      "traffic_rms": 1883693.5568009613
+    },
+    {
+      "batch_loss": 4.029507393342669e+16,
+      "innovation_rms": 440374.734249781,
+      "instruction_rms": 331098.596306034,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8692858651410432.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5595130500200192e+32,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 398267488.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41641024.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 557904256.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39568628.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 616880768.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2772331520.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 398198240.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1791974016.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 398267488.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41641024.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 557904256.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39568628.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3387991.483746393,
+      "step": 36,
+      "teaching_rms": 440374.734249781,
+      "traffic_rms": 3376506.5327720894
+    },
+    {
+      "batch_loss": 1.515085275993211e+17,
+      "innovation_rms": 566168.9215680822,
+      "instruction_rms": 214461.72417149102,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.132191014702285e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.9662555617449678e+32,
+          "nonfinite_indices": [
+            2,
+            3,
+            5
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 673042048.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 80296432.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1064955328.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 75412304.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1123432192.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5065514496.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 722701568.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3245032960.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 673042048.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 80296432.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1064955328.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 75412304.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 6064553.2668123925,
+      "step": 37,
+      "teaching_rms": 566168.9215680822,
+      "traffic_rms": 6070580.35362293
+    },
+    {
+      "batch_loss": 5.0994679281352704e+17,
+      "innovation_rms": 1298205.7873512765,
+      "instruction_rms": 898095.8094674291,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3507144723936051e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 7.592467116239852e+31,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            5
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 605731520.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 143622624.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1888025216.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 124093664.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2048379264.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9249470464.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1314317568.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5916160000.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1626336128.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 143622624.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1888025216.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 161270864.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 10852977.537589325,
+      "step": 38,
+      "teaching_rms": 1298205.7873512765,
+      "traffic_rms": 10834198.090458583
+    },
+    {
+      "batch_loss": 1.6903871149585203e+18,
+      "innovation_rms": 2580481.244099079,
+      "instruction_rms": 1946104.4878205103,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.9810089662336205e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 3.453718083626939e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            5
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1922476416.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 257840896.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3445433600.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 193768576.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3740413952.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16920347648.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2388641792.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10743242752.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1922476416.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 257840896.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3445433600.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 307404896.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 19616150.926117733,
+      "step": 39,
+      "teaching_rms": 2580481.244099079,
+      "traffic_rms": 19553308.792439166
+    },
+    {
+      "batch_loss": 5.021784064378536e+18,
+      "innovation_rms": 4016555.1786649,
+      "instruction_rms": 2597688.4638285968,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.935086050953134e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.1064129334164363e+33,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            5,
+            7,
+            11
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1730209408.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 478963552.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6206426624.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 366789504.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6831440896.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30910269440.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4346770944.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19581292544.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1730209408.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 478963552.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6206426624.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 438424128.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 35285574.49222561,
+      "step": 40,
+      "teaching_rms": 4016555.1786649,
+      "traffic_rms": 35245522.493821
+    },
+    {
+      "batch_loss": 1.6099650686756585e+19,
+      "innovation_rms": 6729123.890427496,
+      "instruction_rms": 3791893.1494202893,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2672601841223598e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 4.780816111966701e+33,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            5,
+            7,
+            9,
+            11
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1557151744.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 868395456.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9106350080.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 522504672.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12477887488.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 56464465920.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7915631104.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35688599552.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9035945984.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 868395456.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9106350080.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 899460224.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 63900374.06818871,
+      "step": 41,
+      "teaching_rms": 6729123.890427496,
+      "traffic_rms": 63885990.45808338
+    },
+    {
+      "batch_loss": 5.029746287782239e+19,
+      "innovation_rms": 32291971.18462265,
+      "instruction_rms": 30675656.788202655,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.892950988417139e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.451466610531847e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            5,
+            7,
+            9,
+            11,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1401384320.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1336779776.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19118104576.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 662643136.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22777966592.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 103000784896.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14444674048.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 65290428416.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15880585216.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1336779776.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19118104576.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1712079744.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 119633767.60450122,
+      "step": 42,
+      "teaching_rms": 32291971.18462265,
+      "traffic_rms": 115803940.76496395
+    },
+    {
+      "batch_loss": 2.0154319056799164e+20,
+      "innovation_rms": 83356331.26823132,
+      "instruction_rms": 81336445.4376472,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0174474883648455e+19,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.5785291630203606e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            7,
+            9,
+            11,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2369611264.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3001660160.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 40000770048.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 788761088.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41588719616.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 188107522048.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26302062592.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 118573891584.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14320762880.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3001660160.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 40000770048.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3028047360.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 224991550.1977497,
+      "step": 43,
+      "teaching_rms": 83356331.26823132,
+      "traffic_rms": 210154323.84208333
+    },
+    {
+      "batch_loss": 5.026570634318431e+20,
+      "innovation_rms": 67314006.34093411,
+      "instruction_rms": 58510094.53730393,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.059688744033714e+19,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.3167100037383314e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            9,
+            11,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3053626368.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5994760192.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 68105404416.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 902259328.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 75986321408.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 343976050688.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 47919730688.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 216176689152.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12935328768.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5994760192.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 68105404416.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4321580032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 385410665.9039606,
+      "step": 44,
+      "teaching_rms": 67314006.34093411,
+      "traffic_rms": 381461843.6252373
+    },
+    {
+      "batch_loss": 2.2993710550741858e+21,
+      "innovation_rms": 109111862.15590881,
+      "instruction_rms": 91099506.87027645,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.7765393152170525e+19,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 6.306078165807595e+31,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            11,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2748229888.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9398620160.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 124369641472.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1004398720.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 138766991360.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 627806699520.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 87256383488.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 393366536192.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 64157241344.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9398620160.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 124369641472.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6356432384.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 697123242.9561172,
+      "step": 45,
+      "teaching_rms": 109111862.15590881,
+      "traffic_rms": 692447495.4749589
+    },
+    {
+      "batch_loss": 7.043883707636449e+21,
+      "innovation_rms": 189349288.78062975,
+      "instruction_rms": 154727389.0745287,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4135076075692936e+20,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 3.108968224077107e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            11,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5462878208.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18127556608.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 241536008192.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1096314112.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 253559308288.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1147923202048.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 158703747072.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 714473668608.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 57786699776.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18127556608.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 241536008192.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12135102464.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1263915295.3259373,
+      "step": 46,
+      "teaching_rms": 189349288.78062975,
+      "traffic_rms": 1256612985.627656
+    },
+    {
+      "batch_loss": 2.4492979695841773e+22,
+      "innovation_rms": 360037179.6555091,
+      "instruction_rms": 300717090.6097152,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.538408562168263e+20,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 3.379661657411953e+31,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4916535296.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34156601344.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 436415791104.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1179027072.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 462871199744.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2093118914560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 289255718912.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1305519849472.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 273681137664.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34156601344.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 436415791104.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27277160448.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2296878938.6894627,
+      "step": 47,
+      "teaching_rms": 360037179.6555091,
+      "traffic_rms": 2281273273.2514668
+    },
+    {
+      "batch_loss": 7.71246300074795e+22,
+      "innovation_rms": 677570497.2357241,
+      "instruction_rms": 574141846.2199012,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.074879541006471e+20,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 8.919857027629502e+31,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9317008384.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 61940568064.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 815574024192.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1481164672.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 844421660672.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3815504347136.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 528922673152.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2396669345792.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 246256697344.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 61940568064.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 815574024192.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 51902828544.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58995.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4174726892.1852612,
+      "step": 48,
+      "teaching_rms": 677570497.2357241,
+      "traffic_rms": 4142309899.502693
+    },
+    {
+      "batch_loss": 2.4919481389192865e+23,
+      "innovation_rms": 1488352898.067969,
+      "instruction_rms": 1337708725.6245377,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7461355067866004e+21,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 4.250693048943523e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9155631104.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 120110514176.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1349066948608.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1998735232.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1549704626176.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7052829130752.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 940796936192.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4118742564864.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2403702931456.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 120110514176.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1349066948608.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 239666741248.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": 4.429140637708763e+17,
+      "predictor_updated": true,
+      "raw_apical_rms": 7614518262.694409,
+      "step": 49,
+      "teaching_rms": 1488352898.067969,
+      "traffic_rms": 7509015563.157561
+    },
+    {
+      "batch_loss": 7.643299599974555e+23,
+      "innovation_rms": 2985905862.5368094,
+      "instruction_rms": 2808099510.4958196,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.6298841296870404e+21,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.6990985459175045e+33,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8892988416.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 171278630912.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2414040055808.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2753382400.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2732871122944.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11831666016256.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1661216358400.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7204194746368.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2165615362048.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 171278630912.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2414040055808.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 456228274176.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 13912553855.773369,
+      "step": 50,
+      "teaching_rms": 2985905862.5368094,
+      "traffic_rms": 13648139024.41573
+    },
+    {
+      "batch_loss": 2.6085843436527636e+24,
+      "innovation_rms": 3725757571.8977556,
+      "instruction_rms": 3251246082.9269857,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.203599364821516e+21,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.4052445407447012e+33,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8003591680.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 340180795392.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4618071310336.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3432537344.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4790135291904.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20572641689600.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2931574964224.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12703585271808.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1951030181888.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 340180795392.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4618071310336.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 651331305472.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 24164564927.36383,
+      "step": 51,
+      "teaching_rms": 3725757571.8977556,
+      "traffic_rms": 23987589957.08591
+    },
+    {
+      "batch_loss": 7.862787463759648e+24,
+      "innovation_rms": 7591530942.732826,
+      "instruction_rms": 6877263818.77342,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.390639957697934e+22,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 8.014242726594584e+24,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9679612928.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 680938242048.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7938722758656.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4043742464.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8408392007680.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36182566109184.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5171825868800.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22402507472896.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11353782747136.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 680938242048.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7938722758656.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1125537677312.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 42335454426.1988,
+      "step": 52,
+      "teaching_rms": 7591530942.732826,
+      "traffic_rms": 41846099942.18983
+    },
+    {
+      "batch_loss": 2.754271828430527e+25,
+      "innovation_rms": 9113753170.222456,
+      "instruction_rms": 7149596953.603826,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.634867891510053e+22,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.3771567382968496e+25,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30222102528.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1101880754176.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13898385195008.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4593786368.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14789197692928.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63808055803904.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9125033410560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39532078563328.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10219173183488.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1101880754176.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13898385195008.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2147455074304.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 73509755944.26274,
+      "step": 53,
+      "teaching_rms": 9113753170.222456,
+      "traffic_rms": 73308152402.15108
+    },
+    {
+      "batch_loss": 1.0501738779860477e+26,
+      "innovation_rms": 11487413839.061468,
+      "instruction_rms": 5873514055.485483,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.912184199964057e+22,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.1238155861713055e+26,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30993381376.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1841041506304.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25140326825984.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5974366720.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26033426268160.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 112442277363712.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16096482557952.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 69714491473920.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9192749400064.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1841041506304.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25140326825984.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3066730119168.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 128205930966.32605,
+      "step": 54,
+      "teaching_rms": 11487413839.061468,
+      "traffic_rms": 128392471176.95818
+    },
+    {
+      "batch_loss": 2.3667097014918653e+26,
+      "innovation_rms": 17859199563.534084,
+      "instruction_rms": 3579859856.9055905,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.076297083103812e+22,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.9085848377989314e+26,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27893745664.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3317374124032.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 45683138625536.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8263052288.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 45903171813376.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 198697434480640.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28426910040064.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 123304266432512.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8263729938432.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3317374124032.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 45683138625536.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3893103165440.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 224678215846.14557,
+      "step": 55,
+      "teaching_rms": 17859199563.534084,
+      "traffic_rms": 225063923767.6041
+    },
+    {
+      "batch_loss": 6.363256019109516e+26,
+      "innovation_rms": 31266496734.427677,
+      "instruction_rms": 4996191095.730709,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6656080435949725e+23,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 5.285649250141692e+26,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25103794176.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6572290342912.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 71161366970368.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10438956032.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 81062080282624.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 351589109858304.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 50168430854144.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 217415220723712.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13385693396992.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6572290342912.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 71161366970368.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4635771273216.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 394710772551.6086,
+      "step": 56,
+      "teaching_rms": 31266496734.427677,
+      "traffic_rms": 395314214498.95026
+    },
+    {
+      "batch_loss": 2.463247453081179e+27,
+      "innovation_rms": 54940308917.543015,
+      "instruction_rms": 10017542861.59355,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.044087471651073e+23,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 8.786629137774917e+26,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22592587776.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9497302007808.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 129833036152832.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12421579776.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 143126027042816.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 620639417270272.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 88625299587072.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 384568720883712.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12094960304128.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9497302007808.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 129833036152832.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5295141027840.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 693654776916.7847,
+      "step": 57,
+      "teaching_rms": 54940308917.543015,
+      "traffic_rms": 694972586475.3656
+    },
+    {
+      "batch_loss": 7.50773098096594e+27,
+      "innovation_rms": 96472561411.09248,
+      "instruction_rms": 16831591468.24844,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.5356340696275834e+23,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.318062512938846e+27,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20332275712.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17721653723136.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 224376976834560.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14205816832.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 252659101597696.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1095330745548800.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 156399690579968.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 677743960260608.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25914962345984.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17721653723136.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 224376976834560.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5874640224256.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1218978382105.5024,
+      "step": 58,
+      "teaching_rms": 96472561411.09248,
+      "traffic_rms": 1221276892356.339
+    },
+    {
+      "batch_loss": 2.67451603546641e+28,
+      "innovation_rms": 169502001680.71426,
+      "instruction_rms": 32146446341.853085,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0020537273361028e+24,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.8983352480116133e+27,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18297792512.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27578431700992.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 397580626296832.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15811488768.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 445869346258944.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1932102312394752.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 276215285940224.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1198155886493696.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 70106717618176.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27578431700992.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 397580626296832.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7003008663552.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2143316300211.9094,
+      "step": 59,
+      "teaching_rms": 169502001680.71426,
+      "traffic_rms": 2147986742241.4607
+    },
+    {
+      "batch_loss": 6.371005068330419e+28,
+      "innovation_rms": 298105551232.3411,
+      "instruction_rms": 46207781064.9337,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8067620248438398e+24,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.5217965333274176e+27,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16466574336.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 57897794928640.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 770289499635712.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17256435712.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 787054971584512.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3411856521691136.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 488344693768192.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2121294011170816.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 140784649109504.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 57897794928640.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 770289499635712.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13872173940736.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3768671726229.25,
+      "step": 60,
+      "teaching_rms": 298105551232.3411,
+      "traffic_rms": 3775085172246.2993
+    },
+    {
+      "batch_loss": 2.38193936438965e+29,
+      "innovation_rms": 523767228994.9787,
+      "instruction_rms": 92643320973.90822,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.251606424951278e+24,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 3.903396863979796e+27,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14818313216.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 108205233405952.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1273358079492096.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18556715008.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1390406944161792.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6033520128425984.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 859997810458624.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3716531032686592.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 277485925171200.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 108205233405952.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1273358079492096.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26914873933824.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 6625947196101.517,
+      "step": 61,
+      "teaching_rms": 523767228994.9787,
+      "traffic_rms": 6639496103560.324
+    },
+    {
+      "batch_loss": 5.934490666855792e+29,
+      "innovation_rms": 922912634352.9395,
+      "instruction_rms": 143373853068.49213,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.838199079134042e+24,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.8424734304124517e+28,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13334729728.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 163619044065280.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2248582748962816.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19726780416.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2457977839681536.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0675709223632896e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1515156014104576.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6551581768024064.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 268708035428352.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 163619044065280.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2248582748962816.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 52736213123072.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 11660140630425.062,
+      "step": 62,
+      "teaching_rms": 922912634352.9395,
+      "traffic_rms": 11679465195425.225
+    },
+    {
+      "batch_loss": 1.8195569711850725e+30,
+      "innovation_rms": 1623465083773.3845,
+      "instruction_rms": 248224117131.73666,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0471690878438812e+25,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 4.701420898571525e+28,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11999371264.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 309622111993856.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4147594083172352.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20779642880.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4345714213650432.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8877362665947136e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2671637430272000.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1564813624803328e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1119606504685568.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 309622111993856.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4147594083172352.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 116086926213120.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 20512146352049.254,
+      "step": 63,
+      "teaching_rms": 1623465083773.3845,
+      "traffic_rms": 20545618593981.08
+    },
+    {
+      "batch_loss": 5.940844629847959e+30,
+      "innovation_rms": 2859247048262.0728,
+      "instruction_rms": 471892247072.92267,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8737238787696311e+25,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.0524112939049246e+29,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10797428736.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 572537259950080.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7214984027176960.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21727010816.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7680697531629568.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.334983264292045e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4719551656755200.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.047914280170291e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5122683192737792.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 572537259950080.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7214984027176960.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 507924845690880.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 943916032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 36088688991233.74,
+      "step": 64,
+      "teaching_rms": 2859247048262.0728,
+      "traffic_rms": 36149889801723.516
+    },
+    {
+      "batch_loss": 1.8546255518067446e+31,
+      "innovation_rms": 5029730257783.29,
+      "instruction_rms": 840084132397.7766,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.3510117782470822e+25,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 3.266960635911617e+30,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9715572736.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 989164455591936.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1397933744259072e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22579425280.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.357118934351872e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.890491382392422e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8331331389882368.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.611779733127168e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4597128982691840.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 989164455591936.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1397933744259072e+16,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 967637743960064.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6667206393856.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": 2.5754232327497582e+25,
+      "predictor_updated": true,
+      "raw_apical_rms": 63482532093786.34,
+      "step": 65,
+      "teaching_rms": 5029730257783.29,
+      "traffic_rms": 63591595748105.21
+    },
+    {
+      "batch_loss": 5.874105759083552e+31,
+      "innovation_rms": 7577226559827.043,
+      "instruction_rms": 1465299717630.1543,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.971547248570525e+25,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.5254808694142445e+31,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8741805056.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1641166193819648.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2574148391796736e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23346372608.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3158727800520704e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.58753781275689e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4399315946504192e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.067984771370189e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4116110093844480.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1641166193819648.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2574148391796736e+16,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1379248753344512.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6667206393856.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 111699032743182.0,
+      "step": 66,
+      "teaching_rms": 7577226559827.043,
+      "traffic_rms": 111894111246207.27
+    },
+    {
+      "batch_loss": 1.502198457196483e+32,
+      "innovation_rms": 13215118265392.182,
+      "instruction_rms": 2422746546713.2803,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0456695520181293e+26,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 3.978105809133734e+31,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7865327104.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3427640794939392.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.368889883610317e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24036390912.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.945805352691302e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6299323578908672e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4403525287018496e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0004208911017574e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.23162742603776e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3427640794939392.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.368889883610317e+16,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2245949128704000.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6667206393856.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 190329108465305.28,
+      "step": 67,
+      "teaching_rms": 13215118265392.182,
+      "traffic_rms": 190619050084061.8
+    },
+    {
+      "batch_loss": 4.6181832867306334e+32,
+      "innovation_rms": 22673223396299.367,
+      "instruction_rms": 4196004143154.863,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8113409563623213e+26,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.1221016376764966e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7076418048.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4532764602793984.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.341995094920397e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24657166336.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.720609776461414e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.774804509669458e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.159173744499098e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.718821172847575e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3881214314151936e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4532764602793984.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.341995094920397e+16,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4262312928083968.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6667206393856.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 323002012265234.7,
+      "step": 68,
+      "teaching_rms": 22673223396299.367,
+      "traffic_rms": 323535822085309.44
+    },
+    {
+      "batch_loss": 1.6636654524005538e+33,
+      "innovation_rms": 37055240941778.47,
+      "instruction_rms": 8234789510640.371,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.110075671837009e+26,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 0.0,
+          "nonfinite_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6366329344.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8861532754542592.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0436655424156467e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25215619072.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1445026127924429e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.724415921966285e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.706112731342438e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6072953032605696e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.160705704820736e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8861532754542592.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0436655424156467e+17,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6076997545492480.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6667206393856.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 524898032629981.6,
+      "step": 69,
+      "teaching_rms": 37055240941778.47,
+      "traffic_rms": 526078441445928.6
+    },
+    {
+      "batch_loss": 4.762813987915413e+33,
+      "innovation_rms": 62267737263059.23,
+      "instruction_rms": 13567037247827.457,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.3223859409644846e+26,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 0.0,
+          "nonfinite_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5727185920.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.324021379497984e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9545583556283597e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25717972992.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.951326135040082e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.068235566073774e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.098357733602099e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5826031769852314e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.944590251930419e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.324021379497984e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9545583556283597e+17,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7710152954216448.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6667206393856.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 879518366719532.0,
+      "step": 70,
+      "teaching_rms": 62267737263059.23,
+      "traffic_rms": 881545345761582.2
+    },
+    {
+      "batch_loss": 1.1772177186243892e+34,
+      "innovation_rms": 105448854263585.6,
+      "instruction_rms": 21158873389790.32,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.092974216607012e+26,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 0.0,
+          "nonfinite_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5151899648.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.68543903924224e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.0506826285973504e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26193645568.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.331835654559498e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3805094164402012e+18,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0668228034848358e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.401429570943713e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.108015046033408e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.68543903924224e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.0506826285973504e+17,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9179915995840512.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6667206393856.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1484377190953996.8,
+      "step": 71,
+      "teaching_rms": 105448854263585.6,
+      "traffic_rms": 1487355648934090.2
+    },
+    {
+      "batch_loss": 3.414349052600583e+34,
+      "innovation_rms": 179815256610566.84,
+      "instruction_rms": 36324376661328.055,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.554544164153604e+27,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 0.0,
+          "nonfinite_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4634089984.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.237856256871629e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.2215113136367206e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26657054720.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.693696582185124e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3618609276256256e+18,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8174513172289946e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.506284793844204e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4447035453210624e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.237856256871629e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.2215113136367206e+17,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0502610230444032e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6667206393856.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2526307426205217.0,
+      "step": 72,
+      "teaching_rms": 179815256610566.84,
+      "traffic_rms": 2531380481157864.0
+    },
+    {
+      "batch_loss": 1.1035653123234685e+35,
+      "innovation_rms": 306491149936605.6,
+      "instruction_rms": 66106915631431.484,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.658010243046741e+27,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 0.0,
+          "nonfinite_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4168015104.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.216375703955046e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.823696948818084e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27073855488.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.731023564785582e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.037327257478365e+18,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.098541964572754e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2810907504229745e+18,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8176167593692365e+17,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.216375703955046e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.823696948818084e+17,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7934014975311872e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6667206393856.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4303030415888376.5,
+      "step": 73,
+      "teaching_rms": 306491149936605.6,
+      "traffic_rms": 4312633895109593.0
+    },
+    {
+      "batch_loss": 3.325870308866356e+35,
+      "innovation_rms": 515739569426776.25,
+      "instruction_rms": 112412323506020.67,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.545171191368264e+27,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 0.0,
+          "nonfinite_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3748506112.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1295779450362266e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.559115322206716e+18,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27448707072.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6631621075431588e+18,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.900597648084959e+18,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.275939200795935e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1773972362231808e+18,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 3748506112.0,
+          "nonfinite_indices": [
+            0
+          ]
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1295779450362266e+17,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.559115322206716e+18,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 27448707072.0,
+          "nonfinite_indices": [
+            0
+          ]
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6667206393856.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": Infinity,
+      "step": 74,
+      "teaching_rms": 515739569426776.25,
+      "traffic_rms": Infinity
+    }
+  ],
+  "trajectory_metrics_finite": false,
+  "validation_evaluations": 0
+}
diff --git a/results/kp_innovation_diagnosis/matched.json b/results/kp_innovation_diagnosis/matched.json
new file mode 100644
index 0000000..e930160
--- /dev/null
+++ b/results/kp_innovation_diagnosis/matched.json
@@ -0,0 +1,7316 @@
+{
+  "first_any_nonfinite_groups": {
+    "bn_running_var": [
+      2
+    ]
+  },
+  "first_any_nonfinite_step": 32,
+  "first_nonfinite_by_group": {
+    "bn_running_var": {
+      "indices": [
+        2
+      ],
+      "step": 32
+    },
+    "forward_momentum": {
+      "indices": [
+        0
+      ],
+      "step": 69
+    },
+    "forward_weight": {
+      "indices": [
+        0
+      ],
+      "step": 69
+    }
+  },
+  "first_nonfinite_by_metric": {
+    "raw_apical_rms": 68,
+    "traffic_rms": 68
+  },
+  "first_nonfinite_metrics": [],
+  "first_training_failure_groups": {
+    "forward_momentum": [
+      0
+    ],
+    "forward_weight": [
+      0
+    ]
+  },
+  "first_training_failure_step": 69,
+  "initial_parameter_state": {
+    "bn_running_mean": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "bn_running_var": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 1.0,
+      "nonfinite_indices": []
+    },
+    "feedback_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "feedback_readout": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.4011031687259674,
+      "nonfinite_indices": []
+    },
+    "feedback_readout_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "feedback_weight": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.8331910371780396,
+      "nonfinite_indices": []
+    },
+    "forward_bn_bias": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_bn_bias_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_bn_scale": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 1.0,
+      "nonfinite_indices": []
+    },
+    "forward_bn_scale_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_readout": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.3956821858882904,
+      "nonfinite_indices": []
+    },
+    "forward_readout_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_weight": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 1.1162850856781006,
+      "nonfinite_indices": []
+    },
+    "predictor_bias": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.027761053293943405,
+      "nonfinite_indices": []
+    },
+    "predictor_slope": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.09883774071931839,
+      "nonfinite_indices": []
+    }
+  },
+  "max_steps": 352,
+  "predictor_warmup": {
+    "first_mse": 0.0002822713358909823,
+    "last_mse": 5.6679267276115794e-06,
+    "post_warmup_traffic_residual_rms_ratio": 0.13994200326198175,
+    "steps": 20,
+    "task_loader_state_restored": true
+  },
+  "protocol": "kp_mixed_traffic_nonfinite_diagnosis_v1",
+  "provenance": {
+    "git_commit": "e9962317b3d976fe25e28a5693f2757e3758480a",
+    "git_tracked_dirty": false
+  },
+  "rule": "matched",
+  "scope": "training_only_no_validation_or_test_evaluation",
+  "split": {
+    "dataset": "cifar10",
+    "input_layout": "NCHW",
+    "input_shape": [
+      3,
+      32,
+      32
+    ],
+    "loader_seed": 0,
+    "normalization_mean": [
+      0.49140000343322754,
+      0.4821999967098236,
+      0.4465000033378601
+    ],
+    "normalization_std": [
+      0.24699999392032623,
+      0.2434999942779541,
+      0.26159998774528503
+    ],
+    "split_from_training_only": true,
+    "split_seed": 2027,
+    "test_examples": 10000,
+    "train_examples": 45000,
+    "training_augmentation": "random_crop_32_padding_4_zero_then_horizontal_flip_p0.5",
+    "validation_class_counts": {
+      "0": 500,
+      "1": 500,
+      "2": 500,
+      "3": 500,
+      "4": 500,
+      "5": 500,
+      "6": 500,
+      "7": 500,
+      "8": 500,
+      "9": 500
+    },
+    "validation_examples": 5000,
+    "validation_index_sha256": "8328b206a97c420e49e54e3eca4abe3274c4756b084355784ea3fb8059e4515b"
+  },
+  "test_evaluations": 0,
+  "traffic_calibration": {
+    "instruction_rms": [
+      0.008471464738249779,
+      0.006617916747927666,
+      0.003988116513937712,
+      0.0043865954503417015,
+      0.002713605063036084,
+      0.003158062230795622,
+      0.002004207344725728,
+      0.003757704282179475,
+      0.0027085733599960804,
+      0.0030200635083019733,
+      0.002056022873148322,
+      0.0023899395018815994,
+      0.0016842411132529378,
+      0.0034010226372629404,
+      0.002549725119024515,
+      0.0030542907770723104,
+      0.0020918569061905146,
+      0.002605808200314641,
+      0.0018916188273578882
+    ],
+    "realized_traffic_instruction_rms_ratio": [
+      4.0,
+      4.0,
+      4.0,
+      4.0,
+      4.0,
+      4.0,
+      3.999999523162842,
+      4.0,
+      4.000000476837158,
+      3.999999761581421,
+      4.0,
+      4.0,
+      3.999999761581421,
+      4.0,
+      4.000000476837158,
+      4.000000476837158,
+      4.0,
+      4.0,
+      4.0
+    ],
+    "target_ratio": 4.0,
+    "traffic_gain": [
+      0.0443243607878685,
+      0.035817068070173264,
+      0.013778294436633587,
+      0.022762222215533257,
+      0.007214079145342112,
+      0.017218582332134247,
+      0.004616298712790012,
+      0.02042512409389019,
+      0.007327491883188486,
+      0.01611751690506935,
+      0.004703814629465342,
+      0.012562948279082775,
+      0.003324519144371152,
+      0.01805141195654869,
+      0.006063441745936871,
+      0.016296451911330223,
+      0.0044285994954407215,
+      0.013874167576432228,
+      0.0036007820162922144
+    ],
+    "unscaled_traffic_rms": [
+      0.7644973993301392,
+      0.7390796542167664,
+      1.157796859741211,
+      0.7708553671836853,
+      1.504616141319275,
+      0.7336404919624329,
+      1.736635684967041,
+      0.7358984351158142,
+      1.478581428527832,
+      0.7495108246803284,
+      1.7483876943588257,
+      0.7609485983848572,
+      2.0264477729797363,
+      0.7536302804946899,
+      1.6820316314697266,
+      0.7496824264526367,
+      1.8894071578979492,
+      0.7512690424919128,
+      2.10134220123291
+    ]
+  },
+  "trajectory": [
+    {
+      "batch_loss": 2.9134020805358887,
+      "innovation_rms": 0.0047711752050758295,
+      "instruction_rms": 0.004215397969971197,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.40285658836364746,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9875131845474243,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8092546463012695,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4058851897716522,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6322960257530212,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8331910371780396,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5760545134544373,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.760544776916504,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.69287109375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.9287109375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8092595338821411,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3794170618057251,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.632293164730072,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.114573359489441,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": 4.976208305399031e-06,
+      "predictor_updated": true,
+      "raw_apical_rms": 0.01731826254949561,
+      "step": 1,
+      "teaching_rms": 0.0047711752485875515,
+      "traffic_rms": 0.016800932589918403
+    },
+    {
+      "batch_loss": 2.8719682693481445,
+      "innovation_rms": 0.0030266363028306265,
+      "instruction_rms": 0.001145980237518358,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.090774655342102,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.9997076988220215,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8277368545532227,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4217219352722168,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6493300199508667,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8331910371780396,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.485217571258545,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.091630935668945,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7201173305511475,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.289445877075195,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8278238773345947,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3897894620895386,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6493247747421265,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1112847328186035,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.03202896469662201,
+      "step": 2,
+      "teaching_rms": 0.0030266363158957077,
+      "traffic_rms": 0.03201675847811013
+    },
+    {
+      "batch_loss": 4.823637008666992,
+      "innovation_rms": 0.0050067298630821185,
+      "instruction_rms": 0.000896361821657787,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1607327461242676,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 42.92159652709961,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.960029125213623,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.44191792607307434,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.919400930404663,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8331910371780396,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.084751605987549,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15.995339393615723,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.417537689208984,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16.97420310974121,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.9600157737731934,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.376192182302475,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9193487167358398,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1031781435012817,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.05582828645223682,
+      "step": 3,
+      "teaching_rms": 0.005006729885788924,
+      "traffic_rms": 0.055838764401403174
+    },
+    {
+      "batch_loss": 12.992831230163574,
+      "innovation_rms": 0.009176180255596859,
+      "instruction_rms": 0.000969024823408796,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.619744300842285,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 299.88580322265625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.021809101104736,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8468732237815857,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.011368751525879,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8331910371780396,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.185171127319336,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31.004196166992188,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.410898208618164,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30.688411712646484,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.021817207336426,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5456865429878235,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.011358737945557,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.137107491493225,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.09631751081446588,
+      "step": 4,
+      "teaching_rms": 0.009176180124162098,
+      "traffic_rms": 0.0963438970013651
+    },
+    {
+      "batch_loss": 40.78102493286133,
+      "innovation_rms": 0.0170937088699459,
+      "instruction_rms": 0.0011944011917212948,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25.037996292114258,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3124.3330078125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.6131322383880615,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4538477659225464,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14.207832336425781,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1268450021743774,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.339202880859375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 61.540321350097656,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.166619300842285,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 57.78557586669922,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.865116119384766,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3744301795959473,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14.207892417907715,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.409882664680481,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.17034175872804092,
+      "step": 5,
+      "teaching_rms": 0.017093708838716058,
+      "traffic_rms": 0.17040170092793797
+    },
+    {
+      "batch_loss": 142.98423767089844,
+      "innovation_rms": 0.03204800427212605,
+      "instruction_rms": 0.002545680320603145,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 60.35129928588867,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23701.150390625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.707615852355957,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.425274610519409,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22.89688491821289,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4470754861831665,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24.551959991455078,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 122.12757873535156,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24.273162841796875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 111.06543731689453,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.945295333862305,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.341252326965332,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22.89682388305664,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3092474937438965,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.30763197201412334,
+      "step": 6,
+      "teaching_rms": 0.03204800457411708,
+      "traffic_rms": 0.30777346909107683
+    },
+    {
+      "batch_loss": 957.0724487304688,
+      "innovation_rms": 0.060249450271203195,
+      "instruction_rms": 0.006955649734660194,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 141.80287170410156,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 149655.578125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30.096025466918945,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.7502431869506836,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 45.717063903808594,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.3656625747680664,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 48.7338981628418,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 241.81936645507812,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 45.16677474975586,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 208.93612670898438,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 67.373291015625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.854081392288208,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 45.717010498046875,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.786167144775391,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.5666255595996695,
+      "step": 7,
+      "teaching_rms": 0.06024944964375605,
+      "traffic_rms": 0.5672159518177116
+    },
+    {
+      "batch_loss": 2516.896728515625,
+      "innovation_rms": 0.11463434948054016,
+      "instruction_rms": 0.010416129946149344,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 676.8906860351562,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5953663.5,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 73.1310043334961,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.238974571228027,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 90.01850128173828,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.138430118560791,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 96.53346252441406,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 477.9956359863281,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 84.43772888183594,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 392.7095031738281,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 123.81836700439453,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.175546646118164,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 90.01856994628906,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.721455574035645,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.065368551233577,
+      "step": 8,
+      "teaching_rms": 0.11463434722593957,
+      "traffic_rms": 1.0661871334601167
+    },
+    {
+      "batch_loss": 12562.6298828125,
+      "innovation_rms": 0.21906842247060332,
+      "instruction_rms": 0.023890223646902072,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2941.41455078125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 73791240.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 66.87805938720703,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.542716979980469,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 166.09530639648438,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.8262357711792,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 191.1508026123047,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 946.17333984375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 159.6721649169922,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 752.3444213867188,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 225.3828582763672,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.429659843444824,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 166.0952911376953,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24.92508888244629,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2.024478018763723,
+      "step": 9,
+      "teaching_rms": 0.2190684279551091,
+      "traffic_rms": 2.026687289958201
+    },
+    {
+      "batch_loss": 41096.91796875,
+      "innovation_rms": 0.42033320389995077,
+      "instruction_rms": 0.042051454657006576,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8817.6533203125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 496269408.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 61.62916564941406,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.990135192871094,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 239.3929901123047,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.989152908325195,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 376.84564208984375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1856.9482421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 302.9001770019531,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1432.2801513671875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 685.1502685546875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22.007396697998047,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 239.39297485351562,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 67.72527313232422,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3.870581658548543,
+      "step": 10,
+      "teaching_rms": 0.4203332023646312,
+      "traffic_rms": 3.874389597577872
+    },
+    {
+      "batch_loss": 124862.0859375,
+      "innovation_rms": 0.8093676972566564,
+      "instruction_rms": 0.06362515517174869,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22300.4375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2603202816.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 133.48045349121094,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 52.99477767944336,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 665.3596801757812,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25.835433959960938,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 742.773193359375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3659.275634765625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 577.0498046875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2741.49609375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 616.894287109375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 52.847938537597656,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 665.3597412109375,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 129.41470336914062,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 7.441567999133049,
+      "step": 11,
+      "teaching_rms": 0.8093676670012973,
+      "traffic_rms": 7.447661733428294
+    },
+    {
+      "batch_loss": 401691.25,
+      "innovation_rms": 1.559591518893618,
+      "instruction_rms": 0.13834452550643744,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 50808.71484375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11259937792.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 774.3778076171875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 102.32852172851562,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 855.5213012695312,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 79.6632308959961,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1463.8865966796875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7211.1337890625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1105.9173583984375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5345.939453125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2050.651123046875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 102.18168640136719,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 855.5213623046875,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 208.39300537109375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 14.300715009784797,
+      "step": 12,
+      "teaching_rms": 1.5595915724301312,
+      "traffic_rms": 14.310771240033576
+    },
+    {
+      "batch_loss": 2402618.0,
+      "innovation_rms": 2.9904745645326396,
+      "instruction_rms": 0.4212605830856891,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 148625.90625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 147667501056.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3336.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 193.5659637451172,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2443.235595703125,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 351.6991882324219,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2879.03515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14151.486328125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2126.80322265625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10208.857421875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3336.000244140625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 193.52435302734375,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2443.23583984375,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 391.2945556640625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 27.351043812320068,
+      "step": 13,
+      "teaching_rms": 2.9904746722602304,
+      "traffic_rms": 27.38133589764681
+    },
+    {
+      "batch_loss": 9321861.0,
+      "innovation_rms": 5.855092715569528,
+      "instruction_rms": 1.1294882783900348,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1116235.875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9418870620160.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7494.45263671875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 360.65252685546875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4648.84375,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 774.2950439453125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5654.7578125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27815.66015625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4100.36181640625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19735.5859375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7494.45263671875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 360.758056640625,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4648.84375,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 774.3240356445312,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 53.08121896419405,
+      "step": 14,
+      "teaching_rms": 5.855092710590196,
+      "traffic_rms": 53.136161440073835
+    },
+    {
+      "batch_loss": 38468964.0,
+      "innovation_rms": 11.327182100279796,
+      "instruction_rms": 2.083964751623628,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13876809.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1243749514477568.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23501.75390625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 712.19140625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9306.037109375,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2351.82763671875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11108.0244140625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 54967.109375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7953.603515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38532.41796875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23501.75390625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 712.138427734375,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9306.037109375,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2351.876220703125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 102.95399884310817,
+      "step": 15,
+      "teaching_rms": 11.327181863943343,
+      "traffic_rms": 103.07497727430156
+    },
+    {
+      "batch_loss": 125100336.0,
+      "innovation_rms": 21.898488531878794,
+      "instruction_rms": 3.7994835882984095,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 65247588.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.014340520817459e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21145.576171875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1586.7659912109375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19559.03515625,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4466.38525390625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21963.986328125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 108559.609375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15210.0947265625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 72564.90625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21145.576171875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1586.98095703125,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19559.03515625,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4466.43408203125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 199.1116505559766,
+      "step": 16,
+      "teaching_rms": 21.898488039576947,
+      "traffic_rms": 199.3150802870113
+    },
+    {
+      "batch_loss": 406332640.0,
+      "innovation_rms": 43.4989279872756,
+      "instruction_rms": 11.814265867935642,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 283797792.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5248195459783066e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19044.82421875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3288.462890625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39172.61328125,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6370.86767578125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 43422.44921875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 214584.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28707.908203125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 134978.125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 37808.76953125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3288.47998046875,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39172.61328125,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6370.91650390625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": 1786.4998506803806,
+      "predictor_updated": true,
+      "raw_apical_rms": 386.07989562731865,
+      "step": 17,
+      "teaching_rms": 43.49892816457474,
+      "traffic_rms": 386.30057467395926
+    },
+    {
+      "batch_loss": 1876102784.0,
+      "innovation_rms": 78.36551715918127,
+      "instruction_rms": 35.43632012032303,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 842817408.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7595716851130696e+18,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17153.703125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5961.1337890625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 70390.171875,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8086.23779296875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 82438.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 390155.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 52523.421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 238155.125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 49051.6484375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5961.15087890625,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 70390.171875,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8086.28662109375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 751.0638513169803,
+      "step": 18,
+      "teaching_rms": 78.36551690097482,
+      "traffic_rms": 751.1796847962376
+    },
+    {
+      "batch_loss": 7756428800.0,
+      "innovation_rms": 155.55207735763057,
+      "instruction_rms": 80.24130903932937,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2082047232.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.09434424128281e+18,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58303.796875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11256.408203125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 145102.84375,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9630.576171875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 155694.796875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 732568.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 95824.796875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 433013.75,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58303.796875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11256.26171875,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 145102.84375,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9994.77734375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1406.7281071438276,
+      "step": 19,
+      "teaching_rms": 155.5520763596266,
+      "traffic_rms": 1406.5372413991738
+    },
+    {
+      "batch_loss": 25476413440.0,
+      "innovation_rms": 352.0562300268429,
+      "instruction_rms": 247.02279571324402,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4634672640.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.528446931144016e+19,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 452802.875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21397.203125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 260395.046875,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 45520.20703125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 294313.8125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1386190.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 175403.9375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 809516.1875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 452802.875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21397.056640625,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 260395.046875,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 45520.37109375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2624.8605174338586,
+      "step": 20,
+      "teaching_rms": 352.0562290318634,
+      "traffic_rms": 2616.9130921938345
+    },
+    {
+      "batch_loss": 75591802880.0,
+      "innovation_rms": 538.3110576557918,
+      "instruction_rms": 256.499206197742,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17902606336.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.590727301251891e+21,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 407473.6875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38487.71875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 469803.28125,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86267.578125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 556835.0625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2625212.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 325503.3125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1525460.75,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 407473.6875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38487.609375,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 469803.28125,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86267.7421875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4898.42575712612,
+      "step": 21,
+      "teaching_rms": 538.3110524500474,
+      "traffic_rms": 4897.572550319883
+    },
+    {
+      "batch_loss": 388053598208.0,
+      "innovation_rms": 918.2520110837629,
+      "instruction_rms": 264.953018332314,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 76294840320.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.54064873706541e+22,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 366310.5625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 73924.09375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 934063.0625,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 122898.6328125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1054451.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4976166.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 611740.375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2862370.5,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 704088.25,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 73924.125,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 934063.0625,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 122898.796875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 9112.967106540247,
+      "step": 22,
+      "teaching_rms": 918.2520124709389,
+      "traffic_rms": 9125.282734985898
+    },
+    {
+      "batch_loss": 1290553917440.0,
+      "innovation_rms": 1850.1333710003532,
+      "instruction_rms": 841.9603339924298,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 221404217344.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3238449777882044e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 855653.3125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 146958.328125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1795860.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 155803.484375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1997824.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9433730.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1150768.375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5390280.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1107408.875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 146958.359375,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1795860.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 155803.65625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 17049.091355620974,
+      "step": 23,
+      "teaching_rms": 1850.1333216593885,
+      "traffic_rms": 17057.213157419923
+    },
+    {
+      "batch_loss": 3708327034880.0,
+      "innovation_rms": 3592.879400879788,
+      "instruction_rms": 1803.5179418880082,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 541278404608.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4978581808474584e+24,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2240393.75,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 251259.203125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1943588.25,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 227409.078125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3786846.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17890222.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2146363.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9955953.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2401297.25,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 251259.09375,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1943588.25,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 245764.578125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 32044.791877423566,
+      "step": 24,
+      "teaching_rms": 3592.879415817244,
+      "traffic_rms": 32038.190557034522
+    },
+    {
+      "batch_loss": 9880979111936.0,
+      "innovation_rms": 6404.828404834937,
+      "instruction_rms": 2625.7189367129095,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1198418558976.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7887215126622784e+25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5847975.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 543186.5,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6718434.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 600033.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7179247.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33924008.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4039358.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18929946.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5847975.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 543186.25,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6718434.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 600033.125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 59989.34354345325,
+      "step": 25,
+      "teaching_rms": 6404.828375666216,
+      "traffic_rms": 59994.730643837574
+    },
+    {
+      "batch_loss": 51177890775040.0,
+      "innovation_rms": 11406.657778535764,
+      "instruction_rms": 3240.587408590824,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5008838361088.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.269158497298987e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5262193.5,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1006085.4375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11107101.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1126252.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13615999.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 64367512.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7605074.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35657160.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5472081.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1006085.1875,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11107101.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1126252.5,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 112724.71282270785,
+      "step": 26,
+      "teaching_rms": 11406.65778627819,
+      "traffic_rms": 112850.36394028923
+    },
+    {
+      "batch_loss": 215583400394752.0,
+      "innovation_rms": 21324.80421108959,
+      "instruction_rms": 5960.228649046364,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21327102607360.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.980302031829473e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4734888.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1642832.375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22819428.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1599741.375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25789916.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 121739160.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14316775.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 67117008.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4852504.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1642832.25,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22819428.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1599741.25,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 211604.99588003894,
+      "step": 27,
+      "teaching_rms": 21324.804192485004,
+      "traffic_rms": 211908.0223644778
+    },
+    {
+      "batch_loss": 621414121996288.0,
+      "innovation_rms": 41268.64956440379,
+      "instruction_rms": 14598.393613235552,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 61453298040832.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.367699175674086e+27,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14747591.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3189767.5,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 42317956.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2024561.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 48855004.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 230650896.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26960114.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 126433384.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18738458.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3189767.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 42317956.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2024561.625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 397754.5371817846,
+      "step": 28,
+      "teaching_rms": 41268.64958126046,
+      "traffic_rms": 398081.2764087075
+    },
+    {
+      "batch_loss": 2447422152245248.0,
+      "innovation_rms": 77244.16429231007,
+      "instruction_rms": 26867.638212735314,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 151000413372416.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.294007528911071e+28,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20010302.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5933548.5,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 82186424.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2807743.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 92529736.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 436747328.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 50822772.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 238626576.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24096832.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5933548.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 82186424.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3557889.5,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 747411.110312879,
+      "step": 29,
+      "teaching_rms": 77244.16438912335,
+      "traffic_rms": 748172.8964952126
+    },
+    {
+      "batch_loss": 8588165065474048.0,
+      "innovation_rms": 174703.42735544467,
+      "instruction_rms": 110193.72896541958,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 338035468140544.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9081184565950123e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 262652112.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11140782.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 152254656.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26274578.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 175166640.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 826368960.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 95748584.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 449258144.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 262652112.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11140782.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 152254656.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26274578.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1402575.6587049598,
+      "step": 30,
+      "teaching_rms": 174703.42967689433,
+      "traffic_rms": 1400577.66314253
+    },
+    {
+      "batch_loss": 2.806604361629696e+16,
+      "innovation_rms": 264924.50688549323,
+      "instruction_rms": 72916.09453158562,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4250259639238656.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.163915969585317e+31,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 623829760.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24135010.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 305817600.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 62872944.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 331736832.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1565701888.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 180599616.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 848510336.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 623829760.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24135010.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 305817600.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 62872944.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2624685.649440279,
+      "step": 31,
+      "teaching_rms": 264924.5180506097,
+      "traffic_rms": 2627727.361342049
+    },
+    {
+      "batch_loss": 1.0868606298069402e+17,
+      "innovation_rms": 491477.6974776303,
+      "instruction_rms": 120402.69522911347,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.829003018318643e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.0629836176722922e+32,
+          "nonfinite_indices": [
+            2
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 586455168.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 44699012.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 570195584.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 119017000.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 628130688.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2963938816.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 337183072.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1565834624.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 586455168.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 44699012.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 570195584.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 119017000.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.87675857543945,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4917847.6396826105,
+      "step": 32,
+      "teaching_rms": 491477.7072157871,
+      "traffic_rms": 4924786.6371714305
+    },
+    {
+      "batch_loss": 4.558429179871232e+17,
+      "innovation_rms": 932941.5113839154,
+      "instruction_rms": 275044.67365834606,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.683223132635136e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 9.02863172403141e+32,
+          "nonfinite_indices": [
+            1,
+            2
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4397747200.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 82793136.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1101971200.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 440058016.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1188713216.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5605825536.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 632294784.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2951117312.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4397747200.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 82793136.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1101971200.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 440058016.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": 828895664443.6073,
+      "predictor_updated": true,
+      "raw_apical_rms": 9230354.221622748,
+      "step": 33,
+      "teaching_rms": 932941.516657185,
+      "traffic_rms": 9245036.648267511
+    },
+    {
+      "batch_loss": 1.4377044994886205e+18,
+      "innovation_rms": 1630905.8140535704,
+      "instruction_rms": 778996.5591467236,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2015425115691418e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 3.592052080173146e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            5,
+            7
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3957928448.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 160324352.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2097789568.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 835850880.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2168863232.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9801499648.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1149754368.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5174596096.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3957928448.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 160324352.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2097789568.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 835850880.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 17328520.098337166,
+      "step": 34,
+      "teaching_rms": 1630905.7658985215,
+      "traffic_rms": 17342804.19631902
+    },
+    {
+      "batch_loss": 4.715965350473105e+18,
+      "innovation_rms": 3071681.6172405006,
+      "instruction_rms": 1563839.6575119062,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.017156019725271e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 7.353650106989244e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            5,
+            7,
+            11
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3562051840.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 286753952.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3678433536.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1192056064.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3935846144.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17669828608.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2080810368.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9310560256.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3562051840.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 286753952.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3678433536.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1192056064.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 31278835.776587892,
+      "step": 35,
+      "teaching_rms": 3071681.6784642367,
+      "traffic_rms": 31296740.980714537
+    },
+    {
+      "batch_loss": 1.169217366442115e+19,
+      "innovation_rms": 6107161.256647392,
+      "instruction_rms": 3701370.7145119845,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.706031706619511e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.0540743759804995e+33,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            5,
+            6,
+            7,
+            9,
+            11,
+            13,
+            15
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3205727232.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 547618624.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6541569536.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1512628736.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7154016256.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 32181702656.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3773047552.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16922372096.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3205727232.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 547618624.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6541569536.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1512628736.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 56555827.00227312,
+      "step": 36,
+      "teaching_rms": 6107161.150270437,
+      "traffic_rms": 56512560.933729336
+    },
+    {
+      "batch_loss": 4.465618537505253e+19,
+      "innovation_rms": 22884459.671562485,
+      "instruction_rms": 21112871.173231423,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.961380300135858e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.8265280908081528e+30,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            9,
+            11,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2885003264.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 911123648.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12269554688.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1801129088.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13023738880.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58697228288.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6843403264.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30703554560.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7166167040.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 911123648.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12269554688.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1801129088.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 104305906.4211801,
+      "step": 37,
+      "teaching_rms": 22884459.649897385,
+      "traffic_rms": 102312342.55797268
+    },
+    {
+      "batch_loss": 1.6445741982710707e+20,
+      "innovation_rms": 36050992.676571175,
+      "instruction_rms": 32294724.55443593,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.013098917927636e+19,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 7.53895957589338e+31,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            9,
+            11,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2596322560.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1790105344.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20802514944.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2060761344.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23738376192.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 107146371072.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12429577216.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 55861739520.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 32268038144.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1790105344.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20802514944.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3230062336.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 187856251.11052728,
+      "step": 38,
+      "teaching_rms": 36050991.8669079,
+      "traffic_rms": 185393605.01866084
+    },
+    {
+      "batch_loss": 4.624345003661229e+20,
+      "innovation_rms": 76675481.30660689,
+      "instruction_rms": 70901461.77702333,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.025385624666479e+19,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.0512674243126724e+33,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            9,
+            11,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2336484096.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2979149056.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39319547904.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2294409728.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 43316940800.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 195785654272.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22614042624.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 101844647936.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29044692992.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2979149056.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39319547904.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6134531584.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 342914968.22547644,
+      "step": 39,
+      "teaching_rms": 76675481.32503909,
+      "traffic_rms": 336020683.30498403
+    },
+    {
+      "batch_loss": 2.186346748500859e+21,
+      "innovation_rms": 65680749.726455815,
+      "instruction_rms": 39341059.124830544,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1486039423468962e+20,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 7.046196866257376e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            11,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2102606080.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5449299456.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63949631488.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2504670464.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 79169110016.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 358521700352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41168171008.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 185541296128.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27971747840.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5449299456.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63949631488.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8746092544.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 609585854.6884223,
+      "step": 40,
+      "teaching_rms": 65680749.61331441,
+      "traffic_rms": 609668932.6573776
+    },
+    {
+      "batch_loss": 6.348771119550322e+21,
+      "innovation_rms": 105919759.12961195,
+      "instruction_rms": 45230820.53852924,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4797160836711213e+20,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.7000454871720327e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1892094976.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9590532096.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 133608169472.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2693880064.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 144703111168.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 655340011520.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 75010285568.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 338421121024.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41843515392.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9590532096.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 133608169472.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11087418368.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1105045016.2017956,
+      "step": 41,
+      "teaching_rms": 105919759.13304187,
+      "traffic_rms": 1106302175.7715905
+    },
+    {
+      "batch_loss": 1.9624367822272526e+22,
+      "innovation_rms": 188750645.33160675,
+      "instruction_rms": 72718352.50208162,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.1583899398804026e+20,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 8.749641045940287e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1702615936.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18135588864.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 245362851840.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2864141568.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 264566013952.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1198629060608.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 135753531392.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 607432409088.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38084112384.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18135588864.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 245362851840.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13186196480.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2005978197.4074543,
+      "step": 42,
+      "teaching_rms": 188750644.1394708,
+      "traffic_rms": 2008378415.9746413
+    },
+    {
+      "batch_loss": 6.90186595849705e+22,
+      "innovation_rms": 343104194.1755184,
+      "instruction_rms": 133636035.70609435,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0460816958530624e+21,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.852642177767109e+33,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1532067840.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33692041216.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 462867038208.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3017348352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 483656925184.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2190909112320.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 242685739008.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1071206236160.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 303714336768.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33692041216.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 462867038208.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30562426880.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3643555534.23608,
+      "step": 43,
+      "teaching_rms": 343104193.5547307,
+      "traffic_rms": 3648303915.502025
+    },
+    {
+      "batch_loss": 1.7492436716989296e+23,
+      "innovation_rms": 665062902.8597152,
+      "instruction_rms": 330469266.29840773,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0789295642010966e+21,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 3.242531509196918e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5199283712.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 68234629120.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 751146303488.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3155204352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 884552433664.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4008955084800.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 440363778048.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1976780455936.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 374071689216.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 68234629120.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 751146303488.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 57989234688.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 6624699095.451093,
+      "step": 44,
+      "teaching_rms": 665062913.9133137,
+      "traffic_rms": 6626466945.365421
+    },
+    {
+      "batch_loss": 8.08633490827765e+23,
+      "innovation_rms": 1182546426.717914,
+      "instruction_rms": 561724426.2461624,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.07595694705381e+21,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 3.088351222492212e+33,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8062496256.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 102501179392.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1443043213312.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3279243264.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1617922359296.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7333699649536.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 801693237248.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3613294854144.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 437574762496.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 102501179392.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1443043213312.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 82840969216.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 12019784905.352076,
+      "step": 45,
+      "teaching_rms": 1182546476.8745863,
+      "traffic_rms": 12032101840.046946
+    },
+    {
+      "batch_loss": 2.3439986581782814e+24,
+      "innovation_rms": 2038196397.8221831,
+      "instruction_rms": 745738997.7654105,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.887741747165321e+21,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 8.739012599090855e+22,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8670181376.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 197775720448.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2715475509248.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3390845440.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2957072400384.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13391500410880.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1462042624000.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6603493343232.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 863094374400.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 197775720448.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2715475509248.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 105448300544.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 21837507161.30762,
+      "step": 46,
+      "teaching_rms": 2038196384.971172,
+      "traffic_rms": 21865384267.01417
+    },
+    {
+      "batch_loss": 7.855109582999719e+24,
+      "innovation_rms": 3553559734.945827,
+      "instruction_rms": 868244453.4951788,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5138343492447843e+22,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.0797704143990243e+23,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10974387200.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 390880559104.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5104688168960.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3491253504.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5401906511872.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24448342163456.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2667921801216.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12058791772160.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2246027771904.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 390880559104.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5104688168960.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 213666430976.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 39664893553.34978,
+      "step": 47,
+      "teaching_rms": 3553559734.903124,
+      "traffic_rms": 39730032199.1302
+    },
+    {
+      "batch_loss": 2.2973291450707742e+25,
+      "innovation_rms": 6402711945.945457,
+      "instruction_rms": 1272611182.8580687,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.8807261505676496e+22,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 9.839208661420225e+23,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9876857856.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 760697126912.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8912675799040.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3581585664.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9888673562624.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 44867669458944.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4867270639616.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21993487335424.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2806785507328.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 760697126912.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8912675799040.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 414114775040.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734350.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1025448739528656,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 72062335392.84178,
+      "step": 48,
+      "teaching_rms": 6402711976.273842,
+      "traffic_rms": 72171629936.54803
+    },
+    {
+      "batch_loss": 7.350969406129424e+25,
+      "innovation_rms": 11520202117.415573,
+      "instruction_rms": 1613939528.6153402,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.451433059626406e+22,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 4.284910882303284e+24,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8888983552.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1255557562368.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15863276109824.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3662849024.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18071131521024.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 81824587972608.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8877832667136.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 40105616080896.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3118330806272.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1255557562368.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15863276109824.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 591650750464.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": 1.3507288403250502e+20,
+      "predictor_updated": true,
+      "raw_apical_rms": 130926686259.11829,
+      "step": 49,
+      "teaching_rms": 11520202119.526861,
+      "traffic_rms": 131130416276.4666
+    },
+    {
+      "batch_loss": 2.9509546108595127e+26,
+      "innovation_rms": 17882632778.5258,
+      "instruction_rms": 3296595429.4604616,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0272283808787427e+23,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.0982356706317332e+25,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7999807488.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2155823497216.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29070572126208.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3820654848.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31901458890752.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 138303282085888.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15479623122944.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 66149693784064.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28049458331648.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2155823497216.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29070572126208.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2816138018816.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 237850947392.67108,
+      "step": 50,
+      "teaching_rms": 17882632775.724625,
+      "traffic_rms": 238300059231.5897
+    },
+    {
+      "batch_loss": 8.036544678770688e+26,
+      "innovation_rms": 32126992994.75963,
+      "instruction_rms": 5083706377.036581,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8902985737559947e+23,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.284886744698554e+26,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7199469056.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4121473056768.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 54616251170816.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4419167744.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 56032395001856.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 241309365305344.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26805425668096.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 115532799934464.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25233977245696.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4121473056768.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 54616251170816.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5339535638528.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 418180373247.3266,
+      "step": 51,
+      "teaching_rms": 32126992160.08077,
+      "traffic_rms": 418879771682.8226
+    },
+    {
+      "batch_loss": 2.871050362188439e+27,
+      "innovation_rms": 56834209207.51279,
+      "instruction_rms": 9963569430.28195,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.4316395134088735e+23,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 5.781939714831011e+26,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6479092224.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7876136402944.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 95757629652992.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5039434240.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 98496443056128.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 424640463765504.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 47145629515776.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 203402051059712.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23639472734208.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7876136402944.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 95757629652992.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7608550817792.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 730837205123.7609,
+      "step": 52,
+      "teaching_rms": 56834207915.74035,
+      "traffic_rms": 732257241462.028
+    },
+    {
+      "batch_loss": 1.0431297895366657e+28,
+      "innovation_rms": 100253947710.13574,
+      "instruction_rms": 20196949753.126984,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.177814590285329e+23,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.4391149174909634e+27,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5830688256.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11649742274560.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 168526295859200.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5597623808.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 173407509413888.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 749110613245952.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 83069558063104.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 359239319027712.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 51376499458048.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11649742274560.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 168526295859200.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9648919805952.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1279107993806.6724,
+      "step": 53,
+      "teaching_rms": 100253947702.4989,
+      "traffic_rms": 1282025981652.1418
+    },
+    {
+      "batch_loss": 2.3914876116286937e+28,
+      "innovation_rms": 176567217492.92838,
+      "instruction_rms": 28811875108.219406,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1082767090111739e+24,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.863492214632606e+27,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5247066624.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22268983902208.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 302314929782784.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6099938304.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 305807543500800.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1324000173096960.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 146550021947392.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 634804689174528.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 47077564350464.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22268983902208.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 302314929782784.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11457151369216.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2242629095755.0547,
+      "step": 54,
+      "teaching_rms": 176567217328.19217,
+      "traffic_rms": 2246461639803.4985
+    },
+    {
+      "batch_loss": 7.605432139189193e+28,
+      "innovation_rms": 310829627793.7019,
+      "instruction_rms": 52319848004.14343,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9841807917321452e+24,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 4.867099341724065e+27,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4721754624.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 46362926252032.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 503153069916160.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6551960576.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 539864638423040.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2340570882113536.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 258929418305536.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1123793896472576.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 72574268604416.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 46362926252032.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 503153069916160.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14483177078784.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3937520992124.9077,
+      "step": 55,
+      "teaching_rms": 310829627852.5968,
+      "traffic_rms": 3944421361720.805
+    },
+    {
+      "batch_loss": 2.8427010399125598e+29,
+      "innovation_rms": 546592211630.80585,
+      "instruction_rms": 104115912248.1194,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.543606854423054e+24,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 7.010979347674779e+27,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4248926720.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 67208306753536.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 888640678920192.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6958714880.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 953703930527744.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4138393189482496.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 457944982355968.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1990155774722048.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 591060782809088.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 67208306753536.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 888640678920192.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 59054693023744.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 6919523846933.549,
+      "step": 56,
+      "teaching_rms": 546592211025.75885,
+      "traffic_rms": 6934193406454.941
+    },
+    {
+      "batch_loss": 9.348485847767946e+29,
+      "innovation_rms": 961513169068.5793,
+      "instruction_rms": 190078842253.97253,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.325263450891755e+24,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 6.034806103552129e+28,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3823339264.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 109699965911040.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1542020631887872.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7330688512.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1686112355483648.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7324083712688128.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 807062137208832.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3491171414310912.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 532012465127424.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 109699965911040.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1542020631887872.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 112255941214208.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 12168444805305.852,
+      "step": 57,
+      "teaching_rms": 961513169934.4867,
+      "traffic_rms": 12196017919582.797
+    },
+    {
+      "batch_loss": 2.7159174379899785e+30,
+      "innovation_rms": 1690621009380.8713,
+      "instruction_rms": 313534251355.0122,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1265537286278865e+25,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.7079383009069308e+29,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3440272128.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 197442817490944.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2843641743671296.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7674715648.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2979045419515904.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2929331177193472e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1424153575948288.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6170914521612288.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 478293866315776.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 197442817490944.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2843641743671296.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 160085326168064.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 21399779349622.55,
+      "step": 58,
+      "teaching_rms": 1690621008917.6174,
+      "traffic_rms": 21445443085127.484
+    },
+    {
+      "batch_loss": 8.52247428556369e+30,
+      "innovation_rms": 2970566259885.1143,
+      "instruction_rms": 544759606758.4581,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.004693225265659e+25,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 7.222456812135463e+29,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3095477248.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 395365396250624.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5328727806312448.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7984263168.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5262898305695744.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2838527251185664e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2513916030615552.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0897624815108096e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 978527566430208.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 395365396250624.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5328727806312448.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 202967252729856.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 37631006264012.195,
+      "step": 59,
+      "teaching_rms": 2970566176039.0728,
+      "traffic_rms": 37713366598632.64
+    },
+    {
+      "batch_loss": 2.3413394494613402e+31,
+      "innovation_rms": 5230416958847.151,
+      "instruction_rms": 888261408253.7384,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.5612892115392094e+25,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.379381035634288e+30,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2785131008.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 791768228429824.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8833743242395648.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8262776320.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9299680085147648.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.036782208948634e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4440718621802496.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9268024301256704e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1552358215516160.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 791768228429824.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8833743242395648.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 250502726025216.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 66199022123638.266,
+      "step": 60,
+      "teaching_rms": 5230416958547.659,
+      "traffic_rms": 66327352946130.055
+    },
+    {
+      "batch_loss": 6.503336173942475e+31,
+      "innovation_rms": 9219152848496.922,
+      "instruction_rms": 1528860527138.3665,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.32307485999956e+25,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.2167256729345215e+30,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2505791488.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1192044819120128.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.543866796736512e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8513355264.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6478851064397824e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.179170549753446e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7741034842816512.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.300316006265651e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4094891084611584e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1192044819120128.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.543866796736512e+16,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2404507912765440.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 116444718517472.53,
+      "step": 61,
+      "teaching_rms": 9219152870246.0,
+      "traffic_rms": 116645360476532.25
+    },
+    {
+      "batch_loss": 2.2974526858584605e+32,
+      "innovation_rms": 16220192144376.395,
+      "instruction_rms": 2923296378750.67,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1216384092020059e+26,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 3.7686937443787767e+31,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2254360832.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2209484587925504.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.916575925501952e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8738791424.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.914425005880115e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2665399853396787e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3635241264545792e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.894206099606733e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1688807455719424e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2209484587925504.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.916575925501952e+16,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4573388550963200.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 204975231950662.84,
+      "step": 62,
+      "teaching_rms": 16220192181462.807,
+      "traffic_rms": 205377585030774.4
+    },
+    {
+      "batch_loss": 6.967999288825816e+32,
+      "innovation_rms": 28563121813911.613,
+      "instruction_rms": 5137832953927.019,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9879938029074712e+26,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.723388167428323e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2028050816.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4131801219989504.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.602425536361267e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8941596672.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.151658223232614e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2372331744028262e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.40951960272896e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0459953903750349e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.952313891618816e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4131801219989504.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.602425536361267e+16,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6525702442582016.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 360760987326826.5,
+      "step": 63,
+      "teaching_rms": 28563121728797.934,
+      "traffic_rms": 361452069837501.06
+    },
+    {
+      "batch_loss": 2.392130959033026e+33,
+      "innovation_rms": 49894981592420.02,
+      "instruction_rms": 9427128438646.328,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.522569219027332e+26,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 0.0,
+          "nonfinite_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1824351488.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6364986376978432.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.823271575257088e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9124031488.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.103665079202611e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.952006941869343e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.261185660230042e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8516662292997734e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7570172189540352e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6364986376978432.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.823271575257088e+16,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8282719876284416.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9421697024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10354210436344147,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 631025438736166.4,
+      "step": 64,
+      "teaching_rms": 49894981311117.86,
+      "traffic_rms": 632386391882606.2
+    },
+    {
+      "batch_loss": 8.303959420421813e+33,
+      "innovation_rms": 87337939462290.6,
+      "instruction_rms": 17647430197379.742,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.237914820347238e+26,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 0.0,
+          "nonfinite_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1641003904.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1552852140883968e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.577964627998802e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9288131584.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.607795065655132e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.974285405550019e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.525347061294694e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.264161401064653e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.820519384167219e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1552852140883968e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.577964627998802e+17,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9863952888496128.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 76020409434112.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": 7.818302689677862e+27,
+      "predictor_updated": true,
+      "raw_apical_rms": 1105716951359009.1,
+      "step": 65,
+      "teaching_rms": 87337938794220.56,
+      "traffic_rms": 1108417709186258.9
+    },
+    {
+      "batch_loss": 2.042142779418334e+34,
+      "innovation_rms": 131884035089892.84,
+      "instruction_rms": 26799978249040.1,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1042741200705416e+27,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 0.0,
+          "nonfinite_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1475974656.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1152172766920704e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.8107697068127027e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9435728896.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7554357533527245e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.147640756417069e+18,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2906705718974874e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.381359001277563e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.51363812145234e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1152172766920704e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.8107697068127027e+17,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1286963600490496e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 76020409434112.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1937949817799875.0,
+      "step": 66,
+      "teaching_rms": 131884038246111.05,
+      "traffic_rms": 1941750730643242.8
+    },
+    {
+      "batch_loss": 6.258915899903426e+34,
+      "innovation_rms": 226466291266054.62,
+      "instruction_rms": 48493876383630.95,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9231185224616184e+27,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 0.0,
+          "nonfinite_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1327433600.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.243508433833165e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.959200137403433e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9568472064.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.6989450479140864e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.94350912276267e+18,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2024320851954893e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.117614961181327e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4580151249534976e+17,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.243508433833165e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.959200137403433e+17,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4159937018003456e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 76020409434112.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3251712622963390.5,
+      "step": 67,
+      "teaching_rms": 226466284523035.6,
+      "traffic_rms": 3258621291669245.5
+    },
+    {
+      "batch_loss": 1.913553639506425e+35,
+      "innovation_rms": 384802294754270.94,
+      "instruction_rms": 87140158410673.97,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.314895552884581e+27,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 0.0,
+          "nonfinite_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1193733376.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.327312991223808e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.517415974685901e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9687845888.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.020228826820772e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.3212837789066854e+18,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.761082588861563e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.558650641105027e+18,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8435391203835904e+17,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.327312991223808e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.517415974685901e+17,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.728193076441907e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 76020409434112.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": Infinity,
+      "step": 68,
+      "teaching_rms": 384802294620885.9,
+      "traffic_rms": Infinity
+    },
+    {
+      "batch_loss": 6.060198595670836e+35,
+      "innovation_rms": 656291305267410.9,
+      "instruction_rms": 157200467366231.16,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.685550385744514e+27,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 0.0,
+          "nonfinite_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1073391296.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.864699783308902e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3529375131062764e+18,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9795184640.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3712299506186322e+18,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.692070541926597e+18,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.430752827475231e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6696703073331446e+18,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1073391296.0,
+          "nonfinite_indices": [
+            0
+          ]
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.864699783308902e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3529375131062764e+18,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 9795184640.0,
+          "nonfinite_indices": [
+            0
+          ]
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 76020409434112.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1044396162033081,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": Infinity,
+      "step": 69,
+      "teaching_rms": 656291305228304.9,
+      "traffic_rms": Infinity
+    }
+  ],
+  "trajectory_metrics_finite": false,
+  "validation_evaluations": 0
+}
diff --git a/results/kp_innovation_diagnosis/raw.json b/results/kp_innovation_diagnosis/raw.json
new file mode 100644
index 0000000..fd05c11
--- /dev/null
+++ b/results/kp_innovation_diagnosis/raw.json
@@ -0,0 +1,2596 @@
+{
+  "first_any_nonfinite_groups": {
+    "bn_running_var": [
+      2
+    ]
+  },
+  "first_any_nonfinite_step": 12,
+  "first_nonfinite_by_group": {
+    "bn_running_var": {
+      "indices": [
+        2
+      ],
+      "step": 12
+    },
+    "forward_momentum": {
+      "indices": [
+        0
+      ],
+      "step": 23
+    },
+    "forward_weight": {
+      "indices": [
+        0
+      ],
+      "step": 23
+    }
+  },
+  "first_nonfinite_by_metric": {
+    "raw_apical_rms": 23,
+    "teaching_rms": 23,
+    "traffic_rms": 23
+  },
+  "first_nonfinite_metrics": [
+    "teaching_rms"
+  ],
+  "first_training_failure_groups": {
+    "forward_momentum": [
+      0
+    ],
+    "forward_weight": [
+      0
+    ]
+  },
+  "first_training_failure_step": 23,
+  "initial_parameter_state": {
+    "bn_running_mean": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "bn_running_var": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 1.0,
+      "nonfinite_indices": []
+    },
+    "feedback_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "feedback_readout": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.4011031687259674,
+      "nonfinite_indices": []
+    },
+    "feedback_readout_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "feedback_weight": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.8331910371780396,
+      "nonfinite_indices": []
+    },
+    "forward_bn_bias": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_bn_bias_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_bn_scale": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 1.0,
+      "nonfinite_indices": []
+    },
+    "forward_bn_scale_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_readout": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.3956821858882904,
+      "nonfinite_indices": []
+    },
+    "forward_readout_momentum": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.0,
+      "nonfinite_indices": []
+    },
+    "forward_weight": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 1.1162850856781006,
+      "nonfinite_indices": []
+    },
+    "predictor_bias": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.027761053293943405,
+      "nonfinite_indices": []
+    },
+    "predictor_slope": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.09883774071931839,
+      "nonfinite_indices": []
+    }
+  },
+  "max_steps": 352,
+  "predictor_warmup": {
+    "first_mse": 0.00028227133620732826,
+    "last_mse": 5.667926742440297e-06,
+    "post_warmup_traffic_residual_rms_ratio": 0.13994200329263365,
+    "steps": 20,
+    "task_loader_state_restored": true
+  },
+  "protocol": "kp_mixed_traffic_nonfinite_diagnosis_v1",
+  "provenance": {
+    "git_commit": "e9962317b3d976fe25e28a5693f2757e3758480a",
+    "git_tracked_dirty": false
+  },
+  "rule": "raw",
+  "scope": "training_only_no_validation_or_test_evaluation",
+  "split": {
+    "dataset": "cifar10",
+    "input_layout": "NCHW",
+    "input_shape": [
+      3,
+      32,
+      32
+    ],
+    "loader_seed": 0,
+    "normalization_mean": [
+      0.49140000343322754,
+      0.4821999967098236,
+      0.4465000033378601
+    ],
+    "normalization_std": [
+      0.24699999392032623,
+      0.2434999942779541,
+      0.26159998774528503
+    ],
+    "split_from_training_only": true,
+    "split_seed": 2027,
+    "test_examples": 10000,
+    "train_examples": 45000,
+    "training_augmentation": "random_crop_32_padding_4_zero_then_horizontal_flip_p0.5",
+    "validation_class_counts": {
+      "0": 500,
+      "1": 500,
+      "2": 500,
+      "3": 500,
+      "4": 500,
+      "5": 500,
+      "6": 500,
+      "7": 500,
+      "8": 500,
+      "9": 500
+    },
+    "validation_examples": 5000,
+    "validation_index_sha256": "8328b206a97c420e49e54e3eca4abe3274c4756b084355784ea3fb8059e4515b"
+  },
+  "test_evaluations": 0,
+  "traffic_calibration": {
+    "instruction_rms": [
+      0.008471464738249779,
+      0.006617916747927666,
+      0.003988116513937712,
+      0.0043865954503417015,
+      0.002713605063036084,
+      0.003158062230795622,
+      0.002004207344725728,
+      0.003757704282179475,
+      0.0027085733599960804,
+      0.0030200635083019733,
+      0.002056022873148322,
+      0.002389939734712243,
+      0.0016842411132529378,
+      0.0034010226372629404,
+      0.002549725119024515,
+      0.0030542907770723104,
+      0.0020918569061905146,
+      0.002605808200314641,
+      0.0018916188273578882
+    ],
+    "realized_traffic_instruction_rms_ratio": [
+      4.0,
+      4.0,
+      4.0,
+      4.0,
+      4.0,
+      4.0,
+      3.999999523162842,
+      4.0,
+      4.000000476837158,
+      3.999999761581421,
+      4.0,
+      4.0,
+      3.999999761581421,
+      4.0,
+      4.000000476837158,
+      4.000000476837158,
+      4.0,
+      4.0,
+      4.0
+    ],
+    "target_ratio": 4.0,
+    "traffic_gain": [
+      0.0443243607878685,
+      0.035817068070173264,
+      0.013778294436633587,
+      0.022762222215533257,
+      0.007214079145342112,
+      0.017218582332134247,
+      0.004616298712790012,
+      0.02042512409389019,
+      0.007327491883188486,
+      0.01611751690506935,
+      0.004703814629465342,
+      0.01256294921040535,
+      0.003324519144371152,
+      0.01805141195654869,
+      0.006063441745936871,
+      0.016296451911330223,
+      0.0044285994954407215,
+      0.013874167576432228,
+      0.0036007820162922144
+    ],
+    "unscaled_traffic_rms": [
+      0.7644973993301392,
+      0.7390796542167664,
+      1.157796859741211,
+      0.7708553671836853,
+      1.504616141319275,
+      0.7336404919624329,
+      1.736635684967041,
+      0.7358984351158142,
+      1.478581428527832,
+      0.7495108246803284,
+      1.7483876943588257,
+      0.7609485983848572,
+      2.0264477729797363,
+      0.7536302804946899,
+      1.6820316314697266,
+      0.7496824264526367,
+      1.8894071578979492,
+      0.7512690424919128,
+      2.10134220123291
+    ]
+  },
+  "trajectory": [
+    {
+      "batch_loss": 2.9134020805358887,
+      "innovation_rms": 0.004771175180211988,
+      "instruction_rms": 0.004215397890235521,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.40285658836364746,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9875131845474243,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.4835381507873535,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4058851897716522,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6322960257530212,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8855142593383789,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.92843496799469,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.28434944152832,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.5336530208587646,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25.336528778076172,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.4835429191589355,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3794170618057251,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.632293164730072,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1108489036560059,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": 4.976208320227749e-06,
+      "predictor_updated": true,
+      "raw_apical_rms": 0.017318262248096506,
+      "step": 1,
+      "teaching_rms": 0.017318262248096506,
+      "traffic_rms": 0.01680093259933294
+    },
+    {
+      "batch_loss": 4.37880802154541,
+      "innovation_rms": 0.006642734878556084,
+      "instruction_rms": 0.000853714235416182,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.809155464172363,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 721.392822265625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 40.836830139160156,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5069015026092529,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.928332328796387,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.310676097869873,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15.794450759887695,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 138.66015625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18.23194122314453,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 147.47482299804688,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 40.836849212646484,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.558488130569458,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.9283623695373535,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.150968551635742,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.07167517478648427,
+      "step": 2,
+      "teaching_rms": 0.07167517478648427,
+      "traffic_rms": 0.07168706967924272
+    },
+    {
+      "batch_loss": 62.39194869995117,
+      "innovation_rms": 0.04320631538853019,
+      "instruction_rms": 0.0010472341542311814,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 356.8426208496094,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1791921.125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 213.0579071044922,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5058348178863525,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25.9843692779541,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.809722900390625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 110.76314544677734,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 949.6869506835938,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 106.9458999633789,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 887.1395874023438,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 213.05795288085938,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5706584453582764,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25.98439598083496,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.629249572753906,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.4096882489245497,
+      "step": 3,
+      "teaching_rms": 0.4096882489245497,
+      "traffic_rms": 0.4097436226513377
+    },
+    {
+      "batch_loss": 2009.0380859375,
+      "innovation_rms": 0.2674370110966999,
+      "instruction_rms": 0.004233299259793086,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14210.78125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1380064384.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5443.09619140625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22.39139747619629,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 228.98875427246094,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 544.3182373046875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 755.9462890625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6451.8310546875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 625.8128051757812,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5188.6689453125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5443.09619140625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22.356828689575195,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 228.98876953125,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 544.3165893554688,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2.461484879487376,
+      "step": 4,
+      "teaching_rms": 2.461484879487376,
+      "traffic_rms": 2.4617708411099555
+    },
+    {
+      "batch_loss": 112892.0234375,
+      "innovation_rms": 1.6685627603860504,
+      "instruction_rms": 0.06684887920874732,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 248347.546875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 623219245056.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5791.7373046875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 138.59193420410156,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1434.5382080078125,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1034.2034912109375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5113.11767578125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 43571.7109375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3740.436767578125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31441.947265625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5791.7373046875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 138.73252868652344,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1434.5382080078125,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1034.201904296875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 15.27535965785013,
+      "step": 5,
+      "teaching_rms": 15.27535965785013,
+      "traffic_rms": 15.278027105225625
+    },
+    {
+      "batch_loss": 2455932.0,
+      "innovation_rms": 10.530254946194479,
+      "instruction_rms": 0.3377461375021928,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9263209.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 785249072054272.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 73129.2734375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 919.689697265625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9476.0703125,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7309.23193359375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34009.1953125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 288960.78125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23957.794921875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 202173.578125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 106296.703125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 919.8031005859375,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9476.0703125,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10629.2666015625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 96.2812310170606,
+      "step": 6,
+      "teaching_rms": 96.2812310170606,
+      "traffic_rms": 96.29033498247702
+    },
+    {
+      "batch_loss": 219007568.0,
+      "innovation_rms": 67.6596506605722,
+      "instruction_rms": 12.255536666155361,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 179776320.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5904562669879296e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 862161.5625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6503.6923828125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 66552.15625,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86710.046875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 223160.828125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1891516.375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 149912.125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1259543.25,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 862161.5625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6503.53173828125,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 66552.15625,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86709.921875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 609.3589883430672,
+      "step": 7,
+      "teaching_rms": 609.3589883430672,
+      "traffic_rms": 609.3716608713507
+    },
+    {
+      "batch_loss": 6313313792.0,
+      "innovation_rms": 635.8658596420023,
+      "instruction_rms": 474.3939766601709,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 55527071744.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4739118525779865e+22,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2244147.25,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 40447.48828125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 415675.28125,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 225028.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1455079.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12319188.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 947911.875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7979997.5,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5845979.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 40447.62109375,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 415675.28125,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 583864.8125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3903.5996254501265,
+      "step": 8,
+      "teaching_rms": 3903.5996254501265,
+      "traffic_rms": 3875.288730373385
+    },
+    {
+      "batch_loss": 300282970112.0,
+      "innovation_rms": 3441.5039919434207,
+      "instruction_rms": 2131.803786759769,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 680749367296.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.705536237315299e+24,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25990798.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 203435.75,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2105182.25,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2594033.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9466755.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 80116752.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5939632.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 49917208.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26899548.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 203435.640625,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2105182.25,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2688694.5,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 24820.6189692148,
+      "step": 9,
+      "teaching_rms": 24820.6189692148,
+      "traffic_rms": 24733.48540903183
+    },
+    {
+      "batch_loss": 13719861985280.0,
+      "innovation_rms": 23182.18261027393,
+      "instruction_rms": 15498.95686617451,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 40207386673152.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.55772600614539e+27,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 220189120.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1721460.875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17673788.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22032622.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 62041124.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 525743680.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36923960.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 310883904.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 281836384.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1721460.75,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17673788.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28155576.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 158582.15373303735,
+      "step": 10,
+      "teaching_rms": 158582.15373303735,
+      "traffic_rms": 157856.055207326
+    },
+    {
+      "batch_loss": 435026600656896.0,
+      "innovation_rms": 114996.38814138998,
+      "instruction_rms": 32270.95626218856,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 713759609847808.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0750616147244287e+30,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 246987664.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10917683.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 112190920.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41869920.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 406617184.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3445760512.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 235566416.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1986424576.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 327553152.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10917683.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 112190920.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 53509244.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1010584.7613561256,
+      "step": 11,
+      "teaching_rms": 1010584.7613561256,
+      "traffic_rms": 1010232.4381931124
+    },
+    {
+      "batch_loss": 2.537886511529984e+16,
+      "innovation_rms": 825126.9735699444,
+      "instruction_rms": 421032.7663593107,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9343225684819968.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 4.72694492673369e+31,
+          "nonfinite_indices": [
+            2
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2473533440.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 69656936.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 720843200.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 247591504.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2665948160.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22608721920.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1496079616.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12605131776.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6821169152.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 69656936.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 720843200.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 682591040.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 6510937.902281908,
+      "step": 12,
+      "teaching_rms": 6510937.902281908,
+      "traffic_rms": 6498775.339094464
+    },
+    {
+      "batch_loss": 1.2513632545427948e+18,
+      "innovation_rms": 5248991.747072194,
+      "instruction_rms": 2637024.464266133,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.133763382406021e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.0061653946146105e+32,
+          "nonfinite_indices": [
+            1,
+            2
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 46665199616.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 472121920.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4865332224.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4661702144.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17467561984.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 148016136192.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9630923776.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 81348444160.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 49905045504.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 472121920.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4865332224.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4989012480.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 41654088.66667222,
+      "step": 13,
+      "teaching_rms": 41654088.66667222,
+      "traffic_rms": 41581860.81261145
+    },
+    {
+      "batch_loss": 3.5509394323072352e+19,
+      "innovation_rms": 42108402.89419961,
+      "instruction_rms": 30342980.296724346,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1136874112905904e+19,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.7558590928736146e+31,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            7,
+            9,
+            11,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41998213120.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2984294400.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30612709376.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8861523968.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 114593660928.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 971260952576.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 60705337344.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 510744133632.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 73283395584.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2984294400.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30612709376.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9468108800.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 268952578.3895035,
+      "step": 14,
+      "teaching_rms": 268952578.3895035,
+      "traffic_rms": 267285264.4603581
+    },
+    {
+      "batch_loss": 1.6333594611427322e+21,
+      "innovation_rms": 800222405.8456465,
+      "instruction_rms": 777877793.1214755,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3903862854120479e+20,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.1975129795930737e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            9,
+            11,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 156586229760.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20057102336.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 206726234112.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15606534144.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 750709440512.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6361157468160.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 390359187456.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3296538394624.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1652640448512.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20057102336.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 206726234112.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 164999217152.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1886884668.1288352,
+      "step": 15,
+      "teaching_rms": 1886884668.1288352,
+      "traffic_rms": 1719438087.4833093
+    },
+    {
+      "batch_loss": 8.245772963581686e+22,
+      "innovation_rms": 2108910074.0224924,
+      "instruction_rms": 1728075433.7621841,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2572378198256764e+21,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 3.104902480817458e+33,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            15,
+            17
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2421397913600.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 136395907072.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1405160914944.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 242128666624.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4917122564096.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41664131366912.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2516013023232.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21256537636864.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3639316578304.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 136395907072.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1405160914944.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 368572334080.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.0261825043708086,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1002056896686554,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 11204605422.40656,
+      "step": 16,
+      "teaching_rms": 11204605422.40656,
+      "traffic_rms": 11073351064.429298
+    },
+    {
+      "batch_loss": 2.1036254772220476e+24,
+      "innovation_rms": 8866334746.490744,
+      "instruction_rms": 4222872690.0720706,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.703865584698561e+21,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.6989238139892423e+24,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2179233873920.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 930190196736.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9565184720896.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 460052070400.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 32078815035392.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 271616919470080.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16319751651328.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 138161489444864.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 49354521968640.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 930190196736.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9565184720896.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4919401119744.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6824764928.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": 6.095242611435654e+19,
+      "predictor_updated": true,
+      "raw_apical_rms": 71465235459.55798,
+      "step": 17,
+      "teaching_rms": 71465235459.55798,
+      "traffic_rms": 71351454965.3111
+    },
+    {
+      "batch_loss": 1.5382899447402633e+26,
+      "innovation_rms": 71572086254.81004,
+      "instruction_rms": 56051603285.53418,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.464551528439866e+22,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 5.262129282512831e+26,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1961264414720.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5703475920896.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 59140839833600.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 656178544640.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 209302262382592.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1772234368614400.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 106315347132416.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 899955971588096.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 80104285798400.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5703475920896.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 59140839833600.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9405515956224.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6824764928.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 462967654001.47754,
+      "step": 18,
+      "teaching_rms": 462967654001.47754,
+      "traffic_rms": 459687978147.4927
+    },
+    {
+      "batch_loss": 6.846591999518655e+27,
+      "innovation_rms": 290501401397.3123,
+      "instruction_rms": 8111214202.742734,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.012106781543683e+23,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.3933527397858936e+27,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1765072306176.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39067479703552.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 400467045646336.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 832685801472.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1371602100944896.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1622998620504064e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 685208479203328.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5788931404595200.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 923942827065344.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39067479703552.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 400467045646336.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 91756725010432.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6824764928.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2962792032789.946,
+      "step": 19,
+      "teaching_rms": 2962792032789.946,
+      "traffic_rms": 2963650282562.4404
+    },
+    {
+      "batch_loss": 1.3576159507461967e+29,
+      "innovation_rms": 1877553196044.7314,
+      "instruction_rms": 34424872130.591675,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.410367779057763e+24,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 2.1133126471674116e+29,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1588481753088.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 245241693077504.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2514232247582720.0,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 991533989888.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8981958805684224.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.610356489990963e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4389450402496512.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.704241789075456e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4248136318910464e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 245241693077504.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2514232247582720.0,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1425324491407360.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6824764928.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 19091212825035.754,
+      "step": 20,
+      "teaching_rms": 19091212825035.754,
+      "traffic_rms": 19093848356143.99
+    },
+    {
+      "batch_loss": 1.0967124787899256e+31,
+      "innovation_rms": 12096769567901.818,
+      "instruction_rms": 307990333308.78033,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.01565603767337e+25,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.5221337164974775e+31,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1429534408704.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1570421103132672.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6285735376125952e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1134487404544.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.876230372982784e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.978034427989852e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.8376574046765056e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.398712262421381e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.084097286668288e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1570421103132672.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6285735376125952e+16,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4068768548388864.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6824764928.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 123100819405555.34,
+      "step": 21,
+      "teaching_rms": 123100819405555.34,
+      "traffic_rms": 123133928159666.89
+    },
+    {
+      "batch_loss": 4.232289909690759e+32,
+      "innovation_rms": 77974794031922.83,
+      "instruction_rms": 1869629700006.385,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.8954553677044687e+26,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1.4056117640516673e+32,
+          "nonfinite_indices": [
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1286467485696.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0861270064431104e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.119012840884142e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1263134179328.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.84298452281983e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.2553614597517476e+18,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8502196843144806e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5664538751273533e+18,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0989330097818829e+18,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0861270064431104e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.119012840884142e+17,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.09862042206208e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6824764928.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 793368631685639.2,
+      "step": 22,
+      "teaching_rms": 793368631685639.2,
+      "traffic_rms": 793566805355304.6
+    },
+    {
+      "batch_loss": 1.2669363088255617e+34,
+      "innovation_rms": 479773464857308.44,
+      "instruction_rms": 9649244758151.72,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0454542301054557e+27,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 0.0,
+          "nonfinite_indices": [
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+            12,
+            13,
+            14,
+            15,
+            16,
+            17,
+            18
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1157694357504.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.932769134975386e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.135186562769224e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1378903654400.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5077735423349883e+18,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.123474931998209e+19,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.203644037580456e+18,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0186219866856358e+19,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1157694357504.0,
+          "nonfinite_indices": [
+            0
+          ]
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.932769134975386e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.135186562769224e+17,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 1378903654400.0,
+          "nonfinite_indices": [
+            0
+          ]
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6824764928.0,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.10143683850765228,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": Infinity,
+      "step": 23,
+      "teaching_rms": Infinity,
+      "traffic_rms": Infinity
+    }
+  ],
+  "trajectory_metrics_finite": false,
+  "validation_evaluations": 0
+}