Hacktoberfest 2026:メンテナが10月に向けて印を付けた、オープンで初心者向けの issue。 Hacktoberfest の issue を見る

[Feature Request] Neighbor-graph builders hardcode n_node = nloc, so the ragged node axis is never exercised

オープン
#5,938 コメント 1 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 4 日以内に返信

まだ誰も着手していません。

評価

難易度
4/5
見積もり時間
3〜5日
初心者へのやさしさ
48/100
issue の種類
機能追加
明瞭さ
おおむね明確
活発さ
静か
技術スタック
python

調査の方向性

deepmd/dpmodel/utils/neighbor_graph/builder.py と from_ijs.py にある n_node の4箇所から始め、次に nv_graph_builder.py と test_graph_ragged.py および test_edge_force_virial.py の ragged ケースを調べます。完了の条件は、フレームごとのカウントが本番用ビルダーに到達でき、矩形フォールバックが維持され、非一様バッチがカバーされていることです。データ層と dense-builder への影響は、引き続き設計作業とします。

索引モデルが issue の本文から書いたものです。

説明

enhancement
Summary

NeighborGraph already models a ragged node axis — n_node is (nf,), the node count of each frame, over a flat axis N = sum(n_node):

https://github.com/deepmodeling/deepmd-kit/blob/devel/deepmd/dpmodel/utils/neighbor_graph/graph.py#L46-L47

The consumer side is built out and tested against non-uniform values. But every builder in the tree hardcodes a uniform count, so no production path ever produces a ragged graph, and batches of unequal atom count still pay full rectangular padding.

The four producers
deepmd/dpmodel/utils/neighbor_graph/builder.py:173    n_node = xp.full((nf,), nloc, ...)
deepmd/dpmodel/utils/neighbor_graph/builder.py:241    n_node = xp.full((nf,), nloc, ...)
deepmd/dpmodel/utils/neighbor_graph/builder.py:454    n_node = xp.full((nf,), nloc, ...)
deepmd/dpmodel/utils/neighbor_graph/from_ijs.py:123   n_node = xp.full((nf,), nloc, ...)

ase_builder, vesin_graph_builder and nv_graph_builder do not compute n_node themselves; each derives a scalar nloc from a rectangular coord tensor and delegates to neighbor_graph_from_ijs. The nv builder additionally assumes the uniform stride when it forms its batch pointer:

batch_ptr = torch.arange(nf + 1, dtype=torch.int32, device=device) * nloc

Every non-uniform n_node in the repository is in a test (test_graph_ragged.py uses [3, 2], test_edge_force_virial.py uses [3, 0, 5], and so on).

Why it matters

A rectangular batch of k frames whose widest holds W atoms occupies k * W slots on the node axis, not sum_f N_f. The padded slots are real tensor rows: they flow through the type embedding, the descriptor's node-level operators and the fitting net, and are only masked at the end, so compute and activation memory scale with k * W.

batch contents k * W slots real atoms padding
16 frames of 20 atoms + 6 of 24 + 11 of 30 990 794 19.8%
500 frames of 1 atom + 1 frame of 1000 atoms 501,000 1,500 99.7%

Existing upstream mechanisms that produce such batches: the mixed_type npy format (real_atom_types.npy, where -1 marks a virtual atom, documented in doc/data/system.md as the way to merge frames of unequal Natoms) and batch_size: "mixed:N", whose DeepmdDataSystem._merge_batch_data pads type up to max_natoms with -1.

What is already in place

The gap is narrow because the whole consumer half is done:

  • frame_id_from_n_node, node_ownership_mask, node_validity_mask in graph.py derive the node→frame map and the offsets from n_node.
  • segment_sum in neighbor_graph/segment.py and fit_output_to_model_output_graph in dpmodel/model/edge_transform_output.py reduce per frame without a rectangular axis.
  • forward_common_atomic_graph in base_atomic_model.py and call_lower_graph in make_model.py take the flat node axis directly.
  • build_edge_csr takes a scalar node count and is layout-neutral.
  • The .pt2 export declares nframes and n_node_total as independent symbolic dims, with no N == nf * nloc relation, so the artifact would accept a ragged graph today.
  • The C++ api_cc graph path is single-frame and already sets n_node from nall vs nloc rather than assuming a stride.
Proposal

Thread the real per-frame atom counts through to n_node instead of synthesizing a uniform vector:

  1. Let the builders accept an optional per-frame node count and fall back to full((nf,), nloc) when absent, so existing rectangular callers are unaffected.
  2. Replace the uniform-stride index arithmetic with the offsets already derivable from n_node (cumulative_sum(n_node) - n_node), which is what node_ownership_mask does. The affected sites are few: builder.py:164-171, 230-235, 447-449; from_ijs.py:104-106; nv_graph_builder.py:106-108,171.
  3. Have the data layer supply the counts. They are typically already known — a loader that pads to a batch-wide maximum necessarily computed the per-frame counts first.

Per-backend cost is uneven:

  • nv: nearly free. nvalchemiops already batches through explicit offsets; batch_ptr becomes a cumulative sum of real counts.
  • vesin / ase: easy. Both already loop per frame and concatenate; the uniform nloc is used only to compute the offset.
  • dense: the awkward one, and the default on non-CUDA devices. build_neighbor_graph materializes an (nf, nloc, nall) all-pairs tensor through the rectangular extend_coord_with_ghosts. It could keep padding internally (correct, but forfeits the saving on that path) or block-diagonalize the search.
Scope

Filing this to record the gap rather than to claim it is urgent. Steps 1 and 2 are self-contained and would make the ragged axis reachable; step 3 and the dense builder are where the real work is.

Related: #5861 applies the same direction to the edge axis.

主要言語
Python
スター
2.1k
フォーク
662
平均マージ
5日 13時間
マージ済み PR(30日)
11

環境構築

Codespaces で開く

このプロジェクトの開発コンテナを、あなたの GitHub アカウントでブラウザ上に起動します。

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

deepmodeling/deepmd-kit のほかの issue

deepmodeling/deepmd-kit の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。