Shorten backward flow site constants

dlwh · dlwh · commit 265073a929ca · 2026-04-22T14:34:19.000-07:00
diff --git a/docs/design/grug-backward-flow-logging.md b/docs/design/grug-backward-flow-logging.md
@@ -49,8 +49,7 @@ canonical Grug base template.
 - `trace_backward_activation(x, name, site=...)`: a convenience wrapper for
   identity-only stream anchors that adds a `jax.named_scope(name)` around
   `log_backward_activation(...)`
-- `BACKWARD_FLOW_SITE_IN` / `BACKWARD_FLOW_SITE_OUT`: named constants for the
-  metric-key site labels
+- `BWD_IN` / `BWD_OUT`: named constants for the metric-key site labels
 - `normalize_name_stack(...)`: removes transform wrappers such as `jvp(...)` and
   `transpose(...)` so metric keys stay stable
 
@@ -65,20 +64,20 @@ def _tagged_identity(metric_prefix: str, site: BackwardFlowSite, x: jax.Array) -
 
 def _tagged_identity_fwd(metric_prefix: str, site: BackwardFlowSite, x: jax.Array):
     levanter.tracker.jit_log(
-        _tensor_metrics(metric_prefix, x, site=site, kind=BACKWARD_FLOW_KIND_ACTIVATION),
+        _tensor_metrics(metric_prefix, x, site=site, kind=_BWD_KIND_ACTIVATION),
         step=None,
     )
     return x, None
 
 def _tagged_identity_bwd(metric_prefix: str, site: BackwardFlowSite, _residual: None, cotangent: jax.Array):
     levanter.tracker.jit_log(
-        _tensor_metrics(metric_prefix, cotangent, site=site, kind=BACKWARD_FLOW_KIND_GRADIENT),
+        _tensor_metrics(metric_prefix, cotangent, site=site, kind=_BWD_KIND_GRADIENT),
         step=None,
     )
     return (cotangent,)
 
 def log_backward_activation(
-    x: jax.Array, *, site: BackwardFlowSite = BACKWARD_FLOW_SITE_OUT
+    x: jax.Array, *, site: BackwardFlowSite = BWD_OUT
 ) -> jax.Array:
     context = _ACTIVE_CONTEXT.get()
     if context is None:
