stage2-grad-flow-viz

Stage 2 · Sanity · Gradient flow visualization

A few-second-per-step diagnostic that surfaces dead layers, exploding gradients, and depth-scaling pathologies before they cause silent training failures.

Stage question

"Are gradients flowing through every layer at the right scale?"

What to measure

For each parameter (or each named module), at the first training step (or every K steps as a watchdog), record:

Quantity	What it tells you
‖p.grad‖₂ (per parameter)	Whether the layer is learning at all. Zero → dead layer.
‖p.grad‖₂ / ‖p‖₂ (per parameter)	Relative update size. Healthy = O(1e-3 to 1e-1).
mean(activation), std(activation) (per module forward output)	Whether activations stay O(1) across depth.
`max(	activation

Plot or table these by layer index. Two visual signals are diagnostic:

1. Monotonic vanishing: ‖p.grad‖ shrinks by orders of magnitude with depth. Likely missing residual connections, or wrong init.
2. Monotonic exploding: ‖p.grad‖ grows with depth. Likely missing layer norms, or activation function with no bounded output.

A healthy modern transformer has activation std stable across depth (the "coord check" property in muP).

The coord check (muP-style)

For models that intend to use muP for scale transfer (stage3-mup-coord-check), the per-coordinate statistics must be width-invariant. Specifically: training a width-256 model and a width-1024 model with muP should give the same activation std at every layer index, after the first few steps.

If this is not the case, muP transfer will not work. The grad flow visualization is the diagnostic that surfaces the problem.

Recommended hook

class GradFlowProbe:
    def __init__(self, model):
        self.records = []
        for name, p in model.named_parameters():
            p.register_hook(lambda g, n=name: self._record_grad(n, g))
        for name, m in model.named_modules():
            m.register_forward_hook(
                lambda mod, inp, out, n=name: self._record_act(n, out)
            )

    def _record_grad(self, name, g):
        self.records.append(("grad", name, g.norm().item()))

    def _record_act(self, name, out):
        if isinstance(out, torch.Tensor):
            self.records.append(("act", name, out.std().item(), out.mean().item()))

Implementations live at template/curry_train/infra/preflight.py:probe_grad_flow.

Procedure when assisting a user

1. Wire up the probe to a single training step (forward + backward only — no optimizer step needed).

2. Group results by layer index. For a transformer, "layer index" is the depth of each block.

3. Plot grad-norm vs layer index on a log scale. Plot activation std vs layer index on a linear scale.

4. Diagnosis:

   • Flat-line at zero for some layer → dead layer. Check that the layer was actually included in optimizer.param_groups. See stage1-preflight-asserts assert 8.
   • Vanishing toward output → missing residual connection in deeper blocks.
   • Exploding toward output → missing post-block layer norm.
   • Activation std growing 10× per layer → init scale too large; consider Xavier or Kaiming with nn.init overrides.

5. If user is targeting muP, confirm activation std is flat across depth within a single width, and flat across widths after a few steps. If not, the model is not muP-compliant.

Boundaries

• This is a one-shot or low-frequency diagnostic, not a continuous logger. Running this every step adds non-trivial overhead.
• Activation hooks may not be representative if the model uses gradient checkpointing — hooks fire during recomputation as well. Disable checkpointing for the probe pass.
• This skill does not fix the architecture; it only surfaces problems. Architecture changes are a Stage 1 concern.

Common failure modes surfaced

• Missing nn.LayerNorm after attention → activation std drift with depth.
• Residual connection accidentally removed → vanishing gradients.
• Too-large embedding init → first-layer activations dominate.
• Frozen layer (requires_grad=False) leaked into a "learning" config → dead layer.

Related

• skills/stage1-preflight-asserts (assert 5).
• skills/stage2-overfit-single-batch — should pass if grad flow is healthy.
• skills/stage3-mup-coord-check — formal muP coord-check, builds on this skill.
• Maxim Naumov's blog post on activation diagnostics.