results: close stability-margin training screen

1 files changed, 33151 insertions, 0 deletions

diff --git a/results/kp_innovation_stability_grid/margin_0p003.json b/results/kp_innovation_stability_grid/margin_0p003.json
new file mode 100644
index 0000000..ac15d15
--- /dev/null
+++ b/results/kp_innovation_stability_grid/margin_0p003.json
@@ -0,0 +1,33151 @@
+{
+  "first_any_nonfinite_groups": {
+    "bn_running_var": [
+      16
+    ]
+  },
+  "first_any_nonfinite_step": 301,
+  "first_nonfinite_by_group": {
+    "bn_running_var": {
+      "indices": [
+        16
+      ],
+      "step": 301
+    }
+  },
+  "first_nonfinite_by_metric": {},
+  "first_nonfinite_metrics": [],
+  "first_training_failure_groups": {},
+  "first_training_failure_step": null,
+  "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.011661699041724205,
+      "nonfinite_indices": []
+    },
+    "predictor_slope": {
+      "all_finite": true,
+      "max_abs_over_finite_tensors": 0.11585475504398346,
+      "nonfinite_indices": []
+    }
+  },
+  "max_steps": 352,
+  "predictor_every": 0,
+  "predictor_mode": "closed_form",
+  "predictor_warmup": {
+    "closed_form_fit": {
+      "max_absolute_residual_soma_slope": 0.003337553469464183,
+      "max_applied_stability_margin": 0.0033375518396496773,
+      "max_positive_residual_soma_slope": 0.0,
+      "min_residual_soma_slope": -0.003337553469464183,
+      "mse": 7.272036461235194e-06,
+      "observations": 64,
+      "residual_traffic_rms_ratio": 0.15995903060998887,
+      "stability_margin": 0.003
+    },
+    "first_mse": 7.272036461235194e-06,
+    "last_mse": 7.272036461235194e-06,
+    "post_warmup_traffic_residual_rms_ratio": 0.15995903060998887,
+    "steps": 1,
+    "task_loader_state_restored": true
+  },
+  "protocol": "kp_mixed_traffic_nonfinite_diagnosis_v1",
+  "provenance": {
+    "git_commit": "0ae690e7b3144da2c5b3384cd7169e86529fd26f",
+    "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"
+  },
+  "stability_margin": 0.003,
+  "test_evaluations": 0,
+  "traffic_calibration": {
+    "instruction_rms": [
+      0.008471464738249779,
+      0.006617916747927666,
+      0.003988116513937712,
+      0.0043865954503417015,
+      0.002713605063036084,
+      0.003158062230795622,
+      0.002004207344725728,
+      0.003757704282179475,
+      0.0027085733599960804,
+      0.0030200635083019733,
+      0.002056022873148322,
+      0.002389939734712243,
+      0.0016842411132529378,
+      0.0034010226372629404,
+      0.002549725119024515,
+      0.0030542907770723104,
+      0.0020918569061905146,
+      0.002605808200314641,
+      0.0018916188273578882
+    ],
+    "realized_traffic_instruction_rms_ratio": [
+      4.0,
+      4.0,
+      4.0,
+      4.0,
+      4.0,
+      4.0,
+      3.999999523162842,
+      4.0,
+      4.000000476837158,
+      3.999999761581421,
+      4.0,
+      4.0,
+      3.999999761581421,
+      4.0,
+      4.000000476837158,
+      4.000000476837158,
+      4.0,
+      4.0,
+      4.0
+    ],
+    "target_ratio": 4.0,
+    "traffic_gain": [
+      0.0443243607878685,
+      0.035817068070173264,
+      0.013778294436633587,
+      0.022762222215533257,
+      0.007214079145342112,
+      0.017218582332134247,
+      0.004616298712790012,
+      0.02042512409389019,
+      0.007327491883188486,
+      0.01611751690506935,
+      0.004703814629465342,
+      0.01256294921040535,
+      0.003324519144371152,
+      0.01805141195654869,
+      0.006063441745936871,
+      0.016296451911330223,
+      0.0044285994954407215,
+      0.013874167576432228,
+      0.0036007820162922144
+    ],
+    "unscaled_traffic_rms": [
+      0.7644973993301392,
+      0.7390796542167664,
+      1.157796859741211,
+      0.7708553671836853,
+      1.504616141319275,
+      0.7336404919624329,
+      1.736635684967041,
+      0.7358984351158142,
+      1.478581428527832,
+      0.7495108246803284,
+      1.7483876943588257,
+      0.7609485983848572,
+      2.0264477729797363,
+      0.7536302804946899,
+      1.6820316314697266,
+      0.7496824264526367,
+      1.8894071578979492,
+      0.7512690424919128,
+      2.10134220123291
+    ]
+  },
+  "trajectory": [
+    {
+      "batch_loss": 2.9134020805358887,
+      "innovation_rms": 0.005032587581149961,
+      "instruction_rms": 0.004215397904306522,
+      "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": 2.814837694168091,
+          "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.11819249391555786,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1819249391555786,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.007530927658081,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.8693108558654785,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.814858913421631,
+          "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.1177096366882324,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.017318262558628915,
+      "step": 1,
+      "teaching_rms": 0.005032587581149961,
+      "traffic_rms": 0.01680093259933294
+    },
+    {
+      "batch_loss": 2.383181571960449,
+      "innovation_rms": 0.009944468177359623,
+      "instruction_rms": 0.009731096474187725,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6955170631408691,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.993224859237671,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4888882637023926,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4080387353897095,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4674604535102844,
+          "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.23896048963069916,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4423701763153076,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0017153024673462,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.358752965927124,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4888834953308105,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3808126151561737,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4674549698829651,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1187843084335327,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.016668615108117832,
+      "step": 2,
+      "teaching_rms": 0.009944468177359623,
+      "traffic_rms": 0.013539687767423866
+    },
+    {
+      "batch_loss": 2.533369541168213,
+      "innovation_rms": 0.025687857543152576,
+      "instruction_rms": 0.025618338704629883,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9180524349212646,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.720404624938965,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.244967222213745,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4115122854709625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3183801472187042,
+          "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.4106180965900421,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.046145439147949,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9962931275367737,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.129619598388672,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2449605464935303,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3904096484184265,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.31837230920791626,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.130269169807434,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.027394911406576326,
+      "step": 3,
+      "teaching_rms": 0.025687857543152576,
+      "traffic_rms": 0.00968280049857625
+    },
+    {
+      "batch_loss": 2.6459836959838867,
+      "innovation_rms": 0.1106321972152315,
+      "instruction_rms": 0.11061645484900672,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0729758739471436,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.718188762664795,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.021787166595459,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4124864339828491,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.1635015904903412,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9752329587936401,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7716590762138367,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.2077956199646,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.988315999507904,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.492952346801758,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0217785835266113,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3930494785308838,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.16349166631698608,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2539069652557373,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.11082192485751999,
+      "step": 4,
+      "teaching_rms": 0.1106321972152315,
+      "traffic_rms": 0.006702016893912953
+    },
+    {
+      "batch_loss": 2.6751513481140137,
+      "innovation_rms": 0.7429880349451337,
+      "instruction_rms": 0.7429848569871106,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1745615005493164,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.818680763244629,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.787534713745117,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4101289212703705,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.18702761828899384,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.250150442123413,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.94925856590271,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28.376998901367188,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4496920108795166,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29.81790542602539,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.787456512451172,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3871249854564667,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.18698816001415253,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4381338357925415,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.7430013826722146,
+      "step": 5,
+      "teaching_rms": 0.7429880349451337,
+      "traffic_rms": 0.005618643776325004
+    },
+    {
+      "batch_loss": 2.683047294616699,
+      "innovation_rms": 0.08270723545345797,
+      "instruction_rms": 0.08267924381914223,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.950665831565857,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.09476375579834,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.508668899536133,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.40692755579948425,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.29002413153648376,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3476288318634033,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.346168518066406,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23.969100952148438,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.116243362426758,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26.66551399230957,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.508580207824707,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3920237720012665,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.2899615466594696,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5596683025360107,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.08395363441473758,
+      "step": 6,
+      "teaching_rms": 0.08270723545345797,
+      "traffic_rms": 0.014575276203898108
+    },
+    {
+      "batch_loss": 2.7092790603637695,
+      "innovation_rms": 0.010724507775494012,
+      "instruction_rms": 0.010385726737930407,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.0963454246521,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 69.8814926147461,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.357487678527832,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4044579267501831,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4839426875114441,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.3721067905426025,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.588037490844727,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22.418691635131836,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.370516777038574,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22.542736053466797,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.357389450073242,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.40037307143211365,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4841077923774719,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.5690042972564697,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.025308910364666513,
+      "step": 7,
+      "teaching_rms": 0.010724507775494012,
+      "traffic_rms": 0.023088053632045007
+    },
+    {
+      "batch_loss": 2.4888246059417725,
+      "innovation_rms": 0.005352727587323237,
+      "instruction_rms": 0.004661166044975754,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.164149284362793,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 255.35482788085938,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.321355819702148,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4056166112422943,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5972133874893188,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.296573638916016,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.74666690826416,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.83891487121582,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.163341522216797,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17.92824363708496,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.321249008178711,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4037473797798157,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5972205400466919,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.477365970611572,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.025638521187578867,
+      "step": 8,
+      "teaching_rms": 0.005352727587323237,
+      "traffic_rms": 0.025213286355138007
+    },
+    {
+      "batch_loss": 2.556396007537842,
+      "innovation_rms": 0.0037166655900639466,
+      "instruction_rms": 0.002835970222294144,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.527647018432617,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 536.90576171875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.387642860412598,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4078368842601776,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3040487766265869,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.128400802612305,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.849884033203125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.221620559692383,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.51327133178711,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15.075358390808105,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.387527465820312,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4061666429042816,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3042426109313965,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.294744968414307,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.021799951879044654,
+      "step": 9,
+      "teaching_rms": 0.0037166655900639466,
+      "traffic_rms": 0.021616941786698086
+    },
+    {
+      "batch_loss": 2.666769504547119,
+      "innovation_rms": 0.0032631036359407235,
+      "instruction_rms": 0.002419598076165226,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.101994514465332,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 863.05419921875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.529796123504639,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4218451976776123,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8099203109741211,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.874650478363037,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.904777526855469,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22.464523315429688,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.447300910949707,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15.334668159484863,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.5296735763549805,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.40610271692276,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8099014759063721,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.027317523956299,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.016490816394568994,
+      "step": 10,
+      "teaching_rms": 0.0032631036359407235,
+      "traffic_rms": 0.016315742550184573
+    },
+    {
+      "batch_loss": 2.3774166107177734,
+      "innovation_rms": 0.003254795229790769,
+      "instruction_rms": 0.0024884142594770236,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14.59959888458252,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1155.075439453125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.761796951293945,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.45294374227523804,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7180604338645935,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.544996738433838,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15.891387939453125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20.33870506286621,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.050105094909668,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.484246253967285,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.761667728424072,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4054393768310547,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7180405259132385,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.69282341003418,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.013826597158824514,
+      "step": 11,
+      "teaching_rms": 0.003254795229790769,
+      "traffic_rms": 0.013608536210805893
+    },
+    {
+      "batch_loss": 2.4111337661743164,
+      "innovation_rms": 0.0031869042785004236,
+      "instruction_rms": 0.0025210489253120524,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15.021065711975098,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1339.161376953125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.077770709991455,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4718649983406067,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.544173538684845,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.147280216217041,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17.775156021118164,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18.837677001953125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.424899101257324,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.1935453414917,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.077635288238525,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.42968493700027466,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5441611409187317,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.300587177276611,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.01165223683891572,
+      "step": 12,
+      "teaching_rms": 0.0031869042785004236,
+      "traffic_rms": 0.011374599667729949
+    },
+    {
+      "batch_loss": 2.450648546218872,
+      "innovation_rms": 0.003780906933157393,
+      "instruction_rms": 0.0033006992742589535,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14.828142166137695,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1401.019775390625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.46356201171875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.47267282009124756,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.44520533084869385,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.6884613037109375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.535444259643555,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17.602876663208008,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.647623062133789,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.047343254089355,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.463449954986572,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.45287495851516724,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.44519245624542236,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.846377849578857,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.009651415897167608,
+      "step": 13,
+      "teaching_rms": 0.003780906933157393,
+      "traffic_rms": 0.009075196499151889
+    },
+    {
+      "batch_loss": 2.4470107555389404,
+      "innovation_rms": 0.0035289725567818186,
+      "instruction_rms": 0.0029231062633938874,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14.207108497619629,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1382.914794921875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.913687229156494,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.47075238823890686,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.575528621673584,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.174945831298828,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.160446166992188,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16.250011444091797,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.775080680847168,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.156169414520264,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.913571357727051,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.47289296984672546,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.575525164604187,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.336942672729492,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.009280910750068113,
+      "step": 14,
+      "teaching_rms": 0.0035289725567818186,
+      "traffic_rms": 0.008808542472500226
+    },
+    {
+      "batch_loss": 2.486706495285034,
+      "innovation_rms": 0.005009498775737296,
+      "instruction_rms": 0.00454384550693962,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.5067720413208,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1307.7867431640625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.554566383361816,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4720124900341034,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6068675518035889,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.612659454345703,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22.64488410949707,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14.844382286071777,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.849983215332031,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.620048522949219,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.554633617401123,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.48807209730148315,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6068639755249023,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.77833080291748,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.009401125077112646,
+      "step": 15,
+      "teaching_rms": 0.005009498775737296,
+      "traffic_rms": 0.008228382187026281
+    },
+    {
+      "batch_loss": 2.389864683151245,
+      "innovation_rms": 0.010189759976338881,
+      "instruction_rms": 0.010007501429911015,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.904996871948242,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1207.4957275390625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.133140563964844,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5007818937301636,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5417394042015076,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.006515502929688,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23.984519958496094,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.662263870239258,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.910775184631348,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.996420860290527,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.133209705352783,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5011630058288574,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5417357087135315,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.175492286682129,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.01192692363443505,
+      "step": 16,
+      "teaching_rms": 0.010189759976338881,
+      "traffic_rms": 0.006489520017431994
+    },
+    {
+      "batch_loss": 2.4523613452911377,
+      "innovation_rms": 0.009482778145883587,
+      "instruction_rms": 0.009274878623483999,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.425341606140137,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1111.7847900390625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.7196762561798096,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5249875783920288,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.46890562772750854,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.360896110534668,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25.1901912689209,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.54744815826416,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.965487480163574,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.196778774261475,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.7197470664978027,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5134652853012085,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4689018428325653,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.532845497131348,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.011125065507331696,
+      "step": 17,
+      "teaching_rms": 0.009482778145883587,
+      "traffic_rms": 0.00615003864721607
+    },
+    {
+      "batch_loss": 2.6268060207366943,
+      "innovation_rms": 0.01434612529134289,
+      "instruction_rms": 0.014185628263882463,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.06570053100586,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1014.4560546875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.3475210666656494,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5426435470581055,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3953171670436859,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.679744720458984,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26.27529525756836,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.542301177978516,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.014728546142578,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.44973087310791,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.3475935459136963,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5262220501899719,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3953900933265686,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.854368209838867,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.015440351673115032,
+      "step": 18,
+      "teaching_rms": 0.01434612529134289,
+      "traffic_rms": 0.006118550390429117
+    },
+    {
+      "batch_loss": 2.4115519523620605,
+      "innovation_rms": 0.01144809614229914,
+      "instruction_rms": 0.011264935800902173,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.753857612609863,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1063.0247802734375,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.0125482082366943,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5551297068595886,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6080416440963745,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.966611862182617,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27.251888275146484,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.530037879943848,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.059045791625977,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.271183013916016,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.0126218795776367,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.535926103591919,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.607783317565918,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.14363956451416,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.012712281415732034,
+      "step": 19,
+      "teaching_rms": 0.01144809614229914,
+      "traffic_rms": 0.005898093843307749
+    },
+    {
+      "batch_loss": 2.4810826778411865,
+      "innovation_rms": 0.004195158037537646,
+      "instruction_rms": 0.003694536228830557,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.476689338684082,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1185.2945556640625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7110424041748047,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5622638463973999,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7733966112136841,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.224692344665527,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28.130823135375977,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.605904579162598,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.098931312561035,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.66251277923584,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7111170291900635,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5386669635772705,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7732793092727661,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.40388298034668,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.006836159200346915,
+      "step": 20,
+      "teaching_rms": 0.004195158037537646,
+      "traffic_rms": 0.005759391464533623
+    },
+    {
+      "batch_loss": 2.6561341285705566,
+      "innovation_rms": 0.00390517687661701,
+      "instruction_rms": 0.003363046179421683,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.251701354980469,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1314.4453125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.43966007232666,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.565222978591919,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7746868133544922,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.456862449645996,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28.921863555908203,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.697919845581055,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.134828567504883,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.017308235168457,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4397358894348145,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5438457727432251,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7745678424835205,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.6379976272583,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.006650494817575689,
+      "step": 21,
+      "teaching_rms": 0.00390517687661701,
+      "traffic_rms": 0.005746952117364737
+    },
+    {
+      "batch_loss": 2.47926926612854,
+      "innovation_rms": 0.0033035772169142897,
+      "instruction_rms": 0.0026776878386386136,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.89649772644043,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1434.8912353515625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1953916549682617,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5722928643226624,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5818239450454712,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.665711402893066,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29.633800506591797,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.854836940765381,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.167135238647461,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.315577030181885,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1954681873321533,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5493986010551453,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5815547704696655,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.848594665527344,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.006495282038977906,
+      "step": 22,
+      "teaching_rms": 0.0033035772169142897,
+      "traffic_rms": 0.005920911677052447
+    },
+    {
+      "batch_loss": 2.735363483428955,
+      "innovation_rms": 0.0032266704958424917,
+      "instruction_rms": 0.0025560871729717423,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.325183868408203,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1553.015380859375,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9755281209945679,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5773526430130005,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5241467356681824,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.853568077087402,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30.27454376220703,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.077362537384033,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.196211814880371,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.1098103523254395,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9756053686141968,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5623118281364441,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5238746404647827,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.038022994995117,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.007377701353366132,
+      "step": 23,
+      "teaching_rms": 0.0032266704958424917,
+      "traffic_rms": 0.006930519050243815
+    },
+    {
+      "batch_loss": 2.553990602493286,
+      "innovation_rms": 0.008727213111804818,
+      "instruction_rms": 0.008506640753066669,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.129941940307617,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1666.7464599609375,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.777631163597107,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6202973127365112,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6593894958496094,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.023048400878906,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30.851213455200195,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.369729518890381,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.222380638122559,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.397741794586182,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7777092456817627,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5774179100990295,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6597332954406738,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.208398818969727,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.010022916862048438,
+      "step": 24,
+      "teaching_rms": 0.008727213111804818,
+      "traffic_rms": 0.005308798462991069
+    },
+    {
+      "batch_loss": 2.545802354812622,
+      "innovation_rms": 0.0053513383751626915,
+      "instruction_rms": 0.004984574680429985,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.479226112365723,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1795.797607421875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.599506139755249,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6824257373809814,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7300738096237183,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.176535606384277,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31.370216369628906,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.732756614685059,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.245932579040527,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.757967472076416,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5995848178863525,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5920913219451904,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7304205894470215,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.361624717712402,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.006984398454455372,
+      "step": 25,
+      "teaching_rms": 0.0053513383751626915,
+      "traffic_rms": 0.004897199514744301
+    },
+    {
+      "batch_loss": 2.595956802368164,
+      "innovation_rms": 0.004965520183061101,
+      "instruction_rms": 0.004550657108521816,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.53748607635498,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1913.9364013671875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4391776323318481,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.722762942314148,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5471009016036987,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.31456184387207,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31.837318420410156,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.275420665740967,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.267128944396973,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.182170391082764,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4392569065093994,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5953822731971741,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5470007061958313,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.499414443969727,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.006820707747921053,
+      "step": 26,
+      "teaching_rms": 0.004965520183061101,
+      "traffic_rms": 0.005088053747863802
+    },
+    {
+      "batch_loss": 2.580071210861206,
+      "innovation_rms": 0.003677458884481769,
+      "instruction_rms": 0.003093431621302441,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.559106826782227,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2003.5703125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3878536224365234,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7527413368225098,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8083950281143188,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.43867301940918,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 32.25770950317383,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.209517002105713,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.286206245422363,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.976695537567139,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3880224227905273,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6172155737876892,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8081746101379395,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.623310089111328,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.005979670141626867,
+      "step": 27,
+      "teaching_rms": 0.003677458884481769,
+      "traffic_rms": 0.00512438602532275
+    },
+    {
+      "batch_loss": 2.62861967086792,
+      "innovation_rms": 0.004175261545819965,
+      "instruction_rms": 0.0036337112731510933,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14.305243492126465,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2095.245361328125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6300779581069946,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7833973169326782,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7774322032928467,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.55025863647461,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 32.63606262207031,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.374323844909668,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.303375244140625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.682098865509033,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6300151348114014,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6541928648948669,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7775590419769287,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.734700202941895,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.0063388557045451705,
+      "step": 28,
+      "teaching_rms": 0.004175261545819965,
+      "traffic_rms": 0.0052042514186117494
+    },
+    {
+      "batch_loss": 2.4183197021484375,
+      "innovation_rms": 0.00471325181862454,
+      "instruction_rms": 0.004248675970188885,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14.963194847106934,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2198.588623046875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8994343280792236,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.813633918762207,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7130439281463623,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.650569915771484,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 32.97658157348633,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.738422393798828,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.318827629089355,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.269129753112793,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8993858098983765,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6884320974349976,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7128601670265198,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.834834098815918,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.006701212373846554,
+      "step": 29,
+      "teaching_rms": 0.00471325181862454,
+      "traffic_rms": 0.00519390097019411
+    },
+    {
+      "batch_loss": 2.4105985164642334,
+      "innovation_rms": 0.004175176733556675,
+      "instruction_rms": 0.0035474943300073704,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16.9818115234375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2275.47314453125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.190159320831299,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8379935026168823,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4350247383117676,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.74073314666748,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33.28304672241211,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.214066505432129,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.332735061645508,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10.478825569152832,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1904077529907227,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7201142907142639,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.43505463004112244,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.924836158752441,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.006525183513879746,
+      "step": 30,
+      "teaching_rms": 0.004175176733556675,
+      "traffic_rms": 0.005481093553232488
+    },
+    {
+      "batch_loss": 2.6399600505828857,
+      "innovation_rms": 0.003999589406219912,
+      "instruction_rms": 0.003156792587712283,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.092191696166992,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2363.721923828125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.106741428375244,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8572355508804321,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.475011944770813,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.821763038635254,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33.55886459350586,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.710826396942139,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.345251083374023,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.854059219360352,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.106692314147949,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7513348460197449,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.47524237632751465,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.005718231201172,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.006379198959272609,
+      "step": 31,
+      "teaching_rms": 0.003999589406219912,
+      "traffic_rms": 0.005543542777430112
+    },
+    {
+      "batch_loss": 2.4447336196899414,
+      "innovation_rms": 0.004918227153974521,
+      "instruction_rms": 0.004315210293283674,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20.831573486328125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2452.54150390625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.146794319152832,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.874858021736145,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6424421668052673,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.894571304321289,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33.80710220336914,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.171438217163086,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.356515884399414,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17.076404571533203,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.146854400634766,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7794262766838074,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6426757574081421,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.078392028808594,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.006852093191662753,
+      "step": 32,
+      "teaching_rms": 0.004918227153974521,
+      "traffic_rms": 0.005318871018738088
+    },
+    {
+      "batch_loss": 2.553338050842285,
+      "innovation_rms": 0.02616920418026258,
+      "instruction_rms": 0.026039330658840575,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.813968658447266,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2522.400634765625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.4427080154418945,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8791693449020386,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.881649374961853,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.959980010986328,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.030513763427734,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.493828296661377,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.366654396057129,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17.8394775390625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.4427080154418945,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.803958535194397,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8817770481109619,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.143677711486816,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.02661582130762899,
+      "step": 33,
+      "teaching_rms": 0.02616920418026258,
+      "traffic_rms": 0.005514740707339414
+    },
+    {
+      "batch_loss": 2.6185712814331055,
+      "innovation_rms": 0.031050996179704648,
+      "instruction_rms": 0.030947408566878977,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.970197677612305,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2615.5654296875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.151602268218994,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8896123766899109,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4784935712814331,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.018728256225586,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.23158645629883,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.356487274169922,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.375779151916504,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18.73900604248047,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.151602268218994,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8187700510025024,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4786216616630554,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.202313423156738,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.03144923191132248,
+      "step": 34,
+      "teaching_rms": 0.031050996179704648,
+      "traffic_rms": 0.0056014033133125465
+    },
+    {
+      "batch_loss": 2.544434070587158,
+      "innovation_rms": 0.014617836008768538,
+      "instruction_rms": 0.01436311545717594,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.694189071655273,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2683.395751953125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.420884132385254,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9065348505973816,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4133628308773041,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.071481704711914,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.41255187988281,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.778570175170898,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14.389724731445312,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20.220232009887695,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.420854568481445,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8333011269569397,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.41359639167785645,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.254963874816895,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.015546733016270646,
+      "step": 35,
+      "teaching_rms": 0.014617836008768538,
+      "traffic_rms": 0.0059539606003954325
+    },
+    {
+      "batch_loss": 2.718416929244995,
+      "innovation_rms": 0.008368202311008719,
+      "instruction_rms": 0.007669086642150511,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.059831619262695,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2750.409912109375,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.560401916503906,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9299095273017883,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5931472182273865,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.118839263916016,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.57542037963867,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.981626510620117,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16.887874603271484,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24.981494903564453,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.560372352600098,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8356553316116333,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5932558178901672,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.302227020263672,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.010251201359474977,
+      "step": 36,
+      "teaching_rms": 0.008368202311008719,
+      "traffic_rms": 0.006813581664422419
+    },
+    {
+      "batch_loss": 2.492528200149536,
+      "innovation_rms": 0.007014625944474726,
+      "instruction_rms": 0.006064850905964013,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20.429508209228516,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2789.216064453125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.86216926574707,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9638930559158325,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7254465222358704,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.16133975982666,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.72200012207031,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.021899700164795,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.83858871459961,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29.507144927978516,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.862170219421387,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8296882510185242,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7255553603172302,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.344640731811523,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.008917383006707061,
+      "step": 37,
+      "teaching_rms": 0.007014625944474726,
+      "traffic_rms": 0.006536073098953408
+    },
+    {
+      "batch_loss": 3.0039165019989014,
+      "innovation_rms": 0.00809082984944477,
+      "instruction_rms": 0.006609173621579656,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.908029556274414,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2837.84326171875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.631509780883789,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9599717259407043,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.46405690908432007,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.199468612670898,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.85392379760742,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.9968156814575195,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23.762418746948242,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39.238304138183594,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.631510734558105,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8473935127258301,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.46402406692504883,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.382689476013184,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.010029943424533465,
+      "step": 38,
+      "teaching_rms": 0.00809082984944477,
+      "traffic_rms": 0.007556575177304682
+    },
+    {
+      "batch_loss": 2.547844409942627,
+      "innovation_rms": 0.020643433635608225,
+      "instruction_rms": 0.019932509260651852,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.49517250061035,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2888.824951171875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.095413208007812,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9309350252151489,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.026196002960205,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.233662605285645,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34.97265625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.4121904373168945,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28.760578155517578,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 49.98160171508789,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.095414161682129,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.853959858417511,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.026163101196289,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.41680908203125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.021634273637308913,
+      "step": 39,
+      "teaching_rms": 0.020643433635608225,
+      "traffic_rms": 0.008413454233111862
+    },
+    {
+      "batch_loss": 2.845684766769409,
+      "innovation_rms": 0.028718062531508076,
+      "instruction_rms": 0.027989387478145444,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.39763069152832,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2929.493896484375,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.669869422912598,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.916434109210968,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8384279608726501,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.264314651489258,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.07951354980469,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.3630852699279785,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.11180114746094,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63.512229919433594,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.66988754272461,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8569803833961487,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8381552696228027,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.447392463684082,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.029535415850416497,
+      "step": 40,
+      "teaching_rms": 0.028718062531508076,
+      "traffic_rms": 0.009444651166420567
+    },
+    {
+      "batch_loss": 2.7636501789093018,
+      "innovation_rms": 0.008757881151688062,
+      "instruction_rms": 0.006151898563154535,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23.589563369750977,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2953.781005859375,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.073614120483398,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9199950695037842,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7200359106063843,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.291778564453125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.17568588256836,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.121078491210938,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41.7818717956543,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 66.70069885253906,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.073583602905273,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8850176930427551,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7200313806533813,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.474793434143066,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.014651962461163523,
+      "step": 41,
+      "teaching_rms": 0.008757881151688062,
+      "traffic_rms": 0.013303055100489751
+    },
+    {
+      "batch_loss": 3.132093906402588,
+      "innovation_rms": 0.008763771286865655,
+      "instruction_rms": 0.005342682194028539,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33.09446716308594,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3012.553955078125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.565156936645508,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9422749280929565,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.43133485317230225,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.316373825073242,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.26224136352539,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.1399507522583,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 48.852176666259766,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 70.70303344726562,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13.56515884399414,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9266519546508789,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.43165266513824463,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.499329566955566,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.015272772715104892,
+      "step": 42,
+      "teaching_rms": 0.008763771286865655,
+      "traffic_rms": 0.014311905141427368
+    },
+    {
+      "batch_loss": 2.7953288555145264,
+      "innovation_rms": 0.013003319265499042,
+      "instruction_rms": 0.009711171322517797,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41.80187225341797,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3067.1171875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20.382736206054688,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9605444669723511,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3666638433933258,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.338386535644531,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.34014129638672,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.163256645202637,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 56.86860275268555,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 80.16426086425781,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20.38273811340332,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9286640286445618,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.36684727668762207,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.521286964416504,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.018005228004698712,
+      "step": 43,
+      "teaching_rms": 0.013003319265499042,
+      "traffic_rms": 0.015165793875102017
+    },
+    {
+      "batch_loss": 3.047285795211792,
+      "innovation_rms": 0.015248060271256409,
+      "instruction_rms": 0.009798885216410707,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 44.63763427734375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3097.8984375,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29.778675079345703,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9089886546134949,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3508516550064087,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.358074188232422,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.41025161743164,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.300484657287598,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 67.35635375976562,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 104.87750244140625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29.778675079345703,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9135449528694153,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.350823163986206,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.540923118591309,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.018461526970091856,
+      "step": 44,
+      "teaching_rms": 0.015248060271256409,
+      "traffic_rms": 0.01565445254678963
+    },
+    {
+      "batch_loss": 3.4712986946105957,
+      "innovation_rms": 0.0350536770321577,
+      "instruction_rms": 0.03369678200340746,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 61.5442008972168,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3120.3623046875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24.009557723999023,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8964021801948547,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2081328630447388,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14.450727462768555,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.473350524902344,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.67985725402832,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 77.94634246826172,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 105.89990234375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24.0095157623291,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9009580612182617,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.208330512046814,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14.450504302978516,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.03997319398827862,
+      "step": 45,
+      "teaching_rms": 0.0350536770321577,
+      "traffic_rms": 0.02151618173668997
+    },
+    {
+      "batch_loss": 4.247214317321777,
+      "innovation_rms": 0.011450093084856174,
+      "instruction_rms": 0.006076391093252892,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 87.59587097167969,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3767.36279296875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22.20289421081543,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9782421588897705,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8597445487976074,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16.34429931640625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.5301399230957,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.064506530761719,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 88.146728515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 102.00385284423828,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22.20285415649414,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9827976226806641,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8599424362182617,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16.34407615661621,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.025617565572914107,
+      "step": 46,
+      "teaching_rms": 0.011450093084856174,
+      "traffic_rms": 0.02490023012375797
+    },
+    {
+      "batch_loss": 4.390831470489502,
+      "innovation_rms": 0.016020220020732983,
+      "instruction_rms": 0.011749720086692065,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 117.75164031982422,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6572.27978515625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20.075485229492188,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0271756649017334,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.48933476209640503,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17.871002197265625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.58124923706055,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.57795238494873,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 98.10897064208984,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 99.62242889404297,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20.0754451751709,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0317306518554688,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.48933032155036926,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17.870779037475586,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.029971894319349292,
+      "step": 47,
+      "teaching_rms": 0.016020220020732983,
+      "traffic_rms": 0.027575469804105746
+    },
+    {
+      "batch_loss": 4.365326404571533,
+      "innovation_rms": 0.026489274499608324,
+      "instruction_rms": 0.023347044432245886,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 147.81744384765625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9323.5595703125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23.569843292236328,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0501924753189087,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.49318927526474,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.04045295715332,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.62724685668945,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.126519203186035,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 108.184326171875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 100.75357818603516,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23.56984519958496,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0547469854354858,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.49350282549858093,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.04022979736328,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.03780134805166639,
+      "step": 48,
+      "teaching_rms": 0.026489274499608324,
+      "traffic_rms": 0.02973449043987266
+    },
+    {
+      "batch_loss": 3.654695749282837,
+      "innovation_rms": 0.02025591886206901,
+      "instruction_rms": 0.012342594427494427,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 173.17286682128906,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11053.31640625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33.86771774291992,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0509788990020752,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6182851791381836,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.70368194580078,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.668643951416016,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.071314811706543,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 119.60890197753906,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 114.2457275390625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33.867774963378906,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0555330514907837,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6181297898292542,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.703458786010742,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.033998354434612,
+      "step": 49,
+      "teaching_rms": 0.02025591886206901,
+      "traffic_rms": 0.03168320152381434
+    },
+    {
+      "batch_loss": 3.135939598083496,
+      "innovation_rms": 0.03203510762112668,
+      "instruction_rms": 0.02318868879347348,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 185.4970703125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11062.89453125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 54.686859130859375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.054850697517395,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7116785049438477,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.480403900146484,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.705902099609375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.264183044433594,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 135.0679473876953,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 154.59039306640625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 54.68682098388672,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0594043731689453,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7119318246841431,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19.480180740356445,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.04062616923935462,
+      "step": 50,
+      "teaching_rms": 0.03203510762112668,
+      "traffic_rms": 0.03336024297286709
+    },
+    {
+      "batch_loss": 3.1825554370880127,
+      "innovation_rms": 0.026335784781854606,
+      "instruction_rms": 0.01932522605208414,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 170.15451049804688,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11015.5830078125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 48.1868896484375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0796067714691162,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8944852352142334,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.264347076416016,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.739437103271484,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.537685871124268,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 151.65248107910156,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 165.84535217285156,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 48.18695831298828,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0841600894927979,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8947386145591736,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21.20984649658203,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.035038851045862685,
+      "step": 51,
+      "teaching_rms": 0.026335784781854606,
+      "traffic_rms": 0.029229298197341172
+    },
+    {
+      "batch_loss": 2.593303918838501,
+      "innovation_rms": 0.03235772925231196,
+      "instruction_rms": 0.008389012281147057,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 172.63140869140625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10738.8583984375,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 75.77613067626953,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.103185772895813,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9540000557899475,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25.853120803833008,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.769615173339844,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11.004308700561523,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 175.52871704101562,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 238.7623748779297,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 75.77608489990234,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1077386140823364,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9542535543441772,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25.85289764404297,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.041351080479929136,
+      "step": 52,
+      "teaching_rms": 0.03235772925231196,
+      "traffic_rms": 0.04049206065724255
+    },
+    {
+      "batch_loss": 2.3610973358154297,
+      "innovation_rms": 0.030370638932891093,
+      "instruction_rms": 0.009814777668682268,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 209.51454162597656,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14353.541015625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 57.35948944091797,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1214303970336914,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8795926570892334,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31.08207130432129,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.796775817871094,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.173280715942383,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 202.1959228515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 266.6719970703125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 57.359466552734375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1259827613830566,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8798462748527527,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31.122419357299805,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.0514848293531622,
+      "step": 53,
+      "teaching_rms": 0.030370638932891093,
+      "traffic_rms": 0.050541691262624484
+    },
+    {
+      "batch_loss": 3.0322768688201904,
+      "innovation_rms": 0.032739887412616596,
+      "instruction_rms": 0.00944852166278402,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 263.0397644042969,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22168.59375,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63.442787170410156,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.130142092704773,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7038282752037048,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.697998046875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.82122039794922,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12.137044906616211,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 231.2354736328125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 290.39544677734375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63.44276428222656,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1346940994262695,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7036235332489014,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.643516540527344,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.05831377206481453,
+      "step": 54,
+      "teaching_rms": 0.032739887412616596,
+      "traffic_rms": 0.05754843517583339
+    },
+    {
+      "batch_loss": 2.339446783065796,
+      "innovation_rms": 0.03778258873897507,
+      "instruction_rms": 0.0034282084950123435,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 322.3948974609375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30872.703125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 76.12713623046875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1364855766296387,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6796102523803711,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 43.31071090698242,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.84321975708008,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14.0341215133667,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 264.1524353027344,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 329.16961669921875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 76.12718963623047,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.141037106513977,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6794053912162781,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 43.2562370300293,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.06512990184244395,
+      "step": 55,
+      "teaching_rms": 0.03778258873897507,
+      "traffic_rms": 0.0650398040349616
+    },
+    {
+      "batch_loss": 2.439615249633789,
+      "innovation_rms": 0.051558904542795755,
+      "instruction_rms": 0.007403607238409676,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 376.95648193359375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36019.21875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 119.00940704345703,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1473190784454346,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.46987974643707275,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 52.944664001464844,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.86302185058594,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23.991199493408203,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 306.5439758300781,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 423.9153137207031,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 119.00930786132812,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1518701314926147,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.46967488527297974,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 52.890193939208984,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.0736180653057624,
+      "step": 56,
+      "teaching_rms": 0.051558904542795755,
+      "traffic_rms": 0.07324420511832026
+    },
+    {
+      "batch_loss": 2.7178752422332764,
+      "innovation_rms": 0.05909676037910553,
+      "instruction_rms": 0.015153413590911997,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 401.0545959472656,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38416.5390625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 107.94449615478516,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1616315841674805,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.40635520219802856,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 51.540470123291016,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.88084411621094,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33.15848159790039,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 359.2872009277344,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 527.43212890625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 107.944580078125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.166182279586792,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.40637457370758057,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 51.48600387573242,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.07710771881407207,
+      "step": 57,
+      "teaching_rms": 0.05909676037910553,
+      "traffic_rms": 0.07560368319165729
+    },
+    {
+      "batch_loss": 2.6427817344665527,
+      "innovation_rms": 0.08083165945815908,
+      "instruction_rms": 0.0261889891640484,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 469.53363037109375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 45240.75390625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 113.75796508789062,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.175717830657959,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4581460654735565,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58.310428619384766,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.896881103515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 51.7728271484375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 429.1938781738281,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 699.0668334960938,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 113.75806427001953,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1802680492401123,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4581654369831085,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58.33393859863281,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.10373912468528985,
+      "step": 58,
+      "teaching_rms": 0.08083165945815908,
+      "traffic_rms": 0.10037922764687267
+    },
+    {
+      "batch_loss": 2.6655642986297607,
+      "innovation_rms": 0.09066593570016225,
+      "instruction_rms": 0.023250037772072504,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 550.097900390625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 55907.66796875,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 156.7481689453125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1841485500335693,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.49948132038116455,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63.08900451660156,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35.91131591796875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 68.36222839355469,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 516.2156372070312,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 870.2175903320312,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 156.74827575683594,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1886982917785645,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.49950072169303894,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63.087730407714844,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.12378176318859907,
+      "step": 59,
+      "teaching_rms": 0.09066593570016225,
+      "traffic_rms": 0.12157914504757174
+    },
+    {
+      "batch_loss": 2.8395540714263916,
+      "innovation_rms": 0.11448457102973987,
+      "instruction_rms": 0.05977616737698357,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 585.0537109375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 62926.8984375,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 184.2338104248047,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1831183433532715,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6064733266830444,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 71.85709381103516,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38.465702056884766,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 82.981689453125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 619.9638671875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1037.4820556640625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 184.23374938964844,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.187667727470398,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6064602136611938,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 71.88059997558594,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.142221618890736,
+      "step": 60,
+      "teaching_rms": 0.11448457102973987,
+      "traffic_rms": 0.12905092468485937
+    },
+    {
+      "batch_loss": 2.537301540374756,
+      "innovation_rms": 0.11368813010972298,
+      "instruction_rms": 0.029219442842263806,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 769.910400390625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 102194.53125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 207.13821411132812,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.175146222114563,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7259321808815002,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 82.11797332763672,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 47.18775939941406,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 87.2205581665039,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 737.0541381835938,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1170.9029541015625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 207.13821411132812,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1796951293945312,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7259190678596497,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 82.08907318115234,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.18031512077817372,
+      "step": 61,
+      "teaching_rms": 0.11368813010972298,
+      "traffic_rms": 0.17793233146304502
+    },
+    {
+      "batch_loss": 2.640613317489624,
+      "innovation_rms": 0.14194659224143316,
+      "instruction_rms": 0.04476266441054064,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1002.3278198242188,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 149818.625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 351.614501953125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1643964052200317,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.520012378692627,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 104.76869201660156,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 57.15672302246094,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 99.68963623046875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 874.3414916992188,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1372.87353515625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 351.614501953125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1689448356628418,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5199437141418457,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 104.76741790771484,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.197273150795289,
+      "step": 62,
+      "teaching_rms": 0.14194659224143316,
+      "traffic_rms": 0.1921279892945194
+    },
+    {
+      "batch_loss": 2.8962724208831787,
+      "innovation_rms": 0.19458345233438518,
+      "instruction_rms": 0.020553067624707027,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1124.315185546875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 208728.265625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 254.76300048828125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.175925612449646,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5247342586517334,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 124.71875762939453,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 72.92265319824219,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 157.65928649902344,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1058.423828125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1840.8232421875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 254.7630615234375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1804735660552979,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5246873497962952,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 124.58049774169922,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.24746602430336212,
+      "step": 63,
+      "teaching_rms": 0.19458345233438518,
+      "traffic_rms": 0.2466111701561511
+    },
+    {
+      "batch_loss": 2.5692477226257324,
+      "innovation_rms": 0.1730873242763517,
+      "instruction_rms": 0.04489582304734378,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1193.8414306640625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 294824.09375,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 504.21295166015625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.184571385383606,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.42213016748428345,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 150.18853759765625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 88.73253631591797,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 158.0988311767578,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1255.2955322265625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1968.717041015625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 504.2130432128906,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1891188621520996,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.422328919172287,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 150.1180419921875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.29892834655506717,
+      "step": 64,
+      "teaching_rms": 0.1730873242763517,
+      "traffic_rms": 0.2955377272368427
+    },
+    {
+      "batch_loss": 2.780534505844116,
+      "innovation_rms": 0.23683305415728928,
+      "instruction_rms": 0.08057459512777704,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1396.0635986328125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 440109.5,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 416.70928955078125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1828025579452515,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5946528315544128,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 191.85946655273438,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 106.79360961914062,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 180.6107635498047,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1481.216064453125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2259.204833984375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 416.7093811035156,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.187349557876587,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5947977304458618,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 191.7889862060547,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.3600993399183776,
+      "step": 65,
+      "teaching_rms": 0.23683305415728928,
+      "traffic_rms": 0.35096911215355187
+    },
+    {
+      "batch_loss": 2.564851999282837,
+      "innovation_rms": 0.31164300753587487,
+      "instruction_rms": 0.011725671278669615,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1563.8875732421875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 711740.25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 644.6356201171875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1742421388626099,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.36937159299850464,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 216.79757690429688,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 131.64797973632812,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 248.54373168945312,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1773.3336181640625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2921.17529296875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 644.6356201171875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.178788661956787,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3695164918899536,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 216.7270965576172,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.43658183533522293,
+      "step": 66,
+      "teaching_rms": 0.31164300753587487,
+      "traffic_rms": 0.43642448870155437
+    },
+    {
+      "batch_loss": 2.9277164936065674,
+      "innovation_rms": 0.12901590503433677,
+      "instruction_rms": 0.012726671557968573,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1305.14013671875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 835252.8125,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 575.3596801757812,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.150208830833435,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5136908292770386,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 259.9652099609375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 144.36927795410156,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 127.21295928955078,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2035.3824462890625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2620.48828125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 575.3597412109375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.154754877090454,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5133702158927917,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 259.91082763671875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.34342994680669703,
+      "step": 67,
+      "teaching_rms": 0.12901590503433677,
+      "traffic_rms": 0.34319461112315824
+    },
+    {
+      "batch_loss": 2.7110092639923096,
+      "innovation_rms": 0.11731177767800663,
+      "instruction_rms": 0.06590872812180482,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 994.9151611328125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1122645.25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 559.735107421875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1411356925964355,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.41180431842803955,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 315.938720703125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 142.0532989501953,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23.15973663330078,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2257.914794921875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2225.32275390625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 559.7351684570312,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1456812620162964,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.41191309690475464,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 315.88433837890625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.4280361332968367,
+      "step": 68,
+      "teaching_rms": 0.11731177767800663,
+      "traffic_rms": 0.4229316077624109
+    },
+    {
+      "batch_loss": 2.6919002532958984,
+      "innovation_rms": 0.311625413753151,
+      "instruction_rms": 0.2882152622464638,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 846.1014404296875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1451099.75,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 541.810546875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.16278874874115,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.701188325881958,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 370.1197814941406,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 126.68783569335938,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 153.65460205078125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2446.888427734375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1889.736328125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 541.8106079101562,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1673338413238525,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7012971639633179,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 370.0653991699219,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.5367136747137968,
+      "step": 69,
+      "teaching_rms": 0.311625413753151,
+      "traffic_rms": 0.4527622472226373
+    },
+    {
+      "batch_loss": 2.683516263961792,
+      "innovation_rms": 0.22311774942885151,
+      "instruction_rms": 0.0660014415509109,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 975.3611450195312,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1699930.625,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 736.3604736328125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1743135452270508,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.38738229870796204,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 418.9798889160156,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 105.2354507446289,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 214.52381896972656,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2627.25048828125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1803.62060546875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 736.3605346679688,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1788581609725952,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.387515127658844,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 418.9255065917969,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.4765346005592224,
+      "step": 70,
+      "teaching_rms": 0.22311774942885151,
+      "traffic_rms": 0.47194173336348644
+    },
+    {
+      "batch_loss": 2.88777494430542,
+      "innovation_rms": 0.49723513059474794,
+      "instruction_rms": 0.054381296179214736,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1162.4237060546875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2412714.5,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1481.525390625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.164742350578308,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7209041714668274,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 431.9229431152344,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 96.66265106201172,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 85.72799682617188,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2906.490966796875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2792.404296875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1481.525390625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1692866086959839,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7207953929901123,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 431.8685607910156,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.7436776155983003,
+      "step": 71,
+      "teaching_rms": 0.49723513059474794,
+      "traffic_rms": 0.7416866945711909
+    },
+    {
+      "batch_loss": 2.7985482215881348,
+      "innovation_rms": 0.5069711965763046,
+      "instruction_rms": 0.02103685515817409,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1643.0394287109375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3597376.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 938.61572265625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.135682463645935,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5232484936714172,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 485.0625915527344,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 98.93904876708984,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22.76397705078125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3262.56396484375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3560.728759765625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 938.61572265625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1402262449264526,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5233424305915833,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 485.00823974609375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.8930689405778567,
+      "step": 72,
+      "teaching_rms": 0.5069711965763046,
+      "traffic_rms": 0.8928212457743013
+    },
+    {
+      "batch_loss": 2.6678414344787598,
+      "innovation_rms": 0.5959988062709746,
+      "instruction_rms": 0.0056863490005857826,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2120.9765625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4855907.5,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1609.1025390625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1163495779037476,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5835896134376526,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 645.5194091796875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 116.0494384765625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 171.10389709472656,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3720.547607421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4579.8369140625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1609.1025390625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1006327867507935,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5834567546844482,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 645.5280151367188,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.9622977698112016,
+      "step": 73,
+      "teaching_rms": 0.5959988062709746,
+      "traffic_rms": 0.9622810412972747
+    },
+    {
+      "batch_loss": 2.5630953311920166,
+      "innovation_rms": 0.6901401666902398,
+      "instruction_rms": 0.024318545207142165,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2793.4609375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5895401.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1759.958740234375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1320587396621704,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.478259414434433,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 732.5122680664062,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 147.11021423339844,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 310.6077880859375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4302.12744140625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5815.79638671875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1759.958740234375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.067610740661621,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4780871868133545,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 732.5208740234375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.9734595800197348,
+      "step": 74,
+      "teaching_rms": 0.6901401666902398,
+      "traffic_rms": 0.973155823294966
+    },
+    {
+      "batch_loss": 2.5907602310180664,
+      "innovation_rms": 0.9369858055135685,
+      "instruction_rms": 0.018531792159115316,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3593.6318359375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7748803.5,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1705.7100830078125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.138120412826538,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.30350011587142944,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 825.2273559570312,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 204.67347717285156,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 575.6326293945312,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5111.5654296875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8094.3779296875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1705.7100830078125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0669828653335571,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.303387314081192,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 825.2359619140625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.2012941359768132,
+      "step": 75,
+      "teaching_rms": 0.9369858055135685,
+      "traffic_rms": 1.2011512222015557
+    },
+    {
+      "batch_loss": 2.5650691986083984,
+      "innovation_rms": 0.9331135419132125,
+      "instruction_rms": 0.007899093330351166,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3954.60205078125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9463860.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2658.0849609375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1536518335342407,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5123666524887085,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 980.4601440429688,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 274.0263977050781,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 693.5291748046875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6071.9072265625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9603.41796875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2658.085205078125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0631496906280518,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.512526273727417,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 980.46875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.2666657814697966,
+      "step": 76,
+      "teaching_rms": 0.9331135419132125,
+      "traffic_rms": 1.2666412033899417
+    },
+    {
+      "batch_loss": 2.474125862121582,
+      "innovation_rms": 0.8986322104740131,
+      "instruction_rms": 0.02285929180078884,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6155.033203125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12920832.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1776.378662109375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.192914366722107,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4762081205844879,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1138.751220703125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 336.7297058105469,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 627.032958984375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7073.86279296875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10019.5546875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1776.378662109375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.060442566871643,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4763677418231964,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1138.7598876953125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.6037263563398982,
+      "step": 77,
+      "teaching_rms": 0.8986322104740131,
+      "traffic_rms": 1.6035634280712907
+    },
+    {
+      "batch_loss": 2.766940116882324,
+      "innovation_rms": 0.9660500106823019,
+      "instruction_rms": 0.04215730618124908,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9126.4052734375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17395740.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2178.343017578125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2023004293441772,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5079782009124756,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1305.6983642578125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 385.1786193847656,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 484.4891357421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8096.2646484375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10224.01953125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2178.343017578125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0661022663116455,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5080909729003906,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1305.70703125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.8546986399966847,
+      "step": 78,
+      "teaching_rms": 0.9660500106823019,
+      "traffic_rms": 1.8542194638643927
+    },
+    {
+      "batch_loss": 2.5057625770568848,
+      "innovation_rms": 1.2204885803735013,
+      "instruction_rms": 0.015047596895545183,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12064.4580078125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24590448.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2635.80908203125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1991353034973145,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5108559131622314,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1499.6605224609375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 428.6344909667969,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 434.55865478515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9210.6240234375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11143.58984375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2635.80908203125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.062853217124939,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5110607743263245,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1499.669189453125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2.1669538210531085,
+      "step": 79,
+      "teaching_rms": 1.2204885803735013,
+      "traffic_rms": 2.1669015729521197
+    },
+    {
+      "batch_loss": 2.6486597061157227,
+      "innovation_rms": 1.5814475330401088,
+      "instruction_rms": 0.09548419596560503,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13943.478515625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34032764.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4789.9697265625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2042156457901,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6315016746520996,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1757.1917724609375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 485.41204833984375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 567.7756958007812,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10530.6171875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13199.935546875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4789.9697265625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0408507585525513,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6313889026641846,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1757.200439453125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2.626749705590453,
+      "step": 80,
+      "teaching_rms": 1.5814475330401088,
+      "traffic_rms": 2.6250137247132304
+    },
+    {
+      "batch_loss": 2.3567957878112793,
+      "innovation_rms": 1.7934331271398425,
+      "instruction_rms": 0.062078883027100226,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13465.4736328125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 40928520.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5013.529296875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2126015424728394,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.40827643871307373,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2014.78076171875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 562.94482421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 775.3277587890625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12118.8701171875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15882.5244140625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5013.529296875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0307551622390747,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4081636667251587,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2014.7894287109375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2.5137616564096876,
+      "step": 81,
+      "teaching_rms": 1.7934331271398425,
+      "traffic_rms": 2.512995018301379
+    },
+    {
+      "batch_loss": 2.5931732654571533,
+      "innovation_rms": 1.8455166882201968,
+      "instruction_rms": 0.0275920675708542,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17028.7734375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 54994608.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3688.41259765625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2186397314071655,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.43427979946136475,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2290.197265625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 634.2454223632812,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 713.0059204101562,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13848.0859375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17292.154296875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3688.41259765625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0562362670898438,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.43431413173675537,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2290.2060546875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3.218825119673509,
+      "step": 82,
+      "teaching_rms": 1.8455166882201968,
+      "traffic_rms": 3.2187068536050893
+    },
+    {
+      "batch_loss": 2.9114420413970947,
+      "innovation_rms": 1.9477108722791987,
+      "instruction_rms": 0.12357752947523508,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21993.4375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 73442088.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4304.357421875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2225544452667236,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6787253022193909,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2578.307373046875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 686.2507934570312,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 520.053955078125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15666.6962890625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18186.107421875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4304.357421875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0746501684188843,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6786982417106628,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2578.316162109375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3.6334470187954637,
+      "step": 83,
+      "teaching_rms": 1.9477108722791987,
+      "traffic_rms": 3.6313449111654443
+    },
+    {
+      "batch_loss": 2.615722417831421,
+      "innovation_rms": 2.4308334518247614,
+      "instruction_rms": 0.04157323506658644,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26772.755859375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 103159592.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4503.525390625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2261089086532593,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.49284911155700684,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2913.008544921875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 736.318359375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 500.67535400390625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17719.1171875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20524.212890625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4503.525390625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.062441349029541,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4926351308822632,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2913.017333984375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4.1998719226151975,
+      "step": 84,
+      "teaching_rms": 2.4308334518247614,
+      "traffic_rms": 4.199666180308627
+    },
+    {
+      "batch_loss": 3.027332305908203,
+      "innovation_rms": 3.3589303705497167,
+      "instruction_rms": 0.13175662553591774,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29609.796875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 137953632.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5893.384765625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2293510437011719,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.978291392326355,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3426.423583984375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 848.3775024414062,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1120.59130859375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20380.685546875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26615.6796875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5893.384765625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0268675088882446,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9784379005432129,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3426.432373046875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 5.0816989048025,
+      "step": 85,
+      "teaching_rms": 3.3589303705497167,
+      "traffic_rms": 5.079990562811194
+    },
+    {
+      "batch_loss": 2.981295347213745,
+      "innovation_rms": 3.2969547450945846,
+      "instruction_rms": 0.013782826565211497,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28382.291015625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 158722992.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7709.556640625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2323791980743408,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8663213849067688,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4003.778076171875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 984.2557983398438,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1358.783203125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23461.33984375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30806.53515625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7709.556640625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9890250563621521,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8665199875831604,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4003.786865234375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4.758571479788445,
+      "step": 86,
+      "teaching_rms": 3.2969547450945846,
+      "traffic_rms": 4.758551548945416
+    },
+    {
+      "batch_loss": 3.2111639976501465,
+      "innovation_rms": 4.901612549118627,
+      "instruction_rms": 0.05055850845593573,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30915.275390625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 232109328.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7182.76318359375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2352887392044067,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1398338079452515,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4432.71923828125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1250.3021240234375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2660.462890625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27625.638671875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41642.9921875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7182.76318359375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0923678874969482,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1397523880004883,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4432.72802734375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 6.646987482737883,
+      "step": 87,
+      "teaching_rms": 4.901612549118627,
+      "traffic_rms": 6.646795222606981
+    },
+    {
+      "batch_loss": 3.0289487838745117,
+      "innovation_rms": 5.887815518004567,
+      "instruction_rms": 0.009915300413684203,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34536.5,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 315397568.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9392.796875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.237768292427063,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7911112308502197,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4933.0361328125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1663.5152587890625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4132.1318359375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33116.33203125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 54906.93359375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9392.796875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1626100540161133,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7910298109054565,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4933.044921875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 7.635853904494076,
+      "step": 88,
+      "teaching_rms": 5.887815518004567,
+      "traffic_rms": 7.635847482274969
+    },
+    {
+      "batch_loss": 3.605638265609741,
+      "innovation_rms": 6.551307961503503,
+      "instruction_rms": 0.24420775446916132,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36446.88671875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 393605632.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12309.41796875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2393265962600708,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.569900631904602,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5640.14892578125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2208.85888671875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5453.43603515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39865.140625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 67488.1015625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12309.41796875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1775425672531128,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5699295997619629,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5640.1650390625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 8.470810126128148,
+      "step": 89,
+      "teaching_rms": 6.551307961503503,
+      "traffic_rms": 8.467289244633154
+    },
+    {
+      "batch_loss": 3.0846385955810547,
+      "innovation_rms": 7.117811497906396,
+      "instruction_rms": 0.020939364666174718,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 46501.58203125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 556572800.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12062.044921875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2401031255722046,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6615093350410461,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6742.40869140625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2774.4951171875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5656.3623046875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 47536.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 76716.09375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12062.044921875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1608370542526245,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6615383625030518,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6742.4248046875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 11.214017447418323,
+      "step": 90,
+      "teaching_rms": 7.117811497906396,
+      "traffic_rms": 11.213997907752432
+    },
+    {
+      "batch_loss": 2.7253620624542236,
+      "innovation_rms": 7.903008611611424,
+      "instruction_rms": 0.096101390605909,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 56616.3984375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 734328640.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14995.02734375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2407326698303223,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.46995601058006287,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8241.9111328125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3370.286865234375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5957.9169921875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 56100.48828125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 85637.3984375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14995.02734375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1198810338974,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.46998506784439087,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8241.927734375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 12.17793342974747,
+      "step": 91,
+      "teaching_rms": 7.903008611611424,
+      "traffic_rms": 12.177554239677573
+    },
+    {
+      "batch_loss": 2.9531383514404297,
+      "innovation_rms": 11.392468824122318,
+      "instruction_rms": 1.4902839060801483,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 60363.75,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1016573888.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41083.93359375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2420693635940552,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7920073866844177,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9054.0400390625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4202.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8323.380859375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 66942.3828125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 108418.9609375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41083.93359375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0541988611221313,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7918230295181274,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9054.056640625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 15.09439533422483,
+      "step": 92,
+      "teaching_rms": 11.392468824122318,
+      "traffic_rms": 15.020646606512098
+    },
+    {
+      "batch_loss": 2.5661420822143555,
+      "innovation_rms": 12.045317469762441,
+      "instruction_rms": 0.7502247595709868,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 73096.3984375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1653744896.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38966.21875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2438013553619385,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5715024471282959,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10603.904296875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5135.8251953125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9332.001953125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 79284.234375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 123418.546875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38966.21875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0290112495422363,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5713180899620056,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10603.9208984375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 20.21337703640122,
+      "step": 93,
+      "teaching_rms": 12.045317469762441,
+      "traffic_rms": 20.199449843544492
+    },
+    {
+      "batch_loss": 2.6598050594329834,
+      "innovation_rms": 14.368176816272697,
+      "instruction_rms": 0.5450545014147837,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 88803.3828125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2483663360.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63169.5,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2455008029937744,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.34767842292785645,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15079.345703125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6191.384765625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10555.595703125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 93549.859375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 142656.25,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63169.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0407719612121582,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3474952280521393,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15079.431640625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 23.42282741673072,
+      "step": 94,
+      "teaching_rms": 14.368176816272697,
+      "traffic_rms": 23.416484789409076
+    },
+    {
+      "batch_loss": 2.915207624435425,
+      "innovation_rms": 20.472754024696453,
+      "instruction_rms": 0.7617634350210508,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 100301.8671875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3450461696.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 76501.1015625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2474591732025146,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5926806926727295,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22729.455078125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7793.4384765625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16020.53515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 112688.5703125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 191387.125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 76501.1015625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0346332788467407,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5923617482185364,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22729.541015625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 26.179454728404107,
+      "step": 95,
+      "teaching_rms": 20.472754024696453,
+      "traffic_rms": 26.168369587098105
+    },
+    {
+      "batch_loss": 3.017101764678955,
+      "innovation_rms": 9.00437041954693,
+      "instruction_rms": 0.5054029967488326,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 95379.5625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3939945216.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 53066.5625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2483724355697632,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.699354350566864,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28036.111328125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8705.994140625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9125.5556640625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 130186.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 174974.265625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 53066.5625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0158299207687378,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6993163824081421,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28036.197265625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 20.895911519719785,
+      "step": 96,
+      "teaching_rms": 9.00437041954693,
+      "traffic_rms": 20.889798616715346
+    },
+    {
+      "batch_loss": 2.9514708518981934,
+      "innovation_rms": 11.432937077694667,
+      "instruction_rms": 0.5506219801015574,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 93893.609375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4641087488.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34912.6015625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2454643249511719,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6687343716621399,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30612.646484375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9004.287109375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2982.9296875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 146498.09375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 163120.90625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34912.6015625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9850057363510132,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6690537929534912,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30612.732421875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 25.019672252021802,
+      "step": 97,
+      "teaching_rms": 11.432937077694667,
+      "traffic_rms": 25.013612596442236
+    },
+    {
+      "batch_loss": 2.939070701599121,
+      "innovation_rms": 30.845611102786833,
+      "instruction_rms": 0.5454424422865625,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 120036.5234375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7281275904.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 78466.6953125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.243788480758667,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8707284331321716,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33268.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10072.048828125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10677.61328125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 170060.0625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 235619.71875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 78466.6953125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.957617461681366,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.870966374874115,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33268.8359375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 41.86190551346288,
+      "step": 98,
+      "teaching_rms": 30.845611102786833,
+      "traffic_rms": 41.85835192613411
+    },
+    {
+      "batch_loss": 3.0843405723571777,
+      "innovation_rms": 31.843570927450234,
+      "instruction_rms": 0.3954096550532135,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 153414.53125,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12147561472.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 53353.3515625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2455326318740845,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5681056976318359,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 37574.60546875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11741.07421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16690.25390625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 198624.9375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 285648.6875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 53353.3515625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9421226978302002,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5683436393737793,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 37574.69140625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 52.80773826095765,
+      "step": 99,
+      "teaching_rms": 31.843570927450234,
+      "traffic_rms": 52.80625788408374
+    },
+    {
+      "batch_loss": 2.925293207168579,
+      "innovation_rms": 43.12883298357078,
+      "instruction_rms": 0.3301051041996635,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 189384.640625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16263103488.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 66175.984375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.252729058265686,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.679929792881012,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 43831.15234375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14655.310546875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29142.365234375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 237580.46875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 389555.3125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 66175.984375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9391531944274902,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.680103600025177,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 43831.23828125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 54.61619068472097,
+      "step": 100,
+      "teaching_rms": 43.12883298357078,
+      "traffic_rms": 54.615193085721444
+    },
+    {
+      "batch_loss": 2.8266639709472656,
+      "innovation_rms": 43.82488468235351,
+      "instruction_rms": 0.49239893715795174,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 222474.421875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19405441024.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 155023.109375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2432429790496826,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5547952651977539,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41449.87109375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18243.478515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35881.6875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 283786.65625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 462061.90625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 155023.109375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9470289945602417,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5548332333564758,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41449.95703125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 58.23632937951089,
+      "step": 101,
+      "teaching_rms": 43.82488468235351,
+      "traffic_rms": 58.234247680066915
+    },
+    {
+      "batch_loss": 3.662792205810547,
+      "innovation_rms": 41.05551177473504,
+      "instruction_rms": 0.26391551913514294,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 317712.34375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 32044472320.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 123698.859375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2395730018615723,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.37808501720428467,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 49024.05859375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21442.548828125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31990.701171875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 330650.875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 468642.0625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 123698.859375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9466376304626465,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3781552314758301,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 49023.9921875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 88.2200018281759,
+      "step": 102,
+      "teaching_rms": 41.05551177473504,
+      "traffic_rms": 88.21960707019826
+    },
+    {
+      "batch_loss": 3.6556308269500732,
+      "innovation_rms": 41.696474343330316,
+      "instruction_rms": 0.4467335927147005,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 450994.84375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 64248741888.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 103541.3828125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2252906560897827,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6527115106582642,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 59378.1953125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23834.080078125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23915.3125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 376003.96875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 453531.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 103541.3828125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9355630874633789,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6527817249298096,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 59378.12890625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 104.65440279118381,
+      "step": 103,
+      "teaching_rms": 41.696474343330316,
+      "traffic_rms": 104.65344931203877
+    },
+    {
+      "batch_loss": 3.1498186588287354,
+      "innovation_rms": 45.50528795083991,
+      "instruction_rms": 0.4140605477513846,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 607876.5,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 129895030784.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86757.484375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1837741136550903,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.977409839630127,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 68053.9453125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25272.814453125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14387.3369140625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 419638.21875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 436342.46875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86757.484375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9258806705474854,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9774856567382812,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 68053.875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 118.53876064101895,
+      "step": 104,
+      "teaching_rms": 45.50528795083991,
+      "traffic_rms": 118.53803747396184
+    },
+    {
+      "batch_loss": 3.0283455848693848,
+      "innovation_rms": 51.43574998891074,
+      "instruction_rms": 0.4793589208688736,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 781005.0625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 210207916032.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 80793.7109375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1631567478179932,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8725119829177856,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 75072.8359375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25833.951171875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5611.375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 462311.125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 426729.15625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 80793.7109375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9444546103477478,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8726747035980225,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 75072.765625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 132.2762268494219,
+      "step": 105,
+      "teaching_rms": 51.43574998891074,
+      "traffic_rms": 132.27535826775366
+    },
+    {
+      "batch_loss": 3.4575417041778564,
+      "innovation_rms": 57.23377795047765,
+      "instruction_rms": 1.3989412439348492,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 962211.375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 297896738816.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 87820.9453125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1588451862335205,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6345753073692322,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 80473.8046875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25579.2890625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2546.6181640625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 504757.3125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 424461.8125,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 87820.9453125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9483126401901245,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6343911290168762,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 80473.734375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 144.67875957883524,
+      "step": 106,
+      "teaching_rms": 57.23377795047765,
+      "traffic_rms": 144.67199603762444
+    },
+    {
+      "batch_loss": 2.9369454383850098,
+      "innovation_rms": 62.89025803453876,
+      "instruction_rms": 6.429681223701708,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1134753.25,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 379065008128.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86952.109375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.162691354751587,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6942891478538513,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 84296.703125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24519.44921875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10598.4052734375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 547533.6875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 427763.4375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86952.109375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9460301995277405,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.694040060043335,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 84296.6328125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 155.83307352062286,
+      "step": 107,
+      "teaching_rms": 62.89025803453876,
+      "traffic_rms": 155.700372517713
+    },
+    {
+      "batch_loss": 2.698108434677124,
+      "innovation_rms": 70.49683540573764,
+      "instruction_rms": 4.340719362257332,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1294861.625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 444540780544.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 100176.7109375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1683412790298462,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.525611162185669,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86429.484375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22863.767578125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16556.82421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 592327.375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 447936.71875,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 100176.7109375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9756701588630676,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5253621339797974,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86429.4140625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 165.3991697267064,
+      "step": 108,
+      "teaching_rms": 70.49683540573764,
+      "traffic_rms": 165.3422012126818
+    },
+    {
+      "batch_loss": 2.931814432144165,
+      "innovation_rms": 82.41276084519632,
+      "instruction_rms": 2.124608629272954,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1429593.5,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 480454836224.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 155229.46875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.170103907585144,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.397666335105896,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 90750.71875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21110.74609375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17530.20703125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 642617.5625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 502901.625,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 155229.46875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0012820959091187,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.3979412913322449,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 90750.765625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 173.9956360833988,
+      "step": 109,
+      "teaching_rms": 82.41276084519632,
+      "traffic_rms": 173.9826641203834
+    },
+    {
+      "batch_loss": 3.071408748626709,
+      "innovation_rms": 107.2342055622969,
+      "instruction_rms": 1.4477689859761262,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1523230.375,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 478138105856.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 261698.515625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1578727960586548,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4557751417160034,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 103605.5234375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20751.646484375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3591.00390625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 708495.125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 658775.75,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 261698.515625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0109890699386597,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4560500383377075,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 103605.53125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 181.57028933754336,
+      "step": 110,
+      "teaching_rms": 107.2342055622969,
+      "traffic_rms": 181.56451728037354
+    },
+    {
+      "batch_loss": 2.640561819076538,
+      "innovation_rms": 135.52137693599906,
+      "instruction_rms": 0.47361342308153753,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1524808.25,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 449764261888.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 222424.1875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1476079225540161,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5703049898147583,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 123739.375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24056.05078125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33044.03515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 803385.0625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 948899.25,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 222424.1875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0196739435195923,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5701068043708801,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 123739.234375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 180.22101131323535,
+      "step": 111,
+      "teaching_rms": 135.52137693599906,
+      "traffic_rms": 180.22038899411436
+    },
+    {
+      "batch_loss": 2.9016144275665283,
+      "innovation_rms": 164.1159447457429,
+      "instruction_rms": 0.3125165117375963,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1525871.625,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 433082105856.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 190438.265625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1486350297927856,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5346590876579285,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 141060.515625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 32166.421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 81103.703125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 935709.375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1323243.25,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 190438.265625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.019037127494812,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5344641804695129,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 141060.375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 214.04812002980754,
+      "step": 112,
+      "teaching_rms": 164.1159447457429,
+      "traffic_rms": 214.0478918888926
+    },
+    {
+      "batch_loss": 3.340075731277466,
+      "innovation_rms": 187.5714741794353,
+      "instruction_rms": 0.27244906874282765,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1664354.25,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 452221763584.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 272124.34375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1572747230529785,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3947086334228516,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 166969.703125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 44053.07421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 118866.5234375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1105689.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1699798.375,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 272124.34375,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0144473314285278,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3943862915039062,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 166969.5625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 255.50549494550614,
+      "step": 113,
+      "teaching_rms": 187.5714741794353,
+      "traffic_rms": 255.50534968848314
+    },
+    {
+      "batch_loss": 3.2119524478912354,
+      "innovation_rms": 228.85881485712105,
+      "instruction_rms": 0.26285197009796335,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1708019.875,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 541990289408.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 373658.6875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1767817735671997,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0319122076034546,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 174118.984375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 61732.9921875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 176799.1875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1325638.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2199494.25,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 373658.6875,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0349206924438477,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0315899848937988,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 174118.84375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 289.877660825223,
+      "step": 114,
+      "teaching_rms": 228.85881485712105,
+      "traffic_rms": 289.8775416522569
+    },
+    {
+      "batch_loss": 3.2312405109405518,
+      "innovation_rms": 230.70748287098,
+      "instruction_rms": 0.2494260461961365,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1967793.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 768522977280.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 327854.0625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.204316258430481,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6053696870803833,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 206394.9375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 80941.0390625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 192080.453125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1570637.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2449984.25,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 327854.0625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.049362301826477,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6050474643707275,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 206394.796875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 375.5269591130475,
+      "step": 115,
+      "teaching_rms": 230.70748287098,
+      "traffic_rms": 375.5268762788022
+    },
+    {
+      "batch_loss": 3.0326597690582275,
+      "innovation_rms": 287.73571070848885,
+      "instruction_rms": 0.2262082878100995,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2254399.5,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1103229878272.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 570543.5625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.203927755355835,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7200848460197449,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 263449.28125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 103210.890625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 222698.53125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1855630.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2849937.5,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 570543.5625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.053753137588501,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7199393510818481,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 263449.15625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 454.4980115453296,
+      "step": 116,
+      "teaching_rms": 287.73571070848885,
+      "traffic_rms": 454.49795525257156
+    },
+    {
+      "batch_loss": 2.943819046020508,
+      "innovation_rms": 323.0424756820146,
+      "instruction_rms": 0.411525512279398,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2355734.75,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1293031309312.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1056063.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.155564546585083,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5260601043701172,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 290037.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 128927.3359375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 257164.46875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2187669.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3320388.5,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1056063.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.047479510307312,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5260794162750244,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 290037.375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 437.8298208045258,
+      "step": 117,
+      "teaching_rms": 323.0424756820146,
+      "traffic_rms": 437.82962740391855
+    },
+    {
+      "batch_loss": 3.189453601837158,
+      "innovation_rms": 323.7121120374936,
+      "instruction_rms": 1.0879591204593757,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3173176.75,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1847120101376.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 716080.3125,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1220262050628662,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0477097034454346,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 335718.875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 151619.796875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 226924.671875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2532467.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3447982.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 716080.3125,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0204190015792847,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0476012229919434,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 335718.75,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 608.2618581059389,
+      "step": 118,
+      "teaching_rms": 323.7121120374936,
+      "traffic_rms": 608.2608851240227
+    },
+    {
+      "batch_loss": 3.4490089416503906,
+      "innovation_rms": 330.93853607215163,
+      "instruction_rms": 1.8764830120710352,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4396273.5,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3617495973888.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 782625.5,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.12716805934906,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8706505298614502,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 390660.03125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 168305.171875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 166853.8125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2874701.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3422340.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 782625.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9737936854362488,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8706313371658325,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 390659.90625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 698.8718545176064,
+      "step": 119,
+      "teaching_rms": 330.93853607215163,
+      "traffic_rms": 698.8693353189916
+    },
+    {
+      "batch_loss": 3.0174272060394287,
+      "innovation_rms": 349.0306058289314,
+      "instruction_rms": 0.12639401672740527,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5825698.5,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7044931256320.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 913711.5625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1585121154785156,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4215455055236816,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 449716.03125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 178124.03125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 98188.5703125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3210690.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3359885.75,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 913711.5625,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9850819706916809,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4215261936187744,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 449715.90625,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 770.6509492526926,
+      "step": 120,
+      "teaching_rms": 349.0306058289314,
+      "traffic_rms": 770.6509388877591
+    },
+    {
+      "batch_loss": 3.795043468475342,
+      "innovation_rms": 416.0778064550746,
+      "instruction_rms": 1.0538432649394032,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7379731.5,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11106994094080.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1031365.25,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2091517448425293,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0361628532409668,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 517440.9375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 181919.15625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 37951.32421875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3553274.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3425839.5,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1031365.25,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9943071007728577,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0361247062683105,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 517440.96875,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 880.6504894230608,
+      "step": 121,
+      "teaching_rms": 416.0778064550746,
+      "traffic_rms": 880.6498588744729
+    },
+    {
+      "batch_loss": 4.717951774597168,
+      "innovation_rms": 478.8642224296497,
+      "instruction_rms": 1.2283072140896445,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8854510.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14881072873472.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1130369.5,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2648406028747559,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4654737710952759,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 630477.875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 181420.609375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4985.4140625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3915157.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3618826.25,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1130369.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.98363196849823,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4654459953308105,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 630477.9375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 973.5306587589234,
+      "step": 122,
+      "teaching_rms": 478.8642224296497,
+      "traffic_rms": 973.5298838793262
+    },
+    {
+      "batch_loss": 3.527984619140625,
+      "innovation_rms": 584.5586153300495,
+      "instruction_rms": 1.069727660280836,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10038373.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17184561561600.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1182384.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.317795753479004,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9430664777755737,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 748716.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 181969.734375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5491.25244140625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4328161.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4130039.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1182384.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0049411058425903,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9430387020111084,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 748716.3125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1089.5535248308906,
+      "step": 123,
+      "teaching_rms": 584.5586153300495,
+      "traffic_rms": 1089.5529996996943
+    },
+    {
+      "batch_loss": 4.550325393676758,
+      "innovation_rms": 782.1404094351735,
+      "instruction_rms": 0.5618852669652941,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10676231.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17416732016640.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2499392.75,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2618032693862915,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8488470315933228,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 847915.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 195406.46875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 134367.296875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4872857.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5446959.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2499392.75,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.948979914188385,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8488138914108276,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 847915.3125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1254.4493207571018,
+      "step": 124,
+      "teaching_rms": 782.1404094351735,
+      "traffic_rms": 1254.4491949190906
+    },
+    {
+      "batch_loss": 4.378256797790527,
+      "innovation_rms": 697.7643735969441,
+      "instruction_rms": 0.40563086193188735,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10310451.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16863341838336.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2234585.5,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2266045808792114,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0241602659225464,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 954818.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 208659.34375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 132528.6875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5485332.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6124755.5,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2234585.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.913812518119812,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0240211486816406,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 954818.3125,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1041.756789294219,
+      "step": 125,
+      "teaching_rms": 697.7643735969441,
+      "traffic_rms": 1041.7567103239635
+    },
+    {
+      "batch_loss": 4.766293048858643,
+      "innovation_rms": 1084.1148705143933,
+      "instruction_rms": 0.3097509164050097,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11065359.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16917727281152.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1703530.25,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2042546272277832,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9427011609077454,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1079797.375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 246598.890625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 379395.4375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6322131.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8367986.5,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1703530.25,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9174820184707642,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9425619840621948,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1079797.5,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1504.625826769009,
+      "step": 126,
+      "teaching_rms": 1084.1148705143933,
+      "traffic_rms": 1504.625794885595
+    },
+    {
+      "batch_loss": 3.576258897781372,
+      "innovation_rms": 1299.6611165013762,
+      "instruction_rms": 0.24742918860958846,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11743141.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19346441109504.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2321468.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1854966878890991,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0059874057769775,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1157776.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 317112.9375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 705140.4375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7436350.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11142193.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2321468.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9541501998901367,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0058287382125854,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1157776.375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1828.0269407286098,
+      "step": 127,
+      "teaching_rms": 1299.6611165013762,
+      "traffic_rms": 1828.0269239835095
+    },
+    {
+      "batch_loss": 3.1089963912963867,
+      "innovation_rms": 1427.6747052630135,
+      "instruction_rms": 0.32790850260104637,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12040637.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22493538549760.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3327499.5,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.168602705001831,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9882718324661255,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1393534.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 418353.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1012406.8125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8822716.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13863656.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3327499.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9794673919677734,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9882419109344482,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1393534.75,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1837.4939748333213,
+      "step": 128,
+      "teaching_rms": 1427.6747052630135,
+      "traffic_rms": 1837.4939455750548
+    },
+    {
+      "batch_loss": 2.863762378692627,
+      "innovation_rms": 1512.5866656318385,
+      "instruction_rms": 0.40021730477894774,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14303887.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 32226483896320.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2414631.25,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1533863544464111,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9787696599960327,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1564162.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 526387.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1080341.375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10377717.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15550005.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2414631.25,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.996431827545166,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9790319204330444,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1564162.75,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2511.889118334292,
+      "step": 129,
+      "teaching_rms": 1512.5866656318385,
+      "traffic_rms": 2511.88908645111
+    },
+    {
+      "batch_loss": 3.2466955184936523,
+      "innovation_rms": 1687.86703147428,
+      "instruction_rms": 0.7122092345113342,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17127094.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 45873337729024.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2437092.5,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1396801471710205,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7392107248306274,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1672925.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 635631.125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1092433.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12085590.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17078734.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2437092.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0023584365844727,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7394729256629944,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1672925.75,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2964.1203671143758,
+      "step": 130,
+      "teaching_rms": 1687.86703147428,
+      "traffic_rms": 2964.120281550778
+    },
+    {
+      "batch_loss": 3.177769184112549,
+      "innovation_rms": 2257.9336797587443,
+      "instruction_rms": 0.7086531163800945,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19559270.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 65252901257216.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4957898.5,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.127333164215088,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8873834609985352,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1730434.375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 775695.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1400644.875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14157098.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20715082.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4957898.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.00437593460083,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8872808218002319,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1730434.375,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3468.6123246526445,
+      "step": 131,
+      "teaching_rms": 2257.9336797587443,
+      "traffic_rms": 3468.6122522622354
+    },
+    {
+      "batch_loss": 2.7658350467681885,
+      "innovation_rms": 2062.28536588471,
+      "instruction_rms": 0.6392194292494019,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19992560.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 72288292569088.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6054416.5,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.116209626197815,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6252777576446533,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2020140.125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 904197.875,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1285022.625,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16394130.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22370316.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6054416.5,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.009324550628662,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6251751184463501,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2020140.25,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3098.636399944063,
+      "step": 132,
+      "teaching_rms": 2062.28536588471,
+      "traffic_rms": 3098.6363340116172
+    },
+    {
+      "batch_loss": 3.1688928604125977,
+      "innovation_rms": 3126.0051665997885,
+      "instruction_rms": 0.3573416802369595,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24216718.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 103930885308416.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5575075.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1061872243881226,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5565117597579956,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2556467.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1079347.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1751494.125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19170912.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27767822.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5575075.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0138046741485596,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5564819574356079,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2556467.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4611.473493352158,
+      "step": 133,
+      "teaching_rms": 3126.0051665997885,
+      "traffic_rms": 4611.473479507024
+    },
+    {
+      "batch_loss": 3.387921094894409,
+      "innovation_rms": 3881.1557399696976,
+      "instruction_rms": 0.3725153147314392,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28579254.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 157574666125312.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8560228.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.097156047821045,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8782595992088318,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3108192.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1330345.375,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2509981.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22710686.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35397740.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8560228.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0150551795959473,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8781499862670898,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3108192.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 5616.80560303352,
+      "step": 134,
+      "teaching_rms": 3881.1557399696976,
+      "traffic_rms": 5616.805590680659
+    },
+    {
+      "batch_loss": 2.7914018630981445,
+      "innovation_rms": 3400.7590086641953,
+      "instruction_rms": 0.30963927953255643,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29327394.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 175685100371968.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9676144.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0890170335769653,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7301984429359436,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3661424.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1574226.125,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2438807.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26549712.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38390268.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9676144.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0170533657073975,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7300887703895569,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3661424.5,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4921.459248980768,
+      "step": 135,
+      "teaching_rms": 3400.7590086641953,
+      "traffic_rms": 4921.459239240113
+    },
+    {
+      "batch_loss": 3.1908466815948486,
+      "innovation_rms": 4760.608532769913,
+      "instruction_rms": 0.45932801060772327,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 37841524.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 259857550671872.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8497278.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0816810131072998,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8392083644866943,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4511152.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1853824.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2795986.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31052300.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 45025888.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8497278.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0156853199005127,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.83940589427948,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4511152.5,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 7409.732932032644,
+      "step": 136,
+      "teaching_rms": 4760.608532769913,
+      "traffic_rms": 7409.732917795835
+    },
+    {
+      "batch_loss": 2.8859994411468506,
+      "innovation_rms": 5629.8612713798,
+      "instruction_rms": 0.44254681816183145,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 47331004.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 384756759920640.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10702314.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.075063705444336,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.975231409072876,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5581384.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2185859.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3320350.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36358304.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 53060048.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10702314.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0181710720062256,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.975412130355835,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5581384.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 8886.91874147676,
+      "step": 137,
+      "teaching_rms": 5629.8612713798,
+      "traffic_rms": 8886.91873045789
+    },
+    {
+      "batch_loss": 3.1533501148223877,
+      "innovation_rms": 7943.879345153944,
+      "instruction_rms": 0.5174531458391697,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 51802568.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 564704112017408.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12196251.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.069207787513733,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9082887768745422,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6244521.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2754732.25,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5688725.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 43568184.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 72098816.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12196251.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0146607160568237,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.908076286315918,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6244521.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 10313.70270912099,
+      "step": 138,
+      "teaching_rms": 7943.879345153944,
+      "traffic_rms": 10313.702696140346
+    },
+    {
+      "batch_loss": 3.428781270980835,
+      "innovation_rms": 5413.832228992532,
+      "instruction_rms": 0.522727456516206,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 54333440.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 637635139731456.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12477298.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0645660161972046,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6204981207847595,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7294934.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3249822.75,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4950906.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 50829032.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 72608472.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12477298.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0057693719863892,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6202856302261353,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7294934.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 9037.484851836896,
+      "step": 139,
+      "teaching_rms": 5413.832228992532,
+      "traffic_rms": 9037.484836719654
+    },
+    {
+      "batch_loss": 2.9520466327667236,
+      "innovation_rms": 10369.494271664567,
+      "instruction_rms": 0.5653726803226765,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 62715068.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 913062936707072.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27802752.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0625202655792236,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7831572890281677,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9653736.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3961157.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7113347.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 60231900.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 94028664.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27802752.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.000232458114624,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7830610871315002,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9653736.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 14212.013545424248,
+      "step": 140,
+      "teaching_rms": 10369.494271664567,
+      "traffic_rms": 14212.01353417862
+    },
+    {
+      "batch_loss": 2.7612345218658447,
+      "innovation_rms": 8232.069155025027,
+      "instruction_rms": 0.6972479393764787,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 68187848.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1431808012976128.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21981382.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0595561265945435,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6589882969856262,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11472180.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4639994.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6788364.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 69792560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 95606576.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21981382.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.003562331199646,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6587651968002319,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11472180.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 17252.212131860073,
+      "step": 141,
+      "teaching_rms": 8232.069155025027,
+      "traffic_rms": 17252.212117770443
+    },
+    {
+      "batch_loss": 2.8146321773529053,
+      "innovation_rms": 9102.237495428895,
+      "instruction_rms": 0.6831310431935164,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 76399864.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2150636825083904.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20845816.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.053944706916809,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9093703627586365,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12914483.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5278476.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6384822.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 79450896.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 96583384.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20845816.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.996519148349762,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9093953967094421,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12914483.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 20314.47457691821,
+      "step": 142,
+      "teaching_rms": 9102.237495428895,
+      "traffic_rms": 20314.474565432116
+    },
+    {
+      "batch_loss": 2.9687376022338867,
+      "innovation_rms": 12658.701427103299,
+      "instruction_rms": 0.6551643283751594,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 90758120.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3055337997336576.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34892384.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.050175428390503,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6933075785636902,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14221451.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6022324.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7438480.5,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 90447688.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 109967912.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34892384.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9885625839233398,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6929895877838135,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14221451.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 22945.34373327101,
+      "step": 143,
+      "teaching_rms": 12658.701427103299,
+      "traffic_rms": 22945.34372391747
+    },
+    {
+      "batch_loss": 2.936171293258667,
+      "innovation_rms": 17835.417787676895,
+      "instruction_rms": 0.5971872270554114,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 112889512.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4021335734878208.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 32973472.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0473620891571045,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.543472170829773,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17257334.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7144984.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11226602.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 104850560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 144028704.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 32973472.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0109281539916992,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.543350100517273,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17257334.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 26012.84756172801,
+      "step": 144,
+      "teaching_rms": 17835.417787676895,
+      "traffic_rms": 26012.847554873082
+    },
+    {
+      "batch_loss": 2.702181816101074,
+      "innovation_rms": 18730.864488898456,
+      "instruction_rms": 0.5723295864178028,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 132201936.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4569380809605120.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 71005016.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0436909198760986,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6382465362548828,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21707718.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8568626.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14236424.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 122373280.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 175227168.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 71005016.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0310471057891846,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6381244659423828,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21707718.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 25075.889718954124,
+      "step": 145,
+      "teaching_rms": 18730.864488898456,
+      "traffic_rms": 25075.88971242273
+    },
+    {
+      "batch_loss": 2.699154853820801,
+      "innovation_rms": 23559.18038976243,
+      "instruction_rms": 0.6807716092635292,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 154845504.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7455964978479104.0,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 56437672.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0391085147857666,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4066677391529083,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25022714.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10450150.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18815236.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 143781152.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 214078720.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 56437672.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0491437911987305,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4065456986427307,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25022714.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 38239.80033299012,
+      "step": 146,
+      "teaching_rms": 23559.18038976243,
+      "traffic_rms": 38239.80032693034
+    },
+    {
+      "batch_loss": 2.5505611896514893,
+      "innovation_rms": 30637.864114983713,
+      "instruction_rms": 0.5575046339387586,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 181237024.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0540588478758912e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 113208016.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0345971584320068,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.49642443656921387,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28541990.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13057035.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26068848.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 171932896.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 281517504.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 113208016.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.065420389175415,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4962425231933594,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28541990.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 41947.326510211744,
+      "step": 147,
+      "teaching_rms": 30637.864114983713,
+      "traffic_rms": 41947.32650650696
+    },
+    {
+      "batch_loss": 2.550424337387085,
+      "innovation_rms": 23805.245288110378,
+      "instruction_rms": 0.5201724330862664,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 205583360.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1391016699428864e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 65665544.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.03141188621521,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6073606014251709,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33308798.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15564548.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25075132.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 201434000.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 295011104.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 65665544.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0800586938858032,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6071771383285522,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33308798.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 35921.222254686734,
+      "step": 148,
+      "teaching_rms": 23805.245288110378,
+      "traffic_rms": 35921.22225092046
+    },
+    {
+      "batch_loss": 2.9406521320343018,
+      "innovation_rms": 43442.77137778443,
+      "instruction_rms": 0.5355287339562754,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 252593952.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.570086498336768e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 135887040.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0282312631607056,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0814695358276367,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 40509800.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19018956.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34544072.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 240909696.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 394756928.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 135887040.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0931788682937622,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0816149711608887,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 40509800.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 56952.91485253286,
+      "step": 149,
+      "teaching_rms": 43442.77137778443,
+      "traffic_rms": 56952.91485001508
+    },
+    {
+      "batch_loss": 2.759784460067749,
+      "innovation_rms": 37712.069807045504,
+      "instruction_rms": 0.6601641956906482,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 275287648.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5096069436145664e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 122440680.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0379722118377686,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0910569429397583,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 46235068.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22630340.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36113844.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 282838400.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 419287136.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 122440680.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1040865182876587,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0912023782730103,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 46235068.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 72786.79670344673,
+      "step": 150,
+      "teaching_rms": 37712.069807045504,
+      "traffic_rms": 72786.79670045295
+    },
+    {
+      "batch_loss": 2.796752452850342,
+      "innovation_rms": 47874.282212063765,
+      "instruction_rms": 0.7903619814445358,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 306323648.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.805823349083341e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 276571648.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0572330951690674,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8778440952301025,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 71897952.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26739932.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41095928.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 329995936.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 471575328.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 276571648.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.112921953201294,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8775703310966492,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 71897952.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 85854.77229711761,
+      "step": 151,
+      "teaching_rms": 47874.282212063765,
+      "traffic_rms": 85854.77229347965
+    },
+    {
+      "batch_loss": 2.586143732070923,
+      "innovation_rms": 65319.618330886406,
+      "instruction_rms": 0.6766725559341822,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 377485952.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.860996740972544e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 148014224.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0616217851638794,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.087349772453308,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86040768.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31907908.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 51679764.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 389320224.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 593242880.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 148014224.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1245537996292114,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0875331163406372,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86040768.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 93352.40969352824,
+      "step": 152,
+      "teaching_rms": 65319.618330886406,
+      "traffic_rms": 93352.40969107578
+    },
+    {
+      "batch_loss": 2.7055630683898926,
+      "innovation_rms": 77793.32118657764,
+      "instruction_rms": 0.6070106106221511,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 429817728.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.894343109063475e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 221098752.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.163169503211975,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0154776573181152,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 108150640.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38529632.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 66217248.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 463760736.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 744405120.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 221098752.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1676771640777588,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0154731273651123,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 108150640.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 102847.41906348773,
+      "step": 153,
+      "teaching_rms": 77793.32118657764,
+      "traffic_rms": 102847.41906169642
+    },
+    {
+      "batch_loss": 2.7074077129364014,
+      "innovation_rms": 74064.8392431188,
+      "instruction_rms": 0.5319253705928304,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 548550016.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.373026717643571e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 205302560.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2284220457077026,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.861937940120697,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 128680896.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 44966560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 64369292.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 543491392.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 797306304.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 205302560.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2329293489456177,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8617547154426575,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 128680896.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 130442.2534701464,
+      "step": 154,
+      "teaching_rms": 74064.8392431188,
+      "traffic_rms": 130442.25346906185
+    },
+    {
+      "batch_loss": 3.639218807220459,
+      "innovation_rms": 78027.43662162602,
+      "instruction_rms": 0.5920673178148144,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 700345216.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.166827680998359e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 239563424.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2534295320510864,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5401358604431152,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 147711024.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 50673500.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 57069400.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 625673792.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 821824128.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 239563424.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2579363584518433,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5402812957763672,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 147711024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 152162.66159722157,
+      "step": 155,
+      "teaching_rms": 78027.43662162602,
+      "traffic_rms": 152162.6615960697
+    },
+    {
+      "batch_loss": 2.8792483806610107,
+      "innovation_rms": 92480.86751815495,
+      "instruction_rms": 0.7501366273409257,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 870676736.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9704646811097498e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 329359680.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2549266815185547,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1331758499145508,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 165294384.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 55935572.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 52620724.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 712994432.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 873206400.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 329359680.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2594330310821533,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1333651542663574,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 165294384.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 177740.0604024514,
+      "step": 156,
+      "teaching_rms": 92480.86751815495,
+      "traffic_rms": 177740.06040086847
+    },
+    {
+      "batch_loss": 4.163220405578613,
+      "innovation_rms": 120028.2538815274,
+      "instruction_rms": 0.6760951588768587,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1040421248.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.585905147842396e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 418208832.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2207814455032349,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9207028150558472,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 182971312.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 62038384.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 61028104.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 814660928.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1016664640.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 418208832.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2252873182296753,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9208730459213257,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 182971312.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 209180.39634919175,
+      "step": 157,
+      "teaching_rms": 120028.2538815274,
+      "traffic_rms": 209180.39634809914
+    },
+    {
+      "batch_loss": 4.599479675292969,
+      "innovation_rms": 172622.12655340284,
+      "instruction_rms": 0.7174052492386327,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1167665664.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.327250691071672e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 446315520.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2628313302993774,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0646450519561768,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 214450800.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 72638560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 106001744.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 955597504.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1409365504.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 446315520.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2673367261886597,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.064371347427368,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 214450800.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 243809.009557449,
+      "step": 158,
+      "teaching_rms": 172622.12655340284,
+      "traffic_rms": 243809.00955639352
+    },
+    {
+      "batch_loss": 3.570925235748291,
+      "innovation_rms": 113449.6749449934,
+      "instruction_rms": 0.8813763439517646,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1220240384.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.581189086259446e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 366182464.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4304286241531372,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.67597234249115,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 251069040.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 80575680.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 79371216.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1096299520.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1407019520.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 366182464.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4349335432052612,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.675967812538147,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 251069040.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 199146.48190649357,
+      "step": 159,
+      "teaching_rms": 113449.6749449934,
+      "traffic_rms": 199146.48190454318
+    },
+    {
+      "batch_loss": 4.348639011383057,
+      "innovation_rms": 179365.99715837569,
+      "instruction_rms": 1.289155181956319,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1353953408.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.3342927037569434e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 524834400.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5439229011535645,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.09785532951355,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 284995680.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 90479960.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 99042800.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1260532608.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1642331264.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 524834400.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5484273433685303,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0977611541748047,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 284995680.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 265897.19737213774,
+      "step": 160,
+      "teaching_rms": 179365.99715837569,
+      "traffic_rms": 265897.1973690127
+    },
+    {
+      "batch_loss": 5.017866611480713,
+      "innovation_rms": 214704.22149254772,
+      "instruction_rms": 1.075081455871528,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1591414784.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.769332314963968e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 419435904.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6133723258972168,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6132946014404297,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 307182752.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 101950208.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 114702448.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1449952896.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1894202624.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 419435904.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6178762912750244,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6132004261016846,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 307182752.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 377355.64373669896,
+      "step": 161,
+      "teaching_rms": 214704.22149254772,
+      "traffic_rms": 377355.64373516745
+    },
+    {
+      "batch_loss": 5.641810417175293,
+      "innovation_rms": 227101.21708148337,
+      "instruction_rms": 0.765951304671334,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1851379456.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0046905009108419e+18,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 429646816.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6605955362319946,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3211666345596313,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 321366112.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 113926464.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 119762552.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1657462272.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2075093888.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 429646816.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.665099024772644,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3210724592208862,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 321366112.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 438654.6107240368,
+      "step": 162,
+      "teaching_rms": 227101.21708148337,
+      "traffic_rms": 438654.61072336807
+    },
+    {
+      "batch_loss": 5.3410773277282715,
+      "innovation_rms": 287818.4870332555,
+      "instruction_rms": 0.739642273835525,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2111778816.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3970631137521828e+18,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 484209824.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6989010572433472,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1403881311416626,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 331401664.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 129006544.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 150800784.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1903331968.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2458696704.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 484209824.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.703404188156128,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1402939558029175,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 331401664.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 489192.4018612799,
+      "step": 163,
+      "teaching_rms": 287818.4870332555,
+      "traffic_rms": 489192.40186072077
+    },
+    {
+      "batch_loss": 6.7927751541137695,
+      "innovation_rms": 411458.68370164325,
+      "instruction_rms": 1.068226689076736,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2333336576.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7411681969263084e+18,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 792055232.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7326725721359253,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.53554105758667,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 363014720.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 155879200.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 268726592.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2249161728.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3458297856.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 792055232.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7371752262115479,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.535637855529785,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 363014720.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 530436.9635501202,
+      "step": 164,
+      "teaching_rms": 411458.68370164325,
+      "traffic_rms": 530436.9635490446
+    },
+    {
+      "batch_loss": 4.422138214111328,
+      "innovation_rms": 455889.12914921314,
+      "instruction_rms": 0.7757927356215367,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2547029504.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1351093318241485e+18,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 966302208.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7630314826965332,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7989866733551025,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 406000352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 192203344.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 363241408.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2688241408.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4390797312.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 966302208.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7675336599349976,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7990833520889282,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 406000352.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 597876.3958556128,
+      "step": 165,
+      "teaching_rms": 455889.12914921314,
+      "traffic_rms": 597876.3958551096
+    },
+    {
+      "batch_loss": 3.5410966873168945,
+      "innovation_rms": 465693.96065000346,
+      "instruction_rms": 0.9184202144357839,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3080942336.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.290092008416215e+18,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 788321152.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7903368473052979,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.397480845451355,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 461515616.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 230852976.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 386496352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3175646464.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4874049536.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 788321152.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.794838547706604,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3975775241851807,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 461515616.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 823789.799567072,
+      "step": 166,
+      "teaching_rms": 465693.96065000346,
+      "traffic_rms": 823789.79956656
+    },
+    {
+      "batch_loss": 3.9242472648620605,
+      "innovation_rms": 528349.9604818338,
+      "instruction_rms": 0.9906903298949374,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3704106496.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.771456053059518e+18,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 934764928.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8148937225341797,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2454994916915894,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 534828672.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 270600864.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 397478720.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3710667776.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5350214144.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 934764928.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8193949460983276,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2455363273620605,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 534828672.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 963988.4667920627,
+      "step": 167,
+      "teaching_rms": 528349.9604818338,
+      "traffic_rms": 963988.4667915538
+    },
+    {
+      "batch_loss": 3.4696388244628906,
+      "innovation_rms": 699437.5958149625,
+      "instruction_rms": 0.6770531889732131,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4229876992.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.612180354956526e+18,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1719664128.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8369767665863037,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.036548137664795,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 654951936.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 319363584.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 487627296.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4356820992.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6461531136.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1719664128.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8414775133132935,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0362416505813599,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 654951936.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1092938.9003840208,
+      "step": 168,
+      "teaching_rms": 699437.5958149625,
+      "traffic_rms": 1092938.900383811
+    },
+    {
+      "batch_loss": 3.864351987838745,
+      "innovation_rms": 782836.476384801,
+      "instruction_rms": 0.5569073316531125,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4366181888.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.646815848891417e+18,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3108187136.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8568331003189087,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7010958194732666,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 820532352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 376448832.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 570852544.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5127845888.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7710249984.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3108187136.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8613334894180298,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7009638547897339,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 820532352.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1064458.4129822347,
+      "step": 169,
+      "teaching_rms": 782836.476384801,
+      "traffic_rms": 1064458.4129820892
+    },
+    {
+      "batch_loss": 3.7031941413879395,
+      "innovation_rms": 826187.8521383369,
+      "instruction_rms": 0.6607697056298107,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5629796352.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.169922923053659e+19,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2103923968.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.874685287475586,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5425182580947876,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1011561216.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 434059136.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 576103104.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5966742528.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8388965376.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2103923968.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8791851997375488,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5424242615699768,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1011561216.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1539287.7657199271,
+      "step": 170,
+      "teaching_rms": 826187.8521383369,
+      "traffic_rms": 1539287.7657197854
+    },
+    {
+      "batch_loss": 3.3824687004089355,
+      "innovation_rms": 866715.7019456401,
+      "instruction_rms": 0.7176838178269197,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7370009600.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6854259810550415e+19,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2533628928.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8907334804534912,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8646517992019653,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1175828352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 485805280.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 517461440.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6840427008.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8736843776.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2533628928.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.895232915878296,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8648638725280762,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1175828352.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1774049.4480587998,
+      "step": 171,
+      "teaching_rms": 866715.7019456401,
+      "traffic_rms": 1774049.4480586548
+    },
+    {
+      "batch_loss": 2.7276573181152344,
+      "innovation_rms": 1000846.9983585689,
+      "instruction_rms": 0.47602002757716144,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9182450688.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3032798277070225e+19,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3278541312.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9051579236984253,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0117155313491821,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1410176640.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 531581664.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 457763680.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7766657024.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9262302208.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3278541312.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9096570014953613,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0116872787475586,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1410176640.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1980801.4645777275,
+      "step": 172,
+      "teaching_rms": 1000846.9983585689,
+      "traffic_rms": 1980801.4645776702
+    },
+    {
+      "batch_loss": 2.780879497528076,
+      "innovation_rms": 1276140.1648976463,
+      "instruction_rms": 0.5442700226342261,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10682336256.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.0356805758383817e+19,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3962405376.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9181208610534668,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8837785124778748,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1783849344.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 582819008.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 512373504.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8834797568.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10681400320.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3962405376.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9226195812225342,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8836311101913452,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1783849344.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2179844.8005258143,
+      "step": 173,
+      "teaching_rms": 1276140.1648976463,
+      "traffic_rms": 2179844.8005257463
+    },
+    {
+      "batch_loss": 3.2321841716766357,
+      "innovation_rms": 1837860.5406107956,
+      "instruction_rms": 0.5413722154492605,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11081073664.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.765541891914229e+19,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3298635008.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9297683238983154,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9395666122436523,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1868238848.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 687963136.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1051441216.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10330745856.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14959486976.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3298635008.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9342666864395142,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9397211074829102,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1868238848.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2359442.5461733886,
+      "step": 174,
+      "teaching_rms": 1837860.5406107956,
+      "traffic_rms": 2359442.546173326
+    },
+    {
+      "batch_loss": 3.0960404872894287,
+      "innovation_rms": 1639131.3289531595,
+      "instruction_rms": 0.5215618101345679,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11132831744.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.229489619347335e+19,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7828737024.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9402318000793457,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6358999013900757,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2287434240.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 793134144.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1051709760.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12006164480.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16754182144.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7828737024.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9447296857833862,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6360543966293335,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2287434240.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2367825.432249436,
+      "step": 175,
+      "teaching_rms": 1639131.3289531595,
+      "traffic_rms": 2367825.4322493784
+    },
+    {
+      "batch_loss": 3.183842658996582,
+      "innovation_rms": 1854092.854128519,
+      "instruction_rms": 0.592428678519642,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14164254720.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.831753818533213e+19,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5052045824.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.949629545211792,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6788562536239624,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2648185600.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 892030080.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 988959360.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13805940736.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17997758464.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5052045824.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9541269540786743,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6787090301513672,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2648185600.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3362835.6307217614,
+      "step": 176,
+      "teaching_rms": 1854092.854128519,
+      "traffic_rms": 3362835.6307217097
+    },
+    {
+      "batch_loss": 3.358607053756714,
+      "innovation_rms": 1943565.1828939416,
+      "instruction_rms": 0.5859087333442974,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18306666496.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.801386842791608e+19,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3440061952.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9580680131912231,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5680499076843262,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2962589184.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 971111232.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 790811200.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15665970176.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18600294400.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3440061952.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9625649452209473,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.567902684211731,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2962589184.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3834982.957368763,
+      "step": 177,
+      "teaching_rms": 1943565.1828939416,
+      "traffic_rms": 3834982.9573687185
+    },
+    {
+      "batch_loss": 3.2987051010131836,
+      "innovation_rms": 2305951.4419510947,
+      "instruction_rms": 0.7116027289342894,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22609078272.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.11349319355666e+20,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5315692544.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.965643048286438,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4545413255691528,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3230385920.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1039109120.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 679979072.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17664575488.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19986042880.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5315692544.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.970139503479004,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4543942213058472,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3230385920.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4382184.041995555,
+      "step": 178,
+      "teaching_rms": 2305951.4419510947,
+      "traffic_rms": 4382184.041995498
+    },
+    {
+      "batch_loss": 3.1356048583984375,
+      "innovation_rms": 3023313.210343838,
+      "instruction_rms": 0.9773978770207608,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26276284416.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.417733514539953e+20,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7765560320.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9724409580230713,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0622755289077759,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3472760320.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1131137664.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 920284800.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20063473664.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23988973568.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7765560320.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.976936936378479,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0625121593475342,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3472760320.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 5090851.24047124,
+      "step": 179,
+      "teaching_rms": 3023313.210343838,
+      "traffic_rms": 5090851.240471146
+    },
+    {
+      "batch_loss": 3.033240795135498,
+      "innovation_rms": 4184381.306149236,
+      "instruction_rms": 1.048884237130945,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28109981696.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9891407354802576e+20,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8103837696.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9785393476486206,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0205912590026855,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4283144192.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1344366592.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2132289280.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23411337216.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33478633472.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8103837696.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9830348491668701,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0208278894424438,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4283144192.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 5832898.591771189,
+      "step": 180,
+      "teaching_rms": 4184381.306149236,
+      "traffic_rms": 5832898.591771095
+    },
+    {
+      "batch_loss": 2.666253089904785,
+      "innovation_rms": 3152821.682957248,
+      "instruction_rms": 0.9227439257829219,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29112506368.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1628834518390958e+20,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8643696640.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9840080738067627,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8666769862174988,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5147513856.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1537025280.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1926586368.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26935027712.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35236896768.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8643696640.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.988503098487854,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8669135570526123,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5147513856.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 5008168.445878427,
+      "step": 181,
+      "teaching_rms": 3152821.682957248,
+      "traffic_rms": 5008168.445878341
+    },
+    {
+      "batch_loss": 2.5911030769348145,
+      "innovation_rms": 5349940.230447374,
+      "instruction_rms": 0.08459414992674905,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33714173952.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.942697560285737e+20,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13084850176.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9889100790023804,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7674687504768372,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5387266560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1841021952.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3039967232.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31492526080.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 45574979584.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13084850176.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9934046268463135,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7677052617073059,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5387266560.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 7727985.097877432,
+      "step": 182,
+      "teaching_rms": 5349940.230447374,
+      "traffic_rms": 7727985.097877431
+    },
+    {
+      "batch_loss": 3.487748384475708,
+      "innovation_rms": 5385400.352739672,
+      "instruction_rms": 0.9628475250107043,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38401032192.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.387199521032734e+20,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10867269632.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9933019876480103,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6924802660942078,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6285279744.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2200202240.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3591801856.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36705738752.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 52132110336.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10867269632.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9977960586547852,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6927167773246765,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6285279744.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 9045234.522533907,
+      "step": 183,
+      "teaching_rms": 5385400.352739672,
+      "traffic_rms": 9045234.522533856
+    },
+    {
+      "batch_loss": 3.193814754486084,
+      "innovation_rms": 7059924.17957883,
+      "instruction_rms": 0.8703194653298297,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 44566437888.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.323255835123966e+20,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14823329792.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.99723482131958,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6298978924751282,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7449515008.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2685039104.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4848367616.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 43161669632.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 64559316992.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14823329792.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0017282962799072,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6300777792930603,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7449515008.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 10601266.562844766,
+      "step": 184,
+      "teaching_rms": 7059924.17957883,
+      "traffic_rms": 10601266.56284473
+    },
+    {
+      "batch_loss": 2.975266695022583,
+      "innovation_rms": 6824557.031294944,
+      "instruction_rms": 0.6273490481900119,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 51652931584.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.708210533094801e+20,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16160788480.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.000800609588623,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7555835247039795,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8279827456.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3197897728.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5128585728.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 50277502976.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 71158341632.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16160788480.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.005293607711792,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7555705904960632,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8279827456.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 10904814.05638237,
+      "step": 185,
+      "teaching_rms": 6824557.031294944,
+      "traffic_rms": 10904814.05638235
+    },
+    {
+      "batch_loss": 4.172522068023682,
+      "innovation_rms": 9994342.180518318,
+      "instruction_rms": 0.8323262494865011,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 59626323968.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0152753151832438e+21,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30291210240.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.003061532974243,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6753022074699402,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10378665984.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3909737984.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7118401536.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 59377172480.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 90996686848.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30291210240.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.007554054260254,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.675193190574646,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10378665984.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 13689529.516119162,
+      "step": 186,
+      "teaching_rms": 9994342.180518318,
+      "traffic_rms": 13689529.516119136
+    },
+    {
+      "batch_loss": 3.2824931144714355,
+      "innovation_rms": 8009353.068320669,
+      "instruction_rms": 0.8209764735341593,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 65386291200.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.501141094432529e+21,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21007974400.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0034351348876953,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7182347774505615,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12479463424.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4554801664.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6450637312.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 68562030592.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 91848589312.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21007974400.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0133097171783447,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7183926105499268,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12479463424.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 17346901.810050577,
+      "step": 187,
+      "teaching_rms": 8009353.068320669,
+      "traffic_rms": 17346901.81005056
+    },
+    {
+      "batch_loss": 3.3454737663269043,
+      "innovation_rms": 8830598.011431491,
+      "instruction_rms": 1.206234155503292,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 74245554176.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2620136961651476e+21,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19391252480.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.99970543384552,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5881356596946716,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14156803072.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5132328448.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5775265280.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 77689208832.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 91271782400.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19391252480.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.047192335128784,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5880512595176697,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14156803072.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 21383643.090844147,
+      "step": 188,
+      "teaching_rms": 8830598.011431491,
+      "traffic_rms": 21383643.090844113
+    },
+    {
+      "batch_loss": 4.9404168128967285,
+      "innovation_rms": 11936563.514105875,
+      "instruction_rms": 1.8044368058063383,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 88868880384.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.2781549691142417e+21,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36298731520.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9932845830917358,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6639219522476196,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15959463936.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5777604096.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6452754432.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 87738408960.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 100491968512.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36298731520.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.075835704803467,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.663790225982666,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15959463936.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 24724676.02487237,
+      "step": 189,
+      "teaching_rms": 11936563.514105875,
+      "traffic_rms": 24724676.024872303
+    },
+    {
+      "batch_loss": 4.24451208114624,
+      "innovation_rms": 16767581.233488455,
+      "instruction_rms": 1.5033725386466987,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 110583627776.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.2693043603561733e+21,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39255060480.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9849169254302979,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6822370290756226,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17622470656.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6676951040.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8993470464.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 100811710464.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 130732982272.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39255060480.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.101811170578003,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.682342529296875,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17622470656.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 26015085.11688197,
+      "step": 190,
+      "teaching_rms": 16767581.233488455,
+      "traffic_rms": 26015085.116881926
+    },
+    {
+      "batch_loss": 5.0928521156311035,
+      "innovation_rms": 16257463.557894828,
+      "instruction_rms": 1.6930290484747048,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 124154732544.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.876201564841273e+21,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 46831538176.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.978594422340393,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2720658779144287,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19928580096.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7730395648.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10534448128.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 115899957248.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 150882484224.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 46831538176.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1255149841308594,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2718631029129028,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19928580096.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 24294626.5249897,
+      "step": 191,
+      "teaching_rms": 16257463.557894828,
+      "traffic_rms": 24294626.524989642
+    },
+    {
+      "batch_loss": 5.231077671051025,
+      "innovation_rms": 24636969.7304788,
+      "instruction_rms": 1.2106342784870598,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 137142788096.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.071801431316897e+21,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 47883071488.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0003137588500977,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9135897159576416,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22653597696.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9534157824.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18037618688.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 136913764352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 210138038272.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 47883071488.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.147219657897949,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9137012362480164,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22653597696.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 30992277.383593526,
+      "step": 192,
+      "teaching_rms": 24636969.7304788,
+      "traffic_rms": 30992277.383593503
+    },
+    {
+      "batch_loss": 4.52946138381958,
+      "innovation_rms": 24867158.516558565,
+      "instruction_rms": 1.3983948560794175,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 170926931968.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.813108747500113e+21,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 42560036864.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0202906131744385,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7927453517913818,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25231923200.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11524419584.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19902621696.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 161333690368.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 244199243776.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 42560036864.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.167181968688965,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7927974462509155,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25231923200.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 41128453.82522598,
+      "step": 193,
+      "teaching_rms": 24867158.516558565,
+      "traffic_rms": 41128453.82522595
+    },
+    {
+      "batch_loss": 5.088406085968018,
+      "innovation_rms": 32712636.572059043,
+      "instruction_rms": 1.2211526370922416,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 202799562752.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2043249152137097e+22,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 85570158592.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0386130809783936,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2534949779510498,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30249666560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14051542016.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25271228416.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 192187318272.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 308536246272.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 85570158592.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1854896545410156,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2535470724105835,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30249666560.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 46337049.14856257,
+      "step": 194,
+      "teaching_rms": 32712636.572059043,
+      "traffic_rms": 46337049.14856256
+    },
+    {
+      "batch_loss": 5.750738620758057,
+      "innovation_rms": 15885394.598761236,
+      "instruction_rms": 1.656600509221617,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 188111814656.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2619133443688217e+22,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 62310109184.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0556488037109375,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.262545108795166,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36480675840.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15446736896.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13951952896.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 220511338496.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 283240136704.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 62310109184.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2025105953216553,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.262605905532837,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36480675840.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 36045356.922179796,
+      "step": 195,
+      "teaching_rms": 15885394.598761236,
+      "traffic_rms": 36045356.92217976
+    },
+    {
+      "batch_loss": 3.81648588180542,
+      "innovation_rms": 17602028.989314694,
+      "instruction_rms": 1.6719932959800894,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 173791723520.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3336100874866066e+22,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 77986217984.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.072373151779175,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1279652118682861,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41594159104.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15657697280.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2109601792.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 246578167808.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 260668342272.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 77986217984.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2192203998565674,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1276495456695557,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 41594159104.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 40637875.93319327,
+      "step": 196,
+      "teaching_rms": 17602028.989314694,
+      "traffic_rms": 40637875.93319323
+    },
+    {
+      "batch_loss": 4.054687023162842,
+      "innovation_rms": 45298553.0331185,
+      "instruction_rms": 1.5260000616572307,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 198643974144.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8592155172578183e+22,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 238991409152.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0823888778686523,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2834362983703613,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 56911917056.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16864366592.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12066689024.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 281178898432.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 346007437312.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 238991409152.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2292213439941406,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.283200740814209,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 56911917056.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 63949722.60212393,
+      "step": 197,
+      "teaching_rms": 45298553.0331185,
+      "traffic_rms": 63949722.602123916
+    },
+    {
+      "batch_loss": 4.341590881347656,
+      "innovation_rms": 39127756.57554863,
+      "instruction_rms": 1.418881298326907,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 282567737344.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.2274763002774135e+22,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 214356131840.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0910089015960693,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4758448600769043,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 78347534336.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17820735488.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9563698176.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 316693053440.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 355141386240.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 214356131840.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2378265857696533,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.475609302520752,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 78347534336.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 90574042.55420594,
+      "step": 198,
+      "teaching_rms": 39127756.57554863,
+      "traffic_rms": 90574042.55420592
+    },
+    {
+      "batch_loss": 3.7518668174743652,
+      "innovation_rms": 38228908.6359847,
+      "instruction_rms": 1.739536897768505,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 401751244800.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.490756239772792e+22,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 201354985472.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.116183280944824,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4364571571350098,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 98483036160.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18114312192.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2935772160.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 350890229760.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 341971861504.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 201354985472.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.254004955291748,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4362215995788574,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 98483036160.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 103269442.89339584,
+      "step": 199,
+      "teaching_rms": 38228908.6359847,
+      "traffic_rms": 103269442.89339583
+    },
+    {
+      "batch_loss": 3.738774299621582,
+      "innovation_rms": 41663675.129963316,
+      "instruction_rms": 2.4203917108476563,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 541450076160.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0061147852496894e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 193856585728.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.154536247253418,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5486483573913574,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 117868691456.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17667911680.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4464006144.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 383894814720.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 330045915136.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 193856585728.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.27488374710083,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5483322143554688,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 117868691456.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 114074475.80079384,
+      "step": 200,
+      "teaching_rms": 41663675.129963316,
+      "traffic_rms": 114074475.80079381
+    },
+    {
+      "batch_loss": 3.021869421005249,
+      "innovation_rms": 46903814.07606112,
+      "instruction_rms": 2.5075527093860015,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 690914131968.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5107418624886578e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 190847074304.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1976661682128906,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2359204292297363,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 136953397248.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16551118848.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11167927296.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 416663240704.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 327684358144.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 190847074304.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2663207054138184,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.235748291015625,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 136953397248.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 123890177.102542,
+      "step": 201,
+      "teaching_rms": 46903814.07606112,
+      "traffic_rms": 123890177.10254197
+    },
+    {
+      "batch_loss": 3.0605738162994385,
+      "innovation_rms": 54590241.16327603,
+      "instruction_rms": 2.2843633565306516,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 833750237184.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9346138139453353e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 193077477376.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2601377964019775,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.291755199432373,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 156261138432.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15019890688.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15312281600.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 451242786816.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 345795461120.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 193077477376.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.259310245513916,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2915830612182617,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 156261138432.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 132589999.0550902,
+      "step": 202,
+      "teaching_rms": 54590241.16327603,
+      "traffic_rms": 132589999.05509017
+    },
+    {
+      "batch_loss": 3.336050271987915,
+      "innovation_rms": 65142996.12737751,
+      "instruction_rms": 1.9729906281450607,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 966279823360.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2569479684912878e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 200165146624.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.294661521911621,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0847113132476807,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 176277651456.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13507064832.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15128254464.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 491026055168.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 397832683520.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 200165146624.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.261547327041626,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0845390558242798,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 176277651456.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 139162324.08998427,
+      "step": 203,
+      "teaching_rms": 65142996.12737751,
+      "traffic_rms": 139162324.08998427
+    },
+    {
+      "batch_loss": 2.9007728099823,
+      "innovation_rms": 80151796.5864981,
+      "instruction_rms": 1.4906833251598925,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1074488475648.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4290746455294626e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 202549379072.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3238792419433594,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9746956825256348,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 196532584448.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12806137856.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7009264640.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 541586489344.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 505604210688.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 202549379072.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.284952163696289,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9745234847068787,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 196532584448.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 143035147.98475286,
+      "step": 204,
+      "teaching_rms": 80151796.5864981,
+      "traffic_rms": 143035147.98475283
+    },
+    {
+      "batch_loss": 3.026761770248413,
+      "innovation_rms": 98324361.5411293,
+      "instruction_rms": 1.030194259206255,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1146093174784.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.507554732923991e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 218282393600.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.358031988143921,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.878250241279602,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 214853238784.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13879653376.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10735155200.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 609393508352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 678070059008.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 218282393600.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.347010374069214,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8780781030654907,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 214853238784.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 150314722.43362042,
+      "step": 205,
+      "teaching_rms": 98324361.5411293,
+      "traffic_rms": 150314722.43362042
+    },
+    {
+      "batch_loss": 3.2664833068847656,
+      "innovation_rms": 116456812.16440998,
+      "instruction_rms": 0.8056251855062885,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1152166264832.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4822547712253193e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 192981958656.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3909358978271484,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8498039245605469,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 233767927808.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17704681472.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38250283008.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 700890021888.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 914964938752.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 192981958656.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.431976079940796,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8496570587158203,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 233767927808.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 155999790.21255654,
+      "step": 206,
+      "teaching_rms": 116456812.16440998,
+      "traffic_rms": 155999790.21255654
+    },
+    {
+      "batch_loss": 4.131566047668457,
+      "innovation_rms": 129515008.559384,
+      "instruction_rms": 0.9844689342052526,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1128488370176.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4745745125647868e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 278365339648.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3803634643554688,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0039820671081543,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 261585551360.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24703475712.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 69987934208.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 816683876352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1157938544640.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 278365339648.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3315632343292236,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0041289329528809,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 261585551360.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 166940534.20616004,
+      "step": 207,
+      "teaching_rms": 129515008.559384,
+      "traffic_rms": 166940534.20616004
+    },
+    {
+      "batch_loss": 4.7382893562316895,
+      "innovation_rms": 156132791.49435118,
+      "instruction_rms": 0.936033430309457,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1256290254848.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.185346656987236e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 220496216064.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4170033931732178,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1190948486328125,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 283635187712.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34245644288.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 95421685760.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 959500582912.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1428167065600.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 220496216064.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.365377426147461,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1190662384033203,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 283635187712.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 231662012.50849628,
+      "step": 208,
+      "teaching_rms": 156132791.49435118,
+      "traffic_rms": 231662012.50849628
+    },
+    {
+      "batch_loss": 3.552975654602051,
+      "innovation_rms": 188138321.66036478,
+      "instruction_rms": 1.060792106346288,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1415911309312.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.240110586419677e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 404839596032.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.468893527984619,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5299534797668457,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 295992197120.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 47207522304.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 129618788352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1136624140288.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1771236098048.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 404839596032.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4233667850494385,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5299248695373535,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 295992197120.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 270174230.855016,
+      "step": 209,
+      "teaching_rms": 188138321.66036478,
+      "traffic_rms": 270174230.855016
+    },
+    {
+      "batch_loss": 4.918725490570068,
+      "innovation_rms": 207525090.98284486,
+      "instruction_rms": 1.6266448324218774,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1440197640192.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.939728215285055e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 631915085824.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.298830509185791,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8674957752227783,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 351615090688.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63238500352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 160309772288.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1347975380992.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2113512013824.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 631915085824.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4705870151519775,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8672599792480469,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 351615090688.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 272752267.26072323,
+      "step": 210,
+      "teaching_rms": 207525090.98284486,
+      "traffic_rms": 272752267.26072323
+    },
+    {
+      "batch_loss": 5.358460903167725,
+      "innovation_rms": 222597754.72368473,
+      "instruction_rms": 1.5148088609686745,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1762747482112.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.073606496903174e+23,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 399326347264.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.342219352722168,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3552141189575195,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 390164480000.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 79844655104.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 166061572096.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1579416289280.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2314409082880.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 399326347264.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.513958692550659,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3549782037734985,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 390164480000.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 376858735.75738055,
+      "step": 211,
+      "teaching_rms": 222597754.72368473,
+      "traffic_rms": 376858735.75738055
+    },
+    {
+      "batch_loss": 5.824455261230469,
+      "innovation_rms": 245852708.54526594,
+      "instruction_rms": 1.1775537243599628,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2215445528576.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0157448143184972e+24,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 545092599808.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3824431896209717,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9171805381774902,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 415652970496.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 96088784896.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 162441297920.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1827608985600.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2481927487488.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 545092599808.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5541651248931885,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.916944682598114,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 415652970496.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 442788130.03626543,
+      "step": 212,
+      "teaching_rms": 245852708.54526594,
+      "traffic_rms": 442788130.03626543
+    },
+    {
+      "batch_loss": 5.206492900848389,
+      "innovation_rms": 307374368.5939395,
+      "instruction_rms": 1.2893712286703851,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2677765308416.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.434744087151142e+24,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 691009945600.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4198477268218994,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9422031044960022,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 421377048576.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 114075402240.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 179866206208.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2110188290048.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2825792520192.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 691009945600.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.591552495956421,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.942101240158081,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 421377048576.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 521905602.09523714,
+      "step": 213,
+      "teaching_rms": 307374368.5939395,
+      "traffic_rms": 521905602.09523714
+    },
+    {
+      "batch_loss": 5.425127029418945,
+      "innovation_rms": 441330087.720237,
+      "instruction_rms": 1.6250140314551373,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3006495719424.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.023919834268122e+24,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1139445989376.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4556350708007812,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9072262048721313,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 384345636864.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 142734950400.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 286595416064.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2489713950720.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3795255820288.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1139445989376.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6273226737976074,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9074636697769165,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 384345636864.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 616335334.4373357,
+      "step": 214,
+      "teaching_rms": 441330087.720237,
+      "traffic_rms": 616335334.4373357
+    },
+    {
+      "batch_loss": 6.03110408782959,
+      "innovation_rms": 226385076.6840312,
+      "instruction_rms": 1.44924311206578,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2951317553152.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1115859678657936e+24,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 891118747648.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4892520904541016,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2239842414855957,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 445358473216.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 158949883904.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 162149400576.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2845990977536.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3562769743872.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 891118747648.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6609225273132324,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2240798473358154,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 445358473216.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 456821305.86181784,
+      "step": 215,
+      "teaching_rms": 226385076.6840312,
+      "traffic_rms": 456821305.86181784
+    },
+    {
+      "batch_loss": 5.717033863067627,
+      "innovation_rms": 229560637.88555336,
+      "instruction_rms": 1.3175253908144784,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2974206918656.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5320597552452365e+24,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 871262715904.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5206832885742188,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6640596389770508,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 524440535040.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 160762085376.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18122047488.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3175949795328.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3299587129344.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 871262715904.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6923365592956543,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6641552448272705,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 524440535040.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 504755644.5746933,
+      "step": 216,
+      "teaching_rms": 229560637.88555336,
+      "traffic_rms": 504755644.5746933
+    },
+    {
+      "batch_loss": 7.14595890045166,
+      "innovation_rms": 487893738.03183144,
+      "instruction_rms": 1.3640218760862832,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3386848575488.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.875957505095009e+24,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2602090627072.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5526130199432373,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3095171451568604,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 759724834816.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 167252312064.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 64902320128.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3567883124736.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3919333818368.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2602090627072.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7242491245269775,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3092572689056396,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 759724834816.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 710466064.847457,
+      "step": 217,
+      "teaching_rms": 487893738.03183144,
+      "traffic_rms": 710466064.847457
+    },
+    {
+      "batch_loss": 4.9911017417907715,
+      "innovation_rms": 584674143.3081429,
+      "instruction_rms": 1.759347317374562,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4687969910784.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.344158379351953e+24,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2268621176832.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5824036598205566,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4874950647354126,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 986586939392.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 178111004672.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 108586876928.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4026093273088.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4582101483520.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2268621176832.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7540225982666016,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4876835346221924,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 986586939392.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1035519911.7689819,
+      "step": 218,
+      "teaching_rms": 584674143.3081429,
+      "traffic_rms": 1035519911.7689819
+    },
+    {
+      "batch_loss": 7.716791152954102,
+      "innovation_rms": 560520537.1047194,
+      "instruction_rms": 2.7319532749223825,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6369966030848.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1359530364863443e+25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2100424474624.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.521545886993408,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5968016386032104,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1196629426176.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 185463898112.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 73528983552.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4503733010432.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4776394752000.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2100424474624.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.693147659301758,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5968942642211914,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1196629426176.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1198687187.5971336,
+      "step": 219,
+      "teaching_rms": 560520537.1047194,
+      "traffic_rms": 1198687187.5971336
+    },
+    {
+      "batch_loss": 7.0226545333862305,
+      "innovation_rms": 595485450.2133509,
+      "instruction_rms": 2.338924540925135,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8194112356352.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9154407293237234e+25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2061152550912.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4691197872161865,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3548108339309692,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1402744733696.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 186740998144.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12770930688.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4992831848448.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4890986283008.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2061152550912.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.640704393386841,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3549035787582397,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1402744733696.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1327633475.6317885,
+      "step": 220,
+      "teaching_rms": 595485450.2133509,
+      "traffic_rms": 1327633475.6317885
+    },
+    {
+      "batch_loss": 5.864837646484375,
+      "innovation_rms": 648489569.3373861,
+      "instruction_rms": 2.237917262822703,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9983195348992.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.636520726984817e+25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2064939220992.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4257047176361084,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.091233730316162,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1609238708224.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 181357756416.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 53832355840.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5494507831296.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5016758255616.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2064939220992.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5972721576690674,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0913264751434326,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1609238708224.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1449264178.3518083,
+      "step": 221,
+      "teaching_rms": 648489569.3373861,
+      "traffic_rms": 1449264178.3518083
+    },
+    {
+      "batch_loss": 5.3016581535339355,
+      "innovation_rms": 758488517.3479275,
+      "instruction_rms": 2.4690656486878515,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11568592977920.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.167196891081594e+25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2180349034496.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.390669107437134,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.083780288696289,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1827273572352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 173329317888.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 80284434432.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6034747293696.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5402392526848.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2180349034496.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5622193813323975,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0838854312896729,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1827273572352.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1581315731.716637,
+      "step": 222,
+      "teaching_rms": 758488517.3479275,
+      "traffic_rms": 1581315731.716637
+    },
+    {
+      "batch_loss": 4.847405910491943,
+      "innovation_rms": 940600854.1860993,
+      "instruction_rms": 2.3306674410397274,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12792326258688.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.3983276767961472e+25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2519111565312.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.363691568374634,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7287805080413818,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2079184781312.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 171615731712.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17135886336.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6672962027520.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6382145241088.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2519111565312.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.535224676132202,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7288856506347656,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2079184781312.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1738331282.182063,
+      "step": 223,
+      "teaching_rms": 940600854.1860993,
+      "traffic_rms": 1738331282.182063
+    },
+    {
+      "batch_loss": 4.7808918952941895,
+      "innovation_rms": 1303226811.6690555,
+      "instruction_rms": 1.642452385483869,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13458272681984.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.312033346850732e+25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3104798932992.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3477394580841064,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8540975451469421,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2389664595968.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 201543843840.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 299281088512.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7582056972288.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9090950496256.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3104798932992.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5192553997039795,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.854192852973938,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2389664595968.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1869848596.660936,
+      "step": 224,
+      "teaching_rms": 1303226811.6690555,
+      "traffic_rms": 1869848596.660936
+    },
+    {
+      "batch_loss": 4.077389240264893,
+      "innovation_rms": 846289945.8163916,
+      "instruction_rms": 1.515490124414656,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13133849559040.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.3085040235408296e+25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1815232118784.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3366169929504395,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8060410022735596,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2571187781632.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 212123893760.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 105800433664.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8493738754048.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9116819914752.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1815232118784.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.508115768432617,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8059633374214172,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2571187781632.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1479678324.3055472,
+      "step": 225,
+      "teaching_rms": 846289945.8163916,
+      "traffic_rms": 1479678324.3055472
+    },
+    {
+      "batch_loss": 3.469071626663208,
+      "innovation_rms": 1314225616.185107,
+      "instruction_rms": 1.2017972565719879,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13409836859392.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.2835188310301937e+25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4436018069504.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.334935426712036,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.047567367553711,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2664480899072.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 231653916672.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 195300147200.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9570532458496.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10767935995904.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4436018069504.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5064170360565186,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.047675609588623,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2664480899072.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1921578738.906363,
+      "step": 226,
+      "teaching_rms": 1314225616.185107,
+      "traffic_rms": 1921578738.906363
+    },
+    {
+      "batch_loss": 3.945005178451538,
+      "innovation_rms": 1490178322.1209784,
+      "instruction_rms": 0.9803442082337497,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16251573239808.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.624310970071435e+25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2985177120768.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.347898483276367,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9918085336685181,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2804757037056.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 258458796032.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 268048711680.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10784138592256.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12136064483328.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2985177120768.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5193629264831543,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9919259548187256,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2804757037056.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2751997945.6873674,
+      "step": 227,
+      "teaching_rms": 1490178322.1209784,
+      "traffic_rms": 2751997945.6873674
+    },
+    {
+      "batch_loss": 3.2627031803131104,
+      "innovation_rms": 1455639730.060778,
+      "instruction_rms": 0.9741723529718793,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20557170999296.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.405784669391858e+25,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2814574592000.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3702094554901123,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5308635234832764,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2960443834368.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 275014418432.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 165556125696.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12036856610816.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12527180185600.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2814574592000.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.541656732559204,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5307259559631348,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2960443834368.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3114246558.6030393,
+      "step": 228,
+      "teaching_rms": 1455639730.060778,
+      "traffic_rms": 3114246558.6030393
+    },
+    {
+      "batch_loss": 3.193086624145508,
+      "innovation_rms": 1550990393.1699104,
+      "instruction_rms": 1.0617513709351747,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25465077104640.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3454439928448889e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2792428666880.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.398181676864624,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3120445013046265,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3122416058368.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 276072562688.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10581508096.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13311054708736.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12741984124928.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2792428666880.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5696117877960205,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3119069337844849,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3122416058368.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3451602518.864972,
+      "step": 229,
+      "teaching_rms": 1550990393.1699104,
+      "traffic_rms": 3451602518.864972
+    },
+    {
+      "batch_loss": 3.31718111038208,
+      "innovation_rms": 1753757019.9883425,
+      "instruction_rms": 1.438335961668968,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30188180078592.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7784790041322114e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2962858967040.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.385268449783325,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1305094957351685,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3322444513280.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 264982249472.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 110903074816.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14641583357952.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13305281249280.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2962858967040.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5566813945770264,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1303718090057373,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3322444513280.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3748327779.222374,
+      "step": 230,
+      "teaching_rms": 1753757019.9883425,
+      "traffic_rms": 3748327779.222374
+    },
+    {
+      "batch_loss": 2.5810706615448,
+      "innovation_rms": 2070875654.0279431,
+      "instruction_rms": 1.303723829295474,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34592358662144.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0816899811461827e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3752753364992.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3576483726501465,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9941970705986023,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3609680674816.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 252546580480.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 124356632576.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16112540450816.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14709574074368.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3752753364992.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5290441513061523,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9940593838691711,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3609680674816.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4088549739.9937615,
+      "step": 231,
+      "teaching_rms": 2070875654.0279431,
+      "traffic_rms": 4088549739.9937615
+    },
+    {
+      "batch_loss": 3.202209949493408,
+      "innovation_rms": 2753276290.5968795,
+      "instruction_rms": 1.3630364078277226,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38153582804992.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.164982195727322e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6775646453760.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3684561252593994,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8767777681350708,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4119680778240.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 278985441280.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 264388640768.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17993211838464.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18806713876480.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6775646453760.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5192694664001465,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8766400814056396,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4119680778240.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4573489846.361334,
+      "step": 232,
+      "teaching_rms": 2753276290.5968795,
+      "traffic_rms": 4573489846.361334
+    },
+    {
+      "batch_loss": 3.475125789642334,
+      "innovation_rms": 3363004942.3412833,
+      "instruction_rms": 1.5016590750411138,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 39394069184512.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0804929719232396e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5783198629888.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.376692056655884,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7453941106796265,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4127102337024.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 390934364160.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1119489097728.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20557066141696.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25638534643712.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5783198629888.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.519538640975952,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7452564239501953,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4127102337024.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4392476598.104629,
+      "step": 233,
+      "teaching_rms": 3363004942.3412833,
+      "traffic_rms": 4392476598.104629
+    },
+    {
+      "batch_loss": 3.118391513824463,
+      "innovation_rms": 3962030807.3712354,
+      "instruction_rms": 1.6886672929334587,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 42000732651520.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0950845309529846e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6241964261376.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.37907075881958,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6342218518257141,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4544054427648.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 581063475200.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1901291372544.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 23877231050752.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33201649090560.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6241964261376.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5174736976623535,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.634084165096283,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4544054427648.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 5624000342.025578,
+      "step": 234,
+      "teaching_rms": 3962030807.3712354,
+      "traffic_rms": 5624000342.025578
+    },
+    {
+      "batch_loss": 3.1699094772338867,
+      "innovation_rms": 4967882906.878688,
+      "instruction_rms": 1.573253249510787,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 42699273011200.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.632207232584341e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6199538876416.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.382962703704834,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6588848829269409,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4802149351424.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 909611302912.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3285478014976.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28308620181504.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 44313887113216.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6199538876416.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.478372812271118,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6591969132423401,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4802149351424.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 6538656565.7894745,
+      "step": 235,
+      "teaching_rms": 4967882906.878688,
+      "traffic_rms": 6538656565.7894745
+    },
+    {
+      "batch_loss": 2.5539095401763916,
+      "innovation_rms": 6014736214.074882,
+      "instruction_rms": 1.148592995681319,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 44066528034816.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.379943052133106e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10337249132544.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3880984783172607,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7623021602630615,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5320332017664.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1386436100096.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4768248365056.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 34086823395328.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 57782036332544.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10337249132544.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4789507389068604,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7626141905784607,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5320332017664.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 7766310830.083111,
+      "step": 236,
+      "teaching_rms": 6014736214.074882,
+      "traffic_rms": 7766310830.083111
+    },
+    {
+      "batch_loss": 2.86086368560791,
+      "innovation_rms": 6091665216.80286,
+      "instruction_rms": 0.9954042298233352,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 43465626877952.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.5608556461974e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10614615310336.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.394059181213379,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9205632209777832,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6068376698880.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1932108890112.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5456727638016.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 40711470710784.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 66246464765952.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 10614615310336.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5059704780578613,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9203174114227295,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6068376698880.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 8989441497.59849,
+      "step": 237,
+      "teaching_rms": 6091665216.80286,
+      "traffic_rms": 8989441497.59849
+    },
+    {
+      "batch_loss": 2.9549288749694824,
+      "innovation_rms": 8888086500.098581,
+      "instruction_rms": 1.0401976260446673,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 46615066836992.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.449117496205242e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21670616104960.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3971025943756104,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8288367390632629,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7858465275904.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2690718892032.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7586098708480.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 49376554647552.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 86650830979072.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21670616104960.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.529541254043579,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.828589677810669,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7858465275904.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 11733185792.062904,
+      "step": 238,
+      "teaching_rms": 8888086500.098581,
+      "traffic_rms": 11733185792.062904
+    },
+    {
+      "batch_loss": 3.2119646072387695,
+      "innovation_rms": 3116910612.6897373,
+      "instruction_rms": 1.4036500700095005,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 38974697504768.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.762650164398216e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19553060913152.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.398188829421997,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5884023904800415,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9813771681792.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3106289221632.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4155703820288.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 57138214862848.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 77616593764352.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19553060913152.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.548218011856079,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5883064270019531,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9813771681792.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 8558743733.271369,
+      "step": 239,
+      "teaching_rms": 3116910612.6897373,
+      "traffic_rms": 8558743733.271369
+    },
+    {
+      "batch_loss": 3.0140650272369385,
+      "innovation_rms": 2624853764.283045,
+      "instruction_rms": 1.6341747243614175,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 28260633149440.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.24502387733365e+26,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17782030729216.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4009175300598145,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6207819581031799,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11591975174144.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3096033886208.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 102553878528.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 63908870094848.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 67706569097216.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17782030729216.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.562734603881836,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6203612089157104,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11591975174144.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 10075668536.010647,
+      "step": 240,
+      "teaching_rms": 2624853764.283045,
+      "traffic_rms": 10075668536.010647
+    },
+    {
+      "batch_loss": 3.3371634483337402,
+      "innovation_rms": 2917226694.1963983,
+      "instruction_rms": 1.9739351559140146,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18224106176512.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3682888649808845e+27,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15568851697664.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4062659740448,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1894067525863647,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13148859924480.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2660704190464.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4353297743872.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 69758410031104.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 58495382585344.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15568851697664.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5732855796813965,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1891758441925049,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13148859924480.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 11724648630.30694,
+      "step": 241,
+      "teaching_rms": 2917226694.1963983,
+      "traffic_rms": 11724648630.30694
+    },
+    {
+      "batch_loss": 3.3384459018707275,
+      "innovation_rms": 3935517146.70016,
+      "instruction_rms": 1.8642606890493638,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21680680337408.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8833175323002775e+27,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13569796079616.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.409973621368408,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6968897581100464,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14422170927104.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1873495064576.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7872090996736.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 74861615513600.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 51032050630656.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13569796079616.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5804784297943115,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6971354484558105,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14422170927104.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 13866436209.484833,
+      "step": 242,
+      "teaching_rms": 3935517146.70016,
+      "traffic_rms": 13866436209.484833
+    },
+    {
+      "batch_loss": 3.817507743835449,
+      "innovation_rms": 6039998353.816764,
+      "instruction_rms": 1.736622532651552,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24659831554048.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3253367904810273e+27,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16984634818560.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3940722942352295,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.558039665222168,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15328381763584.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 894030905344.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9794641461248.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 79587866312704.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 47262495408128.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16984634818560.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5848405361175537,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5578789710998535,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15328381763584.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 17370676908.951775,
+      "step": 243,
+      "teaching_rms": 6039998353.816764,
+      "traffic_rms": 17370676908.951775
+    },
+    {
+      "batch_loss": 4.625663757324219,
+      "innovation_rms": 10457405808.47075,
+      "instruction_rms": 1.3219881525110715,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 26489122717696.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6246219314212317e+27,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33716422311936.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.381760358810425,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2489410638809204,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15838108188672.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 77256933376.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8167739621376.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 85125027069952.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 55371569823744.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33716422311936.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5969674587249756,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.248780369758606,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15838108188672.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 22381994876.774254,
+      "step": 244,
+      "teaching_rms": 10457405808.47075,
+      "traffic_rms": 22381994876.774254
+    },
+    {
+      "batch_loss": 4.555985927581787,
+      "innovation_rms": 14661185099.710457,
+      "instruction_rms": 1.2840761786753843,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 42101098151936.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.951038986266906e+27,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31836103245824.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.37414813041687,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9100722074508667,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17676541886464.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 443725185024.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5209821020160.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 93200119234560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 80750921646080.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31836103245824.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6074538230895996,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.909911572933197,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17676541886464.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 22962947594.1218,
+      "step": 245,
+      "teaching_rms": 14661185099.710457,
+      "traffic_rms": 22962947594.1218
+    },
+    {
+      "batch_loss": 4.641958236694336,
+      "innovation_rms": 16314338691.835798,
+      "instruction_rms": 1.1803972344983948,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 66879775309824.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.646538496539361e+27,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33412893114368.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.371702194213867,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9216773509979248,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21017831407616.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 661245132800.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2175199608832.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 104081980915712.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 108818650365952.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33412893114368.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.613609790802002,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9218560457229614,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21017831407616.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 24533699655.480564,
+      "step": 246,
+      "teaching_rms": 16314338691.835798,
+      "traffic_rms": 24533699655.480564
+    },
+    {
+      "batch_loss": 4.207598686218262,
+      "innovation_rms": 20958447787.57356,
+      "instruction_rms": 1.1725772225076232,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 77009929633792.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.408841488522205e+27,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36331329683456.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3738901615142822,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9747524261474609,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22953267822592.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 181666414592.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4795787116544.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 119912064352256.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 158300800811008.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 36331329683456.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6139633655548096,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9748608469963074,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22953267822592.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 26326534310.906525,
+      "step": 247,
+      "teaching_rms": 20958447787.57356,
+      "traffic_rms": 26326534310.906525
+    },
+    {
+      "batch_loss": 3.3513238430023193,
+      "innovation_rms": 22374408649.476772,
+      "instruction_rms": 1.1927043543899123,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 103241895378944.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.448609841808333e+27,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33391915302912.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.379103899002075,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7926591634750366,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25722351517696.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 694892167168.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8765585752064.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 139419981971456.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 195079142637568.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 33391915302912.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.609853744506836,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7924171686172485,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25722351517696.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 35343713928.99611,
+      "step": 248,
+      "teaching_rms": 22374408649.476772,
+      "traffic_rms": 35343713928.99611
+    },
+    {
+      "batch_loss": 3.8315584659576416,
+      "innovation_rms": 26021528093.86996,
+      "instruction_rms": 1.3425449543232812,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 132947877298176.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.872951193176675e+27,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 53750894952448.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.389530658721924,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8848695158958435,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29927592165376.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1984264011776.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12893718315008.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 163139064692736.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 237190860767232.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 53750894952448.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.615872383117676,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8846741914749146,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 29927592165376.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 39948335652.44792,
+      "step": 249,
+      "teaching_rms": 26021528093.86996,
+      "traffic_rms": 39948335652.44792
+    },
+    {
+      "batch_loss": 5.269939422607422,
+      "innovation_rms": 33581947024.02712,
+      "instruction_rms": 1.8879461094301389,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 145037438484480.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0897899579847938e+28,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 84507004764160.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3483526706695557,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8742738962173462,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35728058744832.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4216035475456.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 22317713326080.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 194296368070656.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 311572966670336.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 84507004764160.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.611508369445801,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.87446928024292,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 35728058744832.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 42479251756.70039,
+      "step": 250,
+      "teaching_rms": 33581947024.02712,
+      "traffic_rms": 42479251756.70039
+    },
+    {
+      "batch_loss": 4.97896671295166,
+      "innovation_rms": 34167288631.721638,
+      "instruction_rms": 1.647582200453476,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 175610961854464.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.556575814758906e+28,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 75842008907776.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.319918632507324,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1048455238342285,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 43312257957888.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6681994461184.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 24659590381568.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 229610260267008.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 353138955517952.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 75842008907776.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6093852519989014,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1049880981445312,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 43312257957888.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 55351684785.12176,
+      "step": 251,
+      "teaching_rms": 34167288631.721638,
+      "traffic_rms": 55351684785.12176
+    },
+    {
+      "batch_loss": 4.158825397491455,
+      "innovation_rms": 44200283289.92419,
+      "instruction_rms": 1.6401720056506044,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 207306981638144.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.230690796768658e+28,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 127931422605312.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.317570924758911,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4348928928375244,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 52553127559168.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9769473540096.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30874785546240.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 272448566067200.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 428383024447488.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 127931422605312.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6318159103393555,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.435035467147827,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 52553127559168.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 65721928277.07544,
+      "step": 252,
+      "teaching_rms": 44200283289.92419,
+      "traffic_rms": 65721928277.07544
+    },
+    {
+      "batch_loss": 6.091696262359619,
+      "innovation_rms": 29616514024.663853,
+      "instruction_rms": 2.095460186980915,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 219990154280960.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3040884139469837e+28,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 97330065833984.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3271889686584473,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9106519222259521,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 62286135820288.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 11885205258240.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 21157319278592.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 314046297407488.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 415977447620608.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 97330065833984.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.64140248298645,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9107944965362549,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 62286135820288.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 52824297822.686356,
+      "step": 253,
+      "teaching_rms": 29616514024.663853,
+      "traffic_rms": 52824297822.686356
+    },
+    {
+      "batch_loss": 6.626300811767578,
+      "innovation_rms": 58762211978.39095,
+      "instruction_rms": 1.166524714708832,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 262481272373248.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1207669177080876e+28,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 198286073921536.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.323005437850952,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4978218078613281,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 73799542243328.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 15030876635136.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 31456711671808.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 365920140656640.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 518738499600384.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 198286073921536.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6371874809265137,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4979642629623413,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 73799542243328.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 82508563963.41739,
+      "step": 254,
+      "teaching_rms": 58762211978.39095,
+      "traffic_rms": 82508563963.41739
+    },
+    {
+      "batch_loss": 7.429737091064453,
+      "innovation_rms": 50501816848.26978,
+      "instruction_rms": 1.6730923047076562,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 340384379043840.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.827967569922928e+28,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 155526352601088.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3120474815368652,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1634938716888428,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 85854752079872.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17733496340480.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 27026199150592.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 418630864994304.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 527107109158912.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 155526352601088.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6261980533599854,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.163636326789856,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 85854752079872.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 106722606213.25696,
+      "step": 255,
+      "teaching_rms": 50501816848.26978,
+      "traffic_rms": 106722606213.25696
+    },
+    {
+      "batch_loss": 8.355598449707031,
+      "innovation_rms": 50378685867.913025,
+      "instruction_rms": 2.0209367305499524,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 439704289804288.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.218963115126639e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 134515213205504.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.319915533065796,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.006071925163269,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 97683737935872.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19519550521344.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 17860550197248.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 469715877101568.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 510850121072640.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 134515213205504.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6142609119415283,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.006105899810791,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 97683737935872.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 122630290052.00587,
+      "step": 256,
+      "teaching_rms": 50378685867.913025,
+      "traffic_rms": 122630290052.00587
+    },
+    {
+      "batch_loss": 5.642224311828613,
+      "innovation_rms": 55833508727.10588,
+      "instruction_rms": 1.7448215453956957,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 552235184422912.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.351955548432636e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 116471946543104.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.340750217437744,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8748375773429871,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 109122468970496.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20457000534016.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9374500126720.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 519456899989504.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 497410094661632.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 116471946543104.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6021981239318848,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8750388622283936,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 109122468970496.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 140124051448.22815,
+      "step": 257,
+      "teaching_rms": 55833508727.10588,
+      "traffic_rms": 140124051448.22815
+    },
+    {
+      "batch_loss": 5.446631908416748,
+      "innovation_rms": 62232485472.44458,
+      "instruction_rms": 1.932859899967959,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 675464272150528.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.5948036376509924e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 113017509380096.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3756184577941895,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6588162183761597,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 120194584281088.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20568703238144.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1117037002752.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 568522203725824.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 490653037363200.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 113017509380096.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5964910984039307,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6586854457855225,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 120194584281088.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 157677362746.15762,
+      "step": 258,
+      "teaching_rms": 62232485472.44458,
+      "traffic_rms": 157677362746.15762
+    },
+    {
+      "batch_loss": 3.661562204360962,
+      "innovation_rms": 69303976980.20546,
+      "instruction_rms": 2.0643835740612326,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 798481732075520.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.840843857717109e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 124096578846720.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4203522205352783,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.410906195640564,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 130916944969728.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 19933035495424.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6356669562880.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 617883122008064.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 493609384148992.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 124096578846720.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.61080002784729,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4107753038406372,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 130916944969728.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 172129577323.2865,
+      "step": 259,
+      "teaching_rms": 69303976980.20546,
+      "traffic_rms": 172129577323.2865
+    },
+    {
+      "batch_loss": 3.473273754119873,
+      "innovation_rms": 76824601507.44531,
+      "instruction_rms": 2.4904398815160387,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 914905947439104.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.069332363829015e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 122001825988608.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4701080322265625,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.202252984046936,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 141393091624960.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18681014779904.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12520215543808.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 668595478593536.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 507123800735744.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 122001825988608.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6203737258911133,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2024544477462769,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 141393091624960.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 184215077975.66595,
+      "step": 260,
+      "teaching_rms": 76824601507.44531,
+      "traffic_rms": 184215077975.66595
+    },
+    {
+      "batch_loss": 3.0373620986938477,
+      "innovation_rms": 84107552672.55267,
+      "instruction_rms": 2.197618891638174,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1022164165722112.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.296566987192389e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 125356086394880.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5225634574890137,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1828148365020752,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 152373242626048.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16856139169792.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 18248751906816.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 721531084734464.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 529356363399168.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 125356086394880.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6225600242614746,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.183016300201416,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 152373242626048.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 194223773487.52783,
+      "step": 261,
+      "teaching_rms": 84107552672.55267,
+      "traffic_rms": 194223773487.52783
+    },
+    {
+      "batch_loss": 2.905885934829712,
+      "innovation_rms": 96300167604.67647,
+      "instruction_rms": 2.5000495725909624,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1134564063838208.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.425825350873051e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 144349891395584.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.556818962097168,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.918532133102417,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 164853310291968.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 14837945991168.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 20181927591936.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 779542134259712.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 580110495252480.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 144349891395584.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.613913059234619,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.918733537197113,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 164853310291968.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 208486707473.4352,
+      "step": 262,
+      "teaching_rms": 96300167604.67647,
+      "traffic_rms": 208486707473.4352
+    },
+    {
+      "batch_loss": 3.059303045272827,
+      "innovation_rms": 112035039852.61487,
+      "instruction_rms": 3.6953822332260002,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1254877908959232.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.3838454168344e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 178214299238400.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5553247928619385,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6180866956710815,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 179736529600512.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 13206128951296.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16318174593024.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 846775753637888.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 672336462217216.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 178214299238400.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.592906951904297,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6180574297904968,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 179736529600512.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 224269031570.81308,
+      "step": 263,
+      "teaching_rms": 112035039852.61487,
+      "traffic_rms": 224269031570.81308
+    },
+    {
+      "batch_loss": 2.6933038234710693,
+      "innovation_rms": 134388350334.98827,
+      "instruction_rms": 3.5456475160155065,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1380324105781248.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.976426852570814e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 272170517266432.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5596272945404053,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5568531155586243,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 198436095787008.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 12962465054720.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2436638965760.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 930188581928960.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 834127947366400.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 272170517266432.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5684714317321777,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5568238496780396,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 198436095787008.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 242725033333.84155,
+      "step": 264,
+      "teaching_rms": 134388350334.98827,
+      "traffic_rms": 242725033333.84155
+    },
+    {
+      "batch_loss": 2.8206863403320312,
+      "innovation_rms": 171257210961.70883,
+      "instruction_rms": 2.7670975109724765,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1504626767560704.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.986722367082107e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 430303109709824.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5635266304016113,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6023237705230713,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 225537322647552.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 16061147643904.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 30986827988992.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1044565574287360.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1143769923584000.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 430303109709824.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5302488803863525,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6025251150131226,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 225537322647552.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 264883524114.05557,
+      "step": 265,
+      "teaching_rms": 171257210961.70883,
+      "traffic_rms": 264883524114.05557
+    },
+    {
+      "batch_loss": 2.46866774559021,
+      "innovation_rms": 215261499272.33505,
+      "instruction_rms": 2.563263018379393,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1579097205506048.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.506262756485851e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 330458139197440.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.555485486984253,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7810490131378174,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 257552948396032.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 25569559314432.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 95084108316672.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1209464938037248.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1648993234845696.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 330458139197440.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.485705852508545,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7809367179870605,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 257552948396032.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 279383794706.7367,
+      "step": 266,
+      "teaching_rms": 215261499272.33505,
+      "traffic_rms": 279383794706.7367
+    },
+    {
+      "batch_loss": 2.8190884590148926,
+      "innovation_rms": 245015504981.55734,
+      "instruction_rms": 2.868306225394445,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1612839005454336.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.156100415149023e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 485100382846976.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.554853916168213,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6152262687683105,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 269686398779392.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 42158237155328.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 165886770020352.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1428389890097152.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2189249252163584.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 485100382846976.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5036590099334717,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6153700947761536,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 269686398779392.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 309075585191.1526,
+      "step": 267,
+      "teaching_rms": 245015504981.55734,
+      "traffic_rms": 309075585191.1526
+    },
+    {
+      "batch_loss": 2.7104287147521973,
+      "innovation_rms": 261069371483.3258,
+      "instruction_rms": 2.630471013508244,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1903486153261056.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.96325636134475e+29,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 341861881348096.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5495033264160156,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.43845969438552856,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 303872593625088.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 61365737750528.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 192074997563392.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1685024252362752.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2566343354220544.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 341861881348096.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5015883445739746,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.43837687373161316,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 303872593625088.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 400106587440.26337,
+      "step": 268,
+      "teaching_rms": 261069371483.3258,
+      "traffic_rms": 400106587440.26337
+    },
+    {
+      "batch_loss": 2.8868520259857178,
+      "innovation_rms": 320692775826.5796,
+      "instruction_rms": 2.890343753925133,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2181285174837248.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2353323107346078e+30,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 590473446031360.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5405216217041016,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5472670197486877,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 362919938228224.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 84878645264384.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 235129075138560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1994449534058496.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3094253085392896.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 590473446031360.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4785943031311035,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5473370552062988,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 362919938228224.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 476798407664.05634,
+      "step": 269,
+      "teaching_rms": 320692775826.5796,
+      "traffic_rms": 476798407664.05634
+    },
+    {
+      "batch_loss": 2.8182077407836914,
+      "innovation_rms": 321470776219.6614,
+      "instruction_rms": 2.655034778282445,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2178624442597376.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4098517086862198e+30,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 899846584139776.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5197155475616455,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5669795870780945,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 374984702689280.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 109714125881344.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 248354772615168.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2341445713264640.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3469960718319616.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 899846584139776.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.436824321746826,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5670547485351562,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 374984702689280.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 445773349057.20764,
+      "step": 270,
+      "teaching_rms": 321470776219.6614,
+      "traffic_rms": 445773349057.20764
+    },
+    {
+      "batch_loss": 2.7423083782196045,
+      "innovation_rms": 392973021002.6427,
+      "instruction_rms": 3.3067094676221385,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2784225971404800.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0077335285630072e+30,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 689826105917440.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5026533603668213,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5839838981628418,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 406988349505536.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 135375725002752.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 256615957659648.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2730962202918912.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3895164359671808.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 689826105917440.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.434230089187622,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5840513110160828,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 406988349505536.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 648418237907.8325,
+      "step": 271,
+      "teaching_rms": 392973021002.6427,
+      "traffic_rms": 648418237907.8325
+    },
+    {
+      "batch_loss": 2.633946418762207,
+      "innovation_rms": 435109454885.83856,
+      "instruction_rms": 3.0802899493364992,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3517024164642816.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.770033793379208e+30,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 781231935455232.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4905688762664795,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4553600549697876,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 473170205212672.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 160617147138048.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 252414238130176.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3158145287323648.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4271831917789184.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 781231935455232.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.459500551223755,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4551309645175934,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 473170205212672.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 748470318279.1506,
+      "step": 272,
+      "teaching_rms": 435109454885.83856,
+      "traffic_rms": 748470318279.1506
+    },
+    {
+      "batch_loss": 2.599663734436035,
+      "innovation_rms": 594900980205.4583,
+      "instruction_rms": 2.0843406400410696,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4074653190455296.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.916359318434007e+30,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1427205284429824.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4785211086273193,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.47080233693122864,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 557058399141888.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 193938392612864.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 333212488302592.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3685781482766336.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5276361954426880.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1427205284429824.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4789538383483887,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.4708421230316162,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 557058399141888.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 884224641750.5643,
+      "step": 273,
+      "teaching_rms": 594900980205.4583,
+      "traffic_rms": 884224641750.5643
+    },
+    {
+      "batch_loss": 2.895754098892212,
+      "innovation_rms": 533758261393.82294,
+      "instruction_rms": 1.5245285307053151,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3960510676467712.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.4179087910691676e+30,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1556746061479936.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.45917010307312,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5334426760673523,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 643541390852096.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 222682595459072.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 287442062016512.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4254592455933952.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5688110805417984.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1556746061479936.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.478808879852295,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5334137678146362,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 643541390852096.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 812884845398.0223,
+      "step": 274,
+      "teaching_rms": 533758261393.82294,
+      "traffic_rms": 812884845398.0223
+    },
+    {
+      "batch_loss": 2.3637542724609375,
+      "innovation_rms": 924097357120.2264,
+      "instruction_rms": 1.5508676393630367,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4415469851246592.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.028938918255016e+30,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1903438505967616.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.442807197570801,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5105206370353699,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 760330879762432.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 277377192034304.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 546946032861184.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5044441773506560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7898494786338816.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1903438505967616.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.477027416229248,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5104377865791321,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 760330879762432.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1222914735697.6338,
+      "step": 275,
+      "teaching_rms": 924097357120.2264,
+      "traffic_rms": 1222914735697.6338
+    },
+    {
+      "batch_loss": 2.6887388229370117,
+      "innovation_rms": 1052835812140.0055,
+      "instruction_rms": 1.1983358524287249,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4954456839946240.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.451093124476283e+30,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2776275986939904.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4282398223876953,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.44247275590896606,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 987940725456896.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 357273499074560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 798963103956992.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6061507209068544.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0170656503103488e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2776275986939904.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.486109733581543,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.44250020384788513,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 987940725456896.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1371106689537.711,
+      "step": 276,
+      "teaching_rms": 1052835812140.0055,
+      "traffic_rms": 1371106689537.711
+    },
+    {
+      "batch_loss": 2.64987850189209,
+      "innovation_rms": 1051405783794.9537,
+      "instruction_rms": 1.1295926717954636,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5246307383902208.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.385466828656519e+31,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2727401138159616.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4124605655670166,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5621763467788696,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1163485836738560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 446810950729728.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 895374516551680.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7206770439094272.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1452631226515456e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2727401138159616.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.493664503097534,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.5621195435523987,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1163485836738560.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1665498011030.9807,
+      "step": 277,
+      "teaching_rms": 1051405783794.9537,
+      "traffic_rms": 1665498011030.9807
+    },
+    {
+      "batch_loss": 2.5996406078338623,
+      "innovation_rms": 1454083280428.344,
+      "instruction_rms": 1.0314111680549072,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6475128665800704.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0233311499339661e+31,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3213578819600384.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.398977756500244,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7118824124336243,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1415430497370112.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 562256383311872.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1154454460039168.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8624466665930752.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4176963342106624e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3213578819600384.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.497809886932373,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7120504975318909,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1415430497370112.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2064826709015.2087,
+      "step": 278,
+      "teaching_rms": 1454083280428.344,
+      "traffic_rms": 2064826709015.2087
+    },
+    {
+      "batch_loss": 3.0978729724884033,
+      "innovation_rms": 1152173335634.5598,
+      "instruction_rms": 1.26700296295527,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6816534005547008.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.236882697587866e+31,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4602675932954624.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.406496286392212,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.191275954246521,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1589436668182528.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 657817258164224.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 955608480088064.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0081332994506752e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4568660064534528e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4602675932954624.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4987728595733643,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1913166046142578,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1589436668182528.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1803457383526.771,
+      "step": 279,
+      "teaching_rms": 1152173335634.5598,
+      "traffic_rms": 1803457383526.771
+    },
+    {
+      "batch_loss": 3.8058114051818848,
+      "innovation_rms": 1715914864377.4136,
+      "instruction_rms": 1.5327025451026466,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9743046337888256.0,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.5631840921695353e+31,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3094122357325824.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4282915592193604,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0749881267547607,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1898848930758656.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 762488362631168.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1046711178887168.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1728370274402304e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6470370651471872e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3094122357325824.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.499631643295288,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0749599933624268,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1898848930758656.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2934296302609.815,
+      "step": 280,
+      "teaching_rms": 1715914864377.4136,
+      "traffic_rms": 2934296302609.815
+    },
+    {
+      "batch_loss": 4.068282604217529,
+      "innovation_rms": 1970952388213.9148,
+      "instruction_rms": 2.1725431396277215,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3238142547001344e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.442327747631302e+31,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4332899910287360.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4586100578308105,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7519958019256592,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2332139056005120.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 876495719366656.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1140073433137152.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3583743532924928e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8553736880193536e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4332899910287360.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5019984245300293,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7522295713424683,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2332139056005120.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3483290201097.091,
+      "step": 281,
+      "teaching_rms": 1970952388213.9148,
+      "traffic_rms": 3483290201097.091
+    },
+    {
+      "batch_loss": 3.69356369972229,
+      "innovation_rms": 2933921679396.9565,
+      "instruction_rms": 3.3490699140878517,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5531832250466304e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.8285722630919e+31,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6844932329308160.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3239917755126953,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.724652647972107,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2882694270681088.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1071638497984512.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1951428054614016.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.612829475995648e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5445516565282816e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6844932329308160.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.502347707748413,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7243497371673584,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2882694270681088.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3932904686313.7485,
+      "step": 282,
+      "teaching_rms": 2933921679396.9565,
+      "traffic_rms": 3932904686313.7485
+    },
+    {
+      "batch_loss": 5.8331708908081055,
+      "innovation_rms": 1325013678385.2354,
+      "instruction_rms": 4.012837766143188,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5269975610621952e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.967484126179883e+31,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4402849358282752.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.264617443084717,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6168303489685059,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3216830378278912.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1179170654650368.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1075322170638336.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.847259206307021e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.344298162107187e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4402849358282752.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4994070529937744,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6169259548187256,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3216830378278912.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2960033243835.289,
+      "step": 283,
+      "teaching_rms": 1325013678385.2354,
+      "traffic_rms": 2960033243835.289
+    },
+    {
+      "batch_loss": 7.504837512969971,
+      "innovation_rms": 1794370778997.6638,
+      "instruction_rms": 4.152443101326175,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5936804716806144e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.46106290910116e+31,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5670000002072576.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.242121696472168,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.549020767211914,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3380119397728256.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1222777189171200.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 436065076772864.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.073592171646157e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.263328794397901e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5670000002072576.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4927241802215576,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.548984408378601,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3380119397728256.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3442432043338.061,
+      "step": 284,
+      "teaching_rms": 1794370778997.6638,
+      "traffic_rms": 3442432043338.061
+    },
+    {
+      "batch_loss": 6.359536647796631,
+      "innovation_rms": 4134529752046.8037,
+      "instruction_rms": 3.87344741958381,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.101034650435584e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3181856337694323e+32,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2745979023327232e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2252871990203857,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6441797018051147,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4004365782220800.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1359015430848512.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1362382819426304.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.385946459976499e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1235424538066944e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2745979023327232e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5157461166381836,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6440629959106445,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4004365782220800.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 5901644666540.01,
+      "step": 285,
+      "teaching_rms": 4134529752046.8037,
+      "traffic_rms": 5901644666540.01
+    },
+    {
+      "batch_loss": 8.541296005249023,
+      "innovation_rms": 3898856185544.41,
+      "instruction_rms": 3.6988169338624584,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.596841165370163e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2857243645550183e+32,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8193770566713344.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2120397090911865,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1012301445007324,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4444539800518656.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1535891445121024.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1768759874289664.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7378756507140096e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.51929083363328e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8193770566713344.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5217885971069336,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1013731956481934,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4444539800518656.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 7292007391244.225,
+      "step": 286,
+      "teaching_rms": 3898856185544.41,
+      "traffic_rms": 7292007391244.225
+    },
+    {
+      "batch_loss": 9.148499488830566,
+      "innovation_rms": 4404672532042.585,
+      "instruction_rms": 2.484277333065391,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.136802735836365e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.536903398848031e+32,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9847408674471936.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.20345139503479,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.446024179458618,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4900806792839168.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1760857537118208.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2249661322625024.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1322376150974464e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.943619643834368e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9847408674471936.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.510789155960083,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.446167230606079,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4900806792839168.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 8390995364066.092,
+      "step": 287,
+      "teaching_rms": 4404672532042.585,
+      "traffic_rms": 8390995364066.092
+    },
+    {
+      "batch_loss": 8.919160842895508,
+      "innovation_rms": 6504532675220.747,
+      "instruction_rms": 1.9813968887954687,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.751261802541875e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.770762878278917e+32,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.340409364086784e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.198662281036377,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8820202350616455,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6004969870196736.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2136884339802112.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3760268295274496.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.669581402996736e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.373437019999437e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.340409364086784e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.481365203857422,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8821632862091064,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6004969870196736.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 9037902641816.133,
+      "step": 288,
+      "teaching_rms": 6504532675220.747,
+      "traffic_rms": 9037902641816.133
+    },
+    {
+      "batch_loss": 7.122723579406738,
+      "innovation_rms": 4068224534637.469,
+      "instruction_rms": 1.8756921015685601,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.126373938253005e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.0291120667837435e+32,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9460555601412096.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1968603134155273,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5155644416809082,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6904242069569536.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2456200796962816.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3193165376913408.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.202716552442675e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.331349346975744e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9460555601412096.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.425567150115967,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5157074928283691,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6904242069569536.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 7056456745540.297,
+      "step": 289,
+      "teaching_rms": 4068224534637.469,
+      "traffic_rms": 7056456745540.297
+    },
+    {
+      "batch_loss": 7.7410125732421875,
+      "innovation_rms": 4873390438377.807,
+      "instruction_rms": 2.9744154477299034,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.650921730087322e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.783258175303676e+32,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0243366277021696e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2044050693511963,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5637519359588623,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7611435278401536.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2733395570327552.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2771948538953728.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.742051057683661e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.393343763919667e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0243366277021696e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3488426208496094,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.563657283782959,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7611435278401536.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 8725100048580.206,
+      "step": 290,
+      "teaching_rms": 4873390438377.807,
+      "traffic_rms": 8725100048580.206
+    },
+    {
+      "batch_loss": 8.13406753540039,
+      "innovation_rms": 9534462662799.725,
+      "instruction_rms": 3.796153437294842,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.55541091886039e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.307782334005631e+32,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4285684269318144e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.219452381134033,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.5726685523986816,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9168287371886592.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3227629268238336.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4942336979107840.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.480696956269363e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.386460703843942e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4285684269318144e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.276510000228882,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.572673797607422,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9168287371886592.0,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 13366863720878.73,
+      "step": 291,
+      "teaching_rms": 9534462662799.725,
+      "traffic_rms": 13366863720878.73
+    },
+    {
+      "batch_loss": 7.107175350189209,
+      "innovation_rms": 7145439951179.104,
+      "instruction_rms": 4.878034986152451,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.728239973354701e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.10987552907757e+33,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5426139547762688e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.160581350326538,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.9802539348602295,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0710901111914496e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3634063634071552.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4064343658332160.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.238862622221926e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.581654941538714e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5426139547762688e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2723376750946045,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.980259418487549,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0710901111914496e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 14147338932279.238,
+      "step": 292,
+      "teaching_rms": 7145439951179.104,
+      "traffic_rms": 14147338932279.238
+    },
+    {
+      "batch_loss": 4.3914031982421875,
+      "innovation_rms": 8622523860282.973,
+      "instruction_rms": 3.9426137594041717,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.029595065188352e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6060698278452868e+33,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.935244215602381e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.153576612472534,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4211466312408447,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1966617479020544e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4019499904794624.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3854363244101632.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.038521537724416e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.996587437037978e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.935244215602381e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.265321731567383,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.421151876449585,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1966617479020544e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 17701099216666.18,
+      "step": 293,
+      "teaching_rms": 8622523860282.973,
+      "traffic_rms": 17701099216666.18
+    },
+    {
+      "batch_loss": 4.1768975257873535,
+      "innovation_rms": 13653612246640.762,
+      "instruction_rms": 2.5921120669455378,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.876082868112589e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.378213572883865e+33,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.4111334639468544e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1500535011291504,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.010742664337158,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3308774357925888e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4669131425054720.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6496317350084608.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.094719708810445e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0561977415892992e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.4111334639468544e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2617874145507812,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0107481479644775,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3308774357925888e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 20993669334356.605,
+      "step": 294,
+      "teaching_rms": 13653612246640.762,
+      "traffic_rms": 20993669334356.605
+    },
+    {
+      "batch_loss": 3.5078752040863037,
+      "innovation_rms": 12069228248765.887,
+      "instruction_rms": 1.777379337171332,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.59835635136594e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.800748969260305e+33,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.222867625489203e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2195053100585938,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6943225860595703,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5177828161028096e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5375377593597952.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7062461685432320.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.25657708185518e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.161857373044736e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.222867625489203e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2601797580718994,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6943281888961792,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5177828161028096e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 19853480432765.457,
+      "step": 295,
+      "teaching_rms": 12069228248765.887,
+      "traffic_rms": 19853480432765.457
+    },
+    {
+      "batch_loss": 3.3462252616882324,
+      "innovation_rms": 17396040662464.398,
+      "instruction_rms": 1.4000164490990457,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.349869971970458e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.4641591794032814e+33,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.209963263269274e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2964396476745605,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.521535873413086,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.810371450188595e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6402926507458560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0275488064864256e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0738864784985293e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4822873595327283e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.209963263269274e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2560927867889404,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5215413570404053,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.810371450188595e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 24432298353941.574,
+      "step": 296,
+      "teaching_rms": 17396040662464.398,
+      "traffic_rms": 24432298353941.574
+    },
+    {
+      "batch_loss": 3.659221887588501,
+      "innovation_rms": 18094178817346.508,
+      "instruction_rms": 1.698861143154008,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.092991226478592e+16,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.268008652367502e+33,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.092863072914637e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.369978189468384,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3692469596862793,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.035994886930432e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7656534045622272.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2536076455378944e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2428761758275994e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6898973168880845e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.092863072914637e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.300248622894287,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.369252324104309,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.035994886930432e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 32045779362862.887,
+      "step": 297,
+      "teaching_rms": 18094178817346.508,
+      "traffic_rms": 32045779362862.887
+    },
+    {
+      "batch_loss": 3.812203884124756,
+      "innovation_rms": 20800575371405.074,
+      "instruction_rms": 1.585817909674392,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.028437478169641e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.73426189022276e+33,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.05471436891095e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4304144382476807,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.232173204421997,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.191365114114867e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9125803702878208.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4692698720043008e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4347667246782874e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.91890548850688e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.05471436891095e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.360691785812378,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2321785688400269,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.191365114114867e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 37683933447103.375,
+      "step": 298,
+      "teaching_rms": 20800575371405.074,
+      "traffic_rms": 37683933447103.375
+    },
+    {
+      "batch_loss": 4.287206172943115,
+      "innovation_rms": 29079114473854.957,
+      "instruction_rms": 1.4663416012390826,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2185174406070272e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0368248575760951e+34,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.01657229485015e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4675121307373047,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1087944507598877,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3138045385506816e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1155336811511808e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0295335381303296e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6856724434583552e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.509057703196754e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.01657229485015e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.397796392440796,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1087998151779175,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3138045385506816e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 41443024493624.09,
+      "step": 299,
+      "teaching_rms": 29079114473854.957,
+      "traffic_rms": 41443024493624.09
+    },
+    {
+      "batch_loss": 3.404052972793579,
+      "innovation_rms": 28813940244934.703,
+      "instruction_rms": 1.2802243311221682,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3732934590660608e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1692273108468177e+34,
+          "nonfinite_indices": []
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.974747340845875e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.511240243911743,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1189416646957397,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.9749395918422016e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.351157519745024e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.356238171190067e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9777805566030643e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.921080444252324e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.974747340845875e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4415314197540283,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1187609434127808,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.9749395918422016e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 39440055189320.83,
+      "step": 300,
+      "teaching_rms": 28813940244934.703,
+      "traffic_rms": 39440055189320.83
+    },
+    {
+      "batch_loss": 3.7177019119262695,
+      "innovation_rms": 43215332344378.836,
+      "instruction_rms": 1.1601879782573492,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6976916517381734e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 5.061835000069689e+33,
+          "nonfinite_indices": [
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.935529806974157e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.562462091445923,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.09159517288208,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.968492658438963e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6938600469889024e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.427025272438784e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3718965482094592e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.941159572466565e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.935529806974157e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.492760181427002,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.091774344444275,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.968492658438963e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 56381130198530.76,
+      "step": 301,
+      "teaching_rms": 43215332344378.836,
+      "traffic_rms": 56381130198530.76
+    },
+    {
+      "batch_loss": 3.1790716648101807,
+      "innovation_rms": 41990742007180.27,
+      "instruction_rms": 1.3106395675635756,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.803576854875996e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 6.8762846441053e+33,
+          "nonfinite_indices": [
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2453761044919091e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6102960109710693,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0921401977539062,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.347955166524211e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0712913576656896e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.774312462522778e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.8238394640931226e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.5194291588366336e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2453761044919091e+17,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5406012535095215,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0923192501068115,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.347955166524211e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 62129347275250.11,
+      "step": 302,
+      "teaching_rms": 41990742007180.27,
+      "traffic_rms": 62129347275250.11
+    },
+    {
+      "batch_loss": 3.702026128768921,
+      "innovation_rms": 1311758775136.3054,
+      "instruction_rms": 1.6657931486666855,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5021601575180698e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 27209.826171875,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1208016088635802e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6243250370025635,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9091581106185913,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.80383502123008e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3738728774107136e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.0258145532051456e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.230588157107896e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.06748624295297e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1208016088635802e+17,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5546371936798096,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9090908169746399,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.80383502123008e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 6587632519136.7295,
+      "step": 303,
+      "teaching_rms": 1311758775136.3054,
+      "traffic_rms": 6587632519136.7295
+    },
+    {
+      "batch_loss": 2.856395721435547,
+      "innovation_rms": 1503385135554.5295,
+      "instruction_rms": 1.6017039338180512,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.6856633067411866e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 29012.794921875,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0086734130629837e+17,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6319046020507812,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9451942443847656,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.81250852019241e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6036688748806144e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.297959974699008e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.5966619636413235e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.6607373781395046e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0086734130629837e+17,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5622236728668213,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.945127010345459,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.81250852019241e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 7549977400305.455,
+      "step": 304,
+      "teaching_rms": 1503385135554.5295,
+      "traffic_rms": 7549977400305.455
+    },
+    {
+      "batch_loss": 2.756110668182373,
+      "innovation_rms": 1648915933687.7974,
+      "instruction_rms": 1.5828848259719943,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.2234883729824154e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 30503.814453125,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.0774786635989e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6454672813415527,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7811537384986877,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.720256300652954e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.763841304998707e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6017234421874688e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.926128217722716e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.2946635716060774e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.0774786635989e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.575793504714966,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7810865640640259,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.720256300652954e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 8280830993833.627,
+      "step": 305,
+      "teaching_rms": 1648915933687.7974,
+      "traffic_rms": 8280830993833.627
+    },
+    {
+      "batch_loss": 3.0536513328552246,
+      "innovation_rms": 1750353998181.8367,
+      "instruction_rms": 1.6109747668347518,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.999356332247613e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 31552.802734375,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.169058720756531e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.671168088912964,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6108850240707397,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.53716251632599e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.8584828421013504e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9464160152715264.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.2226479494751846e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.965196973927301e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.169058720756531e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.601501226425171,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6107044816017151,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.53716251632599e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 8790251152231.012,
+      "step": 306,
+      "teaching_rms": 1750353998181.8367,
+      "traffic_rms": 8790251152231.012
+    },
+    {
+      "batch_loss": 3.494131088256836,
+      "innovation_rms": 1810290946448.5408,
+      "instruction_rms": 1.7437503337564042,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.918584115100385e+17,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 32082.279296875,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.351399167819776e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.708743095397949,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7781113386154175,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.272302862604698e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.892451308948685e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3396845610991616.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.489515570613453e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6686772421748326e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.351399167819776e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6390833854675293,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7782567143440247,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.272302862604698e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 9091253856357.625,
+      "step": 307,
+      "teaching_rms": 1810290946448.5408,
+      "traffic_rms": 9091253856357.625
+    },
+    {
+      "batch_loss": 3.6805436611175537,
+      "innovation_rms": 1831803387554.9172,
+      "instruction_rms": 1.5092261582339774,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0855722942369628e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 32033.927734375,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.615432040336589e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.619041919708252,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9646860361099243,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.933846367286067e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.8712053942255616e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2124590398570496.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.729696670156063e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4018094492378726e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.615432040336589e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.549389123916626,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9643762111663818,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.933846367286067e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 9199288314302.055,
+      "step": 308,
+      "teaching_rms": 1831803387554.9172,
+      "traffic_rms": 9199288314302.055
+    },
+    {
+      "batch_loss": 4.06059455871582,
+      "innovation_rms": 1818348150545.9204,
+      "instruction_rms": 1.077047202586497,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2703744977818092e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 31609.888671875,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.952995182457651e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5530970096588135,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7019176483154297,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.529146014380851e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.800647241990144e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7055821129121792.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.945859625384673e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1616285214939546e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.952995182457651e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4834511280059814,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7016735672950745,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.529146014380851e+16,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 9131716878165.963,
+      "step": 309,
+      "teaching_rms": 1818348150545.9204,
+      "traffic_rms": 9131716878165.963
+    },
+    {
+      "batch_loss": 4.295765399932861,
+      "innovation_rms": 1773663403463.889,
+      "instruction_rms": 1.232870998547259,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4502573488650322e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 30935.669921875,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.356742482119885e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.512589693069458,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9892271757125854,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0064820348492186e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6869719130701824e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1367526449545216e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.140406182011208e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.945465566265344e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.356742482119885e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.44295072555542,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.9891583919525146,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0064820348492186e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 8907310352431.25,
+      "step": 310,
+      "teaching_rms": 1773663403463.889,
+      "traffic_rms": 8907310352431.25
+    },
+    {
+      "batch_loss": 3.8537802696228027,
+      "innovation_rms": 1701672259398.7644,
+      "instruction_rms": 1.3666382753509014,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6352456447528796e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 30079.271484375,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.820061837171098e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.493788003921509,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4506874084472656,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0546826489259622e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5365277986258944e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5044413591912448e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.3154982547737805e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7509190096388096e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.820061837171098e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4340054988861084,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4506187438964844,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0546826489259622e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 8545772257801.934,
+      "step": 311,
+      "teaching_rms": 1701672259398.7644,
+      "traffic_rms": 8545772257801.934
+    },
+    {
+      "batch_loss": 3.379138708114624,
+      "innovation_rms": 1606395407773.369,
+      "instruction_rms": 1.3910418105135907,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.823407543515349e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 29143.642578125,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.337000852435763e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4922752380371094,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2139723300933838,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0980526273855488e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3556869121376256e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8084095091277824e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.473080948461404e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5758271086749286e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.337000852435763e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4620237350463867,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2141317129135132,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0980526273855488e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 8067292225774.687,
+      "step": 312,
+      "teaching_rms": 1606395407773.369,
+      "traffic_rms": 8067292225774.687
+    },
+    {
+      "batch_loss": 3.3933894634246826,
+      "innovation_rms": 1491867839259.123,
+      "instruction_rms": 1.7961017341901608,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0142914207513313e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 28232.228515625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.902202415205581e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.523283004760742,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9909117221832275,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1370746386527027e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.150728495792128e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.0495837339582464e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.614905235341312e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.418244329087959e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.902202415205581e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.487215518951416,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9910710453987122,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1370746386527027e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 7492137671135.153,
+      "step": 313,
+      "teaching_rms": 1491867839259.123,
+      "traffic_rms": 7492137671135.153
+    },
+    {
+      "batch_loss": 3.9649484157562256,
+      "innovation_rms": 1362066644465.2925,
+      "instruction_rms": 2.383113676544935,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2006645515343626e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 27411.04296875,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.510844995194061e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.524592161178589,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7418479919433594,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1721830757197414e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.927736156015821e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2299227535179776e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.742547196612444e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2764198704093594e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.510844995194061e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5098633766174316,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7420918941497803,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1721830757197414e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 6840278211907.43,
+      "step": 314,
+      "teaching_rms": 1362066644465.2925,
+      "traffic_rms": 6840278211907.43
+    },
+    {
+      "batch_loss": 3.8712611198425293,
+      "innovation_rms": 1220844638654.8481,
+      "instruction_rms": 2.9671257874600694,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.371936852182565e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 26769.486328125,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1585883132002304e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5357213020324707,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2355642318725586,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2037689524092928e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6925082060324864e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.352279285084979e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.857424858677248e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1487778232388813e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1585883132002304e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5302212238311768,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2351467609405518,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2037689524092928e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 6131063084286.562,
+      "step": 315,
+      "teaching_rms": 1220844638654.8481,
+      "traffic_rms": 6131063084286.562
+    },
+    {
+      "batch_loss": 3.7337825298309326,
+      "innovation_rms": 1071873568531.4731,
+      "instruction_rms": 2.521951837806514,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.52508727705967e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 26334.7890625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.841525666971648e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5540175437927246,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.133857011795044,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.232184198341591e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4504822891872256e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4202595979493376e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.960814685816095e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0338999893753856e+17,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.841525666971648e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.548517942428589,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1334396600723267,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.232184198341591e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 5382932837750.842,
+      "step": 316,
+      "teaching_rms": 1071873568531.4731,
+      "traffic_rms": 5382932837750.842
+    },
+    {
+      "batch_loss": 3.2854020595550537,
+      "innovation_rms": 918597433123.386,
+      "instruction_rms": 1.9248413458376215,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6690771208099594e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 26064.169921875,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5561407667830784e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5704586505889893,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.809615969657898,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2577456103043891e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2066739609141248e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4380837122277376e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.05386566768001e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9.30509990437847e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5561407667830784e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5649595260620117,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8091986775398254,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2577456103043891e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 4613182244754.958,
+      "step": 317,
+      "teaching_rms": 918597433123.386,
+      "traffic_rms": 4613182244754.958
+    },
+    {
+      "batch_loss": 3.318406820297241,
+      "innovation_rms": 764192437566.5042,
+      "instruction_rms": 0.7449933209783343,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.8244609286809846e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 25986.5,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.299268865982464e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5852298736572266,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.712548017501831,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2807382882267955e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9656291219734528.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.4104479599099904e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.137611345199104e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8.374590085739315e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.299268865982464e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5797312259674072,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.7128598690032959,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2807382882267955e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3837762713133.6562,
+      "step": 318,
+      "teaching_rms": 764192437566.5042,
+      "traffic_rms": 3837762713133.6562
+    },
+    {
+      "batch_loss": 4.367030143737793,
+      "innovation_rms": 611537579612.6832,
+      "instruction_rms": 0.6789742133770322,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.9905556290739896e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 26088.525390625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.068061112762368e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5984981060028076,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5056228637695312,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3014188993544192e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7313898197745664.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.3423929146146816e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.212982867283149e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7.537130733568e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.068061112762368e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.5930001735687256,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.505850911140442,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3014188993544192e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 3071132353098.192,
+      "step": 319,
+      "teaching_rms": 611537579612.6832,
+      "traffic_rms": 3071132353098.192
+    },
+    {
+      "batch_loss": 4.783128261566162,
+      "innovation_rms": 463192685126.2008,
+      "instruction_rms": 0.6762415101069177,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1602787178386555e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 26352.3359375,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.859953626395443e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6104135513305664,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.354646921157837,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.320018427028439e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5074718071717888.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.239180126027776e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.280817237158789e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.7834176602112e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.859953626395443e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6049160957336426,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3544948101043701,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.320018427028439e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2326146722831.202,
+      "step": 320,
+      "teaching_rms": 463192685126.2008,
+      "traffic_rms": 2326146722831.202
+    },
+    {
+      "batch_loss": 4.381250381469727,
+      "innovation_rms": 321384340252.371,
+      "instruction_rms": 0.6883420379311149,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.3173703421455237e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 26627.21484375,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6726381840826368e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6211113929748535,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1022981405258179,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3367448335653274e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2968543687606272.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1061742230503424e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.341868307485819e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.105075679441715e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6726381840826368e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.615614414215088,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.102146029472351,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3367448335653274e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1613987223858.9526,
+      "step": 321,
+      "teaching_rms": 321384340252.371,
+      "traffic_rms": 1613987223858.9526
+    },
+    {
+      "batch_loss": 4.2254958152771,
+      "innovation_rms": 187999306399.16638,
+      "instruction_rms": 0.7657056286929292,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.456688086111486e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 26914.337890625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5040376322654208e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6307132244110107,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8095412254333496,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.351785207740498e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1019806161567744.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.948737472351437e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.396813652304855e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5.494567982648525e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5040376322654208e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6252167224884033,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8098962306976318,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.351785207740498e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 944129630938.7628,
+      "step": 322,
+      "teaching_rms": 187999306399.16638,
+      "traffic_rms": 944129630938.7628
+    },
+    {
+      "batch_loss": 4.357901573181152,
+      "innovation_rms": 64584815070.50162,
+      "instruction_rms": 1.086741100641491,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.575120606707843e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 27133.603515625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3522820495572992e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6393284797668457,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8275766372680664,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3653079981712998e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 752327141097472.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7721332489781248e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.446264874958848e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.945111098484326e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3522820495572992e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6338326930999756,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8279987573623657,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3653079981712998e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 324343952313.9717,
+      "step": 323,
+      "teaching_rms": 64584815070.50162,
+      "traffic_rms": 324343952313.9717
+    },
+    {
+      "batch_loss": 3.6836187839508057,
+      "innovation_rms": 1.288813169265937,
+      "instruction_rms": 1.2888124546817934,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.676929885881762e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 27224.05859375,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2156885603975168e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6470558643341064,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9609711170196533,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3774649031026278e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2347247140339712.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5949198918680576e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.490771044066918e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.450599730937856e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2156885603975168e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6415605545043945,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9611742496490479,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3774649031026278e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.2888220749446628,
+      "step": 324,
+      "teaching_rms": 1.288813169265937,
+      "traffic_rms": 0.004987924046494246
+    },
+    {
+      "batch_loss": 3.0870285034179688,
+      "innovation_rms": 2.1805320104092933,
+      "instruction_rms": 2.1805310077208606,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.7731148882016993e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 27163.345703125,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.092742222446592e+16,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6539840698242188,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9446864128112793,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3883923306958029e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3782675139657728.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4354278919438336e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.530826252666798e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.00553975784407e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.092742222446592e+16,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.648489236831665,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9448895454406738,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3883923306958029e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2.1805368464732537,
+      "step": 325,
+      "teaching_rms": 2.1805320104092933,
+      "traffic_rms": 0.005091514919230705
+    },
+    {
+      "batch_loss": 3.809065103530884,
+      "innovation_rms": 2.5611753161662834,
+      "instruction_rms": 2.561174966199351,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.8755357707185684e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 27147.275390625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9820795983364096.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6601929664611816,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6160850524902344,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3982131170154906e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5074560231669760.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2918850383249408e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.566875802967736e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.604985524361626e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9820795983364096.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.654698610305786,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6160545349121094,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3982131170154906e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2.561179960077052,
+      "step": 326,
+      "teaching_rms": 2.5611753161662834,
+      "traffic_rms": 0.0050796373955823695
+    },
+    {
+      "batch_loss": 5.620018005371094,
+      "innovation_rms": 2.459811170569769,
+      "instruction_rms": 2.459809511600436,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.9809311070664786e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 27213.7421875,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8824734011621376.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6657543182373047,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2677507400512695,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.40703781451989e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6237256733949952.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.162696502280192e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.599320329519104e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.244486800126771e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8824734011621376.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6602604389190674,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.2677910327911377,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.40703781451989e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 2.4598145501571533,
+      "step": 327,
+      "teaching_rms": 2.459811170569769,
+      "traffic_rms": 0.005090772453762295
+    },
+    {
+      "batch_loss": 4.8975605964660645,
+      "innovation_rms": 1.9774144803161828,
+      "instruction_rms": 1.977412691895821,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.0746441323692687e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 27228.337890625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7928190190223360.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6707329750061035,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6884660720825195,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4149659805509222e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7283683639689216.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0464267983650816e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.628520609573765e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.9200380986392576e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7928190190223360.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6652395725250244,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.6885063648223877,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4149659805509222e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.9774190306189625,
+      "step": 328,
+      "teaching_rms": 1.9774144803161828,
+      "traffic_rms": 0.005122406138830753
+    },
+    {
+      "batch_loss": 5.090121269226074,
+      "innovation_rms": 1.2132298119507232,
+      "instruction_rms": 1.213231835267189,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.146998594546696e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 27081.25390625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7121221133008896.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6751868724823,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3893728256225586,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.422087208126382e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8225467962228736.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9417841077911552.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.654800999061914e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6280342243508224e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7121221133008896.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.669694185256958,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3891856670379639,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.422087208126382e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.2132428283776222,
+      "step": 329,
+      "teaching_rms": 1.2132298119507232,
+      "traffic_rms": 0.005123541527144432
+    },
+    {
+      "batch_loss": 5.204372882843018,
+      "innovation_rms": 0.6984657323130641,
+      "instruction_rms": 0.6984601009588121,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.2158206007419863e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 26879.931640625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6394878275616768.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.679168701171875,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2704479694366455,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4284820708327424e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9073073315643392.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8476056755372032.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.678453555759677e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.365230758966067e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6394878275616768.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6736764907836914,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.270228385925293,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4284820708327424e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.69847814463874,
+      "step": 330,
+      "teaching_rms": 0.6984657323130641,
+      "traffic_rms": 0.005159413828791277
+    },
+    {
+      "batch_loss": 7.2479753494262695,
+      "innovation_rms": 0.7694590034720585,
+      "instruction_rms": 0.769458138017497,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.263100975125889e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 26439.263671875,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5741105333141504.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.682725667953491,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0017136335372925,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4342231536173056e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9835918563213312.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7628450865086464.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.699740788068188e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.1287075542204416e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5741105333141504.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.677233934402466,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0017958879470825,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4342231536173056e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.7694749194340567,
+      "step": 331,
+      "teaching_rms": 0.7694590034720585,
+      "traffic_rms": 0.00515263687724699
+    },
+    {
+      "batch_loss": 6.125957012176514,
+      "innovation_rms": 0.76342239269215,
+      "instruction_rms": 0.7634215750064157,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.2856038549778596e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 25749.01171875,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5152652937658368.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6858999729156494,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2216097116470337,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4393758258823168e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0522478749155328e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6865605617516544.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.718899090987418e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.915836734373888e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5152652937658368.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6804089546203613,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2216919660568237,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4393758258823168e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.7634380828713512,
+      "step": 332,
+      "teaching_rms": 0.76342239269215,
+      "traffic_rms": 0.005081960050256244
+    },
+    {
+      "batch_loss": 5.977670669555664,
+      "innovation_rms": 0.927103791469379,
+      "instruction_rms": 0.9271008725500851,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.28365194696065e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 24802.1328125,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4622993812619264.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.688730001449585,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9775581359863281,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4439987568810394e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.114038280912896e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6179044894703616.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.736141494895247e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.7242530609364992e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4622993812619264.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.683239459991455,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9776403903961182,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4439987568810394e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.9271137969824076,
+      "step": 333,
+      "teaching_rms": 0.927103791469379,
+      "traffic_rms": 0.005015180527409155
+    },
+    {
+      "batch_loss": 4.730621814727783,
+      "innovation_rms": 1.1270446560844627,
+      "instruction_rms": 1.1270395774902942,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.2534640305863393e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 23663.36328125,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4146254321811456.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6912500858306885,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0904147624969482,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.448144946509906e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1696497322098688e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5561140297859072.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.751659727131771e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5518277333680128e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4146254321811456.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.685760259628296,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0904440879821777,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.448144946509906e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.1270499420409956,
+      "step": 334,
+      "teaching_rms": 1.1270446560844627,
+      "traffic_rms": 0.004955452755976563
+    },
+    {
+      "batch_loss": 4.047018051147461,
+      "innovation_rms": 1.1522882680674469,
+      "instruction_rms": 1.1522873846906998,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.198003839447597e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 22407.634765625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3717147360493568.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.693491220474243,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0672166347503662,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.451862154805248e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2196999524777984e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5005026321760256.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.765626273583596e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3966448741318656e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3717147360493568.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.688002109527588,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0671088695526123,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.451862154805248e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.152297738854542,
+      "step": 335,
+      "teaching_rms": 1.1522882680674469,
+      "traffic_rms": 0.00489558817220692
+    },
+    {
+      "batch_loss": 3.5104641914367676,
+      "innovation_rms": 1.3612572084537276,
+      "instruction_rms": 1.3612569504993548,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.1406744785416356e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 21172.212890625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3330913836466176.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.695481300354004,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0027934312820435,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.455192987842642e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2647451399815168e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4504523582210048.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.778195753073377e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.2569803759812608e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3330913836466176.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.689992904663086,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0026856660842896,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.455192987842642e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.3612658409548353,
+      "step": 336,
+      "teaching_rms": 1.3612572084537276,
+      "traffic_rms": 0.004920497919811268
+    },
+    {
+      "batch_loss": 4.536196708679199,
+      "innovation_rms": 1.1568673142634316,
+      "instruction_rms": 1.1568668271896962,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.101625323081171e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 20040.44140625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2983270593593344.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6972432136535645,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9430446624755859,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4581762721264435e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3052858731593728e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4054071170301952.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.789508353333658e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1312823598579712e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2983270593593344.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.691755533218384,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9430726766586304,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4581762721264435e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.1568776528200302,
+      "step": 337,
+      "teaching_rms": 1.1568673142634316,
+      "traffic_rms": 0.005003780651115537
+    },
+    {
+      "batch_loss": 3.903050184249878,
+      "innovation_rms": 1.115619462543129,
+      "instruction_rms": 1.115616711691902,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.0966250190759526e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 19152.759765625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2670361824985088.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.698962450027466,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1659021377563477,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4608467109924045e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.341772479332352e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3648663838523392.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.799689831006863e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0181541023973376e+16,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2670361824985088.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6934752464294434,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.165930151939392,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4608467109924045e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 1.1156280585308462,
+      "step": 338,
+      "teaching_rms": 1.115619462543129,
+      "traffic_rms": 0.005077874459792762
+    },
+    {
+      "batch_loss": 4.258658409118652,
+      "innovation_rms": 0.929005831079437,
+      "instruction_rms": 0.9290045912649745,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.1248929132702597e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 18539.900390625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2388716928630784.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7005956172943115,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4109281301498413,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.463235400003748e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3746104034131968e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3283797239922688.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.808852886034842e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 9163386814201856.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2388716928630784.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6951088905334473,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4109561443328857,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.463235400003748e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.9290183099565517,
+      "step": 339,
+      "teaching_rms": 0.929005831079437,
+      "traffic_rms": 0.005119375900198495
+    },
+    {
+      "batch_loss": 3.5544490814208984,
+      "innovation_rms": 0.8278712507507313,
+      "instruction_rms": 0.8278674291883269,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.1819812062518313e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 18171.15234375,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2135212792217600.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7020437717437744,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3816676139831543,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4653706859446272e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4041646102478848e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2955417462243328.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.817099910437929e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 8247048025407488.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2135212792217600.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6965577602386475,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3816957473754883,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4653706859446272e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.8278831971008139,
+      "step": 340,
+      "teaching_rms": 0.8278712507507313,
+      "traffic_rms": 0.005173980067696939
+    },
+    {
+      "batch_loss": 3.7154829502105713,
+      "innovation_rms": 0.8308643602310706,
+      "instruction_rms": 0.8308612797935344,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.2614489589051556e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 18018.96484375,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1907037688561664.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.703252077102661,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9007612466812134,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4672776514240512e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4307633427120128e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2659875662331904.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.824522301120184e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 7422342900744192.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1907037688561664.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.6977665424346924,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.9005942344665527,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4672776514240512e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.8308769840785208,
+      "step": 341,
+      "teaching_rms": 0.8308643602310706,
+      "traffic_rms": 0.005204013959066228
+    },
+    {
+      "batch_loss": 4.065582752227783,
+      "innovation_rms": 0.8688396148864103,
+      "instruction_rms": 0.8688427071360576,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.348662770976162e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 17997.875,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1701661076619264.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7043190002441406,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1878403425216675,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4689793174667264e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.454702201929728e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2393888069255168.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.831202521453691e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6680108664356864.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1701661076619264.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.69883394241333,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.1876733303070068,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4689793174667264e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.8688584973902671,
+      "step": 342,
+      "teaching_rms": 0.8688396148864103,
+      "traffic_rms": 0.005170583604762769
+    },
+    {
+      "batch_loss": 4.02751350402832,
+      "innovation_rms": 0.9288514519783345,
+      "instruction_rms": 0.9288482584558048,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.4601573732012524e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 18299.630859375,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1516805147328512.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.705252170562744,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.165120840072632,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4704961281169818e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4762471537508352e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2154499208642560.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.837214788473324e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6012097905295360.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1516805147328512.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.699767589569092,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.165116310119629,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4704961281169818e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.928862122357023,
+      "step": 343,
+      "teaching_rms": 0.9288514519783345,
+      "traffic_rms": 0.005142927808215437
+    },
+    {
+      "batch_loss": 3.1495614051818848,
+      "innovation_rms": 0.89925628716287,
+      "instruction_rms": 0.8992518954560195,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.5661016407293297e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 18785.845703125,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1350419590676480.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7060649394989014,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8506603240966797,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.471846465834844e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4956376426020864e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1939049287778304.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.842625760071516e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 5410887846330368.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1350419590676480.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7005810737609863,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.8506557941436768,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.471846465834844e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.8992662313063784,
+      "step": 344,
+      "teaching_rms": 0.89925628716287,
+      "traffic_rms": 0.0051176283864204745
+    },
+    {
+      "batch_loss": 3.198390483856201,
+      "innovation_rms": 0.9162048225248365,
+      "instruction_rms": 0.916202637838497,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.658398220177637e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 19448.6015625,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1200659181338624.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7067620754241943,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3072474002838135,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.473047166892114e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.513089061093376e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1745144265048064.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.847495222193029e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4869798739574784.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1200659181338624.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7012786865234375,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3072429895401,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.473047166892114e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.9162167541227076,
+      "step": 345,
+      "teaching_rms": 0.9162048225248365,
+      "traffic_rms": 0.005122446695848414
+    },
+    {
+      "batch_loss": 3.358837604522705,
+      "innovation_rms": 0.9807668884644559,
+      "instruction_rms": 0.9807617610851475,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.737354150168232e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 20270.484375,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1065862773604352.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.707362413406372,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8856156468391418,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4741130059762893e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5287953269981184e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1570629811699712.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.851878150419251e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4382818704556032.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1065862773604352.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.7050750255584717,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8854486346244812,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4741130059762893e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.9807745215734704,
+      "step": 346,
+      "teaching_rms": 0.9807668884644559,
+      "traffic_rms": 0.005086486559686008
+    },
+    {
+      "batch_loss": 3.1672120094299316,
+      "innovation_rms": 0.9325493952199383,
+      "instruction_rms": 0.9325489137782271,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.793747551801049e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 21098.01953125,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 944535316201472.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.763143539428711,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8598737716674805,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4750575551840256e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.542931030736896e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1413566749999104.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.855822648383898e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3944536619352064.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 944535316201472.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.767564535140991,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8597067594528198,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4750575551840256e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.9325624204673469,
+      "step": 347,
+      "teaching_rms": 0.9325493952199383,
+      "traffic_rms": 0.005019390622966128
+    },
+    {
+      "batch_loss": 3.176997184753418,
+      "innovation_rms": 0.8663215338617005,
+      "instruction_rms": 0.8663228960680956,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.832408029901095e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 21842.552734375,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 835331142189056.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.802187919616699,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0768303871154785,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4758928404237517e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5556531533643776e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1272209981046784.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.85937269655208e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3550082796355584.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 835331142189056.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.8066084384918213,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0766633749008179,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4758928404237517e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.86633761176817,
+      "step": 348,
+      "teaching_rms": 0.8663215338617005,
+      "traffic_rms": 0.005036069477023855
+    },
+    {
+      "batch_loss": 3.7028298377990723,
+      "innovation_rms": 0.8680103288687832,
+      "instruction_rms": 0.868011349973286,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.863447792909025e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 22524.423828125,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 737039070789632.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.9399969577789307,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.378089189529419,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4766298568117453e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5671029993046016e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1144988888989696.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.862567465025536e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3195074489876480.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 737039070789632.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2.9444169998168945,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.3780848979949951,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4766298568117453e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.868026105642608,
+      "step": 349,
+      "teaching_rms": 0.8680103288687832,
+      "traffic_rms": 0.0050931430414438905
+    },
+    {
+      "batch_loss": 3.554030179977417,
+      "innovation_rms": 0.8077744327429904,
+      "instruction_rms": 0.8077713302288739,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.878906926395556e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 22990.53125,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 648568851398656.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.0438220500946045,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.038251280784607,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4772783968734413e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5774079143378944e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1030489959825408.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.865443375126938e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2875567040888832.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 648568851398656.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.04824161529541,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.0382468700408936,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4772783968734413e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.8077875012188334,
+      "step": 350,
+      "teaching_rms": 0.8077744327429904,
+      "traffic_rms": 0.005138505512529699
+    },
+    {
+      "batch_loss": 3.2409234046936035,
+      "innovation_rms": 0.8653150152326538,
+      "instruction_rms": 0.8653067097844316,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.885059243708776e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 23233.673828125,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 568939151097856.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.124375820159912,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8055378198623657,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4778473941408154e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5866823593426944e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 927440943710208.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.868031350620815e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2588010390487040.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 568939151097856.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1287949085235596,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.8055334091186523,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4778473941408154e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.8653216344876623,
+      "step": 351,
+      "teaching_rms": 0.8653150152326538,
+      "traffic_rms": 0.00517447232292819
+    },
+    {
+      "batch_loss": 2.6862165927886963,
+      "innovation_rms": 0.6994373255213455,
+      "instruction_rms": 0.6994327128574601,
+      "parameter_state": {
+        "bn_running_mean": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 4.874316465349591e+18,
+          "nonfinite_indices": []
+        },
+        "bn_running_var": {
+          "all_finite": false,
+          "max_abs_over_finite_tensors": 23222.82421875,
+          "nonfinite_indices": [
+            15,
+            16
+          ]
+        },
+        "feedback_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 497266750128128.0,
+          "nonfinite_indices": []
+        },
+        "feedback_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.1864097118377686,
+          "nonfinite_indices": []
+        },
+        "feedback_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6203396320343018,
+          "nonfinite_indices": []
+        },
+        "feedback_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4783445795550003e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.5950293061599232e+16,
+          "nonfinite_indices": []
+        },
+        "forward_bn_bias_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 834696829206528.0,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 6.870360253687398e+17,
+          "nonfinite_indices": []
+        },
+        "forward_bn_scale_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 2329209351438336.0,
+          "nonfinite_indices": []
+        },
+        "forward_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 497266750128128.0,
+          "nonfinite_indices": []
+        },
+        "forward_readout": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 3.190828323364258,
+          "nonfinite_indices": []
+        },
+        "forward_readout_momentum": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.6203352808952332,
+          "nonfinite_indices": []
+        },
+        "forward_weight": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 1.4783445795550003e+17,
+          "nonfinite_indices": []
+        },
+        "predictor_bias": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.011661699041724205,
+          "nonfinite_indices": []
+        },
+        "predictor_slope": {
+          "all_finite": true,
+          "max_abs_over_finite_tensors": 0.11585475504398346,
+          "nonfinite_indices": []
+        }
+      },
+      "predictor_mse": null,
+      "predictor_updated": false,
+      "raw_apical_rms": 0.6994511373069514,
+      "step": 352,
+      "teaching_rms": 0.6994373255213455,
+      "traffic_rms": 0.005097079923921986
+    }
+  ],
+  "trajectory_metrics_finite": true,
+  "validation_evaluations": 0
+}