Commit adbe67a
feat(dpa4): fold the focus competition norm into the fused value kernel
The fused SO(2) value-path training operator previously required an
identity competition norm; a real norm (the focus entry of edge_norm)
dropped the block to the unfused paths at roughly twice the training
memory. The per-focus RMS norm now runs inside the operator: its
learnable scales enter as the norm_scale input and follow the same
input/gradient/second-gradient pattern as the projection weight.
The forward folds the scales into the head projection and rescales the
logit by the inverse RMS of the gate row in the same lane-strided pass;
the backward pushes the logit gradient through the norm Jacobian in the
per-edge competition kernel and contracts the scale gradient on the
host; the second order derives the logit cotangent, the parameter
curvature and the Jacobian's own gate dependence in closed form, with
the head-Hessian term mapped onto the rotation operands by one rotation
backward. The identity norm keeps the original code paths.
Verified by the parity suite (19 cases, including the float64
structural bound on the norm shapes) and a Neo benchmark: the
norm-enabled block now matches the film-only fused step time within
run-to-run noise at 43% less peak memory than the unfused fallback.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>1 parent d92b1ac commit adbe67a
5 files changed
Lines changed: 479 additions & 145 deletions
File tree
- deepmd/pt_expt/kernels/cuda/dpa4
- source
- op/pt/dpa4
- so2_conv_train
- tests/pt_expt/kernels
0 commit comments