@@ -87,7 +86,7 @@ def log_backward_activation(
     return _tagged_identity(f"{context.prefix}/{name_stack}", site, x)
 
 def trace_backward_activation(
-    x: jax.Array, name: str, *, site: BackwardFlowSite = BACKWARD_FLOW_SITE_OUT
+    x: jax.Array, name: str, *, site: BackwardFlowSite = BWD_OUT
 ) -> jax.Array:
     with jax.named_scope(name):
         return log_backward_activation(x, site=site)
diff --git a/docs/recipes/add_grug_backward_flow_logging.md b/docs/recipes/add_grug_backward_flow_logging.md
@@ -41,14 +41,14 @@ explicitly when a variant should opt out. Positive intervals sample that often.
 
 At each named module boundary you want in the graph, wrap the returned activation with
 `log_backward_activation(...)`. For modules where you want to see what backward is
-sending *into* the module, mark the input with `BACKWARD_FLOW_SITE_IN`:
+sending *into* the module, mark the input with `BWD_IN`:
 
 ```python
-from levanter.analysis.backward_flow import BACKWARD_FLOW_SITE_IN, log_backward_activation, trace_backward_activation
+from levanter.analysis.backward_flow import BWD_IN, log_backward_activation, trace_backward_activation
 
 @named_call
 def __call__(self, x):
-    x = log_backward_activation(x, site=BACKWARD_FLOW_SITE_IN)
+    x = log_backward_activation(x, site=BWD_IN)
     out = ...
     return log_backward_activation(out)
 ```
diff --git a/experiments/grug/base/model.py b/experiments/grug/base/model.py
@@ -15,7 +15,7 @@
 from jaxtyping import Array, Float, Int, PRNGKeyArray
 
 from levanter.analysis.backward_flow import (
-    BACKWARD_FLOW_SITE_IN,
+    BWD_IN,
     is_backward_flow_active,
     log_backward_activation,
     trace_backward_activation,
@@ -82,7 +82,7 @@ def init(cfg: GrugModelConfig, *, key: PRNGKeyArray) -> "CausalSelfAttention":
 
     @named_call
     def __call__(self, x: Float[Array, "B S D"], mask: AttentionMask | jax.Array) -> Float[Array, "B S D"]:
-        x = log_backward_activation(x, site=BACKWARD_FLOW_SITE_IN)
+        x = log_backward_activation(x, site=BWD_IN)
         head_dim = self.cfg.inferred_head_dim
         seq_len = x.shape[1]
 
@@ -111,7 +111,7 @@ def init(cfg: GrugModelConfig, *, key: PRNGKeyArray) -> "MLP":
 
     @named_call
     def __call__(self, x: Float[Array, "B S D"]) -> Float[Array, "B S D"]:
-        x = log_backward_activation(x, site=BACKWARD_FLOW_SITE_IN)
+        x = log_backward_activation(x, site=BWD_IN)
         up = jnp.einsum("bsh,hm->bsm", x, self.mlp_up)
         activated = jax.nn.relu(up)
         out = jnp.einsum("bsm,mh->bsh", activated, self.mlp_down, out_sharding=Pbatch)
diff --git a/lib/levanter/src/levanter/analysis/__init__.py b/lib/levanter/src/levanter/analysis/__init__.py
@@ -5,14 +5,11 @@
     "BackwardFlowConfig",
     "BackwardFlowEdge",
     "BackwardFlowGraph",
-    "BACKWARD_FLOW_KIND_ACTIVATION",
-    "BACKWARD_FLOW_KIND_GRADIENT",
     "BackwardFlowPlate",
     "BackwardFlowRenderHints",
-    "BACKWARD_FLOW_SITE_IN",
-    "BACKWARD_FLOW_SITE_OUT",
     "BackwardFlowSite",
-    "BackwardFlowTensorKind",
+    "BWD_IN",
+    "BWD_OUT",
     "SummaryStats",
     "cb_compute_entropies",
     "cb_compute_top2_gap",
@@ -38,14 +35,11 @@
     BackwardFlowConfig,
     BackwardFlowEdge,
     BackwardFlowGraph,
-    BACKWARD_FLOW_KIND_ACTIVATION,
-    BACKWARD_FLOW_KIND_GRADIENT,
     BackwardFlowPlate,
     BackwardFlowRenderHints,
-    BACKWARD_FLOW_SITE_IN,
-    BACKWARD_FLOW_SITE_OUT,
     BackwardFlowSite,
-    BackwardFlowTensorKind,
+    BWD_IN,
+    BWD_OUT,
     SummaryStats,
     backward_flow_graph_from_jaxpr,
     backward_flow_node_stats,
diff --git a/lib/levanter/src/levanter/analysis/backward_flow.py b/lib/levanter/src/levanter/analysis/backward_flow.py
@@ -42,13 +42,13 @@
 _NAME_STACK_PART_RE = re.compile(r"^(?P<wrapper>[A-Za-z_][A-Za-z0-9_]*)\((?P<inner>.*)\)$")
 _STAT_NAMES = ("norm", "rms", "rms_scaled", "mean_abs", "max_abs", "max_abs_scaled", "finite_fraction")
 BackwardFlowSite: TypeAlias = Literal["in", "out"]
-BackwardFlowTensorKind: TypeAlias = Literal["activation", "gradient"]
-BACKWARD_FLOW_SITE_IN: BackwardFlowSite = "in"
-BACKWARD_FLOW_SITE_OUT: BackwardFlowSite = "out"
-BACKWARD_FLOW_KIND_ACTIVATION: BackwardFlowTensorKind = "activation"
-BACKWARD_FLOW_KIND_GRADIENT: BackwardFlowTensorKind = "gradient"
-_FLOW_SITES = (BACKWARD_FLOW_SITE_IN, BACKWARD_FLOW_SITE_OUT)
-_TENSOR_KINDS = (BACKWARD_FLOW_KIND_ACTIVATION, BACKWARD_FLOW_KIND_GRADIENT)
+_BackwardFlowTensorKind: TypeAlias = Literal["activation", "gradient"]
+BWD_IN: BackwardFlowSite = "in"
+BWD_OUT: BackwardFlowSite = "out"
+_BWD_KIND_ACTIVATION: _BackwardFlowTensorKind = "activation"
+_BWD_KIND_GRADIENT: _BackwardFlowTensorKind = "gradient"
+_FLOW_SITES = (BWD_IN, BWD_OUT)
+_TENSOR_KINDS = (_BWD_KIND_ACTIVATION, _BWD_KIND_GRADIENT)
 _FLOW_DIRECTIONS = ("tb", "lr")
 _DEFAULT_PREFIX = "backward_flow"
 _DEFAULT_RESIDUAL_GAIN_HORIZON = 50
@@ -202,7 +202,7 @@ def normalize_name_stack(name_stack: str) -> str:
     return "/".join(parts)
 
 
-def log_backward_activation(x: jax.Array, *, site: BackwardFlowSite = BACKWARD_FLOW_SITE_OUT) -> jax.Array:
+def log_backward_activation(x: jax.Array, *, site: BackwardFlowSite = BWD_OUT) -> jax.Array:
     """Return ``x`` unchanged while logging activation and backward-gradient scale when enabled."""
     context = _ACTIVE_CONTEXT.get()
     if context is None:
@@ -219,9 +219,7 @@ def log_backward_activation(x: jax.Array, *, site: BackwardFlowSite = BACKWARD_F
     return _tagged_identity_with_scale(f"{context.prefix}/{name_stack}", site, context.gradient_scale, x)
 
 
-def trace_backward_activation(
-    x: jax.Array, name: str, *, site: BackwardFlowSite = BACKWARD_FLOW_SITE_OUT
-) -> jax.Array:
+def trace_backward_activation(x: jax.Array, name: str, *, site: BackwardFlowSite = BWD_OUT) -> jax.Array:
     """Return ``x`` unchanged while logging under an extra JAX named scope."""
     if not name:
         raise ValueError("name must be non-empty")
@@ -236,18 +234,14 @@ def _tagged_identity(metric_prefix: str, site: BackwardFlowSite, x: jax.Array) -
 
 
 def _tagged_identity_fwd(metric_prefix: str, site: BackwardFlowSite, x: jax.Array) -> tuple[jax.Array, None]:
-    levanter.tracker.jit_log(
-        _tensor_metrics(metric_prefix, x, site=site, kind=BACKWARD_FLOW_KIND_ACTIVATION), step=None
-    )
+    levanter.tracker.jit_log(_tensor_metrics(metric_prefix, x, site=site, kind=_BWD_KIND_ACTIVATION), step=None)
     return x, None
 
 
 def _tagged_identity_bwd(
     metric_prefix: str, site: BackwardFlowSite, _residual: None, cotangent: jax.Array
 ) -> tuple[jax.Array]:
-    levanter.tracker.jit_log(
-        _tensor_metrics(metric_prefix, cotangent, site=site, kind=BACKWARD_FLOW_KIND_GRADIENT), step=None
-    )
+    levanter.tracker.jit_log(_tensor_metrics(metric_prefix, cotangent, site=site, kind=_BWD_KIND_GRADIENT), step=None)
     return (cotangent,)
 
 
@@ -264,9 +258,7 @@ def _tagged_identity_with_scale(
 def _tagged_identity_with_scale_fwd(
     metric_prefix: str, site: BackwardFlowSite, gradient_scale: jax.Array, x: jax.Array
 ) -> tuple[jax.Array, jax.Array]:
-    levanter.tracker.jit_log(
-        _tensor_metrics(metric_prefix, x, site=site, kind=BACKWARD_FLOW_KIND_ACTIVATION), step=None
-    )
+    levanter.tracker.jit_log(_tensor_metrics(metric_prefix, x, site=site, kind=_BWD_KIND_ACTIVATION), step=None)
     return x, gradient_scale
 
 
@@ -281,7 +273,7 @@ def _tagged_identity_with_scale_bwd(
             metric_prefix,
             cotangent,
             site=site,
-            kind=BACKWARD_FLOW_KIND_GRADIENT,
+            kind=_BWD_KIND_GRADIENT,
             gradient_scale=gradient_scale,
         ),
         step=None,
@@ -519,12 +511,12 @@ def _tensor_metrics(
     tensor: jax.Array,
     *,
     site: BackwardFlowSite,
-    kind: BackwardFlowTensorKind,
+    kind: _BackwardFlowTensorKind,
     gradient_scale: jax.Array | None = None,
 ) -> dict[str, jax.Array]:
     summary = SummaryStats.from_tensor(tensor)
     metrics = summary.to_metrics(f"{metric_prefix}/{site}_{kind}")
-    if kind == BACKWARD_FLOW_KIND_GRADIENT and gradient_scale is not None:
+    if kind == _BWD_KIND_GRADIENT and gradient_scale is not None:
         gradient_scale = jnp.asarray(gradient_scale, dtype=jnp.float32)
         metrics[f"{metric_prefix}/{site}_{kind}_rms_scaled"] = summary.rms * gradient_scale
         metrics[f"{metric_prefix}/{site}_{kind}_max_abs_scaled"] = summary.max_abs * gradient_scale
@@ -1116,57 +1108,47 @@ def _is_supported_metric_name(metric_name: str) -> bool:
 
 
 def _metric_value(
-    stats: Mapping[str, float], site: BackwardFlowSite, kind: BackwardFlowTensorKind, metric: str
+    stats: Mapping[str, float], site: BackwardFlowSite, kind: _BackwardFlowTensorKind, metric: str
 ) -> float | None:
     return stats.get(f"{site}_{kind}_{metric}")
 
 
 def _preferred_gradient_rms(stats: Mapping[str, float]) -> float | None:
-    scaled = _preferred_metric(stats, BACKWARD_FLOW_KIND_GRADIENT, "rms_scaled", preferred_site=BACKWARD_FLOW_SITE_IN)
+    scaled = _preferred_metric(stats, _BWD_KIND_GRADIENT, "rms_scaled", preferred_site=BWD_IN)
     if scaled is not None:
         return scaled
-    return _preferred_metric(stats, BACKWARD_FLOW_KIND_GRADIENT, "rms", preferred_site=BACKWARD_FLOW_SITE_IN)
+    return _preferred_metric(stats, _BWD_KIND_GRADIENT, "rms", preferred_site=BWD_IN)
 
 
 def _has_scaled_gradient_rms(stats: Mapping[str, float]) -> bool:
-    return any(
-        _metric_value(stats, site, BACKWARD_FLOW_KIND_GRADIENT, "rms_scaled") is not None for site in _FLOW_SITES
-    )
+    return any(_metric_value(stats, site, _BWD_KIND_GRADIENT, "rms_scaled") is not None for site in _FLOW_SITES)
 
 
 def _preferred_gradient_max_abs(stats: Mapping[str, float]) -> float | None:
-    scaled = _preferred_metric(
-        stats, BACKWARD_FLOW_KIND_GRADIENT, "max_abs_scaled", preferred_site=BACKWARD_FLOW_SITE_IN
-    )
+    scaled = _preferred_metric(stats, _BWD_KIND_GRADIENT, "max_abs_scaled", preferred_site=BWD_IN)
     if scaled is not None:
         return scaled
-    return _preferred_metric(stats, BACKWARD_FLOW_KIND_GRADIENT, "max_abs", preferred_site=BACKWARD_FLOW_SITE_IN)
+    return _preferred_metric(stats, _BWD_KIND_GRADIENT, "max_abs", preferred_site=BWD_IN)
 
 
 def _has_scaled_gradient_max_abs(stats: Mapping[str, float]) -> bool:
-    return any(
-        _metric_value(stats, site, BACKWARD_FLOW_KIND_GRADIENT, "max_abs_scaled") is not None for site in _FLOW_SITES
-    )
+    return any(_metric_value(stats, site, _BWD_KIND_GRADIENT, "max_abs_scaled") is not None for site in _FLOW_SITES)
 
 
 def _preferred_activation_rms(stats: Mapping[str, float]) -> float | None:
-    return _preferred_metric(stats, BACKWARD_FLOW_KIND_ACTIVATION, "rms", preferred_site=BACKWARD_FLOW_SITE_OUT)
+    return _preferred_metric(stats, _BWD_KIND_ACTIVATION, "rms", preferred_site=BWD_OUT)
 
 
 def _preferred_finite_fraction(stats: Mapping[str, float]) -> float | None:
-    gradient_fraction = _preferred_metric(
-        stats, BACKWARD_FLOW_KIND_GRADIENT, "finite_fraction", preferred_site=BACKWARD_FLOW_SITE_IN
-    )
+    gradient_fraction = _preferred_metric(stats, _BWD_KIND_GRADIENT, "finite_fraction", preferred_site=BWD_IN)
     if gradient_fraction is not None:
         return gradient_fraction
-    return _preferred_metric(
-        stats, BACKWARD_FLOW_KIND_ACTIVATION, "finite_fraction", preferred_site=BACKWARD_FLOW_SITE_OUT
-    )
+    return _preferred_metric(stats, _BWD_KIND_ACTIVATION, "finite_fraction", preferred_site=BWD_OUT)
 
 
 def _preferred_metric(
     stats: Mapping[str, float],
-    kind: BackwardFlowTensorKind,
+    kind: _BackwardFlowTensorKind,
     metric: str,
     *,
     preferred_site: BackwardFlowSite,
diff --git a/lib/levanter/tests/test_backward_flow.py b/lib/levanter/tests/test_backward_flow.py
@@ -8,12 +8,12 @@
 
 import levanter.tracker
 from levanter.analysis.backward_flow import (
-    BACKWARD_FLOW_SITE_IN,
-    BACKWARD_FLOW_SITE_OUT,
     BackwardFlowConfig,
     BackwardFlowEdge,
     BackwardFlowGraph,
     BackwardFlowPlate,
+    BWD_IN,
+    BWD_OUT,
     SummaryStats,
     backward_flow_graph_from_jaxpr,
     capture_backward_flow,
@@ -46,8 +46,8 @@ def test_normalize_name_stack_strips_jax_transform_wrappers():
 def test_log_backward_activation_records_activation_and_gradient_metrics():
     @jax.named_call
     def inner(x):
-        x = log_backward_activation(x, site=BACKWARD_FLOW_SITE_IN)
-        return log_backward_activation(x * 2, site=BACKWARD_FLOW_SITE_OUT)
+        x = log_backward_activation(x, site=BWD_IN)
+        return log_backward_activation(x * 2, site=BWD_OUT)
 
     @jax.jit
     def compute_grad(x):
@@ -74,7 +74,7 @@ def compute_grad(x):
 def test_log_backward_activation_records_scaled_gradient_rms_when_configured():
     @jax.named_call
     def inner(x):
-        return log_backward_activation(x * 2, site=BACKWARD_FLOW_SITE_OUT)
+        return log_backward_activation(x * 2, site=BWD_OUT)
 
     @jax.jit
     def compute_grad(x):
@@ -119,8 +119,8 @@ def compute_grad(x):
 def test_log_backward_activation_allows_callers_to_skip_checkpoint_when_active():
     @jax.named_call
     def inner(x):
-        x = log_backward_activation(x, site=BACKWARD_FLOW_SITE_IN)
-        return log_backward_activation(jnp.tanh(x * 2), site=BACKWARD_FLOW_SITE_OUT)
+        x = log_backward_activation(x, site=BWD_IN)
+        return log_backward_activation(jnp.tanh(x * 2), site=BWD_OUT)
 
     def maybe_checkpointed_inner(x):
         if is_backward_flow_active():
@@ -151,8 +151,8 @@ def init(weight):
 
         def step(self, carry: jax.Array) -> jax.Array:
             with jax.named_scope("ArrayBlock"):
-                carry = log_backward_activation(carry, site=BACKWARD_FLOW_SITE_IN)
-                return log_backward_activation(jnp.tanh(carry * self.weight), site=BACKWARD_FLOW_SITE_OUT)
+                carry = log_backward_activation(carry, site=BWD_IN)
+                return log_backward_activation(jnp.tanh(carry * self.weight), site=BWD_OUT)
 
     def apply_layers(stack: ArrayStacked[Layer], carry: jax.Array) -> jax.Array:
         if is_backward_flow_active():
