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.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions