Skip to content

refactor(loss): extract masked per-frame reduction idioms to cut nested branching #5768

Description

@wanghan-iapcm

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 existing has_* / mse/mae / use_huber conditions, this produces deeply nested, duplicated code. The energy term alone is now four levels deep:

if self.has_e:
    if self.loss_func == "mse":
        if maskf is not None:
            if not self.use_huber: ...
            else: ...
        else:
            if not self.use_huber: ...
            else: ...
    elif self.loss_func == "mae":
        if maskf is not None: ... else: ...

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,
  • preserves exact numerics - the else branch still evaluates the original expression, so the bit-identical for non-mixed guarantee from fix(loss): exclude mixed_type padding atoms from the training loss #5738 holds,
  • 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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions