mad

[xke123502]

Summary by CodeRabbit

• New Features

  • Added MADGap regularization: new loss type “ener_mad” with tunable coefficient, plus DPA3 descriptor options (enable_mad, mad_cutoff_ratio). Training now supports this loss.
  • Enhanced RepFlows descriptor with dynamic neighbor indexing, improved angle features, and parallel execution support.

• Bug Fixes

  • Training output printing is resilient to missing validation metrics.

• Documentation

  • Extensive Chinese docstrings and comments across descriptors, layers, and MLP; no behavior changes.

• Chores

  • Updated water/dpa3 example: reduced training steps to 200.

[coderabbitai]

Caution

Review failed

The pull request is closed.

📝 Walkthrough

Walkthrough

Adds MAD-related regularization to DPA3 descriptor and energy loss (new EnergyStdLossMAD), wires it into training and config parsing, and significantly refactors RepFlows, network MLPs, and graph-index utilities. Several modules receive extensive documentation/comments. An example reduces training steps.

Changes

Cohort / File(s)	Summary
Loss export surface deepmd/pt/loss/__init__.py	Exposes EnergyStdLossMAD by importing from ener and adding to __all__.
Energy loss (MAD) deepmd/pt/loss/ener.py	Adds EnergyStdLossMAD subclass with mad_reg_coeff, applies MAD regularization using descriptor’s last_mad_gap, extends serialize/deserialize (includes mad_reg_coeff); note duplicate class insertion in patch.
Descriptor DPA3 (MAD support) deepmd/pt/model/descriptor/dpa3.py	Adds MAD computation and config (enable_mad, mad_cutoff_ratio), stores last_mad_gap; updates forward to compute/store MAD; serialize includes new fields; adds _compute_mad.
Env matrix comments deepmd/pt/model/descriptor/env_mat.py	Adds inline comments; no logic changes.
RepFlow layer docs deepmd/pt/model/descriptor/repflow_layer.py	Adds extensive Chinese docstrings/comments; no functional or API changes.
RepFlows block (functional overhaul) deepmd/pt/model/descriptor/repflows.py	Major forward-path refactor: environment matrix construction for edges/angles, masking, dynamic neighbor selection via get_graph_index, local/global mapping, optional parallel communication (border_op), iterative layer application, returns rotation matrix; adds Union import; defines fallback torch.ops.deepmd.border_op.
Model wiring comments deepmd/pt/model/model/__init__.py	Adds comments in _get_standard_model_components; no behavior change.
MLP/network stack deepmd/pt/model/network/mlp.py	Expands MLPLayer with timestep gating, precision/dtype handling, multiple inits, residual modes, new helpers (dim_in, dim_out, check_type_consistency), revamped serialization; ensures ModuleList registration; adds empty_t and NetworkCollection.NETWORK_TYPE_MAP.
Network utils (aggregation/indexing) deepmd/pt/model/network/utils.py	Reworks aggregate (safer counts, padding, avg option) and get_graph_index (vectorized edge/angle index construction, local/global mapping via nall, adds type annotation).
Training integration deepmd/pt/train/training.py	Imports EnergyStdLossMAD; get_loss supports loss_type == "ener_mad"; single-task printing handles missing validation keys without KeyError.
Arg parsing/config deepmd/utils/argcheck.py	Adds DPA3 args enable_mad, mad_cutoff_ratio; changes DPA3 default activation_function to "silu"; adds mad_reg_coeff to energy loss; registers new loss variant ener_mad with MAD-related prefs.
Example config examples/water/dpa3/input_torch.json	Reduces training.numb_steps from 1,000,000 to 200.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant User
  participant Trainer as training.get_loss
  participant Loss as EnergyStdLossMAD
  participant Model as DPA3/RepFlows Model
  participant Desc as DPA3 Descriptor

  User->>Trainer: loss_type="ener_mad", loss_params
  Trainer->>Loss: construct(**loss_params)
  Note right of Loss: mad_reg_coeff stored
  loop Training Step
    Trainer->>Model: forward(input)
    Model->>Desc: forward(...)
    alt enable_mad
      Desc->>Desc: compute MAD (store last_mad_gap)
    else disable
      Desc->>Desc: last_mad_gap=None
    end
    Trainer->>Loss: forward(outputs, model, labels, natoms, lr)
    Loss->>Loss: base energy/force/virial losses
    alt mad_reg_coeff > 0 and last_mad_gap available
      Loss->>Desc: read last_mad_gap
      Loss->>Loss: mad_reg_loss = coeff * |MAD - 1|
      Loss->>Trainer: total_loss += mad_reg_loss
    else no MAD reg
      Loss-->>Trainer: total_loss (base)
    end
  end

