diagnostic: record mixed-traffic failure onset

1 files changed, 7778 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
+}