research: align AttnRes residual precision

experiments/k3/attnres_gradient/train.py

@@ -602,6 +602,22 @@ def main() -> None:
         )
         gradient_gate = None
     diagnostics = [{"step": 0, **initial_diagnostic}]
+    print(
+        json.dumps(
+            {
+                "event": "diagnostic",
+                "step": 0,
+                "architecture": args.architecture,
+                "depth": args.depth,
+                "validation_bpc": evaluations[0]["bits_per_byte"],
+                "activation_gradient_cv": initial_diagnostic[
+                    "activation_grad_statistics"
+                ]["population_cv"],
+            },
+            sort_keys=True,
+        ),
+        flush=True,
+    )
     model.zero_grad(set_to_none=True)
 
     training_history: list[dict[str, float | int]] = []
@@ -664,6 +680,22 @@ def main() -> None:
                     ),
                 }
             )
+            print(
+                json.dumps(
+                    {
+                        "event": "diagnostic",
+                        "step": step,
+                        "architecture": args.architecture,
+                        "depth": args.depth,
+                        "validation_bpc": evaluations[-1]["bits_per_byte"],
+                        "activation_gradient_cv": diagnostics[-1][
+                            "activation_grad_statistics"
+                        ]["population_cv"],
+                    },
+                    sort_keys=True,
+                ),
+                flush=True,
+            )
             model.zero_grad(set_to_none=True)
             model.train()