Skip to content

feat(models): add pipeline neighbor adaptation policy#119

Open
zubatyuk wants to merge 2 commits into
NVIDIA:mainfrom
zubatyuk:feat/pipeline-neighbor-adaptation
Open

feat(models): add pipeline neighbor adaptation policy#119
zubatyuk wants to merge 2 commits into
NVIDIA:mainfrom
zubatyuk:feat/pipeline-neighbor-adaptation

Conversation

@zubatyuk

Copy link
Copy Markdown
Collaborator

ALCHEMI Toolkit Pull Request

Description

Add configurable neighbor-list source planning to PipelineModelWrapper so composed models can share, split, or exactly group source neighbor lists through neighbor_adaptation and max_cutoff_ratio. This avoids forcing every short-cutoff sub-model to adapt from the longest source list, because filtering a long-cutoff neighbor list down to a much shorter cutoff can be more expensive than building the short-cutoff list directly.

Type of Change

  • Bug fix (non-breaking change that fixes an issue)
  • New feature (non-breaking change that adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Performance improvement
  • Documentation update
  • Refactoring (no functional changes)
  • CI/CD or infrastructure change

Related Issues

None.

Changes Made

  • Add neighbor_adaptation ("auto", "always", "never") and max_cutoff_ratio options to PipelineModelWrapper.
  • Build per-step neighbor source plans that select captured source data and filter or convert tensors only when the step plan requires it.
  • Update make_neighbor_hooks() to return either a single NeighborListHook or a composite hook for multiple planned source lists.
  • Document the pipeline neighbor-list policy in the user guide and composable model example, update the changelog, and add focused pipeline tests.

Testing

  • Unit tests pass locally (make pytest)
  • Linting passes (make lint)
  • New tests added for new functionality meets coverage expectations?

Focused pipeline neighbor tests pass locally.

Checklist

  • I have read and understand the Contributing Guidelines
  • I have updated the CHANGELOG.md
  • I have performed a self-review of my code
  • I have added docstrings to new functions/classes
  • I have updated the documentation (if applicable)

Additional Notes

None.

Tip

This repository uses Greptile, an AI code review service, to help conduct
pull request reviews. We encourage contributors to read and consider suggestions
made by Greptile, but note that human maintainers will provide the necessary
reviews for merging: Greptile's comments are not a qualitative judgement
of your code, nor is it an indication that the PR will be accepted/rejected.
We encourage the use of emoji reactions to Greptile comments, depending on
their usefulness and accuracy.

Add PipelineModelWrapper neighbor_adaptation and max_cutoff_ratio controls for sharing, splitting, or exactly grouping neighbor-list source cutoffs.

Build per-step neighbor source plans, support composite neighbor hooks for multiple planned source lists, and cover the policy in docs, examples, changelog, and focused tests.

Signed-off-by: Roman Zubatyuk <rzubatiuk@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Jun 20, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@greptile-apps

greptile-apps Bot commented Jun 20, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds configurable neighbour-list source planning to PipelineModelWrapper via two new keyword-only parameters: neighbor_adaptation ("auto" / "always" / "never") and max_cutoff_ratio (default 1.5). When the cutoff gap between sub-models is large, the pipeline now builds separate source lists instead of forcing every short-cutoff model to adapt from the longest list.

  • _build_neighbor_plan groups steps into one or more _NeighborSource objects; in "auto" mode the ratio guard avoids unnecessary long-to-short filtering, while "never" skips filtering entirely and "always" preserves the previous single-max-source behaviour.
  • make_neighbor_hooks() returns a plain NeighborListHook for single-source pipelines and a composite _PipelineNeighborListHook that runs all planned hooks and caches results on the batch for multi-source pipelines.
  • Documentation, the example script, and the changelog are updated; focused tests cover planning logic, hook factory variants, multi-source selection, and the missing-hook RuntimeError.

Important Files Changed

Filename Overview
nvalchemi/models/pipeline.py Core implementation of the new neighbour-list adaptation policy; logic and planning are correct, but save/load silently drops the two new parameters, so round-tripped pipelines will use wrong defaults.
test/models/test_pipeline.py Good coverage of planning, hook factory, and multi-source selection paths; existing adaptation tests updated with explicit neighbor_adaptation="always" to preserve previous behaviour.
CHANGELOG.md Changelog entry accurately describes the new neighbor_adaptation and max_cutoff_ratio options and their defaults.
docs/userguide/models.md Documentation updated to explain the three adaptation modes, the ratio condition, and the multi-hook registration pattern; note about make_neighbor_hooks() being required for multi-source pipelines is accurate.
examples/advanced/07_composable_model_composition.py Example comments updated to reflect that make_neighbor_hooks() may return multiple hooks; no functional code changes.

Comments Outside Diff (1)

  1. nvalchemi/models/pipeline.py, line 1277-1285 (link)

    P1 save/load silently drops neighbor_adaptation and max_cutoff_ratio

    Neither parameter is persisted in save, and load always reconstructs the pipeline using the defaults ("auto", 1.5). A pipeline saved with neighbor_adaptation="never" or max_cutoff_ratio=2.0 will load silently with a different neighbour-list plan — potentially grouping models that should have separate source lists, or filtering cutoffs that should not be filtered — with no warning or error.

    Fix: add "neighbor_adaptation": self._neighbor_adaptation and "max_cutoff_ratio": self._max_cutoff_ratio to the torch.save dict, and read them back with checkpoint.get(...) in load, passing them to cls(groups=groups, additive_keys=..., neighbor_adaptation=..., max_cutoff_ratio=...).

Reviews (2): Last reviewed commit: "fix(models): require hooks for multi-sou..." | Re-trigger Greptile

Comment thread nvalchemi/models/pipeline.py
Comment thread nvalchemi/models/pipeline.py Outdated
Raise an error when a pipeline planned for multiple neighbor-list sources is
called without captured pipeline source data instead of falling back to the
canonical batch neighbor tensors.

Type make_neighbor_hooks() as returning Hook instances and add coverage for
missing multi-source hook registration.

Signed-off-by: Roman Zubatyuk <rzubatiuk@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant