experiment: enumerate Transformer crossover work

2 files changed, 95 insertions, 5 deletions

diff --git a/TRANSFORMER_CROSSOVER.md b/TRANSFORMER_CROSSOVER.md
index 2bcff12..df0cf0f 100644
--- a/TRANSFORMER_CROSSOVER.md
+++ b/TRANSFORMER_CROSSOVER.md
@@ -186,7 +186,13 @@ allocated/reserved memory, forward/feedback parameters, ordinary training
 and validation tokens, positive/negative or perturbed presentations,
 candidate-token presentations, relaxation-token passes, local-VJP token
 evaluations, first nonfinite step, and completed optimizer steps.  The
-accuracy--cost plot and the within-family depth plot retain timeouts and
+hardware-independent ledger additionally enumerates full-forward token
+passes, extra local block/head evaluations, block affine MACs per token,
+causal-attention MACs per token, and whole-forward MACs per token.  These
+primitive counts remain separate for DP/EP because a local VJP is not
+silently priced as an ordinary forward.
+
+The accuracy--cost plot and the within-family depth plot retain timeouts and
 failures.  A low-cost method is not declared scalable from depth 4 alone, and
 an accurate DP/EP point is not declared local in the weight-transport-free
 sense merely because its VJPs were evaluated at detached edges.
diff --git a/experiments/transformer_crossover_native.py b/experiments/transformer_crossover_native.py
index 43d3e62..7ccb8fb 100644
--- a/experiments/transformer_crossover_native.py
+++ b/experiments/transformer_crossover_native.py
@@ -141,10 +141,11 @@ def train_step(net, optimizer, tokens, targets, args, step, generators):
         metrics = {"loss": float(loss.detach())}
     elif args.method == "dfa":
         with torch.no_grad():
-            loss = net(tokens, targets)["loss"]
+            cache = net(tokens, targets, return_cache=True)
+            loss = cache["loss"]
         metrics = {
             "loss": float(loss),
-            **net.dfa_gradients(tokens, targets),
+            **net.dfa_gradients(tokens, targets, cache=cache),
         }
     elif args.method == "pepita":
         local = net.pepita_gradients(tokens, targets)
@@ -270,7 +271,9 @@ def hardware_report(device):
     return report
 
 
-def work_report(net, args, train_tokens, validation_evaluations):
+def work_report(
+        net, args, train_tokens, completed_steps,
+        validation_evaluations):
     ordinary_validation = sum(
         row["tokens"] for row in validation_evaluations)
     presentations = train_tokens
@@ -291,6 +294,70 @@ def work_report(net, args, train_tokens, validation_evaluations):
             train_tokens * (args.ep_free_steps + args.ep_nudge_steps)
             + sum(row["relaxation_token_passes"]
                   for row in validation_evaluations))
+    width = args.width
+    block_affine_macs_per_token = (
+        (4 + 2 * args.mlp_ratio) * width * width)
+    block_attention_macs_per_token = (
+        2 * args.context_length * width)
+    forward_affine_macs_per_token = (
+        args.depth * block_affine_macs_per_token
+        + width * net.config.vocab_size)
+    forward_attention_macs_per_token = (
+        args.depth * block_attention_macs_per_token)
+    training_forward_multiplier = {
+        "bp": 1,
+        "fa": 1,
+        "dfa": 1,
+        "pepita": 2,
+        "ff": 2,
+        "ep": 2,
+        "dualprop": 1,
+        "clean_kp": 1,
+        "sdil": 1,
+    }[args.method]
+    if args.method == "ff":
+        validation_forward_tokens = sum(
+            row["candidate_token_presentations"]
+            for row in validation_evaluations)
+    else:
+        validation_forward_tokens = ordinary_validation
+    full_forward_token_passes = (
+        training_forward_multiplier * train_tokens
+        + validation_forward_tokens)
+
+    local_block_token_evaluations = 0
+    local_head_token_evaluations = 0
+    if args.method in {"dfa", "pepita"}:
+        local_block_token_evaluations = args.depth * train_tokens
+        local_head_token_evaluations = train_tokens
+    elif args.method == "ff":
+        steps_per_layer = math.ceil(
+            args.train_steps / net.ff_num_layers)
+        block_steps = 0
+        head_steps = 0
+        for step in range(completed_steps):
+            layer = min(
+                step // steps_per_layer, net.ff_num_layers - 1)
+            block_steps += int(1 <= layer <= args.depth)
+            head_steps += int(layer == args.depth + 1)
+        tokens_per_step = args.batch_size * args.context_length
+        local_block_token_evaluations = block_steps * tokens_per_step
+        local_head_token_evaluations = head_steps * tokens_per_step
+    elif args.method == "dualprop":
+        multiplier = 2 * args.dp_inference_passes + 1
+        local_block_token_evaluations = (
+            args.depth * train_tokens * multiplier)
+        local_head_token_evaluations = train_tokens * multiplier
+    elif args.method == "ep":
+        train_multiplier = (
+            2 * (args.ep_free_steps + args.ep_nudge_steps) + 2)
+        validation_multiplier = 2 * args.ep_free_steps
+        local_block_token_evaluations = args.depth * (
+            train_tokens * train_multiplier
+            + ordinary_validation * validation_multiplier)
+        local_head_token_evaluations = (
+            train_tokens * train_multiplier
+            + ordinary_validation * validation_multiplier)
     return {
         "ordinary_training_tokens": train_tokens,
         "ordinary_validation_tokens": ordinary_validation,
@@ -300,6 +367,22 @@ def work_report(net, args, train_tokens, validation_evaluations):
             for row in validation_evaluations),
         "relaxation_token_passes": relaxation,
         "local_vjp_token_evaluations": local_vjps,
+        "full_forward_token_passes": full_forward_token_passes,
+        "local_block_token_evaluations":
+            local_block_token_evaluations,
+        "local_head_token_evaluations": local_head_token_evaluations,
+        "block_affine_macs_per_token":
+            block_affine_macs_per_token,
+        "block_attention_macs_per_token":
+            block_attention_macs_per_token,
+        "forward_affine_macs_per_token":
+            forward_affine_macs_per_token,
+        "forward_attention_macs_per_token":
+            forward_attention_macs_per_token,
+        "enumerated_full_forward_macs": (
+            full_forward_token_passes
+            * (forward_affine_macs_per_token
+               + forward_attention_macs_per_token)),
         "forward_parameter_count": net.n_forward_parameters,
         "feedback_parameter_count": net.n_feedback_parameters,
         "completed_optimizer_steps": len(
@@ -434,7 +517,8 @@ def run(args):
         },
         "final": final,
         "work": work_report(
-            net, args, train_tokens, validation_evaluations),
+            net, args, train_tokens, final["step"],
+            validation_evaluations),
         "hardware": {
             **hardware_report(args.device),
             "peak_memory_allocated_bytes": peak_allocated,