[Feature Request] Neighbor-graph builders hardcode n_node = nloc, so the ragged node axis is never exercised
メンテナーはふだん 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 の本文から書いたものです。
説明
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):
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_maskingraph.pyderive the node→frame map and the offsets fromn_node.segment_suminneighbor_graph/segment.pyandfit_output_to_model_output_graphindpmodel/model/edge_transform_output.pyreduce per frame without a rectangular axis.forward_common_atomic_graphinbase_atomic_model.pyandcall_lower_graphinmake_model.pytake the flat node axis directly.build_edge_csrtakes a scalar node count and is layout-neutral.- The
.pt2export declaresnframesandn_node_totalas independent symbolic dims, with noN == nf * nlocrelation, so the artifact would accept a ragged graph today. - The C++
api_ccgraph path is single-frame and already setsn_nodefromnallvsnlocrather than assuming a stride.
Proposal
Thread the real per-frame atom counts through to n_node instead of synthesizing a uniform vector:
- 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. - Replace the uniform-stride index arithmetic with the offsets already derivable from
n_node(cumulative_sum(n_node) - n_node), which is whatnode_ownership_maskdoes. 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. - 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_ptrbecomes a cumulative sum of real counts. - vesin / ase: easy. Both already loop per frame and concatenate; the uniform
nlocis used only to compute the offset. - dense: the awkward one, and the default on non-CUDA devices.
build_neighbor_graphmaterializes an(nf, nloc, nall)all-pairs tensor through the rectangularextend_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
環境構築
このプロジェクトの開発コンテナを、あなたの GitHub アカウントでブラウザ上に起動します。
- Dockerfile・Docker Compose ファイルなし
- プルリクエストのテンプレートなし
- コントリビューションガイドを読む
はじめの一歩
- issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
- 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
- リポジトリをフォークし、ブランチを切って変更します。
- issue 番号を参照したプルリクエストを送ります。
deepmodeling/deepmd-kit のほかの issue
-
難易度 2/5 1〜3時間 初心者へのやさしさ 73/100
deepmodeling/deepmd-kit#6086 ·
メンテナーはふだん 4 日以内に返信
-
Fine-tune descriptor mismatch warning should ignore seed対応中かも @SiyuanChenToDo が 10 日前に担当しました。 オープン
難易度 1/5 1時間未満 初心者へのやさしさ 88/100
deepmodeling/deepmd-kit#6038 · コメント 2 件 ·
メンテナーはふだん 4 日以内に返信
-
Docs enhancement
難易度 2/5 1〜3時間 初心者へのやさしさ 68/100
deepmodeling/deepmd-kit#5766 · コメント 1 件 ·
メンテナーはふだん 4 日以内に返信
-
[Code scan] Freeze Paddle DPA3 parameters with stop_gradient対応中かも @njzjz が 107 日前に担当しました。 オープンbug
難易度 2/5 1〜3時間 初心者へのやさしさ 78/100
deepmodeling/deepmd-kit#5686 · コメント 1 件 ·
メンテナーはふだん 4 日以内に返信
-
branch: master lang: py todo comment type: todo
難易度 2/5 1〜3時間 初心者へのやさしさ 72/100
deepmodeling/deepmd-kit#5480 ·
メンテナーはふだん 4 日以内に返信
deepmodeling/deepmd-kit の issue をすべて見る
似ている issue
-
難易度 2/5 1〜3時間 初心者へのやさしさ 86/100
-
Validate finite output clamping limits in StepPredictor.prepare_clamping_params to prevent silent NaN/Inf propagation対応中かも @muskan-creates352 が今日担当しました。 オープン
難易度 2/5 1〜3時間 初心者へのやさしさ 84/100
mllam/neural-lam#848 ·
メンテナーはふだん 2 日以内に返信
-
infra
難易度 1/5 1時間未満 初心者へのやさしさ 88/100
fab-geocommuns/RNB-coeur#1166 ·
メンテナーはふだん 1 日以内に返信
-
難易度 2/5 1〜3時間 初心者へのやさしさ 78/100
deepspeedai/DeepSpeed#8951 ·
メンテナーはふだん 2 日以内に返信
-
難易度 2/5 1〜3時間 初心者へのやさしさ 76/100
aws/context-ontology-accelerator#343 ·
メンテナーはふだん 1 日以内に返信