deepmodeling / deepmodeling/deepmd-kit
[Feature Request] Neighbor-graph builders hardcode n_node = nloc, so the ragged node axis is never exercised
- Dominant language
- Python
- Stars
- 2k
- Forks
- 649
- Avg merge
- 6d 18h
- Merged PRs (30d)
- 15
Description
### 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:
```python
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.
Contributor guide
Assessment
This issue has not been assessed yet.