You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
The mixed_type padding fix in #5738 added a masked (maskf is not None) branch to every term of the shared losses. Combined with the existing has_* / mse/mae / use_huber conditions, this produces deeply nested, duplicated code. The energy term alone is now four levels deep:
The same shape repeats for force, virial, atom_ener, atom_pref, gen_force, and again across ener_spin. The maskf fork roughly doubles each term, and the per-frame arithmetic is copy-pasted throughout. It also forces the top-level inv/_nf/_nloc locals whose else-branch dead stores CodeQL flagged.
Proposed refactor (numerics-preserving)
Every masked block is one of three reduction "idioms":
idiom 1 - per-atom masked mean over ncomp components (force, atom_ener, atomic dos/tensor, spin real-force, gen-force)
idiom 2 - extensive per-frame normalization by 1/real_natoms**norm_exp (energy, virial)
idiom 3 - global plain mean (already-reduced global dos/tensor)
Extract these as staticmethods on TaskLoss (or a deepmd/*/loss/_reduction.py shared util for the dpmodel + pt backends). Each term then collapses to masked_reduce(...) if maskf else <original expr>, which:
removes the duplicated per-frame arithmetic and flattens the nesting,
lets ener_spin reuse the same helpers as ener, cutting the largest duplication between those files,
cleans up the inv/_nf/_nloc dead-store dance.
Scope: deepmd/dpmodel/loss/{ener,ener_spin,dos,tensor}.py and their deepmd/pt/loss/ mirrors.
Not in scope
Collapsing the maskf fork entirely by defaulting the mask to all-ones (so every term always runs the per-frame idiom) is cleaner still, but changes the reduction order and therefore shifts non-masked results at the ULP level - breaking the bit-identical-for-non-mixed guarantee. That is a deliberate non-goal unless a ULP-level change to existing trainings is explicitly accepted.
This is a pure readability/maintainability refactor with no intended behavior change; the grad-accumulation-invariant tests added in #5738 are the safety net.
Follow-up to #5738.
Problem
The mixed_type padding fix in #5738 added a masked (
maskf is not None) branch to every term of the shared losses. Combined with the existinghas_*/mse/mae/use_huberconditions, this produces deeply nested, duplicated code. The energy term alone is now four levels deep:The same shape repeats for force, virial, atom_ener, atom_pref, gen_force, and again across
ener_spin. Themaskffork roughly doubles each term, and the per-frame arithmetic is copy-pasted throughout. It also forces the top-levelinv/_nf/_nloclocals whoseelse-branch dead stores CodeQL flagged.Proposed refactor (numerics-preserving)
Every masked block is one of three reduction "idioms":
ncompcomponents (force, atom_ener, atomic dos/tensor, spin real-force, gen-force)1/real_natoms**norm_exp(energy, virial)Extract these as staticmethods on
TaskLoss(or adeepmd/*/loss/_reduction.pyshared util for the dpmodel + pt backends). Each term then collapses tomasked_reduce(...) if maskf else <original expr>, which:elsebranch still evaluates the original expression, so thebit-identical for non-mixedguarantee from fix(loss): exclude mixed_type padding atoms from the training loss #5738 holds,ener_spinreuse the same helpers asener, cutting the largest duplication between those files,inv/_nf/_nlocdead-store dance.Scope:
deepmd/dpmodel/loss/{ener,ener_spin,dos,tensor}.pyand theirdeepmd/pt/loss/mirrors.Not in scope
Collapsing the
maskffork entirely by defaulting the mask to all-ones (so every term always runs the per-frame idiom) is cleaner still, but changes the reduction order and therefore shifts non-masked results at the ULP level - breaking the bit-identical-for-non-mixed guarantee. That is a deliberate non-goal unless a ULP-level change to existing trainings is explicitly accepted.This is a pure readability/maintainability refactor with no intended behavior change; the grad-accumulation-invariant tests added in #5738 are the safety net.