Loading

sequenceDiagram
  autonumber
  participant Model as DescrptBlockRepflows.forward
  participant Env as prod_env_mat
  participant Utils as get_graph_index
  participant Layer as RepFlowLayer[xN]
  participant Border as torch.ops.deepmd.border_op

  Model->>Env: build edge env (dmatrix, diff, sw)
  Model->>Env: build angle env (a_dist_mask, a_diff, a_sw)
  Model->>Utils: get_graph_index(nlist, masks, a_masks, nall, use_loc_mapping)
  alt parallel_mode
    Model->>Border: exchange/gather features
  end
  loop layers
    Model->>Layer: forward(node_ebd_ext, edge_ebd, h2, angle_ebd, indices, masks, sw)
    Layer-->>Model: updated embeddings
  end
  Model->>Layer: _cal_hg / _cal_hg_dynamic
  Model-->>Caller: node_ebd, edge_ebd, h2, rot_mat, sw

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Possibly related PRs

• feat: dpmodel energy loss & consistent tests #4531 — Modifies energy loss classes and serialization, overlapping with adding EnergyStdLossMAD and its serialize/deserialize.
• pd: support dpa3 with paddle backend #4701 — Adds DPA3-related support and loss exports in another backend, aligning with DPA3 MAD and loss integration here.
• Chore: refactor get standard model #4205 — Touches _get_standard_model_components, which this PR comments on; same function surface.

Suggested labels

Python, Docs, Examples

Suggested reviewers

• iProzd
• njzjz
• wanghan-iapcm

✨ Finishing touches

• 📝 Generate Docstrings

🧪 Generate unit tests

• Create PR with unit tests
• Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 15c2b7d and a9bb0be.

📒 Files selected for processing (12)

• deepmd/pt/loss/__init__.py (2 hunks)
• deepmd/pt/loss/ener.py (1 hunks)
• deepmd/pt/model/descriptor/dpa3.py (7 hunks)
• deepmd/pt/model/descriptor/env_mat.py (2 hunks)
• deepmd/pt/model/descriptor/repflow_layer.py (26 hunks)
• deepmd/pt/model/descriptor/repflows.py (11 hunks)
• deepmd/pt/model/model/__init__.py (2 hunks)
• deepmd/pt/model/network/mlp.py (4 hunks)
• deepmd/pt/model/network/utils.py (3 hunks)
• deepmd/pt/train/training.py (3 hunks)
• deepmd/utils/argcheck.py (4 hunks)
• examples/water/dpa3/input_torch.json (1 hunks)

Tip

👮 Agentic pre-merge checks are now available in preview!

Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.

• Built-in checks – Quickly apply ready-made checks to enforce title conventions, require pull request descriptions that follow templates, validate linked issues for compliance, and more.
• Custom agentic checks – Define your own rules using CodeRabbit’s advanced agentic capabilities to enforce organization-specific policies and workflows. For example, you can instruct CodeRabbit’s agent to verify that API documentation is updated whenever API schema files are modified in a PR. Note: Upto 5 custom checks are currently allowed during the preview period. Pricing for this feature will be announced in a few weeks.

Please see the documentation for more information.

Example:

reviews:
  pre_merge_checks:
    custom_checks:
      - name: "Undocumented Breaking Changes"
        mode: "warning"
        instructions: |
          Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).

Please share your feedback with us on this Discord post.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}