From 3f7cc1f54420a76a61b9d011691582acee0dcd76 Mon Sep 17 00:00:00 2001
From: wuyang <5700876+banisherwy@user.noreply.gitee.com>
Date: Thu, 30 Jul 2026 07:54:38 +0800
Subject: [PATCH] research: align AttnRes residual precision

---
 experiments/k3/attnres_gradient/train.py | 32 ++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/experiments/k3/attnres_gradient/train.py b/experiments/k3/attnres_gradient/train.py
index f95e7c4..9e85504 100644
--- a/experiments/k3/attnres_gradient/train.py
+++ b/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()